WILD-Diffusion: A WDRO Inspired Training Method for Diffusion Models under Limited Data
Authors: Xianglu Wang
Wanlin Zhang
Hu Ding
Venue: ICLR 2026 Poster (2026)
Keywords: Diffusion model; Wasserstein Distributionally Robust Optimization; Limited Data
Primary area: generative models

Abstract: Diffusion models have recently emerged as a powerful class of generative models and have achieved state-of-the-art performance in various image synthesis tasks. 
However, training diffusion models generally requires large amounts of data and suffer from overfitting when the dataset size is limited. To address these limitations, we propose a novel method called WILD-Diffusion, which is inspired by Wasserstein Distributionally Robust Optimization (WDRO), an important and elegant mathematical formulation from robust optimization area. Specifically, WILD-Diffusion utilizes WDRO to iteratively generate new training samples within a Wasserstein distance based uncertainty set centered at the limited data data distribution. This carefully designed method can progressively augment the training set throughout the training process and effectively overcome the obstacles caused by the limited data issue. Moreover, we establish the convergence guarantee for our algorithm even though the mixture of diffusion process and WDRO brings significant challenges to our analysis in theory. Finally, we conduct a set of experiments to verify the effectiveness of our proposed method. With WILD-Diffusion, we can achieve more than a $10$% reduction in FID using only $20$% of the training data across different datasets. Moreover, our method can attain state-of-the-art FID with as few as $100$ images, both in pretrained and non-pretrained settings.

Published as a conference paper at ICLR 2026
WILD-DIFFUSION: A WDRO INSPIRED TRAINING METHOD FOR DIFFUSION MODELS UNDER LIMITED DATA
Xianglu Wang', Wanlin Zhang".3, Hu Ding
"School of Artificial Intelligence and Data Science, University of Science and Technology of China School of Computer Science and Technology, University of Science and Technology of China
i "Shanghai Innovation Institute
{wxlu, ideven, huding)@maill esstc.dd.cc
ABSTRACT
Diffusion models have recently emerged as a powerful class of generative models and have achieved sfate-of-the-art performance in various image synthesis tasks. However, training diffusion models generally requires large amounts of data and suffer from overfilting when the dataset size is limited. To address these limi- tations, we propose a novel method called WILD-Diffusion, which is inspired by Wasserstein Distributionally Robust Optimization (WDRO), an important and elegant mathematical formulation from robust optimization area. Specifically, WIL.D-Diffusion utilizes WDRO iteratively generate new training samples within a Wasserstein distance based uncertainty set centered at the limited data data distribution. This carefully designed method can progressively augment the training set throughout the training process and effectively overcome the obstacles caused by the limited data issue. Moreover, we establish the convergence guar- antee for algorithm even though the mixture of diffusion process and WDRO brings significant challenges to our analysis in theory. Finally, we conduct set of experiments verify the effectiveness of proposed method. With WILD- Diffusion, we can achieve more than a 10% reduction in FID using only 20% the training data across different datasets. Moreover, method can attain state- of-the-art FID with as few as 100 images, both in pretrained and non-pretrained settings. The code available at githubu repo
INTRODUCTION
Song et al., 2021b) have become leading family of deep generative models. Unlike generative adversarial net-
Diffusion models (Ho et al. 2020 Sohl-Dickstein et al. 2015 Song & Ermon 2019
works (GANs) (Goodfellow et al. 2014) and variational autoencoders (VAEs) (Kingma et al. 2013 Rezende al. 2014), which generate samples by decoding from low dimensional latent variable, diffusion models learn to iteratively denoise noise corrupted signal through forward-reverse dif- fusion process (Yang et al. 2023b). Recent studies show that diffusion models have been shown to outperform GANs in many image generation tasks, including image editing (Huang et al. 2025b
Gu et al. allaal.22220 2023 Kawar et al. 2023 Yang et al. 2023a). image restoration (Xia et al. |023 |
Fei ct al.
2023
Zhu al. 2023 Lin al. 2024), style transter (Zhang et al. 2023b Wang al. [223d]] Yang
et al.
2023c), and text-to-image generation (Zhang et al. 2023a Saharia et al. 2022 Ruiz et al.
2023
However, the increasingly impressive results of diffusion models are fueled by the seemingly un- limited supply of images. In other words, diffusion models require large amounts of data for stable training (Wang et al. 2023a; Li et al.) [025] Zhang et al. 2025). which hinders the application of diffusion models in limited data settings. For example, training vanilla diffusion model (Ho al. 2020) on only 2,000 samples from the FFHQ dataset (Karras et al. 2019) (about 4% of the full dataset) leads to about 30. To address this limitation, recent studies have explored fine-tuning for image gener-
sharp performance drop, with the FID increasing from about 2.5 (full dataset)
ation under limited data (Ruiz al.) 2023 Moon et 2022
Zhu et al. 2022 Hur et al. 2024
"Corresponding author.

---

