Constrained Diffusion Models via Dual Training
Authors: Shervin Khalafi
Dongsheng Ding
Alejandro Ribeiro
Venue: NeurIPS 2024 poster (2024)
Keywords: Constrained diffusion model
constrained optimization
Lagrangian method
dual algorithm
Primary area: diffusion_based_models

Abstract: Diffusion models have attained prominence for their ability to synthesize a probability distribution for a given dataset via a diffusion process,  enabling the generation of new data points with high fidelity. However, diffusion processes are prone to generating samples that reflect biases in a training dataset. To address this issue, we develop constrained diffusion models by imposing diffusion constraints based on desired distributions that are informed by requirements. Specifically, we cast the training of diffusion models under requirements as a constrained distribution optimization problem that aims to reduce the distribution difference between original and generated data while obeying constraints on the distribution of generated data. We show that our constrained diffusion models generate new data from a mixture data distribution that achieves the optimal trade-off among objective and constraints. To train constrained diffusion models, we develop a dual training algorithm and characterize the optimality of the trained constrained diffusion model. We empirically demonstrate the effectiveness of our constrained models in two constrained generation tasks: (i) we consider a dataset with one or more underrepresented classes where we train the model with constraints to ensure fairly sampling from all classes during inference; (ii) we fine-tune a pre-trained diffusion model to sample from a new dataset while avoiding overfitting.

Constrained Diffusion Models via Dual Training
Shervin Khalafi
Dongsheng Ding"
Alejandro Ribeiro
{shervink, donggshed aribeiro) upenn. ude
University of Pennsylvania
Abstract
Diffusion models have attained prominence for their ability to synthesize probabil- ity distribution for a given dataset via a diffusion process, enabling the generation of new data points with high fidelity. However, diffusion processes are prone generating samples that reflect biases in a training dataset. To address this issue, develop constrained diffusion models by imposing diffusion constraints based on desired distributions that are informed by requirements. Specifically, we cast constrained distribution
the training of diffusion models under requirements optimization problem that aims to reduce the distribution difference between origi- nal and generated data while obeying constraints on the distribution of generated data. We show that our constrained diffusion models generate new data from mixture data distribution that achieves the optimal trade-off among objective and constraints. To train constrained diffusion models, we develop a dual train- ing algorithm and characterize the optimality of the trained constrained diffusion model. We empirically demonstrate the effectiveness of our constrained models two constrained generation tasks: (i) we consider dataset with one or more underrepresented classes where train the model with constraints to ensure fairly sampling from all classes during inference; (ii) we fine-tune a pre-trained diffusion model to sample from new dataset while avoiding overfitting.
Introduction
Diffusion models have become a driving force modern generative modeling, achieving ground- breaking performance in tasks ranging from image/video/audio generation [64, 6, 43] molecular design for drug discovery [75, 74]. Diffusion models learn diffusion processes that produce a probability distribution for a given dataset (e.g., images) from which
data (e.g., classic models [66, 34, 68, 78]). As diffusion models are used to generate data with societal impacts, e.g., art generation and content creation for media, they must comply with requirements from specific domains, e.g., social fairness in image generation [54, 57], aesthetic properties of images [12, 13), bioactivity in molecule generation [39], and more [40, 81, 8. 19, 79, 14].
generate
Classic diffusion models [66, 34, 68] have been extended to generate data under different requirements through either trstncrinciple methods fine-tuning, In image generation with fairness, for instance, first principle methods directly mitigate biases towards social/ethical identities by revising the training loss functions per biased/unbiased data [49, 41]; while fine-tuning methods align biased diffusion models with desired data distributions by optimizing the associated metrics [24, 72, 71, 70]. Although these methods allow to equip diffusion models with specific requirements. they are often designed for particular generation tasks and do not provide transparency on how these requirements are satisfied. Since diffusion models are trained by minimizing the loss functions of diffusion processes, is natural to incorporate requirements into diffusion models by imposing constraints on these optimization
Corresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

