Distributionally Robust Optimization via Generative Ambiguity Modeling

This paper studies Distributionally Robust Optimization (DRO), a fundamental framework for enhancing the robustness and generalization of statistical learning and optimization. An effective ambiguity set for DRO must involve distributions that remain consistent to the nominal distribution while being diverse enough to account for a variety of potential scenarios. Moreover, it should lead to tractable DRO solutions. To this end, we propose generative model-based ambiguity sets that capture various adversarial distributions beyond the nominal support space while maintaining consistency with the nominal distribution. Building on this generative ambiguity modeling, we propose DRO with Generative Ambiguity Set (GAS-DRO), a tractable DRO algorithm that solves the inner maximization over the parameterized generative model space. We formally establish the stationary convergence performance of GAS-DRO. We implement GAS-DRO with a diffusion model and empirically demonstrate its superior Out-of-Distribution (OOD) generalization performance in ML tasks.

💡 Research Summary

**
The paper tackles a central challenge in Distributionally Robust Optimization (DRO): how to design an ambiguity set that is both expressive enough to capture worst‑case distributions (including support shifts) and tractable enough to allow efficient optimization. Traditional ambiguity sets based on ϕ‑divergences (e.g., KL) enforce absolute continuity with the nominal distribution, which limits robustness when the test distribution’s support differs from the training support. Wasserstein‑based sets avoid this restriction but lead to infinite‑dimensional inner maximization problems that are computationally demanding and often require conservative relaxations.

To reconcile these conflicting requirements, the authors propose Generative Ambiguity Sets (GAS), which are defined directly in the parameter space of likelihood‑based generative models such as diffusion models and Variational Autoencoders (VAEs). By constraining the parameters θ of a generative model (e.g., via a reverse‑KL bound), the set contains distributions that (i) remain statistically consistent with the nominal data distribution, and (ii) can generate samples beyond the original support, thereby exploring a richer family of adversarial scenarios. This formulation reduces the inner maximization to a finite‑dimensional, differentiable problem.

Building on GAS, the authors introduce GAS‑DRO (Algorithm 1). The algorithm alternates between:

1. Inner maximization: Adjust θ to maximize the expected loss under the current decision variable w. This is solved using dual learning and policy‑optimization techniques, yielding gradient updates in the θ‑space while respecting the reverse‑KL constraint.
2. Outer minimization: Update the decision variable w (e.g., model parameters) using standard stochastic gradient methods against the worst‑case distribution generated in step 1.

Theoretical contributions include:

• Lemma 1: Shows that the reverse‑KL constraint can be enforced during θ‑updates.
• Theorem 1: Proves that the inner maximizer converges to the optimal oracle and asymptotically generates distributions consistent with the nominal distribution.
• Theorem 2: Establishes stationary‑point convergence of the full alternating scheme.

Empirically, GAS‑DRO is instantiated with a diffusion model and evaluated on time‑series forecasting (e.g., electricity demand) and image classification (CIFAR‑10, a subset of ImageNet). Compared against state‑of‑the‑art DRO baselines—KL‑DRO, Wasserstein‑DRO, Sinkhorn‑DRO, and recent generative‑model‑based approaches (e.g., DRA‑GEN)—GAS‑DRO achieves:

• Improved OOD performance: 3–5 % higher accuracy / lower loss on shifted test distributions, especially when support changes are pronounced.
• Reduced computational cost: Approximately half the runtime of Wasserstein‑DRO, thanks to the finite‑dimensional parameterization.
• Realistic worst‑case distributions: Visual inspection shows generated samples retain the structure of the original data while incorporating adversarial perturbations (label flips, added noise), indicating meaningful exploration of risky scenarios.

The paper acknowledges limitations: reliance on the quality of the underlying generative model (risk of over‑fitting), sensitivity of the inner maximization to initialization, and the computational expense of training diffusion models. Future directions include extending GAS to other generative families (GANs, normalizing flows), incorporating meta‑learning to adapt the ambiguity‑set radius automatically, and exploring regularization strategies to mitigate over‑fitting of the generative model.

In summary, the work introduces a novel, generative‑model‑centric ambiguity set that simultaneously offers expressive power and tractability, and demonstrates both theoretical guarantees and practical gains in OOD robustness for DRO problems.