Distributionally Robust Optimization via Generative Ambiguity Modeling
Authors: JIAQI WEN
Jianyi Yang
Venue: ICLR 2026 Poster (2026)
Keywords: Distributionally Robust Optimization
Generative Models
OOD Generalization
Primary area: optimization

Abstract: 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.

Published as a conference paper at ICLR 2026
DISTRIBUTIONALLY ROBUST OPTIMIZATION VIA GENERATIVE AMBIGUITY MODELING
Jiaqi Wen and Jianyi Yang
Department of Computer Science
University of Houston
ABSTRACT
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 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 sion model and empirically demonstrate its superior generalization performance in ML tasks. []
We implement GAS- GAS-DRO with a diffu- (OOD)
INTRODUCTION
Distributionally Robust Optimization (DRO) is fundamental framework for enhancing the robust- ness of statistical learning and optimization problems, particularly under out-of-Distributicn (OOD) scenarios Liu et al. (2021); Arjovssy (2020); Ren & Majumdarr (2022). DRO formulates minimax optimization problem, where the inner maximization identifies the worst-case distribution within an ambiguity set, and the outer minimization optimizes the decision variable against this worst-case scenario Chen et al. (2020): [iu et al. (002b)); Chen & Paschalidis (2018); Kuhn et al. (2019); Smirnova et al. (2019). Unlike non-probabilistic robust optimization, DRO leverages probabilistic ambiguity modeling to enable improved generalization performance. This property has made DRO increasingly important in statistical learning and optimization tasks for addressing distribution shifts, noisy data, and adversarial conditions Quiñonero- Canrro--Cnnllaaa et al. (2022); Muratore et al. (2022 Ren
Majumdar (2022); Li et al. 2019); Veit al. (2017).
The performance of DRO algorithms significantly depends on the design of the ambiguity set, which must contain diverse distributional variations around the nominal distribution. common approach to model the ambiguity set using dddivergences, such as the Kullback-Leibler (KL) divergence Hu Hong (013); Husain et al. (2023); Kirschner et al. (2020); Kocuk (2020); Liu et al. (2022b). Al- though such Φ d-divergeece- based DRO formulations can sometimes yield closed-form solutions Hu & Hong 2013); Husain et al. (2023), they require any distribution P' in the ambiguity set to be absolutely continoous with respect to the nominal distribution Po (denoted P < Po). meaning that for any measurable set A, if (AA) = 0, then P(AA) = 0. This implicit constraint limits robustness
of DRO in scenarios with support shifts Staib & Jegelka (2019) Lu et al. (2024): Gao & Kleywegt
2023a)
While the ambiguity defined by Wasserstein distance allows for support shifts, solving Wasserstein-DRO over the infinite probability space is challenging and computationally inefficient. Some approaches
MoppraachesMMhhaeiinn Mohajerin Esfahani & Kuhn (2018) Chen et al. 2020); Gao & Kleywegt (202bb)
"Correspondence to Jianyi Yang hjyang66@@hccccll Source Code Link: http://gthhaaa.oomm.....aaaaaaa

---

