Adaptive Scenario Generation for KL-DRO Around a Learned Density Model

Adaptive Scenario Generation for KL-DRO Around a Learned Density Model§

Problem Setting§

Two-stage stochastic programs optimize decisions under uncertainty, but the true distribution is unknown. Distributionally robust optimization (DRO) hedges against distributional ambiguity by considering a set of plausible distributions. Common choices are:

• Divergence ball around empirical distribution: only reweightings of observed data.
• Transport ball: allows off-support distributions, including impossible realizations.

This paper proposes a KL-divergence ball centered at a learned density model $\hat{p}$ (the learned nominal), not at the empirical distribution $\hat{p}_\text{emp}$. The ambiguity set is: $$\mathcal{P} = \left\{ q : D_{\mathrm{KL}}(q \| \hat{p}) \leq \rho \right\}.$$ This lets the adversary up-weight costly unobserved realizations in proportion to the model's assigned probability, hedging only against plausible rare events.

Challenges§

1. Tractability: The inner maximization over $q$ is infinite-dimensional. The worst-case expectation for a fixed decision $x$ is:

$$\sup_{q \in \mathcal{P}} \mathbb{E}_q[\ell(x,\xi)] = \min_{\lambda \geq 0} \left\{ \lambda \log \mathbb{E}_{\hat{p}}\left[ e^{\ell(x,\xi)/\lambda} \right] + \lambda \rho \right\},$$ which still involves an expectation over $\hat{p}$.

1. Scenario generation: The worst-case often concentrates on rare events. A naive Monte Carlo sample from $\hat{p}$ may miss them. Need to build a scenario set adaptively, steering draws toward the current worst case.

Methodology§

Alternating optimization:

• Given a scenario set $\{\xi_i\}_{i=1}^N$ with weights $w_i$, solve the finite convex problem:

$$\min_{x} \max_{q \in \Delta_N} \sum_{i=1}^N q_i \ell(x,\xi_i) \quad \text{s.t.} \quad \sum_i q_i \log(q_i / w_i) \leq \rho.$$

• Then, given current decision $x$, generate new scenarios from the learned nominal $\hat{p}$ but steered toward high-loss regions using importance sampling. Draw from a proposal $g(\xi) \propto \hat{p}(\xi) \exp(\ell(x,\xi)/\lambda)$ for some temperature $\lambda$, and attach importance weights $w_i \propto \hat{p}(\xi_i)/g(\xi_i)$ to correct back to the model.
• Iterate until convergence.

The learned nominal $\hat{p}$ can be any density model (e.g., Gaussian, normalizing flow, VAE). The key is that it captures structure and can extrapolate beyond the observed record.

Experimental Results§

On a multi-item newsvendor problem and a capacity expansion problem using realistic grid data, centering the ball at the learned nominal (instead of the empirical distribution) lowers cost under distribution shift while maintaining comparable in-distribution cost. The advantage is most pronounced when the record is scarce and the model class captures structure that enables extrapolation. The benefit diminishes as the record grows or the shift departs from the learned structure.

Implementation Sketch§

Below is a simplified Python/PyTorch snippet illustrating the adaptive scenario generation loop:

import torch
import torch.distributions as dist

def adaptive_kl_dro(loss_fn, p_hat, rho, n_scen, n_iter=100):
    # loss_fn(x, xi) -> tensor of losses
    # p_hat: learned density model (e.g., torch.distributions.Distribution)
    x = torch.zeros(decision_dim, requires_grad=True)
    scenarios = p_hat.sample((n_scen,))
    weights = torch.ones(n_scen) / n_scen
    optimizer = torch.optim.Adam([x], lr=1e-3)
    for it in range(n_iter):
        # Solve inner finite problem (convex in q) via bisection on lambda
        losses = loss_fn(x, scenarios)
        # Find lambda such that KL constraint tight (simplified)
        lam = find_lambda(losses, weights, rho)
        q = weights * torch.exp(losses / lam)
        q = q / q.sum()
        # Outer minimization over x
        optimizer.zero_grad()
        loss = (q.detach() * losses).sum()
        loss.backward()
        optimizer.step()
        # Resample scenarios with importance sampling steered by current losses
        if it % 10 == 0:
            with torch.no_grad():
                # Proposal: p_hat * exp(loss/lam) approximated via rejection sampling
                proposals = p_hat.sample((n_scen * 10,))
                prop_losses = loss_fn(x, proposals)
                accept_prob = torch.exp(prop_losses / lam)
                accept_prob = accept_prob / accept_prob.max()
                accepted = proposals[torch.rand_like(accept_prob) < accept_prob][:n_scen]
                if len(accepted) < n_scen:
                    accepted = torch.cat([accepted, p_hat.sample((n_scen - len(accepted),))])
                scenarios = accepted
                # Recompute importance weights to correct to p_hat
                w_new = p_hat.log_prob(scenarios).exp() / (p_hat.log_prob(scenarios).exp() * torch.exp(loss_fn(x, scenarios)/lam))
                weights = w_new / w_new.sum()
    return x.detach()

Key insight: The learned nominal $\hat{p}$ provides a principled way to extend the ambiguity set beyond the empirical support, while the adaptive importance sampling ensures that rare but costly scenarios are represented in the finite sample, yielding a practical algorithm for KL-DRO.