Improved Diffusion-based Generative Model with Better Adversarial Robustness
Authors: Zekun Wang
Mingyang Yi
Shuchen Xue
Zhenguo Li
Ming Liu
Bing Qin
Zhi-Ming Ma
Venue: ICLR 2025 Poster (2025)
Keywords: Generative Model; Adversarial Robustness; Diffusion Model; Distributional Robustness Optimization
Primary area: generative models

Abstract: Diffusion Probabilistic Models (DPMs) have achieved significant success in generative tasks. However, their training and sampling processes suffer from the issue of distribution mismatch. During the denoising process, the input data distributions differ between the training and inference stages, potentially leading to inaccurate data generation. To obviate this, we analyze the training objective of DPMs and theoretically demonstrate that this mismatch can be alleviated through Distributionally Robust Optimization (DRO), which is equivalent to performing robustness-driven Adversarial Training (AT) on DPMs. Furthermore, for the recently proposed Consistency Model (CM), which distills the inference process of the DPM, we prove that its training objective also encounters the mismatch issue. Fortunately, this issue can be mitigated by AT as well. Based on these insights, we propose to conduct efficient AT on both DPM and CM. Finally, extensive empirical studies validate the effectiveness of AT in diffusion-based models. The code is available at https://github.com/kugwzk/AT_Diff.