Published as a conference paper at ICLR 2026
Yang et al. 2024 Lu et al. 2023 Zhang et al. 2025). For example, Ruiz et al. (2023) applied fine- tuning to transfer knowledge from models pre-trained on large-scale external datasets, which allows the model to synthesize high-quality images using only few target examples. However, these ap- proaches heavily rely on the similarity between the source (i.e., large-scale external datasets) and the target dataset (i.e., the limited dataset) (Hur al. 2024). This reliance hinders the broader adoption of generative diffusion models in data-sensitive fields such as medicine (Kazerouni et al.) 2022). More critically, |Moon et al. (2022) observed that when limited data are used fine-tune pretrained diffusion backbone, the model suffers from overfitting, which means that it memorizes individual training examples rather than captures the underlying data distribution, and this results in near-duplicate outputs and reduced diversity (Webster et al.] 2019] Karras et al. (020). This prob- lem is particularly severe under limited data settings, where the scarcity of training samples makes the model prone to memorization rather than generalization.
To further illustrate this overfitting phenomenon, we conduct an empirical study to examine how training data size influences their convergence behavior. Specifically, we investigate the performance dynamics by training a denoising diffusion probabilistic model (DDPM) (Ho et al. (020) on subsets of FFHQ (64 x 64) (Karras al., 2018). We measure the quality by computing Fréchet inception distance (FID) (Heusel et al. 2017) between 50k generated images and all available training images. As shown in Fig. Dh. the FID curve exhibits "U-shaped" trend: it decreases in the early stages, reaches a minimum FID, and then worsens as training continues; smaller datasets yield an earlier turning point and higher final FID, which clearly indicates overfitting. It is worth noting that previous work (Karras et al., 2020) reported similar convergence behavior for GANs. Furthermore, we also evaluated DDPM on the CelebA-HQ (64x64) (Liu al., 2015) dataset. The results, shown in Fig. [p. are consistent with the above findings that the FID curves also exhibit a U-shaped trend. For completeness, we also illustrate that the training loss decreases monotonically in all cases (as shown in Fig. [b), while at the same time the FID curve exhibits U-shaped pattern, which indicates that overfitting indeed exists.
40
35
FFHO
FFHQ
30-
Coldbn (106)
1.2
CalabA-11 (20%)
Cd
30
FFHQ (59%)
25
1.0-
Caldia--OO
CataNA-1Q
20-
0.8
1911 (195)
0.6
1911
FFHO
0.4
FFHO (100%)
0.2
000-
0-
50
Namber of iimagee shown lo model (M)
Number of images shown model (M)
Number of images shown model (M)
Convergence of FFHO (64x64)
Convergence of CelebA-HQ (644664)
(c) Training loss
Figure I: Evidence of overfitting in diffusion models with limited data. (a, b) FID curves of DDPM on FFHQ (64x64) and CelebA-HQ (64x64) datasets, both exhibiting a "U-shaped" trend where smaller datasets yield earlier turning points and higher final FID. Percentages (e.g., 50%) indicate the fraction of training data used. (c) Training loss decreases monotonically across all cases.
For classification models, a wide range of methods have been developed to address the problem of overfitting. These approaches can be broadly divided into two categories: (I) regularization- based techniques, such as L1/L2 penalties ( (Tibshirani) 1996; Ng. 2004); and (2) data augmentation strategies, such as Cutout (De Vrees & Taylor] 2017). Mixup (Zang et al. 2018), and CutMix (Yun) et al. 2019). However, most of these methods are tailored to classification objectives and cannot be directly transferred to handle diffusion models due to the following two main reasons. (R1) Regularization-basse techniques are primarily designed to improve the generalization of decision houndaries in classification models; they provide limited benefit when the goal is to capture the un- derlying data distribution, in diffusion models. (R2) ingmentation-basel techniques
typically static and rule-driven. These methods can not adaptively constrain distributional shift, and may even exacerbate the discrepancy by pushing the training marginal distribution further away from the true data distribution. As a result, the model could learn off-distribution artifacts and reproduce them at generation time, where this problem is called "augmentation leakage" in (Karras et al. 2020).
In this paper, our proposed method is inspired by Wasserstein Distributionally Robust Optimization (WDRO) (Gao & Kleywegt 2023 Sinha et al. 2018 Huang & Ding 2025). elegant and pow-
2

---

Published as a conference paper at ICLR 2026
erful mathematical framework from the field of robust optimization. A major advantage is that it operates directly on data distribution and adaptively expands the support of the training distribution while remaining close to the true data distribution. Specifically, WDRO replaces empirical risk min- imization (ERM) on the limited data data distribution Pdata with optimization against the worst-case distribution in Wasserstein uncertainty set
Up(Poms) {p: (c Pdat) S pt,
p-neighborhood of the distribution Pdota under the Wasserstein metric We(..) (see Section 2.2
for a formal definition). WDRO has been proven to effectively mitigate overfitting in supervised learning (e.g., adversarial training (Liu et al. 2025) and continual learning (Wagg al. 2023c ) by dynamically adjusting the data distribution. Conceptually, WDRO can be viewed as an adaptive method for support expansion: rather than fitting only the narrow support of Plata (a key source of overfitting in limited data settings), the learner is trained to perform well over a neighborhood of distributions within transportation budget p. Therefore, under the WDRO perspective, natural question arises:
Can the idea of "adaptive support expansion" in WDRO be applied to diffusion models to enlarge the effective training support, with the goal of improving generative quality while mitigating overfittig in limited data settings?
1.1 OUR MAIN CONTRIBUTIONS
To address the above question, we propose "WDRO Inspired training method for Diffusion model under Limited Data (WILD-Diffusion",, a plug-and-play training framework that leverages WDRO to dynamically expand the support of the lim- ited data distribution, which can mitigate overfitting and enhance generative performance. It is worth noting that the idea of DRO has recently been introduced into diffusion models (Wang et al., 2025a); however, this work addresses a different problem about diffusion models, which focuses on the training and sampling distribution mismatch issue rather than limited data generation. Specifically, we apply WDRO to the diffusion problem, where the objective can be formulated
Data distribution
Figure 2: Illustration of support expansion in a 1D setting. Yel- low (solid): true distribution; Red (dotted): limited data dis- tribution; Blue (dashed): dis- tribution induced by WILD- Diffusion, which expands the support of the limited data dis- tribution toward the true distri- bution and narrows the gap.
minimize sup
leple(0: x. (Ptn))
(2)
where the uncertainty set (P(maa) is defined in Eq. (1. A denotes the model parameters, (x, t) are the diffusion training inputs (data and time point i), and E(0; x, ()) represents the diffusion training loss function, which will formally defined a later section (see Eq. (6). The solution of the problem 2 guarantees reliable performance against data distributions that are distance P away from the limited data distribution Pama- Roughly speaking, the solution of problem (2) is expected to expand the support toward the underlying data distribution and narrow the gap (as illustrated in Figure D). which in turn mitigates overfitting and improves sample quality under limited data settings.
Nevertheless, efficiently implementing this idea within diffusion training is not straightforward, as it involves two major challenges. (CI) Because both diffusion training and the computation of the Wasserstein distance are computationally expensive, the first difficulty is ensure the inner max- imization tractable while preserving overall training efficiency. (C2) Since WDRO is inherently min-max optimization problem with notoriously difficult convergence, another critical challenge is to establish theoretical convergence guarantee for the WILDD-iiffusion framework. To tackle these challenges, we build on the surrogate loss idea (Blanchet Murthy, 2019) and reformulate problem (2) as an approximate optimization problem that is tractable in Euclidean space. This re- formulation ensures the otherwise intractable inner maximization can be efficient computation, and further propose "Bi-level Interval Update" strategy to derive practical approximate solution. Specifically, the strategy alternates between parameter updates on the current mixed training set (i.e., original samples and their adversarial counterparts) and distribution interval updates through
3

---

Published as a conference paper at ICLR 2026
worst-case sample generation. Furthermore. establish convergence guarantee for the proposed WILD-Diffusion method. Unlike prior work (Lee et al. 2022). which analyzes the convergence of standard diffusion models, the incorporation of WDRO requires additional technical step: prove an upper bound for the worst-case objective (Leema 3.5). which is essential for achieving the convergence of our proposed WILDD-Diffusion.
The experiments on a variety of diffusion architectures (DDPM++, NCSN++ and ADM) and datasets (CIFAR-10,, LSUN-Church, CelebA-HQ. and FFHQ) suggest the effectiveness of our method. With WILD-Diffusion, we can achieve more than a 10% reduction in FID using only 20% of the training data across all datasets. addition, our method achieves ootte--f-thearrr FID with as few as 100 images, in both pretrained and non con-pretrained settings.
BACKGROUND
In this section, we first review the background of crffusionn-basee generative models, outlining their key formulations and training objectives. We then introduce the concept of Wasserstein distance, which plays a central role in the formulation of our WILDD-Diffusion framework. Due to space limitations, additional related work provided in the Appendix
2.1 RIFFUSION-BASER GENERATIVE MODELS
Suppose we are given dataset {x; [x]] where each data point is independently drawn from an underlying data distribution with positive density Pdata (x). We slightly abuse notation by using measure and its density interchangeably when the context is clear. The forward process is to construct a process {x(t)){= indexed by a continuous time variable E [0, Note that the process starts from x(0) ~ Pawa(x) and evolves to x(T) (x()), where pr typically denotes simple prior distribution, such as a standard Gaussian (Ho et al.) 2020). According to (Song et al.] 2021b), the forward diffusion process can be modeled as stochastic differential equation (SDE):
=x=f(xx f(x, () dl + g({) dw,
where f(., f) _d + Rd is called the drift coefficient of x(f), (()): : I is scalar function known as the diffusion coefficient of x(t), is the standard Wiener process (a.k.a., Brownian mo- tion), and dt represents negative infinitesimal timestep. Importantly, for any forward diffusion process the form of Eq. 3).. Anderson (982) showed that it could be reversed by solving the following reverse-time SDE:
dx [f(x, - g(l)2Vx (xg] di + g(l)
(4)
where w a standard Wiener process when time flows backwards, and the gradient of the log probability density with respect to the data, Vx log (x(),, is the (Stein) score (Liu et al.) 2016 Moreover, Song et al. (2021b) proved the existence of an ordinary differential equation (ODE), namely the probability flow ODE, whose trajectories have the same marginals as the reverse-time SDE (4). The probability flow ODE is expressed as:
dx f(x, t) g())² log (x)) dt.
(5)
Note that if the score of the marginal distributions, Vx (x)), is known for all / E [o, then the reverse diffusion process can be derived from Eq. (5) and subsequently simulated to generate samples from (x). Specifically, a iime-dependent score model sø(x, t) is trained to estimate the score function, which yields the following training objective:
E(A, x) X(t) [|se (x, t) Vx log Pi (x)
where A(t): : J0, T) + R,, is positive weighting function (Yang et al. 2023b
2.2 WASSERSTEIN DISTANCE
The Wasserstein distance, which originates from the theory of optimal transport (Peyré al. 2019 Villani et al.) 2008), has been widely adopted in machine learning (Sinha et al. 2018 Kolouri al. 2017). Let X C Rd denote the sample space. For x, x' E X. the transportation cost e associated
with moving mass from x to x' is defined as (Volpi al. 2018
c(x, x') = 2 [|x x 112.
(7)
4

---

Published as a conference paper at ICLR 2026
As the Le norm is the standard choice in optimal transport, confine analysis to this setting. Given two probability measures P and Q supported on X, let [I, Q) denote the set of couplings between P and Q, ie., measures M on X' x X' with marginals P and Q. Then, the Wasserstein distance between P and Q is defined as
Qc(P.Q) inf Mell(..)
EM [c(x,x')].
(8)
3 METHOD
In this section, we present WILD-Diffusion, WDRO inspired framework designed to enable effec- tive training of diffusion models in limited data settings. highlight of WIL.D-Diffusion is that dynamically leverages WDRO to construct worst-case distributions that lie close to the limited data distribution (in Wasserstein distance), which expands the support of the training distribution and im- proves sample diversity, and consequently relieves the negative impact of overfitting. Moreover, framework is flexible and can be combined with a wide range of baseline methods. We first present Next, we provide the convergence analysis of our
our WILD-Diffusion framework in Section 3.1
proposed approach in Section 3.2
3.1 WILD FRAMEWORK
Wasserstein Distributionally Robust Optimization (WDRO) (Kuhn et al. 2019 Rahimian & Mehro- tra 2019b] Sinha et al. 2018 formulates robust decision decision-making under uncertainty by optimizing for the worst-case over all probability distributions within Wasserstein ball. The Wasserstein ball consists of all distributions whose distance from the limited data distribution does not exceed a given framework, assume that the true data dis- tribution lies in a Wasserstein uncertainty set I i.e., (Ptura = {p : Wc(p, Paam) S p)- This formulation captures the distributional uncertainty arising from limited data, which particularly severe when the sample size is small because the limited data distribution poorly approximates the true underlying distribution (see Fig. 2 Recall the optimization objective (2). the inner sup over the Wasserstein uncertainty set enforces the model to cope with increasingly harder perturbations of the limited data distribution. Namely, this strategy can guide the model to learn some new samples
threshold (recall p in Eq. ) In
and therefore prevents memorization and thus mitigates overfitting.
In general, the worst-case optimization that involves the sup operator within the Wasserstein ball computationally challenging for two main reasons: (i) the Wasserstein ball encompasses rich family of probability distributions, making the inner maximization problem inherently infinite- dimensional; and (ii) computing the Wasserstein distance itself is computationally expensive even in approximate forms. While these challenges already arise for relatively simple models, they become particularly severe in the context of diffusion models. To handle the inner maximization problem in 2 we adopt the strong duality property given in (Gao & Kleywegt 2023] Theorem 1) and obtain its dual formulation. Suppose X' IRd is the sample space. Given a fixed penalty parameter y 2 0, the worst-case loss in Eq. (2) can be reformulated as
minimize (()) sup [E,]((0;x, t)] Wel( Pdaa)) = Epun (o) (0: (00xx,))} t
(9a)
where Φ (0;x,t := sup (E(0;x',t) - 7c(x'x x)} x'EX'
(9b)
is the surrogate loss (Blancher & Murthy 2019 Volpi et al. 2018) that replace the usual diffusion
loss f(0; x, 1) (i.e., Eq. (6)) Here, the penalty parameter 7 controls the degree of support expansion; it balances fidelity to the training data and robustness to distributional shifts. Since Paaaa is unknown, the penalty problem (9a) is solved by replacing Pdata with the empirical distribution Pn. where n is the sample size.
Remark 3.1 Eq. (9a) gives the dual formulation of Eq. 2 i.e., both problems share the same optimal value. The advantage of this reformulation is that we can ignore the complicated uncertainty set Up (P Instead, only add a surrogate loss (0; (,:x,t x,t) t) the Eq. (9a). which yields more succinct formulation for optimizing the problem. However, the solution to Eq. (a)) is non -iriiiiaal; we provide further details ity optimization in the following discussion.
5

---

Published as a conference paper at ICLR 2026
In order to solve the duality formulation (aa), we now perform stochastic gradient descent on the surrogate loss Φy. Specifically, suppose that the loss E(0) x, t) satisfies the Lipschitz smoothness conditions (Boyd & Vandenberghe 2004) and that the surrogate loss is strongly concave. Under
these conditions, we have
V Vodl, (0: Vol(O; (x" .0) where
x =
argmaa () ,x) x'EX
(10)
Computing the gradient of the surrogate loss bby for a given sample x requires solving the inner maximization problem to obtain x". Notably, we observe that x" is similar to an adversarial per- turbation of x under the current model 0. Following the intuition of adversarial training (Madry et al. 2018), propose a "Bi-level Interval Update" strategy for from adversarial training is that, while adversarial training typically generates adversarial exam- ples within fixed norm ball, our approach imposes soft constraint via the penalty parameter % which governs distributional robustness at the support level. The strategy couples two updates. (I) Parameter update level. The model parameters 0 are updated every training iteration using the current training set. (II) Distribution (sample) update level. Every m epochs we refresh the WDRO-induced "worst-case" samples via gradient ascent and mix them with the real data to form the augmented training distribution. Between distribution updates, the worst-case samples are kept fixed. Specifically, at the sample update level, for each training example we first draw an initial point x from the data distribution Pdata- We then iteratively update it through the injection of adversarial perturbations, which produces an adversarial variant as defined by the following update rule:
The difference
x +
x ",t) --(-----_-------- xx))
(11)
where Ç denotes the step size and k 1, ..... K indexes the iterations. At the parameter update level, the model parameters 0 are updated at every training step by performing stochastic gradient descent on the loss E(0; x, t), where the training sets is a mixture of the original samples and their ad- versarial counterparts. Algorithm presents the proposed WILD-Diffusion algorithm, which offers the flexibility to incorporate a variety of baseline methods, since it operates on the data distribution without requiring changes to the model architectures. In addition, we take Sw epochs to train the model on the limited dataset as a warmup stage. The warmup stage yields a stable initialization before incorporating worst-case samples. Starting from well-initialized state enables the model to produce more informative gradients used in Eq. (1 In practice, We allocate 20% of the total training epochs to the warmup stage.
3.2 CONVERGENCE ANALYSIS
In this section, we establish the convergence guarantee for the proposed WILD-Diffusion method. In contrast to prior work (Lee et al. 2022), which focuses on standard diffusion models, our analysis must account for the additional complexity introduced by WDRO. To this end, we establish a upper bound for the worst-case objective (Lemma B.5). which enables the convergence proof of WILD- Diffusion. We first make the following assumptions (i.e., Assumption an.33) on the probability density Paata and the score estimate so(x, 1) (defined in Section 2p. which will be used throughout the analysis.
Assumption 3.2 Assume that Pauta satisfies the log log-Sobolev inequality with constant Cis > 1; log Prisa is L- L-Lipschitz for some L 2 1: Paota has finite first and second moments.
Assumption 3.3 Suppose that so (x)x)t t) L-pLipschitz in its first argument with Lg 2 1, and the error in score estimate {(0; x, t) is uniformly bounded by given parameter & >>0 0.
Remark 3.4 Assumptions 3.2 and 3.3 also adopted in Lee et al. 2022), are standard assumptions
in analyses of score-based diffusion models. In particular, the Lipschitz assumption on Pdata is used
to ensure the existence of a unique strong solution to the reverse-time SDE (Eq. 4 (Block et al.
2020 [esendal 2003).. The detailed definition of the log-Sobolev inequality is given in Appendix D Building on the above assumptions, we derive an upper bound for the optimization objective in
Eq. (2). stated in Lemma 3.3 which is essential condition for the convergence analysis of WILD-Diffusion.

---

Published as a conference paper at ICLR 2026
Algorithm 1 WILD-Diffusion
Input: Training datasets [x,}==) Initialized model parameter 8. learning rate 7), step size num- ber of iterations K in inner optimization, interval parameter m, total diffusion steps T, the number of epochs S. and the number of warmup epochs Sw.
Output: Final diffusion model parameter
0 - 0o 1 Initialize model */
2: for 1....,, Sw do
3: / Take S, epochs train the model the warmup */
4:
for io ....... do
5: Sample Uniform ((1...., T})
6: 0+ 1 Vo
7: end for
8: end for
9: for s Sw L...., S do
10: if (s mod m)=0 then
11:
/0 Support Expansion via WDRO */
12:
13:
for ,,..,, do
x X;, 1~ Uniform (()...,,))
15:
for k = =1, 1...., K do
16:
x +- x;
k-1 + +C (E(0; x k-1 7o(x; ,rcxxx-1xx]1
Distribution (sample) up-
date via Eq. end for
71
*/
17:
D DU fx*1
1 Save worst-case samples */
19:
end for
end if
21:
for imri=la 1i n do
22:
Sample x ~ D, 1. Uniform(11..... 23: 0- V ((0x2,,)
1 Parameter update */
24:
end for
25:
end for
for any fixed > O, the following inequality holds with proba- bility at least 1 e-T, uniformly over all 2 O and 2 0
Lemma 3.5 Under Assumption 3.3
sup (P.
Ep(e(0; x, ()] < 7p Ep, (0; + O(1
(12)
Here, n is the sample size, is the confidence parameter, and pa. is the empirical distribution of the samples from Pdata- We adopt the total variation distance Drv(..)) to quantify conver- gence. Given two distributions p and q. the total variation distance is defined as DTv g) ! [p()) (xx) which measures the maximum discrepancy between two distributions. Be- fore presenting the convergence result of WILD-Diffusion Namely, let qi denote the reverse process with the estimated score. We define the bad set B, as B, =
first provide outline of the proof.
sup supp.Wa....emu Ep ||s0 t) Vx log Pe (x)] > EB for some EB to be chosen, and define T the reverse process with the estimated score except in B,. Hence, the convergence proof can be divided into two parts by applying the triangle inequality
Drv (ge: pt) S DTv (ge, p.) + DTV (t).
(13)
Since (Lee et al. 2022 established the bound Drv(qe, Pt) S $ < with =² denoting the cor- responding error term, the main task is therefore to control second term Drv (g,, (c). This established in Theorem 3.6
Theorem 3.6 (Convergence of WILD-Diffusian
Suppose Assumptions 3.2 and 3.3 hold, and
Lemma 3.5 applies. If we run the SDE (Eq. 4) starting from a Gaussian distribulton for time
+ max log(Cisd), Cis log ( ))}) with step size h = θ Mexx then
=
7

---

Published as a conference paper at ICLR 2026
the final sampling distribution go satisfies
Drv To) S o Esn. Ie (0; V
I (Cs+t))(++23)
1+1oz
-
(14)
For simplicity, we denote the upper bound in Eq. 14 by Dub. Thus, Drv (q0. Pattta S ox + Dub-
The complete proof is provided in Appendix B.3 Theorem 3.6 establishes convergence guarantee for WILD-Diffusion under standard assumptions. Specifically, the total variation distance between the generated distribution qo and the limited data distribution Pdata is bounded by the sum of two terms: an estimation error term ex which arises from the approximation of the score function, and sampling error term Drv (1o, To which is due the numerical computation of the reverse SDE. Notably, when the robustness budget - O and the sample size n + oo, the bound recovers the result of (Lee et al. (022). which showed that
Drv Pdat)) xx + O (VE C (Cis+ +.2) (1 log (
This suggests that our convergence guarantee can be regarded as generalization of the result in (Lee) al. 2022) to the more complicated distributionally robust setting (see Appendix B.4 for details).
4 EXPERIMENTS
In this section, we first present a hyper-parameter sensitivity analysis to investigate the key fac- tors influencing the performance of our method, detailed in Section 4.1 Next, we compare our approach with state-of-the-ar diffusion model baselines on widely-used benchmark datasets in Sec- tion 4.2 In Section 4.3 we further demonstrate that our method performs well on few-shot datasets. is worth noting that in generative modeling, the few-shot setting differs from the limited data regime: the former typically involves adapting a pretrained model to new distribution with only handful of samples (tens to hundreds), whereas the latter refers to training on a small dataset of only thousands of samples without access to large-scale pretraining Abdollahzadeh et al.
2023). Finally,
we conduct the ablation studies in Section
Experimental Setting. In line with previous works (Wang al. 2023a Zhao al. 2020 Karras
let al. 22020), we conduct experiments on standard benchmarks, where subsets of the training data randomly selected. For the limited data setting, adopt CIFAR-10 (32 x 32) (Krizhevsky et al. 2009), FFHQ (64 x 64) (Karras et al. 2019), CelebA-HQ (64 x 64) (Karras et al. 2018), and LSUN-Church (256 x 256) (Yu al., 2015p. For the few-shot setting, we adopt the 100- shot datasets (256 x (936)-Ohmma, Grumpy Cat, and Panda (Zhao et al. 2020)- and AnimalFace (256 256; cats and dogs) (Si Zhuj 2011]. We implement our method on the current start-of- the-art diffusion framework EDM (Karras et al.] 2022). which integrates DDPM++ (Song et al.) 2021b), NCSN++ (Song et al., 2021b), and ADM (Dhariwal & Nichol 2021). DDPM++ is our default backbone model for training low-resolution (i.e., 32x32 and 64x64) datasets, while ADM coupling with Stable Diffusion (Rombach et al. 2022 is our backbone model for training high- resolution (i.e., 256x256) datasets. We evaluate image generation quality using Fréchet Inception Distance (FID) (Heusel et al. (017). Following Karras et (2022) 2020). FID is computed between 50k generated samples and the full set of training images. The detailed experimental settings provided in Appendix C
4.1 SENSITIVITY OF HYPER-PARAMETER
In this section, we investigate the sensitivity of our method key hyper-parameters. In particu- lar, the interval parameter m (in Algorithm ) plays crucial role, as it can substantially influence both generation quality and training efficiency. To assess effect, we conduct a series of experi- ments by varying over {5, 10, 20, 30, 40, 50, 100} on the FFHQ dataset with 50% training data. Figure 3 shows that increasing m reduces the total training time while degrading generative per- formance (higher FID). This reveals clear trade-off between efficiency and quality. Taking both training efficiency and generative quality into account, we set m 20 as the default choice in all experiments.
In addition, to better interpret the influence of injected adversarial perturbations (see Eq. II we examine method on the FFHQ dataset with 50% of the training data across the number of steps,
8

---

Published as a conference paper at ICLR 2026
step size, and penalty strength. When studying one factor, the others are fixed at their best values. From Figure 4 several observations can be drawn: (1) Increasing the number of steps K improves performance up to K 5, after which the gains diminish (Figure 4a); (2) The step size nj 0.01 achieves the best balance, while both smaller and larger values degrade performance (Figure(p)) (3) The penalty parameter y is relatively stable, with 7 performing best (Figure 4b)) In summary, we adopt these configurations as the default in all experiments.
4.2 EXPERIMENTS ON LIMITED DATA GENERATION
In this section, compare our method with state-of-the-art diffusion approaches both low- resolution (Table D and high-resolution (Table 5 benchmarks. For the low-resolution setting, we evaluate on CIFAR-I0 10, CelebA, CelebA-HQ, and FFHQ. Specifically, the basclines include EDM- DDPM++ (Karras et al. 2022), EDM-NCSN+ (Karras et al., 2022), EDM-ADM (Karras et al., 2022). (Wang et al. 2025a). Patch Diffusion (Wang et 2023a). and DeepCache (Ma al. 2024). We include Patch Diffusion and DeepCache to illustrate that WILD-Diffusion serves plug-and- play framework that "orthogonal" to these methods; moreover, they can be seamlessly combined to achieve more promising performance in practice. We incorporate diffusion methods to assess its performance across datasets. For completeness, we also compare with data-efficient GAN-based approaches, including BigGAN (Brock al.) (219)) StyleGAN- v2 Karras et al. 2019), DiffAugment (Zhao et al. 2020), and CR-BigGAN (Zhng et al. 2020).
WILD-Diffusien into these
Table 1: FID results low-resolution datasets. FID (lower is better) is computed with 50k samples. The numerical results of the baseline methods are taken from the original papers. "-" indicates that the result not reported in the original paper (Zhao al. 2020). The notation "(-A%)" indicates percentage decreases compared to the baseline. "A% data" refers to randomly selecting "A%" of
the training data from the dataset, and "cond." " denotes the class-conditional setting. The best-
performing results are highlighted in bold.
Dataset
Method
20% data
30% data
100% data
BigGAN Brook et al. 2019
21.58
9.59
StyleGAN-v2 KKrras
2019
23.08
11.07
CR-BigGAN Zhang et 2020
20.62
9.06
BigGAN--DirAugmenn Zhao al.
2020
14.04
8.70
EDM-DDPM++ Karras et ul. 2022
13.91
6.62
1.97
WWag et al. 2025a
13.63
6.49
-WILD-Diruston
12.14 (-12.72%)
6.02
(-9.08%)
1.93 (-2.03%)
EDM-DDPM++ (cond.) Karras et
2022
12.33
6.03
1.79
CIFAR-10
+ EDM-NCSN++ (Karras et al. 2022
(cond.)
10.89 (-11.68%)
5.37
(-10.95%)
1.71 (-4.47%)
(32 x x2)
13.68
6.53
2.02
WWLLD-Dffussie
12.08 (-11. (1%%
5.97
(-8.58%)
1.98 (-1.98%)
Patch Diffusion (Wang 2023a
12.53
6.42
2.47
-WILD-Diffusion
11.78 (-5.99%)
6.07 (-5.45%)
2.38(-3.6499)
DeepCache (Ma et al. 2024
15.33
9.31
4,35
+ WILLD-Difussoon
13.96 (-8.94%)
8.72 (-6.34%)
4.21 (-3.37%)
HDM-DDPM++ KKaraa al. 2022
10.02
5.21
2.60
FFHQ (64 x 64)
WILLD-Difusion
8.57 (-14.47%)
4.68 (-10.17%)
2.53 (-2.70%)
EDM-NCSN++ (Karras et
2022
9.38
5.04
2.57
+
7.89 (-15.88%)
4.60
(-8.73%)
2.54 (-1.16%)
EDM- EDM-DDPM++ Karras al.
2022
11.86
6.11
3,73
CelebA-HQ (64 x 64)
WILD-Diffusoon
10.22
(-13.83%)
5.55 (-9.17%)
3.63 (-2.68%)
DDM--CSN++ (Karras et al.
2022
11.63
5.81
3.70
WIL.-
10.07 (-13.41%)) %)
5.36 (-7.75%)
3.65 (-1.35%)
The results for the low-resolution benchmarks are summarized in Table I The following two obser- vations can be drawn: (1) with the same amount of training data (from 20% to 100%), our method consistently outperforms the baseline model; and (2) the performance gains are larger when the amount of training data is smaller. This phenomenon is understandable, as limited training data makes models more susceptible to overfitting (see Figure []. which leads to poor generative perfor- mance. For example, the 20% FFHQ training set, method yields 15.88% improvement in FID compared with the baseline EDM-NCSN++ method. However, as the training data increases, the performance gain diminishes to 1.16%. We further evaluate our method the high-resolution

---

Published as a conference paper at ICLR 2026
benchmark CUN-Church, with results reported in Table 5 Appendix C.5 consistent with those drawn from the low low-resolution benchmarks.
The conclusions
4.3 EXPERIMENTS ON FEW-SHOT GENERATION
In practice, it is often impossible to collect large-scale dataset for specific images of interest.
To address this few-shot image generation problem, researchers recently exploit few-shot learn- ing (Gharoun et al. 2024, Wang et al. 2020a) in the setting of image generation, including LD- Diffusion (Zang et al. 2025). LPDM-8 (Wang et al. 2023a). FreezeD (Mo et al. 2020), Transfer-
GAN (Wang et al., 2018bp, MineGAN (Wang et al., 2020b), and DiftAugment (Zhao et al. 2020)
We compare these transfer learning approaches with our data-efficient training scheme. Note that these diffusion-bain -based transfer learning methods start from pre-trained EDM-NCSN++ (Karras) et al. 2022) model on the FFHQ dataset, while these GAN-based methods start from pre-trained StyleGAN-v2 Karras al. 2019) model on the same dataset. Our comparison experiments are conducted on the 100-shot datasets (Obama, Grumpy Cat, and Panda) (Zhao et al., (020),, and AnimalFace (160 cats and 389 dogs) (Si & Zhu, 2011). The results in Table 2 show that WILD- Diffusion achieves consistent gains on all datasets, with without pre-training. For example, our method achieves the lowest FID score 34.52 (representing an improvement at least 7%) on the 100-shot Obama dataset when trained from scratch.
Table The FID results on few-shot generation. Following the setting used in (Zhao et al. 2020 we calculate the FID with 5k samples and the training dataset is adopted as the reference distribution transfer learning methods have their pre-training: from the FFHQ dataset. The numerical results of the baseline methods are quoted from their papers. We highlight the best results in bold.
100-shot
Animal-Face
Methods
Architecture
Pre-training?
Ohama
Grumpy
Panda
Cat
Dog
StyleGAN-v2 Karras et al. 2019
GAN
No
80.20
48.90
34.27
71.71
130.19
+CM (Karaas et al.
2022
Diffusion
No
37.10
29.94
10.81
36.88
57.14
MineGAN Wang 2020h
GAN
50.63
35.54
14.84
$4.45
93.03
TransferGAN (Wang
GAN
48.73
34.06
23.20
$2.61
82.38
FreezeD Mo et al. 2020
GAN
41.87
31.22
17.95
47.70
70.46
LPDM-8 Wang el al. 20233
Diffusion
Yes
14.27
14.56
5.13
14.92
15.95
LD-Diffusion (Zhang ct al.
2025
Diffusion
Yes
13.00
13.31
4.70
12.77
12.48
WILDDDifflsion (ours)
Diffusion
Yes
12.54
12.83
4.66
12.93
12,21
DiffAugment (Zhao al.
2020
GAN
No
46.87
27.08
12.06
42.44
58.85
Patch Diffusion (Wang 2023a
Diffusion
41.47
30.89
13.25
43.71
72.17
WII.D-Diflusion (ours)
Diffusion
34.52
26.33
9.96
34,21
53.18
4.4 ABLATION EXPERIMENTS
Given that proposed method incorporates the "Wasserstein distance" (Eq. it is natural to compare with other distributional divergences that are commonly used in distributionally robust op- timization (DRO). To this end, we perform ablation study by replacing the Wasserstein distance with alternative divergences: (1) KL-divergence, x2)x² -dvvvergenee, and (3) a-divergence (see de- tailed definition in Appendix D). The results are summarized in Figure [ in Appendix C.6]. which demonstrate that our method achieves superior performance compared to these alternatives. Further- more, as our method can be regarded as a novel data augment method with theoretical guarantee, we perform the ablation experiments comparing it against representative augmentation techniques, including Mixup (Zhang al., 2018), CutMix (Yun al.) 2019), and CutOut (De Vries & Taylor 2017]. The results are presented in Table in Appendix C.6, which show that our method achieves better performance than other methods.
CONCLUSION
In this paper, we introduced WILD-Diffusion. novel diffusion training framework based on WDRO. Our method dynamically expands the support the training distribution, which miti- gates overfitting and improves generation quality under limited data. We proposed an efficient algorithm with theoretical convergence guarantee, and extensive experiments demonstrated that WIL.DDiiffsion can improve state-of-the-art diffusion models across diverse datasets and architec- tures.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
The authors would like thank the anonymous reviewers for their valuable comments and sugges- tions. This work was partially supported by the National Key Research and Development Program of China (No. 1221100000000 the National Natural Science Foundation of China (No. 62272432, No. 62432016), and the Natural Science Foundation of Anhui Province (No. 2208085MF163).
REFERENCES
Milad Abdollahzadeh, Touba Malekzadeh, Christopher TH Teo, Keshigeyan Chandrasegaran, Guimeng Liu, and Ngai-Man Cheung. survey on generative modeling with limited data, few shots, and zero shot. CoRR, 2023.
Mohammed Amin Abdullah, Hang Ren, Haitham Bou Ammar, Vladimir Milenkovic, Rui Luo, Mingtian Zhang, and Jun Wang. Wasserstein robust reinforcement learning. arXiv preprint arXiv:1907.13190 2019.
Brian DO Anderson.
Reverse-time diffusion equation models. Stochastic Processes and their Ap- plications, 12(3):3133325 1982.
Peter I. Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of machine learning research, 3(Nov):463-483 2002.
Aharon Ben-Tal, Dick Den Hertog, Anja De Waegenaere, Bertrand Melenberg, and Gijs Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Sci- ence, 59(2):341-33577 2013.
Jose Blanchet and Yang Kang. Semi-supervised learning based on distributionally robust optimiza- tion. Data Analysis and Applications 3: Computational, Classification, Financial, Statistical and Stochastic Methods, 5:1-33, 2020.
Jose Blanchet and Karthyek Murthy. Quantifying distributional model risk via optimal transport. Mathematics of Operations Research, 44(2):565-600, 2019.
Jose Blanchet, Yang Kang, and Karthyek Murthy. Robust wasserstein profile inference and applica- tions to machine learning. Journal of Applied Probability, 56(3):830-857, 2019a.
Jose Blanchet, Yang Kang, Karthyek Murthy, and Fan Zhang. Data-driven optimal transport cost selection for distributionally robust optimization. In 2019 winter simulation conference (WSC), pp. 3740-3751. IEEE, 2019b.
Adam Block, Youssef Mroueh, and Alexander Rakhlin. Generative modeling with denoising auto- encoders and langevin sampling. arXiv preprint 17Xiv:2002.00107, 2020.
Stéphane Boucheron, Olivier Bousquet, and Gábor Lugosi. Theory of classification: A survey of some recent advances. ESAIM: probability and statistics, 9:323-375, 2005.
Stephen P Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. International Conference on Learning Representations, 2019.
Ruidi Chen and Ioannis Ch. Paschalidis. A robust learning approach for regression models based on distributionally robust optimization. J. Mach. Learn. Res., 19:13:1-13:48, 2018.
Tianlong Chen, Yu Cheng, Zhe Gan, Jingjing Liu, and Zhangyang Wang. Data-efficient gan train- ing beyond (just) augmentations: A lottery ticket perspective. Advances in Neural Information Processing Systems, 34:20941-20953 2021.
Xin Chen, Melvyn Sim, and Peng Sun. A robust optimization perspective stochastic program- ming. Operations research, 1556) 2007.
Erick Delage and Yinyu Ye. Distributionally robust optimization under moment uncertainty with application to data-driven problems. Operations rexearch, 58())595-612, 2010.

---

Published as a conference paper at ICLR 2026
Terrance De and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint WrXiv:170..04522 2017.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in neural information processing systems, 34:8780-88 1188880-7744 2021.
John C Duchi, Peter W Glynn, and Hongseok Namkoong. Statisties of robust optimization: A generalized empirical likelihood approach. Mathematics of Operations Research, 46(3):946-99) 96999999, 2021.
Ben Fei, Zhaoyang Lyu, Liang Pan, Junzhe Zhang, Weidong Yang, Tianyue Luo, Bo Zhang, and
Bo Dai. Generative diffusion prior for unified image restoration and enhancement. In Proceedings of the JEEE/CVE conference on computer vision and pattern recognition, pp. 9935-9946, 2023.
Rui Gao and Anton Kleywegt. Distributionally robust stochastic optimization with wasserstein distance. Mathematics of Operations Research, 48(2):603-655, 2023.
Hassan Gharoun, Fereshteh Momenifar, Fang Chen, and Amir H Gandomi. Meta-learning ap- proaches for few-shot learning: survey of recent advances. ACM Computing Surveys, 56(12): 1-41, 2024.
Joel Goh and Melvyn Sim. Distributionally robust optimization and its tractable approximations. Operations research, 58(4-patt-):90229111 2010.
lan J Goodfellow, Jean Pouget-Abadie, Abadie, Mehdi Mirza, Bing Xu, David Warde-Farles Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
Jing Gu. Yilin Wang, Nanxuan Zhao, Tsu-Jui Fu, Wei Xiong, Qing Liu, Zhifei Zhang, He Zhang, Jianming Zhang, HyunJoon Jung, et al. Photoswap: Personalized subject swapping in images. Advances in Neural Information Processing Systems, 36:35202 16:35202-35217, 2023.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840- 6851, 2020.
Jiawei Huang and Hu Ding. An effective manifold- based optimization method for distributionally robust classification. In The Thirteenth International Conference on Learning Representations, 2025.
Jiawei Huang, Minming Li, and Hu Ding. Bootstrap your uncertainty: Adaptive robust classifica- tion driven by apiimal-transport. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025a.
Ruomin Huang, Jiawei Huang, Wenjie Liu, and Hu Ding. Coresets for wasserstein distributionally robust optimization problems. Advances in Neural Information Processing Systems, 35:26375- 26388, 2022.
Yi Huang, Jiancheng Huang, Yifan Liu, Mingfu Yan, Jiaxi Lv, Jianzhuang Liu, Wei Xiong, He Zhang, Liangliang Cao, and Shifeng Chen. Diffusion model-based image editing; A survey. IEEE Transactions Pattern Analysis and Machine Intelligence, 2025b.
Jiwan Hur, Jaehyun Choi, Gyojin Han, Dong- Jae Lee, and Junmo Kim. Expanding expressiveness of diffusion models with limited data via -elf-distillation based fine- tuning. In Proceedings the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 5028-50080 5037, 2024.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for im- proved quality, stability, and variation. In International Conference on Learning Representations, 2018.
12

---

Published as a conference paper at ICLR 2026
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401-4410, 2019.
Tero Karras, Miika Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Training generative adversarial networks with limited data. Advances in neural information processing systems, 1111214--22144 2020.
Tero Karras, Miia Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in neural information processing systems, 35:26565 265777 2022.
Bahjat Kawar, Shiran Zada, Oran Lang, Omer Tov, Huiwen Chang, Tali Dekel, Inbar Mosseri, and Michal Irani. Imagic: Text-based real image editing with diffusion models. In Proceedings the IEEE/CVF conference on computer vision and pattern recognition, pp. 6007-6017, 2023.
Amirhossein Kazerouni, Ehsan Khodapanah Aghdam, Moein Heidari, Reza Azad, Mohsen Fayyaz, lkker Hacihaliloglu, and Dorit Merhof. Diffusion models for medical image analysis: A compre- hensive survey, arXiv preprint arXiv:2211.078M4 2022.
Diederik P Kingma, Max Welling, et al. Auto-encoding variational bayes, 2013.
Donald E Knuth. The Art of Computer Programming: Fundamental Algorithms, Volume 1. Addison- Wesley Professional, 1997.
Soheil Kolouri, Se Rim Park, Matthew Thorpe, Dejan Slepcev, and Gustavo K Rohde. Optimal mass transport: Signal processing and machine-leamingg applications. IEEE signal processing magazine, 34(4):43-59, 2017.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Daniel Kuhn, Peyman Mohajerin Esfahani, Viet Anh Nguyen, and Soroosh Shafieezadeh- Abadeh.. 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.
Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score-based generative modeling with polynomial complexity. Advances in Neural Information Processing Systems, 15:22870-22882 2022.
Jaeho Lee and Maxim Raginsky. Minimax statistical learning with wasserstein distances. Advances in Neural Information Processing Systems, 31, 2018
Alexander Levine and Soheil Feizi. Wasserstein smoothing: Certified robustness against wasserstein adversarial attacks. In International Conference on Artificial Intelligence and Statistics, pp. 3938- 3947. PMLR, 2020.
Yijun Li, Richard Zhang, Jingwan Cynthia Lu, and Eli Shechtman. Few-shot image generation with elastic weight consolidation. Advances in Neural Information Processing Systems, 33:15885- 15896, 2020.
Yize Li, Yihua Zhang, Sijia Liu, and Xue Lin. Pruning then reweighting: Towards data-efficient training of diffusion models. ICASSP 2025-2025 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2025.
Xinqi Lin, Jingwen He, Ziyan Chen, Zhaoyang Lyu, Bo Dai, Fanghua Yu, Yu Qiao, Wanli Ouyang, and Chao Dong. Diffbir: Toward blind image restoration with generative diffusion prior. In European Conference on Computer Vision, pp. 430-448 Springer, 2024.
Bingchen Liu, Yizhe Zhu, Kumpeng Song, and Ahmed Elgammal. Towards faster and stabilized gan training for fighltidelity few-shot image synthesis. iclr, 2021.
Qiang Liu, Jason Lee, and Michael Jordan. A kernelized stein discrepancy for goodness-of-fr tests. In International conference machine learning, pp. 284-284. PMLR, 2016.
13

---

Published as a conference paper at ICLR 2026
Shuang Liu, Yihan Wang, Yifan Zhu, Yibo Miao, and Xiao-Shan Gao. Provable robust overfitting mitigation in wasserstein distributionally robust optimization. In The Thirteenth International Conference on Learning Representations, 2025.
Zijian Liu, Qinxun Bai, Jose Blanchet, Perry Dong, Wei Xu, Zhengqing Zhou, and Zhengyuan Zhou. Distributionally robust g-learning. In International Conference on Machine Learning, pp- 13623-13645 PMLR, 2022.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Haoming Lu, Hazarapet Tumanyan, Kai Wang. Shant Navasardyan, Zhangyang Wang, and Humphrey Shi. Specialist diffusion: Plug-and-play sample-efficient fine-tuning of text-to-image diffusion models to learn any unseen style. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 14267-14276, 2023.
Xinyin Ma, Gongfan Fang, and Xinchao Wang. Deepeache: Accelerating diffusion models for free. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 15762-15772, 2024.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. In International Conference on Learning Reprexentations, 2018
Sangwoo Mo, Minsu Cho, and Jinwoo Shin. Freeze the discriminator: a simple baseline for fine- tuning gans. arXiv preprint arXiv:2002. 200444000...044,, 2020.
Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization
using the wasserstein metric: Performance guarantees and tractable reformulations. Mathematical Programming, 175(1):11--66,, 2018.
Taehong Moon, Moonseok Choi, Gayoung Lee, Jung- Woo Ha, and Juho Lee. Fine-tuning diffusion models with limited data. In NeurlPS 2022 Workshop on Basee Methods, 2022.
Amir Najafi, Shin-ichi Maeda, Masanori Koyama, and Takeru Miyato. Robustness to adversar- ial perturbations in learning from incomplete data. Advances in Neural Information Processing Systems, 32, 2019.
Hongseok Namkoong and John C Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. Advances in neural information processing systems, 29, 2016.
Andrew Ng. Feature selection, 1 vs. regularization, and rotational invariance. In Proceedings of the twenty-first international conference on Machine learning, pp. 78, 2004.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pp. 8162-8171. PMLR, 2021.
Utkarsh Ojha, Yijun Li, Jingwan Lu, Alexei A Efros, Yong Jae Lee, Eli Shechtman, and Richard Zhang. Few-shot image generation via cross-domain correspondence. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10743-10752, 2021.
Bernt Øksendal. Stochastic differential equations. In Stochastic differential equations: an introduc- tion with applications, pp. 38-50. Springer, 2003.
Gabriel Peyré, Marco Cuturi, al. Computational optimal transport: With applications data science. Foundations and Trends@ in Machine Learning, 11(-66):35-607, 2019.
Hamed Rahimian and Sanjay Mehrotra. Distributionally robust optimization: A review. CoRR, 0e///988.055577 2019a.
Hamed Rahimian and Sanjay Mehrotra. Distributionally robust optimization: A review. arXiv preprint arXiv:/908. 190,0559,,, 2019b.
Hamed Rahimian and Sanjay Mehrotra. Frameworks and results in distributionally robust optimiza- tion. Open Journal Mathematical Optimization, 3:1-85, 2022.
14

---

Published as a conference paper at ICLR 2026
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and ap- proximate inference in deep generative models. In International conference on machine learning, pp. 1278-1286 PMLR, 2014.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and puttern recognition, pp. 10684-10695, 2022.
Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. Dreambooth: Fine tuning text-to-image diffusion models for subject-driver generation. In Pro- ceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 22500- 22510, 2023.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in neural informa- tion processing systems, 35:36479-36494, 2022.
Zhangzhang Si and Song-Chin Zhu. Learning hybrid image templates (hit) by information projec- tion. IEEE Transactions on pattern analysis and machine intelligence, 34(7):354-1136, 2011.
Aman Sinha, Hongseok Namkoong, and John Duchi. Certifying some distributional robustness with principled adversarial training. In International Conference on Learning Represenlutions, 2018.
Jascha Sohl-Dickstein Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing. pp. 2256-2265, pmlr, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- fional Conference on Learning Representations, 2021a.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Yang Song and Stefano Ermon. Improved techniques for training score-bas bare-based generative models. Advances in neural information processing systems, 33:124384 15:188122244, 2020.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 202
Matthew Staib and Stefanie Jegelka. Distributionally robust deep learning as a generalization of adversarial training. In NIPS workshop Machine Learning and Computer Security, volume 3, pp. 4, 2017.
Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society Series B: Statistical Methodology, 58(1):267-288, 1996.
Ngoc-Trung Tran, Viet-Hung Tran, Ngoc-Bao Nguyen, Trung-Kien Nguyen, and Ngai-Man Che- ung. On data augmentation for gan training. IEEE Transactions on Image Processing, 30:1882 1882- 2021.
Santosh Vempala and Andre Wibisono. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. Advances in neural information processing systems, 32, 2019.
Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2008
Riccardo Volpi, Hongseok Namkoong, Ozan Sener, John C Duchi, Vittorio Murino, and Silvio Savarese. Generalizing to unseen domains via adversarial data augmentation. Advances in neural information processing systems, 31, 2018.
Dilin Wang, Hao Liu, and Qiang Liu. Variational inference with tail-adaptive f-divergence. Ad- vances Neural Information Processing Systems, 31, 2018a.
15

---

Published as a conference paper at ICLR 2026
Yaqing Wang, Quanming Yao, James T Kwok, and Lionel Mi.. Generalizing from few examples: A survey on few-shot learning. ACM computing surveys (csur), 53(3):1-34, 2020a.
Yaxing Wang, Chenshen Wu, Luis Herranz, Joost Van de Weijer, Abel Gorialez-Garcia, and Bog- dan Raducanu. Transferring gans: generating images from limited data. In Proceedings of the European conference on computer vision (ECCV). pp. 218-2344 2018b.
Yaxing Wang, Abel Gonzalez-Garcia, David Berga, Luis Herranz, Fahad Shahbaz Khan, and Joost van de Weijer. Minegan: effective knowledge transfer from gans to target domains with few images. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp.
2020b.
Zekun Wang, Mingyang Yi, Shuchen Xue, Zhenguo Li, Ming Liu, Bing Qin, and Zhi-Ming Ma. Improved diffusion-based generative model with better adversarial robustness. In The Thirteenth International Conference on Learning Representations, 2025a.
Zhendong Wang, Yifan Jiang, Huangjie Zheng, Peihao Wang, Pengcheng He, Zhangyang Wang, Weizhu Chen, Mingyuan Zhou, et al. Patch diffusion: Faster and more data-efficient training of diffusion models. Advances in neural information processing systems, 36:72137-72155 2023a.
Zhendong Wang, Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Diffusion- gan: Training gans with diffusion. In The Eleventh International Conference on Learning Repre- sentations, 2023b.
Zhenyi Wang, Li Shen, Tiehang Duan, Qiuling Suo, Le Fang. Wei Liu, and Mingchen Gao. Dis- tributionally robust memory evolution with generalized divergence for continual learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45312):14377-4435, 2023c.
Zhizhong Wang, Lei Zhao, and Wei Xing. Stylediffusion: Controllable disentangled style transfer via diffusion models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 7677-7689, 2023d.
Zitao Wang. Ziyuan Wang. Molei Liu, and Nian Si. Knowledge-guided wasserstein distributionally robust optimization. In Forty-second International Conference on Machine Learming, 2025b.
Ryan Webster, Julien Rabin, Loic Simon, and Frédéric Jurie. Detecting overfitting of deep generative networks via latent recovery. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11273-11288 1282, 2019.
Jon Wellner et al. Weak convergence and empirical processes: with applications to statistics. Springer Science Business Media, 2013.
David Wozabal. Robustifying convex risk measures for linear portfolios: A nonparametric approach. Operations Research, 62(6): 2014.
Bin Xia, Yulun Zhang, Shiyin Wang, Yitong Wang. Xinglong Wu, Yapeng Tian, Wenming Yang, and Luc Van Gool. Diffir: Efficient diffusion model for image restoration. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pp. 13095-13105, 2023.
Binxin Yang, Shuyang Gu, Bo Zhang, Ting Zhang, Xuejin Chen, Xiaoyan Sun, Dong Chen, and Fang Wen. Paint by example: Exemplar-bar based image editing with diffusion models. In Proceed- ings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 18381-18391, 2023a.
Ling Yang, Zhilong Zhang, Yang Song. Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: A comprehensive survey of methods and applications. ACM Computing Surveys, 56(4):1-39, 2023b.
Ruofeng Yang, Bo Jiang, Cheng Chen, Baoxiang Wang, Shuai Li, al. Few-shot diffusion models escape the curse of dimensionality. Advances Neural Information Processing Systems, 37: 68528-68558 2024.
16

---

Published as a conference paper at ICLR 2026
Serin Yang, Hyunmin Hwang, and Jong Chul Ye. Zero-shot contrastive loss for text gext-gguied diffusion image style transfer. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 228822 2023c.
Chaojian Yu, Bo Han, Li Shen, Jun Yu, Chen Gong, Mingming Gong, and Tongliang Liu. Under- standing robust overfitting of adversarial training and beyond. In International Conference on Machine Learning. pp. 25595-25610 PMLR. 2022.
Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. Lsun: Construction of large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv::506. 03365, 2015.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy train strong classifiers with localizable features. In Proceed- ings of the IEEE/CVF international conference on computer vision, pp. 6023-6032, 2019.
Han Zhang, Zizhao Zhang, Augustus Odena, and Honglak Lee. Consistency regularization for generative adversarial networks. In International Conference on Learning Representations, 2020.
Zhagy, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical
risk minimization. In International Conference on Learning Representations, 2018.
Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control text-to-image diffusion models. In Proceedings of the IEEF/CVF international conference on computer vision, pp. 3836-3847, 2023a.
Yuxin Zhang. Nisha Huang, Fan Tang, Haibin Huang, Chongyang Ma, Weiming Dong, and Chang- sheng Xu. inversion-basee style transfer with diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 101 10146-10156, 2023b.
Zhaoyu Zhang, Yang Hua, Guanxiong Sun, Hui Wang, and Seán McLoone. Training diffusion- based generative models with limited data. In Forty-second International Conference on Machine Learning. 2025.
Shengyu Zhao, Zhijian Liu, Ji Lin, Jun-Yan Zhu and Song Han. Differentiable augmentation for data-efficient gan training. Advances neural information processing systems, 33:7559-7570, 2020.
Yunqing Zhao, Keshigeyan Chandrasegaran, Milad Abdollahzadeh, and Ngai-Man Man Cheung. Few-shot image generation via adaptation-aware kernel modulation. Advances in Neural Infor- mation Processing Systems, 15: 2022.
Zhengli Zhao, Sameer Singh, Honglak Lee, Zizhao Zhang, Augustus Odena, and Han Zhang. Im- proved consistency regularization for gans. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 11033-11041, 2021.
Jingyuan Zhu, Huimin Ma, Jiansheng Chen, and Jian Yuan. Few-shot image generation with diffu- sion models. arXiv preprint arXiv:221 1. ..3264, 2022.
Yuanzhi 7hu, Kai Zhang, Jingyun Liang, Jiezhang Cao, Bihan Wen, Radu Timofte, and Luc Van Gool. Denoising diffusion models for alu--and-play image restoration. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 1219-1229, 2023.
17

---

Published as a conference paper at ICLR 2026
A RELATED WORK
Recent advances in generative modeling have been driven by diffusion models, which have achieved state-of-the-art performance across a wide range of applications. However, their effectiveness in limited data settings remains a major challenge, as models often suffer from overfitting. To address this issue, prior works have explored strategies such as data augmentation and few-shot adaptation. In parallel. the framework of Wasserstein distributionally robust optimization (WDRO) has emerged as a powerful tool for mitigating overfitting by optimizing against worst-case perturbations of data distributions. In this section, we review related work along three directions: diffusion models, generative modeling under limited data, and WDRO.
Denoising diffusion probabilistic models (DDPM). In recent years, diffusion models (Ho et al., 2020 [Sohl-Dickstein et al., 2015] Karras et al., 2022 Song et al. (2021b) have emerged as a state- of-the-art family of generative models. They work by sequentially corrupting training data with gradually increasing levels of noise (i.e., the forward process), and then learning to reverse this conrppiion to construct generative model of the data (i.e., the reverse process). Current research on diffusion models has primarily focused on two main formulations: denoising diffusion prob- abilistic models (DDPM)(Ho et al. 2020; Nichol & Dhariwal 2021) and score-based stochastic differential equations (Score (DEs)(Song et al., 2021b; Karras et al. 2022) (where score-based generative models (SGMs) (Song & Ermon 2019 (020)) be viewed as their discrete coun- terparts). Given data point x(0) ~ Pdata the forward process generates a sequence of ran- dom variables {x(1),.... (..)...,,(())) with the transition kernel p(x(t) | x(t - 1)) for all timesetp t E 10. T}. A common choice for the transition kernel is Gaussian kernel (Yang et al,, 2023b), i.., p(x(t) x(t - 1)) N(x(t): Exx(t 1), 3,I), where 8, (0. is se- quence of positive noise scales. Following SShl-DDikstein et al. (2015); 5): Ho et al. (2020). with setting 0; == 1 3e and ā₁ : II' as we have p(x(t) x(0)) = ((x(t); vā,x (1 I.... Therefore, we can easily obtain a sample of x(t) by sampling a Gaussian vector E ~ I0,
and applying the transformation x(t) (a)x(0) + /1 â,e. Since the noise scales ā are pre- scribed (Song et al., 2021b), so that x(T) is almost Gaussian in distribution, i.e., x(T) ~ NN I).
The reverse process
variational Markov chain and parameterized with px(x(t - 1) x(t))
N(x(t 1):
vI-3, (x(t) + (x))( f)), 8,I). Thus, the loss takes the following form (see Song
et al. (2021b) for details):
((0,x, = x(t) ||5 t) 100 log Pe (x)
(15)
where X(t) is a positive weighting function YYng et al. 2023b).
Generative models with limited data. Prior to the rise of diffusion models, a large body of work
studied training schemes for generative models in limited data settings, primarily in the context of
GANs (Abdollahzadeh et al. 2023). A significant challenge in this scenario is "overfitting" (Kaatii
et al. 2020 Liu et al.) (2211), where the model may memorize the training data (Li et al. 2020 Ojha et al. 2021) and reproduce training examples rather than learn the real data distribution Zhao iet al.) (022) Moreover, under limited data regimes, generative models are more prone to mode collapse (Tran et al. (022). i.e., the models learn only a limited set of modes and fail to capture other modes of the data distribution, resulting in limited diversity in generated samples Yu et al.) 2022]. Various strategies have been proposed to mitigate this phenomenon, primarily focusing on data augmentation (Zhang et al. 2020 Zhao et al., 2022 Karras al. 2020 Chen al. |02 Wang iet al. (202b)). which increases the quantity and diversity of the training data. For example. the ADA method (Karras al. 2020) applies an adaptive augmentation strategy (i.c., with augmentation probability < 1) in the limited data setting to prevent information leakage. Mcanwhile, recent works have also begun exploring the few shot adaptation of diffusion models (Lu et al., 2023 Ruiz et al. 2023 Zhang et al. [2025]. For instance, DreamBoott (Ruiz et al. (2023) finetunes a pretrained text-to-image model on few images of a specific subject and introduces a special identifier token in the prompt, enabling the finetuned model to generate diverse images that preserve the subject's identity. However, these works do not fully explore training diffuusion models from scratch under limited data, and they differ drastically from our proposed method.
Wasserstein distributionally robust optimization (WDRO). WDRO (Rahimian & Mchrotra 2019a; Wang et al.] (225b) is an effective optimization framework for learning and decision- making under uncertainty (Wozabal 2014| Rahimian & Mchrotra 2022] Kuhn et al.) 2019] Huang et al.) 2025a). The core idea of WDRO is optimize the worst-case expected loss over Wasserstein
18

---

Published as a conference paper at ICLR 2026
uncertainty set (also known an ambiguity set) of plausible distributions, rather than single em- pirical distribution (Rahimian Mehrotra 2019a). Previous approaches to distributional robustness have considered fintte-dimeniinal parametrizations for the uncertainty set, such constraint sets for moments, support, or directional deviations (Chen et al. 2007 Delage & Yel
2010; Goh & Sim
2010), well non-parametric distances for probability measures, such as
-divergences (e.g.-
x divergence, a-divergence, and Sullback-Leible divergence) (Ben-Tal et al.
2013: Duchi et al.
2021 Namkoong & Duchi 2016) and Wasserstein distances (Blanchet et al. 201 9a b Mohajerin Es-
fahani & Kuhn 2018 Gao & Kleywegt 2023 Huang et al. (022). WDRO has been successfully
applied to numerous problems machine learning. including jeppe-suup learning (Blanchet
& Kang 2020; Chen & Paschalidis 20188) adversarial training (Levine & Feizi] 2020 Najafi et al.
2019
Sinha et all||||00000 (20])) Staib & Jegelka 2017] Liu et al. 2025). ), reinforcement learning (Liu al.,
2022 Abdullah et 2019), and transfer learning (Volpi et al. 2018 Lee & Raginsky. 2018). Re-
cent work has also investigated the incorporation of DRO into diffusion models. For instance, Wang et al. (2025a) employ DRO to mitigate the distribution mismatch that arises between the training and sampling procedures. In contrast, Our work differs substantially in both problem setting and DRO formulation: we focus on limited data diffusion training, and we design WDRO method (i.e., implemented via a "Bi-level Interval Update" strategy) on the original data distribution to expand support and mitigate overfitting.
B PROOF
In this section, we provide a detailed proof of the convergence result in Section 3.2 In Section B.
we establish an upper bound on the worst-case objective (2) (i.e., Lemma 3.5). In Section B.2
we present several auxiliary lemmas that are directly used in the proof of Theorem 3.6 Finally, in
Section B3.3] we establish the convergence result of the WILDDDiiffusio algorithm (Theorem 3.6
In addition, we provide the convergence result from Lee et al. (2022) for comparison in Section B.4
B. PROOF OF LEMMA 3,5
Lemma B.1 Under Assumption 3.3 for any fixed r > 0, the following inequality holds with proba- bility at least 1 - e uniformly over all 2 0 and 7 2 0
sup p: PS
Eple(0:x, S Yp + Ep.. 1 (0:x, +
(16)
Proof The proof follows (Sinha et al. following duality result holds for problem
2018 For any data distribution Pdata and p > 0, the
2):
sup W. (p.Paa
Ep [(0: x, t)]
inf {yp + Epan 1o (0: x, f)]} 720
(17)
SSp
From the above duality result 17 for all
0, data distributions Paltas and > 0, we have
sup Ep x, t) S 7p + Epman (6; x, (f;,,t)]. t)]. p: DSp
(18)
Let dx denote the point mass x. We first present the empirical result for Eq. (ab:
minimize (0) = sup
(E,]((0:xx (
(P. E,p., [o,, x, 11]
(19)
1 8x, denotes the empirical distribution of the samples Xim- Next, we show that E. (o_ (0: x, t)] concentrates around its population counterpart at the standard rate (Boucheron et al. 2005).
where Pn En
Since we assume that the loss function {(e: x, t) uniformly bounded by Assumption 3.3 i.e., |e(0; x 0) 5 Together with the definition of the surrogate loss, we have that
- < {(0;x, x, t) < Φ,, (0: t) I sup {E(0; ((0;xxt)) t)} S
and hence |d, (0: x, t) (Boucheron al. 2005)
< E. Thus, the functional 0 + (,) satisfies the bounded differences
19

---

Published as a conference paper at ICLR 2026
Note that bound relies the usual covering numbers for the model class ((0:-) : O₂ measure of complexity (Wellner et al. 2013). where 21 denotes the parameter space. Recall the definition of covering numbers: for set V. collection {Ui,-- ( UN is an e-cover of V in norm I| if for each v e V, there exists e2 such that ||u v/ S 6. Then the covering number of V with respect || is
e,,,,, Il (I) == inf there exists an e-cover of with respect to ||.
For our problem, let == E((:.) : 0 E θ₁ denote the loss function class equipped with the Lo (2) norm, i.e.
| == sup [((x)), & C xeX
therefore the covering number of C is Et, (l
By applying standard results on Rademacher complexity (Bartlett & Mendelson 2002) and entropy integrals (Wellner al.] 2013), have that for any fixed > 0, the following inequality holds with probability at least
c'
/logN(C, (lr..) de + bo V T (20)
y (0: t)] S p.. (f:x,t))
n
where b,, b2 > 0 are absolute constants.
Substituting Eq. into Eq. 18
Ep(((0;x,0) & 7 Ep, (0:x,)) +0()))
sup PP
B.2 AUXILIARY LEMMAS
For analytical convenience, we consider the following discretization and approximation of Eq. (5) which can be expressed as
(f+1)h
((+1)A = Xih
f(xin: T - () (T - =12 so (xin, T -ihh) dt,
(T))
ih
where h denotes the step size with T = kh (and is the number of steps), and time is reversed such that in the reverse process corresponds to (T - 1) in the forward process. Following Lee et al. (2022), proof method is construct "bad set",, which is formalized in lemma Specifically, we define bad set By as the set of X for which the worst-case error is large (see Eq. 24). Let denote the discretized process ( with the estimated score, and ÿN denote the discretized process (D)
that also uses the estimated score except in B₁. The following lemma formalizes this construction and provides the key bound needed for our analysis.
Lemma B.2 Let F,F.) be probability space and {Fi} filtration of F. Suppose X ~ Pk. 2g 9k. and 2い The are Fx-adapted stochastic processes taking values in Ω. Assume further that
if 2; E B for all 1 S < k - 1, then 2 = 2g. Under these conditions, the following results hold
DrV gx) Σ
(x + 112 6//2
(21)
Drv( S x²x pk) 1/2 1/2
(x² pb) 11) 1/22
(22)
where 0g satisfies P(ag E B&) S 8y for every k N.
Note that (-1-) denotes the x² diverrgne, and its detailed definition is provided in Section
20

---

Published as a conference paper at ICLR 2026
Proof By the definition of the total variation distance (Section D
obtain
(yy (k)
P( + Ei)
k-1
k-1 = UP k-0
U Bi)
(2) E Bi)
1m0
k-1 Σ
P(z: E B:)
By, Ln:
1=
Bp
)
i pc.
2)
1/2
(1p, LB,)
(q: ll p,) + +1) 1/2 81/2.
==0
The second inequality 22 then follows from the triangle inequality and Schcay-Schwarz:
Drv Pk) s Drv (Pt. (k) + Drv (Āk: 9k)
(xe ll pl) 1/2 + Drv Qk).
Notice that x² convergence bounds directly yield bounds on the total variation distance between the real distribution Pdata and the sampling distribution go (with i 0). We therefore recall the convergence result of Lee et al. 2022) as follows.
Lemma B.3 Lee et al. (2022 Theorem 4.3)) Let p : Rd + RR be a probability density satisfying Assumption 3.2 and let so (x, Idd 10, T - Rd be score estimator with error bounded in L for each t E [0,
||V
- (x(x l) max XERd
VInp.(x) so (x(,1)) S S1.
Let = O (max {1, log(((())) and 0 ( max max(2 =f < 1280is' then
x" (To Pdom)) = exp 16C1s
x² x20 (ddma) + O(Cise)) + Q + Ci)
(23)
Lemma B.4 Suppose that distribution p has log-Soboler constant at most Cis and satisfy Assump- tion 3.2 Then for T
x2 Pdata) = 0(1).
For detailed proof plsese see (Lee et al. 2022 Lemma E.9).
B.3 PROOF OF THEOREM 3.6
We first define a sequence of "bad" sets BiE(O,T) where the worst-case error in the score estimate is large,
B =
x R
sup P.W..PPPPa Se
Expp [||s0 (x, f) - -V pl > EB
(24)
for some EB to be chosen. Define t_ := 1 for 0. recall discretization sampling process I and define an interpolated process
[[(x_,T T -) -) 1g(T - t)² b(x,_, T t) dt,
X, = X
where
[s((x))) t),
x B.,
b(x,t)=
Vlnp((x),
x e B.
21

---

Published as a conference paper at ICLR 2026
Specifically, simulate the ODE using the score estimator Sø whenever the point lies in the good set the "previous" discretization step (i.e., at time t-), and replace it with the true gradient V In p. otherwise. Note that this interpolated process is introduced purely for analysis, since V In Pr not available in practice.
Then, applying Chebyshev's inequality (Knuth) 1997) and Lemma 3.5 to the Eq. 24
obtain
7p Sp lo (0x.0)) 1
P(B:)
(25)
Recall Lemma B.3 (Eq.
22 ) we have
x (0o
0o Pdata)
exp
I6C
x2 Páata) + ( O(L) + L3)(ish).
(26)
To ensure that this quantity bounded by ex it suffices to require
x² |l Pdam) - 4
exp
16c1s
2
CisE S 2
(L2 + L2d)Oishh S 2
Consequently, we obtain
T 32C1s log
2x p(m
s²
EB
101s
To satisfy the condition in Lemma B.3. we choose h CissC +d) max (12,12) Note that
Eq. 26 also satisfies S 22 since Cis Furthermore, by Lemma B.3 (Eq. ), we have
Drv go) S
Σ
(1 + (9 Pdaaa)) 1/2 P(BB))) PP²B))))))
1=0
1/2
Npt (0xx))),,,, /1) h)
Σ
x2 x² (ddma) | + 32C
exp
1=0
1/2
0
7pt 1 10x.1)))00
Σ 1=0
ih 320's
x2 x² Petnaa ) )/22))22++)))
exp
o
1/2
[1
7p+ 10 (0:x.)))))
G1C18 x2 Pdmaa) 1/2 O(K))
(
V11
1/2
Npt ( (0:x,1)]-00
O (max {k; Cisx (umm) 1/22]
<
67
By Lemma we obtain that x (400 | Praam)) (1) when ((log(C))))). Thus, if chosen such that
T = @ (max { (og(Ci),)) Cis log ())
22