problems. Therefore, it is imperative to develop diffusion models under constraints by generalizing constrained learning methods and theory (e.g., [9, 10, 22, 36)) for diffusion models.
In this work, we formulate the training of diffusion models under requirements as a constrained distribution optimization problem in which the objective function measures training loss induced by the original data distribution, and the constraints require other training losses induced by some desirable data distributions to be small. This constrained formulation can be instantiated for several critical requirements. For instance, to promote fairness for unrepresented groups, the constraints can encode the closeness of the model the distributions of underrepresented data. Compared with the typical importance re- e--wegghting method [41], our constrained formulation provides an optimal trade-off between matching given data distribution and following reference distribution. Not limited to constraints with other desirable data distributions, our constrained formulation captures more general requirements. For instance, when adapting pretrained diffusion model to new data, the constraints require the model to be close to the pretrained model, not degrading the original generation ability. Specifically, our main contribution is three-fold.
(i)
We propose and analyze constrained diffusion models from the perspective of constrained optimization in an KL divergences between the model and our desired data distributions to be under some thresholds. We exploit the strong duality in convex optimization to show that our constrained diffusion models generate new data from mixture data distribution that achieves the optimal trade-off among objective and constraints.
distribution space, where the constraints require
(ii) To train constrained diffusion models, we introduce parametrized constrained diffusion models and develop a Lagrangian- based dual training algorithm. We exploit the relation between un/parametrized problems to show that constrained diffusion models generate new data from the optimal mixture distribution, up to some
errors.
(iii) We empirically demonstrate the merit of our constrained diffusion models in two afore- mentioned requirements. In the fair generation task, we show that our constrained model promotes sampling more from the minority classes compared to unconstrained models, lead- ing to fair sampling across all classes. In the adaptation task, we show that our fine-tuned constrained model learns to generate new data without significantly degrading the original generation ability, compared to the unconstrained model which tends to overfit to new data.
Related work. As a first trst-principle method, our constrained diffusion approach is more relevant to diffusion models that incorporate requirements in distribution space [20, 35, 49, 21, 60, 28). rather than those applied in sample space [37, 53, 29, 27, 26, 17, 48, 25]. In comparison with conditional diffusion models that restrict generation through conditional information [20, 35, 1J. our constrained diffusion models impose distribution constraints within the constrained optimization framework. Compared to compositional generation [49, 21, 60) and fair diffusion [28], our work provides a constrained learning approach to balance different distribution models using Lagrangian multipliers, which is different from equal weights [49, 21], hyperparameter [60] or fair guidance [28]. Our constrained approach is also relevant to the importance re- re-ieighting method for diffusion models [41] and GANs [16], which reduces bias in a dataset by re-weighting it with pre-trained debiased density ratio. In contrast, we design our diffusion models to mitigate the bias in dataset by imposing distribution constraints without pre-training In addition to being distinct from existing methods, we provide systematic study of our constrained diffusion models, covering duality analysis, dual-based algorithm design, and convergence analysis, which is absent in the diffusion model literature.
Our work is also pertinent to recently surging fine-tuning and alignment methods that aim to improve pre-trained diffusion models by optimizing their downstream performance, e.g., aesthetic scores of images. In reward-guided fine-tuning methods, such as supervised learning [45, 80, 76), control-based feedback learning [77, 65, 71, 70, 18], and reinforcement learning [23, 32, 24, 72, 5, 82], reward functions have to be pre-trained from an authentic evaluation dataset, and the trade-off for reward functions in pre-trained diffusion models is often regulated heuristically. In contrast, our constrained approach directly minimizes the gap between tne-tuning model and high-quality dataset with the desired properties, while ensuring the generated outputs being close to that of pre-trained models.
Compared to other generative models under requirements (e.g., VAEs [61], GANs [16]), and classical sampling methods (e.g., Langevin dynamics and Stein variational gradient descent [50]), our work different because we focus on diffusion-based generative models.
2

---

2
Preliminaries
We overview diffusion models from the perspective of variational inference [55] by presenting forward/backward processes in Section 2.1, and the evidence lower bound in Section 2.2.
2.1 Forward and backward processes
The forward process begins with a true data sample zn E R and generates latent variables (#, [x+)7- from = 1 to T by adding white noise recursively. The joint distribution of latent variables results from conditional probabilities ((21:1 [2)) II-= g Ø(-1), where the distribution of latent variable conditioned on the previous latent 2i- is given by Gaussian distribution ((Tx 1x-1) = (xx [t, (t)T), where 14 = yoi 1 and (1) = a₁ Since the forward process is linear Gaussian model with pre-selected mean and variance, it is solely determined by the data distribution ((r)). We often refer to ((2)) as a forward process.
The backward process begins with the latent æт sampled from the standard Gaussian p(ær) NxT; 0, I), and decodes latent variables from = T
=
0 with a joint distribution
T
p(ZOT) = p(æT) II p p xx)).
(I)
t=1
Here, is our distribution model that can be used to generate new samples. We denote by P the set of all joint distributions over Zo:T in form of (1), where p(z;-1 2x) is conditional Gaussian with
fixed variance (see Appendix A). Throughout the paper, we work in the convergent regimes of the
backward process (c.g.. [15, 11, 2J). Without loss of generality, we adopt the convergent regime in [47] by taking the scheduling parameter
1/7 and = 1 cT min ((1 a1)( + +er)', 1) for > and the variance as of 1/a, 1, where er == ai log(T) /T and co. C1 are some constants. Hence, oer implies q(=x) = ((rr;; 0, I). Thus, (( E P. Also, âj implies ((x1) ) = g(zo).. It is ready to generalize our results to other diffusion processes (e.g., [46, 38, 3)).
2.2 The evidence lower bound (ELBO)
Denote the KL divergence of distribution q from distribution p by DxI. (q [) Ex~g(z) log( g(x)
p(zj/
Generative diffusion modeling aims to generate samples whose distribution close that of an observed dataset of samples. Formally, we express this objective as maximizing the log-likelihood of an observation generated by the diffusion model: maximizc, Eng( log p(#)) ]. where
(g( [log p = E(p: q) + Eg(zo) DKL (4(21r)
xx) p(T1r xo))]
(2)
p(rpr)) is known as ELBO in variational inference [7, 55]. and E(p; q) == (y( Ey(aur (=)) log q(wi 20] Alternatively, aim to minimize the KL divergence between the forward/hackward processes,
DEL (g(#orr) I((uz))
E(p: q) + Ey(ra) [log qxx))
(3)
Thus, connect the likelikel maximization the KL divergence minimization via ELBO.
Lemma (Equivalent formulations). The ELBO maximization and the KL divergence minimization are equivalent over the distribution space P, and the unique solution of these two problems is solution for the log-likelihood maximization problem, i.e.,
maximize E(p: q) P
minimize DRL (q(zor)) p(Fo:T)) PeP
maximize pe
Eg(zo) logo((()))
See Appendix B.1 for proof. Lemma states that improving the ELBO score increases the likelihood of backward process that generates the data, together with the fit of backward process to the forward process. Hence, finding the best backward process becomes optimizing one of three equivalent objectives. In practice, ELBOO serves loss function approximated by
T
E(p:g) x
Σ
By(na) Ey(sr 2n) DKI. (g(xx-1 [x,, x0) [| [x))]]
(4)
3

---

With the variance schedule described in Section 2.1, it is known that this approximation is almost exact (see Appendix A and also [42]), which our focal setting. Using standard diffusion derivations [55]. the ELBO maximization can be shown to equal to quadratic loss minimization,
minimize Eu,,,, 2t [[(( V log {(x))) RES
(5)
where Exo,t, is an expectation over the data distribution (zo)), a discrete distribution (t) from 2 to T, and the forward process g( ro) at time given the data sample zo; see Appendix for details. The minimization is done to find a function 8 E S that can best predict the gradient of the forward process over data V log g(x4), commonly called the (Stein) score function, where S is a of valid score functions mapping from Rd N to Rd, In practice, however, we parametrize the estimator $(z,,t, t) as S0(22, t) with parameter 6, which gives our focal objective of generative mod- cling: minimizeg e e Eeg, a, (0t,
" A parametrized form of p(æ;-1 xx)) and the backward process has parametrized joint distribution py (Zorr). ). We remark that the prediction problem (5) can be also be formulated as data or noise prediction instead [55], with our results directly transferable to these formulations.
t) v log g(2+)]l
associated with So (zz, t) is denoted by po ( 24)
3 Variational constrained diffusion models
We introduce constrained diffusion models by considering the unparametrized set of joint distributions in Section 3. and illustrating constraints via two examples in Section 3.2.
3.1 KL divergence-conssrainee diffusion model: unparametrized case
The standard diffusion model specifies single data distribution, denoted by in Lemma 1. To account for other generation requirements, we introduce m additional data distributions {q p 1 that represent desired properties on generated data. To incorporate properties into the diffusion model, we formulate an unparametrized KL divergence-constrrinn optimization problem.
minimize P
DRL. (4(20.7) (((or)
(U-KL)
subject to DEL (4' (morr p(xo:r)) S b, for i 1, m.
Let an optimal solution to Problem (U-KL) be p². Then the optimal value of the objective function is
F == DKI.6 (q || p*). Problem (U-KL) aims to find a model p" that generates data from the original distribution q while staying close m distributions {q " that encode our desired properties, e.g..
unbiasedness towards minorities. Let the Lagrangian for Problem (U-KL) be
m
C(p, A) = DKL. (q(x0r)) ll (((oTT)) +
Σ ==1
Ae (DKL. (zo:T) [I((ror)) b:)
(6)
for A 2 0. The dual function g(A) is given by g(A) == minpep L(p, A), which is always concave.
To make Problem (U-KL) meaningful, we assume the constraints are strictly satisfied by some model. Assumption 1 (Strict feasibility). There exists model E P and C > 0 such that Dxi (xor)) || (((o.))) S b, for all i i= 1,,...,m
Let optimal dual variable of Problem (U-KL) be X" argmaxx 2 o g(A), and the optimal value of the dual function be D* : 9(>*). From weak duality, the duality gap is non-negative, i.c., F+ - D* Z O. Moreover, due to the convexity KL divergence, Problem (U-KL) convex optimization problem, and thus it satisfies strong duality; Appendix B.2 for proof.
Lemma 2 (Strong duality). Let Assumption 1 hold. Then, Problem (U-KL) has zero duality gap, i.e.
F+ = D Moreover, (p" (p*,>*) is an optimal primal-dual pair of Problem (U-KL).
(A) : := ET (11+A11 for X 2 0. We denote by Let mixture data distribution be Amix (X)
Umix. (Iro.T) joint distribution of the forward process with data distribution 4mix Leveraging strong duality, we show that an optimal model can obtained by solving an equivalent unconstrained problem in Theorem and its proof is deferred to Appendix B.3.
(A)

---

Theorem (Optimal constrained model), Let Assumption I hold. Then, Problem (U-KL) equals
minimize
DRI. 2 Ip(oorr)
Imix
(U-MIX)
PEP
(2)) (2") (1OTT) is the joint distribution of the forward process at optimal dual variable X*. Umix
where
Theorem states that the KL divergence-constrane constrained problem reduces to an unconstrained KL diver- gence minimization problem. We notice that the KL divergence is zero if and only if two probability
(X*) distributions match each other. Hence, Qmix (Zo:T) the optimal solution Problem (U-MIX).
Corollary 1. Let Assumption I hold. Then, the solution of Problem (U-MIX), ie., p' (fo.r) = (2") Imix (wo:r), is the solution of Problem (U-KL).
Let D' := b' == based constrained optimization problem.
Ei (mo). log (zo) J. Application of Equality (3) to Problem (U-KL) yields an ELBO-
minimize
E(p:
PEP
(U-ELBO)
subject to - E(ppg')) S B fori=1,.. I....,
Recall the model representation in Section 2.2, can characterize each joint distribution p P with function 3 E S. Moreover, ELBO reduces to the denoising matching term that has a simplified quadratic form given in Section 2.2. With this reformulation mind, we cast Problem (U-ELBO) into a convex optimization problem over the function space S,
minimize SES subject to (xo()
[ 8) log log (1x=)))
(U-LOSS)
[|((z,,)) - V log (xx))
S b' for: =ori=L......
where 8 = (b' r) Here, the notation v a constant shift due to the variance mismatch term; see it in Appendix A. We note that scaling or shifting objective and constraints from both sides with some constants doesn't it alter the solution to a constrained optimization problem. Thus, the key difference between Problems (U-KL) and (U-LOSS) is the optimization variable (respectively, p and 3). Let the Lagrangian (3, X) for Problem (U-LOSS) be
By(x)) t, [18(21,)) ()- - V Voog (()g(()))) + Σ As (By (aa). [| V Vlog (xx)) T
Let the associated dual function be (A) == minges s (8, A) for 2 0. Hence, g(A) and g. (A) have the same maximizer At, and the partial minimizer 8* argminge (3 )*) is the solution to Problem (U-LOSS). Hence, an optimal primal-dual pair (3", x*) to Problem (U-LOSS) gives an optimal primal-dual pair (p* X²) for Problem (U-KL), where p" is joint distribution of the backward process induced by 3 By this dual property, we take a dual perspective to train constrained diffusion models: we maximize the dual function go (A) obtain the optimal dual variable X", and then recover the solution S by minimizing the Lagrangian C. (²)
3.2 Examples of KL divergence constraints
To illustrate our KL-divergence constraints, we provide two generation tasks of exemplary.
(i) Fairness to underrepresented classes. We consider a fair image generation task in which some classes are underrepresented in the available training dataset. An example of this would be the Celeb-A dataset [51] which contains pictures of celebrity faces with those labeled as male being underrepresented (42% Male vs 58% Female). To promote representation of the under-represented classes, we can pose it as an instance of Problem (U-KL), where each q' denotes the distribution of an under-represented subset minority class of q.
(ii) Adapting pretrained model to new data. Given a pretrained diffusion model over some original dataset that is no longer accessible, we aim to fine-tune the pretrained model for generating data from new data distribution. Similarly, we can pose this as an instance of Problem (U-KL.), where g² denotes the distribution of the new data and g is the distribution of samples generated by the pre-trained model.
5

---

Let h; == (zo) (20) [log g' (zo) be the differential entropy of data distribution q². We relate the KL divergence constraints with the optimal dual variable through entropy in Theorem
Theorem Let Assumption hold, and the supports of data distributions q and {q [Q=) be disjoint Then, the optimal dual variables X to Problem (U-ELBO) are given by
X
(X+)T e bb for i 1, .,..., m.
See Appendix B.4 for proof. Theorem characterizes the mixture weights in the target distribution (A"): (i) the tighter a constraint (i.e., smaller threshold b,), the more the model will sample from the associated distribution; (ii) for the same constraint thresholds, the model will sample more often from the associated distributions that have higher entropy h; Assumption can be relaxed to a feasibility condition for Problem (U-KL); see Lemma 6 in Appendix B.4.
Omix
Parametrization and dual training algorithm
Having introduced unparametrized models, we move to parametrization for constrained diffusion models in Section 4.1, provide optimality analysis of Lagrangian- baased dual method in Section 4.2, and present a practical dual training algorithm in Section 4.3.
4.1 KL diverrencee-constrained diffusion model: parametrized case
With the parametrized model pa for 0 E O, we present a parameterized constrained problem,
minimize 0e
DKL (g(=0.r) (x(:)))
(P-KI.)
subject to
DEL (g' ( po( S for 1,...,m.
Let the Lagrangian for Problem (P-KL) be [(0, A) := L(po, A). The associated dual function g(A) is given by g(A) == ming e L(0, X). Let optimal solution to Problem (P-KL) be 0*. We denote p == Pe and p == pa*, and the optimal objective by F == Dx l1 Let an optimal dual variable be X argmax 2 o ((A) and the optimal value of the dual function be D g(X*).
Problem (P-KL) is non-convex in parameter space, and strong duality does not hold any more. Thus, unparametrized results in Section 3.1 don't 't directly apply to Problem (P-KL). For in- stance, it's invalid find an optimal solution via an unconstrained problem as in Theorem 1, ie., (x')) argmin, e L(0 X*) doesn't 't equal p* The effect of parametrization has to be characterized. However, regardless of parametrization weak duality always holds, i.c.. F*
- DD 2 0.
To quantify the optimality of 5 (X*) (closeness of it to ) study a practical representation of model Po as a parametrized function Kg E So, where So is the set of all parametrized score functions.
Problem (U-LOSS) is in parametrized form of
minimize
E (,
|| t) log ((=:))]²]
subject to By((s)), [|50(112,)) 1) (x=)) < & for 1,...,m. (P-LOSS)
where D' = (b' - v)/a. We note that Problem (P-LOSS) is equivalent to Problem (P-KL). Thus, let the Lagrangian of Problem (P-LOSS) be (0. = L. A), and the dual function (A) == minimizeg e C(B Since g(A) and (e) have the same maximizer X+, 8" argming e (0, which naturally gives dual training algorithm in Algorithm 1.
Denote 3* = Sz.. Thus, s" --inuuued diffusion model is given by 5a (X*). Algorithm works as dual ascent method with two natural steps: (i) find a diffusion model with fixed dual variable X(h): and (ii) update the dual variable using the (sub)gradient of the Lagrangian (3 A). It is known that Algorithm converges to X since the dual function Øs (A) is concave. However, such convergence in the dual domain doesn't 't provide optimality guarantee on the primal solution / (X*) due to the
(h),
non-convexity in parameter space. We next exploit the optimization properties of unparametrized diffusion models in Section 3.1 to characterize the optimality of the dual training algorithm.

---

Algorithm Constrained Diffusion Models via Dual Training
1: Input: total diffusion steps T, diffusion parameter Oz,, total iterations H, stepsize
2: Initialize: X(1) O.
3: for h == 1, 1, .H do
4: Compute model Sy(h) argmin (Se, ((h)).
8ee
5: Update the dual variable
x.(h+1) [Ar(h)) (Exo gt [1aa())) ( 8)) b)] for all i.
6:
end for
4.2
Optimality analysis of dual training algorithm
We analyze the optimality of (X*) as measured by its distance to mix i.c., TV pr (X*)), where we denote the total variation distance between two probability distributions and by TV p) := 2 Ip(x) (dr) We first exploit the convergence analysis of diffusion models, and then characterize the additional error induced by parametrization.
Let us begin with the difference between p (A*) and 4mix at X Denote a partial minimizer of the Lagrangian by p (A) E argming e (6, A) for À 2 0. Noting that minimizeg (0, A) is an unconstrained diffusion problem, we are ready to quantify the difference between p² (A) and g for any À Z 0 using the convergence theory of diffusion models. To do so, we assume the boundedness of samples from a mixed data distribution Imix for 2 0, and a small score estimation error.
(X*)
Assumption 2 (Boundedness of data). The data samples generated from Qmix are bounded, i.e.. (Iiz|| T Ing Omix
for any X and some large constant >> O.
Assumption 3 (Boundedness of score estimation error). The score estimator Bo (o( estimates the data samples from mix with bounded score matching error Excore.
E 4mm (2o), Re
[[ t) V log q(x))1²
22
score
for any X 2 O, where (N) Smn (zo), t, xt is an expectation over the mixed data distribution qmix (zo)). uniform distribution over from 2 to T. and forward process q(Ix (o) given the data sample To-
(N)
Since data samples are bounded, Assumption is mild in practice. Assumption 3 the typical score matching error that is near zero if the function class Sp is sufficiently rich.
(X) With Assumptions 2 and 3, below we bound the TV distance between I and pa (A) using the convergence theory of diffusion models from [47]; see Appendix B.5 for proof.
Lemma 3 (Convergence of diffusion model). Let Assumptions 2 and 3 hold. Then, the TV distance from p (A) to Zmix is bounded by
(A)
TV m (*()) I 1 DEL (1)) (((( S
d log VT
+ Vd (log* T) Excore-
Lemma 3 states that the TV distance between Imix and 5a (A) decays to zero with a sublinear rate
(X)
0
Vi up to a score matching error (Escorc). When the diffusion time T is large, the TV distance
(A) between 4mix and pa (A) is dominated by the matching error. Substitution of À x into (7)
yields an upper bound on TV (X*) (A+)),, which is the second term of the inequality
(A*) + TV TV (gmix (X))) S TV "Vginic 9
(X*) A
(*((*)
(8)
Next, we quantify the gap between X and X*, which lets us bound the first term on the RHS of (8) and complete the optimality analysis. To analyze the parametrized optimal dual variable X*, we introduce the richness of the parametrized class Sp and redundancy of constraints at 3* below.
7

---

Assumption 4 (Richness of parametrization). For any function 3 E S, there exists parameter & E Q
such that 6g S2 S v, where Il-Ilz, is with respect the forward process.
Assumption (Redundancy of constraints). There exists a > 0 such that
m
inf ||
Σ
d Va (o (ae,t) log [z+)]
2o
(9)
=1
I/2
where V3 is the Fréchet derivative over the function á and S* is a solution to Problem (U-LOSS).
Assumption 4 is mild since the gap is small for expressive neural networks [56,31]. Assumption 5 captures the linear independence of constraints, which is similarly used in optimization [4].
Due to Assumption 2, we set the function class to be bounded || $ 11 I T := R. Using Problem (U-LOSS), we prove that the unparametrized dual function (A is differentiable, and strongly-conean over H. with parameter μ. where H. x^r* + (1 7),,, [o, 1]} and = (a/ (1 + max ( (|*x* ,which leads to Lemma 4; see Appendix B.6 for their proofs.
Lemma 4. Let Assumptions and 5 hold. Then, - ||*+- R (1 ||*1|1) v.
(x*) is bounded by |パ-- - X-11, (see Appendix B.6), application of Lemma 3 and Lemma 4 to (8) leads to Theorem 3; see Appendix 7 for proof.
Since TV Amiix
Omix
Theorem (Optimality of constrained diffusion model). Let Assumptions 1-5 hold. Then, the total mmix is upper bounded by
variation distance between 5² (X*) and
d² log T
8 V
TV (aixx ('()))
+
R (1
11+1))) + Vd (log* T) Escore-
VT
Theorem 3 states that the TV distance between pa (X+) and fixx decays to with sublinear rate O( O( (T)))) ). up to parametrization gap O(v2) and a prediction error ((Esore). When the parametriza- tion rich enough, the parametrization gap v and the prediction error (Escor) are nearly zero. In this case, if the diffuusin time T is large, then (X*) is close to Laus in TV distance, which recovers the ideal optimal constrained model the unparametrized case in Section 3.1.
4.3 Practical dual training algorithm
Having established the optimality of our dual training method, we futher turn Algorithm into a practical algorithm. First, we relax the computation of a diffusion model @a(h) in line 4 of Algorithm I to be approximate: (.)) A(h)) I min L. (Sa, X(h)) + = Epppoxx where 22 Epppox is
approximation error of training diffusion model given A(h). Second, replace the gradi- in line 5 of Algorithm by stochastic gradient E [[ t) V log l(;)();))1 which enables Algorithm to be a stochastic algorithm, where E. of pfge,,,aa w. is unbiased es- To analyze this approximate and stochastic variant of Algorithm 1, timate of Ee f, it is usefuul to introduce the maximum parametrized dual function in history up to step h by Bees(()) (h) := max h Ds (A(h"), and an upper bound of the second-order moment of stochastic gradient 5 = ΣT= R (R wigi, w. (H)(()) (=( )² X(M)]-
Denote the dual variable that achieves (h) by Abest- To bound the TV distance between pr (Anes) (Amai) and 5" (Avea) using Lemma 3. The rest is to analyze the convergence of Avest to X" via application of martingale convergence. We defer their proofs to Appendix B.8 and present the optimality of px (Area) in Theorem 4.
and 4mix. we check the TV distance between Jmix
Theorem 4 (Optimality of approximate constrained diffusion model). Let Assumptions 1-5 hold. Then, the total variation distance between 5e (Nnes) and tinixx is upper bounded by
d log T
8R (1 + | (eeuull))) v Vd (log" T) Excure + Sppprox + ?
TV (fiix: D² (Abea)) S
VT
8

---

Theorem 4 states that the TV distance between 10 (Abea) and Imix decays to zero with a sublinear rate O (T ) to a parametrization gap (v²), a score matching error (E an approximation error (p and stepsize O()). When the parametrization is rich enough, the parametrization gap and the score matching error O (E are near zero. Thus, if the diffusion time T is large, then the closeness of 5² (Abest to mixx in TV distance is governed by the appproximation error and stepsize.
5 Computational experiments
We demonstrate the effectiveness of constrained diffusion models trained by our dual training algorithm in two constrained settings in Section 3.2; see Appendix C for experimental details.
Fairness underrepresented classes. We train constrained diffusion models over three datasets: MNIST digits [44], Celeb-A faces [51], and Image-Net' [63]. For MNIST and Image-Net, we create dataset for the distribution q in (P-LOSS) by taking subset of the dataset with equal number of samples from each class. Then we make some classes under-represented by removing their samples. For each distribution g', we use samples from the associated underrepresented class. For Celeb-A, our approach is similar to MNIST except we don't tt remove any samples due to the existence of class imbalance in the dataset (58% female vs 42% male). For Image-Net, since the images are of high resolution, we employ the latent diffusion scheme [62] by imposing distribution constraints in latent space. Figures 1-3 show that our constrained model samples more often from the underrepresented classes (MNIST: 4, 5, 7; Celeb-A: male; Image-Net: 'Cassette player', 'French horn', and "Golf ball'), leading to a more uniform sampling over all classes. This reflects our theoretical insights on promoting fairness for minority classes (see Section 3.2). Quantitatively, we observe fairly lower FID scores when training over the same dataset but with constraints (see Appendix C for further discussion on FID scores). Furthermore, Image-Net experiment shows that approach extends to the state-of-the trrt---rrt diffusion models in latent space.
3
2
0
8
8
6
9
9
0
9
8
9
2
~
3
0
8
3
4
8
4
6
7
frequencies
338
1
3
8
9
1
3
8
6
6
9
2
5
/
6
3
6
3
6
3
5
5
8
7
6 5
6
8
9
1
8
6 0 0
7
70
6
8
7
135
7
3
1
9
9
2
8
8 2
3
9 4
9
2
6
2
/
0
4
8
5
7
3
3
/
2
8
8
0
9
7
2
9
O
6
3
4
classes
Figure 1: Generation performance comparison of constrained and unconstrained models that trained on MNIST with three minorities: 4, 5, 7. (Left) ) Frequencies of ten digits that are generated by an unconstrained model and our constrained model (m)); (Middle) Generated digits from unconstrained model (FID 15.9): ); (Right) ) Generated digits from our constrained model (FID 13.4) ).
Adapting pretrained model to new data. Given a pretrained diffusion model over some original dataset Pretrain- we fine-tune the pretrained model for generating data that resemble D Dnew. To cast this problem into (P-KL), we let the data distribution be Doow. i.c., ( = Qnew (Zo.T) and the constrained distribution be the pre-trained model, i.e., g (#x.T) = pe,.. (ID.T). ). our experiments, we pretrain diffusion model on subset of MNIST digits excluding a class of digits (MNIST: 9), and fine-tune this model using samples of the excluded digit. Figure 4 shows that our constrained fine-tuned model samples from the new class as well all previous classes, whereas the model fine-tuned without the constraint quickly overfits to the new dataset (see Appendix C for details). Our constrained model generates much better high-quality samples than the unconstrained model.
Conclusion
We have presented a constrained optimization framework for training diffusion models under distri- bution constraints. We have developed Lagrangian- dual algorithm to train such constrained
'ee use subset of ten classes from Image-Net: "Tench Fish', 'English Springer Dog", 'Cassette Player', 'Chain saw", 'Church', 'French Horn', 'Garbage Truck', 'Gas Pump', "Golf Ball', 'Parachute.'

---

030
frequencies
670
030
000
classes
Figure 2: Generation performance comparison of constrained and unconstrained models that are trained on Celeb-A with male minority. (Left) Frequencies of two genders that are generated by an unconstrained model and our constrained model ); (Middle) Generated faces from unconstrained model FID 19.6); Righht) ) Generated faces from our constrained model (FID 11.6).
030
frequencies
338
8:6
0:4
classes
Figure 3: Generation performance comparison of constrained and unconstrained models that are trained on Image-Net with minority classes: "Cassette player' (2), 'French horn' (5), and 'Golf ball" (8). (Left) ) Frequencies of ten classes that are generated by an unconstrained model and our constrained model (m);; (Middle) Generated images from unconstrained model (FID 36.0). ); (Right) Generated images from our constrained model (FID 27.3).
9
8
2
9
0
4
K 7
0
1
2 4
2
9 P
1
7
6
6
frequencies
f
8
5
0
f 4
4
7 7 8 A
7
3
/
4
0
R
7
C
6
6
0
0
4
0
a
4
4
8
4 9
3
5
9
7
Q
6 3 7
4
8
6
9
0
0
7
7
83d
1 9
8
0
8
0
2
9
3
6
classes
Figure 4: Fine-tuming performance comparison of constrained and unconstrained models that are trained MNIST. (Left) Frequencies of ten digits that generated by pre-trained model without digit 9 and our fine-tuned constrained model (-): ); (Middle) Generated digits from unconstrained model (FID 45.9). Right) Generated digits from our constrained model (FID 25.2). ).
diffusion models. Our theoretical analysis shows that our constrained diffusion model generates new data from an optimal mixture data distribution that satisfies the constraints, and we have demonstrated the effectiveness of our distribution constraints in reducing bias across three widely-used datasets.
This work directly stimulates several research directions: (i) extend distribution constraints other domain constraints, e.g., mirror diffusion [48]; (ii) incorporate conditional generations, e.g., text-to-image generation [28, 65], into our constrained diffusion models; (iii) conduct experiments with text-to-image datasets to identify and address biases; (iv) improve the convergence theory using more advanced diffusion processes.
10

---

Acknowledgments
We thank reviewers and program chairs for providing helpful comments.
References
[1] A. Bansal, H.-M. Chu, A. Schwarzschild, S. Sengupta. M. Goldblum, Geiping, and T. Gold- stein. Universal guidance for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 853-852, 2023.
[2] Benton, V. Bortoli, A. Doucet, and G. Deligiannidis. Linear convergence bounds for diffusion models via stochastic localization. In Proceedings of the International Conference on Learning Representations, 2024.
[3] J. Benton, V. De Bortoli, A. Doucet, and G. Deligiannidis. Nearly d-linear convergence bounds for diffusion models via stochastic localization. In Proceedings of the International Conference on Learning Representations, 2024.
[4] D. P. Bertsekas. Nonlinear programming. Athena Scientific, 2016.
[5] K. Black, M. Janner, Y. Du, I. Kostrikov, and S. Levine. Training diffusion models with reinforce- ment learning. In Proceedings of the International Conference on Learning Representations, 2024.
[6] A. Blattmann, R. Rombach, H. Ling, T. Dockhorn, S. W. Kim, S. Fidler, and K. Kreis. Align your latents: High-resolution video synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 22563-22575, 2023.
[7] D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518):859 877, 2017.
[8] H. Cao, C. Tan, Z. Gao, Y. Xu, G. Chen, P.-A. Heng, and S. Z. Li. A survey on generative diffusion models. IEEE Transactions Knowledge and Data Engineering, 2024.
[9] I. Chamon and A. Ribeiro. Probably approximately correct constrained learning. In Proceedings of the Advances in Neural Information Processing Systems, volume 33, pages 16722-16735, 2020.
[10] L. Chamon, S. Paternain, M. Calvo-Fullana, and Ribeiro. Constrained learning with non-convex losses. IEFE Transactions Information Theory, (739):173-1706 2022.
[11] H. Chen, H. Lee, and J. Lu. Improved analysis of score-based generative modeling: User- friendly bounds under minimal smoothness assumptions. In Proceedings of the International Conference Machine Learning, pages 4763, 2023.
[12] Chen, Z. Shao, X. Zheng, K. Zhang, and J. Yin. Integrating aesthetics and efficiency: Al- driven diffusion models for visually pleasing interior design generation. Scientific Reports, 14(1):3496, 2024.
[13] Chen, R. Zhang, Y. Zhou, and C. Chen. Towards aligned layout generation via diffusion model with aesthetic constraints. arXiv preprint arXiv:2402.0475: 04444444444 2024.
[14] M. Chen, S. Mci, J. Fan, and M. Wang. An overview of diffusion models: Applications, guided generation, statistical rates and optimization. arXiv preprint arXiv:2404.07771, 2024.
[15) S. Chen, S. Chewi, J. Li, Y. Li, A. Salim, and A. R. Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In Proceedings of the International Conference on Learning Representations, 2023.
[16) K. Choi, A. Grover, T. Singh, R. Shu, and S. Ermon. Fair generative modeling via weak supervision. In Proceedings of the International Conference on Muchine Learning, pages 1887-1898, 2020.
11

---

[17] J. K. Christopher. S. Back, and F. Fioretto. Projected generative diffusion models for constraint satisfaction. arXiv preprint arXiv:2402.03559, 2024.
[18] K. Clark, P. Vicol, K. Swersky, and D. J. Fleet. Directly fine-tuning diffusion models on differ- entiable rewards. In Proceedings of the International Conference on Learning Representations, 2024.
[19] F.AA. Croitoru, V. Hondru, R. Ionescu, and M. Shah. Diffusion models in vision: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
[20] P. Dhariwal and A. Nichol. Diffusion models beat gans on image synthesis. In Proceedings of the Advances Neural Information Processing Systems, volume 34, pages 8780-8794, 2021.
[21] Y. Du, Durkan, R. Strudel, J. B. Tenenbaum, Dieleman, R. Fergus, J. Sohl-Dickstein, A. Doucet, and W. S. Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and MCMC. In Proceedings of the International conference on machine learning, pages 8489-85 10, 2023.
[22] J. Elenter, L. F. Chamon, and A. Ribeiro. Near-optimal solutions of constrained learning problems. In Proceedings of the International Conference on Learning Representations, 2024.
[23] Y. Fan and K. Lee. Optimizing DDPM sampling with shorteut fine-tuning. In Proceedings of the International Conference on Machine Learning, pages 9623-9639, 2023.
[24] Y. Fan, O. Watkins, Y. Du, H. Liu, M. Ryu, C. Boutilier, P. Abbeel, M. Ghavamzadeh, K. Lee, and K. Lee. Reinforcement learning for fine-tuning text-to-image diffusion models. In Proceed- ings of the Advances in Neural Information Processing Systems, volume 36, 2023.
[25] B. T. Feng, R. Baptista, and K. L. Bouman. Neural approximate mirror maps for constrained diffusion models. arXiv preprint arXiv:2406. 12816, 2024.
[26] N. Fishman, L. Klamer, V. De Bortoli, Mathieu, and M. Hutchinson. Diffusion models for constrained domains. Transactions on Machine Learning Research, 2024.
[27] N. Fishman, I.. Klarner, E. Mathieu, M. Hutchinson, and V. De Bortoli. Metropolis sampling for constrained diffusion models. In Proceedings of the Advances in Neural Information Processing Systems, volume 36, 2024.
[28] F. Friedrich, M. Brack, L. Struppek, D. Hintersdorf, P. Schramowski, S. Luccioni, and K. Ker- sting. Fair diffusion: Instructing text-to-image generation models on fairness. arXiv preprint arXiv:2302. 10893, 2023.
[29] G. Giannone, A. Srivastava, O. Winther, and F. Ahmed, Aligning optimization trajectories with diffusion models for constrained design generation. In Proceedings of the Advances Neural Information Processing Systems, volume 36, 2023.
[30] S. Gugger, L. Debut, T. Wolf, P. Schmid, Z. Mueller, S. Mangrulkar, M. Sun, and B. Bossan. Accelerate: Training and inference at scale made simple, efficient and adaptable. https: wtwtthu..mmmhuuee http///ffaaaaaaaeeeaaeeeeaeaeeee
acc.eereeeeeeeeeeee 2022.
[31] Y. Han, M. Razaviyayn, and R. Xu. Neural network-based score estimation in diffusion models: Optimization and generalization. arXiv preprint arXiv:2401.15604 2024.
[32] Y. Hao, Z. Chi, L. Dong, and Wei. Optimizing prompts for text-to-image generation. In Proceedings of the Advances in Neural Information Processing Systems, volume 36, 2023.
[33] M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. In Proceedings of the Advances in Neural Information Processing Systems, volume 30, 2017.
[34] J. Ho, A. Jain, and P. Abbeel. Denoising diffusson probabilistic models. In Proceedings of the Advances in Neural Information Processing Systems, volume 33, pages 6840-6851, 2020.
[35] Ho and Salimans. Classifier-free diffusion guidance. In NewrIPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.
12

---

[36] I. Hounie, A. Ribeiro, and L. F. Chamon. Resilient constrained learning. In Proceedings of the Advances Neural Information Processing Systems, volume 36, 2023.
[37] C.-W. Huang, M. Aghajohari, J. Bose, P. Panangaden, and A. C. Courville. Riemannian diffusion models. In Proceedings of the Advances in Neural Information Processing Systems, volume 35, pages 2750-2761, 2022.
[38] D. Z. Huang. J. Huang, and Z. Lin. Convergence analysis of probability flow ODE for score- based generative models. arXiv preprint arXiv:24044009730, 2024.
[39] L. Huang, T. Xu, Y. Yu, P. Zhao, K.-C. Wong, and H. Zhang. A dual diffusion model enables 3D binding bioactive molecule generation and lead optimization given target pockets. bioRxiv, pages 2023-01, 2023.
[40] A. Kazerouni, E. K. Aghdam, M. Heidari, Azad, M. Fayyaz, I. Hacihaliloglu, and D. Merhof. Diffusion models in medical imaging: A comprehensive survey. Medical Image Analysis, page 102846, 2023.
[41] Y. Kim, B. Na, M. Park, Jang, D. Kim, W. Kang, and I.-C. Moon. Training unbiased diffusion models from biased dataset. In Proceedings of the International Conference Learning Representations, 2024.
[42] D. Kingma and R. Gao. Understanding diffusion objectives the ELBO with simple data augmentation. Proceedings of the Advances in Neural Information Processing Systems, volume 36, 2023.
[43] Z. Kong, W. Ping. J. Huang, Zhao, and Catanzaro. DiffWave: A versatile diffusion model for audio synthesis. In Proceedings of the International Conference Learning Representations, 2020.
[44] Y. LeCun, C. Cortes, and C. J. Burges. The MNIST database of handwritten digits. http: //yann. lecunn htn//xdbb/mnin
[45] K. Lee, H. Liu, M. Ryu, O. Watkins, Y. Du, C. Boutilier, P. Abbeel, M. Ghavamzadeh, and S. S. Gu. Aligning text-to-image models using human feedback. arXiv preprint 12111 2023.
[46] G. Li, Y. Huang, T. Etimov, Y. Wei, Y. Chi, and Y. Chen, Accelerating convergence of score- based diffusion models, provably. arXiv preprint arXiv:2403.03852, 2024.
[47] G. Li, Y. Wei, Y. Chen, and Y. Chi. Towards faster non-asymptotic convergence for diffusion- based generative models. In Proceedings of the International Conference on Learning Repre- sentations, 2024.
[48] G.-H. Liu, T. Chen, E. Theodorou, and M. Tao. Mirror diffusion models for constrained and watermarked generation. In Proceedings of the Advances in Neural Information Processing Systems, volume 36, 2023.
[49] N. Liu, S. Li, Y. Du, A. Torralba, and J. B. Tenenbaum. Compositional visual generation with composable diffusion models. In Proceedings of the European Conference Computer Vision, pages 423 439,449,220 2022.
[50] X. Liu, X. Tong, and Q. Liu. Sampling with trusthworthy constraints: A variational gradient framework. In Proceedings of the Advances in Neural Information Processing Systems, pages 23557-23568, 2021.
[51] Z. Liu, P. Luo, X. Wang, and X. Tang. Large-scale celebfaces attributes (celeba) dataset. https://mmlett ie. cuhk. edu. hk/ pttppooss/eelllelllllppp html.
[52] I. Loshchilov and F. Hutter. Decoupled weight decay regularization, 2019.
[53] A. Lou and S. Ermon. Reflected diffusion models. In Proceedings of the International Conference Machine Learning, pages 22675-22701, 2023,
13

---

[54] A. S. Luccioni, C. Akiki, M. Mitchell, and Y. Jemiite. Stable bias: Analyzing societal represen- tations in diffusion models, arXiv preprint arX1v:2303.11408, 2023.
[55] C. Luo. Understanding diffusion models: A unified perspective. arXiv preprint arXiv:2208. 11700, 2022.
[56] S. Mei and Y. Wu. Deep networks as denoising algorithms: Sample-efficient learning of diffusion models in high-dimensional graphical models. arXiv preprint arXiv:2309. 1420, 2023.
[57] R. Naik and B. Nushi. Social biases through the text-to-image generation lens. Proceedings of the AAAT/ACM Conference on AI, Ethics, and Society, pages 786-808, 2023.
[58] G. Parmar, R. Zhang, and J.-Y. Zhu. On aliased resizing and surprising subtleties in gan evaluation. In Proceedings of the IEEF/CVF Conference on Computer Vision and Puttern Recognition, pages 11410-11420, 2022.
[59] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Yang. Z. VeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. PyTorch: An imperative style, high-p geghhperffrmance deep learning library, 2019.
[60] T. Power, R. Soltani-Zarrin, S. Iba, and D. Berenson. Sampling constrained trajectories using composable diffusion models. In Proceedings of the IROS 2023 Workshop on Differentiable Probabilistic Robotics: Emerging Perspectives on Robot Learning, 2023.
[61] D. J. Rezende and F. Viola. Generalized ELBO with constrained optimization, GECO. In Workshop Bayesian Deep Learning, NeurlPS, 2018.
[62] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 10695, 2022.
[63] O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein, A. C. Berg, and L. Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJC)), 2111:222,- 2015.
[64] C. Saharia, W. Chan, S. Saxena, L. Li, J. Whang, E. L. Denton, K. Ghasemipour, R. Gon- tijo Lopes, Karagol Ayan, T. Salimans, al. Photorealistic text-to-image diffusion models with deep language understanding. In Proceedings of the Advances in Neural Information Processing Systems, volume 35, pages (6479-36494, 2022.
[65] X. Shen, C. Du, T. Pang, M. Lin, Y. Wong, and M. Kankanhalli. Finetuning text-to-image diffusion models for fairness. In Proceedings of the International Conference on Learning Representations, 2024.
[66] Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Proceedings of the International Conference on Machine Learning, pages 2256-2265. 2015.
[67] V. Solo and X. Kong. Adaptive signal processing algorithms: stability and performance. Prentice- Hall, Inc., 1994.
[68] J. Song, C. Meng, and S. Ermon. Denoising diffusion implicit models. In Proceedings of the International Conference on Learning Representations, 2021.
[69] C. Szegedy, W. Liu, Jia, P. Sermanet, S. Reed, D. Anguelov, D. Erhan, V. Vanhoucke, and A. Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1-9, 2015.
[70] W. Tang. Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond. arXiv preprint 27310:2403.0627, 2024.
14

---

[71] M. Uehara, Y. Zhao, K. Black, E. Hajiramezanali, G. Scalia, N. L. Diamant, A. M. Tseng, Biancalani, and S. Levine. Fine-tuning of continuous- timm diffusion models as entropy- regularized control. arXiv preprint arXiv:2402.1 1tt.:2222...19444 15/94, 2024.
[72] M. Uehara, Y. Zhao, E. Hajiramezanali, G. Scalia, G. Eraslan, A. Lal, S. Levine, and T. Bian- calani. Bridging model- based optimization and generative modeling via conservative fine-tuning of diffusion models. arXiv preprint arXiv:2405.1966 2024.
[73] P. von Platen, S. Patil, A. Lozhkov, P. Cuenca, N. Lambert, K. Rasul, M. Davaadorj, D. Nair, S. Paul, W. Berman, Y. Xu, S. Liu, and T. Wolf. Diffusers: State-of-the-art diffusion models. https://8s tig...tt..... cww http///gggnggcc.ecc ffseess, 2022.
[74] Watson, Juergens, N. R. Bennett, B. L. Trippe, J. Yim, H. Eisenach, W. Ahern, A. J. Borst, R. J. Ragotte, L. F. Milles, et al. De novo design of protein structure and function with RFdiffusion. Nature, 620(7976):1089 1100, 2023.
[75] Weiss, E. Mayo Yanes, Chakraborty, L.. Cosmo, A. M. Bronstein, and R. Gershoni-Poranne. Guided diffusion for inverse molecular design. Nature Computational Science, 3(10):873-882 2023.
[76] X. Wu, K. Sun, F. Zhu, R. Zhao, and H. Li. Human preference score: Better aligning text- to-image models with human preference. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 2096-2105, 2023.
[77] J. Xu, Liu, Y. Wu, Y. Tong, Q. Li, M. Ding, J. Tang, and Y. Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. In Proceedings of the Advances in Neural Information Processing Systems, volume 36, 2023.
[78) L. Yang, Z. Zhang, Y. Song, S. Hong, R. Xu, Y. Zhao, W. Zhang, B. Cui, and M.-H. Yang. Diffusion models: A comprehensive survey methods and applications. ACM Comput. Surv., 56(4), nov 2023.
[79] L. Yang, Z. Zhang, Y. Song. S. Hong, R. Xu, Y. Zhao, W. Zhang, B. Cui, and M.-H. Yang. Diffusion models: A comprehensive survey of methods and applications. ACM Computing Surveys, 56(4):1-39, 2023.
[80] H. Yuan, K. Huang, C. Ni, M. Chen, and M. Wang. Reward-directed conditional diffusion: Provable distribution estimation and reward improvement. In Proceedings of the Advances in Neural Information Processing Systems, volume 36, 2023.
[81] C. Zhang, C. Zhang, Zheng, M. Zhang, M. Qamar, S.-H. Bae, and I. Kweon. Audio diffusion model for speech synthesis: A survey on text to speech and speech enhancement in generative AI. arXiv preprint arXhv:2303.13336 2023.
[82] Z. Zhang, S. Zhang, Y. Zhan, Y. Luo, Y. Wen, and D. Tao. Confronting reward overoptimization for diffusion models: A perspective of inductive and primacy biases. In Proceedings of the International Conference on Machine Learning, 2024.
15

---

Supplementary Materials for "Constrained Diffusion Models via Dual Training"
A Details on ELBO
Recall the evidence lower bound (ELBO).
p(
g) Bute mo) log q wr))
we can utilize conditionals to expand it into
E(p: q) = Bg((o) By(zi (ap) p(zo [1)]
E Bg(w)) [Dx (g(xT)] xo) [1) [[(()))
reconstruction likelihood
final latent mismatch
By(zn) By(e (2)
DRi. (g(xx-1 xr. 2o) || p(r--1
xx))
denoising matching term
where the first term the reconstruction likelihood of the original data given the first latent 21, the second term is the mismatch between the final latent distribution and the Guassian prior, and the last summation measures the mismatch between the denoising transitions from forward/backward processes. With the variance schedule described in Section 2. it is known that the reconstruction likelihood and final latent mismatch are negligible, and thus the approximation (4) is almost exact, which is our focal setting of this paper.
We next focus on one summand of the denoising matching term.
( Butn mo) DKL (g(21-1 vu) p(1:p(1111)) [+))] [2:))
Application of the reparametrization trick leads to 2t = Vatdi11 + /1 - 0 where et- ~ (O0 /) is white noise sample. Using Bayes rule, we can express ( 2e, Tu) as Guassian distribution
where Me (xx))
1-a1 V log g(x,) is the mean and (t) (1-a))(1-a-1) is the variance. Vaa
Var"
1-d.
To stay close to the ground-trute backward conditional q(x;-1 Its xo) as much as possible, take p(x,-1 x.) to be the same as q(x;-1 I2, zo) except replacing V log p(x,) by 3(2,, t) and (t) by of(().
N N(m.i:((()))
where ((xx))
vor 04 + 1-ar S(1,, t). Thus,
foi
DKL (q(Ir-1 20) 1l p( (x))
og (t)
=
dlog (t
+ d (t
|| (2,, 20) - fi(z,, 2000
(t
(t())
(t)
2 doog (t
+ d (t))
(1 - -a))
|8( 1)
V log ((xx))²
=
20(1)
(1
(1) ar
variance mismatch
prediction loss
where the second equality is due to the KL Divergence between Gaussians. Since (t) and 01) are constants, the variance mismatch term is irrelevant to optimization. Denote
(1-a,)² and T == Ei= wi. We can define a discrete distribution over the set (2...., T) as Wi == 203(t)a;
16

---

()) (1) == - Also denote U == 2i=22 dlog
d d
Hence, the ELBO maximiza-
(f())
tion: maximizez Elp; q). is equivalent to the quadratic loss minimization,
minimize v + @ Era, t, [ - V log {2
where t, a, an expectation over the data distribution (xx0), the discrete distribution Pa (1) from 2 to I', and forward process ((', zo) given the data sample Zo. Since shifting objective function by constant and scaling an objective function by a constant don't change the solution of an optimization problem, omit constants and w for brevity, and only emphasize them whenever it necessary. Hence, the ELBO maximization equals the quadratic loss minimization,
minimize En [[((zr,)) () log ((x+)]]*
up some scaling and shifting constants, where Exo,, is expectation over the data distribution ((ro), the discrete distribution (t) from 2 to T, and a forward process g(zz (o)) given the data sample zo- In practice, however, have to parametrize the estimator ((x,, t) as Se S0( () with parameter 0 E O
minimize é
E, (0((,, - V log {(x+))]
which is our focal objective of generative modeling. A parametrized representation of p(æ;-1 xx,)) associated with So (ate, 1) is denoted by (= | and the backward process has parametrized joint distribution po (xoor)) We remark that the above prediction problem can be reformulated as data noise predictions [55], with results directly transferrable these formulations.
B
Proofs
We provide proofs of all lemmas and theorems in the main paper.
B.1 Proof of Lemma
Proof. The ELBO maximization has the same optimal solution with the KL divergence minimization because of the equality (3). This directly proves the second equivalence. Next, we relate these two problems to the likelihood maximization problem.
We note that the KL divergence is non-negative and is zero if and only if two distributions the same. Since q(FoTT) E P for large T' the solution of the ELBO maximization and KL divergence minimization is given by p* = q. For the KL divergence minimization, the optimal value is zero. For the optimal value of the ELBO maximization, from (3) it follows that:
Epp": q) = Bg(zo) log {( DKL
]| (æo.r) = Eg(()) log {.
(10)
It is clear that the likelihood maximization problem maximize, (g(()) [logp(æ)) to minimize,, ( (g(x)) p(ro)). Therefore, any distribution p" (xxoor) whose marginal satisfies
p (2o) is equivalent
p" (xn) q(20). will solution of the likelihood maximization problem. This includes the solution of the KL divergence minimization and ELBO maximization which is p* Therefore,
maximize E(p; 4) 3 maximize Ey(ss)) [log p(xo) PEp
(11)
PEP
which concludes the proof.
B.2 Proof of Lemma 2
Proof. It straightforward check the zero duality gap in convex optimization; see e.g., [4, Proposition 5.3.2]. Furthermore, for convex optimization problem, an optimal dual variable X* that maximizes the dual function geometric multiplier. Hence, (p*. X*) is an optimal primal-dual pair of the convex optimization problem.
D
17

---

B.3 Proof of Theorem 1
Proof. From the strong duality in Lemma 2, it is known from [4, Proposition 5.1.4] that is also geometric multiplier. Thus, Problem (U-KL) reduces to an unconstrained problem,
minimize PEP
C(p,
(12)
where the objective function results from plugging optimal dual variable X into Lagrangian [p, A).
By the definition of Lagrangian,
A(v,)) DKa(I((FO:)) p( + Σ A (Dei. (g'())))) p((oor)) - bs)
m Σ
E((:g)
A E(p:
+ g(0)) [log a( (<0))] + Σ A (Bo((r) [log [ogg((x))) - bi)
By taking À Xt, Problem (12) is equivalent to
maximize PEP
E(p 4)
Σ
x E(v:g))
(13)
From the definition of ELBO, know that
m
122
By (m) log p(Zo:))
E(p:q) + Σ XE(p:g'
(gg(o))
12 (1o)
q xr)
where we use the fact that the forward processes have the same marginal distribution given any initial data samples. Normalization of initial data distributions leads to Imix Thus, Problem (13) is equivalent
(X*).
maximize E(p: (A7)
PEP
which, together with Lemma 1, completes the proof.
B.4 Proof of Theorem 2 and Feasibility Criterion
start with the proof of Theorem 2.
Proof. Similar the proof of Theorem 1, we begin with the Lagrangian of Problem (U-ELBO),
A(p,)) - F((55)
Xe (F((sg") 8.)
_1
Σ i=1
plosdt)
AT
(14)
X
X
E (zo) Ealaur (m)) log p(Zo:T)
q(x1:1 220) ) - X
X
(xo)
XT1
p(Zor)) ATb
XT1
-
EEg (wo) Egttner 22) log q( (o))
1A AT6
18

---

where from (14) onwards use notation: Ao 1, = 0, x=
Am_= " b= [60 ..... 6m
and use d2 to represent 9. which will be used in the rest of proof. To formulate the dual problem, check the minimum of the Lagrangian,
g(A)
minimize C(p, A)
PEP
minimize PEP
- 1)pp)))( A
b + (1 minimize E(p: (1)))
(15)
PEP
where the only term that depends on is the ELBO. Recall that:
(x((T() ()(oo:))
-F(p; q) Bgg(s)) [log ((x0))
Since the minimum value of the KL divergence is zero (attained when = q). the minimum of -p(; q) is likewise attained when p - q. Thus, is straightforward that the minimum is equal to the entropy of the distribution g, denoted by h(g) == -E(ao) [log ((zo) ]- With this in mind, from (15) we have
g(A) 6 + (A"1) h(9
Thus, the dual problem reads
maximize A
9(A) = - - + (A"11 (()(a)).
We first reformulate the entropy of the mixture distribution g(a),
h(g(2)) = (x))
[Log ())
/
Σ
m
96 (zo) log
Σ
m
A
q' ('(20) dro
(16)
X=1
XT1
A
A (x11 r² (20))
Σ 0
o (xo) log
dro
(17)
X11
/ d' (20) log (9 (2mo) dro -
m
A
A
log
x
XT1
ATI
he
m
m
A
log
X
X
A X h4
Σ
Σ
log(As) log(A?1)
X'1
:=0
where going from (16) (17) utilize the assumption that the distributions (m have disjoint supports; see Remark on when this is the case.
Now, we can compute the gradient of the dual function over Ao. = I,..., m,
2 (-xTb+ + (1 (())))) = Ox,
0
172
m Σ
x
+
Xyh1
\3 log A, + 1) log(A"1)
Oi
==0
j ==0 0
A
h4
b
log
XT1
Setting the gradient be zeros allows us to find the optimal dual variables X*
h4 b, - log ((+)=I)
= for i 1,... m.
19

---

Hence,
Remark 1. We remark on the assumption of the distributions {q J==u having disjoint supports. In the setting of adapting model to new data in Section 3.2, this is reasonable ussumption. Since we often finetune pre-truined diffussio model on new data not seen in the original pre-truining dataset, the new data distribution and the pre-training data distribution have mostly disjoint supports. In the minority class setting in Section 3.2. the constrained distributions {q' Ji=1 and the objective distribution dº usually are not disjoint. However, since the distributions {q" 1m often are often e 6 for i 1,....
x;
(18)
(X+)²1
We clarify that in (18), X* = PAi.... Am
T with its first element being A = 1. Finally, if we return
back to notation X [A].... Xm 1"
then.
X 1 (A-)T1
chi for 1,..., m
which completes the proof.
restrictions of d subsets of the support of q", i.e., the minority classes, the derivation of optimal dual variables is similar to what we have provided in this section, so we omit the repeated details. Extending these results
where the distributions are neither disjoint nor restrictions of the objective distributions, is challenging and has been left to future work.
To prove feasibility criterion, first show that the dual function is finite in Lemma 5.
Lemma 5 (Boundedness of the optimal dual function). Let the differential entropy h4 of each distribution q' be finite. Then, the optimal value of the dual function D* max xaxx20g(X) finite if and only if
Σ
hi by
Proof. (+) From En= e < 1, the otpimal dual variable X given by (18) is finite. Thus, the optimal value of the dual function g(X²) becomes
,(A²) = - 6 +
Σ X4 - Σ A log + 1=0
==0
and X ch > 0, and also {h+ are all finite. Therefore, D* finite.
(+) We prove by contradiction. Assume Σm= chi-h 0º 2 for some 8 2 0. For any À 2 there exists direction in which g(A) increases, i.e.,
0g
A
h4 - log
> 8 li.
(19)
aλ;
XT1
To see (19) by contradiction, we check that
A;
h4 b, llog
< 8 for i 1, 1,..., m
A
A;
- a4 - b, 8
<
for i = 1 m
XT1
Σm
-i=1 A
Σn
A
e by
- bi
<
1 -+=
1+> di
h - bc
60
which contradicts our assumption that >"" k b. ed By contradiction, we have (19). Further- more, (19) implies that (()) is unbounded above, which contradicts the finiteness of D* Therefore, must have that 2m ei b. < 1.
20

---

+ 7 0. Lemma 6 (Feasibility criterion). Let the differential entropy h2 of each distribution q' be finite. Suppose that there exists a feasible solution to Problem (U-ELBO) such that its objective function bounded from below. Then, Problem (U-ELBO) is feasible if and only if
h -
Proof. (=) Since the primal problem is feasible, the optimal objective function F* bounded from below and it is attained at a feasible point. For the sake of contradiction, we assume Ei e b, 2 which is equivalent D* oo according to Lemma 5. However, this violates weak duality, i.e., D < F*. By contradiction, we must have Ex=1 eh <
(+) We consider a set A.
A : [(u),..., Um t) E(p, - S u, for I,..., m and E(p.4") S por P}.
The set A convex since is the intersection of m + epigraphs of convex functions. We also introduce another convex set B,
B = {(0....,0,4 0,4) [
We utilise proof by contradiction. Assume that the primal problem is infeasible. Then there doesn't t exist any p E P such that - Flp, q¹) b, 0 for all i 1....; m. Hence, A and B disjoint convex sets. From the separating hyperplane theorem, there exists hyperplane that separates them, i... 3u g=+1 and IR
x'v 2 for all x E yf v S for all y B.
(20)
(21)
Am, r Then, (21) reduces to Let = [A1,
-T v X" + 76 71 S for all (u, E B
This is because yt S c for any t E R. Note that 0 means the separating hyperplane is vertical, ie., being parallel the t-axis. Furthermore, from (20) we can write
IT
XTu + 7t = XT 2 c for all (u, t) E
A 0.
The above is true because the set of values that each u, can take in A is unbounded above. Thus, since Xuu u 2 c, necessarily every X has to be son-negative. Now, consider
g(A)
inf (u,
+ ATu
g(a))
inf + axTu for R4
(u
lim g(a)) = lim inf t + axTu
0
a 00 (u.))A
lim a
inf (utt)
XT
lim ac
00
which shows that D* = do. This contradicts D* being finite (D* is finite due to Em= hh- <1 and Lemma 5). Because of the contradiction, the primal problem has to be feasible.
B.5 Proof of Lemma 3
Proof. The proof is an application the convergence theory of DDPM [47, Theorem 3]. We next
check all assumptions of [47, Theorem 3]. is easy that cast 9 target distribution
(2)
21

---

of diffusion model. By the definition,
[(0,)) DKL. (((#T.:) pa( +
Σ
Ne (DRL. (T'( pl (2(0:)))) bi)
==1
m
-)(:::) (lgz))loogg(
(-E(po; bi). Σ >
Thus, the partial minimization of E(0, X) over is equivalent to a weighted EBLO minimization,
minimize - E(po: q) 6
Σ A E(po:
or, equivalently,
minimize 8EB
Eppo (2))
(22)
where we normalize the weighted ELBO objective by introducing mixed data distribution Tmix We note that 5" (A) is also minimizer of Problem (22).
On the other hand, using Problem (P-LOSS), we can rewrite Problem (22) as
minimize E
E (A) (oo), f. [[s)(x),)) t) V log (xx))]
which is equivalent to the score matching objective in DDPM. Therefore, the score matching as- sumption [47, Assumption 1] satisfied with the error bound & from Assumption Viewing Assumption 2, and using appropriate stepsize and variance, all assumptions in [47, Theorem 3] are satisfied. Therefore, application [47, Theorem 3] completes the proof.
B.6 Proof of Lemma
()) (x)
Lemma 7, The TV distance between two mixture data distributions quix Imix is bounded by
TV
mmx
9mix
I 1
Proof. By the definition,
TV (2)
(A')
1
+ Eis Ntg
+ Em Xing
Jmux Amnix
2
30
1+AT1
(++)
1
(1 (At) 1)(g E. N'9) -1 E (
2 300
1
Σ X + (1+) - -
T
2
(1+AT1)(1 + (1+
(X x5*1a* + (X* (x)
2 20
1
a
111
where the first inequality is due to (1 + X the second incquality.
2 1, and we use triangle inequality in
We recall the Lagrangians for Problems (U-LOSS) and (P-LOSS),
A3,
Eg(zo))
|| 1) V log (ox))()))))
m
X (ggt(s [|3( 1) log 1ogg(1))1 -B²
==1
22

---

(30, = Lo(s8,
and their associated dual functions,
(A) minimize (8, and (s(A) minimize (z, 2). BE8
REB
For brevity, we use shorthand B, and E. for Bogwo) s, and (mo), t, respectively.
Lemma (Parametrization gap). Let Assumption 4 hold. Then, 0 I 9. (A) go(A) S 4R(1+ }1 for any À O.
Proof. Let the partial minimizer C, A) over @ be (A) = argming L. (. A) for any À 2 0. For any À 2 0, there exists O E θ such that |1$ - S v for any x 2 0, according to
Assumption 4. Thus,
(3g. A) (A), = Ey ( - xoll ||] - E2 || (A) t) - xxol?]
A Be [| z01l -Ey [13 (A3 xool²)
1
4R B, [||3g (Tee t) -3 (A) O1]
+ 4R Σ
A E. [1)g)(0),)) - (A),,,,) |I]
==1 4Rv + |
where the first inequality is due to that the quadratic function is locally Lipschitz continuous with parameter 4R, and the second inequality is because that there exists 0 E O such that (A) Tاll L S for any À 2 O, according to Assumption 4.
By the definition (7 E argminge e (, (A))))), & Ao
Therefore,
0 < 2.(()))))) S -. which gives our desired result by the definition of dual functions.
< 4R(1 + (1l1))
Lemma 9 (Differentiability). The dual function 9, (A) is differentiable with gradient (5 (A), A).
Proof. For any À 2 0, the Lagrangian L, (3, A) is strongly convex in function 3 @ S. Since S is convex and compact, any partial minimizer 3 (A) is unique. By Danskin's 's theorem [4], gs (A) is differentiable and its gradient is the gradient of (8, X) over À at 3* (X)-
Lemma (Convexity). The dual function ga(A) is p-strongly concave in À E H. where
(1 1+ (11~11 [x+
a
Proof. For any \1. A2 E H. we denote 8i = (A1) and 3) == (A2),, which are
unique partial minimizers of the Lagrangians (,(8,\1) and L. A2). \²) Denote (o(8) == E.y | t) Vlog
(() == E. [|s(,,, log (x:))) for 1,..., m, and ((3) = (3). e.. (3) TT. By the convexity f..
(87) > ((3)) 2 2V (V 3)
4(3) > 4(3) + 2 (V 3) -
we multiply the first inequality above by 22, 2 0 and the second inequality above by A 2 0. and add them up from both sides, then,
-e((2) ((51), AD 2 v X Vo((s,) -
23

---

We apply Lemma 9 to the LHS of the inequality above,
(V) Vg.(>i))* (A2 - -1)) 2
-VV -3). (23)
On the other hand, by the optimality of 3 and 32.
Volo(3)) - L V 0
(24a)
Volo(3)) M (z) =
(24b)
which allows us to simplify the right-hand side of (23) and obtain
(Vg.(te)-Vg(A)))) (22- - A1) 2 (V Vslo(31), 81)
(25)
W 2 2 - 121L2
where the last inequality results from the strong convexity of quadratic functionals. By the smoothness of quadratic functionals with parameter 1,
| $2 I L2
2 [V Valo(33)
= |A V2(()) 13 (>((())))
= 1(A2 - V.((()) - -
N 1(2--A1) V [V
-
v L2
where the equality is due to the optimality condition (24) and the last inequality is due to triangle inequality. By Assumption 5,
11O22 x1) Vz((3)) i. 2 l Aull-
We also notice that
m
A [z(((51) Vsli l6)
L2
VI
A1,
8g و
where the first inequality is due to triangle inequality and the second inequality is due to the smooth- ness of quadratic functionals. Hence,
|o lll.. 1122- -h1l-111-
||il | Ei
or, equivalently,
|| 321l22 2 1 |
1A - A11l
Therefore, (25) becomes
(T+iAIT))
)
2
- -VVg.)A2) - Vai( (A2 (A2-A1) x1) >
1122 2112
which completes the proof by choosing the smallest modulus over λ1 H.
Proof. By Lemmas 9 and 10, for any X E H.
(A) I 9^(A²) + (a) -x)- # I| -11122
Thus, if we choose == A², then
m
g.(A²) s gx(X*)+ + Σ -x))))) (Egi [1=+())(())) V llog (ogg - - # -1 - 1-1
1=1
24

---

Optimality of (8* (A*), leads to the complementary slackness,
x Rey
(A=)( - log (ogg -b') -. = 0
and the feasibility,
Egj
(x2),,) t) V Vlog Vloga(r)))() s B¹.
Therefore,
9x(x*) s 9x(2²)
1 2 |ス~ x11
According to Lemma 8, (.*) 4R (1 - --11~11 11*+11)) v S gg((x*) Hence,
gx(x* - 4R (1+ + | v - go(2²) N/2 - - x=12².
Thus,
2 | x+12 < (x.()*) 2x(())) + 8 RR(1 + 11*1))))
11
,
8
R(1 ||^ 11))
#
where the last inequality is due to that g, (A) (A) for any A 0. and the optimality of X*
(x()') gx(X²) s (x()..
B.7 Proof of Theorem 3
Proof. By the triangle inequality for TV distance,
TV (mix (()))
1 TV (miss (3-)
+ TV Omix
p¹ (x*))
d log T
5 ||^ 1-111
Vd (log" T) Escore
VT
dl logs T
M mR (1+ 115111))
+ Va (log² T) Excore
vT
where the second inequality is due to Lemma and Lemma 3, and the last inequality is due to ||a|] S Vm ||a|| and Lemma 4.
B.8 Proof of Theorem 4
Lemma 11. For a stochastic variant of Algorithm 1 in Section 4.3, we have
2 1) {{m)] s ||A(h)) 1-1 + 52 - -2n (D+ - g(((h)) =
Proof. For brevity, we let the stochastic gradient be f(1) = [F(()),.... ... with
f.(h)) = Eno ~ gi. t, [11o(()((zn))) vo(x:))?] - B..
By the definition of X(h + 1),
|| [[(AA]+ + +nF(H)] - X 1
|A(h 1) x-112
S |A(A)) x- + + 2
= ||A(K) - x-112 + 1,
| 2 + 2n)( (A(h) Xxx)
25

---

where the inequality is due to the non expanssiiiesssss of projection. Application of the conditional expectation over both sides of the inequality above yields.
E
[[1 + xx [200]
||^(M)) x=11 - = ||/cco| 1x(1)) (A(h) - x²)
2n E [(()] ((())
which gives our desired result when we use the fact that (lf(h) X(h) is an approximate descent direction of the dual function g.
IB (((m)1(())) (A(h) xx) - E uppprxx S o(\(h)) -(X)).
Lemma 12. In the stochastic variant of Algorithm I in Section 4.3, the maximum prarametrized dual function in history up to step h satisfies
11822 2
lim h-30
Trest (h) 2
D
+
aappeox
Proof. The proof is based on the supermartingale convergence theorem [67, Theorem E7.4]. We introduce two processes,
a(h) =11A(h) x-1121 (D- [hest(h)
1752
+ e2
2
Supprox
?7522
8(h) = (2n (D* g(((h)) - 2 <mpnnx) - -7² 1 D* -
+ Fapprxx
2
where (h) measures the gap between X(h) and X when the optimality gap D Tres (h) is below the threshold, and (h) measures the gap bewteen D* and @(A(h)) (up to some optimization errors) when the optimality gap D* Trea below the threshold. Clearly, a(h) is non-negative. Also, ((h) is non-negative due to
17.5
115
D Trest (h)
S D g(A(h))
2 aapprox
aapprox
Let Fr be the o-algebra generated by sequences: a(h'), Bhh',, and A(h') for H < h. Thus, [FAJa2 natural filtration. We notice that a(h + 1) and 8(h + 1) are determined by X(h) in each step.
Hence,
E [a + 11) = [[a(h + A(h))
E|a(h 1) 1)) = 0] Pr (r(hh) = 0) +B[a(h+ 11) (a)r(((()))))) 0).
We first show that
E [[ + 1) < a(h) B(h).
(26)
A simple case is when a(h) = 0.
E [aa(h + 1) Fr
E [a + 1) A(h), a(h) = 0].
There two situations for ae (h) = O. First. if D Joes (h) s ns² + + = then a(h) ((h) 0.
Due to Toeat (h + 1) 2 (h), we have ((h) 0 and D Toes +
n
+ 2 Eappoox Thus,
(h 1) O, and (26) holds. Second, if X(h) X but D Toest (h)
2
e approx then
D. g()())). Hence, 3(h) O, which contradicts the non-negativeness of B(h). Therefore,
D Aosat( (h) 2 + Eopprox has to hold, which is the first situation.
nsa
26

---

We next show (26) when a(h) > 0.
E [a(h + 1) = E[a(h + 1) ((h), ah) 00]
1182
x(h), > 0)
= [|(()) x=11²1 D
-Sber(h) DD-
22
Eapprox
= E [|a0h) - xx11 M()..(0 > =0]
= ||a(k) =x + 12
2n (D' -G(((h) - =upprox)
S (h) B(h)
where the last inequality is due to Lemma 11, and the last equality is due to D Ibea(h) > nsa
2 + Epprox Therefore, (26) holds.
Finally, application of the supermartingale convergence theorem [67, Theorem E7.4] to the processes (hh) and 8(h) for h 2 1 concludes that B(t) is almost surely summable,
lim inf ((h) 0 almost surcly.
This means that either
lim inf 2n (ñ' - (((h)) g(\h)))\h))
2 ipprox 52
or D* Ibeat (h) S n 2 + t approxx which concludes our desired result.
Denote the step when (he achieves D' D'-(01 -(15" + @)pprox) by Area and the associated dual variable be Abest A(hoess).
Lemma 13. For stochastic variant of Algorithm I in Section 4.3, we have
| 101
2
+ 6 approx
4R(1 + "Sbeall))u
1 2
Proof. We denote the segment between Apea and X by B. By Lemma 10, the dual function g()) is strongly concave on B with parameter μ. Thus.
|| / x-11
2
(g(2*) - 9 (Abea))
2
(g(A") ((beas) + +4R(11 + || (aaeull)))
2 1)52
+ 52 Eapprox + 4R(1 + || )0
11
where the second inequality is due to Lemma We note that g ((oes) = g(A(hoess) and D g(Noes), and the third inequality is due to Lemma 12.
Proof. By the triangle incquality for TV distance,
TV (9mix Px (Rnes))
< TV
amixs Imix
+ TV (Are))
9 F (Aress)
d² log T
S
||**
Abeal
|l
+
+ Vd (log² T) Excore
VT
2
S2
2
+ 2
approx
+ 4R (1 |nesa (1) v
+
d² Plog T
VT
+ Vd (log2 T)
Excure
where the second inequality is due to Lemma and Lemma and the last inequality is due to ||ll 5 /m ||||| and Lemma 13.
27

---

C Experimental details
We provide implementation details of our computational experiments in Section 5. The source code is available here.
We train our constrained diffusion models by replacing the exact primal mini- mization step in Algorithm with N steps of gradient descent with the Lagrangian as loss function. Without loss of generality, we take the noise prediction formulation of diffusion rather than the score-matching formulation used in our theory. Since these two formulations are equivalent, this has main results. Algorithm depicts practical implementation of Algorithm 1.
Algorithm details:
bearing
Algorithm 2 Practical Implementation of Algorithm
1: Input: total diffusion steps T, diffusion parameter ag total dual iterations H, number of primal descent steps per dual update N, dual step size Ta. primal step size Tpp initial model parameters (()).
2: Initialize: X(1) =
3: for h = 1, H do
4:
for n 1,... .N do
5:
01 =0(h- 1).
6:
Ox+1 = 8n Ve
Ero Ex
11e (2), -oll?]
+ Σ>
A Ezo ~ ti
|Ee, (1x, - coll cool)
==1
7:
0(h) = NN+1-
8:
end for
9:
Update the dual variable
Achh 11 (bc(h)) + na
(x)) - sol1"] - [,)], for all i.
10:
end for
In Algorithm 2, the unbiased estimate of the noise prediction loss evaluated via
Exo g.
||ea (2 call?
Σ (0010000 (()) -0²)
where ( pm is a randomly chosen batch of samples from 9. {t() 1/- are time steps randomly
sampled from the interval [2, 7). and Zy(a) is the noisy version of xo at time step z(0), Then, the noise is sampled from the standard Gaussian, and is derived from x, VoLaa + v1 Gr
(t)
We remark important implementation detail in the fine-tuning experiment. In the fine-tuning constraints, we have to evaluate the KL divergence
DKL (pe.. (IO:T) || (o(o)))) S 6;
(27)
where pe (XorT) is the joint distribution of the samples and latents generated by the backward process of the pre-trained model. The KI. divergence constraint in (27) further reduces to
DEL (po,.. (Forr) (xooT)) Ez,, ~ pu || (z., t) (x t)](
constant.
(28)
To estimate the expectation in (28), we need to sample latents 2, from the backward distribution Pe,. (Xorr). ). In practice, this is computationally inefficient, since it requires running inference each time wants to sample latent. This is why implement this with sampling 2; random Gaussian noise instead. This still ensures that the predictions the new model pe don't ' differ too much from the pre-trained distribution Pag.. while making sampling batches much faster.
Resilient constrained learning. The choice of the constraint thresholds {b.) ye has noticeable effect on the training of constrained diffusion models. To avoid an exhaustive hyperparameter tuning process, in the minority class experiment, we use the resilient constrained learning technique [36] to
hutp://gthaa..cmmsberriaaaaaaaae Diffusson Dual Training
28

---

Table 1: Parameters of U-net Model used as noise predictor
# Res-Net layers per U-Net block
2
Res-Net down/upsampling blocks
6
Output channels for U-Net blocks
(128, 128, 256, 256, 512, 512)
adjust the thresholds [b,) during training. essence, the resilient constrained learning adds constraint relaxation cost to the loss and relaxes the thresholds by updating them through gradient descent each time we update the dual variable. can further shown theoretically that an equivalent formulation of resilience is achieved by adding quadratic regularizer of the dual variable into the
loss objective and setting the constraint thresholds to be zero, i.e., by 0. This is the approach we used in our experiments since it has fewer hyperparameters. We note that the only difference between Algorithm 2 and Algorithm 3 is the additional term in the dual variable update step (line of Algorithm 3).
Algorithm 3 Resilient Constrained Diffusion Models via Dual Training
Input: total diffusion steps T', diffusion parameter ar total dual iterations H. number of primal descent steps per dual update N, dual step size 77d primal step size Tpp constraint step size πe. initial model parameters (0)), constraint relaxation cost ?.
2: Initialize: ((1) = O.
3: for h 1,... do
4:
for = -I.... NN do
5:
0, -(hh- 1).
O,+1 On Up Vo a
[ (st:1) -col?
Σ
X ~ 2,, [ (1e, ca coll*)
6:
7: ((h) ON+1-
8:
end for
9:
Update the dual variable
(xh++)) [a((t)) + na (Bzo ([eon))Z11))) -col]]] -b²(h) - 2Td((t))] for all i,
end for
Model architecture. We use a time-conditioned U-net model as is common in image diffusion tasks for all three datasets. The time conditioning is done by adding positional embedding of the time to the input image. The parameters of the model are summarized in Table The fifth downsampling block and the corresponding upsampling block are Res-Net blocks with spatial self-attention.
Hyperparameters. We summarize the important hyperparameters in our experiments in Table 2. In the unconstrained models that we train for comparison, we use the same hyperparameters as the constrained version, disregarding the parameters related the dual and relaxation updates. For models trained on Image-Net, when training the constrained models, we initialized to the parameters of the unconstrained model make training times shorter.
Hyperparameter sensitivity. We remark the sensitivity of the dual training algorithm to the number of dual iterations, primal/dual batch sizes, and primal/dual learning rates.
Number of dual iterations: In our implementation this shows up as the number of primal GD steps per dual update, N. Experimentally, we have observed that long as N is greater than 1, the results are not sensitive to this value. Additionally, the dual updates add a negligible computational overhead. Hence, updating the dual nearly as many times update model parameters doesn't 't reduce training efficiency.
Primal/dual batch sizes: We have included results of training a constrained model on an unbalanced subset of MNIST, using different primal/dual batch sizes (See Table 3.) The results suggest that for the minority class experiments, when the ratio between Primal and Dual batch sizes larger, the model performs better (lower FID and more evenly distributed
29

---

Table 2: Hyperparameter values used in the main experiments. MC denotes Minority Class ex- periments and FT denotes Fine-Tuning experiments. denotes that resilience was not used for the experiment.
MNIST MC
Celeb-A MC
Image-Net MC
MNIST FT
#training epochs (N H)
250
1000
2000
500
primal steps per dual step (N)
2
2
2
2
Primal batch size
128
256
128
256
Dual batch size
64
128
64
256
Primal learning rate
0.0001
0.0001
5e-5
5e-6
Dual learning rate
0.1
0.05
le5
Resilience Relaxation cost
0.09
0.005
0.025
main dataset size
31000
12500
2000
200
constraint dataset(s) size
5000
500
64
samples). This is in line with the heuristic we used in the included experiments in the paper where we chose the batch sizes such that the ratio of primal to dual batch size is close to size ratio of entire dataset constraint datasets (which much smaller). Howver for the fine-tuning task, the batch sizes did not seem to affect the final result as much.
Primal/dual learning rate: For the primal learning rate, we followed the best practice used to train standard diffusion models. For the dual learning rate 7, refer to Theorem 8 in the paper, showing a smaller error bound for smaller 77 while slowing convergence. practice, long 1, observed that the model converges to similar results reliably.
Efficiency of constrained diffusion: We note that the complexity of sampling from constrained diffusion model does not increase with the number of constraints, as our trained diffusion model functions like standard diffusion model generate samples. Importantly, remark that training our constrained diffusion model has comparable efficiency to training standard diffusion models detailed next.
The additional computational cost of our dual-based training (Algorithm 1) arises from: (i) updating the dual variables; (ii) updating the diffusion model in the primal update.
Cost of updating the dual variables: We note that our dual-based training has the same number of dual variables as the number of constraints. Thus, the cost for the dual update is linear in the number of constraints. To update each dual variable, ELBO loss over the batches sampled from each constrained dataset (already computed for the Lagrangian). Therefore, the cost of updating dual variables negligible.
directly the
Cost of updating the diffusion model in the primal update: We note that the primal update trains standard diffusion model based on the Lagrangian with updated dual variables. In experiments, this primal training often requires as few as 2-3 updates per dual update. Thus, when training our constrained model, we can train for the same number of epochs unconstrained model but update the dual variables after every few primal steps. As result, training our constrained diffusion model is almost as efficient as training standard unconstrained models.
The only concern we encountered regarding efficiency that batches need be sampled from every constrained dataset at each step to estimate the Lagrangian. This introduces small GPU memory overhead that increases with additional constraints. However, this somewhat mitigated by the fact that constrained datasets are often much smaller than the original dataset, allowing us to choose smaller batch size for the constrained datasets without degrading performance (see discussion on batch sizes in hyperparameter sensitivity section).
30

---

Table 3: Constrained model trained on MNIST with different Primal/Dual Batch sizes
(Primal batch size, Dual batch size) (64, 16)
(64, 64)
(128, 16)
(128, 64)
FID score
16.6
20.6
16.7
20.24
FID scores. As a quantitative means of evaluating our constrained diffusion models, we use the FID (Frechet Inception Distance) as metric to gauge the quality of the samples generated by diffusion models. The FID score was first introduced in [33] in form of
(m,, C), (m Cw)) ||m - m= + Tr(C+ -((CC,))
(29)
where and C' represent mean and variance, respectively, of the distribution of the features of the data samples which have been extracted by an inception model [69]. Similarly, mhe and C. represent the mean and variance of some reference distribution that we are computing the distance to. In our experiment, we compute the FID scores by generating 15000 samples from the diffusion model we are evaluating and comparing them to a balanced version of the original dataset. In the experiment with MNIST, this is the actual dataset. In the experiments with Celeb-A and Image-Net, since there an imbalance in the original dataset, we consider balanced subset of each with an equal number of samples from each class as reference. We use the clean-FID library [58] for standard computation of the FID scores.
We note that our FID scores are somewhat larger compared to typical baselines in the literature. This is expected as consequence of our experimental setup. We train both the unconstrained and constrained models, on biased subset of the dataset wherein some of the classes have significantly fewer samples than the rest. We then compute the FID scores for these models compared to the actual dataset itself which unbiased (i.e., every class has the same number samples). These FID scores approximate how close the learned distribution of the model trained on biased data, is to the underlying unbiased distribution.
This setup contrasts with existing results in the literature, where the FID is computed with respect to unbiased data, and the models are also trained on unbiased data. Therefore, it is expected that such models will achieve better FID) scores than constrained or unconstrained models trained with biased data. Our purpose in reporting the FIDs was not to compare them to existing results (as such comparison would be uninformative) but to demonstrate that, when trained on biased data, the constrained model achieves better FID scores than the unconstrained model.
Compute resources. We run all experiments on two NVIDIA RTX 3090-Ti GPUs in parallel. The amount of GPU memory used was 16 Gigabytes per GPU. For experiments with MNIST and Celeb-A datasets, training each model took between 2-3 hours. This increased to 7-8 hours for latent diffusion models trained for the Image-Net experiments.
Assets and libraries. We use the PyTorch [59] and Diffusers [73] Python librarics for train- ing constrained diffusion models, and Adam with decoupled weight decay [52] op- timizer. The accelerate library [30] is used for the parallelization of the training processes across multiple GPUs. For classifiers used in evaluating the generated samples of the mod- els, we use the following pretrained models accessible on the Huggingface model database: A Vision transformer-based classifier for MNIST digits with %99.5 validation accuracy (https: / / Thuggingf ace. co/ //ffrre he/ffrrleyle wawalyknigghttte.. ceatt-orrgusernaeeee www.naseevii........ bttp://viiittaemeiimmiaaaa A classifier for images of male/female faces with %98.6 validation accuracy (https: httpss:///ggiggaaccc
co/cledoxx22/ GenderNew_ 2_0222)). For classifying the image-net data, a zero-shot classifier based on a CLIP model (https: h/th///ugiigg mouggingface ace. co/ co/openna openai/ cllp-viittticccccccc saap-vit-baaa base-aaeeeeeeeeee patchh2)))) was used. The Autoen- coder for the latent diffusion model was the stable diffusion VAE with KL regularization found on co/st htts/staaiittttt
bectps:/hhugginn
fmmee).
31

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We summarize our constrained diffusion models in Section 3, optimality guar- antees of unparametrized constrained diffusion models in Section 3.1, optimality guarantees of parametrized constrained diffusion models and training algorithms in Section 4, and experimental results in Section
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We mention two potential limitations of this work several future directions in Section 6.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We provide all assumptions,, lemmas, and theorems in Sections 3 and and provide proof details in Appendix B.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We summarize our experimental results in Scction 5, and provide imple- mentation details of experiments in Appendix C, together with additional experimental results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: A link to the code for replicating our main experiments has been provided in Appendix C. The datasets open source machine learning datasets that accessible online.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide experimental details in Appendix C including the hyperparameters used for each experiment. Our training details can be found in the code in supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
32

---

While the histogram plots that showcase our main results do not include error bars, we do provide Frechet Distance metries (that are measure of statistical distance between sample distributions) which emphasize our results.
Justification:
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Answer: [Yes]
Justification: We include compute details in Appendix C.
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https:/ /ttppp::///uurps...mmm wttppppcccpppppccccccccccccc hegbbiic/thicssudddlleeesss
caGuideliness
Answer: [Yes]
Justification: We have reviewed the NeurIPS Code of Ethics and fully comply with it during the preparation of this paper.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
The potential impacts of the work are discussed in Section 1.
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: We release no models and the supplemental code for training image generation model, trains the model on MNIST and Celeb-A A datasets neither of which have potential for misuse.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The libraries and assets used have been noted and credited Appendix
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [NA]
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
33

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
34


Reviews:
[{"review_id": "qcMKc0yYyb", "reviewer": "Reviewer_EijC", "summary": "This paper aims to address the issue of generating biased data based on the training dataset for diffusion models. The authors introduce a constrained diffusion model by imposing the diffusion constraints based on desired diffusions that are informed by requirements/constraints. They propose a dual training algorithm to train the model, and characterize the convergence of the trained model. Two constrained generation tasks are explored: ensuring fairness to underrepresented classes and adapting a pretrained model to new data.", "questions": "- How efficient the model is as the number of constraints increases? Each time the model updates the duality variables, it needs to train a diffusion model, which is slow.\n- how the proposed method differs from existing constrained generative models, for example [1].", "limitations": "The paper can be improved by:\n- a discussion of the efficiency of the proposed constrained diffusion models should be included.\n- a discussion of how the proposed method differs from existing constrained generative models should be added.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is well written and easy to follow.\n- Its key idea is using the quadratic loss formulation of score matching, which is equivalent to the ELBO of diffusion models. The authors employ Lagrangian duality to update both the parameters of the diffusion model and dual variables of constraints. While Lagrangian duality is commonly used for constrained optimization, particularly for constrained generative models, see [1,2,3], this paper first introduces this approach for diffusion models.\n- Thanks to the quadratic loss minimization of diffusion models, the authors can establish the convergence of the proposed constrained models. I think this is the most novel part of the paper.\n\n[1] Liu et al., Sampling with trustworthy constraints: a variational gradient framework\n, NeurIPS 2021.\n\n[2] Danilo et al., Generalized ELBO with Constrained Optimization, GECO,  Bayesian Deep Learning (NeurIPS 2018).\n\n[3] Ferdinando et al., Lagrangian Duality for Constrained Deep Learning, ECMLPKDD 2020.", "weaknesses": "- This paper shares a close connection to [1] regarding its motivation: formulating fairness as a constrained distributional optimization. Both of them employs a methodology based on primal-dual optimization derived  from Lagrangian duality to address the constraints. While [1] focuses on the constrained SVGD, this paper focuses on constrained diffusion models.\n- Some minor corrections:\n   +) Algorithm 1, line 4: x_{\\theta}(h) should be s_{\\theta}(h)\n   +) Section 5, line 314, q(x_0:T) should be q_{i}(x_0:T) (constraint part)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive evaluation and the valuable feedback. We believe that we have fully addressed your concerns/questions, and will incorporate all points mentioned below in the final version. We would be happy to address any further questions you might have. \n\n---\n> **Weakness** 1... connection to [1] ...\n\n> [1] Liu et al., Sampling with trustworthy constraints: a variational gradient framework , NeurIPS 2021.\n\n> **Question** 2... differs from existing constrained generative models, for example [1] ...\n\n> **Limitation** 2 a discussion of how the proposed method differs from existing constrained generative models should be added.\n\n> [2] Danilo et al., Generalized ELBO with Constrained Optimization, GECO, Bayesian Deep Learning (NeurIPS 2018).\n\n> [3] Ferdinando et al., Lagrangian Duality for Constrained Deep Learning, ECMLPKDD 2020.\n\n**Response:** We completely agree that [1] and our paper share the motivation of using distribution constraints. It is worth mentioning that our constrained diffusion model differs from the constrained sampling in [1] in four key aspects.\n\n- In methodology, we model *unknown* data distribution, while [1] works on sampling from *known* distribution. \n\n- We have different optimization problems: [1] optimizes a *reverse KL divergence* subject to a *moment constraint*, while we use the *forward KL divergence* to form both objective and constraints. \n\n- We have different algorithms: we develop a training algorithm (see Algorithm 1) in a dual space, while [1] poses a sampling method working in the primal-dual fashion. \n\n- We have different theory: Regarding distributions, our theory only has the mild assumption that the samples should be bounded, while [1] assumes analytical properties of the target distribution, e.g., the Log-Sobolev conditions. \n\nThank you for pointing out the additional references [2, 3]. Reference [2] studies the classical VAE under moment constraints, and [3] studies constrained deep learning, being not directly applicable to diffusion models. Both references empirically develop primal-dual training algorithms without guarantees. Therefore, our constrained diffusion model distinguishes itself in terms of problem, algorithm, and theory. \n\n---\n> **Weakness** 2 Some minor corrections ...\n\n**Response:** Thank you for pointing out typos. We will remove them upon revision, and double-check the paper's writing in the final version. \n\n---\n> **Question** 1 How efficient the model is as the number of constraints increases? ...\n\n> **Limitation** 1 a discussion of the efficiency of the proposed constrained diffusion models should be included.\n\n**Response:** Thank you for bringing our attention to the efficiency. First, we note that the complexity of sampling from our constrained diffusion model does not increase with the number of constraints, as our trained diffusion model functions like a standard diffusion model to generate samples. Importantly, we remark that training our constrained diffusion model has comparable efficiency to training standard diffusion models detailed next. \n\nThe additional computational cost of our dual-based training (Algorithm 1) arises from: (i) updating the dual variables; (ii) updating the diffusion model in the primal update.  \n\n- (Cost of updating the dual variables) We note that our dual-based training has the same number of dual variables as the number of constraints. Thus, the cost for the dual update is linear in the number of constraints. To update each dual variable, we can directly use the ELBO loss over the batches sampled from each constrained dataset (already computed for the Lagrangian). Therefore, the cost of updating dual variables is negligible.\n\n- (Cost of updating the diffusion model in the primal update) We note that the primal update trains a standard diffusion model based on the Lagrangian with updated dual variables. In our experiments, this primal training often requires as few as 2-3 updates per dual update. Thus, when training our constrained model, we can train for the same number of epochs as an unconstrained model but update the dual variables after every few primal steps. As a result, training our constrained diffusion model is almost as efficient as training standard unconstrained models.\n    \nThe only concern we encountered regarding efficiency is that batches need to be sampled from every constrained dataset at each step to estimate the Lagrangian. This introduces a small GPU memory overhead that increases with additional constraints. However, this is somewhat mitigated by the fact that constrained datasets are often much smaller than the original dataset, allowing us to choose a smaller batch size for the constrained datasets without degrading performance.", "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": "WGeZpIzBHw", "reviewer": "Reviewer_Rjho", "summary": "This paper studies the constrained diffusion models, motivated by customizing the generation by specific tasks. The idea is to formulate KL divergence-constrained optimization problem (U-KL), which is shown to have zero duality gap. The convergence (rate) of the sampling process is given by assuming a mixture initial distribution. Several experiments have been conducted, e.g., on MNIST and on the fairness.", "questions": "See weakness", "limitations": "NA", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper provides a thorough study on the KL-constrained diffusion models: both optimization and sampling aspects are considered. The paper is overall interesting (methodologically). I also checked most proofs, and they are correct.", "weaknesses": "There are several weaknesses (and questions):\n\n(1) The problem (U-KL) is still abstract. It is not clear how to choose $b_i$ to represent desired properties, or what are the right interpretations of the KL constraints (e.g., Section 3.2 on the fairness).\n\n(2) Theorem 7: the authors may emphasize that Theorem 7 is valid under the assumption that the initial distribution is a mixture. Also I think DDPM is used. The authors may also compare Theorem 7 with the results in [28].\n\n(3) Experiments: the authors consider MNIST, CIFAR 10 and Celeb-A. However, the FID obtained seems to be large compared to the existing results. For instance, for MNIST the FID is typically less than 1 while the paper reports ~20-40; for CIFAR 10 the FID is less than 3 (the best is around 1.8) while the paper reports ~50. The authors may explain why this happens, otherwise the experiment results are not so convincing. \n\n(4) Literature and comments: There have also been a line of work on the SDE-based diffusion models (see ref. a), which show good empirical results. Related to (3), ref b. reports the FID 1.8 for CIFAR 10, and ref c. reports the FID 0.8 for MNIST (both without constraints). The authors may refer to these results on the continuous-time (SDE-based) models, and explain why there is so large difference in the empirical results. Another suggestion is that the authors may consider fine-tuning in the continuous framework, as in [45] and ref d.\n\na.  Score-Based Generative Modeling through Stochastic Differential Equations, Song et al., arXiv:2011.13456.\n\nb.  Elucidating the Design Space of Diffusion-Based Generative Models, Karras, Aittala, Aila and Laine, arXiv:2206.00364.\n\nc.  Contractive Diffusion Probabilistic Models, Tang and Zhao, arXiv:2401.13115.\n\nd.  Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond, Tang, arXiv:2403.06279.\n\nI would be happy to raise the score if some (or all) of the above concerns are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive evaluation and the valuable feedback. We believe that we have fully addressed your concerns/questions and will incorporate all points mentioned below in the final version. If you have any further questions, please feel free to post them, and we would be glad to address them.\n\n---\n\n> **Weakness** (1) ... how to choose $b_i$ ... interpretations of the KL constraints ....\n\n**Response:** Our KL constraint thresholds $(b_i, i=1,\\ldots m)$ function as a set of balancing weights over constrained distributions $\\{q^i, i=1,\\ldots m\\}$. A smaller $b_i$ (a tighter constraint) causes the model to give more weight on the constrained distribution $q^i$ (i.e., sampling more often from $q^i$). This ties into our desired properties for each setting in different ways:\n\n- In the minority class setting, a smaller $b_i$ leads to sampling more often from the minority classes, which is our desired property. \n\n- In the fine-tuning setting, a smaller $b_i$ means sampling more often from the pre-trained model, ensuring the new model does not *forget* the pre-trained model.\n\nThe distribution-balancing role of constraint thresholds $(b_i, i=1,\\ldots,m)$ can be shown by analysing the optimal dual variable $\\lambda^\\star$; Recall from Section 3.2: the dual function $g(\\lambda) = h(q_{\\text{mix}}^{(\\lambda)}) + \\sum_{i = 1}^m \\lambda_i b_i $, where $h(q_{\\text{mix}}^{(\\lambda)})$ is the differential entropy of a mixture distribution $q_{\\text{mix}}^{(\\lambda)}$. The maximizer $\\lambda^\\star$ of the dual function determines the optimal weights in the final learned mixture $q_{\\text{mix}}^{(\\lambda^\\star)}$ i.e., how often the trained model samples from each distribution $q^i$. Setting the gradient of the dual function to be zero leads to \n$$\\frac{\\lambda^\\star_i}{1 + (\\lambda^\\star)^T 1} \n= e^{h_i - b_i}  \\text{ for }  i = 1,\\ldots, m$$ \nwhere $\\frac{\\lambda^\\star_i}{1 + (\\lambda^\\star)^T 1}$ is the weight of $q^i$. A smaller $b_i$ leads to a larger weight for $q^i$ and vice versa. Another interesting implication of $\\lambda^\\star$ is based on its dependence on the entropy $h_i$ of each individual distribution. When the constraint thresholds are equal, the model learns to sample more often from a distribution $q^i$ that has high entropy, meaning that a constrained model can generate more 'diverse' samples than an unconstrained model. \n\nIn practice, we choose the constraint thresholds $(b_i, i=1,\\ldots,m)$ by starting with a value that is close to the minimum loss achieved by an unconstrained model. Based on the insight above, we increase/decrease $(b_i, i=1,\\ldots,m)$ depending on if we're sampling too rarely/often from the constrained distributions. Another tuning method that we found useful in practice, is resilient constrained learning in which the constraint thresholds are updated adaptively during training; see it in Appendix E.2 and Algorithm 3.\n\n---\n\n> **Weakness** (2) Theorem 7 ... the initial distribution is a mixture ...  \n\n**Response:** Theorem 7 shows that the trained diffusion model from our dual-based training (Algorithm 1) converges to a distribution that is close to a mixture of data distributions weighted by an optimal dual. Importantly, Theorem 7 does not assume an initial mixture distribution, rather, Theorem 3 asserts that the optimal solution to the constrained problem is a mixture distribution and we use this in Theorem 7. Compared to [28], we have advanced unconstrained diffusion models to address constrained problems, and we additionally characterize the effect of constraints on the diffusion model through duality analysis in constrained optimization.\n\n---\n\n> **Weakness** (3) ... FID obtained seems to be large ...\n\n**Response:** We note that our FID scores being larger compared to existing results is not a weakness of our approach, but rather a consequence of our experimental setup. We train both the unconstrained and constrained models, on a *biased* subset of the dataset wherein some of the classes have significantly fewer samples than the rest. We then compute the FID scores for these models compared to the actual dataset itself which is *unbiased* (i.e., every class has the same number of samples). These FID scores approximate how close the learned distribution of the model trained on biased data, is to the underlying unbiased distribution.\n\nThis setup contrasts with existing results in the literature, where the FID is computed with respect to unbiased data, and the models are also trained on unbiased data. Therefore, it is expected that such models will achieve better FID scores than constrained or unconstrained models trained with biased data. Our purpose in reporting the FIDs was not to compare them to existing results (as such a comparison would be uninformative) but to demonstrate that, when trained on biased data, the constrained model achieves better FID scores than the unconstrained model.\n\n---\n\n> **Weakness** (4) ... SDE-based diffusion models ... good empirical results ... fine-tuning in the continuous framework, as in [45] and ref d. \n\n> d. Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond, Tang, arXiv:2403.06279.\n\nFor the difference in empirical results, please refer to our previous response. We note that [45] and ref. d study how to fine-tune pre-trained diffusion models to respect certain desired generation properties. In contrast, we focus on training new diffusion models to respect desired generation properties by imposing KL-divergence constraints. Despite having different problem setups, our constrained formulation can be used in fine-tuning problems. For instance, if a high-quality dataset that satisfies our desired properties is available, we can impose the KL divergence between the fine-tuning model and the underlying distribution of the high-quality dataset as a constraint in [45]. We will discuss this direction as future work in the final version.", "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": "yAO5gOQgqy", "reviewer": "Reviewer_c1z8", "summary": "This paper proposes Dual Training, a method designed to constrain the distributions that denoising diffusion models can learn. The authors propose an extension to the standard diffusion model training objective that minimizes the Kullback-Leibler (KL) divergence of the learned distribution with respect to two components: (1) the original data distribution and (2) a set of auxiliary distributions representing relevant constraints. The authors derive a tractable algorithmic approximation of their approach and apply it to two scenarios: the fair generation of underrepresented image classes and the fine-tuning of pre-trained models on new data while preserving image classes from the pre-training dataset.", "questions": "* Can you provide more details on the computational overhead of dual training compared to standard diffusion model training?\n* How does the method compare to other approaches that explicitly aim to fit mixture distributions, e.g. [1]?\n* How sensitive are the results to the choice of hyperparameters (number of dual iterations, primal/dual batch sizes, primal/dual learning rates, etc.)? Is there a principled way to set these values?\n\n--- \n\n[1] Du, Yilun, et al. \"Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and MCMC.\" International conference on machine learning. PMLR, 2023.", "limitations": "The authors discuss some limitations, including the need to experiment with more datasets/attributes and explore converting other types of constraints into the KL divergence formulation. It would be good to also discuss and compare the computational overhead of the proposed formulation as this would be beneficial to evaluate the method's practical implications and trade-offs", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "* The paper proposes an original approach to addressing two critical challenges in diffusion-based image generation: fairness and the ability to integrate new data with pre-trained diffusion models.\n* The method is well-motivated and the theoretical analysis of both the constrained optimization problem and the proposed solution embeds it in a rigorous mathematical framework.\n* Experimental results demonstrate the efficacy of the proposed approach on two relevant tasks (fair generation and fine-tuning), showcasing quantitative and qualitative improvements over unconstrained and pre-trained baselines.", "weaknesses": "* The experimental evaluation in the main text is limited, focusing primarily on the MNIST and CelebA datasets. While some results on CIFAR-10 are provided in the appendix, they are not referenced in the main text. A more comprehensive empirical evaluation using more challenging datasets would strengthen the experimental section and better demonstrate the method's utility.\n* The empirical evaluation lacks relevant baselines. For instance, the fine-tuning experiments do not provide quantitative results for models fine-tuned without constraints. More generally, it would be beneficial to see how the model compares with standard conditioning approaches capable of enforcing similar constraints, for example, training a model with classifier-free guidance on datasets with underrepresented classes and providing balanced conditioning information at inference time.\n* The empirical results do not include uncertainty estimates or measures of statistical significance, which would enhance the robustness of the findings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and the valuable feedback. We believe that we have fully addressed your concerns and will incorporate the points mentioned below into the final version. We would be happy to address any further questions you might have.\n\n---\n> **Weakness** 1 ... more challenging datasets ... \n\n**Response:** We have expanded the scope of our experiments to include constrained latent diffusion models to tackle more challenging datasets. Initial results for the image-net dataset using a latent space diffusion model are included in the attached PDF in the global rebuttal. The constrained model samples more often from the minority classes, demonstrating the utility and effectiveness of our method even when applied to a more modern diffusion paradigm and on a much more challenging dataset. We note that the image-net dataset is significantly more challenging than MNIST, CELEB-A, both because of the much higher resolution (resized to 256*256 vs 32*32) and the classes being much more diverse compared to each other. \n\n---\n> **Weakness** 2 ...  quantitative results for models fine-tuned without constraints ... baselines ...\n    \n**Response:** We refer to Figure (2b) in the paper which in fact provides a quantitative baseline for models fine-tuned without constraints, both as generated samples and its FID score. Regarding the second point, we believe there is no meaningful metric of comparing a constrained unconditional model to existing conditional models since the conditional model can sample from any class if it is conditioned to do so. The suggested approach of using a conditional model with balanced conditioning information at inference time is interesting. However, we were unable to find any existing baselines using the suggested approach. Hence, we believe implementing such a novel approach to use as a baseline is beyond the scope of this paper.\n\n---\n> **Weakness** 3 ...  statistical significance ... \n\n**Response:** We thank the reviewer for this suggestion. It is straightforward to get uncertainty estimates and error bars for the results/plots in the paper as they are reproducible from our shared code. We will include them in the final version. \n\n---\n> **Question** 1 ... computational overhead of dual training ...\n\n**Response:** See our response to **Question** 1 of **Reviewer EijC**.\n\n---\n> **Question** 2 ... compare to other approaches that explicitly aim to fit mixture distributions, e.g. [1]?\n\n> [1] Du, Yilun, et al. \"Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and MCMC.\" 2023.\n\n**Response:** Thank you for pointing out this interesting reference [1]. We summarize three main differences below.\n\n- (Task) One of the compositional generation tasks in [1] is to learn a mixture of several distributions (or experts) using the energy-based model method. However, they assume that the weights of different distributions comprising the mixture are the same. In contrast, our approach aims to find weights such that the mixture satisfies certain constraints.\n\n- (Loss) In the energy-based model method [1], the training loss is the Fisher divergence between the model and a smoothed version of a data distribution. When the data distribution is a mixture of several experts, we have to individually train an energy-based model for each expert, and finally mix them through a specific sampling method. In contrast, our constrained diffusion model method optimizes the standard diffusion ELBO loss (e.g., score matching) over an dynamically-mixed distribution, where the mixing weight is determined by the dual update. Sampling from our trained model during inference is the same as any standard diffusion model: we iteratively refine a white noise through our trained reverse diffusion process.\n\n- (Theory) Although an energy-based model method is outlined in [1], performance of a trained energy-based model is not analyzed in theory. In contrast, we show that our trained constrained diffusion model converges to a mixed data distribution.  \n\n---\n> **Question** 3 ... sensitive ... to the choice of hyperparameters ... \n\n**Response:** We thank the reviewer for bringing up this sensitivity question. We will include a more thorough discussion of the sensitivity to different hyperparamters in the supplementary material. A brief discussion of each hyperparameter is given below:\n\n- Number of dual iterations: In our implementation this shows up as the number of primal GD steps per dual update (\\#primal\\_per\\_dual). Experimentally, we have observed that as long as \\#primal\\_per\\_dual  is greater than 1, the results are not sensitive to this value. Also, as discussed in our response to **Question** 1, the dual updates add a negligible computational overhead. Hence, updating the dual nearly as many times as we update model parameters doesn't reduce training efficiency.\n        \n- Primal/Dual batch sizes: We thank the reviewer for bringing this to our attention. We have included (see the attached PDF in the global rebuttal) results of training a constrained model on an unbalanced subset of MNIST, using different primal/dual batch sizes. The results suggest that when the ratio between Primal and Dual batch sizes is larger, the model performs better (lower FID and more evenly distributed samples).  This is in line with the heuristic we used in the included experiments in the paper where we chose the batch sizes such that the ratio of primal to dual batch size is close to size ratio of entire dataset to constraint datasets (which are much smaller).\n\n- Primal/Dual learning rate: For the primal learning rate, we followed the best practice used to train standard diffusion models. For the dual learning rate $\\eta$, we refer to Theorem 8 in the paper, showing a smaller error bound for smaller $\\eta$ while slowing convergence. In practice, as long as $\\eta \\leq 0.1$, we observed that the model converges to similar results reliably.", "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": "pJvvby7KlS", "reviewer": "Reviewer_Zh95", "summary": "In this paper, the authors introduce a novel approach termed `Dual Training` for training constrained diffusion models, particularly focusing on scenarios involving `biased data generation`. Initially, the authors adeptly derive the learning objectives for diffusion models within a constrained optimization framework. Subsequently, they develop corresponding learning methodologies based on the principles of Lagrangian duality and propose an innovative training procedure. The effectiveness of their approach is rigorously validated through experiments on two distinct generative tasks.", "questions": "Please see weaknesses.", "limitations": "1. **Organization:** The manuscript lacks a dedicated section for related work, which is essential for contextualizing the study within the existing literature. The absence of this section diminishes the foundational motivation of the manuscript, potentially limiting the reader's understanding of its contributions relative to prior research.\n\n2. **Baseline Comparison:** The authors did not include experiments comparing their methods against established baselines. Such comparisons are crucial for demonstrating the efficacy and advancements of the proposed approach over existing techniques.\n\n3. **Hyper-parameter Sensitivity Analysis:** The manuscript does not address hyper-parameter sensitivity. Including such analysis is important to evaluate the robustness and reliability of the model across various parameter settings.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. **Robust Derivation Procedure**: The authors present a thorough and well-founded derivation procedure throughout the manuscript, which stands out as particularly compelling. This rigorous approach significantly strengthens the theoretical foundation of their work.\n\n2. **Relevant Application Scenario**: The authors explore a critical yet underrepresented field in their research. This focus not only highlights the relevance of their study but also underscores its potential impact in areas that have previously received limited attention.", "weaknesses": "**Major Issues:**\n\n1. The authors state, \"Compared with the loss re-weighting method [9], our constrained formulation provides an optimal trade-off between matching given data distribution and following reference distribution.\" It would be helpful if the authors could clarify whether the re-weighting-based approach can be considered a special case of the proposed method in this manuscript.\n\n2. The discussion on diffusion models appears to focus predominantly on the DDPM model. Could the authors explore whether this proposed approach can be extended to other diffusion processes, such as the Variance Exploding Stochastic Diffusion?\n\n3. There seems to be an unclear transition from (U-KL) to (U-LOSS). Can the authors directly justify this conversion? It would be beneficial to include a derivation or a more detailed explanation, especially how it relates to equations (6) and (U-MIX).\n\n4. The manuscript lacks comparative results with baseline models. Given that the authors suggest this approach provides a more general framework for generative models, could comparisons be made with other methods mentioned in reference [1]?\n\n5. Could the authors discuss whether this approach can be interpreted from the perspective of partial optimal transport? The constraints discussed appear to share similarities with those in unbalanced optimal transport approaches, such as those outlined in equation 3 of reference [2].\n\n6. The paper discusses a dual training procedure involving primal-dual optimization. In the context of augmented direction multiplier methods used in directed acyclic graph learning, batch size is known to significantly influence model efficacy. Could the authors provide a sensitivity analysis regarding batch size effects?\n\n**Minor Issue:**\n\n1. Under Assumption 1, should the range of $\\zeta \\in (0, b_i)$, considering that KL divergence cannot be negative? \n\n2. Regarding the term (U-LOSS), should the subject be revised to $\\nabla{\\log{q^i(x_t)}}$? Additionally, it would be helpful if the authors could provide a detailed explanation or derivation of how the constraint term is converted into the (U-LOSS) term.\n\n---\n\n**References**  \n[1] Choi, Kristy, et al. \"Fair generative modeling via weak supervision.\" International Conference on Machine Learning. PMLR, 2020.  \n[2] Duque, Andrés F., Guy Wolf, and Kevin R. Moon. \"Diffusion transport alignment.\" International Symposium on Intelligent Data Analysis. Cham: Springer Nature Switzerland, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and the valuable feedback. We believe that we have fully addressed your concerns and will incorporate the points mentioned below into the final version. We would be happy to address any further questions you might have.\n\n---\n> **Major Issue** 1 ... the loss re-weighting method [9] ... can be considered a special case of the proposed method ...\n\n**Response:** First of all, [9] and our method both utilize a biased dataset to train a generative model with fairness. However, we have different fairness-promoting objectives: (i) [9] aims to remove the bias of a generative model by adding an importance sampling weight based on a fair reference dataset; (ii) our diffusion model aims to eliminate the bias of a diffusion model against certain minorities by using minority datasets to constrain the model. Hence, it seems *unfair* to claim that one is a special case of another. However, it is useful to check their target distributions. For instance, if we choose our original and constrained data distributions as $q= p_{\\text{bias}}$ and $q^1 = p_{\\text{ref}}$ from [9], then our optimal target distribution is in the mixed form of $q_{\\text{mix}}^{(\\lambda^\\star)} = (q+\\lambda^\\star q^1)/(1+\\lambda^\\star) = \\frac{1}{1+\\lambda^\\star} (p_{\\text{bias}} + \\lambda^\\star p_{\\text{ref}})$, where the optimal dual  $\\lambda^\\star$ provides a tradeoff between $p_{\\text{ref}}$ and $p_{\\text{bias}}$. In comparison, Algorithm 1 of [9] applies the importance sampling weight $w(x) = \\frac{p_{\\text{ref}}(x)}{p_{\\text{bias}}(x)}$ for $p_{\\text{bias}}$ to completely eliminate $p_{\\text{bias}}$, which can be viewed as an extreme case of $q_{\\text{mix}}^{(\\lambda^\\star)}$ for very large $\\lambda^\\star\\to\\infty$. In this case, we see that our constrained diffusion model generalizes the re-weighting approach [9] to be a *soft* re-weighting mechanism.\n\n---\n> **Major Issue** 2 ... extended to other diffusion processes ... Variance Exploding Stochastic Diffusion\n\n**Response:** Our constrained distribution optimization formulation doesn't rely on a particular diffusion process, and we use DDPM mainly for clean exposition. We note that the variance exploding diffusion and DDPM (variance-preserving) share share the same variational structure (e.g., ELBO loss), with only difference in scheduling noise parameters [R1]. Hence, we can apply our constrained diffusion model to the variance exploding diffusion [R1], and establish non-asymptotic convergence [R2]. In the final version, we will discuss further inclusions of other diffusion processes [R3]. \n\n**Reference**\n\n[R1] Variational Diffusion Models\n\n[R2] The Convergence of Variance Exploding Diffusion Models under the Manifold Hypothesis\n\n[R3] Score-based Diffusion Models via Stochastic Differential Equations--a Technical Tutorial \n\n\n---\n> **Major Issue** 3 ... unclear transition from (U-KL) to (U-LOSS) ... how it relates to equations (6) and (U-MIX).\n\n**Response:** Reformulation of Problem (U-KL) as Problem (U-LOSS) has two key steps: (i) apply the ELBO representation of each KL divergence in Equation (3); (ii) represent EBLO in form of denoising matching (e.g., score matching) in Appendix B. This derivation doesn't rely on the Lagrangian (6) and the optimal constrained model (U-MIX). \n\n---\n> **Major Issue** 4 ... comparisons be made with other methods mentioned in reference [1]?\n\n> [1] Choi, Kristy, et al. \"Fair generative modeling via weak supervision.\" 2020.\n\n**Response:** We cite [1] above as [9] in our paper. Related works on fairness & generative modeling in [1] study the classical VAE and GAN, being not directly comparable to diffusion models in methodology. \n\n---\n> **Major Issue** 5 ... interpreted from the perspective of partial optimal transport? ... in equation 3 of reference [2].\n\n> [2] Duque, Andrés F., Guy Wolf, and Kevin R. Moon. \"Diffusion transport alignment.\" 2023.\n\n**Response:** [2] aims to find a coupling between data samples from two domains, under constraints on total/individual masses. Analogously, our constrained diffusion model can be viewed as a transportation from a white noise to data samples via a reverse diffusion process. However, our constrained distribution optimization problem is a nonlinear optimization, while the constrained optimal transport is a linear optimization (see (3) in [2]). Therefore, our method can't be viewed as a case of partial optimal transport [2] or vice versa.\n\n---\n> **Major Issue** 6 ... a sensitivity analysis regarding batch size effects? \n\n**Response:** We have included the results of training a constrained model on an unbalanced subset of MNIST, using different primal and dual batch sizes (see the attached PDF in the global rebuttal). A larger ratio between the primal and dual batch sizes leads to better performance, as indicated by lower FID scores and more evenly distributed samples. This finding aligns with the heuristic used in the experiments in the paper, where we selected batch sizes so that the ratio of primal to dual batch sizes approximates the ratio of the entire dataset to the constrained datasets. We will include this empirical sensitivity analysis, along with similar analyses for other hyperparameters, in the final version. \n\n---\n> **Minor Issue** 1 ... range of $\\zeta$ ...\n\n**Response:** It is correct. \n\n---\n> **Minor Issue** 2 ... the term (U-LOSS) ... $\\nabla \\log q^i(x_t)$ ... \n\n**Response:** It's worth noting that the forward processes for original and constrained datasets are the same, except for initial data points. Thus, we use the notation $\\nabla \\log q(x_t)$ in (U-LOSS), and differentiate initial data points in two expectations: $E_{q(x_0)}$ and $E_{q^i(x_0)}$. \n\n---\n> **Limitations** Existing Literature ... Baseline ... Sensitivity ...\n\n**Response:** See our global rebuttal on **Related work**. For the sensitivity analysis, please refer to our response to **Question** 3 of **Reviewer c1z8**.", "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:
We thank the reviewers for recognizing the strengths of our contribution and providing valuable feedback. We believe that we have addressed your concerns and hope that you will reconsider our paper in light of our rebuttal. To better assist your evaluation, we summarize three shared concerns in this global rebuttal and outline our responses below. Please see our individual rebuttals directly following your review. 

- **Related work:** We thank the reviewers for bringing some related works to our attention. In each individual rebuttal, we have clarified how our work differs from these references in several key ways. We will incorporate these points into an expanded related work section in the final version, alongside our existing discussion of related work, to better contextualize how our approach relates to (and differs from) previous methods tackling similar problems.

- **Significance of experiments:** We believe that our current experiments in the paper have validated the utility and effectiveness of our constrained diffusion framework. To further strengthen this message, we have prepared new results from training constrained latent diffusion models on the more challenging *Image-Net* dataset; see them in the attached PDF. In our individual rebuttals, we have also discussed how the efficiency of training constrained models is comparable to that of training unconstrained models, and further discussed the choices of different hyperparameters and their effects. In particular, we have included a table in the attached PDF showing how the primal/dual batch sizes affect the performance of the constrained model. Based on the reviewers' concerns, we will include these new results and an expanded hyperparameter sensitivity analysis in the appendix of the final version.

- **Other clarifications:** **Reviewer Rjho** brings up an important clarifying question regarding the choice of constraint thresholds and their relation to the desired properties of the constrained model. We have addressed this by pointing out the relationship between the constraint thresholds $b_i$ and the weights that the model learns for each constrained distribution $q^i$. **Reviewer Zh95** raises an important question about the generality of our constrained diffusion model. We have discussed how other diffusion processes can be incorporated into our framework, and we will discuss such inclusions in the final version.

If you have any further questions, please feel free to post them, and we would be glad to address them.



Best,

The authors of Submission21568

Meta-review: The paper proposes a method for training diffusion models to satisfy constraints based on the KL divergence to certain reference distributions. The problem is reformulated using Lagrangian duality and a dual ascent algorithm is proposed to solve it. A rigorous optimality analysis is given under clearly stated assumptions.

Reviewers felt the paper addressed an important problem and found the paper to be sound, thorough, methodologically interesting, and rigorous. Weaknesses included questions about the running-time overhead of dual training and questions about the experiments, including comparisons to baselines, comparisons of FID score to published results, and performance on more challenging data sets. 

During the rebuttal, the authors clarified a number of points:
* They used biased subsets of existing benchmarks, so their FID scores are not comparable to published results. 
* Quantitative comparisons a baseline of models fine-tuned without constraints already appeared in the paper.   
* They added results on a more challenging data set, ImageNet.
* They clarified that the running time overhead is small in practice because the primal variables only need a few training iterations each time the dual variables are updated.

Based on these responses, two reviewers raised their scores. In the end, the reviewers unanimously recommend accept. Reviewer Rjho did not raise their score, but the meta-reviewer feels that their weaknesses were addressed by the rebuttal. 

The meta-reviewer also looked at the paper and was impressed with the scope and rigor of the theoretical results, which go beyond other papers in this area. For example, in addition to the convex analysis in the space of distributions and score functions, the authors give a convergence/optimality analysis that considers the score function parameterization under clearly stated assumptions about the parameterization error.
Decision: Accept (poster)