Published as a conference paper at ICLR 2025
IMPROVED DiFFUSION-BASED GENERATIVE MODEL WITH BETTER ADVERSARIAL ROBUSTNESS
Zekun Wang"), Mingyang Y Shuchen Xue Zhenguo Li Ming Liu Bing Qin!
Zhi-Ming Ma
1 Harbin Institute of Technology Renmin University of China
3 Academy of Mathematics and Systems Science, Chinese Academy of Sciences
*University of Chinese Academy of Sciences 'Huawei Noah's Ark Lab zkwang@ir www.hiitt.u........m............ edu. on yimingyange ruc. ruc.edu.co enu.cn xueshuchen1 1 7@ma 1s. accas.ac.om
ABSTRACT
Diffusion Probabilistic Models (DPMs) have achieved significant success in gener- ative tasks. However, their training and sampling processes suffer from the issue of distribution mismatch. During the denoising process, the input data distributions differ between the training and inference stages, potentially leading to inaccurate data generation. To obviate this, we analyze the training objective of DPMs and the- oretically demonstrate that this mismatch can be alleviated through Distributionally Robust Optimization (DRO), which is equivalent to performing robustness-drig driven Adversarial Training (AT) on DPMs. Furthermore, for the recently proposed Con- sistency Model (CM), which distills the inference process of the DPM, we prove that its training objective also encounters the mismatch issue. Fortunately, this issue can be mitigated by AT as well. Based on these insights, we propose to conduct efficient AT on both DPM and CM. Finally, extensive empirical studies validate the effectiveness of AT in diffusion- based models. The code is available at https: /tte::///ittttt giggtut.uu... com wwwwkk/wwwkkckccckk..k. ATT Diff
INTRODUCTION
Diffusion Probabilistic Models (DPMs) (Ho et al.
2020 Song et al. 2020 Yi et al. 2024) have
achieved remarkable success across wide range of generative tasks such as image synthesis (Dhari- wal & Nichol, 2021 Rombach et al. 2022 Ho et al. (022a).. video generation (Ho et al. 2022b
Blattmann et al. 2023), text-to-image generation (Nichol et al. Ramesh et al.) 2022] [aharia et al. 2022), etc. The core mechanism of DPMs involves a forward diffusion process that progressively injects noise into the data, followed by a reverse process that learns to generate data by denoising the noise. Unlike traditional generative models such GoNssGoodiellow et al.) 2014) or VAEs (Kingma & Welling 2013). which directly map an easily sampled latent variable (c.g., Gaussian noise) to the target data through single network function evaluation (NFE), DPMs adopt gradual denoising
approach that requires multiple NFES (Song et al., 2022, Salimans & Ho 2022 Lu et al. 2022b Ma et al. 2024). However, this thennngthnn-ddnoisinng process introduces distribution mismatch between the training and sampling stages, potentially leading to inaccuracies in the generated outputs.
Concretely, during the training stage, the model is learned to predict the noise in ground-truth noisy data derived from the training set. In contrast, during the inference stage, the input distribution is obtained from the output generated by the DPM in the previous step. which differs from the training phase, caused by the inaccurate estimation of the score function due to training (Song et al. 2021 Yi al. 2023a) and the discretization error (Chen et al./ 2022 Li et al. 2023) Xue et al.) 2024b a) brought by sampling. Such distribution mismatches are referred to as Exposure Bias, which has been discussed in auto-regressive language models (Bengio et 2015 Ranzato al. 2016).
Recently, the aforementioned distribution mismatch problem in diffusion has been also recognized
by (Ning et al. 2023 Li & van der Schaar 2024 Ren et al. 2024 Ning et al. 2024 Li al.
"qqual contribution
*Corresponding author

---

Published as a conference paper at ICLR 2025
2024 Lou & Ermon 2023). However, these studies are either rely on strong mismatch distributional
assumptions (e.g.. Gaussian) (Ning et al. 2023| 2024 Ren et al. 2024) or incur significant additional computational costs (Li & van der Schaar (024). This indicates that more practical solution to this problem has been overlooked until now. To bridge this gap, we begin with the discrete DPM introduced in (Ho et af.] 2020). Intuitively, although there is mismatch between training and inference, the distributions of intermediate noise generated during the inference stage are close to the ground-trute trutt distributions observed during training. Therefore, improving the distributional robustness (Yi et al. Namkoong 2019) Shapiro 2017) (which measures the robustness of the model to distributional perturbations in training data) of DPM mitigates the distribution mismatch problem. To achieve this, we refer to Distribution Robust Optimization (DRO) (Shapiro) 2017] Namkoong, 2019), which aims to improve the distributional robustness of models. We then prove that applying DRO to DPM is mathematically equivalent to implementing robusiness-driven Adversarial Training (AT) (Madry et al. 2018 Shafahi et al. 2019 et ct al. 2021) on DPM. Following the DRO framework, we also analyze yze the recently proposed diffusion based Consistency Model (CM) (Song) et al. 2023 Luo et al. 2023) which distills the trajectory of DPM into a model with one NHE generation. We first prove that the training objective of CM similarly suffers from the mismatch issue as in multi-step DPM. Moreover, the issue can also be mitigated by implementing AT. Therefore, for both DPM and CM, we propose to apply efficient AT (e.g.. "Free-AT" (Shafahi et al. 2019)) during their training stages to mitigate the distribution mismatch problem Finally, we summarize our contributions as follows.
We conduct an in-depth analysis of the diffusion-based models (DPM and CM) from theoretical perspective and systematically characterize its distribution mismatch problem.
For both DPM and CMM, we theoretically show that their mismatch problem is mitigated by DRO, which is equivalent to implementing AT with proved error bounds during training.
We propose to conduct efficient AT on both DPM and CM in various tasks, including image generation on CIFAR10 32x32(Krizhev)) 64 64 (Deng et al. 2009). and zero-shot Text-to-Image (T2I) generation on MS-COCO 512 (Lin et 2014b). Extensive experimental results illustrate the effectiveness of the proposed AT training method in alleviating the distribution mismatch of DPM and CM.
& Hinton| 2009) and Ima Net
al.
2 RELATED WORK
Distribution Mismatch in DPM. The problem is analogous to the exposure bias in auto rego-segressive
language models (Bengio et al. 2015 Ranzato et al. 2016 Shen et al. 2016 Rennie et al.) 2017 Zhang et al. 2019c), whereas the next word prediction (Radford et | 2019) relies tokens predicted by the model in the inference stage, which may be mismatched with the ground-truth one taken in the training stage. The similarity DPMs becomes evident due their gradual denoising generation process. [ing et al. (023) and Ning et al. (2024) propose adding extra Gaussian perturbation during the training stage or detendepe perturbation during the inference stage, to mitigate this issue. Following this line of work, several methods are further proposed. For instance, to reduce the accumulated discrepancy between the intermediate noisy data in the training and inference stages.[i [i et al. (2024) search for suboptimal mismatched input time step of the model to conduct inference. Similarly, [i & van der Schaar (2024) and Ren et al. (2024) directly minimize the difference between the generated intermediate noisy data and the ground-truth data. However, these methods either rely on strong assumptions (Ning et al. 2023 2024 Li et al. 2024 Ren et al. 2024) or are computationally expensive (Li & van der Schaar 2024). In contrast, distribution mismatch problem from the perspective of DRO. Meanwhile, our proposed AT with strong theoretical foundations is both simple and efficient, compared with the existing methods.
are the first to explore the
In this paper, we leverage the Distributionally Robust Opti- mization (DRO) (Shapiro 2017; Namkoong 2019 Yi et al., 2021 [Sinha et al.) 2018 WWag et al., 2022 Yi et al. 2023b) to improve the distributional robustness of DPM and CM, thereby mitigating
Adversarial Training and DRO.
Note that the "adversarial" here refers perturbation input training data, instead of the adversarial of gencrator-disccrmmiaato in GAN (Goodfellow et al. 2014).
'Notably, the standard AT (Madry et al. 2018) solves minimax problem that slows the training process.
The efficient AT has no extra computational cost compared the standard training ones (Shafahi al. 2019).
2

---

Published as a conference paper at ICLR 2025
the distribution mismatch problem. As demonstrated in (Sinha et al. 2018 Yi et al. 2021 Lee &
Raginsky 2018). we link the DRO with AT (Madry et al. 2018: Goodfellow et al. 2015). which
designed to improve the input (instead of distributional) robustness of the model. In supervised learning, the adversarial examples generated by efficient AT methods (Shafahi et al. 2019 Zhang
al. 2019a b Zhu al. 2020 Jiang al.] 2020) have been proven to efficient augmented data
to improve the robustness and generalization performance of models (Rebuffi et al. 2021 Wu et al.
2020 Yi et al., 2021). In this paper, we further verify that the AT generated adversarial augmented examples are also beneficial for generative models DPM and CM.
In addition, recent studies (Nie et al.) 2022; | Wang et al. 2023 Zhang al. 2023 utilize DPM to generate examples in adversarial training to improve the robustness of the classification model. This is quite different from the method in this paper, as we focus on employing AT during training of diffusion- based model to improve its distributional robustness to alleviate the distribution mismatching.
3 PRELIMINARY
Diffusion Probabilistic Models. DPM (Sohl-Dickstein al. 2015 Ho et al. 2020 constructs the Markov chain æ, by transition kernel g Øt) = =W(aat+ (V/a++13, R2 (1 Qx+1 ( where "1,""", @r gg are in [0, 1]. Let O2 11 0 and be ground-truth data. Then, for HG, it holds
Vā vI 6,6; 1 1= T,
(1)
with E ~ ~N I). The reverse process pe x++1) parameterized as
pa(x; xx+1) = N ((ue(x:1:: + 1), 1)
(2)
where 0/41 = a++1- To learn po (x; xz+1), a standard method is to minimize the following evidence lower bound of negative log- ligllikllhhodd (NLL) (Ho et al. 2020).
pe (wo.r)
Eng [log po [ S Fgg log
(3)
Here, minimizing the ELBO in the r.s.. of above incquality links to pe(x;; x++1) since it is equivalent to minimizing the following rewritten objective
min
(IcL((( (o(ar)) +
Dia((g(a m441) |I (ar mm41))
(4)
as in (Ho et al. 2020 Bao 2022 et al. 2023a). Here, the conditional Kullback- (KL) divergence D (q(2; x++1)
p(æ, ( (++1)) = q(x,, xx41) log p(a. (2++))
g( daydar (Duchi
2016), and minimizing L2 is equivalent to solve the following noise prediction problem
| (eeoVvaao + VI- ãges; 1) <all"]
min
E
We use || Ilp denote no-nomm. Unless specified, the norm || refers to the 22-norm || ||2. Since ōe + for f. + T, Dg is obtained by conducting the reverse diffusion process (a x0+11) starting from 27 ~O, I) and A'0 I), under the learned model co with
at+1
co (2 (1)) + +I- axtie.
Qr+1
#141
VI ā,-1
(6)
Wasserstein Distance. For integer p >> 0, I'(M, v) as the set of union distributions with marginal and v, the Wasserstein p-distance (Villani et al. 2009) between distributions μ and v with finite p-moments is
(H, inf vET(u,))
B(ma))yll -
ROBUSTNESS- ADVERSARIAL TRAINING OF DIFFUSION MODELS
In this section, we formally show that the success of DPM relies specific conditions, i.e., is close to It+1- Next, to mitigate the drawbacks brought by the restriction, we propose to consider the distribution mismatch problem discussed in Section and connect the problem to rewritten ELBO. Finally, we apply DRO for this ELBO to mitigate the distribution mismatch problem and finally link it to AT to be implemented in practice.
3

---

Published as a conference paper at ICLR 2025
(x4 I xx+1)
min Dxcla(Xt, Mxess Mlieoo 1x
(X I (x++1)
Xr
X2+1
Xe
Xg
Training
Standard Training
(xx xx+1)
Xp+1
1 Matching
Xo
Inference
X1 q (x))
x Pa (x+)
min (a (xt xx-1)llPa (Xte (x+1))
2 q (x+)
Adv Training
Adversarial Training
XT
Xt
20 Zo
BD (g(x)), To)
Inference
X1
Po (x.). (x₂) O
Figure 1: A comparison between standard training and the proposed distributional robust optimization
in (12). When minimizing D (ār (z++1) ) po po(#; 1)) the # is sampled from dr (3x++1)) ) such both Ge ( training stage and pe (T(41) inference stage are BD (1(*+41), No). so that po(x)) ) tends to locates in DD (g(x)), 17) as well as do (xx)). Then, the distributional
robustness captured by 12) guarantees the generated po(x)) always locates around q(x;) for all t.
4.1 How DOES DPM WORKS IN PRACTICE?
Notably, minimizing 4 potentially obtains sharp NLL under target distribution q (a))). However, in the following proposition, we show that 4 also implicitly minimizes the NLL of each æt.
Proposition 1. The minimization problem [
is equivalent to minimizing an upper bound of
lolooge po (zx for any < t
The proof is provided Appendix A shows that though 4 is proposed to generate 2n ~ g(xo), it also guides the model to generate ia such that (a,) approximates the ground-truth distribution (xx)). The conclusion is nontrivial as minimizing the ELBO of NLL Eg | log pe (2o) does not necessarily impose any restrictions æ; for
Next, we will further explain why (4) leads to small NLL of x;, In L of po xt xz+1) approximates q(x, E++1) with x++1 (( representing ground-trev -ouut data. Consequently, (e();) approximates g(x;) by recursively applying such a relationship as the following proposition.
Proposition 2. Suppose po(a; (141)) matches g(x, a141) well such that
La = Dxu(g(an m+41) po m(41)) S T
2
(8)
and the discrepancy satisfies DKI(g(wT) || (a())) S 70. for any 0 S t S T, we have
Dxı(g(Zs) po(x))) S Dxz((((x)))) po(2))1)) + L. S 70 (T - f)y
(9)
The results is similarly obtained in CCen et al. 2023 while their result is applied for D (g(()
pe.). which narrowed compared with Proposition 2 The proof is provided in Appendix A which formally explains why (4) results in po(x)) approximating (.x,). However, this proposition is built upon small L, and notably, the error introduced by L, will accumulated on the r.h.s. of (9) as it increases w.r.t. t. This phenomenon caused by the distribution mismatch problem discussed in Section Concretely, in (4. minimizing L, learns the transition probability pa (x, Øe++) based on 2041 ((#441). while in practice, æ; in (6) is generated from # ~ pe (z141)). ). The error between po(x++1) and g(E+41) will propagates into the error between po (xx) and g(x,) as in (9).
Therefore, owing the existence of distribution mismatch, only if L, is minimized, the gap between po(x;) and (xx;) can be guaranteed. However, the following proposition proved in Appendix A indicates that L, theoretically minimized with restrictions.
Proposition Lt in (4) is well minimized, only if ((((+4)) Gaussian or ||044] - xz + 0.
In practice, the (xx+11) is usually non-Gan Gaussian.. Besides, the gap | - x is not necessarily small, especially for samplers with few sampling steps, e.g., DDIM (Song al. 2022). DPM-Solver (Lu et al.) (002a)). Therefore, in practice, the accumulated error in (9) caused by the distribution mismatch problem may become large, and degenerate the quality of 200
4

---

Published as a conference paper at ICLR 2025
4.2 DISTRIBUTIONAL ROBUSTNESS IN DPM
Inspired by the discussion above, we propose a new training objective as the sum of NLLs under æt.
min (())
Σ Eg 1- po))
(10)
Then the following proposition constructs ELBOs for each of Egy log po Proposition For any distribution q satisfies ((x+) g(x4) for specific we have
Bg (ogooooo(zz)) I DxI(Ā(æ x++1) (o (e+1)) +C,
(II)
1
for a constant C independent of
The proof is in Appendix
A.2 This proposition generalizes the results in Proposition I since can be
taken q in Proposition During minimizing L, the transition probability pa #e++)) matches i(a,, 1). while # ~ ((2+++) ((xx+1) = Po (0(++1)) then in Lţ (a zz-1) matches q(æ, Ze4+1) leads pa (ax,) = g(x,) = q(x,), which mitigates the distribution mismatch problem, when minimizing such L..
1 in the training stage has no restriction. Thus, one may take
Unfortunately, for each t, obtaining such specific (xx-1) = po(x+++) is computationally expensive (Li & van der Schaar 2024). which prevents using desired (22+1)) However, know pe (22+1) is around q(z++1). Therefore, by borrowing the idea from DRO (Shapiro 2017). for each t, we propose to minimize the maximal value of L over all possible qe (x(++))) around q This leads to a small I.P", as (x(+1) locates around ((x++1), so that included in the "maximal range". Technically, the DRO-based EBLO of (11) is formulated as follows. Here po (2(44)) ) is supposed in DD (g(#:+1), mp). and it capatures the distributional robustness of pe (xx, x++1) w.r.t. input #1+1-
min
Lt
DRO (0) min
sup == d ) 13-m0)
DKu x=-1) Pa(as w(41)); (12)
s.t. (xx)) = ((x))-
Here (x++1) E BDK (9( no) means DKL (g(#)41) || (c(x)+1)) S no- By solving problem (2), the desired qt (zt+1) = po (2:+1) is in BD (( 1)),, then the conditional probability in (2) transfers Fz+11 Po po (0+++)) to target 0t ~ ((x;)) is learned, which mitigates the distribution mismatch problem. The theoretical clarification is in the following Proposition proved in Appendix
A.2 which indicates that small DRO loss guarantees the quality of generated æo-
Proposition 5. If (O) < ?oo in for all 1, and DKI(ạ(æT)) (T(T)) S Yo- then
Dкz(g(xo) pe (xo)) S 1/0-
Up to now, we do not know how to compute the DRO-based training objective 12 we derived.
Fortumately, the following theorem corresponds (12) to "perturbed" noise prediction problem similar to (5) The theorem proved in Appendix A.2
Theorem 1. There exists 6, depends on aa and e, makes 13) equivalent to problem
12
6. vI
min
Eg(wa).er
calvaroo
vI cie; 8., es
(13)
4=0
This theorem connects the proposed DRO problem (2) with noise prediction problem
13 Naturally, we can solve (13), if know the exact d.. Fortunately, we have the following proposition to characterize the range of ô,, and it is proved in Appendix A.2
Proposition 6. For 1) 0 and 6 in 13 ||6,11
S 1) holds with probability at least 2(1 ā.)/n
The proposition indicates that for any 5, depends Zo. E; in 13 it is likely in small range (measured under any f, no-norm, since they can bound each other in Euclidean space). Thus, to resolve 13 (so that (2)), we propose to directly consider the following adversarial training (Madry et al.
5

---

Published as a conference paper at ICLR 2025
2018 objective with the perturbation ô is taken over its possible range as proved in Proposition which captures the imput (instead of distribution) robustness model Eg.
8
min
Bytmo)
Engan (m))
sup 6:166|
co + v1 6. + 8) - Et
(14)
v1
We present fine- tme-grained connection between (14) and classical AT in Appendix C Notably, our objective (14) is different from the ones in (Ning et al. 2023 whereas in is a Gaussian, and co
predicts instead of e, + 6/vVI O2 as ours.
To make it clear, we summarize the rationale from DRO objective (12 to AT our objective (14)
Since Theorem 1 shows solving (12) is equivalent to (13). which conducts noise prediction ( with
perturbation ó; in a small range added (Proposition 6). Thus, we propose to minimize the maximal
loss over the possible òG, which is indeed our AT objective (14
5 ADVERSARIAL TRAINING UNDER CONSISTENCY MODEL
Although the DPM generates high-quality target data To. the multi-step denoising process 6 requires numerous model evaluations, which can be computationally expensive. To resolve this, the diffusion- based consistency model (CM) is proposed in (Song et al., (023). Consistency model fo (2;, t) transfers az ~ q(æ+)) into a distribution that approximates the target q(vn). So is optimized by the following consistency distillation (CD) loss
min Con
B2 (= (), 1))],
(15)
where &, (1 a solution specific ordinary differential equation (ODE) (37) Appendix
B which is a deterministic function transfers Zt++1 to æt, i.e., (x(+++1) g(x:), and d(xr, y) is
distance between x and y e.g., L1, &g distance.
Remark In (Song el al.) 2023 Luo al. 2023), the noisy data x, in (15) is described by ODE
(37) in Appendix B However, we use the discrete æ, ( here to unify the notations with Section 4
The two frameworks are mathematically equivalent as all æ; in ① located in the trajectory of ODE
in (Song et al. 2023). More details of this claim refer to Appendix B
Next, we use the following theorem to illustrate that solving problem with distribution close target ((z)). The theorem is proved in Appendix B
15
indeed creates fo t)
Theorem For (CD(O) (0)
15) with ) is le distance, then W Wi x0) S vtCcp 0
Though solving problem
15 creates the desired CM fe. computing the exact $, ( involves solving an ODE as pointed out Appendix B Thus, in practice (Song et al. 2023 Luo et al. 2023)
the t, (x)+1) approximated by a computable numerical estimation + (3 < of it, e.g., Fuler ( ((22) in Appendix B.T) or DDIM (Song et al. [202)), where es is pretrained noise prediction model as in Therefore, the practical training objective of (15) becomes
min Σ (0) = (e (d [ (. 0
(16)
In (16). & (x1+1, EΦ) is an estimation to Dr (Ix+1), which causes an inaccurate training objective LOD in (16). compared with target LCD (15.. Thus, this results in the distribution mismatch problem in CM, as in DPM of Section 4 However, similar to Section 4.2 if we train fo with robustness the gap between & (x cΦ) and (I(+1), the distribution mismatch problem CM mitigated.
Technically, suppose & (x = (2 Es) ) + ò (x)41) we can consider minimizing the follow- ing adversarial training objective of CM, if | 16 (x 17 uniformly over for some constant 7).
In practice, [5) is updated under target model fo ($x(2)+1))) with exponential moving average (EMA) under a stop gradient operation. (Song et al. (023) find that greatly stabilizes the training process. this section, we focus on the theory of consistency model and still use 0 in formulas.
*Here Wt (fo(z,, t), mo) the Wasserstein I-distance between distributions of fo t) and mo-

---

Published as a conference paper at ICLR 2025
Algorithm 1 Adversarial Training for Diffusion Model
1: Input: dataset D, model parameter 0, learning rate K, loss weighting X(-), adversarial steps K, adversarial learning rate a
2: while do not converge do
3: Sample a: ~ 'D and /. [
4: Sample E 10,
5: x-
6:
for
1,2, ..... K do
7:
+ cal 0, + v - c V1-& in
8:
8 + 8 +
[
V&L
maximize perturbation update model
9:
0-0 VoL
10: end for
end while
Algorithm Adversarial Training for Consistency Distillation
1: Input: dataset DD initial model parameter 0, learning rate K, pretrained noise prediction model
EΦ. ODE solver d (., Eg. metric d((,,), r), loss weighting X(-), target model EMA 14, adversarial
steps K, adversarial learning rate
2: 0 <- 0
3: while do not converge do
4: Sample x ~ 'D and t ~ u[0, -
5: Sample #+++1 from (I)
6:
8
x 0
7:
for 1,2..... K do
8:
[ ((t)d(f(()441)) 1), fo (0)( 8,t) in I7
9:
+ +a [
V.C
maximize perturbation update model
10:
- x. VoL
0 stopgrad( + (1-1)00
12:
end for
end while
so that the target 4e is included in the maximal range as well.
T-1
Adde
Bargs
sup d [a]|sn
+ +6,t),
(17)
1=0
By doing so, the learned model can be robust the perturbation brought by 6e( so that
results in a small LCD(O), as well as the small Wi(fe (ær, T"), zo) as proved in Theorem 2 Next,
we use the following theorem to show that [0++( .)l is indeed small, and minimizing results in fo (xT, T') with distribution approximates 26g.
CDD (0)
Theorem 3. Under proper regularity conditions, for 0 < t 1< T, we have Ba (1(6( $ o(1). On the other hand, holds
W1 (fo,
T
T o(1)
(18)
The theorem is proved in Appendix B.I and it indicates that using the proposed adversarial training objective (17) of CM indeed guarantees the learned CM transfers ør into data from g(zo))
EXPERIMENTS
6.1 ALGORITHMS
In the standard adversarial training method like Projected Gradient Descent (PGD) (Madry et al. 2018), the perturbation á is constructed by implementing numbers (3-8) of gradient ascents d
7

---

Published as a conference paper at ICLR 2025
before updating the model, which slows down the training process. To resolve this, we adopt an efficient implementation (Shafahi et al. 2019 in Algorithms | 2 to solve AT (1) and (17) of DPM and CM, which has similar computational cost compared to standard training. and significantly accelerate standard AT. Notably, unlike PGD, in Algorithms| and 2 every maximization step of perturbation 8 follows an update step of the model 9. Thus, the efficient AT' do not require further back propagations to construct adversarial samples as in PGD. We provide a comparison between our efficient AT and standard AT (PGD) with the same update iterations of model 0 in Appendix I.I Moreover, we observe that efficient AT can yield comparable and even better performance than PGD while accelerating the training (2.6x speed-up), further verifying the benefits of our efficient AT.
6.2 PERFORMANCE ON DPM
Settings. The experiments are conducted on the unconditional generation on CIFAR 10 32 32 (Krizhevsky & Hinton 2009) and the class-come caas--oondiiional generation on ImageNet 64 x 64 (Deng et 2009). Our model and training pipelines in adopted from ADM (Dhariwal & Nichol 2021 paper, where ADM is a UNet-type network (Ronneberger et al., 2013), with strong performance in image generation under diffusion model.
To save training costs, our methods and baselines are fine-tuned from pretrained models, rather than training from scratch. By doing so, we can efficiently assess the performance of methods, which is more practical for general scenarios. We also explore training from scratch in Appendix G.2 which also verifies the effectiveness of our method in this regime. During training, we fine-tune the pretrained models (details are in Appendix [EII) with batch size 128 for 150K iterations under learning rate le-4 on CIFAR-10, and batch size 1024 for 50K iterations under learning rate of 3e-4 on ImageNet. For the hyperparameters of AT, we select the adversarial learning rate a from (0.05, 0.1, 0.53 and the adversarial step K from (3,5). 5). More details are in Appendix E.I
We use the Frechet Inception Distance (FID) (Heusel et al. 2017) to evaluate image quality. Unless otherwise specified, 50K images are sampled for evaluation. Other results of metric Classification Accuracy Score (CAS) (Ravuri& & Vinyals 2019). sFID, Inception Score, Precision, and Recall are in Appendix and F.4 for comprehensive evaluation.
Baselines. For experiments on diffusion models, we consider the following baselines. 1): the original pretrained model. Compared with it, we verify whether the models are overfitting during fine-tuning. 2): continue fine-tuning the pretrained model, which is fine-tuned with the standard diffusion objective (5). Compared to it, we validate whether performance improvements come only from more training costs. We also compare with the existing typical method alleviate the DPM distribution mismatch, 3): ADM-IP (Ning et al., 2023), which adds a Gaussian perturbation to the input data to simulate mismatch errors during the training process. The last two fine-tuning baselines are based on the same pretrained model and hyperparameters as in the original literature.
Results. To verify the effectiveness of our AT method, we conduct experiments with four diffusson samplers: IDDPM (Dhariwal & Nichol 2021), DDIM (Song et al. 2022). DPM-Solver (Lu et al.) 2022b), and ES (Ning et al.; 2024) under various NFEs. The sampler choices contain the three most popular samplers: IDDPM, DDIM. DPM-Solver and ES, sampler that scales down the norm of predicted noise to mitigate the distribution mismatch from the perspective of sampling. The experimental results of CIFAR-11 0 and ImageNet are shown in Table 1 and Table respectively. Results of more than hundreds of NFES are shown in Appendix F3
As can be seen, the proposed AT for DPM significantly improves the performance of the original pretrained model and outperforms the other baselines (continue fine-tuning and ADM-IP) overall for all diffusion samplers and NFEs we take. Moreover, we have the following observarions.
1): Fewer (practically used) sampling steps (5,10) will result in larger mismatching errors, while our AT method demonstrates significant improvements in this regime across various samplers, e.g., AT improves FID 27.72 to 17.36 under 5 NFEs DPM-Solver on Ieet. This suggests that our method is indeed effective in alleviating the distribution mismatch of DPM. The results also indicate that our method consistently beats the baseline methods, regardless of stochastic (IDDPM)
Foor the experts in AT, they would recognize that the AT in Algorithms actually constructs the adversarial augmented data improve the performance of the model (7hu et al. 2020 lann et 2020 Yi al. 202T
8

---

Published as a conference paper at ICLR 2025
Table 1: Sample quality measured by FID & of different sampling methods of DPM under different
NFEs on CIFAR10 32x32. All models are trained with same iterations (computational costs).
(a) IDDPM
(b) DDIM
Methods
NFEss
Methods \ NHES
20
ADM (original)
37.99
26.75
22.62
10.52
4.55
ADM (original)
34.28
14.34
11.66
7.00
ADM (finetune)
36.91
26.06
21.94
10.58
4.34
ADM (finetune)
29.30
15.08
12.06
6.80
4.15
ADM-IP
26,91
20.09
7.81
3.42
ADM-IP
43.15
15.72
10.47
4.58
ADM-AT (Ours)
37.15
23.59
15.88
6.60
3,34
ADM-AT (Ours)
26.38
12.98
9.30
4.40
3.07
(c) ES
(d) DPM- PPM--Slver
Methods NFEEs
8
10
20
50
Methods NFEs
5
8
10
20
50
ADM (original)
82.18
29.28
17.73
5.11
2.70
ADM (original)
23.95
8.00
5.46
3.46
3.14
ADM (finetune)
63.46
24.80
17.03
5.19
2.52
ADM (finetune)
22.98
7.61
5.29
3.41
3.12
ADM-IP
91.10
31,44
18.72
5.19
2.89
ADM-IP
43,83
6.70
6.80
9.78
10.91
ADM-AT (Ours)
41.07
21.62
14.68
4.36
2.48
ADM-AT (Ours)
18.40
5.84
4.81
3.28
3.01
Table Sample quality measured by FIID 4 of different sampling methods of DPM under different NFEs on ImageNet. 64x64. All models are trained with the same iterations (computational costs).
(a) IDDPM
(b) DDIM
Methods
NFEs
Methods \ NPEs
5
20
ADM (original)
76.92
33.74
27.63
12.85
5.30
ADM (original)
60.07
20.
14.97
8.41
5.65
ADM (finetune)
78.87
33,99
27.82
12.80
5.26
ADM (finetune)
60.32
20.26
15.04
8.32
5.48
ADM-IP
29.96
22.60
8.
3.83
ADM-IP
76.51
26.25
18.05
8.40
ADM-AT (Ours)
45.65
23.79
19.18
8.28
4.01
ADM-AT (Ours)
43.04
16.08
12.15
6.20
4.67
(c) ES
(d) DPM-Solver
Methods
NFEs
8
10
20
50
Methods
NFEs
5
8
10
20
50
ADM (original)
71.31
28.97
21.10
8.23
3.76
ADM (original)
27.72
10.06
7.21
4.69
4.24
ADM (finetune)
72.30
29.24
21.58
8.25
3.64
ADM (finctune)
27.82
7.22
4.64
4.15
ADM-IP
88.37
33,91
23.32
7.80
3,54
ADM-IP
32.43
9.94
8.87
9.16
9.68
ADM-AT (Ours)
43.95
19.57
14.12
6.16
3.45
ADM-AT (Ours)
17.36
6.55
5.78
4.56
4.34
or deterministic samplers (DDIM, DPM-Solver). 2): orthogonal the sampling- combined to further alleviate the issue. Notably, we further verify in Appendix G.2 that our methods will not slow the convergence unlike AT classification (Madry et al. (018)). We also perform ablation analysis of hyperparameters in our AT framework in Appendix G.3
The ES sampler results show that our AT is method to mitigate the distribution mismatch problem and can be
6.3 PERFORMANCE ON LATENT CONSISTENCY MODELS
Settings. We further evaluate the proposed AT for consistency models text-to-image generation tasks with Latent Consistency Models (Luo et al. 2023) Stable Diffusion (SD) v1.5 (Rombach et al.,
2022) backbone, which generates 512 x512 images. Both AT and the original LCM training (baseline) are trained from scratch with the same hyperparameters (the IP method (Ning et al.) 2023) not applied straightforwardlyy. The training set is LAION-Aestheties-66
(Schuhmann al.
2022) with hyperparameters following Song et al. 2023 : Luo et al. (2023). We select the adversarial
learning rate a from (0.02, 0.05) and adversarial step from {2, 3). The models are trained with
batch size of 64 for 100K iterations. More details are shown in Appendix E..2
Following Luo et al. (022) and Chen al. (2024). we evaluate models MS-COCO 2014 (Lin et 2014a) at a resolution of 512x512 by randomly drawing 30K prompts from its validation set. Then, we report the FID between the generated samples under these prompts and the reference samples from the full validation set following (Saaaria et al. (2022). We also report CLIP scores (Hessel et al.) 2021 to evaluate the text-image alignment by CLIP-ViT-B/16.

---

Published as a conference paper at ICLR 2025
Table Results of LCM on MS-COCO 2014 validation at 112 resolution in terms of FID and CLIP score 1. All models are trained with the same setting (computational costs).
FID
CLIP Scure
Methods
step
2 stop
4 step
8 step
step
step
step
8 step
LCM
25.43
12.61
11.61
12.62
29.25
30.24
30.40
30.47
LCM-AT (Ours)
23.34
11.28
10.31
10.68
29.63
30.43
30.49
30.53
Results. The methods are evaluated under various sampling steps Table which shows that the LCM with AT consistently improves FID under various sampling steps. Besides, though the AT is not specified improve text-image alignment, we observe that has comparable or even better CLIP scores across various sampling steps, which shows that AT will not degenerate text-image alignment.
7 CONCLUSION
In this paper, we novelly introduce efficient Adversarial Training (AT) in the training of DPM and CM to mitigate the issue of distribution mismatch between training and sampling. We conduct an in-depth analysis of the DPM training objective and systematically characterize the distribution mismatch problem. Furthermore, we prove that the training objective of CM similarly faces the distribution mismatch issue. We theoretically prove that DRO can mitigate the mismatch for both DPM and CM, which is equivalent to conducting AT. Experiments on image generation and text-to- image generation benchmarks verify the effectiveness of the proposed AT method in alleviating the distribution mismatch of DPM and CM.
ACKNOWLEDGMENTS
We thank anonymous reviewers for insightful feedback that helped improve the paper. Zekun Wang, Ming Liu, Bing Qin are supported by the National Science Foundation of China (U22B2059) 62276083), the Human-Machine Integrated Consultation System for Cardiovascular Diseases (2023A003). They also appreciate the support from China Mobile Group Heilongjiang Co., Ltd.
REFERENCES
Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang. Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models. In International Conference Learning Representations, 2022.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. In Advances in Neural Information Processing Systems, 2015.
Andreas Blattmann, Robin Rombach, Huan Ling, Tim Dockhorn, Seung Wook Kim, Sanja Fidler, and Karsten Kreis. Align your latents: High-resolution video synthesis with latent diffusion models. In Conference on Computer Vision and Pattern Recognition, 2023.
Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of bcore-based generative modeling: User-friendly bounds under minimal smoothness assumptions. In International Conference on Machine Learning, 2023.
Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Zhongdao Wang, James T. Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-o: Fast training of diffusion transformer for photorealistic -ot-tt--mmae synthesis. In International Conference on Learning Representations, 2024.
Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In International Conference on Learning Representations, 2022.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In Conference on Computer Vision and Pattern Recognition, 2009.
10

---

Published as a conference paper at ICLR 2025
Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs image synthesis. In Advances in Neural Information Processing Systems, volume 34, pp.
2021.
John Duchi. Lecture notes for statistics 1/electria engineering 377. URL: https://stanford (uudcluss/ttass31/eeeuuee
notes. pdf. Last visited on, 2:23, 2016.
Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems, 2014.
lan J. Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. In International Conference on Learning Representations, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition, 2016.
Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Le Bras, and Yejin Choi. Clipscore: A reference- free evaluation metric for image captioning. In Proceedings of the Conference on Empirical Methods in Natural Language Processing., 2021.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
Jonathan Ho, Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. Journal of Machine Learning Research, 23(4)):1-33, 2022a.
Jonathan Ho. Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. In Advances in Neural Information Processing Systems, 2022b.
Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Tuo Zhao. SMART: robust and efficient fine-tuning for pre-trained natural language models through principled regular- ized optimization. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, 2020.
Diederík Kingma and Max Welling. Auto-encoding variational (Bayes). In International Confer- ence on Learning Representations, 2013.
Alex Krizhevsky and Geolfrey Hinton. Learning multiple layers of features from tiny images. 2009.
Jaeho Lee and Maxim Raginsky. Minimax statistical learning with wasserstein distances. In Advances in Neural Information Processing Systems, 2018.
Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Towards faster non-asymptotic convergence for diffusion-based generative models. Preprint arXiv:2306.09251, 2023.
Mingxiao Li, Tingyu Qu, Wei Sun, and Marie-Francihs Moens. Alleviating exposure bias in diffusion models through sampling with shifted time steps. In International Conference Learning Representations, 2024.
Yangming Li and Mihaela van der Schaar. On error propagation of diffusion models. In The Twelfth International Conference on Learning Representations, 2024.
Tsung-Yr Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision- ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, 2014a.
Tsung- Yi Lin, Michael Maire, Serge J. Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft COCO: common objects in context. In ECCV, 2014b.
11

---

Published as a conference paper at ICLR 2025
Iyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019.
Aaron Lou and Stefano Ermon. Reflected diffusion models. In International Conference on Machine Learning, 2023.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. In Advances in Neural Information Processing Systems, 2022a.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan LI, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. In Advances in Neural Information Processing Systems, 2022b.
Simian Luo, Yiqin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models: Synthesizing high-resone reshuttsoottion images with few-step inference, 2023.
Jiajun Ma. Shuchen Xue, Tianyang Hu, Wenjia Wang. Zhaoqiang Liu, Zhenguo Li, Zhi- Zhi-Ming Ma. and Kenji Kawaguchi. The surprising effectiveness of skip-tuning in diffusion sampling. arXiv preprint arXiv:2402. 2024.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. International Conference on Learning Representations, 2018.
Hongseok Namkoong. Reliable machine learning via distributional robustness. PhD thesis, Stanford University, 2019.
Alexander Quinn Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. GLIDE: towards photorealistic image generation and editing with text-guided diffusion models. In International Conference on Machine Learning.
Weili Nie, Brandon Guo, Yujia Huang. Chaowei Xiao, Arash Vahdat, and Animashree Anandkumar. Diffusion models for adversarial purification. International Conference on Machine Learning, 2022.
Mang Ning, Enver Sangineto, Angelo Porrello, Simone Calderara, and Rita Cucchiara. Input perturbation reduces exposure bias in diffusion models. In International Conference on Machine Learning, 2023.
Mang Ning, Mingxiao Li, Jianlin Su, Albert Ali Salah, and Itir Ónal Ertugrul. Elucidating the exposure bias in diffusion models. In International Conference on Learning Representations, 2024.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents, 2022.
Marc' Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba. Sequence level training with recurrent neural networks. In International Conference Learning Representations, 2016.
Suman Ravuri and Oriol Vinyals. Classification accuracy score for conditional generative models. In Advances in Neural Information Processing Systems, 2019.
Sylvestre-Alvise Rebuffi, Sven Gowal, Dan A. Calian, Florian Stimberg, Olivia Wiles, and Timothy A. Mann. Fixing data augmentation improve adversarial robustness. Preprint arXiv:2103.01944 2021.
Zhiyao Ren, Yibing Zhan, Liang Ding, Gaoang Wang, Chaoyue Wang, Zhongyi Fan, and Dacheng Tao. Multi-step denoising scheduled sampling: Towards alleviating exposure bias for diffusion models. In AAAI Conference Artificial Intelligence, 2024.
12

---

Published as a conference paper at ICLR 2025
Steven J. Rennie, Etienne Marcheret, Youssef Mroueh, Jerret Ross, and Vaibhava Goel. Self-critical sequence training for image captioning. In IEEE Conference on Computer Vision and Pattern Recognition, 2017.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Conference Computer Vision and Puttern Recognition, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. Medical Image Computing and Computer-Asiistee Intervention-MICCAN 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp. 234-241. Springer, 2015.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily I. Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, Jonathan Ho, David J Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. In Advances in Neural Information Processing Systems, 2022.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference Learning Representations, 2022.
Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, Patrick Schramowski, Srivatsa Kundurthy, Katherine Crowson, Ludwig Schmidt, Robert Kaczmarczyk, and Jenia Jitsev. LAION-5B: an open large-scale dataset for training next generation image-text models. In Advances in Neural Information Processing Systems, 2022.
Ali Shafahi, Mahyar Najibi, Amin Ghiasi, Zheng Xu, John P. Dickerson, Christoph Studer, Larry S. Davis, Gavin Taylor, and Tom Goldstein. Adversarial training for free! In Advances in Neural Information Processing Systems 32: Annual Conference Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, 2019.
Alexander Shapiro. Distributionally robust stochastic programming, SIAM Journal on Optimization, 27(4):2258 2275,
Shiqi Shen, Yong Cheng, Zhongjun He, Wei He. Hua Wu, Maosong Sun, and Yang Liu. Minimum risk training for neural machine translation. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics, 2016.
Albert N Shiryaev. Probability-I, volume 95. Springer, 2016.
Aman Sinha, Hongseok Namkoong, and John Duchi. Certifying some distributional robustness with principled adversarial training. In International Conference on Learning Representations, 2018.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning. 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- fional Conference on Learning Representations, 2022.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. 2021.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In International Conference on Machine Learning, 2023.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
13

---

Published as a conference paper at ICLR 2025
Martin J Wainwright. High-dimensional statistics: non-asymptotic viewpoint, volume 48, Cam- bridge university press, 2019.
Ruoyu Wang, Mingyang Yi, Zhitang Chen, and Shengyu Zhu, Out-of-distribution generalization with causal invariant transformations. In Conference on Computer Vision and Pattern Recognition, 2022.
Zekai Wang, Tianyu Pang, Chao Du, Min Lin, Weiwei Liu, and Shuicheng Yan. Better diffusion models further improve adversarial training. In International Conference on Machine Learning. 2023.
Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust general- ization. In Advances in Neural Information Processing Systems, 2020.
Shuchen Xue, Zhaoqiang Liu, Fei Chen, Shifeng Zhang, Tianyang Hu, Enze Xie, and Zhenguo Li. Accelerating diffusion sampling with optimized time steps. arXiv preprint arXiv:2402.17376 2024a.
Shuchen Xue, Mingyang Yi, Weijian Luo, Shifeng Zhang, Jiacheng Sun, Zhenguo Li, and Zhi-Ming Ma. Sa-solver: Stochastic adams solver for fast sampling of diffusion models. Advances Neural Information Processing Systems, 36, 2024b.
Mingyang Yi. Lu Hou, Jiacheng Sun, Lifeng Shang. Xin Jiang, Qun Liu, and Zhiming Ma. Improved ood generalization via adversarial training and pretraing. In International Conference on Machine Learning, 2021.
Mingyang Yi, Jiacheng Sun, and Zhenguo Li. On the generalization of diffusion model. Preprint arXiv:2305, 14712, 2023a.
Mingyang Yi, Ruoyu Wang, Jiacheng Sun, Zhenguo Li, and Zhi-Ming Ma. Breaking correlation shift via conditional invariant regularizer. The International Conference on Learning Representations, 2023b.
Mingyang Yi, Aoxue Li, Yi Xin, and Zhenguo Li. Towards understanding the working mechanism of text-to-image diffusion model. Preprint S.Ki:22405.13330 2024.
Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Frédo Durand, William T Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. In CVPR, 2024.
Boya Zhang, Weijian Luo, and Zhihua Zhang. Enhancing adversarial robustness via score-based optimization. In Advances in Neural Information Processing Systems, 2023.
Dinghuai Zhang, Tianyuan Zhang, Yiping Lu, Zhanxing Zhu, and Bin Dong. You only propagate once: Accelerating adversarial training via maximal principle. In Advances in Neural Information Processing Systems, 2019a.
Hongyang Zhang, Yaodong Yu, Jiantao Jiao, Eric Xing. Laurent El Ghaoui, and Michael I. Jordan. Theoretically principled trade-off between robustness and accuracy. In International Conference on Machine Learning, 2019b.
Wen Zhang, Yang Feng, Fandong Meng, Di You, and Qun Liu. Bridging the gap between training and inference for neural machine translation. In Proceedings of the 57th Conference of the Association for Computational Linguistics, 2019c.
Chen Zhu, Yu Cheng, Zhe Gan, Siqi Sun, Tom Goldstein, and Jingjing Liu. Freelb: Enhanced adversarial training for natural language understanding. In International Conference on Learning Representations, 2020.
14

---

Published as a conference paper at ICLR 2025
A PROOFS IN SECTION
In this section, we present the proofs of the results in Section
A.1 PROOFS IN SECTION
Proposition 1. The minimization problem is equivalent to minimizing an upper bound of E log Po (z+) for any S
Proof. We prove the first equivalence, by Jensen's inequality. For any 0 S T, have
Bg Po po (2x:)
log (aFa:T
B
q(<++1:1 xxt)
log po(x))
Σ
log (mr ma+11)
g(xo+1 xto)
ISART
log pa (æt xa+11) g(ar, (x+1)
((x.))
Eg
log po (xT))
(19)
((
pe (o(r)) q(xT))
log pa(a. x.41)
E
log
log g(27))
q 1x+1)
ISAET
=xx (g(TT) (o(T))) + Eg
Dxi(a(a, (+41)) po(ar, m++1l)
Hxx,))
Taking = O, we prove the first equivalence. Besides that, the entropy H Hxx) of æ; is a constant for 0 given data distribution øg for any 0 < t T. The second conclusion holds due the non-negative property Kl.ddieegeece.
Proposition Suppose po (at (++1)) matches g(xx (x41)) well such that
L. = (Kl(g((. m141) pa(x: x(41)) S y T
(8)
and the discrepancy satisfies (KI((((TT) || po(w))) S 70. then for any 0 5 T, have
Dkz(g(æ.) po(x:)) S Dwz(((azr))) || po(2)1)))) S go (T - f)r
(9)
Proof. We have the following decomposition due to the chain rule of KL-divergence
Drilg(z,, w(+1) (o()))+++1)) m++1)) (x m++1) po (e++)) Deu (xz(g(0+11) || (=)+1))
(xI(((#44 wx) pe ax)) + (g(()( (o()))),
(20)
The transition probability po (x #x+1) matches (ax, 1.++)), so that the above equality implies
Dxr(g(mx) || (o)))))
=Dxu(g(mm1) po (()41))) + (Kl(g(a: m441) || (o m(41)) (glmaa ax) || Po xx))
<Dxi(g(ze)) po(x))))) 7
(21)
The proposition holds due to initial condition DKL(g(æT) || (e( S To and simple induction.
Proposition 3. Lap in (4) is well minimized, only if of ((044))) Gaussian or || xx, 0.
15

---

Published as a conference paper at ICLR 2025
Proof. Due Bayes' rule, we have
q m(+1) g(
2 (2)1)
(((++1)
z141 var4iae logg(æz) log logg(z)e+1)
exp
2(1- -0++1)
(22)
x14 vartiazl
(V log g(oeg(()))))))))) set1))
exp
2(1 -0++1)
exp
(m: m441)
V2 logq( m141) O( mx =2211))
As can be seen, the conditional probability be approximated by Gaussian only if V log ((2++1)
zero or || - æt || is extremely small with high probability. The two conditions can be respectively satisfied when g(æ+)) is Gaussian or æ; close to æz+11
A.2 PROOFS IN SECTION 4.2
Proposition 4. For any distribution q satisfies ((æ,) = g(x0e) for specific t, we have B -- log]o(m:)) s (xl(ū( (e+11) (a (e+11) +O,
(I1)
for a constant C independent o. O.
Proof. W.o.l.g. suppose po (+41) po (ate (++1)g((441)) and ((21, x(+1) ((x4+1 ax (x))((x)). By Jensen's inequality, we have Fgg [- po (a::)]
/ ((w.)) loo
pe (w (d41)d dæ:
/ ((w.))
po ((204+1
log
w. (e dæ
0(x041 wi)
(ʃ log (o( 0041))
q(xxa)
g(2+41
dw
g(0041 wi)
/ ((w.) (f q(weet ax,) log Po (x
Re41)
dac
x(x+41
wxs)
[ ((<aa) (f (((0+11 (2 log
g(0c+1)
(23)
dds++) [2.)
dx
[ i( x++11 log (x [1+11)
da, + C1
7(x++1 (x))
i( x1+11 log (x [1+1)
q(x))
do + C1 (((++1)
q(x;
#(+1)
log ( [141)
((1r, x(+1)
do
+ +0++0
(( (++11)
Dx (ā(#; (e411) |l (x [2++1)) +C
= L (0, (),t)= + where C, C₁, C2 are all constants independent of 0.
A.2.1 PROOF OF THFOREM
In this section, we prove the Theorem I To simplify the notation, let pa(a, NH + 1),0(+1 in then the optimal solution (Lemma 9 in (Bao al. minimizing L
Btn1 ~
2022 of
(o = Eg, [xx [2+1).
(24)
'Here 0:-1 can be also optimized as in (Bao et al. 2022). but we find optimizing it in practice does not. improve the empirical results.
16

---

Published as a conference paper at ICLR 2025
For every specific i, consider the following 9: in 12 such that
q: ( am)) # g( m:);
(2+++) # a(zc+1);
(25)
(<(:)))) g(zo:c).
Ge (+) zo, x++1)
g( zo, x++1) = N (1u+ (2o, T++1),
where Me+1 (Io, It+1) = var(1-a.11)
1-a--1 To +
1-6,-1 Te+1- The T can be taken due to the
Bayesian rule. Next, analyze the optimal formulation in 24) Due to the property of conditional expectation, we have
Mo( 15, [Eg, [
2o. B1+1]
[ pa-- (Bgg [2o [141], 2.41)
(26)
As can seen, the optimal transition rule decided by the conditional expectation Beg (a x++1] for some (#x-1) E BDxx (#(*++1), in 12). Then, we have the following lemma to get the desired conditional expectation.
Lemma 1. There exists some yj 2 no which makes
equivalent to problem
12
min (xn)
sup
E ( (200 + xo xol1²]
(27)
1=0
(e(#+41 [ (ec( (mo).))
where Ep co
x++1l x0(
Proof. Let check the training objective ming & (ac11)) DKL(Ā((E) x++1) ll (a (++1)). During this proof, we abbreviate Bax (4++1 (I241), n) as B. Since po(ar, x(+1) ~ N( + 1), 02)1, then
sup DxI(se(a, m+41) ll po(x a++1)) (m44)))
(28)
log 2
sup Egie) [[|æ (0 1
20 d.(m 1)B
JEB
As consider J constant, an analysis of the expectation term enough. Due to
[||as - (a + 1)]l 2 inf Eg (mp. [[| - f(roo a++1)]1
(29)
[| Ez(ar, mo, 1]+1]]
where the last term is invariant q: B so that it is uniform lower bound over all possible ğo and po (z, 1. The above inequality indicates that the optimal 14o (#++1, 1) is achieved when the left in (29 becomes the right in 29
On the other hand, for any q E B, let us compute the gap such that
(q.( [||= 14g ( ( t ]]]
Fq. [||T Tig lay mo, me + Bg [ms ao. me41] 1ho 1o (xret1 ++ 11]1
Eg |||0: Eig [xo. Lo, 21+111
+ Big [1l1g((1441
1) Eg. [w) xo, x111]
(30)
-2E,, [(ar - Eqg [ax [oo 344 (2141 + Bg, [ao wo. 2041])]]
Egg. [||ar- Eg, px. xa, Te+1)1 [1
(aa
V a. of v
ac++
Eq.leo. 1) [[| wa cal(44114 + 1)11a].
where the equality is due to the property of conditional expectation leads to Eq. [(x E,g [x, æo, 2e+1]], 1), Mgg (2++1,
[xr æo. (++ O. and rewriting E,, [ll [[ ++1)++1--1--+)) E,g [x,
2o, x1+1 1]] as in equations (5)-(10) in (Ho et 2020). Due this, know that minimizing the
We can do this since
12 only relates (((+++1)
17

---

Published as a conference paper at ICLR 2025
square enror is equivalent to minimizing the (m)))))))) 20(0141, + 1)]l On the other hand, since qi E B, then we have
Dei(q(w441 wo) |l di (#x+1 (o))
=Dxi(g(æo m++1) |l (#ro m++1)) (x || @ (m++1))
(31)
27m-
Thus, we prove our conclusion.
Theorem 1. There exists d depends zo and c₁ makes 13
equivalent to problem
12
8:
min
E (
colVārao
v1 0eE + 6,,
(13)
Proof. By combining Lemma] with
I suppose the supreme is attained under d:-1 such that æt ~ qt- (xx))
vā,zo + -T-a + 8c: with ò depends zo and os. Then we prove the conclusion.
(32)
A.2.2 PROOF OF PROPOSITION
Proposition 5. If L (P) (0) < in 12 for all t, and DKL(ạ(æT) po(xr)) S To. then
DKr(g(ar) pa (o))) S v/o-
Proof. This theorem proved by induction. Since (g(xr) po (wr)) < 1Jo, then, let @ ( = po (xT) and satisfies 07 (xT)
g(xr-1)- The existence of such distribution is due
to Kolmogorov existence theorem (Shiryaev 2016 Then, have
Dkz(ōr-(ar po
(o())))
+ Dx 1 (0 Φr) pa(#r-1 (r))
(33)
< L. DRO (0)
S na
where the first inequality is due to the definition of L (O) and ---((TT) (DTT)) Then, prove our conclusion by induction over t.
A.2.3 PROOF OF PROPOSITION
Proposition 6. For rj > 0 and 8, in 13 |
S holds with probability at least 1- 2211 ā.)/n
Proof. Due the definition of the first order Wasserstein distance W1 (:..) (Villani al. 2009) for any specific suppose
arg min
T.. [[] xr [r]]11]
(34)
T( (.)a (
so that
E. [| æe i] = (gi wo). (e(e mo))
(35)
Let ò, be the one of (13 under Λ² derived by Lemmal I
then
([16,11 2 n mo) < E,- [|8.1]]
"
Wi (g: ax)), qe (xc, ac))
/2 ax) Dr (gt 2o) de (wr, aF))
(36)
2(1 ā.)
77
where inequality a is due to the Talagrand's inequality (Wainwright 2019). Then we prove our conclusion.
18

---

Published as a conference paper at ICLR 2025
B PROOFS IN SECTION 5
Next, we give the proof of results in Section 5 Firstly, let us check the definition of the (2(+1))) For the variance-preservingg stochastic differential equation in Song et (2022)
8. =,dl B.dW...
dz.
(37)
2
Due to the solution of in Song et al. (2023). know 2s, has the same distribution with 2; in for [s satisfies
exp (-
%." B(u)du ae (80 0).
(38)
In the rest of this section, we use d(x, y) in 15 as la distance || yl whereas the conclusions under other distance can be similarly derived. Owing the the discussion in above, similar (Song et al. 2023), when 2++1 R let 4, (2 V (Z we can rewrite the objective (15) as Tollows.
min min (cD
min
Ex.,
[| (zxes)).)) So S0(2 + D)
(39)
Here 2, follows the following reverse time ODE of (37) with 2o q(x0)).
Bg 2
dz.
+ 2 V= log qa (z.) ds,
(40)
and such 2s has the same distribution with the ones in 37 Song et al. 2022). where is the pett1 (z,)d, which deterministic function of 2a.11*
density of 2.. V (Zatt1) and (o so: 0) s0o 20.
80411
Now, we are ready to prove the Theorem 2 as follows.
Theorem For (c 15) with d(.,.) ) is le distance, then W (f 2o)
VtLcD
Proof. Owing to the definition of W distance, and the discussion in above, we have
Wi xo)
Wi(fo 1), 2x0)
Wi (fo(zep: T), 4 (Ver (... 1 Ver-i (Esy E. [1lfo(2 T) V (Ve, (... ....D. P (#er
(41)
fo t 1) fo (V. t)
<LDD(O0))
where the first inequality is due the definition of Wasserstein distance, the second and last inequalities respectively use the triangle inequality and Schwarz's ' inequality.
B. PROOF OF THEOREM 3
As pointed out in the above, the used i (x(+1, cΦ) numerical estimator of & (Z(41). In the sequel, let us consider & an Euler estimator as follows, whereas our analysis can be similarly generalized to the other estimators.
i. (.(111) W., (Zu))1) = 20111 + (5141 - 81) Bect Ttlee
cp(zmn1 11/1I ar+1
(42)
Here W1 (fffoar,, f), a7) the Wasserstein I-distance between distributions of fo t) and mo-
19

---

Assumption 2 describes the estimation error of d,, (z). which terms out to be the training objective of Published as a conference paper at ICLR 2025
where VI- and the condition 2(+1 = z is hold.
+++ estimates v log gs, (Z)11) as pointed out in (Song et al. 2020
Next, we illustrate the used regularity conditions to derive Theorem 3
Assumption 1. The discretion error of s, (2 Ebb) that says
is smaller than C(8t+1 8 2 for constant C.
Zoets S,a gott @.(z.)da S C((++1 -e)?
Vee (2.(20141) E6)
(43)
Assumption 2. The estimated score v log (z(2) has bounded expected error, i.e..
(x) [1(ba((e) - th. (2) [=)
Se.
(44)
for all 0 <t<T.
Assumption 3. For the learned model fe. it holds || S.D. D.
The [ describees the discretion error of the Euler method under ODE with drift term 6.. which can be satisfied under proper continuity conditions of model cs. On the other hand,
obtaining it, see (Song al. 2020) for more details. The Assumption ii natural, since fe predicts Ep. which is usually an image data with bounded norm. Now, we are ready to prove the Theorem 3
which is presented by proving the following formal version.
Theorem 4. Under Assumptions 7 2 and 3 for all dsy we have Ez,, (1,,,, S 0(42,
EV A..)) for As, S441 St. Besides that, we have
4D2 Wi( s T (0) CA2. +c0(vA.))
(45)
Proof. Noting that (x++1)) = V (2 and ( cØ) (Z Cp). the key problem is upper bound the difference between Vi (2) and V (2) for all t and z. To do so, we note that
Eb)
(2, (=)( s (2, - S.mm (z.)ds + =
(..z)d V., (z)
(46)
where the first one in r.h.s can be upper bounded by C(s141 82)² according Assumption 1 On the other hand, define dž, ds (z())). then when Esc=1 2s141 z and [se, 86+1)
2 2.11² 2
22 nn. 6..2) 00(2z))
ds
($. 2s, (2,)) - -dx(z.) + (.((zs) - -da(z.))
(47)
T1]B, - ==.12 (2.- ms. (.(z)) )e(2.))
$ 2 + +4) | =.11 + 1 2 [ - oz.z)) ||?
Taking expectation over z, by Gronwall's inequality, Assumption! 2 and 20111 =2 1 we have
E [l]a.
x=. S
1914 (122+)((--))
E [10)(x))) 0s (zx))] ds 4 J/a""
B.c (/2 ds.
2
(48)
Plugging this into 46 we know
E V (Z.,, ()) - v (R.1))
S - = 8c).
(49)
By Markov's inequality, we have
Vey (Zo) F6) - D., (zei))
P ((] (1
(2,,,, (Za,, 6x) - V Ver 2 2n) 5
77
(50)
[((s41 - =.)² +00
"
20

---

Published as a conference paper at ICLR 2025
Thus,
E [Il +) - (0(()))))(((((()))))))))
[o 1) So (2
fo 2a E
fo ($
+ de Eb).
1 B, 1120 + 10,, 110, 11 fo( + fo (Za111 ds,,
- I
sup S
fo 2 (Z +) Wo( (2e441 + dee E 0/|
+ 1D2P (116,11 116 2 2n)
AD² [((s411 - s:)² -c0 -s)]
$
sup 18
fo( 1) ( + c., es),
(51)
Taking over t and combining Theorem 2 we prove conclusion.
Therefore, this theorem, by taking A., = 81+1 8t close to zero, we get the results in Theorem
C THE CONNECTION TO STANDARD ADVERSARIAL TRAINING
In this section, clarify why the proposed AT objective (14) is general version of the standard AT objective proposed in (Madry et al. 2018) used for classification problems.
For classification problem, given model So (ar), data az, and label y, it aims minimize the adversarial training objective
min E. (2..)) sup ((folæ +),
(52)
6:1611m0
for some loss function (e.g. entropy) and adversarial radius 7o- However, the objective is not directly generalized to the diffusion model, as its training objective is a regression problem instead of classification (5 Thus, should refer the general version of adversarial training in (Yi et al. 2021 Sinha et al. (018), where the training objective ming Ex [a))):) and the adversarial training objective becomes
min Ex
to + 8)) 6:11611$
(53)
where lg is the parameterized loss function, and æ data. Then, we can conclude our objective [14 follows the above formulation, such that the goal is represented as
min
sup 8:
P (xt + 6)
(54)
compared with the original noise prediction objective ming El= Emo Weg, se that the loss function
(6 (x::))
such
(m.)
Eo (t. mx)
(55)
This clarifies the equivalence of our objective 14 to general adversarial training.
D ADVERSARIAL TRAINING ON CONSISTENCY TRAINING MODEL
In (Sng al. 2023 the consistency model can be even trained without estimator é.. They prove that the empirical consistency distillation loss (cp(O) can be approximated by the following Lor (0)
T-1
Let (0)
Σ 1=0
Ex Em Emp [I]fa(a..)) fa(w41, 1)1
(56)
21

---

Published as a conference paper at ICLR 2025
In our adversarial regime, can also prove that the desired L can be approximated by the following C
(0) with adversarial perturbation
Σ
sup IIfa(x; 8,t) fo( +1]l
(57)
1=
[18|sn
The results can be checked by the following theorem. Theorem 5. Suppose fo( t) is twice continuously differentiable with bounded second derivative. Then
T T
LAde (0) +0
Σ
Vac + Tn
(58)
1=1
where "<" means approximately less than or equal.
Proof. Due the continuity of fo t), for any 8 with ||8| S 7) by Taylor's expansion on 1
from 2t + 6, we have
for + 8,t) fo(x)+1, 1)] [o fo + 1)11
+ E (Jo(ee11,)) fa 1)) V /a( () + 8 m111 + 0 (B [[ - 8811])
(59)
Due to the Taylor's expansion fo + t, = fo(xx-1.t) + V () + x1+1]
O1 - St. - 811²). Then, from the formulation æt we know E [| - St - 8 8 0(1 Vae + 7))) Noting that due to definition st. we have
E|æ; Ø1+11 2o441] = E(z,, #at41]
(Zor41
(1 arr1) Vo log 90143 (21442220111))
va++
F5t4
(Zeee - -e-let 11 Bede) vx log geeps (Zeets)) (60)
exp
2
Bods
544s
+ 1 Fct1
S Byda v log aou a (Zeez)
2
1 +
Bdds
B5441
2 V (Ze441: /1
5i+1 V log 9e.4));
where the first equality is due Tweedie's ' formula i.e., Lemma in (Bao al. 2022), the ">"" is due to e" x 1 + a when a + 0, and the last = is due to Euler-Mayaruma discretion. Due to this, we notice that
E (Jo(aet1, fe( 1)) Vie(a:+1 + x++1) 0e+1 Z0c41
[Ja(0t1,)) - fo( t 1)) v fe(c::11, 1) [E [ + 8 xe+1 [=c-1]
xe+1) Ør+1 = #6
E (Jo(zo(11: fo(z 1)) V fo (Z V2 (og4)(11)
E(d #sc4al z[+1
(61)
where the first equality is due to the property of conditional expectation, and the second "x²" is due to (60). Combining this with (59). have
E [ale: 8,1) fa(w11 1)1l²
x+41 #xtt1]
[[,o(zz. + 6.t) - fo 1)Il²
Z
F [ 1) 20( []]
T. (fo(z 1) fo fo 1)) Vol (Eee Vx (og +E, Z Zec4s)
0(1 Vai+ + n²²)
= [22||1o(Ve.. (V V. log ( ) 5,t) fo(2, t [1]l +0 +
(62)
where the last equality is due to Taylor's expansion from fo (x V. log 9rmm)) ) + 8,t) to fo f2 1). Due to the arbitrariness of ō, prove conclusion.
22

---

Published as a conference paper at ICLR 2025
E IMPLEMENTATION DETAILS
E.I HYPERPARAMETER OF DIFFUSION MODELS
For the diffusion models, all methods adopt the ADM model (Dhariwal & Nichol 2021 as the backbone and follow the same training pipeline. Following existing work (Dhariwal & Nichol 2021 Ning et al. 2023), we train models using the AdamW optimizer (Loshchilov & Hutter 2019) with mixed precision training and the EMA rate is set to 0.9999, For CIFAR-10, 10, the pretrained ADM trained using batch size of 128 for 250K iterations with learning rate set to 1c-4. For ImageNet. the pretrained model is trained with a batch size of 1024 for 400K iterations, employing a learning rate of 3e-4. The models are trained in cluster of NVIDIA Tesla V100s. More hyperparameters reported in Table 4
Table 4: Hyperparameters of diffusion model each datasets.
Hyperparameters
CIFAR10 32 x 32
ImageNet 64 x 64
Channels Batch size
128
192
128
1024
Learning rate
le-4
3e-4
Fine-tuning iterations
200K
200K
Dropout
0.3
0.1
Noise schedule
Cosine
Cosine
F.2 HYPERPARAMETERS OF LATENT CONSISTENCY MODEL.
For experiments on Latent Consistency Models (LCM) (Luo et 2023). we train models on LAIOIN-Aestheti-6.555 (Schuhmann et al. (2022) at the resolution of 512x512, comprising 650K text-image pairs with predicted aesthetic scores higher than 6.5. Stable Diffusion v1.5 (Rombach) et al. 2022) is adopted as the teacher model and initialized the student and target models in the latent consistency distillation framework. We set the range of the guidance scale [Wmin; Wman] = [3. 5] during training and use = 4 in sampling because it performs better in our preliminary experiments, which is similar to DMD (Yin et al. 2024). The models are trained in a cluster of NVIDIA Tesla V100s. Both models of our AT and the original LCM training trained from scratch with the same hyperparameters. We select the adversarial learning rate a from {0.02, 0.05) and adversarial step K from 3). More details of hyperparameters are shown in Table 5 and other details of implementations can be found in the original LCM paper (Luo et al. 2023).
Table 5: Hyperparameters of latent consistency model.
Hyperparameters
LA1OIN-Aesthetic66.55
Batch size
Learning rate
8e-6
Training iterations
100K
EMA rate of targel model
0.95
Conditional guidance scale [Wmin, Wmar]
[3.5]
F ADDITIONAL RESULTS
F.1 RESULTS OF CLASSIFICATION ACCURACY SCORE
Classification Accuracy Score (CAS) Ravuri & Vinyals 2019) is proposed to evaluate the utility of the images produced by the generative model for downstream classification tasks. The underlying motivation for this metric is that if the generative model captures the real data distribution, the real data distribution can be replaced by the model-generatee data and achieve similar results on downstream tasks like image classification.
23

---

Published as a conference paper at ICLR 2025
Table 6: Comparasion of CAS of different methods on CIFAR-10 32x32 dataset.
Methods
CAS
Real
92.5
only using the synthetic data.
ADM
91.0
ADM-IP
89.2
ADM-AT (Ours)
91.6
using the synthetic data with real data.
ADM
95,0
ADM-IP
94.9
ADM-AT (Ours)
95.4
Following the evaluation pipeline in Ravuri & Vinyals (201)), we train the image classifier in two settings: only synthetic data or real data augmented with synthetic data, and use the classifier to predict labels on the test set of real data. Synthetic images are generated with a DDIM sampler under 20 NFEs. We use ResNet-11 (He et al. 2016) as the image classifier and train it for 200 epochs with learning rate of 0.1 and batch size of 128. We report CAS in the CIFAR-I0 10 dataset at a resolution of 32x32 in Table 6 The results indicate that our method consistently performs better than other baseline methods on CAS metric in both settings. Although CAS with synthetic data cannot surpass real data, demonstrates significant potential for enhancing classifier accuracy when employed as an augmentation technique alongside real data.
Table 7: Comparasion of AT with TS-DDIM CIFAR10 0 32x Both models are based on the ADM backbone. The results of TS are taken directly from the original paper.
Methods NFEs
50
10
5
ADM-TS-DDILL
3.52
5.35
10.73
26.94
ADM-AT (Ours)
3.07
4.40
9.30
26.38
F.2 COMPARISON TO TS-DDIM
al 2024 introduces another approach named Time-Shift (TS) to alleviate the DPM distribution mismatch by searching for coupled time steps in sampling. shows the comparison between our AT method with TS on CIF AR-10 10 with the DDIM Sampler. Both methods are based on the ADM pretrained model (Dhariwal & Nichol, 2021) as a backbone, which is the same Section 6.2 We observe our method consistently better than the TS method across various sampling steps.
F.3 RESULTS OF MORE NFES
We present results obtained with various samplers under 100 or 200 NFEs on CIFAR10 32x32 and ImageNet 64x64 in Table 8 and Table 99) respectively. The results show that our method is still effective for samplers under hundreds of NFEs.
F.4 RESULTS OF MORE METRICS
We present the results of more generation quality metrics, including sFID, Inception Score (IS). Precision, and Recall, on CIFAR10 32x32 (Table I0 and Table IT) and ImageNet 64x64 (Table I2 and Table The evaluation is performed following Dhariwal Nichol (02I). We observe that method shows effectiveness across these metrics.
24

---

Published as a conference paper at ICLR 2025
Table Sample quality measured by FID of various sampling methods of DPM under 100 200
NFEs on CIFAR10 10 32x32.
IDDPM
DDIM
ES
DPM-Solver
Methods
100
200
100
200
100
200
100
200
ADM-FT
3.34
3.02
4.02
4.22
2.38
2.45
2.97
2.97
ADM-IP
2.83
2.73
6.69
8.44
2.97
3.12
10.10
10.11
ADM-AT (Ours)
2.52
2.46
3.19
3.23
2.18
2.35
2.83
3.00
Table 9: Sample quality measured by FID 4. of various sampling methods of DPM under 100 200
NFEs on ImageNet 64x64.
IDDPM
DDIM
ES
DPM-Solver
Methods
100
200
100
200
100
200
100
200
ADM-FT
3.88
3.48
4.71
4.38
3.07
2.98
4.20
4.13
ADM-IP
3.55
3.08
8.53
10.43
3.36
3.31
9.75
9.77
ADM-AT (Ours)
3.35
3.16
4.58
4.34
3.05
3.10
4.31
4.10
Table 10: Comparison of sFID and IS 1 on CIFAR10 32x32.
(a) IDDPM
5
8
10
20
50
sFID
IS
sFID
IS
sFID
IS
sFID
IS
sFID
IS
ADM
20.95
8.25
25.03
8.51
23.56
8.50
16.01
9.14
6.81
9.49
ADM-IP
25.81
7.02
24.51
8.04
19.02
8.50
8.99
9.28
5.32
9.66
ADM-AT
19.78
8.71
25.67
8.66
23.09
8.77
6.01
9.30
5.04
9.65
(b) DDIM
5
10
20
50
sFID
IS
sFID
IS
sFID
IS
sFID
IS
sFID
IS
ADM
12.75
7.76
8.53
8.62
8.39
8.70
6.19
4.99
9.19
ADM-IP
15.53
7.55
8.00
8.98
7.12
9.15
5.30
9.41
5.64
9.49
ADM-AT
12.56
7.97
7.93
8.90
7.08
8.90
5.37
9.17
4.66
9.51
(c) ES
8
sFID
sFID
IS
sFID
sFID
sFID
IS
ADM
27.39
6.14
14.91
8.33
10.04
8.79
5,45
9.55
4.12
9.62
ADM-IP
34.70
5.73
16.84
8.23
10.89
8.88
4.94
9.59
4.08
9.70
ADM-AT
16.84
6.97
10.33
8.60
8.00
4.78
9.65
4.04
9.77
(d) DPM-Solver
10
20
50
sFID
IS
sFID
IS
sFID
IS
sFID
IS
sFID
IS
ADM
11.82
8.00
5.79
9.12
5,05
9,41
4.43
9.78
4.32
9,82
ADM-IP
26.46
7.09
5.93
9.19
5.49
9.45
7.53
9.66
8.37
9.75
ADM-AT
11.19
8.43
5.10
9.35
5.29
9.65
4.75
10.03
4.59
9.93
G MORE ANALYSIS
G. EFFICIENT AT VS STANDARD AT
In this section, conduct ablation of the AT method in diffusion model training. We compare the performance of our used efficient AT and a standard AT method PGD on CIFAR-10 dataset at the resolution of 32x The adversarial step K is set to be for both methods. We fine-tune both
25

---

Published as a conference paper at ICLR 2025
Table 11: Comparison of Precision (P) and Recall (R) on CIFARI0 32x32.
(a) IDDPM
8
10
20
50
R
R
R
R
P
R
ADM
0.54
0,47
0.59
0.45
0.61
0,46
0.64
0.52
0.68
0.58
ADM-IP
0.54
0.39
0.59
0.43
0.61
0.46
0.66
0.54
0.68
0.59
ADM-AT
0.52
0.47
0.57
0.45
0.62
0.46
0.68
0.55
0.69
0.59
(b) DDIM
10
20
50
P
R
R
P
R
P
R
P
R
ADM
0.57
0.47
0.59
0.52
0.61
0.52
0.64
0.52
0.63
0.60
ADM-IP
0.57
0.44
0.62
0.53
0.63
0.56
0.65
0.60
0.65
0.61
ADM-AT
0.59
0.46
0.62
0.52
0.63
0.54
0.65
0.58
0.66
0.61
(c) ES
10
20
50
R
R
ADM
0.54
0.37
0.60
0.48
0.61
0.52
0.64
0.52
0.63
0.60
ADM-IP
0.46
0.32
0.58
0.45
0.51
0.58
0.68
0.60
ADM-AT
0.61
0.45
0.64
0.51
0.65
0.54
0.65
0.58
0.66
0.61
(d) DPM-Solver
8
10
20
50
P
R
R
R
R
P
R
ADM
0.61
0.47
0.65
0.58
0.65
0.59
0.66
0.61
0.63
0.62
ADM-IP
0.49
0.32
0.65
0.58
0.65
0.59
0.62
0.58
0.61
0.56
ADM-AT
0.62
0.49
0.65
0.59
0.65
0.61
0.67
0.62
0.65
0.61
models from the same pretrained ADM model with 100K update iterations of the model. The results are shown in Table 14 We report the results of sampler settings (method-NFESs) IDDPM-50, DDIM-50, ES-20, and DPM-Solver-10.
We observe that efficient AT achieves performance comparable to or even better than PGD with the same model update iterations while accelerating the training (2.6x speed-up). Thus, we propose applying the efficient AT method for our adversarial training framework.
G.2 CONVERGENCE OF AT ON DIFFUSION MODELS
DDIM-50
Method
ADM
ADM-IP
ADM-AT (ours)
FID
5
100K
150K
200K
Iterations
Figure 2: The convergence methods trained from scratch on CIFAR-10 32 32, We use the DDIM sampler with 50 NFEs for sampling.
26

---

Published as a conference paper at ICLR 2025
Table 12: Comparison of sFID and IS ImageNet 64x64.
(a) IDDPM
5
8
10
20
50
sFID
IS
sFID
IS
sFID
IS
sFID
IS
sFID
IS
ADM
26.17
12.55
36.34
22.61
40.52
26.55
26.08
39.10
11.35
45.68
ADM-IP
40.90
12.19
47.98
23.47
37.72
27.86
25.06
39.40
6.75
44.87
ADM-AT
24.82
14.50
37.04
23.84
36.50
30.03
22.83
39.12
5.69
46.25
(b) DDIM
5
8
10
20
50
sFID
IS
sFID
IS
sFID
IS
sFID
IS
sFID
IS
ADM
27.74
14,30
14.27
25.88
12.78
28.29
8.84
33.54
6,31
38.08
ADM-IP
52.08
10.21
16.40
22.03
11.70
25.94
9.09
32.04
15.14
31.62
ADM-AT
25.49
14.82
10.68
26.62
9.22
29.29
6.41
34.33
4.66
39.36
(c) ES
20
50
sFID
IS
sFID
IS
sFID
IS
sFID
IS
sFID
ADM
34.55
13.29
42.32
24.98
34,44
29.36
14.44
40.45
6.41
45.36
ADM-IP
44.81
10.07
41.01
22.44
30.12
27.66
10.13
39.50
4.67
44.69
ADM-AT
29.72
16.49
33.58
27.85
27.64
31.94
10.22
42.18
5.10
45.59
(d) DPM-Solver
5
8
10
20
50
sFID
IS
sFID
IS
sFID
IS
sFID
IS
sFID
IS
ADM
25.70
24.34
11.08
34.77
8.05
37.45
5.35
40.54
4.69
41.31
ADM-IP
42.68
16.93
7.47
33.85
7.22
33.57
14.74
31.29
18.99
30.32
ADM-AT
20.79
26.32
7.60
34.89
6.36
36.51
4.51
38.79
4.22
39.10
IDOPM-20
DOIM 50
E5-20
DPM Solver 10
10
4.5
4.0
6.5
3.5
8.0
3.0
4.4
4.6
5.5
5.0
2.5
4.5
50K
150K
250K
50K
150K
250K
50C
150K
250K
50K
150K
250K
teracions
terations
Itarations
Iterations
Methed
Figure 3: The convergence of methods fine-tuned from a same pretrained model on CIFAR- 10 32 32. We compare the performance of methods on various samplers.
In classification tasks, adding adversarial perturbations usually slows the convergence of model training (Zhu et al. 2020). We are interested to see whether AT also affects the convergence of the diffusion training process.
Firstly, we explore the convergence of models trained from scratch. We utilize DDIM the sampler We observe that our AT method and ADM-IP
with 50 NFEs and the results are shown in Figure 2 exhibit slower convergence compared to ADM at the beginning (before 100K iterations), while as training more iterations (200K), our AT method shows a notable advantage.
Moreover, we explore the convergence of models under fine-tuning setting and the results are shown in Figure 3 We observe under this setting, when given a pretrained diffusion model like ADM, fine-tuning with our proposed AT improves performance faster than other baselines. Overall, observe that incorporating AT with diffusion framework does not affect the convergence of the model much, especially in the fine-tuning setting.
27

---

Published as a conference paper at ICLR 2025
Table 13: Comparison of Precision (P) and Recall (R) on ImageNet 64x64.
(a) IDDPM
8
10
20
50
R
R
R
R
P
R
ADM
0.34
0.48
0.46
0.50
0.51
0.48
0.65
0.52
0.73
0.57
ADM-IP
0.39
0.39
0.50
0.45
0.56
0.48
0.68
0.55
0.73
0.60
ADM-AT
0.40
0.50
0.50
0.50
0.55
0.49
0.69
0.52
0.77
0.59
(b) DDIM
10
20
50
P
R
R
P
R
P
R
P
R
ADM
0.42
0.47
0.54
0.56
0.58
0.58
0.65
0.60
0.69
0.61
ADM-IP
0.38
0.40
0.51
0.53
0.55
0.57
0.63
0.61
0.62
0.61
ADM-AT
0.44
0.43
0.58
0.55
0.62
0.56
0.69
0.59
0.72
0.61
(c) ES
10
20
50
R
R
ADM
0.40
0.44
0.52
0.47
0.58
0.48
0.69
0.55
0.73
0.59
ADM-IP
0.37
0.35
0.49
0.44
0.56
0.49
0.57
0.72
0.60
ADM-AT
0.44
0.46
0.58
0.48
0.63
0.49
0.73
0.55
0.76
0.59
(d) 3PM--ollver
8
10
20
50
P
R
R
R
R
P
R
ADM
0.51
0.49
0.65
0.58
0.67
0.60
0.69
0.62
0.69
0.62
ADM-IP
0.39
0.44
0.64
0.60
0.64
0.60
0.59
0.60
0.57
0.59
ADM-AT
0.56
0.50
0.68
0.57
0.69
0.59
0.72
0.60
0.71
0.61
Table 14: Comparison of different AT methods used in our AT framework. All models are trained with the same model-updating iterations while the efficient AT has less training time.
FID
Training Time Specdup
Methods
IDDPM-50
DDIM-50
ES-20
DPM-Solver-10
Standard AT PGD-3
4.02
3.37
6.42
7.60
1.0x
Efficient AT (Ours)
3,97
3,42
6.05
2.6x
Table 15: Comparison of different adversarial learning rate a of our AT framework on CIFAR10 32x32. IDDPM is adopted as the inference sampler.
NFFs
5
8
10
20
50
0.05
51.72
32.09
25.48
10.38
4.36
a 0.1
37.15
23.59
15.88
6.60
3.34
0.5
63.73
40.08
27.57
7.23
3.42
Table 16: Comparison of different adversarial learning rate a of our AT framrwork on ImageNet 64x64, IDDPM adopted as the inference sampler.
NFEs
5
8
10
20
50
0 0.1
56.92
27.39
24.06
10.17
5.82
0.5
45.65
23.79
19.18
8.28
4.01
a 0.8
46.92
28.46
22.47
9.70
4.25
28

---

Published as a conference paper at ICLR 2025
Table 17: Comparison of different perturbation norms (l,, (6)) of our AT framework on CIFARI0
32x32.
Perturbation Norm
IDDPM-50
DDIM-50
ES-20
DPM-Solver- 10
h
4,45
4.91
4.72
5.05
12
3.34
3.07
4.36
4.81
1o0
3,87
3.63
4.48
5.32
G.3 MORE ABLATION STUDY
Ablation We investigate the impact of adversarial learning in framework. The results of various a on CIFAR10 0 32x32 and ImageNet 64x64 are shown in Table 15 and Table[I6 respectively. We observe that set 0.1 better CIFAR1 32x32 and a5 ImageNet 64x64. That says, the image in larger size corresponds to larger optimal perturbation level a. We speculate this because we use the perturbation measured under E2 æ2-norm, where the 22-norm of vector will increase with its dimension.
0.5 is better for
During our experiments, we adopt fz-adversarial perturbation. Actually, perturbations in Euclidean space under different lp norm equivalent with each other, e.g., for vector 8 Rd it holds |, S || < Vd Therefore, we select ll2 as representation in our paper. Next, we explore the proposed ADM-AT under different adversarial perturbations.
Ablation on perturbation norm
The results are in Table 17
We found that our method under Ee-perturbation more stable and indeed has better performance, thus we suggest to use fg-perturbation as in the main body of this paper.
G.4 QUALITATIVE COMPARISONS
Figure 4: The qualitative comparsions of ADM-AT (top, FID 6.60). ADM-IP (middle, FID 7.81), and ADM (bottom, FID 10.58) on CIFAR1 0 32 32. We the IDDPM sampler with 20 NFEs for sampling.
Figure 6 7 show the qualitative comparisons between our proposed AT method and baselines. Our proposed AT method generates more realistic and figher-fidelity samples. We attribute this to our AT algorithm mitigates the distribution mismatch problem.
29

---

Published as a conference paper at ICLR 2025
Figure 5: The qualitative comparsions of ADM-AT (top, FID 6.20), ADM-IP (middle, FID 8.40) and ADM (bottom, FID 8.32) on ImageNet 64 64. We use the DDIM sampler with 20 NFEs for sampling.
Figure 6: The qualitative comparsions of LCM (left) and LCM-AT (right) with one-step generation.
The text prompt is photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece.
Figure 7: The qualitative comparsions of LCM (left) and LCM-AT (right) with one-step generation. The text prompt is Astronaut a jungle, cold color palette, muted colors, detailed, 8k.
30


Reviews:
[{"review_id": "ktVMVdGfdj", "reviewer": "Reviewer_tM1A", "summary": "The paper proposes to introduce DRO to address the distribution matching problem at training diffusion model.", "questions": "as above", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper present theories to show that DRO can help address the distribution matching problem in training and testing diffusion models.\n\n2. The improvement over baselines on Cifar and Imagenet64 show that DRO is useful.", "weaknesses": "1. There is no qualitative comparisons. Authors mainly conduct experiments on Cifar, ImageNet and Laion dataset. It would be better to put some images for more direct comparisons. In addition, the code is not provided. \n\n2. The efficiency comparison. I am wondering how much overhead it brings to adopt eq 14 instead of the classical denoising objective.  I am expecting that it is quite large.\n\nI am giving score of 6 based on the prerequisite that above two concerns are answered during rebuttal.", "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": "t80JQGSNjy", "reviewer": "Reviewer_ouMQ", "summary": "This paper studies the training of unconditional diffusion model. In particular, in order to achieve a better generation quality and enable robust learning of the score network, this paper develops a DRO-based method, and prove the DRO objective in training diffusion models can be formulated as an adversarial learning problem. The paper also identifies a similar mismatch issue in the recently proposed consistency model (CM) and demonstrates that AT can address this problem as well. The authors propose efficient AT for both DPM and CM, with empirical studies confirming the effectiveness of AT in enhancing diffusion-based models.", "questions": "1. some ablation studies for different perturbation levels $\\alpha$ should be given.\n2. Some discussions about different perturbation methods ($\\ell_1$, $\\ell_2$, or $\\ell_\\infty$) should be discussed.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper performs a theoretical analysis of diffusion models and identifies the distribution mismatch problem.\n\n2. This paper further builds a connection between the distribution robust optimization and adversarial learning for diffusion models, and develops an adversarial training method for diffusion models.\n\n3. This paper conducts efficient adversarial training methods on both diffusion models and consistency models in many tasks. Experimental results demonstrate the effectiveness of the developed algorithms.", "weaknesses": "1. In general, the algorithm developed in this paper is motivated by the distribution mismatch along the diffusion path. However, there is no experimental results to justify the motivation, there are also no experimental results to verify that the DRO framework can indeed help mitigate the distribution mismatch problem. \n\n2. Proposition 2 has already been discovered in existing theoretical papers [1], see their section 3.1. The authors should comment on this point around Proposition 2.\n\n3. The advantage of ADM-AT is not that significant compared with the ADM method, a more detailed ablation study or theoretical analysis on using adversarial noise or random Gaussian noise should be added.\n\n4. Some statements are not clearly presented. For instance, the description of ADM is not given, the norm notations $\\|\\|$ are abused, should that be $\\ell_1$, $\\ell_2$, or $\\ell_\\infty$?\n \n\n[1] Chen, Lee, and Lu, Improved Analysis of Score-based Generative Modeling: User-Friendly Bounds under Minimal Smoothness Assumptions, ICML 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3Qs4LKsAko", "reviewer": "Reviewer_eJv3", "summary": "This paper identifies the distribution mismatch problem in the training and sampling processes. Consequently, they propose a distributionally robust optimization procedure in the training to bridge the gap. The authors apply the method to both diffusion models and the consistent model, and demonstrate the effectiveness of the proposed method on several benchmarks.", "questions": "The paper is a good one in general. I like how the problem is formulated and how the solution is derived. However, given the current evaluation (see the weakness), I am not fully convinced the proposed method is an effective way to deal with the problem. I would like to see how the authors respond to my concerns.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Identifying and formulating the distribution mismatch problem in diffusion model is an important problem in practice.\n2. The proposed solution is elegant, supported by sufficient theoretical analysis. The derivations of the solution is clear and sound.\n3. The writing is fairly clear.", "weaknesses": "My main concern on this paper is the evaluation. Currently the proposed method is only evaluated using the ADM model. I wonder whether the effectiveness on more advanced model such as the stable diffusion still holds?\n\nFurthermore, the authors only use FID score as the evaluation metric, while it is easy to evaluate the results using other metrics such as IS, sFID, precision, recall, as done in the ADM paper. Why these metrics are not included?", "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": "XbdBoTH2Kb", "reviewer": "Reviewer_gZkr", "summary": "This paper points out the distribution mismatching problem in traditional training of diffusion-based models (DPM) and proposes to conduct efficient adversarial training (AT) during the training of DPM to mitigate this problem. Theoretical analysis is strong enough to support its argument and experiments also verify the effectiveness of the proposed method.", "questions": "1.  As the weakness above, for Table 1 and Table 2, more NFEs should also be verified.\n\n2. Why not also try generation using consistency models on benchmark datasets such as CIFAR10 and ImageNet, which can be more common and convincing?\n\n3. Derivations in supplementary material should be checked carefully and written with more details. \n\n4. Why efficient AT can improve performance compared with PGD is a bit confusing. Intuitively, PGD should be more accurate to find $\\delta_t$, thus more deep insights should be provided here.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The motivation for mitigating distribution mismatching is clear and important for efficient sampling. \n\n2. This paper provides strong theoretical support for implementing adversarial training to correct distribution mismatching, making this method convincing.", "weaknesses": "1. The experimental results may not be enough, for example, for Table 1 and Table 2, more NFEs should also be verified, although this method can improve efficient sampling, whether is adaptable and robust for more denoising steps should also be verified. \n\n2. Some complex derivations in supplementary material are too brief to understand, such as Eq(30) and Eq(59-62), I'm not sure if there are any typos in them, I suggest checking the equations carefully and modifying them.", "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)