---

Published as a conference paper at ICLR 2026
have
t (1 254
whrt Epn (0:x, (1))
Drv(go: 4o) S0
Therefore, by Eq. (22 we get
Drv (4o: Pdata)
5 (00 Pamaa) 1/2_ 1/2 Drv (go: 7o)
E + DTv(Qo: (Using Eq. (66
B.4 BACKGROUND THEOREMS
For reference, we include the convergence result of the reverse SDE (i.e., Eq. 4 with the estimated
score from Lee et al. (2022), and the detailed result is presented Lemma B.5 below.
Lemma B.5 Lee et al. (2022) Theorem 3.1) Let Patto
IRd + R be probability density satisfying Assumption 3.2 and let Pt be the distribution resulting from evolving the forward SDE according to
DDPM with = 1. Suppose furthermore that V log Pt is L-Lipschitz for every t Z 0, and that each so(,,t) satisfies Assumption 3.3
Then if
( (Cis +
STVE
=
max {log( Cis log log(1/e2)) L,, La))² (maxx
D starting from prior distribution for time T e (0maxx [log(Cis)) Cis log
( and
running
xx
step size h @ Cis(Cis (xx results in a distribution so that Dry (go, Pdiaa) S ex +
eTV.
C MORE EXPERIMENT RESULTS
In this section, we provide additional experimental results further validate the effectiveness of our proposed WILD-Diiffsion method. We begin with detailed implementation settings in Sec- Section C.2 followed
tion C.I Next, we present sensitivity analyses of key hyper-parametent by supplementary results under limited data settings in Section C.3 and few-shot generation tasks in Section 4.3 Furthermore, in Section C.5 we extend our method text-to-image generation. Finally, in Section C.6, we conduct ablation studies to examine the contributions of different components in our method.
C.I EXPERIMENTAL IMPLEMENTATION DETAILS
We developed method top of widely used codebase EDM KKrras et al. 2022). We imple- mented and trained our model with Py PoTch on a 64-bit Linux machine with 8 NVIDIA A100 (80G) GPUs. As described in the experimental setting in Section 4 method is built upon three differ- ent models: DDPM++ (Song et al.) (021b), NCSN++ (Song et al. 2021b), and ADM (Dhariwal & Nichol 2021). Specifically, we highlight the architectural differences among these three models, as illustrated in Table 3 In addition, we provide the detailed training configurations in Table 4
C.2 EXPERIMENT RESULTS FOR SENSITIVITY OF
Following the sensitivity analysis in Section 4.1 we present the FID and computation time across different settings of the interval parameter creasing m reduces the total training time but also degrades generative performance (higher FID), revealing clear trade-off between efficiency and quality. Considering both training efficiency and generative quality, we set m = 20 as the default choice in all experiments.
50% FFHQ datasets. As shown in Figure 3 in-
Additionally, also perform sensitivity analysis on the key hyperparameters of WILD-Diffusion including the number of steps K. the step size 7j, and the penalty parameter y, As shown in Figure observe that the generation performance is sensitive to the choice of the number of steps K, the
23

---

Published as a conference paper at ICLR 2026
Table 3: Details the network architectures used in this paper.
Parameter
+DPM++
NCSN+
ADM
Resampling filter
Box
Bilinear
Box Positional
Noise embedding
Positional
Fourier Residual
Skip connections in encoder
Skip connections in decoder
Residual blocks per resolution
4
3
Attention resolutions
{16}
{16}
{32, 16, 8}
Attention heads
6-9-12
Attention blocks in encoder
Attention blocks in decoder
13
Table Hyperparameters used for the training runs in Section
Datasets
Duration (Mimg)
Minibatch size
CIFAR-10 (20% 100%)
200
1024
10c 5/10e
FFHQ & CelebA-HQ
200
512
26 -1
CUUN-Church
200
256
1e-4
step size 77, and the penalty parameter 7. Specifically, too few steps or a very small step size leads to weak perturbations, which reduces the effectiveness of WIL.D-Diffusion, while overly large values introduce instability and degrade the FID. Similarly, the penalty parameter ? controls the trade-off between perturbation strength and stability, where extreme values yield suboptimal results. Based on this analysis, we set the default configuration as K = 5, 0.01, and 7 =
5.1
5.0
Time
1.6
1.4
4.8
1.2
1.0
20
40
100
parameter
Figure 3: Sensitivity to the interval parameter m. Bars show FID (lower is better); the black line shows training time (normalized). Increasing (less frequent WDRO updates) reduces time but degrades FID, which reveals a trade-off between efficiency and quality.
C.3 MORE EXPERIMENT RESULTS FOR LIMITED DATA SETTING
In this section, we compare our method with state-of-the-art diffusion approaches on the high- resolution benchmark LSUN-Church (256 x 256). Namely, the baselines include DDPM (Ho al, 2020 DDIM (Song al. 2021a), DeepCache (Ma et al.) 2024). and EDM-ADM (Karras et al. The results are summarized in Table 5 and suggest that our method can outperform the baseline models. We further include Figure 5 to separately visualize the off-distribution samples produced by our bi-level update. addition. Figures 6 and present generative samples from WIL Wiff-Diffusion trained CIFAR-10, FFHQ, CelebA-HO and LSUN-Church.
2022
24

---

Published as a conference paper at ICLR 2026
0
0.001
6.0
7.0
5.4
0.001
1000
0.01
3.3
5.0
0.1
0.005100
0.1
7
5
0.05
(a) The number of steps
(b) Step size
(c) Penalty parameter
Figure 4: Sensitivity analysis of WILDDDDiffsson with respect to: (a) the number of steps K, (b) the step size y, and (c) the penalty parameter
Table 5: A comparison of FID between WILD-Diffusion and other diffusion models the LSUN- Church (256 x 256) dataset. The best results are highlighted in bold.
Data size
Methods
20%
50%
100%
DDPM (Ho al.
2020
7.89
DDIM (Song et al.
2021a
10.58
DeepCache (Ma al.
2024
11.31
EDM-ADM (Karras el al. + WILDD-Difrusion
2022
7.74
S.79
4.66
6.98 (-9.82%)
5.13 (-11.40%)
4.47 (-4.07%)
Figure 5: Off-distribution samples generated by the bi-level update.
C.4 MORE EXPERIMENT RESULTS FOR FEW-SHOT GENERATION
In Table 6 we report the FID results of WILD-Diffusion on the 100-shot, Animal-Face, CelebA- HQ. and LSUN-Cat datasets using GAN architecture. As described in Section $ we adopt the pre-trained StyleGAN-v2 (Karras et al. 2019), trained on the FFHQ dataset, as the source model. We compare method with GAN-based approaches for limited data generation, including Dif- fAugment (Zhao et al., 2020), ADA (Karras et al. 2020), and MAFP (Zhang et al.) (025). The results suggest that method achieve the lowest FID scores across all datasets. In addition,
25

---

Published as a conference paper at ICLR 2026
(a) Samples generated on CIFAR-I0 (32 x 32) with 20% of training data. 12.14.
(b) Samples generated on CIFAR-10 (32 x 32) with 50% of the training data. FID 6.02.
(c) Samples generated on CIFAR- 10 (32 x 32) with 100% of the training data. FID = 1.93.
Figure 6: Samples generated on CIFAR-10 (32 32) using different proportions of the training data with EDM-DDPM+ Karras al. 2022 combined with WILD-Diffusion.
26

---

Published as a conference paper at ICLR 2026
(a) Samples generated on FFHQ (64 64) with 20% the training data. FID 8.57.
(b) Samples generated on FFHQ (64 64) with 50% of the training data. FID = 4.68.
(c) Samples generated on FFHQ (64 x 64) with 100% of the training data. FID 2.53.
Figure 7: Samples generated on FFHQ (64 64) using different proportions of the training data with EDM-DDPM+ KKrras al. 2022 combined with WILD-Diffusion.
27

---

Published as a conference paper at ICLR 2026
(a) Samples generated on CelebA-HQ (64 x 64) with 20% of training data. 10.22.
R
(b) Samples generated on CelebA-HE HQ (64 x 64) with 50% of the training data. FID 5.55.
(c) Samples generated on CelebA-Ho HQ (64 x 64) with 100% of the training data. FID = 3.63.
Figure& 8: Samples generated on CelebA-HQ (64 64) using different proportions of the training data with EDM-DDPM++ (Karras et al. 2022) combined with WILD-Diffusion.
28

---

Published as a conference paper at ICLR 2026
(a) Samples generated on LSUN-Church (256 x 256) with 20% of the training data. FID 6.98.
shup
(b) Samples generated on LSUN-Church (256 256) with 50% training data. FID = 5.13.
(c) Samples generated on LSUN-Church (256 x 256) with 100% of the training data. FID 4.47.
Figure 9: Samples generated on LSUN-Church (566x256) using different proportions of the training data with EDM-ADM (Karaas et al. 2022) combined with WILD-Diffusion.
29

---

Published as a conference paper at ICLR 2026
provide generative samples from WILD-Diffusion in both pretrained and non-pretrained settings in Figure 10
Table 6: The FID results few-shot generation with GAN architecture. Following the setting used in (Zhao et al. 2020), we calculate the FID with 5k samples and the training dataset adopted as the reference distribution. When FFHQ and LSUN-CA are used the target datasets, the number of larget domain images is 2k. The numerical results of the baseline methods are quoted from their papers. We highlight the best results in bold.
FFHQ
100-shot
FFHQ
Animal-Fare
CelebA-HQ FFHQ
FFHQ LSUN-CaL
Methods
Ohama
Grumpy
Panda
Dog
DiffAugment (Zhao et al.
2020
46.87
27.08
12.06
42.44
58.85
11.20
20.18
ADA Karras al. 2020
45.69
26.62
12,90
40.77
56.83
10.08
19.34
MAFP Zhang 2025
25.87
9.67
17.93
Ours
40.02
24.97
10.52
37.66
54.03
8.53
16.28
(a) Few-shot generation results of our method without pretraining.
(b) Few-shot generation results of method with pretraining.
Figure 10: Few-shot image generation results of our method on 100-shot and Animal-Face datasets, shown in both (a) non pretraiied and (b) pretrained settings.
C.5 EXPERIMENTAL RESULTS FOR TEXT-TO-IMAGE
In this section, we evaluate whether our method generalizes to conditional diffusion models by test- ing WILD-Diffusion on standard text-to-image personalization task. We adopt DreamBooth (Ruiz et al.) 2023) as the baseline, where the goal is to generate images of target concept from text prompts using only a handful of reference images. Following the experimental setup of Ruiz et al.) 2023 evaluate performance using three standard metrics: PRES (lower is better), DINO simi- larity (higher is better), and CLIP-I similarity (higher better). We consider DreamBooth basclines following the Imagen-based implementation, with and without the prior preservation loss (PPL.),
30

---

Published as a conference paper at ICLR 2026
while keeping all other hyperparameters identical. As shown in Table 7 WILD-Diffusee improves these metrics over DreamBooth under both settings (with and without PPL), which suggests that our method can enhance generation quality in large-scale text-conditioned diffusion models.
Table Text-to-image results on the DreamBooth dataset. We highlight the best results bold.
Methods
PRES 4
DINO
CLIP-I
DreamBooth (w/ PPL.) (Ruiz et al.
2023
0.493
0.684
0.81.
WIl D-Difussion (Ours)
0,478
0.696
0.823
DreamBooth (w/o PPI.) (Ruiz 2023
0.664
0.712
0.828
WIL.D-Diffsiion (Ours)
0.643
0.715
0.830
C.6 EXPERIMENTAL RESULTS OF THE ABLATION STUDY
To further investigate the role distributional divergence, we compare our method based on Wasserstein distance with variants using KL-divergence, >=-divergence, and a-divergence. As illus- trated in Figure III Wasserstein distance consistently outperforms the alternatives across different data sizes (20%, 50%, and 100%). Notably, the improvement seems most evident in the low-data regime, indicating that the Wasserstein distance may play a role in stabilizing training under limited data.
12
Wasscrstein distunce
Ki-diveeggnce
2i-diverpmoee e divvergmmee
FID (FFHQ)
20
100
Data Size (%)
Figure 11: Ablation study different distributional divergences for limited data generation the FFHQ dataset. FID (lower is better) reported under varying data sizes.
In addition, compare our proposed WILD-Diffusion with commonly used augmentation tech- niques, including Mixup (Zhang et al., 2018). CutMix (uun et al. 2019), and CutOut (DeVries & Taylor (017). based on the EDM-DDPM++ baseline (Karras al. (022).. For these methods, we follow the default hyperparameters used in the original papers, which are also the standard config- urations adopted in prior generative modeling work (Zhang et al. 2025). Specifically, we set the
interpolation strength to a 1 for Mixup and CutMix, and use 16 x 16 mask size for Cutout. The results are summarized in Table which shows that WILD-Diffusion achieves the best performance among all compared approaches.
Table 8: Ablation study data augmentation methods for FFHQ generation with 20% training data. Results are reported in terms of FID using 50% samples.
EDM-DDPM++ Karras et
2022
10.02
WILD-Diffusion
8.57
Mixup Zhang et al. 2018
10.21
Cutmix Yun et al.
2019)
10.43
Cutout De Vries Taylor
2017
10.25
We further analyze the computational efficiency of our method by measuring the relative running time under different data sizes (20%, 50%, and 100%). As summarized in Table the training time
31

---

Published as a conference paper at ICLR 2026
of the baseline EDM-DDP DDMD+ (Karras al. 2022) normalized to "1". The results show that our method maintains almost identical running times across all data sizes, with values ranging from 1.20 to 1.22. This indicates that the performance gains of WILD- HLLD-Difusion come at negligible additional computational cost, therehy ensuring both effectiveness and efficiency.
Table 9: Training time analysis under varying data sizes (20%, 50%, and 100%). The training time of the baseline EDM-DDPM++ [Krraa al. 2022) is normalized to "I".
Data size
Training time
20%
1.20
50%
1.22
100%
1.21
To provide a comprehensive analysis of computational overhead, we report wall-clock time, FLOPs, and peak GPU memory for both EDM-DDPM+ and WILD-Diffusion in Table WILD-Diffusion increases training time by 1.21 x and FLOPS by 1.25 x, while peak GPU memory increases by only 3%. These results indicate that our method adds minimal overhead relative standard training.
Table 10: Comparison of training cost between EDM-DDPM++ Karras al. 2022 and WILD-
Diffusion across wall-clock time. FLOPs, and peak GPU memory.
Methods
Wall-clock time (h)
FLOPs (G)
GPU memory (GB)
EDM-DDPM++ Karras et al. 2022
26.4
137
16.32
WILD-DDiflusion (Ours)
31.9 (1.21x)
172 (1.25x)
16.84 (1.03x)
To further understand how the computational cost scales with the hyperparameters, we additionally examine the effects of the interval parameter m and the number of inner ascent steps K. shown in Figure 12] increasing reduces FLOPs rapidly and then stabilizes, since "Bi-level Interval Update" occur less frequently. In contrast, increasing leads "near-linear" growth FLOPs due ad- ditional forward-backward passes. These results highlight the controllable computational behavior of WILD-Diffusion.
1.5
FLOPs (normalized
2.0
TsOPs (normalized)
1.4
1.8
1.6
1.3
1.4
1.2
1.1
1.0
1.0
20
40
60
80
100
10
(a) Effect of the interval parameter m FLOPs.
(b) Effect of the number of inner steps K on FLOPs.
Figure 12: Overall comparison FLOPss under different support expansion configurations. (a) Relationship between FLOPs and interval parameter (b) Relationship between FLOPs and the number inner steps K. The FLOPs of the baseline EDM-DDPM++ (Karras et al. 2022 is nor- malized to "1".
32

---

Published as a conference paper at ICLR 2026
D USEFUL FACTS
In this section, we collect some facts used throughout the paper.
Definition D.I (f-divergence). Let f : (0, oo) + IR be a convex function with f(1) - 0. Let P and Q be two probability distributions on a measurable space F). If is defined as
Q then the f-divergence
dP dQ) ]
(P())) & EQ
where aQ
dP is Radon-Nikodym derivative and f(0) - f(0+). Suppose that Q(dx)= = ((z))(dx)
and P(dz)= = ((x))(dr) for some common dominating measure p. then we have
p(x))
Dy(PIIQ) = ((x))
du
((x)
The following common faaanmonnndvveeee f-divergences that used in this paper:
1. Kullback-Leibler (KL) divergence: f(x) x 2logæ.
DRL. (kl(())))) A p(x))og log p(x)
g(x) u(dx).
2. x2 2 --dvvergence: f(x) = (xr - 11)
(dP
-1) 1) 27 /
dp2
(P(11) & EQ
-1
dQ
dQ
3. Total variation: f(x) = 1/10 - -1/,
1
dP
1
DIv(P.Q)
EQ
1
|P -QQ
dQ
2
4. a-divergence WWag al. 2018at: f(x)
=--1
R {0, 1} and hence
a(a-1)*
(dPP) dQ
dP
-1
(,(I)) a(a - 1) BQ
Definition D.2 (log-Son (log-Sobolev inequality (Vempala & Wibisono 2019 Let P be a probability mea- sure with density p. We say that p satisfies log-Soboles inequality with constant Cis for any probability measure 4.
KL(q p) < Cis
[|
log alw) p(a)
q(x) da.
33


Reviews:
[{"review_id": "Q6k8zoxuge", "reviewer": "Reviewer_vXtY", "summary": "The paper proposes WILD-Diffusion, a novel framework that adapts Wasserstein Distributionally Robust Optimization (WDRO) to diffusion models to promote robust training under limited-data regimes via  constructing worst-case distributions that are close to the limited data distribution which effectively expands the support of the training data and mitigates overfitting. Experiments on various benchmark datasets show consistent performance improvements, especially when data are scarce (e.g., N=100). Unlike traditional adversarial training, WILD-Diffusion is a plug-and-play procedure that requires no architectural changes and can be used with or without pretraining.", "questions": "Refer to the above.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper tackles a timely and relevant problem—improving diffusion model performance with limited data—supported by both theoretical grounding and extensive experiments. \n- The overall presentation is clear, and the method is well motivated with sound empirical results.", "weaknesses": "- **Incomplete reporting of computational overhead** The paper reports a *relative* training-time increase of 1.20–1.22× compared to the baseline (Table 8), but does not provide absolute wall-clock time, FLOPs, or GPU memory usage. Since the method targets data-scarce scenarios where efficiency is important, providing these absolute cost measurements would significantly strengthen the practical relevance. Additionally, the paper should include scaling trends with respect to the inner ascent steps *K* and interval *m*. Table 8 also shows minimal overhead variation across dataset sizes—please clarify what the dominant bottleneck is and which component contributes most to runtime.  \n- **Mismatch between baselines and problem focus** The paper frames the core contribution as a *data-efficient training strategy for diffusion models*, yet some baselines are orthogonal to this goal. For instance, DeepCache (Table 1) focuses on *inference-time acceleration*, not robustness to limited data, and Patch Diffusion primarily aims to *reduce architectural training cost*, rather than address data scarcity. It would help to clarify whether these methods are included to demonstrate orthogonality or to claim direct superiority. In the current form, these comparisons blur the paper’s main message.  \n- **Insufficiently controlled augmentation baselines** In Table 7, WILD-Diffusion is compared with basic augmentation methods such as Mixup and CutMix, which perform even worse than the vanilla model. However, there is no indication that these techniques were tuned (e.g., optimal mixing ratios, label smoothing, class-aware mixing, or dataset-specific configurations). Without proper tuning and control, these results may underestimate augmentation baselines and overstate WILD’s relative advantage.  \n- **Extension to other domain tasks** Experiments are primarily conducted on low-resolution unconditional image generation. To better support claims of broad applicability, the method should also be validated on more practical and diverse settings such as class-conditional generation or text-to-image personalization (e.g., DreamBooth-style few-shot fine-tuning). Evaluating WILD-Diffusion on pretrained text-to-image models under data-scarce adaptation would demonstrate stronger real-world relevance.", "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": "YqXdNghGHD", "reviewer": "Reviewer_hDug", "summary": "The paper proposes WILD‑Diffusion, a WDRO‑inspired, plug‑and‑play training framework for diffusion models under limited data. It extends Wasserstein Distributionally Robust Optimization to train diffusion models. Furthermore, the authors provide a convergence guarantee despite the coupling between the diffusion process and WDRO. Empirically, WILD‑Diffusion shows competitive performance in both fine-tuning and from-scratch setting.", "questions": "1. Overhead of bi‑level interval updates.\n- How does the proposed bi‑level interval update compare to standard full fine‑tuning in terms of wall‑clock time, training steps, and memory/compute overhead? A quantitative breakdown would help practitioners assess the cost–benefit trade‑off.\n\n2. Comparison to attention‑only fine‑tuning under limited data.\n- Prior work on fine‑tuning diffusion models with limited data [1] suggests that attention‑only fine‑tuning can prevent overfitting while remaining competitive. Has WILD‑Diffusion been compared against such parameter‑efficient baselines, and if so, under identical data budgets and evaluation protocols?\n\n3. Large scale model experiments on limited data\n- The proposed method demonstrates strong effectiveness under limited-data conditions, but its impact could be further amplified by including experiments on large-scale text-to-image diffusion models. Showing that WILD-Diffusion maintains or improves few-shot performance in such settings would better highlight its generality and practical relevance.\n\nReferences\n\n[1] Moon et al., Fine-tuning diffusion models with limited data", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Theory‑grounded, practical framework with broad applicability.\n- Beyond heuristic data augmentation, the approach is anchored in WDRO and comes with a convergence guarantee. It shows strong performance in limited‑data regimes in both from‑scratch training and fine‑tuning scenarios, accompanied by ablations that indicate robust gains.", "weaknesses": "1. Clarification of few‑shot comparisons and OOD adaptation.\n- In Section 4.3 (Experiments on Few‑shot Generation), Table 2 compares fine‑tuning settings with pretraining, but the original pretrained model’s performance is not shown alongside. Moreover, the FFHQ‑pretrained fine‑tuning appears in‑distribution for faces; it would be informative to evaluate out‑of‑distribution adaptation, for example transferring to a stylized or non‑face domain such as a Pokémon dataset, to assess robustness under distribution shift.", "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": "ZQ5DAjqzuT", "reviewer": "Reviewer_GEiA", "summary": "This paper addresses overfitting when training diffusion models with limited data, observing U-shaped FID curves and reduced diversity as training proceeds on small subsets.\nConcretely, a training framework inspired by Wasserstein Distributionally Robust Optimization (WDRO) is proposed, termed WILD-Diffusion,  which expands the support of the training distribution by iteratively generating worst-case samples within a Wasserstein ball around the empirical data and mixing them into training.\nTo keep this tractable, the WDRO inner max is dualized into a surrogate loss and optimized with a bi-level interval update (periodically refreshing adversarial samples, then standard parameter updates), and the paper proves a convergence guarantee tailored to the WDRO+diffusion setting.\nExperiments on CIFAR-10, FFHQ, CelebA-HQ, LSUN-Church and several backbones report >10% FID reductions with only 20% data and SoTA FID with as few as 100 images, both with and without pretraining.", "questions": "1. How is the robustness radius selected across datasets and data regimes?\n\n2. How closely does the dual surrogate approximate the true worst-case set?\n\n3. What are the wall-clock/GPU-hour and peak-memory overheads of the adversarial refresh loop relative to standard training?\n\n4. Do the improvements extend to high-resolution generation and text-to-image settings, and what changes (if any) are needed to the loss, model, or optimization?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper investigates an interesting and important problem for the research community. The proposed approach is technically sound, viewing data scarcity through a distributionally robust lens and converting the push for more data into adversarial, Wasserstein-bounded augmentations that expand support without drifting off-distribution. The optimization is practical: a tractable surrogate objective with a simple bi-level refresh of worst-case samples that plugs into standard training loops and works with or without pretraining. Theoretical analysis supports the design, and experiments show meaningful FID gains in very low–data regimes, demonstrating utility beyond incremental tweaks. Overall, the method is conceptually clean, easy to adopt, and delivers strong sample-efficiency improvements in settings where diffusion models typically overfit.", "weaknesses": "1. While the dual surrogate makes the WDRO objective tractable, it can mis-specify the true worst-case set and depends on strong-duality assumptions.\n\n2. The convergence proof requires Lipschitz/log-Sobolev–style assumptions for data and score error bounds; these may not hold for real high-res image manifolds.\n\n3. The bi-level update generates **worst-case** samples and mixes them with real data, but the paper doesn’t quantify how often those samples are usefully hard versus toxic.", "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": "tL5ti1KepB", "reviewer": "Reviewer_rEDT", "summary": "It is known that training diffusion models requires a lot of data. To mitigate this, the paper proposes a training method using Wasserstein Distributionally Robust Optimization (WDRO) called WILD-Diffusion. Using DRO, the paper can train the model to be more fitted on not only given sharp limited training data, but also near smooth distribution by expanding the support of the data. To do this, the paper uses the two-levels update strategy one for updating parameters and the other for updating samples. They also provide convergence analysis for the proposed method.", "questions": "- What are the key differences between this work and previous studies that applied DRO to diffusion models? (e.g. Improved Diffusion-based Generative Model with Better Adversarial Robustness)\n- Is the proposed method robust even when the data manifold is highly complex?\n- How much does the training time increase due to the two-level update strategy?\n- Can the data generated by the proposed method but not belonging to the given data distribution be separately visualized?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "The paper is well written and provides rich theoretical analysis. Applying DRO to diffusion models for few-shot learning is a new idea and technically sounds. To prove the method works well, they show that the proposed method achieved high performance when the data is limited (e.g. 20%, 50% in Table 1). Also, it can be applied to existing diffusion models easily as they applied the method to various existing models such as EDM, Patch Diffusion and DeepCache.", "weaknesses": "- Applying DRO to diffusion models is already studied in previous works (e.g. Improved Diffusion-based Generative Model with Better Adversarial Robustness, ICLR 2025). While the main topics of them and this paper may different, the novelty of the proposed work may decrease.\n- The datasets used in the experiments seem limited. These dataset have relatively narrow in topic, so perturbing them may effectively expand the dataset support. However, for larger and more complex dataset (e.g. ImageNet), expanding the support may introduce degradation by off-manifold samples.\n- The comparisons with other few-shot diffusion model are only conducted on 100-shot and Animal Face datasets. In fact, since the numbers in Table 6 in the appendix show that other few-shot works also achieve comparable performance, it seems necessary to verify whether the performance gap remains consistent across a wider variety of datasets. Comparisons on other datasets such as CelebA, MNIST, and LSUN would strengthen the evaluation.", "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)