Published as a conference paper at ICLR 2026
reformulate Wasserstein as finite-dimensioaal optimization problem based the convex assumptions which typically do not hold in deep learning. Other methods approximate Wasserstein- DRO via adversarial optimization (Staib & Jegelka (2017): Gao & Kleywegt (023b); however, these relaxations are often overly conservative, which restricts their ability to exploit probabilistic model- ing for effectively balancing average-case performance with worst-case robustness.
Recent advances have sought to address the challenges of ambiguity modeling in DRO. Some works Wang et al. (2025a); Yang al. (2025), Blanchet al. (2023) study DRO using ambiguity sets based on the Sinkhorn distance, which preserve computational tractability and enhance the expressiveness of adversarial distributions. Generative models have also been utilized in the traditional ambigu- ity modeling frameworks. Among them, Ren & Majumdar (2022) construct a Wasserstein -ased ambiguity set in the latent space of generative models. Michel et al. (2021 2022) exploit param- eterized generative models to represent the adversarial distributions or the likelihood ratios within the KL.-divergence- based ambiguity modeling framework and derive tractable solutions with addi- tional approximations. Nevertheless, the traditional ambiguity modeling frameworks still introduce tractability issues or restricts the support space of adversarial distributions, limiting the utilization of the distributional representation capacity of generative models.
To the best of our knowledge, this is the first work to model the ambiguity sets in the parameter- ized space of generative models P such as diffusion models and Variational Autoencoders (VAEs), which offers several advantages: (1) Generative models have a strong capability to represent the un- derlying data distribution, ensuring that the distributions in the ambiguity set remain consistent with the nominal distribution. (2) Generative models are capable of producing diverse samples beyond the nominal support space, thereby enabling the discovery of distributions with worst-case opti- mization performance. (3) Generative models provide a finite, parameterized optimization space, avoiding the need to solve problems over an infinite probability space.
Our main contributions are summarized below: (1) Based on the probability modeling of gener- ative models, we introduce novel Generative Ambiguity Set (GAS) for DRO which involves di- verse distributions while preserving consistency with the nominal distribution. While GAS is intro- duced based on the formulations of diffusion models and VAEs, it can be flexibly adapted to other likelihood-based generative models. (2) We design
(Algorithm []. which gives direct solution to DRO with the proposed Generative Ambiguity Sets. We solve the inner maximization in GAS-DRO based on dual learning and policy optimization techniques, enabling tractable iterative optimization within a finite, parameterized space. (3) We prove in Theorem that the inner max- imizer of GAS-DRO converges to the optimal maximization oracle and, asymptotically, generates distributions consistent with the nominal distribution. Building on this inner-maximization error bound, Theorem 2 further establishes the stationary convergence of GAS-DRO. (4) We implement GAS-DRO on series prediction and image classification tasks and demonstrate its superiority by comparing it against state-of-the-art DRO baselines based on generative models.
2 RELATED WORK
Distributionally Robust Optimization. The performance of DRO relies largely on the design of ambiguity sets. Tractable DRO solutions can be obtained based on the ambiguity sets defined by ø-divergence Husain et al. (2023) Kirschner et al. (2020): (uu & Hong (2033) (2cuk((022)) How- ever, the definition of oi-iiergencce requires any distribution P in the ambiguity set to be absolutely continuous with respect to the nominal distribution, which can cause support mismatch problems Staib & Jegelka| (2019). By contrast, Wasserstein-distane -istaace-tateeeddd ambiguity set Kuhn et al. (2019); Mohajerin Esfahani & Kuhn (2018): Gao & Kleywegt (2023b) has no restrictions on the support of the distribution, but it introduces much computational difficulty. Wasserstein based DRO is solved by either conservative ealaattoos/Sinha et al. (2017), Staib Staib & Jegelka (201)) reformulations based on some assumptions on the objective Kuhn et al. (0019); Mohajerin Esfahani & Kuhn 2018): Chen
et al. 2020); Gao & Kleywegt (2023b).
"This claim means designing new ambiguity sets that directly constrain the parameters of generative models without replying on traditional distributional discrepancy measures.
We design concrete algorithms with likelihood based generative models that satisfy the constraints on inclusive KI. divergence in Lemma [ but the proposed framework can employ other generative models which potentially provide other guarantees consistency with the nominal distribution.
2

---

Published as a conference paper at ICLR 2026
A line of recent studies focuses addressing the challenges in ambiguity modeling for DRO Xu
(et al. (2024); Zhu & Xie (2024); Ren & Majumdari (0022; Liu et al. (2022a): Yang et al. 2025
Among them, Liu et al. (2022a); Wang et al. (2025c) inject prior knowledge into the design of ambi-
guity sets to mitigate excessive conservativeness, Ma et al. 2024) develop an end-to-end framework
in which MI. model is trained to predict the ambiguity set for downstream DRO tasks. More- over, recent studies Wang et al. (2025a): Yang et al. ) [lanchet et al. (2023) investigate DRO with bookhorn-distance-bass
ambiguity set which includes continuous adversarial distributions and preserves tractability. Different from these methods, we utilize the strong distribution learning ca- pability of generative models to build the ambiguity set, enabling the discovery of worst-case and realistic distributions. Meanwhile, the proposed algorithm GAS-DRO converts DRO into a finite tractable problem in the parameterized space of generative models.
Generative Models for Robust Learning. While generative models are typically used for gen- erative tasks /Zhang al. (2024): Li et al. (2025); Chen et al. (2024), they have also been used for (robust) discriminative
classification tasks Li et al. (2023); Chen et al. (2023); Tong et al. (2025). A line of works use generative models to synthesize adversarial samples for robust training Dai et al. (2024b a); Xie et al. (2024): Du et al. (2022). The target of these works is to generate adversarial attacking examples which is fundamentally different from the worst-case distribution generation which is studied in this paper and aims to improve OOD generalization. Recently, Waang] framework, where a diffusion model is trained to generate data for downstream statistical optimization with focus on the conditional value-at-risk (CVaR) objective.
eet al. (2025b) introduced a Senneratethen-Optinnee
Prior work has utilized generative models to represent adversarial distributions within traditional ambiguity modeling frameworks. For example, DRAGEN Ren & Majumdar (2022) models adver- sarial distributions by a generative model. However, this approach remains within the Wasserstein- based DRO framework and continues to suffer from intractability and relaxation-relate issues. In addition, Michel et al. (202)) represent adversarial distributions using generative models within a traditional -Il-divergence-baeee ambiguity set and derive an approximated tractable objective. Sim- ilarly, Michel et al. (2022) employ generative models to parameterize the likelihood ratio between the adversarial and nominal probability densities. However, both Michel et al. (2021) and Michel jet al. (2022) remain within the KL-divergence-basee DRO framework, which requires the adver- sarial distributions to share the same support the nominal distribution. This constraint restricts the expressive power of generative models in constructing diverse adversarial distributions and ulti- mately limits the robustness performance. Different from these methods, our generative ambiguity modeling overcomes the limitations of traditional ambiguity sets: It provably constrains the reverse KL divergence that does not restrict the support space of adversarial distributions and enables flexi- bly tractable solutions through policy-optimization techniques.
Our work is also related studies on generative models for worst-case environment generation. In this line of research, Berdica et al. (2024) select adversarial world models from a set of world mod- els trained on a single dataset. addition, Hill (2025) conduct an experimental study on learning adversarial world models aimed at defeating defender agents in multi-agent reinforcement learn- ing. Moreover, [heng et al. (2025) learn a transport map for out-of-sample generation within the Wasserstein space and establish convergence guarantees.
3 AMBIGUITY MODELING IN DRO
Distributionally Robust Optimization. DRO optimizes for the worst-case performance given an ambiguity set constructed on a nominal distribution Po- An empirical dataset So can be drawn from the nominal distribution. Consider an objective function ((w,2) with the decision variable w E W and the random parameter E X E Rd. Given the nominal distribution Pz( of the random parameter x, DRO solves the following minimax optimization problem.
min max En-p(f(w,r) wEW Po.()).e)
where B(Po, is the ambiguity set containing possible testing distributions and is typically modeled as a distribution ball ((P = (P D(P,P)) < e} given a distribution discrepancy measure D and adversary budget
3

---

Published as a conference paper at ICLR 2026
Challenges of Ambiguity Modeling in DRO. The choice of the ambiguity set in DRO has signifi- cant impact on both generalization performance and solution tractability. An effective ambiguity set should satisfy three key properties. First, it should capture broad range of distributions with dif- ferent support spaces, enabling the identification of worst-case scenarios. Second, the distributions within the ambiguity should remain consistent with the nominal distribution since inconsistent distributions can result in overly-conservatte DRO solutions with poor average-case performance. Finally, the ambiguity set should enable a tractable DRO solution despite the complexity of the infinite probability space in DRO.
The ambiguity sets based on classic distribution discrepancy measures do not satisfy one or more of the above properties. Ambiguity sets based on o-divergences Hu & Hong (2013); Husain et al. (2023); Kocuk (2020); Liu et al. (202b)) requires any distribution P in the ambiguity set to share the same support X' as the nominal distribution Po. i.e. P is absolutely continuous with Po (P & Po)- This implicit constraint will restrict the search of worst-case distributions and limit the robustness under testing cases with support shift. In addition, Wasserstein- based ambiguity set Mohajerin Es- Tahani & Kuhn (2018) 8): Chen et al. (2020); Gao & Kleywegt (2023b); Staib Jegelka (2017): Gao & Kleywegt 2023b) (0033b introduces too much complexity in solving the inner maximization over infinite probability space. The relaxation of the Wasserstein- ambiguity set can include a lot of invalid distributions that are not consistent with the nominal distribution Po. leading to overly-conservative
DRO solutions.
We observe that the challenge of ambiguity modeling stems from fundamental tension between expressiveness and tractability. When probability distributions are restricted to the same support as the nominal distribution, in d-divergence-basel ambiguity sets, robustness performance is often degraded due to limited expressiveness. In contrast, when no restrictions are imposed on the support space, as in Wasserstein- based ambiguity sets, the resulting DRO is not easily tractable because of the distribution modeling capabilities of generative models to construct ambiguity sets over a finite, parameterized distribution space, with the goal of achieving both expressiveness and tractability.
probability space. This observation motivates us to leverage the strong
4
AMBIGUITY MODELING BASED ON GENERATIVE MODELS
4.1
PROBABILITY MODELING VIA GENERATIVE MODELS
The foundation of the generative ambiguity sets is the probability modeling based on generative models. Here, we representatively introduce the probability models of diffusion models and VAEs since they are both widely-used generative models.
Diffusion Models. The diffusion models learn data distributions through a forward process and reverse denoising process which are introduced based on the formulations of variational diffusion models Ho al. (2020); Luo (2022).
Forward Process. The forward process in a diffusion model begins with an initial sample xo E Rd drawn from the nominal distribution Pa. and evolves according to Markov chain that gradually adds Gaussian noise to the data:
T q( xo) = II a( x1-1), q(re I xx-1) = Var I1)
(2)
where a E (0, 1) is variance schedule at step With a slight abuse of notation, we also write the nominal distribution Po as q(ro). A sample of the forward process at step can be written as
Jt. 20 (1 ā.).,, where is a standard Gaussian noise and a II==1 0r-
Reverse Process. By reversing the forward process, a diffusion model learns to recover the original data distribution starting from Fy p' where p' is prior distribution usually selected as a standard Gaussian distribution. The reverse process can be represented by transitions parameterized by 0:
T
Po( = 2/ (XT) II pa (xt-1 xx), po po(ar-1 1 wx) = (1, 1), Eo (2ie, 1))
(3)
where no (x), 1) and 2x, 1) are parameterized mean and variance.
4

---

Published as a conference paper at ICLR 2026
The loss function to train the diffusion model is derived to guide the reverse process to approxi- mate the forward process, It holds by Bayesian's rule that the denoising transition corresponding to the forward process is q(x211 zo)
( o,I)
o,I) where By (z t)
Vi-āuva;"" with v being the standard Gaussian noise to sample 2, and o, [1
1-0e
1-a: a.
Thus, to match the denoising mean and variance, express the parameterized denoising mean as pi (Ic. रoi 2t VI-árya, I-ai so(2,, t), where sa(2,, t) is a ML model parameterized by 0 to
predict the noise that determines the mean of rt-1 from He. and Eg I) = of I. can be proved that maximizing the log probability log (T( averaged by Po is equivalent to minimizing the loss below
T Σu t=1
(D, Po) = -EZOMM
(4)
where 2t Vā zo + (1 Vz)r with My being standard Gaussian noise, and 2n? (1- Joe.
(1 (1-a.)
Note that although the forward and reverse processes are formulated based on variational diffusion models Ho et al. 22200), the probability modeling Po and similar loss can also be derived by score- based generative models Song& Ermon (2019); Song et al. (202T) as detailed in Appendix [F]
Variational Autoencoders (VAEs). VAE Kingma et al. (2019) encodes an input into latent variable z and reconstruct the input x from the latent variable. The encoder of VAE is a y-- param- eterized distribution 4 (z(zz) conditioned on the input z. The decoder po(x]z) parameterized by 0 is used to reconstruct æ from the latent variable 2. The probability model VAE is
pa (x(x)= (z) - pa (2, dz EEny [x](2]
(5)
where (=) is the prior distribution.
VAE is trained to maximize the log-likelihood log (ar). Since directly maximizing log (o()) is intractable, we usually minimize the average Evidence Lower Bound (ELBO) loss:
(6) The first term, /VAE(B. 0, Po) E [ [- In (x(]=)]] is the reconstruction loss that en- courages the learned model regenerate data that is consistent with the original data distribution. The second term RVAR(p', ,, Po) Ez- [DxI.(g,<>(z) || p p²(z))] is the prior matching loss which measures how similar the learned latent distribution is prior belief held over latent vari- ables. In practice, the encoder is usually parameterized as Gaussian distribution with mean and variance given by neural networks 4:p (2 ()) = M(= mi (x), (=))) such that the latent vari- able can be sampled as = (p(()) + +p(()) o v, N(O, where is element wise product.
Po p. 2)] = Ez-Po Po [B lo (x]]==]]] + Ez-Po (g)la)
4.2 GENERATIVE AMBIGUITY SET
From the probabilistic modeling perspective of generative models, well-trained generative model with sufficient expressive capacity can closely approximate the nominal distribution, while pertur- bations of such generative model can induce distributions that deviate from the nominal distribu- tion. To build the generative ambiguity set, we need to establish the consistency of the perturbed generative models with respect the nominal distribution. Thus, we show in Lemma [ that for likelihood-basee generative models, the discrepancy between the nominal distribution Po and the generative distribution Pe can be formally bounded through the inclusive KL-divergence.
The inclusive KL-divergence between the nominal distribution Po and the sampling distribution Pg of a likelihood basdd generative model can be bounded as
Lemma 1.
DKL (PollPP) < J(6, Po) + R(p Po) C₁
where [e, Po) is the reconstruction loss which can be JDM (6, Po) in equation or (f, Po) in the first term of equation 6 R(p',Po is prior matching loss relying on a prior distribution p' but not relying on 0, and C2 is constant that does not rely p' or 0.
The proof of this lemma can be found in Appendix C.I Lemma D shows that if the reconstruc- tion loss of diffusion model or VAE is bounded, i.e., J(0, Po)
c, then the KI divergence
5

---

Published as a conference paper at ICLR 2026
(K Po) is bounded by plus additional terms not relying the parameter 0. The prior matching term R(p'. Po) depends on the generative model designs. For example, it holds that
R(y', Po) (p)(w) DKL ((xr(=)))ll'
for diffusion models, and well-designed forward
process can contribute nearly zero prior matching error Ho et al. (2020); Song et al. (2020). The constant C1 relies on the negative entropy -H (Pb) of the nominal distribution.
Importantly, the inclusive KI-divergence DKL. h Po) in Lemma is not the excluvive KL.-
divergence DKL (P& Po) in KL-DRG Hu & Hong (2013 Kocuk (202)) Liu et al. (202bb); Michel al. (2021). The inclusive KL-divergence DKL (Po || Po) Lemma alllows the distribution Pa of generative models to have broader support space than Po (i.e., Po < Pp).
Given likelihood- generative model satisfying Lemma I the induced probability distribution Pg maintains a constrained discrepancy from the nominal distribution Po whenever its reconstruction loss J J0, Po) remains bounded. At the same time, Po is capable of capturing a broader support space than Po- This observation motivates us to define ambiguity sets in DRO using generative models, leading to the formulation of DRO with Generative Ambiguity Sets (GAS).
min max Ez- [f(,, x)] s.t. J( Se
wEW BE
where J(, Po) the reconstruction loss of a likelihood-bant generative model such as diffusion models or VAE, Pa denotes the sampling distribution of the generative model, and c is the adversarial budget.
The generative ambiguity set has the following advantages. Due to the distribution modeling capability of diffusion models, the constraint on the reconstruction loss J(0, Po) ensures that the generative distributions remain consistent with the nominal distribution Po. mitigating over- conservativeness issucs in DRO. In addition, it leverages generative models' ability to generate di- verse samples heyond the nominal support space, enabling the identification of distributions with the worst-case optimization performance. Furthermore, the inner maximization in equation 7 operates finite, parameterized space, leading tractable solutions.
5 ALGORITHM FOR DRO WITH GENERATIVE AMBIGUITY SET
Despite the advantages of generative ambiguity sets, it is challenging to solve DRO with generative ambiguity set in equation given the complexity of generative models. In this section, we provide an algorithm to solve DRO with generative-ambiguiy set in GAS GAS-DRO (Algorithm [ Illustrated in includes an inner maximization (InnerMax)) loop and outer minimization
Figure I loop which are introduced below.
Inner Maximization (InnerMax). The inner maximization of equation (roonn/preee presents the follow- ing challenges. First, the inner maximization is a constrained optimization over the diffu- sion parameter space, which can not be directly solved by gradient-based methods. Second, the objective an expectation over the distri- bution parameterized by the generative model, which cannot be directly differentiated. There- fore, need reformulate the objective into a tractable form.
Aeurs
Model
Generative Ambiguity Set
Optimizaaion
Sa
Pg
Objective Function
Dataset
Optimization
Sg Pa
To solve the constrained maximization prob- lem, we adopt dual learning approach: intro- ducing Lagrangian dual μ > 0 to reformulate it the unconstrained problem:
Figure 1: DRO with generative ambiguity set
max Fa [f .x) - (0, Po)
8
and updating μ via dual gradient descent. As shown in Algorithm [ the Lagrangian dual variable in equation is adaptively updated to balance the maximization of the expected objective with the enforcement of the constraint. Intuitively, Algorithm] I increases μ when J(0, Po) exceeds the bud- get e, thereby placing greater emphasis on reducing constraint violations, and decreases μ otherwise to prioritize objective maximization.

---

corresponding steps in J(A, Po). The performance of this trick to reduce computation complexity is objective Published as a conference paper at ICLR 2026
Algorithm 1: DRO with Generative Ambiguity Set (GAS-DRO)
Input: Nominal distribution Po: Adversary budget > 0; Step size v) > 0, 0,
Init: Initialize decision variable wc sampling model parameter 0 and Lagrangian weight > 0:
for j 1,
; 1, , do //
maximisation (InnerMax) update generative model
for k = 1, K do
Update generative model parameter 02 by solving equation 8 given ii; Update Lagrangian parameter: μ +- max J4 + (0,s,, Po) })};
// minimization update decision variable
Generate dataset Sj with generative model Potn (l1) (.,... 0k}:
Update decision variable w: Wi Wj-1 X - V BueS Wf 1.);;
return w uniformly selected from [wh, ... wz]:
Next, we apply the policy optimization methods transform the objective in equation into tractable form. Here, we give the transformed objectives based on the diffusion model formulations
while deferring the objective reformulations based on VAE formulations Appendix B
Based on the policy gradient trick Sutton et al. (199),, we can transform the objective equation into the following tractable form detailed in Appendix A.I
max E (x [In Pa (Zo.T) f(w, zo)] (0,
(9)
0
where ( is the empirical mean based on the T step samples for the backward process of the diffusion model Po. and In Po(To:)) = [x--1 ( 1)]² + C2 where C2 is a constant.
Proximal Policy Optimization (PPO) [Schulman et al. (017)) is believed to have more stable perfor- mance which transforms the objective equation [ into a differentiable form.
max E Passa (mo [min(ro(xoor))) (w,
K. kk) f(w, æo))] μ JDM (0, Po)
(10)
where Pa is a reference diffusion model which can be selected as a pre-trained diffusion model on the nominal distribution Po. the probability ratio is ro (=(:T) Pa (Rorr
Po (a(morr)
SSp(-EE( (|ا =o(a,,0)) - -Hany (xx.t)))) }} and 20;
(0,1) is the clipping parame- ter to avoid overly-large policy updates. To reduce the training complexity, instead of optimiz-
20
ing all the T steps, can only optimize the last T' steps of the backward process by choosing
ro (o(orr))) = expl- -I----1 [| (I
[| (e.y) }} and only keep the loss terms of 20
207
verified by experiments in Section 7 More derivation details for the
can be found in Appendix A.2
Min-Max Solution. With the inner maximization solution, integrate the design methodology of gradient descent with max-oracle Jin et al. (2020) into the min-max solution in Algorithm Specif- ically, in each iteration, we first run InnerMax. to search for the worst-case generative model Paus that maximizes the expected loss given the optimization variable w. Next, in order to provide diverse samples for the statistical optimization based on the worst-case probability model, we generate an adversarial dataset S₁ based on the diffusion model Poos and use it to update w. The convergence of Algorithm [ in proved in Theorem T Tand Theorem 2
ANALYSIS
We provide the convergence analysis which holds for Algorithm I in this section. We first prove the convergence of the inner maximization, followed by the overall convergence of the DRO with generative ambiguity sets.
Theorem (Convergence Inner Maximization). Let 0" be the optimal diffusion parameter that solves the inner maximization equation 7 given variable w. If the reconstruction loss is bounded
7

---

Published as a conference paper at ICLR 2026
as J(0) J and the step size is chosen 17 O (( the inner maximization ervor holds that
& 1 VK
A' == Pa [[(w, x)] Ey Po ((w, (
max(e, J. J1
(11)
where the outer expectation is taken over the randomness of output selection, In addition, the inclu- sive KL-divergence with respect the nominal distribution is bounded as
En (Pol s maxfe, HHC | + R(T, Po) C
(12)
VK(uc w²)
where PC > 1e" with μe being the optimal dual variable for the constrained optimization in equa- tion R(T, Po) is the prior matching loss with n being the prior distribution, and C₁ a constant.
Proofs of Theorem T are provided in Appendix C.2 The bound shows that when the inner iteration number K sufficiently large, the inner maximization error will be small enough, so GAS-DRO find worst-case distribution in the Diffusion Ambiguity Set. Moreover, the inclusive KI- divergence w.r.t. the nominal distribution
bounded by the budget e, the prior matching error H(T, (e.g. DKL (pT (|=T) in diffusion models), and the constant gap Cj which relying on the negative entropy- -H (Po). This substantiates that inner maximization of case distribution that is consistent with the nominal distribution by adjusting
can find a worst-
Theorem 2 (Convergence of GAS-DRO). Assume that objective f(w,x) is B-smooth and L-piscchitz with respect to w and is upper bounded by f. If each dataset Sş sampled from the generative model contains examples and the step size is X ~ O(V 3L then with probability
1 - 6,0 E 1). the average norm of Moreau envelope of ((w)= == max, Ep, (f(w, x)] satisfies
V1
V/2
E VI (V6=((1) (w)] S 43A
+
(13)
Vn
VH
where the expectation is taken over the randomness of output selection, A' is the error of inner max- imization bounded in Theorem
Vi
8 (206(2/0) and V22 4 (6 (wi) min,. ((w)).
The proof of Theorem is deferred Appendix C.3 It shows that with sufficiently large iteration number H and sampling size n, the average gradient norm of the Moreau envelope for the optimal inner maximization function ((w) bounded by the error A' of the maximization oracle. Since A decreases as the inner iteration number K becomes sufficiently large as proved by Theorem !| the average gradient norm of the Moreau envelope can be small enough. The convergence of the Moreau envelope gradient norm indicates that the algorithm converges to an approximately stationary point as explained in Appendix [.3.T] The intuitive reason that if the gradient norm of Moreau envelope is bounded by small enough value A for a decision variable w, i.e. V/ 20 (w)ll < A, there exists
w which is close enough to w (1ıů wyll < 23 and is approximately stationary solution for minimizing the optimal inner maximization function 6. Thus, GAS-DRO converge to an approximately stable solution of equation T with enough maximization and minimization iterations.
Impact of the expressive capacity of generative models. While the generative models with suf- ficient expressive capacity can define an ambiguity set with diverse adversarial distributions, the real world generative models always have finite capacity. The expressive capacity of generative models determines whether a real worst-case distribution can be effectively identified by the inner maxinization and further affects the robustness performance of DRO. To quantify the impact of the expressivity of generative models, we formally define the I exexprsssiiitty of the generative model.
Definition 1 (I ex-eesssssvviyy of the generative model). A generative model Pa is I eexpresssive (T 0) if for any testing distribution P in the set of all possible testing distributions P. there exists a 0 such that Dw (Po, P) < I where Dw is the Wasserstein metric.
Next, we analyze the impact of the I expprssssv on the DRO performance. Given the set of all possible testing distributions P. GAS DDR has better robustness performance if it can identify the real worst-case distribution P- arg maxp E (w, u, x=]] with smaller error. Thus, bound the error of the inner maximization Á = Bzn p ([ x)) EgErrr, [f(w, x:) in Corollary 1 (in
8

---

Published as a conference paper at ICLR 2026
We observe that I expressivity affects the inner maximization performance by introducing an additional error LyI' where Lg is the Lipschitz continuity parameter of f(w, x)) with respect When the expressivity parameter I' is smaller, the generative model has stronger abil- ity to express the possible testing distributions and GAS-DRO can identify adversarial distributions that are more consistent with the true worst-case testing distribution, as reflected by a smaller inner maximization gap. Consequently, by optimizing the parameter w based on the learned adversarial distribution, GAS-DRG can achieve better robustness across the true testing distribution. In the ideal case when r = 0, the error bound of the inner maximization coincides with the inner maximization error bound Theorem which asymptotically converges zero as the number of inner iterations K increases.
Appendix C.4).
EXPERIMENTS
This section reports the main experiment results on a representative ML fask-time series forecast- ing based on the Electricity Maps (2025) -to verify the effectiveness of the proposed algorithm. Additionally, we evaluate the performance of
on an image classification task based on the MNIST dataset[Lecus et al. (1998) and USPS dataset Hull (1994), with details in Appendix D
Task. Accurate forecasting of renewable is critical for power grid optimization, energy-efficient computing, sustainable power management, and greenhouse gas emission reduction. We evaluate GAS-DRO on time-series prediction with the Electricity Maps (2025) datasets, using an n-step slid- ing window to predict emissions at step (n 1).
Datasets. We conduct experiments using datasets from the Electricity Maps (2025) platform, which provides high-resolution spatio-temporal data electricity operations, including carbon intensity (gCO2eq/kW)) Wh) and energy mix. We construct training and OOD testing datasets based on the hourly records from 2021-2024 across four regions: California BAANC.21-24 Texas (ERCO ~1 24), Queensland (QLD 21~224), and the United Kingdom
(GB 21~24). The training set is constructed by merging BANC 23--4 (BANC 2324) into 438 sequence samples, while evaluation is performed on independent test sets of 312 samples each from other years and regions. To measure distributional shift between training and test sets, we compute the Wasserstein distance, reported in Table !
Baselines.
The baselines which are compared with GAS-DRO in our experiments are introduced method that trains the standard ML to minimize the time series forecasting error augmented datasets.
as below.
without DRO. DML represents an ML model fine-tuned with diffusion-generaane W-DRO represents the FWDRO algorithm proposed in Staib Jegelka (2017), where the ambiguity set characterized by the Wasserstein metric. KL-DRO represents the classic KL-divergence-bice based DRO framework, in which the ambiguity set is characterized by the KL- divergence (uu & Hong 2013). DRAGEN represents Wasserstein DRO framework, in which the ambiguity set is structed as a Wasserstein ball in the latent space of a learned generative model [Ren & Majumdar 2022). P-DRO Michel et al. (2021) is DRO framework that parameterizes the adversarial distri- butions using a neural generative model within a KL-divergence-base ambiguity set.
Setups. Detailed training setups of GAS and other baselines are provided in Appendix D
Main Results. The main results are given in Table T We observe that all DRO methods im- prove testing performance on OOD datasets compared ML by optimizing the worst-case expected loss, while DML improves upon ML by leveraging diffussin-generraeed augmented data. Using ML as the reference baseline, GAS- DRO achieves the highest performance gain of 63.7%, followed by DRAGEN (48.9%),P-Do)) (42.5%), DML (39.7%), KL-DRO (36.1%), and W-DRO (24.0%). Al- though both DML and traditional DRO methods deliver noticeable improvements, GAS-DRO still surpasses them by average margin of 25.4%, likely because it exploits the distribution learning capability of diffusion models to construct ambiguity sets, thereby generating adversarial distribu- tions that are both strong and realistic and ultimately achieving superior OOD generalization.
In addition, our method outperforms the DRO baselines with generative models. Although DRAGEN also employs generative model to construct its ambiguity set, it still relies on relaxing Wasser- stein DRO framework which is limited in overcoming the inherent limitations of W-DRO. Moreover, P-DRO, as a DRO algorithm that uses generative models to model the adversarial distributions, mains restricted by the KL-divergence-bart worst-case distributions. Also, the crude approximations in P-DRO may bring harmful effects. By
ambiguity set, preventing it from modeling diverse

---

Published as a conference paper at ICLR 2026
contrast, GAS-DRO directly defines the ambiguity set the parameterized space of generative models, enabling it to generate worst-case distributions that are both diverse and consistent with the nominal distribution, thereby offering principled and tractable DRO algorithm.
In Appendix D.4 we provide additional results and more detailed analy- sis. Ablation studies are
Table I: Test Result on Different Datasets.
Datasets (Wasserstein Distance)
MSE
GAS-DRO
DRACEN
P-DRO
KI-DRO
W-DRO
conducted to evaluate the impacts of the adversarial budgets and the distribu- tional discrepancies caused by different noise types and strengths. Furthermore, we empirically validate
BANC 22 (0.0240)
110047
0.0069
0.0079
0.0086
0.0073
0.0078
0,0183
BANC,21 (0.1213)
0,0054
0.0087
0.0074
0.0112
0.0121
0.0093
0.0238
QLD .24 (0.2171)
0.0450
0.0618
0.0752
0.0754
0.0823
0.0766
0.0887
QLD 23 (0.2033)
0.0681
0.0820
0.0831
0.0879
0.0834
0.0946
QLD 22 (0.2782)
0.0192
0.0298
0.0331
0.0379
0.0557
0,0822
0.0667
QLD 21 (0.3054)
0.0186
0,0288
0.0522
0.0577
0.0574
0.0339
0.069
GB_24 (0.0419)
0.0119
0.0146
0.0191
0.0178
0.0176
0,0172
0.0285
GB.23 (0.0666)
0.0100
0.0141
0.0154
0.0178
0.0200
0.0164
0.0311
the effect of the expressive capacity by evaluating GAS-DRO with diffusion models with different architectures and a VAE model in Appendix D.4.5 Moreover, we measure and discuss the computational overhead of GAS-DRO and other baselines in Appendix D.4.6
GB (0.1255)
0.0153
0.0158
0.0197
0.0245
0.0172
0.0360
GB 21 (0.1359)
0.0094
0.0140
0.0137
0.0181
0.0229
0.0188
0.0340
ERCO_24 24 (0.1206)
0,0158
0.0201
0.0211
0,0241
0.0266
0.0224
0.0379
ERCO_23 (0.1207)
0,0106
0.0146
0.0148
0.0179
0.0196
0,0162
0.0319
ERCO 22 (0.1581)
0.0076
0.0115
0.0103
0.0146
0.0187
0.0123
0.0312
ERCO (0.1417)
0.0141
0.0143
0.0189
0.0263
Average
0.0230
0.0259
0.0288
0.0342
Worst
0,0509
0,0681
0.0820
0,0831
0.0879
0.0834
0.0946
8 CONCLUSION
In this paper, we introduce a novel moderativeemoel--bueee ambiguity modeling framework for DRO and propose GAS-DRO to solve DRO with generative ambiguity sets. While the method is developed based on the formulations of diffusion models or VAEs, the proposed framework is gen- eral and can be adapted to other likelihood- distributional modeling capacity of generative models, GAS-DRO can effectively identify adversar- ial distributions that induce worst-case performance while remaining consistent with the nominal distribution, thereby yielding tractable DRO solutions with strong generalization performance We establish the stationary convergence guarantees of GAS-DRO and demonstrate through experiments its superior OOD generalization performance by comparing with the state-of-the-an art baselines.
generative models. By exploiting the expressive
Future Directions. would be interesting to formally establish consistency bounds for the training losses of other generative models and to exploit them for constructing ambiguity sets in DRO. In addition, improving the training efficiency of generative models would directly improve the com- putational efficiency of GAS-DRO Moreover, generative ambiguity sets provide a flexible interface for incorporating prior knowledge about the testing environment into the generative model, which can further enhance robustness and generalization under distributional shifts.
Broader Impacts. The proposed algorithm provides an effective generative model based approach to ambiguity modeling that addresses fundamental challenges in DRO. It has the potential to bridge the gap between imperfect training environments and real-world deployment, thereby enhancing the safety and robustness of machine learning models across a broad range of critical real-world tasks.
ACKNOWLEDGMENT
This work were supported in part by University of Houston Start-up Funds 74825 [R05120 (00512039] and 74833 [R05120 2042].
We thank the ICLR reviewers for the constructive comments which helped us to improve the draft.
10

---

Published as a conference paper at ICLR 2026
ETHICS STATEMENT
This work adheres to the ICLR Code of Ethics. Our research does not involve human subjects, per- sonally identifiable information, or sensitive personal data, and thus IRB approval was required. All datasets used in this study are publicly available, widely recognized in the research community, and utilized strictly in accordance with their licenses. For any preprocessing, we provide clear docu- mentation to ensure transparency. The proposed methods are designed to advance the theoretical and practical understanding of distributionally robust optimization and generative modeling. These con- tributions are primarily academic in nature and intended to support future work in reliable machine learning. We do not anticipate direct harmful applications.
REPRODUCIBILITY STATEMENT
We have taken extensive measures to ensure the reproducibility of our work. complete description of the proposed model architecture, training objectives, and optimization procedure is provided in Section The datasets used in our experiments are all publicly available and widely adopted in the research community. We provide detailed account of the data preprocessing pipeline, including normalization and sampling strategies, in Appendix D
To further enhance reproducibility, all hyperparaameters, model configurations, and training sched- ules are reported explicitly in Section 7 and Appendix D along with the evaluation metrics and criteria employed for performance comparison. For theoretical contributions, clearly state all assumptions and provide complete, rigorous proofs our claims in Appendix C
We also include ablation studies and sensitivity analyses illustrate the robustness of our sults with respect to key hyperparameters, which may assist future researchers in replicating or extending our work. Finally, we provide our source code in the link https///ittub.com///LLBBB Houston/GAS-DRO

---

Published as a conference paper at ICLR 2026
REFERENCES
Martin Arjovsky. Out of distribution generalization in machine learning. PhD thesis, New York University, 2020.
Uljad Berdica, Kelvin Li, Michael Beukman, Alexander David Goldie, Perla Maiolino, and Jakob Nicolaus Foerster. Robust offline learning via adversarial world models. In NeurIPS 2024 Workshop on Open-Wote World Agents, 2024.
Jose Blanchet, Daniel Kuhn, Jiajin Li, and Bahar Taskesen. Unifying distributionally robust opti- mization via optimal transport theory. arXiv preprint arX1v:2308.05514, 2023.
Changyou Chen, Han Ding, Bunyamin Sisman, Yi Xu, Ouye Xie, Benjamin Z. Yao, Son Dinh Tran, and Belinda Zeng. Diffusion models for multi-modal generative modeling. arXiv preprint 0rXiv:2407.17571, 2024.
Huanran Chen, Yinpeng Dong, Zhengyi Wang, Xiao Yang, Chengqi Duan, Hang Su, and Jun Zhu. Robust classification via a single diffusion model. arXiv preprint arXiv:2305. /5241, 2023.
Ruidi Chen and Ioannis Ch Paschalidis. A robust learning approach for regression models based on distributionally robust optimization. Journal of Machine Learning Research, 19(13)11-4,, 2018.
Ruidi Chen, Joannis Ch Paschalidis, et al. Distributionally robust learning. Foundations and Trends@ in Optimization, 4(1-2):1-245 2020.
Xiuyuan Cheng, Yao Xie, Linglingzhi Zhu, and Yunqin Zhu. Worst-case generation via minimax optimization in wasserstein space. arXiv preprint arXiv:2512.08176, 2025.
Xuclong Dai, Yanjie Li, Mingxing Duan, and Bin Xiao. Diffusion models strong adversaries. IEEE Transactions on Image Processing, 2024a
Xuelong Dai, Kaisheng Liang, and Bin Xiao. Advdiff: Generating unrestricted adversarial examples using diffusion models. European Conference on Computer Vision, pp. 93-109. Springer, 2024b.
Bingqian Du, Zhiyi Huang, and Chuan Wu. Adversarial deep learning for online resource allocation. ACM Transactions on Modeling and Performance Evaluation of Computing Systems, 6(4):1-25, 2022.
Electricity Maps. The world's most comprehensive electricity data platform. https://www Electricitymappl com, 2025.
Rui Gao and Anton Kleywegt. Distributionally robust stochastic optimization with wasserstein distance. Mathematics of Operations Research, 48(2):603-655 2023a.
Rui Gao and Anton Kleywegt. Distributionally robust stochastic optimization with wasserstein distance. Mathematics of Operations Research, 48(2):603-655, 2023b.
Brennen Hill. Learning an adversarial world model for automated curriculum generation in marl.
arXiv e-prints, pp. arXiv 200v,25599202 2025.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 6840 6851. Curran Associates, Inc.,
2020. URL httppss://procceedngg...... neurips www.eei....ppppppeeeeeeeee cc/paper_......... files//aper/20207 File/Acbbccccc38a880051111111111 4b- .pdf
Zhaolin Hu and L Jeff Hong. Kullback-leibler divergence constrained distributionally robust opti-
mization. Available at Optimization Online, 1(2)9, 2013.
J.J. Hull A database for handwritten text recognition research. IEEE Transactions Pattern Analysis and Machine Intelligence, 16(5):550-554, 1994. doi: 10.1 00.1109/3..294400
Hisham Husain, Vu Nguyen, and Anton van den Hengel. Distributionally robust bayesian optimiza- tion with d-divergences. 2023.
12

---

Published as a conference paper at ICLR 2026
Chi Jin, Praneeth Netrapalli, and Michael Jordan. What is local optimality in Mincconnee--ononnccvve minimax optimization? In International conference on machine learning, pp. pp. 4880-4889, 488.. PMLR, 2020.
Leonid Vasilevich Kantorovich and SG Rubinshtein. On space of totally additive functions. Vestnik of the St. Petersburg University: Mathematics, 13(7):52-59, 1958.
Diederik P Kingma, Max Welling, et al. An introduction to variational autoencoders. Foundations and Trends@ in Machine Learning, 12(4):307-33) 2019.
Johannes Kirschner, Ilija Bogunovic, Stefanie Jegelka, and Andreas Krause. Distributionally robust bayesian optimization. In International Conference Artificial Intelligence and Statistics, pp-
2174-2184. PMLR, 2020.
Burak Kocuk. Conic reformulations for kullback-leibler divergence constrained distributionally robust optimization and applications. arXiv preprint arXiv:2007.05596, 2020.
Daniel Kuhn, Peyman Mohajerin Esfahani, Viet Anh Nguyen, and Soroosh Shafiezzadeh-Abadee Wasserstein distributionally robust optimization: Theory and applications in machine learning. In Operations research & management science in the age of analytics, pp. 130-166. Informs, 2019.
Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recog- nition. Proceedings of the IEEE, 86(11):227--2324, 1998. doi: 10. 109/5.7667.
Alexander C Li, Mihir Prabhudesai, Shivam Duggal, Ellis Brown, and Deepak Pathak. Your dif- fusion model is secretly a zero-shot classifier. In Proceedings of the IEEE/CVF International 2023.
Conference Computer Vision, pp.
Dawei Li, Yue Huang, Ming Li, Tianyi Zhou, Xiangliang Zhang, and Huan Liu. Generative models for synthetic data: Transforming data mining in the genai era. In Proceedings of the 34th ACM International Conference on Information and Knowledge Management, pp. 6833-6836, 2025.
Junnan Li, Yongkang Wong, Qi Zhao, and Mohan S Kankanhalli. Learning to learn from noisy labeled data. In Proceedings of the IEEE/CVF conference on computer vision and pattern recog- nition, pp. 5051-5059, 2019.
Jiashuo Liu, Zheyan Shen, Yue He, Xingxuan Zhang, Renzhe Xu, Han Yu, and Peng Cui. Towards out-of-distribution generalization: A survey, arXiv preprint 2tt.::2001.1322444 2021.
Jiashuo Liu, Jiayun Wu, Bo Li, and Peng Cui. Distributionally robust optimization with data geom- etry. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 33689-33777 33701. Curran Associates, Inc., 2022a. URL
neur cc /aceppaa paper files/paper/20 -aa-seeee--apppr--e Conference pdff
Zijian Liu, Qinxun Bai, Jose Blanchet, Perry Dong, Wei Xu, Zhengqing Zhou, and Zhengyuan Zhou. Distributionally robust q-learning. In International Conference on Machine Learning, pp- 13623-13643. PMLR, 2022b.
Miao Lu, Han Zhong, Tong Zhang, and Jose Blanchet. Distributionally robust reinforcement learn- ing with interactive data collection: Fundamental hardness and near-optimal algorithms. Ad- vances in Neural Information Processing Systems, 37:12528-12580, 2024.
Calvin Luo. Understanding diffusion models: A unified perspective. arXiv:2208.11999 2022.
arXiv preprint
Xutao Ma, Chao Ning, and Wenli Du. Differentiable distributionally robust optimization layers. 2024.
abs 2406.
Paul Michel, Tatsunori Hashimoto, and Graham Neubig. Modeling the second player in distribu- tionally robust optimization. In International Conference on Learning Representations, 2021.
Paul Michel, Tatsumori Hashimoto, and Graham Neubig. Distributionally robust models with para- metric likelihood ratios. In International Conference on Learning Representations, 2022.
13

---

Published as a conference paper at ICLR 2026
Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the wasserstein metric: Performance guarantees and tractable reformulations. Mathematical Programmingg 171(1):115-166) 2018.
Fabio Muratore, Fabio Ramos, Greg Turk, Wenhao Yu, Michael Gienger, and Jan Peters. Robot learning from randomized simulations: A review. Frontiers in Robotics and AI, 9:799893, 2022.
Joaquin Quinonero- Candela, Masashi Sugiyama, Anton Schwaighofer, and Neil D Lawrence. Dataset shift in machine learning. Mit Press, 2022.
Allen Z Ren and Anirudha Majumdar. Distributionally robust policy learning via adversarial envi- ronment generation. IEEE Robotics and Automation Letters, 7(2):1379-1333 2022.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 2017.
Aman Sinha, Hongseok Namkoong, Riccardo Volpi, and John Duchi. Certifying some distributional robustness with principled adversarial training. arXiv preprint 1rX1v:1710.10571, 2017.
Elena Smirnova, Elvis Dohmatob, and Jérémie Mary. Distributionally robust reinforcement learning- arXiv preprint arXiv: arXiv:1902.08800 2019.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint 0ett20000..2502. 2020.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Yang Song, Conor Durkan, Jain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances neural information processing systems, 34:1415- 1428, 2021.
Matthew Staib and Stefanie Jegelka. Distributionally robust deep learning as generalization of adversarial training. In NIPS workshop on Machine Learning and Computer Security, volume 3, pp. 4, 2017.
Matthew Staib and Stefanie Jegelka. Distributionally robust optimization and generalization in ker- nel methods. Advances in Neural Information Processing Systems, 32, 2019.
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. icy gradient methods for reinforcement learning with function approximation. 12. 1999. URL
Pol-
file htte://555222bbbbccceeccccaaaaa.
Yunze Tong, Fengda Zhang, Zihao Tang, Kaifeng Gao, Kai Huang, Pengfei Lyu, Jun Xiao, and Kun Kuang. Latent score-based based reweighting for robust classification on imbalanced tabular data. In Forty-second International Conference Machine Learning, 2025.
Andreas Veit, Neil Alldrin, Gal Chechik, Ivan Krasin, Abhinav Gupta, and Serge Belongie. Learning from noisy large-scale datasets with minimal supervision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 839-847, 2017.
Jie Wang, Rui Gao, and Yao Xie. Sinkhorn distributionally robust optimization. Operations Re- search, 2025a.
Prince Zizhuang Wang, Jinhao Liang, Shuyi Chen, Ferdinando Fioretto, and Shixiang Zhu. Gen- dfl: Decision-focused generative learning for robust decision making. 2025b. URL https:
05468
Zitao Wang. Ziyuan Wang, Molei Liu, and Nian Si. Knowledge-guided wasserstein distributionally robust optimization. arXiv preprint arXiv:2502.081445 2025c.
Weicheng Xie, Zenghao Niu, Qinliang Lin, Siyang Song, and Linlin Shen. Generative impercep- tible attack with feature learning bias reduction and multi-scale variance regularization. /EËE Transactions Information Forensics and Security, 2024.
14

---

Published as a conference paper at ICLR 2026
Chen Xu, Jonghyeok Lee, Xiuyuan Cheng, and Yao Xie. Flow-based distributionally robust opti- mization. IEEE Journal on Selected Areas in Information Theory, 2024.
Yufeng Yang, Yi Zhou, and Zhaosong Lu. Nested stochastic algorithm for generalized sinkhorn distance-regullarized distributionally robust optimization. arXiv preprint arXiv:2503. 229223, 2025.
Hengrui Zhang, Jiani Zhang, Zhengyuan Shen, Balasubramaniam Srinivasan, Xiao Qin, Christos Faloutsos, Huzefa Rangwala, and George Karypis. Mixed-type tabular data synthesis with score- based diffusion in latent space. The Twelfth International Conference on Learning Represen- tations, 2024.
Linglingzhi Zhu and Yao Xie. Distributionally robust optimization via iterative algorithms in tinuous probability spaces. arXiv preprint arXiv:2412.20552 20556, 2024.
15

---

Published as a conference paper at ICLR 2026
A DETAILS OF GAS-DRO (ALGORITHM 1) WITH DIFFUSION MODELS
In this section, we provide details to transform the objective of inner maximization in Algorithm based on the formilations of diffusion models.
A.1 OBJECTIVE DERIVATION POLICY GRADIENT
Let 2o:T be the output vector of each step in the backward process of the diffusion model. De- note Porr,0 as the joint distribution of To:r. Since 2o ~ Po. we can express the first term of the Lagrangian-relaxed objective
Bg (w, x)] = BaoT Por (ww,xo)
(14)
Then, the gradient of Lagrangian-relaxed objective can be expressed as
Vg PE (f PJ
V -_
HJ P.))
(ou ~Vo Po)
(15)
Foor
Po:T,o (f P -oattoor--ooJJo VVJ Po)
BOOT
-Ezor--Prr Porn, [f( (w, 20) Ve o Po:r. Po.r)((zor]]] (o.
The first term E Br Pozz.e be calculated by empirical mean E based the examples
sampled from Po.T. Thus, we can equivalently implement the gradient ascent by optimizing the objective in equation 9
max Ex [n Po.r, f(w, æo)] P()
(16)
8
Next, we derive the expression for In (( Consider discrete-time diffusion backward process as
It-1 (a () aew.
(17)
where w; is a standard multi multi-dimensional Gaussian variable. Given 6, the conditional probability each step t is
P .x)) (( (), o
(18)
We can explicitly express the joint distribution Po:r,e of Forr for all T backward steps:
T
ET.
PO.T,((TOTT) = P(xT) II Pe- xx) 2m
e
202
2
(19)
Thus, we can get the expression of In Po:1,0(0u:2) as
In Po.,,(o0))) = Σ
|| (o0,,, + C2
(20)
where C₂ -I(@Tr)) ||
A.2 OBJECTIVE DERIVATION BY PROXIMAL POLICY OPTIMIZATION
In this PPO method, convert the first term of equation 8 as PPO-like objective given reference diffusion model Po., i.e.
Ex-P, [f x)] = Egg Parr, [f(w,, æo)]
Po.T,
(w, 2n)
(21)
= Exrg - Por..9
Pozroo (xo.r)
16

---

Published as a conference paper at ICLR 2026
where the probability ratio is
ro (x = Po.T,a(VO:TT) exp
T
[ (1 1)1
|
:40 (1 1)12
Po:T, 80 (ZO.T)
20?
20
given the joint probability expression. The expectation can be approximated by empirical mean based on the examples sampled by Po.T, Like PPO, we apply clipping the ratio to avoid overly-large updates, and we can get the objective in Eqn. equation
max E Par,, (Roor)) (mii((o((ooo))) ((w, cli((rooor)l)( K, +x) (o))] p J(0, Po)
(22)
where x E (0, is the clipping parameter.
B DETAILS OF GAS-DRO (ALGORITHM I WITH VAES
B. AMBIGUITY SET FOR VAE
Given the formulations of VAE Section as proved in Appendix restated as below.
Lemma
can be
Lemma (Restated for VAEs). The inclusive KL-divergence between the empirical distribution Po(x) and the VAE distribution Po(x)) is bounded as
(x)( po(x)) S Jvar(s,, 8, Po) + Epo(a) DKalas(z I [1)( H(Po(z))
where JuAK(s, 0, Po) = By (slux) [- poop =)] is the reconstruction loss of VAE, qu x)) denotes the encoder with parameter yp (approximate posterior), po(x 2) denotes the decoder model with parameter 8, and p' (z) is the prior distribution.
To construct a VAE-based ambiguity set for DRO, we first train the a VAE model based on the nominal distribution Po:
(40" 00 = min( RE-ro (xlx)) pog]o(r))) + Egwro [DKI(g ll [[)(])
(23)
8.8
Then, keep o 40 and change 0 in the set {0 JvAE(ee, 0, Po) s c). In this way, get DRO with VAE-based generative ambiguity set as
min max NEW E, [f(w,)),, s.t. 1v se
(24)
where JV, Po) = Pp E (xg() log Po(x |z]].
B.2 ALGORITIM DETAILS GAS -DR WITH VAES
As in Algorithm I to solve the constrained maximization problem, we adopt dual learning ap- proach: introducing a Lagrangian dual μ O to reformulate the unconstrained problem and updating via dual gradient descent:
max Ex [f x)] p.vvell" .00 0.
(25)
log (x =)], we first encode the initial dataset Sy ~ Po into latent variables Zo using the pretrained encoder y0" Then, we can empirically
In order to get JVAE((p", ,,, (x))) = (er
compute the reconstruction loss as JvAEEs0", 6 (x)) = E.-Ro..-22
log Px (x] =)]]
Next, solve the inner maximization in Algorithm [ we apply the policy optimization methods to transform the objective in equation 25 into tractable form. As shown in Appendix A. based on policy gradient SSttonn et al. (999)... we can transform the objective equation 25 into
max (p (o(x, 2) f(w, x)] - μ JvAE(p", (x)))
(26)
where Epo(o..)) the empirical mean based the decoding process of the VAE model.
17

---

Published as a conference paper at ICLR 2026
Similar to diffusion models, PPO Schulman et al. (2017) provides enhanced stability by transforming the objective equation 25 into a differentiable objective function.
max E pe (2,z) [min(ro (w, x)), (l
++) ((w,))) - JVA (o())
(27)
where Pe (z, 2) is reference VAE decoder which can be selected as a pre-trained VAE decoder on the nominal distribution Po (z)). the probability ratio is ro(x, Pa(w.z) and (0, is the clipping parameter to avoid overly-large policy updates.
Pezaa (x,2)*
C THEOREM PROOFS
C.I PROOF LEMMA
Lemma 1.
The inclusive KL-divergence between the nominal distribution Po and the sampling distribution Pg of likelihood-base generative model can bounded as
(Pol I J(e, Po) + R(p'. Po) + C1
where J0, Po) is the reconstruction loss which can be JDM (0, Po) in equation A JvAR(0, Po) in the first term of equation 6 R(y', Po) is a prior matching loss relying on prior distribution but not relying on 8, and C₁ is constant that does not rely on p' or O.
Proof. Likelihood-based generative models are typically trained to maximize the lower bounds of their expected log-likelihood given their sampling distribution Pe and nominal distribution Po. i.e.
Ep,(s) log Pa(x)) Applying the relationship between KI. divergence and cross-entropy, we have
DKL (P((((()))PPP((E))) Pop) log Po(ar) + Epo(i)) log Po(x))
(28)
where the first term can be bounded by the training loss J(0, Po) of likelihood-based generative models plus average prior matching error R(p', Po), and the second term is the entropy of Po- Thus, we can get a general bound of the KL. divergence in Lemma]] []
Next, give the concrete bounds for variational diffusion models and VAEs. Also, the bound for sooe-mmtthiin--bsee diffusion models is given in Lemma based on the SDE formulations of diffusion models.
We first give the bound based on the variational diffusion models. According to Luo 2022) it holds that
Ero(r) log Po(xx)
2 (xr)
Be(er log (#o 21)] - Bro(z) DKL q( (o)
Bpa(z)
(xt [DKL I 20)l me))l|
1=2
-DMpm Po) d
Epo(2) DKL ((orkw0))l(o)))
(29)
where (0, Po) the loss in equation and o is the variance for the last reverse step in equa- tion 3 Based on the relationship between KL divergence and cross-entropy, we have
(Po( - BPo(.))
log (o))
+ +Ep(())) log Po(ar)
(30)
< JDM Po) Ero(r) DKL ((erlz)))l((r))
+
where Ci = dllolll22mn - H(P)). Therefore, get the bound in Lemma 1 with R(p', Po)
EP(())
DKL ((0T) (((er(to)lPla))
18

---

Published as a conference paper at ICLR 2026
Next, prove Lemma for VAEs. According to Kingma al. (2019),
bound the the
log-liken log-likelihood log (e( by its EVIDENCE LOWER BOUND (ELBO):
log po 2 (x x) p(=)) - E (el=) (og po 2)]
(31)
Applying the relationship between KI. divergence and cross-entropy, we have
D Po(z)) x Ba log p0 x)] + (x) 1 x))l)(())) (P)((z))
(32)
Thus, we get the inclusive KL-divergence bound in Lemma for VABS with J(0, Po) JVAE( Po) (n Po 2)]. R(p', Po) Bro(z) 2) (z)).
and C (P((z))..
C.2 PROOF OF THEOREM 1
Theorem 1. Let l" be the optimal diffusion parameter that solves the inner maximization equation 7 given a variable w. If the expected score more-matching loss bounded (()) S J and the step size
chosen as ~ O (RR the inner maximization error holds that
A == EawPo [fww, x)] - Eg (f(W
S
max(c, 111
(33)
VK
where the outer expectation is taken over the randomness of output selection. In addition, the in- clusive Kl-divergence with respect the nominal distribution is bounded as E& [DKI (Pol) s
+ max (ax(e_)(le--p()) DKL (P)l(T) T) C1, where PC μe and C₁ are constants, Pr is the output distribution of the forward process, and π is the initial distribution of the reverse process.
(Kipe-p")
C.2.1 CONVERGENCE OF DUAL GRADIENT DESCENT
To prove Theorem I following lemma.
we need the convergence analysis of general dual gradient descent in the
Lemma 2. Assume that the dual variable is updated following
pi+1 max [uu - 11 by- 0}
where 1j > 0 and by has the same dimension as μ and maxt-,0f 0} is an element-wise non-negativity operation. With as the step size in dual gradient descent, given any μ > 0, we have
Σ (ue
ns bab
1 K
Σ
ball
11 (2
(34)
K
2Kn
Proof. For any dimension j such that jik. 2 nbbb.j we have uk+11 = [kk.j - br.j. so it holds for any p. 0 that
bu.j + 1 (URtt) - mmns)) (ng pk+1.)
(35)
11
For any dimension j such that poc. < qbk.ss we have 1K+11 = and it holds for any μ > that
beet + (untis - (Ms - -K+1.) = (bu..
[kk. )(g 2 (bx.j
nybk.j
) =
(36)
"
17
Combing equation 35 and equation 36 and we have for any μ >> 0 that
T
b/ + 1 - - 1mm) (1e- 1k+1) 2
(37)
19

---

Published as a conference paper at ICLR 2026
Therefore, it holds for any lo E [K] and j 0 that
(cex be bv =(ux 1k-1) by + (MR+1 - -11) -1)"bE bk
S(4x - 1k-1) bb + (uu+1 m1)) T 1K+11)
1
(38)
=Iux pk-1) b +
(Ilx [xxl1 - 11 pRtill? || mxl119
2m
Sal b&ll
(1l - Mxl
-x+11ll)
2m
where the first inequality holds by equation and the last inequality holds because ||a||
37 the second equality holds by three-point property, |
2 2 2a bb
Taking the over & E [K], we have
1 (1l M1ll' m2ll?)
Me b,b S
| 11
+
(39)
2n
A=
k=
which proves Lemma
C.2.2 BOUND EXPECTED OBJECTIVE THEOREM
Proof. Denote F(0) = Pa (w, x)] and the optimal parameter to solve the inner maximization of equation 7 is 0". Define the dual problem the inner maximization equation 7 as
Q()) maaxF((B + plc So))
(40)
Given any y 0, it holds be weak duality that
F(0") < F(O") + ple J(0", (b) S Q())
(41)
where the second inequality holds because Q(z) maximizes F(0) u(e - (()) over 0. Thus, the average gap between the expected loss of 6* and 0₁ is
K (((0)) F(O.)) < (Q())) -((+))
T
1
Σ
K
T
t=1
1
uale S(0)
K
(42)
K
Σ
(e - J(Kk. Sol) + 2n K 112
$ " K
k=1
$ n/mmxxfc f1)2 + 1
11
K
where the equality holds because 0₁ arg maxe F(0) + malE So)) by equation 8 and so Quux)) = F(Ox) + u4(E - J(O,. So)), the second inequality holds by Lemma 2 with the choice of J O, and the last inequality holds by the assumption So) < J.
Choosing VK JJ it holds that
1
Σ
(F(O)) F(OK)) <
1
max ())
(43)
K
VK
which proves the bound of the expected loss given the uniformly selected [K].
20

---

Published as a conference paper at ICLR 2026
C.2.3 BOUND OF KL DIVERGENCE IN THEOREM
For iteration k; denote ((0) = J(0, So) and by J So). Denote the constraint violation on the score matching loss at round k as = J(2n, So) E. Denote the optimal dual variable as
fe² = arg ming Qu)). Choose a dual variable He > p,, we have the following decomposition.
K
K
K
Σ (PE nc)ba= = Σ
pk bk + Σ
(c))
(44)
Am
km
km
For the first term, we have
K
K
Mk
bk
F(Ok) + pk . bb - F((6) = Σ
Q(ux)= - F((kk)
1=1
R=1
K
2 KQ(P))
Σ F(ON)
K
2 KQ(N")
Σ
max
F(0)
(45)
2 -QQW')) max (00) + (e Un (()))
K
-KQ(")-KQ())) p.
Σ
Uk
2k
k=1
where the first inequality is because 1r" minimizes Q(xx), the second inequality holds because
Og E {0] J(0, Se Us the third inequality holds by weak duality for μ². Continuing with equation 44 we have
Σ (uR pb) >
v Uk + uc (C((, - +)) = (-p" + HC) "(-
Σ km
Uk
(46)
1=
where the equality holds because Uk J(Kk, So) c.
By Lemma
2 with the choice of μ HC, we have
K
(kk Hc) S 17
1 K
|bxl
1
(MC (10)2 < ()m
1
(MC 11
K
2Kn
2Kn
(47)
(1)
If the step size is chosen as 1)
VK JJ it holds that
K
1
max(c, J) (ni) + IIc -20
(48)
K
VK
24((1)
Since MC is larger than nt, by equation
46 we have
maxfe, JJ (1(() + luc-u00122
K Σ
2012)
Ca
uk K(uc 11))
(uR uc)bk S
(49)
VK(Hc -1¹ p*)
K
k=1
which means & Σk J(O,, So) S VK
Co
Since J(Ok, (o) is the denoising score matching loss JOSM (0,r(.))) by Lemma 1 the proof by
we complete
K DKL (Po. Po.) S
K J(O,, Su) + DKL + Ci
K
K
k=1
(50)
C₁
< +
+ (ky +
K
21

---

Published as a conference paper at ICLR 2026
where Ci 1 r(1) (p0.s)) [V log Pr(xx)ll
b/ll - log Pej(az mo)]l) dl, C2
max(e.J) ( 1 2nlc)
with Pa /²"
(MC-1²)
C.3 PROOF OF THEOREM 2
Theorem rellie on the following assumptions.
Assumption 1. The objective function f(w, x) is L Lipschitz with respect w, i.e. for any x E X,
w, w E W, it holds that
(f(w,x) - f(w', 2>1 S L- ||w w 22 Moreover, f(w,,) x) is bounded by f, i.e. for any x E X, w E W, ((,,,)))l S f.
Assumption 2. The objective function f(w,z) is 3- 3-smooth with respect w, i.e. for any X, w, w E W, it holds that
V [w x))l2 S B B-1lw wwll2
Assume that the DRO objective [(w,,)) x) is 3 smooth and T.-Lipschitz with respect to w and is upper bounded by f. If each sampled dataset Sy from diffusion model has samples and the
Theorem 2
step size for minimization is chosen as X ~~ Oly BL then with probability 1 5, 5 E (0, 1), the average Moreau envelope of the optimal inner maximization function <(w) == max, Er, [f x)]
satisfies
V1 Vn
V/2
VV (w)) S 43
+
+
(51)
VH
where w is the output ML parameter by Algorithm 7 o i the Moreau envelope of 6. A' is the error of inner maximization bounded in Theorem 7 Vi 833 (log(2/0) and V22 = / (w1) min, (ww)).
(w) == min ((w') + B1Iw - w/l1 is
Proof. The proof of Theorem 2 follows the techniques of Jin et al. (020)) with the difference that the maximization oracle an empirical approximation. The optimization variable is updated as wj = 0i-1 - V WwlEss, (f Define ((w) := maxę Ep, [f xr] as the inner maximization function and define (w)) == min ((w') + 2Φ Ilw w' 1|2 as the Moreau envelope of é.
Denote w; = arg min,, ((w)+ + B|[ wy 1| as the proximal point of the Moreau envelope e (w).
Since f is smooth, we have
3 11 wil 112
f(ūj, fw,, x) + (V v. (Wj, i6 w,)
(52)
Thus, it holds with probability at least 1 5 E (0, 1) that
f log(2/0)
((ūx)) 2 [f(w), = If,
Vn
3
w Folo((2/0)
2 Es, [f( x) (g(u, Sy), 30),00 16 - -w,)
[0
2
nn
(53)
[ū ww, 2 2 2ffllog (2/8)
2 Er, 2)] 2 (g(w,, 8,), ún ww,)
2
Vn
B
2 |0, w. 1| 2 12g(2/8)
2(w)) 2 + (g( Sy), ú w.)
vn
where the second inequality holds by applying McDiarmid's inequality on Ep, [f x)] and f is the upper bound of f. the third inequality holds by equation 52 and g(w,, S2) Es, (V (wy, x)],
22

---

Published as a conference paper at ICLR 2026
the forth inequality holds by applying McDiarmid's inequality Es, [s,]y( x)), and the last in- equality holds by Theorem
Now, it holds that
* (Wy+1)
-in ((w') + Bl1wf+1 - w'12
<ō(ū;) + Blw,, - >g(w;, S,) - úū
=b(ū;; + B||u, - w + 231 < g(w,. Sn), új - wj > +X +x Blg S,)11
(54)
وΦ (w,) + 23A < {g) 0j wy > (
A' + o وقاا wy + 2f /og (2//)
S 23 (w,) + 28A
- )lw,) ((1),)
A2
vn
where the last inequality holds by equation 53
Summing up inequality equation 54 from j 1 to j H, we have
(w)+1)
H
3
2 (VVoo(((8))
((w,)) ((0))) ((w)) + " I11ū wyl12 + <Φ (w1) 2,3
Σ
X2 B1H
28
n2
(55)
and we further have
x(w,) )(ū,) - 3 -- wsy
j=1
(56)
2f (2/0
2 (wi) min d(w)
XI,2
= +
+
Vn
23AH
Also, it holds that
8 - Aū wyll
6(w;) -0(ūr)
2
w;,1² 112 ((w) - =B1 w,51 5 || w
= -ww;) + Bl|u; - wz
(57)
2
3
2 || wy 1|2
v (w;))l2 433
where the inequality holds by the detinition of wj. and the last equality holds by the property of Moreau envelope such that 28 (w) w - wi for any w E W. Therefore, we have
(1))
(58)
2( 23 881 (2/8)
(wi)
min,, ((w))
1 H
<3 Σ
A
+ 2
XH
(wi) min. k(w)
By optimally choosing
have
BLAH
E
Vo
(1
(59)
83f log(2/0)
( ( (w,) min Φ(w))3
<48A'+
/
n
H
23

---

Published as a conference paper at ICLR 2026
C.3.1 EXPLANATION OF CONVERGENCE BY MOREAU ENVELOPE
The gradient bound of Moreau envelope indicates that the algorithm converges approxi- mately stationary point, which is explained as below. The Moreau envelope V 27
(wy satisfies
VØ (my)) = 2 (wg xūz). Since the proximal point is úi = arg min, Φ(w) + B - wyll² I1?
have V (1,) 2 wyy) = 0. Thus, it holds that V (w,) 23 (my viy) (iy 23
and || m0 21 Vo(w))l
(mz))l- Therefore, if the gradient of Moreau envelope V 2F VV (w)l S A, its proximal point w is an approx- iiately stationary point for the optimal inner maximization function (bounded gradient of the inner maximization function VV(()]l S A), and the distance between w and v is close enough:
is bounded for the decision variable w, i.e.
||ū - -wyl = z3 V (myy) $ 28 Thus, Algorithm T apppoxximately converges.
C.4 ANALYSIS ON THE EFFECTS OF THE CAPACITY OF GENERATIVE MODELS
Corollary 1. Let P = arg maxp Ef x)] be the worst-case distribution from the set of possi- ble testing distributions P given a variable w. Assume that the generative model Pa is expressive (Definition D. and is large enough such that there exists Ó which satisfies J(ό, Po) S in equa- tion| and (P P) < r. If the objective function f(w, x) is L Lippchitt continuous with respect to I, the expected score-matching loss is bounded as (() S J and the step size is chosen as 7 ~ O( 1 the overall error of inner maximization is
VR
1 maxfe, JHlp + L1T VK
À = E.ppfl x)] Ex [f(w, x) S
(60)
where Ly is the Lipschitz continuity parameter of fw, with respect to x. In addition, the inclusive KL-divergence with respect to the nominal distribution is bounded as
max(c, (e, JYMC - 1101 + DKL((r((T) C1.
(PoPia( sc
(61)
VK(pc- -C*
where () is initialized dual variable, H.c. 1. and Ci are constants, Pr is the output distribution of the forward process, and π is the initial distribution of the reverse process.
Proof. Let 6" be the optimal diffusion parameter that solves the inner maximization equation given the inner maximization error holds that
variable By Theorem 1
A P. [f(w, x)] E,E., (f x) S 1 maxfe, H1 VK
(62)
where the outer expectation is taken over the randomness of output selection. In addition, the inclu- sive KL-divergence with respect to the nominal distribution is bounded as
[xKt + + max(c, JNAC - (1 + R(T, Po) +
(63)
VK 11*)
That means the KL divergence with respect to the nominal distribution Po is bounded in the same way as Theorem [
remains to bound the overall error in equation 60 Since 0 satisfies J(Ø, Po) S c, within the ambiguity set defined in equation 7 so it holds that
Eg [f xr] 2 Ex [f x)).
(64)
Moreover, since Ø satisfies Dw (Pg. P) < r, it holds by Kantorovich- Rubinsstei Lemma (Kan- torovich & Rubinshtein (1958). Theorem 3.2 in Mohajerin Esfahani & Kuhn (2018) that
E If x)] E,- x)] < L.
(65)
Combing the above inequalities with the inner maximization error in equation 62 can bound the overall maximization error Á in equation 60
24

---

Published as a conference paper at ICLR 2026
Remark 1. Corollary 1 shows the effect of the expressive capacity of the generative model on the performance of GAS DRRC. When the expressivity parameter I' is smaller, the generative model has a stronger ability to approximate the possible testing distributions. In this case, the inner maximization in GAS-DRO can identify adversarial distributions whose objective value is closer to that of the true worst-case distribution P. Consequently, by optimizing the objective parameter wo with respect the parameterized adversarial distribution, GAS-DRO can achieve better robustness under the true worst-case distribution in the testing distribution set P. If the expressivity parameter of the generative model is O, the generative model approximate the possible testing distributions arbitrarily well. In this ideal case, the overall error bound of the inner maximization coincides with the inner maximization error bound in Theorem which asymptotically converges to zero as the number of inner iterations K increases.
D EXPERIMENTS ON TIME SERIES PREDICTION
D.1 DATASETS
The experiments are conducted based on Electricity Maps 2025). a widely utilized global platform that provides high-resolution spatio-temporal data electricity system operations, including carbon intensity (gCO)cq/Wh)) and energy mix, and is actively employed for carbon-aware aware scheduling and carbon footprint estimation in real systems such data centers.
We utilize datasets from Electricity Maps that record hourly electricity carbon intensity over the period 2021~2024 across four representative regions: California, United States Texas, United States (ERCO 21~244, Queensland, Australia (QLD 21~244, and the United King- dom (GB 21-244). These datasets capture fine-grained temporal variations in carbon intensity (mea- sured gCOzeq/kWh) arising from different energy mixes in diverse geographical and regulatory contexts, thereby providing comprehensive benchmark for evaluating carbon-aware forecasting models. For model training. we construct dataset by merging BANC_23 and BANC 24, resulting in 438 sequence samples. Model evaluation is then performed on multiple independent each consisting of 312 sequence samples drawn from other years and regions to ensure heteroge- neous and challenging testing scenarios. To quantify the distributional shift between the training and test sets, we compute the Wasserstein distance, which provides principled measure of discrepancy between probability distributions. The calculated distances are reported alongside the dataset names in Table
sets,
D.2 BASELINES
The baselines which are compared with in our experiments are introduced as below.
ML: This method trains the standard ML to minimize the time series forecasting error without DRO.
DML: DML represents an ML model fine-tuned with diffusion- diffusiion-ennrrated augmented datasets. Com- pared to GAS-DRO, DML performs standard training based on the augmented datasets rather than a distributionally robust training.
W-DRO: W-DRO represents a Wasserstein-basee DRO framework, in which the ambiguity set is characterized by the Wasserstein metric. For our experiments, we employ the FWDRO algorithm proposed in /Saii & Jegelka (017), which transforms the inner maximization of W-DRO into an ad- versarial optimization with mixed norm ball and then alternatively solves the adversarial examples and the ML weights.
KL-DRO: XL-DRO represents a KL-divergence-bar based DRO framework, in which the ambiguity set characterized by the KL- divergence. We employ the standard KL-DRG solution derived in Hu & Hong (2013). which is commonly adopted in practice.
DRAGEN: DRAGEN represents W-DRO framework, in which the ambiguity set is constructed Wasserstein ball in the latent space of a learned generative environment model(VAR Ren & Ma- jumdar (2022).
P-DRO: P-DRO represents a DRO framework that parameterizes the adversarial distributions using neural generative model within Kl--divergence-balt based ambiguity set. Michel et al. (2011)
25

---

Published as a conference paper at ICLR 2026
D.3 TRAINING SETUPS
The predictors in GAS-DRO and all the baselines share the same two-layer stacked LSTM architecture with 128 and 64 hidden neurons.
Predictor:
Diffusion Model: The diffusion models in GAS-DRR and P-DRO are both DDPM Ho et al. (2020) which has T 500 steps in a forward or a backward process. The U-Net architecture that the DDPM used consists of four downsampling layers with channel sizes of 128, 256, 256, and 256, respectively.
VAE: The VAE model used for DRAGEN is built from an encoder- decoder architecture in which both components contain two convolutional layers with 16 channels followed by two MLP layers. The encoder progressively downsamples the input image through stacked convolutions and outputs the mean and log-variance of a 2-dimensional latent distribution, after which a reparameterization step is applied. The decoder then maps the latent variables back to the image space through MLP projection followed by successive upsampling and convolutional layers.
Training: For GAS we adopt the pPO-based reformulation in equation 10 for inner maxi- mization. We train the reference DDPM 9g in equation I0 based on the original training dataset BANC 2324 BANC 23 & BANC 24) and it to generate initial dataset 20 calculate re in equation 10 The sampling variance of DDPM is chosen from a range [0.1, 0.5]. To improve training efficiency, only the last T 15 backward steps of the DDPM model are fine-tuned by equation I0 We choose a slightly higher clipping parameter K = 0.4 in equation to to encourage the maximization while maintaining stability. We choose = 0.015 as InnerMax's adversarial budget which gives the best average performance over all validation datasets. We choose as the rate update the Lagrangian parameter μ Algorithm [] We use the Adam optimizer with a learning rate of x 10-5 5 for both the diffusion training in the maximization and the predictor update in minimization. The diffusion model is trained for 10 inner epochs with a batch size of 64. The predictor is trained for 15 epochs with a batch size of 64.
0.01
For the baseline methods, we choose the same neural network architecture as GAS-DRO. We care- fully tuned the hyperparameters of the baseline algorithms to achieve optimal average performance over all validation datasets. For -DDRO, we consider the Wasserstein distance with respect to 12 12-norm and set the adversarial budget as 0.3 with learning rate of 5 x 10-4, For KL-DRO, we choose the adversarial budget e = 4 with a learning rate of 5 x 10-4. For DRAGEN, we adopt the Wasserstein distance with respect to l2 12-norm and set the adversarial budget
2 with learn- ing rate of 5 x 10-5 It is worth noting that, we pre-trained a well-performing VAE for DRAGEN with a learning rate of x 10-2 20-2, which achieved an average Wasserstein distance of 0.0103 on data reconstruction over the training set. For P-DRO, its generative model uses the exact same DDPM architecture as GAS-DRO. We set the temperature as r 0.1 and set the learning rate as 1 X 10-5 for the inner maximization and the outer minimization. In addition, P-DRO maintains a window of consecutive minibatches for updates of the Normalizer, and the window length is set as 5. Finally, same as GAS-DRC the diffusion model in P-DRO is trained for 10 inner epochs with a batch size of 64 while the predictor is trained for 15 epochs with a batch size of 64.
D.4 EXPERIMENTAL RESULTS ON TIME SERIES PREDICTION
D.4.1 MIAN RESULTS
We tune the hyperparameters of all methods based on validation datasets and present their best performance in Table [ where all the datasets are OOD testing datasets with different discrepancies.
We can find that all DRO methods improve the testing performance for OOD datasets compar- ing to ML. by optimizing the worst-case expected loss while DML improves upon ML by leveraging ceffisiion-generate
augmented datasets. Using ML as the reference baseline, GAS-DRO achieves the highest performance gain of 33,, followed by DRAGEN, P P- DDR and DML with improvements of 48.9 18.9%, 42.5% and 39.7%, respectively. In contrast, KL-DRO and W-DRO yield improvements of 36.1% 1% and 24.0%, respectively.
Notably, DML surpasses KL-DRO by 3.6% only after augmentation training with diffusion-generatee generated datasets, highlighting the significant positive impact of diffusion model. Moreover, DML differs from GAS-DRO only in that disables the DRO component, while their training procedures remain identi-
26

---

Published as a conference paper at ICLR 2026
cal. Thus, GAS-DRO, DML, and ML together form two ablation studies: the performance gap between GAS-DRO and DML confirms that DRO contributes a 39.7% performance gain to GAS-DRO while the gap between DML and ML verifies that the diffusion model also provides a 39.7% performance gain.
Although both DML and the other four DRO methods deliver noticeable performance gains, CAS-DRO still outperforms them by an average margin of 25.4% and by 27.8% in the worst case.
The underlying reason may lie in the fact that GAS-DRO leverages the distribution learning capabil- ity of the diffusion model when constructing the ambiguity set, enabling the generation of adversarial distributions that are both strong and realistic, thereby achieving superior OOD generalization.
In contrast, although DRAGEN also employs a generative model to construct its ambiguity set, it still relies on relaxing Wasserstein DRO framework which is limited in overcoming the inherent limitations of W-DRO. In addition, P-DRO, as a DRO algorithm that uses generative models to model the adversarial distributions, remains restricted by the KL-divergence-be based ambiguity set, preventing it from modeling diverse worst-case distributions. By contrast, GAS- -DR directly defines the ambiguity set on the parameterized space of generative models, enabling it to generate worst- case distributions that are both diverse and consistent with the nominal distribution, thereby offering a principled and tractable DRO algorithm.
Furthermore, the KL-divergence used in KL-DRO requires all distributions to be absolutely contin- uous with respect to the training distribution Po. which constrains the support space of the ambi- guity set and limits the search for strong adversaries. Among the baselines, W-DRO performs the worst, likely due to the fact that solving Wasserstein- based DRO usually involves optimal transport problems their relaxations, which are computationally more demanding and often rely on approx- imate methods such as dual reformulation or adversarial training. resulting in suboptimal solutions. Moreover, since W-DRO allows adversarial distributions with supports different from the training distribution, it is theoretically more flexible and closer to real distribution shifts, but this flexibility can produce overly extreme adversaries and thus lead to overly conservative model training.
D.4.2 EFFECT OF NOISY TYPES
Table 2: Gaussian- Corrupted Test.
In this test, we add certain amount of differ- ent types of noise into each test set to compare the noise robustness of different algorithms. The noise types include Gaussian Noise, Perlin Noise, and Cutout Noise.
MSE
Datacet
GAS-DRO
DRAGEN
KL-DRO
W-DRD
DEC.
ML.
BANC 22
0.0176
0.0188
0.0197
0.0183
0.0189
0.0291
BANC.21
0,0177
0.0199
0.0214
0.0216
0,0199
0.0337
QLD_24
0.0116
0.0911
0.0847
0.099
Gaussian Noise:
In the Gaussian Noise test, we
QLD 23
IL0615
0.0766
0.0925
0.0960
0.0934
0.1034
QLD
I.0377
0.0402
0.0476
0.0639
0.0445
0.0766
add Gaussian Noise with σ 0.1 to each test
QLD.21
0,0297
0.0291
0,0175
0.0649
0,0431
0,0772
set. As shown in Table 2 all algorithms ex- hibit performance degradation compared to the noise-free setting. Nevertheless, GAS-DRO still significantly outperforms the others: taking ML as the reference, our method achieves 48.3% improvement, followed by DRAGEN and DML with gains of 38.1% % and 31.7%, respectively. Among the remaining baselines, KL-DR4 im- proves performance by 29.1 19, while W-DRO is the weakest, surpassing ML by only 20%. In fact, W-DRO is relatively adept at handling Gaussian Noise compared to other noise types, and thus maintains noticeable advantage even under o = 0.1 Gaussian Noise, a trend further confirmed in subsequent experiments. On the other hand, both GAS-DRO and DML are trained with diffusion- generated augmented datasets, which naturally incorporate Gaussian- based noise during the diffu- sion process. This exposure enables them to remain robust under Gaussian noise perturbations. In contrast, DRAGEN 's performance, which falls between GAS-DRC and DML, can be attributed to the property of the VAE: although Gaussian noise in the input space appears random, is often compressed and filtered when mapped into the latent space, thereby preventing severe disruption of latent representations. Consequently, DRAGEN also performs reasonably well under Gaussian noise, though not strongly as GAS-DRO.
GB 24
0.0238
0.0233
0.0280
0.0276
0.0279
0.0381
GB
0.0221
0.2558
0.0305
0.0269
0.0414
GBt 22
0,0227
0.0063
0,0302
0.0343
0.0279
0.0158
0,0211
0.0256
0.0279
0.0334
0.0256
0.0438
ERCO 24
IU0281
0.0317
0.0350
0.0369
0.0340
0.0471
23
IJ0228
0.0251
0.0281
0.0294
0.0267
0.0418
ERCO,22
0,0206
0.0231
0:0256
0.0293
0,0250
0,0416
ERCO.21
0.0217
0.0244
0.0295
0.0357
0.0264
0.0475
Awerage
0.0182
0,0338
0.0373
0.0547
Worst
I.0615
0.0766
0.0925
0.0960
0.0934
0.1034
27

---

Published as a conference paper at ICLR 2026
Perlin Noise: Perlin Noise is smooth pseudo-random gradient noise commonly used to simulate natural textures such as clouds, terrains, and wood grains. By combining multiple Perlin Noise components with different frequencies and amplitudes (known as octaves), more complex fractal noise can be produced. In this experiment, we superimpose 8 layers of Perlin Noise, with the noise amplitude normalized to the range [-1, 1].
Table 3: Perlin- Celli-Corrupted Test.
As shown in Table 3 taking ML as the refer- ence, our method outperforms ML by 76.4%, while DRAGEN achieves a 66.3% improvement, and both perform better than in the Gaussian Noise test. DML consistently ranks third in per- formance, trailing DRAGEN by approximately 6.2%, a gap that is nearly identical to that ob- served in the Gaussian Noise test. Although KL-DRO also shows a larger gain relative to ML, its MSE remains roughly the same as in the Gaussian Noise test; the apparent improve- ment is primarily due to the substantial perfor-
MSR
Dutaset
GAS-DRC
DRAGEN
KL-DRO
W-DRO
XL
BANC 22
0,0117
0.0111
0.0296
0.0620
0,0227
0.0663
BANG,21
0.0156
0.0288
0.0591
0.0211
0.0662
QLD.24
0,0355
0.0506
0,0636
0.0862
0,0588
0.0928
QLD 23
I,0406
0,0400
0.0685
0.0787
0.0669
0.0860
QLD_22
0,0186
0.0321
0,0176
0.0852
0,034]
0.0001
QI.D_21
0,0192
0.0288
0.0475
0.0865
0,0355
0.0940
GB
0,0147
0.0122
0,0272
0.0499
0,0252
0,0563
GB 23
0,0133
0.0184
0.0295
0.0560
0,0239
0.0028
0.0223
0:0311
0.0611
0.0247
0,0666
GBt.21
0,0114
00067
0.0311
0,0586
0,0224
0.0077
24
0.0154
0.0224
0:0267
0.0436
0.0238
0.0525
ERCO,23
1,0142
0.0200
0.0382
0.0767
0,0262
0.0829
mance drop of ML in this experiment. In con- trast, W-DRO outperforms MT. by only 9.6%, a significant decline compared to its performance under Gaussian Noise, with the MSE reduced by as much as 50%. This indicates that W-DRO
ERCO.22
0,0109
0.0195
0.0324
0.0655
0,0223
0,0732
ERCO.21
0.0104
0.0211
0.0283
0.0506
0.0194
0.0005
Average
0
0.0245
0.0379
0.0687
0,0305
0,0727
Worst
0,0406
0.0506
0,0685
0,0865
0,0669
0.0940
is not well- for handling Perlin Noise, likely because the Wasserstein distance measures global distributional transport cost and is more effective in capturing smooth, small perturbations (e.g.. Gaussian Noise), but fails to adequately model the long-range correlated patterns of Perlin Noise within the Wasserstein ball.
Cutout Noise: Cutout Noise is a commonly used perturbation method that randomly selects region of the input data and sets its values to a constant, thereby simulating partial informa- tion loss. In our experiment, we randomly mask 30% of the sequence and set the masked val- ues to constant of 1. As shown in Table the performance of all algorithms is very close to their performance under Perlin Noise. We attribute this to the fact that, although Perlin and Cutout Noises differ in form, both represent structured local perturbations that disrupt the continuity of the input patterns, thereby posing similar challenges to all algorithms and result- ing in comparable performance under these two types of noise at given perturbation level. This is further confirmed in the subsequent gradient- perturbation tests.
Table 4: Cutout- Cotrut-Corrupted Test.
MSE
Dataset
GAS-DRO
DRAGEN
KL-DRO
M-DRO
DEL
ML
BANC 22
0,0063
0.0164
0,0181
0.0297
0,0125
0.0585
0.0173
0.0302
0.0547
0.0192
0.0637
QLD 24
IJ0404
0.044
0.0798
0.1096
0.0697
0.1148
QLD
0,0426
0.0614
0,0802
0.1028
0,0743
0.1092
QLD 22
0,0196
0.0354
0,0507
0.0881
0,0384
0.0971
QLD 21
0,0288
0.0294
0.0534
0.0916
0,0400
0.1005
GB
0,0145
0.0655
0.0351
0,0598
0.0252
0.0688
GB:.23
0,0122
0.0215
0.0343
0.0603
0,0250
0.068]
0.0129
0,0228
0:0361
0.0639
0.0245
0,0720
GB,21
n,0127
0.0225
0.0356
0.0651
0,0224
0,0740
ERCO 24
0.0157
0.0014
0:0375
0.0643
0,0275
0.0732
ERCO_33
0,0134
0.0225
0,0349
0.0599
0,0241
0.0678
ERCO 22
0.0134
0.0102
0.0349
0.0599
0.0241
0.0678
ERCO,21
0.0116
0.0227
0.0354
0.0671
0.0225
0.0758
Average
0,0174
0.0278
0.0:25
0.0699
0,0319
0,0782
Worst
0.0611
0,0802
0,0743
0.1148
D.4.3 EFFECT OF NOISY LEVELS
In this experiment, we progressively increased the intensity of three types of noise. For Gaus- sian Noise, the perturbation range set e [0.05, 0.2); for Perlin Noise. the amplitude is con- trolled within the range [0.05, 1]: and for Cutout Noise, the Cutout Mask Ratio is adjusted between [10%,40 40%). As shown in Figures 2 3 -AS-DRO consistently outperforms all baseline meth- ods across different noise types and intensity levels. With increasing noise strength, we observe that variations in Gaussian Noise have a more pronounced impact on all methods compared to Per- lin and Cutout Noise. In contrast, the impact of stronger Perlin and Cutout Noise on remains limited, and even at higher noise levels, GAS-DRO maintains stable and superior perfor- mance, highlighting its strong robustness. For W-DRO, however, the performance degradation under Perlin and Cutout Noise is much greater. with trends almost identical to ML, indicating that W-DRO is not effective in handling Perlin and Cutout Noise but is relatively better at coping with Gaus-
28

---

Published as a conference paper at ICLR 2026
sian Noise. Interestingly, when the Cutout Mask Ratio is 30%, the performance of all algorithms is nearly identical to their performance under Perlin Noise with amplitude 1, suggesting that at specific noise levels, Perlin and Cutout-thouge different in com--oot represent structured local perturba- tions that disrupt input continuity to a similar degree, thereby producing comparable impacts on the algorithms.
D.4.4 EFFECTS OF ADVERSARIAL BUDGET
In this adversarial budget experiment, we examine the impact of the budget parameter c in equa- tion the performance GAS-DRO. As illustrated in Fig. [ the loss-é curves across all datasets display concave trend, with the best average performance achieved around 0.015. When is smaller than this threshold, the diffusion-modeled distributions are overly restricted to training data, thereby hindering the ability of GAS-DRO to generalize to OOD datasets. In contrast, when e becomes excessively large, the enlarged ambiguity set causes GAS DRO to conservatively opti- mize against irrelevant distributions, which degrades its performance real OOD datasets. Hence, selecting an appropriate value of is essential for constructing effective adversarial distributions, ensuring a proper balance between average-case and worst-case performance.
D-DRO
KL-DRO
DML
D-DRO
KL-DRO
DML
DRAGEN
W-DRO
ML
DRAGEN
W-DRO
ML
0.06
0.06
MSE
MSE
0.04
0.04
0.02
0.02
0.050
0.075
0.100
0.125
0.150
0.175
0.200
1.0
Gaussian Noise Std (a)
Perlin Noise Amplitude
Figure 2: Gaussian perturbation strength
Figure 3: Perlin perturbation strength
D-DRO
KL-DRO
DML
BANC, 2324
QLD_212 GB_2124
ERCO 2124
DRAGEN
W-DRO
ML
BANC 2122
Average
0.14
0.10
0.03
0.08
MSE
MSE
0.02
0.04
0.01
0.02
10
15
20
25
30
35
0.005
0.010
0.015
0.020
0.025
Culout Mask Ratio (%)
Figure 4: Cutout perturbation strength
Figure Effect of budget in CAS-DRC
D.4.5 EFFECT OF THE CAPACITY OF GENERATIVE MODELS
To investigate the influence of diffusion model capacity on the performance of GAS-DRO. we vary the width of the UNet by modifying the channel dimensions at each level. Specifically, we denote the original four-level UNet with channels [128, 256, 256, 256] as DFU-Full. We then scale the model down to [96. 192, 192. 192], denoted DFU-3 3 and further to [64, 128, 128, 128), denoted 4:
Table 5: Comparison of GAS-DRO with Dif- ferent Generative Models.
Model
Avg MSE
Worst MSE
DFU-Full
0.0163
0.0509
DFU-2
0.0171
0.0544
DFU-1 In addition, include VAE-based variant of CAS-DRO for comparison, and the VAE architecture is given in Section D.3
DFU-
0.0178
0.0608
VAE
0.0168
0.0463
To evaluate the effect of generative model choice, we evaluate the four generative model configurations-thee diffusion models of different UNet sizes and one VAE-on all original (noise-
29

---

Published as a conference paper at ICLR 2026
free) OOD datasets in the time-series forecasting task. The results in Table 5 show that, for the diffusion-basee variants GAS-DRO, increasing the UNet width leads to a clear improvement: the average MSE decreases monotonically, and the worst-case MSE is also substantially reduced the model becomes larger. This indicates that a stronger generative model can better characterize the ambiguity set, thereby enhancing both the robustness and accuracy of GAS-DRO. The VAE-based variant achieves performance between DFU-Full and DFU- 4. further illustrating that GAS-DRO does not depend on diffusion models specifically; rather, it is insensitive the choice of generative architecture and exhibits strong generalization across model types.
D.4.6 COMPUTATIONAL OVERHEAD
Finally, we tested all methods on the time series prediction task, specifically evaluating their runtime, memory overhead, and GPU memory overhead during both the training and inference phases. All methods are evaluated on single NVIDIA RTX 6000 Ada GPU, with 48 GB of memory, and TB of system storage.
The training computational overhead is shown in Table 6 GAS-DRO (DFU) represents GAS-DRO using the diffusson model, GAS-DRO (VAE) repre- sents GAS-DRO using the VAE model, and other baselines have their names in Section D.2
Table 6: Training Cost Comparison.
Method
Time (s)
Memory (MB)
GPU Memory (MB)
GAS-DRO (DFU)
1897
2274
12239
GAS-DRO (VAE)
494
2449
5849
DRAGEN
2191
2355
434
P-DRO
1799
2205
11576
KL-DRO
149
2314
417
As shown in Table 6 DML has a higher memory cost than MI due to the need to store the augmented dataset. Compared to ML and DML, GAS-DRO (DFU) and P-DRO both introduce significant additional training time overhead primarily due to sampling processes diffusion models. The memory overhead of GAS and P-DRO is also substantially increased due to the large size of the U-Net and the need to store sampling tra- jectories. However, the higher computational overhead of GAS-DRO (DFU) and P-DRO traded for superior generalization performance across OOD testing datasets, particularly for GAS GAS-DRO, as shown in Table ! GAS-DRO (VAE) and DRAGEN are based on VAEs, so they have less memory footprint. However, due to the complexity to search within the Wasserstein-basse ambiguity set, they have the highest time overhead. KL-DRO is the most efficient DRO algorithm in terms of time and memory due to its efficient closed-form solution.
W-DRO
2896
2304
417
DML
130
2316
417
143
1352
33
In test phase, however, since all methods use the same ML models for time-series prediction, they have the same runtime and memory for inference. To perform inference on a testing dataset (BANC_22) with 8760 samples, the inference runtime is 0. and GPU memory is 409 MB.
E EXPERIMENTS ON IMAGE CLASSIFICATION
E.
TASK
We further compare the performance of GAS DRO with baseline methods on an image classification task. Handwritten digit classification is a moand-stnding and fundamental benchmark in computer vision. It captures the essential structure of real-world image recognition tasks, including high- dimensional pixel dependencies, sensitivity to corruptions, and vulnerability to distributional shifts. As such, the handwritten digit classification has been a standard testbed for evaluating the robustness and generalization ability of machine learning models. Beyond that, handwritten digit recognition supports many real-world applications, including postal mail sorting, bank check processing, form digitization, and optical character recognition (OCR) systems.
In this experiment, we adopt handwritten digit recognition to assess whether GAS-DRO can effec- tively handle distribution shifts and image corruptions in vision domains. The task is particularly relevant for DRO evaluation because even small perturbations such as noise, occlusions, or con- trast changes- can significantly alter the underlying data distribution. A robust DRO method should
30

---

Published as a conference paper at ICLR 2026
therefore maintain high accuracy under these shifts. This makes handwritten digit classification important and well-defined environment for comparing different designs and validating the perfor- mance of GAS-DRO.
E.2 DATASETS
This experiment a standard image classification task, where the training set is the MNIST dataset Lecun et al. (1998), and the evaluation is conducted on both MNIST and an OOD dataset USPS Hull (994) test sets. The USPS dataset, collected by the U.S. Postal Service for automatic mail sorting, serves as a natural OOD benchmark due to its distinct handwriting styles and distributional characteristics. In addition to the original test sets, we construct more OOD datasets by applying various perturbations both MNIST and USPS, including Gaussian noise, Perlin noise, and Cutout noise, consistent with the noise types used in our time-series forecasting experiments. As shown in Table for MNIST, the corrupted variants include M-Gaussian, M-Perlin, and M-Cutout, corresponding to Gaussian noise with a standard deviation of 0.4, Perlin noise with amplitude 2.5, and Cutout with a cover rate of 0.3, respectively. For USPS, we adopt milder perturbations due to its inherent OOD nature, resulting in U-Gaussian, U-Perlin, and U-Cutout, which come with Gaussian noise with a standard deviation of 0.03, Perlin noise with amplitude 0.5, and Cutout with cover rate of 0.05, respectively. Both M-Original and U-Original mean the original dataset of MNIST and USPS. The training is performed on M-Original, and the classifiers are tested on other OOD datasets. We select the first 4096 images from the MNIST training dataset as the training set for classifier training in algorithms. For evaluation, we use the original MNIST test set consisting of 10000 images, together with 7291 images from the USPS dataset. All corrupted test sets used in our experiments are generated by applying various perturbations to these 17291 images.
During data preprocessing, mean of each pixel 0.1307 and variance of each pixel is 0.3081. For the USPS dataset, we use the same data preprocessing method as Ren & Majumdar (2022): Each original 16 x 16 grayscale digit image is reconstructed from the HDF5 file, resized to 30 x 30, and then centered on a 48 x 48 black background to match the spatial layout of MNIST-like inputs. The images are subsequently downsampled to a 14 x 14 resolution, paired with their corresponding labels for training and evalu- ation. Since the classifier used in this experiment datasets are ultimately downsampled to a 14 x 14 resolution to facilitate efficient training and testing.
apply the standard normalization the MNIST dataset so that the
lightweight two-layer DNN, all images across
E.3 BASELINES
The baselines which are compared with GAS-DRO in our vision experiments introduced as below.
ML: This method trains the standard ML model to minimize the classification error without using DRO.
DML: DML represents MI. model fine-tuned with data augmented training Unlike the DML, setting described in Appendix [D. here we apply Gaussian noise with a standard deviation of 0.02 to augment the training set for data augmentation.
W-DRO: We employ the same implementation of WDDRO described in Appendix D with only minor modifications to its input and output formats accommodate training and inference on image data.
KL-DRO: We employ the same implementation of KL-DRO described in Appendix D with only minor modifications its input and output formats accommodate training and inference on image data.
DRAGEN: We employ the same implementation of DRAGEN described in Appendix D with only minor modifications to its input and output formats to accommodate training and inference on image data.
E.4 TRAINING SETUPS
Predictor:
The classifiers in GAS-DRO and all the baselines share the same two-layer DNN archi- tecture with hidden layer of 8 neurons followed by a 10-dimensional output layer. We use MSE
31

---

Published as a conference paper at ICLR 2026
instead of cross-entropy while training the predictor, because the predictor is very small two-layer network, and MSE provides more stable gradients.
VAE For DRAGEN & GAS-DRO: In this experiment, both DRAGEN and GAS-DRO use the same VAE, and its architecture is identical to the VAE employed by DRAGEN in the time-series forecasting task in Section D.3] The only difference is that we increase the latent dimensionality to 8 in order to preserve more image information. Specifically, the encoder is a convolutional neural network that progressively downsamples the input image and maps it into a now-dimensional latent representation. It first applies a stack of 2 convolutional layers, each with a stride of and ReLU activation, where the number of channels doubles at each layer (from 1 an initial low-width channel count and then up to the specified inner channel width). After spatial downsampling, the resulting feature map is flattened and passed through an MI.P. This MLP outputs the parameters of an 8-dimensional latent Gaussian distribution, producing both the mean and log-variance vectors through two separate linear heads.
The decoder first maps the latent vector into low-resolution convolutional feature map using an MLP, effectively expanding the compressed latent representation back into a spatial structure. This feature map is then progressively upsampled through sequence of layers, where each layer doubles the spatial resolution using interpolation and refines the features with a 3x3 convolution followed by ReLU activation. As the resolution increases, the number of channels is gradually reduced to match the structure of natural image. Finally, a 1x1 convolution projects the features into a single output channel, and Sigmoid activation constrains the pixel intensities the [0,1] range. Based on the decoder output, we can model a Gaussian distribution for the image as (=]2)) = N (z; po(), d²1) where ( is the decoder network. Given a pretrained encoder (2(2(zr) the reconstruction loss of VAE JVAE = Bro(s) .) log pe (x 2)] is in a linear relationship with the expected squared error based on the latent distribution, i.e. IE Po)) (x)()) [||0 [=)] Thus, we directly use this expected squared loss to replace the VAE reconstruction loss in equation 25 and obtain the
inner-maximization objective following equation 27
Training For GAS-DRO: Before training GAS-DRO, we freeze the parameters yp of the pretrained VAE encoder. We choose 65 as InnerMax's adversarial budget which gives the best average performance over all validation datasets. We choose = 0.1 as the rate to update the Lagrangian parameter y in Algorithm[] We use the Adam optimizer with a learning rate of 10-4 for the VAE training in the maximization and a learning rate of x 10 for the predictor update in minimization. The VAE model trained for 10 inner epochs with a batch size of 64. The predictor is trained for 10 epochs with a batch size of 64. During the predictor training, we encode the training data into latent variable z using the VAE encoder. Then we introduce a small perturbation to the latent space with the perturbation radius as e, 0.25. Next, we generate new samples by decoding the latent variable Z with the updated decoder Pp-
Training For DRAGEN; For DRAGEN, we adopt the Wasserstein distance with respect to 4 -orm and set the adversarial budget as e = 1.5. Then a W-DRO is applied on the latent space with a learning rate of 5 10
Training For W-DRO: For W- W-DRO, we consider the Wasserstein distance with respect to L2 norm and set the adversarial budget 3 with a learning rate of 1 10-5
Training For KL-DRO: For KL-DRO, we choose the adversarial budget c = with a learning rate of 5 10
Training For DML: The training configuration of DML is identical to that of ML; the only difference is that we replace the original training samples with their Gaussian-noised versions, where noise with a standard deviation of 0.02 is added to each image.
All methods are trained for total of 100 epochs; for GAS CDS-DRO that adopt nested training scheme, the product of the numbers of outer and inner iterations is also set to 100.
E.5 EXPERIMENTAL RESULTS ON IMAGE CLASSIFICATION
We conducted comparative tests on the MNIST and USPS datasets. As shown in Table GAS-DRO achieves the best overall performance with an average accuracy of 71.52%, outperforming DRAGEN by 3.27%, KL-DRO by 6.24%, and standard MI. by 9.49% demonstrating the strongest generaliza-
32

---

Published as a conference paper at ICLR 2026
tion and robustness across all distributions and corruption types. GAS-DRO also achieves the best worst-case performance with an accuracy of 50.58%, clearly outperforming DRAGEN (46.84%), KL-DRO (45.38%), W-DRO (46.11%), DML (43.08%) and ML (42.38%), making it the only method that maintains above 50% accuracy under the most challenging OOD conditions. Across the OOD testing settings, GAS-DRO achieves the best near-hest accuracy on MNIST and consistently ranks first in all USPS OOD scenarios, demonstrating its superiority under stronger distribution shifts.
Compared with DML, which bene- fits only from data augmentation. GAS-DRO leverages both generative modeling and worst-case optimiza- tion, yielding significantly superior OOD generalization. More impor- tantly, we note that the superiority of GAS-DRO persists even when it is implemented with the same VAE backbone used by DRAGEN, which demonstrates that the advantage of GAS-DRO rely on the design of its ambiguity set and optimization ap- proach. Unlike DRAGEN, which still relaxes traditional Wasserstein DRO framework to perturb the latent space of VAE, forms inner maximization directly in the parameterized generative model space, enabling more stable convergence and the construction of more realistic adversarial distributions. Moreover, the reconstruction-loss based ambiguity set used in CAS-DRO avoids the issues in traditional ambiguity modeling, including the restrictions on the support space inherent to KL-DRO and the overly conser- vative adversaries introduced by relaxing W-DRO. Thus, outperforms the traditional DRO methods and the generative DRO method based traditional ambiguity modeling.
Table 7: Image Classification Test.
Accuraey (5)
Dataset
GAS-DRO
DRAGEN
KL-DRO
W-DRO
DML
ML
M-Original
85.31
82.71
86.92
85.70
79.02
85.63
M-Gaussian
83.13
81.23
70,26
74.16
81.70
65.89
M-Perlia
81.52
81.57
71.74
75.11
79.68
67.15
M-Cutout
62.10
65.22
59.62
62.86
61.28
59.31
U-Original
71.81
64.09
63.87
65.49
63.05
60.22
U-Gaussian
70.81
63.49
62.23
63.98
62,94
58.62
U-Perlin
66.86
60.83
62.24
62.95
59.51
57.07
U-Cutout
50.58
46.84
45,38
46.11
42,08
42.38
Average
71.52
68.25
65,28
67.05
66.16
62.03
Worst
50.58
46.84
45.38
46.11
42,08
42.38
per-
F SCORE-MATCHING BASED DIFFUSION MODELS
The score-based diffusion models learn the distributions through forward process and a back- ward process. We introduce boore-based generative modeling by Stochastic Differential Equations (SDEs) Song et al. (2021).
Forward Process. The forward process incrementally injects noise into the data, generating a quence of perturbed samples. begins with an initial sample Xo E Rd drawn from a training dataset Su, and evolves according to SDE defined as:
dr = 6(x, t)dt + g(t)dw
(66)
where b(.,t) Rd - Rd is drift coefficient, g(t) E R, and w is standard Wiener process. By the SDE, we get variable v, which represents the data corrupted by noise at time L. We use Pt to represent the distribution of x, and P. to denote the conditional distribution of x; given Tn- In the framework of score-based diffusion models, the forward diffusion process terminates at a sufficiently large time T. so that the marginal distribution Pr approximates a tractable distribution p' which is typically chosen as the standard Gaussian distribution N(0,, (O)
Backward Process. By reversing the forward process in time, it's possible to recover the original data distribution starting from pure noise. The backward SDE reverses the time evolution of the forward equation in equation 66
(67) where w is a standard Wiener process in the reverse-time direction, V, log (x( is the time- dependent score function relying on the marginal probability density of Iz in the forward process.
dr t) - g(t)² log Pt(x)) dt + g(t)dw
Score Matching. In the backward SDE, the score function V log P.(-) plays critical role in di- recting the generative dynamics. To estimate the score function V log (.).), we train score- based model so (x. t) based on samples generated from the forward diffusion process. The score-based model is optimized by minimizing the following denoising score-mas more-matching loss given a dataset Sy:
((0,So) [" Brow Ena [u(t) ||v log Pro xo)) - so(ar, -s(mxr,))] 1)]]] dt
(68)
33

---

Published as a conference paper at ICLR 2026
where ((t) > 0 is weighting function.
F.1 BOUND OF THE -AACHIGG LOSS BY KL DIVERGENCE
Lemma 3. Given the assumptions in Appendix F. if the score-matching loss satisfies JO,,
JDSM (0, (ri)==)) S e, the output distribution of the diffusion model Po satisfies
Poo(P)))Pa S < E
where Pr is the forward process distribution step T. pf is the prior distribution (by design Pr P p'), ), and C1 is a constant that does not depend on as below
C G-= [==""
( (0
(69)
1=
where (Iu, = (o 1 2, log (1 - 21, log Pajo mu)ll
To prove this lemma, we make the following assumptions for the SDE-based diffusion model in Section
Assumptions
1. (o(x) is a density function with continuous second-order derivatives and Ee~ [|]=1]]] < 20.
2. The prior distribution p' (x) is a density function with continuous second-order derivatives and Enwyy [[]=-[[]] oc.
3. Vt [0,T]: b(.,t) t) is function with continuous first order derivatives. 3C 0, Vr
Rd, t E []: [16(z, t)ll2 I C(1 + ||
4. 3C 0, Va,y Rd [|(,, () - by l)]l2 S Cl| - v/2-
5. g continuous function and A/ [0, [r()]
6. For any open bounded set O, fo
7. 3C > 0, Vr Rd, E [0, | (()) C(1 (|
8. 30 0, yr, Rd: (()) Po s C(| vl12).
9. 30 0, Vr Rd, f [0.
C(+ |
10. 3C > 0, VI, Rd [|5 t) so 12 - C(Ix= - yl12).
11. Novikov's condition: E |exp( f I log (.(x)) s0 1) < co.
Given the assumptions, the following lemma is re-statement of Theorem in Song et al. (0221) Lemma 4. Let Po(x) be the underlining data distribution, (zx be a known prior distribution, and Po be marginal distribution of (0), the output of reverse-time SDE defined as below.
dā
[b( 1) r(1)3 (0 1)]di r(I)du, (o(T) x/
(70)
With the assumptions in Section
F.
we have
(P s (sM r(-)²) + D
(71)
where JSM JSM (8,r())) 1 Em((a)) [r(t]l log Pz( sy (z, 1 dt.
Proof. We denote the path measure the forward outputs {xx ) as p and the path measure of the backward outputs (ejorr as q. By assumptions 1-5, 9, 10, both and y are uniquely given by the forward and backward SDEs, respectively. Consider Markov kernel ((2e)reeoon) v) 6(zo v) given any Markov chain (zeeeo Since xo Po and $ Pe. we have
/ M =
x)dp(( Po(x)
(72)
34

---

Published as a conference paper at ICLR 2026
/ M M(jaaa.eeeuo
Po(x))
(73)
Here the Markov kernel M essentially performs marginalization of path measures to obtain distri- butions at We can use the data processing inequality with this Markov kernel to obtain DKL(Polled)
=DKL (MM(facheto.T) x)do((= /
(
(74)
<DKI(P. Since 2r Pr and žogr ~ p'. Leveraging the chain rule of KL. divergence, we have
p(11T
Pr
DKL(P, = Ex log( 7 ((0,,11: &0,r) (
(75)
DKL((PTT Ilp'') + Eg [DxL (p( .t = =)(l(( दे.т = 2))]
Under assumptions 1,3-8, the SDE in Eqn. equation 66 has a corresponding reverse-time SDE as dr (b () - r(t) log (x)) + r
(76)
Since Eqn. equation 76 is the time reversal of Eqn. equation 66 they share the same path measure p. Thus, Pr [D (p(- 2y = =) å
=))] can be viewed as the KI. divergence between
the path measures induced by the two SDEs in Eqn. equation 66 and Eqn. equation 70 with the same starting points IT = Σ.
The KL divergence between SDEs with shared diffusion coefficients and starting points exists under assumptions 7, and can be bounded by the Girsanov theorem
dp DKL(p(- xT =)]l((- (0. =)) EE, log dg
1 1'
( log Pa()) so(x,t))duu dut =
J.
r(t) log P((x) so(x, so(x,t)līdt
Ep
V V. - s0(x, clt
2 1=0
[" 2
Bx(()) [r log Px(x) de s0 [ (0,((,r()))
(77)
where the second equality holds by Girsanov Theorem II, and the third equality holds because Y, SC (V= log (z( so (a. dd); is continuous-time Martingale process (B Yr,rgo'] Ysr, Vs' < 8) and we have E,,, Y= = 0, Vs <8.
Proof. The score matching loss (0, in equation 68 with i(t) = g())² is actually the denoising score matching loss J (0, i..))),
1 T
So) = JDSM g(-))
EN
|| [(())2|Va, log Pip mo) - so(x,, 1) |1] de
2
(78)
which is used to compute the original score matching loss JsMr g()²)) given dataset So. The gap between JDSM (0, g()) and JsM (0, g())²) is constant C1 that does not depend on 0, which shown as below.
The difference is expressed as
JSM .())) (0.M(0.9(.)3
1
E EPo.SE (9(1) (I| log P.(xt) - s0 (x 1) - I|V., log Pejo(ae xo) - s0(x,, 1)112
2
T
g()2 Epoa(zonN)) (80( 1), Vz, log (((x)) V., log Pejo(are mu))]] (Tro dl,
(1)
(79)
35

---

Published as a conference paper at ICLR 2026
where
Po.e is the joint
distribution of zo and 2t. and (1 V., log (1 xo)
E Bro ((roarr))
2x, x, log (z)))
The first term (1) is zero because
E (ao.ar) ((50 t), Va, log P((xs)) = Br((a)) (so(Fe, t), Va, log Pt(a:)))
1
sa( t) Pa(xt)
Vn V., Pc(ot)) (1(
so(ax, 1), V.. / Pao(ex (00 Po)oo dr
(80)
so( 1), / Prola Po)oo) v,, (a ro))dro) ) dx
Po,e Te) (50 (), Va, log( (a (n))) dundde,
20.21
= (Ro [( 1), V., log P mo))]
Thus, can bound the gap between J /DM ( (-))²) and (0, as
/" g(t) xz)dt
C1 = (=
(81)
-=
where C2 21) = = [ 2,, log 12
114 V log Pajo(re x0))
Therefore, if S0) JnSM (0, e, by Lemma 4 we have
(kL(Pol(Po) SJs (0, g(.))) (Kr(Pr
JDSM (0, (Pr(ll)))') +
(82)
Sc (Pr C2
which completes the proof.
G THE USE OF LARGE LANGUAGE MODELS (LLMs)
LLMSs were used as general-purper assistive tools to improve the clarity of English writing and to check grammar. No parts of the technical contributions, experimental design, or theoretical results were produced directly by LLMs. All mathematical derivations, proofs, algorithmss and final writing decisions were carried out by the authors. The use of LLMs was limited to supportive roles and did not constitute substantive contribution to the research itself.
36


Reviews:
[{"review_id": "Xc0URaz4OZ", "reviewer": "Reviewer_u62m", "summary": "This paper proposes GAS-DRO. It uses diffusion models and VAEs to build ambiguity sets for DRO. The method constrains the reconstruction loss of generative models instead of using KL divergence or Wasserstein distance. The authors provide convergence guarantees and test on time series forecasting tasks.", "questions": "See the weakness", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper identifies real problems with existing DRO methods. φ-divergence restricts support. Wasserstein distance is hard to optimize.\n\n2. Theorems 1 and 2 provide convergence guarantees. Lemma 1 connects reconstruction loss to KL divergence.\n\n3. GAS-DRO achieves 63.7% improvement over baseline ML. It outperforms existing DRO methods by significant margins.", "weaknesses": "1. Paper claims \"first to model ambiguity sets in parameterized space of likelihood-based generative models\" (page 2). This is incorrect. Michel et al. (2021) \"Modeling the Second Player in Distributionally Robust Optimization\" (ICLR 2021) already proposed Parametric DRO (P-DRO) that: (a) uses neural generative models q_ψ for adversary, (b) parameterizes uncertainty set with model weights ψ, (c) uses likelihood-based Transformers (evaluates log q_ψ(x,y) in Equation 8), (d) solves same min-max game structure, (e) tests on NLP tasks. The paper completely omits this citation. The distinction between \"implicit\" (GAN) and \"explicit\" (diffusion) generative models is insufficient—Michel et al. already used likelihood-based models. After removing false claims, only incremental contributions remain: using diffusion instead of Transformers, reconstruction loss instead of KL constraint, and time series application. No direct comparison with P-DRO is provided.\n\n2. Limited experimental scope. Only one domain (time series forecasting). Michel et al. tested sentiment classification and toxicity detection. Need broader evaluation across vision or NLP tasks.\n\n3. No computational cost analysis. Diffusion models are expensive. No runtime, memory, or scalability comparisons provided.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GCwFNryCNc", "reviewer": "Reviewer_Focz", "summary": "This paper proposes a novel approach to Distributionally Robust Optimization (DRO) that leverages generative models to define the ambiguity set of distributions. Instead of using a divergence or Wasserstein ball, the ambiguity set is parameterized by a generative model (e.g. a diffusion model or VAE) which can produce adversarial distributions beyond the support of the nominal data while still remaining similar to it. The authors introduce an algorithm called GAS-DRO (Generative Ambiguity Set DRO) that alternates between an inner loop (updating the generative model to find a worst-case distribution) and an outer loop (updating the primary model/decision variable on data sampled from that worst-case distribution). They enforce a constraint (via a dual Lagrange multiplier) to keep the generative distribution within a certain “distance” (ambiguity budget *ε*) of the nominal distribution. The paper provides theoretical guarantees – proving that the inner maximization converges to the optimal worst-case distribution and that the overall iterative procedure converges to a stationary point. Empirically, GAS-DRO is implemented with a diffusion model as the adversary and tested on a time-series forecasting task under distribution shift (electricity grid emissions data). It achieves significant improvements in out-of-distribution (OOD) generalization, outperforming several baselines like standard DRO with KL or Wasserstein ambiguity sets and a recent generative DRO method (DRAGEN). In summary, the paper’s main contributions are the introduction of generative-model-based ambiguity sets for DRO, a corresponding optimization algorithm with convergence analysis, and demonstration of improved robustness on real-world data.", "questions": "Please see Weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper introduces a generative-model-based ambiguity set for DRO, which allows considering distributions outside the original support while maintaining similarity to the nominal distribution. The approach does not require too much prior condition and addresses the limitations of KL-divergence ambiguity (no support shift) and improves over Wasserstein ambiguity by providing a tractable, parameterized search space. \n2. Empirical OOD Performance Gains: GAS-DRO demonstrates state-of-the-art OOD generalization performance on the Electricity Maps time-series datasets.\n3. The writing is straightforward and easy to understand.", "weaknesses": "1. This paper investigates DRO, which is widespread considered as an approach for discriminative tasks. The method in the paper instead uses a generative method to address the classification problem but determining ambiguity set itself needs the sampling operation. Therefore, I wonder how much extra compuational cost this method has introduced compared to normal plug-and-play baselines? For instance, how much longer (in terms of training time or iterations) does GAS-DRO take compared to standard DRO or ERM on the same task?\n2. How sensitive is GAS-DRO to the choice and quality of the generative model? For example, if one used a smaller diffusion model or a different type of generator (like a Variational Autoencoder, as mentioned, or even a GAN), would the method still perform as well? Any insight into how approximation errors in the generative model impact the robustness of the learned model would be valuable.\n3. The empirical evaluation, while thorough for the chosen task, is focused on a single domain (renewable energy time-series forecasting on Electricity Maps data). It would strengthen the paper to see results in other domains to confirm that GAS-DRO’s benefits hold broadly. The current evidence of superior performance is strong for the scenario presented, but the approach's generality would be more convincing with a wider range of experiments.\n4. Maybe adding more discussions on the use of generative modeling in common classification or OOD problems will be better. For example, a diffusion classifier [1] is adopted for image classification. Although authors have presented related work on IID settings, the description of OOD is somewhat lacking [2, 3].\n\n[1] Li, Alexander C., et al. \"Your diffusion model is secretly a zero-shot classifier.\" *Proceedings of the IEEE/CVF International Conference on Computer Vision*. 2023.\n[2] Tong, Yunze, et al. \"Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data.\" *Forty-second International Conference on Machine Learning*.\n[3] Zhang, Hengrui, et al. \"Mixed-type tabular data synthesis with score-based diffusion in latent space.\" *arXiv preprint arXiv:2310.09656* (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1R2onxXs8s", "reviewer": "Reviewer_WYi2", "summary": "The paper addresses the problem of distributionally robust optimization for enhancing robustness in statistical learning under out-of-distribution scenarios. The main idea is to construct a novel ambiguity set using likelihood-based generative models, such as diffusion models or VAEs, which allows for distributions that are consistent with the nominal distribution while enabling support shifts and tractable solutions. Key results include the GAS-DRO algorithm (Algorithm 1), which solves the minimax problem via dual learning and policy optimization. Theoretical contributions comprise Theorem 1, establishing convergence of the inner maximization to the optimal oracle with bounded KL divergence to the nominal distribution, and Theorem 2, proving stationary convergence of GAS-DRO under smoothness and Lipschitz assumptions. Empirical results on Electricity Maps datasets show superior OOD generalization compared to baselines like W-DRO, KL-DRO, and DRAGEN.\n\n## Clarity\n\nThe paper is very well-written and clearly structured. The introduction provides excellent motivation by clearly outlining the limitations of existing DRO frameworks. The proposed method (GAS) is developed logically, building from the problem setup (Sec 3) to the generative model formulationand the algorithm.\n\n## Originality / Novelty\n\nThis work is highly original. The central idea of defining a DRO ambiguity set using a generative model's reconstruction loss is, to my knowledge, new.", "questions": "Lemma 1 is doing a lot of work: it’s the bridge from small reconstruction loss” to bounded $D_{\\mathrm{KL}}(P_0 | P_\\theta)$.\nBut in the main text, the statement is fairly high-level:\n\n    - It asserts $D_{\\mathrm{KL}}(P_0 | P_\\theta) \\le J(\\theta,P_0)+R(p',P_0)+C_1$, where $R(p',P_0)$ depends on the prior and $C_1$ depends on $-H(P_0)$.\n\n    - It then argues that this allows support shift (since it’s the *reverse* KL, which only needs $P_0 \\ll P_\\theta)$. \nHowever, several details are pushed to Appendix C.1 and are not visible in the main body we have:\n\n- Are there assumptions on model capacity (i.e., is $P_\\theta$ assumed expressive enough to approximate $P_0$ arbitrarily well)?\n- Are there bounds on $R(p',P_0)$ in practice, especially for diffusion, where $p'$ is a fixed Gaussian prior?\n- Is $C_1$ finite when $P_0$ is empirical with finite support (so $-H(P_0)$ is well-defined)?\n  These questions matter because if $R(p',P_0)$ is large or uncontrolled, then “bounded reconstruction loss ⇒ bounded KL” may not actually constrain the adversary meaningfully. Right now, the story is intuitive but not fully airtight in the visible text.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The core idea of formulating the ambiguity set over the parameters of a generative model and, crucially, using the *reconstruction loss* $J(\\theta, P_0) \\le \\epsilon$ as the constraint (Eq. 7) is a highly novel and elegant contribution. It reframes the problem from an intractable search over distributions to a tractable optimization over model parameters.\n  \n- The method provides a compelling solution to the fundamental tension between $\\phi$-divergence DRO (no support shift) and Wasserstein-DRO (intractable/conservative). The use of generative models inherently allows for sampling beyond the nominal support space, while the constraint $J(\\theta, P_0)$ ensures consistency with the nominal distribution.\n  \n- The paper proposes a concrete, end-to-end algorithm (Algorithm 1). The application of dual learning and policy optimization to solve the complex inner loop is a clever technical contribution that makes the framework practical.\n  \n- The authors provide non-trivial convergence guarantees for their proposed algorithm. Theorem 1 establishes the convergence of the inner-loop maximization oracle, and Theorem 2 proves stationary convergence for the full min-max procedure, lending strong theoretical weight to the method.\n  \n- The experiments, summarized in Table 1, show a clear and substantial performance improvement for GAS-DRO over all baselines. It achieves the lowest average error (0.0163) by a significant margin over the next-best baseline, DRAGEN (0.0230), and standard ML (0.0450). The ablation studies in Appendix D.2 are also thorough.", "weaknesses": "- The theoretical link (Lemma 1) between the proposed constraint $J(\\theta, P_0) \\le \\epsilon$ and traditional divergence-based ambiguity sets is a one-way bound ($J \\le \\epsilon \\implies D_{KL}(P_0 || P_\\theta)$ is bounded). This ensures distributions in GAS are \"consistent,\" but it's not clear if the set is sufficiently expressive. It's possible that a \"worst-case\" distribution $P^*$ that is close to $P_0$ (in a standard metric) might not be representable by any $P_\\theta$ in the GAS, even with a powerful generative model. The paper relies on the \"strong capability\" of generative models, but this is not formally characterized.\n  \n- The analysis assumes bounded reconstruction loss and specific conditions on objectives (e.g., $\\beta$-smooth, L-Lipschitz), which may not hold in all deep learning settings; this is a theoretical concern (Sec. 6, Thm. 2).\n  \n- The paper doesn't explore the connection between the adversary budget $\\epsilon$ with regular ambiguity radius in DRO, like as $\\mathcal{B}_\\delta(P_0)$.\n\n## Minor:\n\n1. correct below typos:\n  \n  - “support **shit** testing cases”\n    \n  - “…restrict **the the** search of worst-case distributions…”\n    \n  - “The convergence of the Moreau **envelop** gradient norm…”\n    \n\n2. Repeated references like Staib & Jegelka (2017a)/(2017b) without disambiguation in prose.\n  \n3. \"Ma et al. Ma et al. (2024)\" — The citation \"Ma et al.\" is repeated. Also \"Wang et al. Wang et al. (2025)\" — The citation \"Wang et al.\" is repeated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yc6VjoY4zZ", "reviewer": "Reviewer_4i8h", "summary": "This paper studies the DRO problems with generative model-based ambiguity sets. For the problems with new constraints, a new framework is proposed to solve the problems and numerical results are provided.", "questions": "While the presentation contains some minor issues, my main concern lies in the motivation and conceptual clarity of this paper.\n\n1. The paper argues that for general DRO problems, in the case of $\\varphi$-divergence DRO, the target distribution must be absolutely continuous with respect to the nominal distribution, and that Wasserstein-DRO formulations are computationally complex. The authors then propose an upper bound on the KL divergence in Lemma 1. However, I find this result puzzling: the left-hand side is expressed as a general function of $P_\\theta$, while the right-hand side depends only on $\\theta$. This mismatch suggests that the authors implicitly treat $P_\\theta$  as a parameterized distribution determined by $\\theta$. Consequently, the results in this paper seem to apply only to specific parameterized models, e.g.,  VAEs or diffusion models, rather than to general optimization problems in the DRO framework.\n\n2. Even within the context of VAEs and diffusion models, I am not convinced of the necessity of this approach. Lemma 1 shows that the upper bound on the KL divergence is linear in $J$ and, given $P_0$,  the gap between KL and J is fixed. Therefore, the results produced by the proposed method may not differ substantially from those obtained through traditional KL-based methods, especially since the constraint functions are linear and the KL divergence is typically easier to compute.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper provides both theoretical and numerical results for the proposed methods.", "weaknesses": "The writing can be improved. For example, the $s_\\theta$ used in line 204 is not previously defined.  The definitions of $J_{VAE}, R_{VAE}$ are important and should be highlighted under equation (6). The smooth and Lipschitz assumptions used in the Theorems should be defined.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)