Title: Constrained Diffusion Models via Dual Training
Venue: NeurIPS 2024
Accepted: yes
Keywords: Constrained diffusion model

constrained optimization

Lagrangian method

dual algorithm
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

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.

1 Introduction

Diffusion models have become a driving force of modern generative modeling, achieving groundbreaking performance in tasks ranging from image/video/audio generation [64, 6, 43] to 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 we can generate new 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].

Classic diffusion models [66, 34, 68] have been extended to generate data under different requirements through either first-principle methods or 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 us 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, it is natural to incorporate requirements into diffusion models by imposing constraints on these optimization

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 a 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 to the distributions of underrepresented data. Compared with the typical importance re-weighting 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 a 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 infinite-dimensional distribution space, where the constraints require 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 a mixture data distribution that achieves the optimal trade-off among objective and constraints.

(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 optimization/parametrization errors.

(iii) We empirically demonstrate the merit of our constrained diffusion models in two aforementioned requirements. In the fair generation task, we show that our constrained model promotes sampling more from the minority classes compared to unconstrained models, leading 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-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, 1], 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-weighting method for diffusion models [41] and GANs [16], which reduces bias in a dataset by re-weighting it with a pre-trained debiased density ratio. In contrast, we design our diffusion models to mitigate the bias in a dataset by imposing distribution constraints without pre-training. In addition to being distinct from existing methods, we provide a 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 a fine-tuning model and a 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., V AEs [61], GANs [16]), and classical sampling methods (e.g., Langevin dynamics and Stein variational gradient descent [50]), our work is different because we focus on diffusion-based generative models.

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 x 0 ∈ R d , and generates latent variables { x t } T t =1 from t = 1 to T by adding white noise recursively. The joint distribution of latent variables results from conditional probabilities q ( x 1: T | x 0 ) = ∏ T t =1 q ( x t | x t -1 ) , where the distribution of latent variable x t conditioned on the previous latent x t -1 is given by a Gaussian distribution q ( x t | x t -1 ) = N ( x t ; µ t , σ 2 q ( t ) I ) , where µ t = √ α t x t -1 and σ 2 q ( t ) = 1 -α t . Since the forward process is a linear Gaussian model with pre-selected mean and variance, it is solely determined by the data distribution q ( x 0 ) . We often refer to q ( x 0 ) as a forward process.

The backward process begins with the latent x T sampled from the standard Gaussian p ( x T ) = N ( x T ; 0 , I ) , and decodes latent variables from t = T to t = 0 with a joint distribution

Here, p is our distribution model that can be used to generate new samples. We denote by P the set of all joint distributions over x 0: T in form of (1), where p ( x t -1 | x t ) is a conditional Gaussian with a fixed variance (see Appendix A). Throughout the paper, we work in the convergent regimes of the backward process (e.g., [15, 11, 2]). Without loss of generality, we adopt the convergent regime in [47] by taking the scheduling parameter α 1 = 1 -1 /T c 0 and α t = 1 -c T min((1 -α 1 )(1 + c T ) t , 1) for t > 1 , and the variance as σ 2 p ( t ) = 1 /α t -1 , where c T := c 1 log( T ) /T and c 0 , c 1 are some constants. Hence, ¯ α T ≈ 0 implies q ( x T ) ≃ N ( x T ; 0 , I ) . Thus, q ( x 0: T ) ∈ P . Also, ¯ α 1 ≈ 1 implies q ( x 1 ) ≃ q ( x 0 ) . 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 D KL ( q ∥ p ) := E x ∼ q ( x ) log( q ( x ) p ( x ) ) . Generative diffusion modeling aims to generate samples whose distribution is close to 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: maximize p ∈P E q ( x 0 ) [ log p ( x 0 ) ] , where

and E ( p ; q ) := E q ( x 0 ) E q ( x 1: T | x 0 ) log p ( x 0: T ) q ( x 1: T | x 0 ) is known as ELBO in variational inference [7, 55]. Alternatively, we aim to minimize the KL divergence between the forward/backward processes,

Thus, we connect the log-likelihood maximization to the KL divergence minimization via ELBO.

Lemma 1 (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 a solution for the log-likelihood maximization problem, i.e.,

See Appendix B.1 for proof. Lemma 1 states that improving the ELBO score increases the likelihood of a backward process that generates the data, together with the fit of a backward process to the forward process. Hence, finding the best backward process becomes optimizing one of three equivalent objectives. In practice, ELBO serves as a loss function approximated by

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 is our focal setting. Using standard diffusion derivations [55], the ELBO maximization can be shown to equal to a quadratic loss minimization,

where E x 0 ,t,x t is an expectation over the data distribution q ( x 0 ) , a discrete distribution p ω ( t ) from 2 to T , and the forward process q ( x t | x 0 ) at time t given the data sample x 0 ; see Appendix A for details. The minimization is done to find a function ̂ s ∈ S that can best predict the gradient of the forward process over data ∇ log q ( x t ) , commonly called the (Stein) score function, where S is a set of valid score functions mapping from R d × N to R d . In practice, however, we parametrize the estimator ̂ s ( x t , t ) as ̂ s θ ( x t , t ) with parameter θ , which gives our focal objective of generative modeling: minimize θ ∈ Θ E x 0 , t, x t [ ∥ ̂ s θ ( x t , t ) -∇ log q ( x t ) ∥ 2 ] . A parametrized form of p ( x t -1 | x t ) associated with ̂ s θ ( x t , t ) is denoted by p θ ( x t -1 | x t ) and the backward process has a parametrized joint distribution p θ ( x 0: T ) . 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.

3 Variational constrained diffusion models

Weintroduce constrained diffusion models by considering the unparametrized set of joint distributions in Section 3.1, and illustrating constraints via two examples in Section 3.2.

3.1 KL divergence-constrained diffusion model: unparametrized case

The standard diffusion model specifies a single data distribution, denoted by q in Lemma 1. To account for other generation requirements, we introduce m additional data distributions { q i } i m =1 that represent m desired properties on generated data. To incorporate new properties into the diffusion model, we formulate an unparametrized KL divergence-constrained optimization problem,

Let an optimal solution to Problem (U-KL) be p ⋆ . Then the optimal value of the objective function is F ⋆ := D KL ( q ∥ p ⋆ ) . Problem (U-KL) aims to find a model p ⋆ that generates data from the original distribution q while staying close to m distributions { q i } i m =1 that encode our desired properties, e.g., unbiasedness towards minorities. Let the Lagrangian for Problem (U-KL) be

for λ ≥ 0 . The dual function g ( λ ) is given by g ( λ ) := min p ∈P L ( p, λ ) , 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 a model p ∈ P and ζ > 0 such that D KL ( q i ( x 0: T ) ∥ p ( x 0: T )) ≤ b i -ζ for all i = 1 , . . . , m .

Let an optimal dual variable of Problem (U-KL) be λ ⋆ ∈ argmax λ ≥ 0 g ( λ ) , and the optimal value of the dual function be D ⋆ := g ( λ ⋆ ) . From weak duality, the duality gap is non-negative, i.e., F ⋆ -D ⋆ ≥ 0 . Moreover, due to the convexity of KL divergence, Problem (U-KL) is a convex optimization problem, and thus it satisfies strong duality; see 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 ⋆ , λ ⋆ ) is an optimal primal-dual pair of Problem (U-KL) .

Let a mixture data distribution be q ( λ ) mix := ( q + ∑ m i =1 λ i q i ) / (1 + λ ⊤ 1) for λ ≥ 0 . We denote by q ( λ ) mix ( x 0: T ) a joint distribution of the forward process with data distribution q ( λ ) mix . Leveraging strong duality, we show that an optimal model can be obtained by solving an equivalent unconstrained problem in Theorem 1 and its proof is deferred to Appendix B.3.

Theorem 1 (Optimal constrained model) . Let Assumption 1 hold. Then, Problem (U-KL) equals

where q ( λ ⋆ ) mix ( x 0: T ) is the joint distribution of the forward process at an optimal dual variable λ ⋆ .

Theorem 1 states that the KL divergence-constrained problem reduces to an unconstrained KL divergence minimization problem. We notice that the KL divergence is zero if and only if two probability distributions match each other. Hence, q ( λ ⋆ ) mix ( x 0: T ) is the optimal solution to Problem (U-MIX).

Corollary 1. Let Assumption 1 hold. Then, the solution of Problem (U-MIX) , i.e., p ⋆ ( x 0: T ) = q ( λ ⋆ ) mix ( x 0: T ) , is the solution of Problem (U-KL) .

Let ¯ b i := b i -E q i ( x 0 ) [ log q i ( x 0 ) ] . Application of Equality (3) to Problem (U-KL) yields an ELBObased constrained optimization problem,

Recall the model representation in Section 2.2, we can characterize each joint distribution p ∈ P with a function ̂ s ∈ S . Moreover, ELBO reduces to the denoising matching term that has a simplified quadratic form given in Section 2.2. With this reformulation in mind, we cast Problem (U-ELBO) into a convex optimization problem over the function space S ,

where ˜ b i := ( ¯ b i -v ) / ¯ ω . Here, the notation v is 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 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 ̂ s ). Let the Lagrangian L s ( ̂ s, λ ) for Problem (U-LOSS) be

Let the associated dual function be g s ( λ ) := min ̂ s ∈S L s ( ̂ s, λ ) for λ ≥ 0 . Hence, g ( λ ) and g s ( λ ) have the same maximizer λ ⋆ , and the partial minimizer ̂ s ⋆ = argmin ̂ s ∈S L s ( ̂ s, λ ⋆ ) is the solution to Problem (U-LOSS). Hence, an optimal primal-dual pair ( ̂ s ⋆ , λ ⋆ ) to Problem (U-LOSS) gives an optimal primal-dual pair ( p ⋆ , λ ⋆ ) for Problem (U-KL), where p ⋆ is a joint distribution of the backward process induced by ̂ s ⋆ . By this dual property, we take a dual perspective to train constrained diffusion models: we maximize the dual function g s ( λ ) to obtain the optimal dual variable λ ⋆ , and then recover the solution ̂ s ⋆ by minimizing the Lagrangian L s ( ̂ s, λ ⋆ ) .

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 i denotes the distribution of an under-represented subset or 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 a new data distribution. Similarly, we can pose this as an instance of Problem (U-KL), where q 1 denotes the distribution of the new data and q is the distribution of samples generated by the pre-trained model.

Let h i := -E q i ( x 0 ) [ log q i ( x 0 ) ] be the differential entropy of data distribution q i . We relate the KL divergence constraints with the optimal dual variable through entropy in Theorem 2.

Theorem 2. Let Assumption 1 hold, and the supports of data distributions q and { q i } i m =1 be disjoint. Then, the optimal dual variables λ ⋆ to Problem (U-ELBO) are given by

See Appendix B.4 for proof. Theorem 2 characterizes the mixture weights in the target distribution q ( λ ⋆ ) mix : (i) the tighter a constraint is (i.e., smaller threshold ¯ b i ), 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 i . Assumption 1 can be relaxed to a feasibility condition for Problem (U-KL); see Lemma 6 in Appendix B.4.

4 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 a Lagrangian-based dual method in Section 4.2, and present a practical dual training algorithm in Section 4.3.

4.1 KL divergence-constrained diffusion model: parametrized case

With the parametrized model p θ for θ ∈ Θ , we present a parameterized constrained problem,

Let the Lagrangian for Problem (P-KL) be ¯ L ( θ, λ ) := L ( p θ , λ ) . The associated dual function ¯ g ( λ ) is given by ¯ g ( λ ) := min θ ∈ Θ ¯ L ( θ, λ ) . Let an optimal solution to Problem (P-KL) be θ ⋆ . We denote ¯ p := p θ and ¯ p ⋆ := p θ ⋆ , and the optimal objective by ¯ F ⋆ := D KL ( q ∥ ¯ p ⋆ ) . Let an optimal dual variable be ¯ λ ⋆ ∈ argmax λ ≥ 0 ¯ g ( λ ) and the optimal value of the dual function be ¯ D ⋆ := ¯ g ( ¯ λ ⋆ ) .

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 directly apply to Problem (P-KL). For instance, it's invalid to find an optimal solution via an unconstrained problem as in Theorem 1, i.e., ¯ p ⋆ ( ¯ λ ⋆ ) ∈ argmin θ ∈ Θ ¯ L ( θ, ¯ λ ⋆ ) doesn't equal ¯ p ⋆ . The effect of parametrization has to be characterized. However, regardless of parametrization, weak duality always holds, i.e., ¯ F ⋆ -¯ D ⋆ ≥ 0 .

To quantify the optimality of ¯ p ⋆ ( ¯ λ ⋆ ) (closeness of it to q ⋆ mix ), we study a practical representation of model p θ as a parametrized function ̂ s θ ∈ S θ , where S θ is the set of all parametrized score functions. Problem (U-LOSS) is in a parametrized form of

where ˜ b i := ( ¯ b i -v ) / ¯ ω . We note that Problem (P-LOSS) is equivalent to Problem (P-KL). Thus, we let the Lagrangian of Problem (P-LOSS) be ¯ L s ( θ, λ ) := L s ( ̂ s θ , λ ) , and the dual function ¯ g s ( λ ) := minimize θ ∈ Θ ¯ L ( θ, λ ) . Since ¯ g ( λ ) and ¯ g s ( λ ) have the same maximizer ¯ λ ⋆ , ¯ θ ⋆ ∈ argmin θ ∈ Θ ¯ L s ( θ, ¯ λ ⋆ ) , which naturally gives a dual training algorithm in Algorithm 1.

Denote ¯ s ⋆ := ̂ s ¯ θ ⋆ . Thus, ¯ s ⋆ -induced diffusion model is given by ¯ p ⋆ ( ¯ λ ⋆ ) . Algorithm 1 works as a dual ascent method with two natural steps: (i) find a diffusion model with fixed dual variable λ ( h ) ; and (ii) update the dual variable using the (sub)gradient of the Lagrangian L s ( ̂ s θ ( h ) , λ ) . It is known that Algorithm 1 converges to ¯ λ ⋆ since the dual function ¯ g s ( λ ) is concave. However, such convergence in the dual domain doesn't provide optimality guarantee on the primal solution ¯ p ⋆ ( ¯ λ ⋆ ) due to the 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 1 Constrained Diffusion Models via Dual Training

1: Input : total diffusion steps T , diffusion parameter α t , total iterations H , stepsize η .

2: Initialize : λ (1) = 0 .

3: for h = 1 , · · · , H do

4: Compute model ̂ s θ ( h ) ∈ argmin L s ( ̂ s θ , λ ( h )) .

5: Update the dual variable

6: end for

4.2 Optimality analysis of dual training algorithm

We analyze the optimality of ¯ p ⋆ ( ¯ λ ⋆ ) as measured by its distance to q ⋆ mix , i.e., TV ( q ⋆ mix , ¯ p ⋆ ( ¯ λ ⋆ )) , where we denote the total variation distance between two probability distributions p and q by TV ( q, p ) := 1 2 ∫ | p ( x ) -q ( x ) | dx . 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 ⋆ ( ¯ λ ⋆ ) and q ( ¯ λ ⋆ ) mix at ¯ λ ⋆ . Denote a partial minimizer of the Lagrangian by ¯ p ⋆ ( λ ) ∈ argmin θ ∈ Θ ¯ L s ( θ, λ ) for λ ≥ 0 . Noting that minimize θ ∈ Θ ¯ L s ( θ, λ ) is an unconstrained diffusion problem, we are ready to quantify the difference between ¯ p ⋆ ( λ ) and q ( λ ) mix for any λ ≥ 0 using the convergence theory of diffusion models. To do so, we assume the boundedness of samples from a mixed data distribution q ( λ ) mix for λ ≥ 0 , and a small score estimation error.

Assumption 2 (Boundedness of data) . The data samples generated from q ( λ ) mix are bounded, i.e., P ( ∥ x 0 ∥ ≤ T c | x 0 ∼ q ( λ ) mix ) = 1 for any λ ≥ 0 and some large constant c > 0 .

Assumption 3 (Boundedness of score estimation error) . The score estimator ̂ s θ ( x t , t ) estimates the data samples from q ( λ ) mix with bounded score matching error ε score ,

for any λ ≥ 0 , where E q ( λ ) mix ( x 0 ) , t, x t is an expectation over the mixed data distribution q ( λ ) mix ( x 0 ) , a uniform distribution over t from 2 to T , and a forward process q ( x t | x 0 ) given the data sample x 0 .

Since data samples are bounded, Assumption 2 is mild in practice. Assumption 3 is the typical score matching error that is near zero if the function class S θ is sufficiently rich.

With Assumptions 2 and 3, below we bound the TV distance between q ( λ ) mix and ¯ p ⋆ ( λ ) 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 ⋆ ( λ ) to q ( λ ) mix is bounded by

Lemma 3 states that the TV distance between q ( λ ) mix and ¯ p ⋆ ( λ ) decays to zero with a sublinear rate O ( 1 √ T ) , up to a score matching error O ( ε score ) . When the diffusion time T is large, the TV distance between q ( λ ) mix and ¯ p ⋆ ( λ ) is dominated by the score matching error. Substitution of λ = ¯ λ ⋆ into (7) yields an upper bound on TV ( q ( ¯ λ ⋆ ) mix , ¯ p ⋆ ( ¯ λ ⋆ )) , which is the second term of the inequality

Next, we quantify the gap between ¯ λ ⋆ and λ ⋆ , which lets us bound the first term on the RHS of (8) and complete the optimality analysis. To analyze the parametrized optimal dual variable ¯ λ ⋆ , we introduce the richness of the parametrized class S θ and redundancy of constraints at ̂ s ⋆ below.

θ ∈ Θ

Assumption 4 (Richness of parametrization) . For any function ̂ s ∈ S , there exists parameter θ ∈ Θ such that ∥ ̂ s θ -̂ s ∥ L 2 ≤ ν , where ∥·∥ L 2 is with respect to the forward process.

Assumption 5 (Redundancy of constraints) . There exists σ > 0 such that

where ∇ ̂ s is the Fréchet derivative over the function ̂ s 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 S to be bounded ∥ ̂ s ∥ L 2 ≤ T c := R . Using Problem (U-LOSS), we prove that the unparametrized dual function g s ( λ ) is differentiable, and strongly-concave over H with parameter µ , where H := { γλ ⋆ + (1 -γ ) ¯ λ ⋆ , γ ∈ [0 , 1] } and µ := ( σ/ ( 1 + max ( ∥ λ ⋆ ∥ 1 , ∥ ∥ ¯ λ ⋆ ∥ ∥ 1 ))) 2 ,which leads to Lemma 4; see Appendix B.6 for their proofs.

Since TV ( q ⋆ mix , q ( ¯ λ ⋆ ) mix ) is bounded by ∥ ∥ ¯ λ ⋆ -λ ⋆ ∥ ∥ 1 (see Appendix B.6), application of Lemma 3 and Lemma 4 to (8) leads to Theorem 3; see Appendix B.7 for proof.

Theorem 3 (Optimality of constrained diffusion model) . Let Assumptions 1-5 hold. Then, the total variation distance between ¯ p ⋆ ( ¯ λ ⋆ ) and q ⋆ mix is upper bounded by

Theorem 3 states that the TV distance between ¯ p ⋆ ( ¯ λ ⋆ ) and q ⋆ mix decays to zero with a sublinear rate O ( 1 √ T ) , up to a parametrization gap O ( √ ν ) and a prediction error O ( ε score ) . When the parametrization is rich enough, the parametrization gap ν and the prediction error O ( ε score ) are nearly zero. In this case, if the diffusion time T is large, then ¯ p ⋆ ( ¯ λ ⋆ ) is close to q ⋆ mix in TV distance, which recovers the ideal optimal constrained model in 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 1 into a practical algorithm. First, we relax the computation of a diffusion model ̂ s θ ( h ) in line 4 of Algorithm 1 to be approximate: L s ( ̂ s θ ( h ) , λ ( h )) ≤ min θ ∈ Θ L s ( ̂ s θ , λ ( h )) + ε 2 approx , where ε 2 approx is an approximation error of training a diffusion model given λ ( h ) . Second, we replace the gradient in line 5 of Algorithm 1 by a stochastic gradient ̂ E x 0 ∼ q i , t, x t [ ∥ ̂ s θ ( x t , t ) -∇ log q ( x t ) ∥ 2 ] , which enables Algorithm 1 to be a stochastic algorithm, where ̂ E x 0 ∼ q i , t, x t is an unbiased estimate of E x 0 ∼ q i , t, x t . To analyze this approximate and stochastic variant of Algorithm 1, it is useful to introduce the maximum parametrized dual function in history up to step h by ¯ g best ( h ) := max h ′ ≤ h ¯ g s ( λ ( h ′ )) , and an upper bound of the second-order moment of stochastic gradient S 2 := ∑ m i =1 E [ ( ̂ E x 0 ∼ q i , t, x t [ ∥ ̂ s θ ( h )( x t , t ) -∇ q ( x t ) ∥ 2 ] -˜ b i ) 2 | λ ( h ) ] .

Denote the dual variable that achieves ¯ g best ( h ) by λ best. To bound the TV distance between ¯ p ⋆ ( ¯ λ best ) and q ⋆ mix , we check the TV distance between q ( ¯ λ best ) mix and ¯ p ⋆ ( ¯ λ best ) using Lemma 3. The rest is to analyze the convergence of ¯ λ best to λ ⋆ via application of martingale convergence. We defer their proofs to Appendix B.8 and present the optimality of ¯ p ⋆ ( ¯ λ best ) in Theorem 4.

Theorem 4 (Optimality of approximate constrained diffusion model) . Let Assumptions 1-5 hold. Then, the total variation distance between ¯ p ⋆ ( ¯ λ best ) and q ⋆ mix is upper bounded by

Theorem 4 states that the TV distance between ¯ p ⋆ ( ¯ λ best ) and q ⋆ mix decays to zero with a sublinear rate O ( 1 √ T ) , up to a parametrization gap O ( ν ) , a score matching error O ( ε score ) , an approximation error O ( ε approx ) , and stepsize O ( η ) . When the parametrization is rich enough, the parametrization gap ν and the score matching error O ( ε score ) are near zero. Thus, if the diffusion time T is large, then the closeness of ¯ p ⋆ ( ¯ λ best ) to q ⋆ mix 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 to underrepresented classes. We train constrained diffusion models over three datasets: MNIST digits [44], Celeb-A faces [51], and Image-Net 1 [63]. For MNIST and Image-Net, we create a dataset for the distribution q in (P-LOSS) by taking a 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 q i , we use samples from the associated underrepresented class. For Celeb-A, our approach is similar to MNIST except we don't 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, our Image-Net experiment shows that our approach extends to the state-of-the-art diffusion models in latent space.

frequencies

classes

0.040

0.035

0.030

0.025

0.020

0.015

0.010

0.005

000°0

D

3

2

8

8

9

6

8

8

？0

7

2

a

3

了

4

6

7

3

b

6

3

4

q

8

5

9

G6

3

3

h

3

S

8

6

8

q

8/

の

8

2

39

7

3

3

9

9

2

3

8

8

3

9

3

6

2

0

8

3

3

2

d

8

2

Adapting pretrained model to new data. Given a pretrained diffusion model over some original dataset D pretrain, we fine-tune the pretrained model for generating data that resemble D new. To cast this problem into (P-KL), we let the data distribution be D new, i.e., q ( x 0: T ) = q new ( x 0: T ) and the constrained distribution be the pre-trained model, i.e., q i ( x 0: T ) = p θ pre ( x 0: T ) . In our experiments, we pretrain a diffusion model on a 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 as 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.

6 Conclusion

We have presented a constrained optimization framework for training diffusion models under distribution constraints. We have developed a Lagrangian-based dual algorithm to train such constrained

frequencies

classes

0.30

0.25

0.20

0.15

0.10

0.05

0.00

male

female

male

female

frequencies

classes

0.10

0.08

0.06

0.04

0.02

0.00

0123456789

0123456789

frequencies

classes

0.14

0.12

0.10

0.08

0.06

0.04

0.02

0.00

0

3

952090u

7707224

1777

24

３7乡3么

q

48

4

h

8024

7℃80

76374

63

/0Y~

4

。

6

@LBhhn

y

3

3

S

工

4

6

7

9

G

5

92368

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 our distribution constraints to 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.

Acknowledgments

We thank reviewers and program chairs for providing helpful comments.

References

A. Bansal, H.-M. Chu, A. Schwarzschild, S. Sengupta, M. Goldblum, J. Geiping, and T. Goldstein. Universal guidance for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 843-852, 2023.

J. Benton, V. De 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.

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.

D. P. Bertsekas. Nonlinear programming . Athena Scientific, 2016.

K. Black, M. Janner, Y. Du, I. Kostrikov, and S. Levine. Training diffusion models with reinforcement learning. In Proceedings of the International Conference on Learning Representations , 2024.

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.

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.

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 on Knowledge and Data Engineering , 2024.

L. 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.

L. F. Chamon, S. Paternain, M. Calvo-Fullana, and A. Ribeiro. Constrained learning with non-convex losses. IEEE Transactions on Information Theory , 69(3):1739-1760, 2022.

H. Chen, H. Lee, and J. Lu. Improved analysis of score-based generative modeling: Userfriendly bounds under minimal smoothness assumptions. In Proceedings of the International Conference on Machine Learning , pages 4735-4763, 2023.

J. Chen, Z. Shao, X. Zheng, K. Zhang, and J. Yin. Integrating aesthetics and efficiency: AIdriven diffusion models for visually pleasing interior design generation. Scientific Reports , 14(1):3496, 2024.

J. Chen, R. Zhang, Y. Zhou, and C. Chen. Towards aligned layout generation via diffusion model with aesthetic constraints. arXiv preprint arXiv:2402.04754 , 2024.

M. Chen, S. Mei, J. Fan, and M. Wang. An overview of diffusion models: Applications, guided generation, statistical rates and optimization. arXiv preprint arXiv:2404.07771 , 2024.

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.

K. Choi, A. Grover, T. Singh, R. Shu, and S. Ermon. Fair generative modeling via weak supervision. In Proceedings of the International Conference on Machine Learning , pages 1887-1898, 2020.

J. K. Christopher, S. Baek, and F. Fioretto. Projected generative diffusion models for constraint satisfaction. arXiv preprint arXiv:2402.03559 , 2024.

K. Clark, P. Vicol, K. Swersky, and D. J. Fleet. Directly fine-tuning diffusion models on differentiable rewards. In Proceedings of the International Conference on Learning Representations , 2024.

F.-A. Croitoru, V. Hondru, R. T. Ionescu, and M. Shah. Diffusion models in vision: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2023.

P. Dhariwal and A. Nichol. Diffusion models beat gans on image synthesis. In Proceedings of the Advances in Neural Information Processing Systems , volume 34, pages 8780-8794, 2021.

Y. Du, C. Durkan, R. Strudel, J. B. Tenenbaum, S. 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-8510, 2023.

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.

Y. Fan and K. Lee. Optimizing DDPM sampling with shortcut fine-tuning. In Proceedings of the International Conference on Machine Learning , pages 9623-9639, 2023.

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 Proceedings of the Advances in Neural Information Processing Systems , volume 36, 2023.

B. T. Feng, R. Baptista, and K. L. Bouman. Neural approximate mirror maps for constrained diffusion models. arXiv preprint arXiv:2406.12816 , 2024.

N. Fishman, L. Klarner, V. De Bortoli, E. Mathieu, and M. J. Hutchinson. Diffusion models for constrained domains. Transactions on Machine Learning Research , 2024.

N. Fishman, L. 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.

F. Friedrich, M. Brack, L. Struppek, D. Hintersdorf, P. Schramowski, S. Luccioni, and K. Kersting. Fair diffusion: Instructing text-to-image generation models on fairness. arXiv preprint arXiv:2302.10893 , 2023.

G. Giannone, A. Srivastava, O. Winther, and F. Ahmed. Aligning optimization trajectories with diffusion models for constrained design generation. In Proceedings of the Advances in Neural Information Processing Systems , volume 36, 2023.

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: //github.com/huggingface/accelerate , 2022.

Y. Han, M. Razaviyayn, and R. Xu. Neural network-based score estimation in diffusion models: Optimization and generalization. arXiv preprint arXiv:2401.15604 , 2024.

Y. Hao, Z. Chi, L. Dong, and F. Wei. Optimizing prompts for text-to-image generation. In Proceedings of the Advances in Neural Information Processing Systems , volume 36, 2023.

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.

J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. In Proceedings of the Advances in Neural Information Processing Systems , volume 33, pages 6840-6851, 2020.

J. Ho and T. Salimans. Classifier-free diffusion guidance. In NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications , 2021.

I. Hounie, A. Ribeiro, and L. F. Chamon. Resilient constrained learning. In Proceedings of the Advances in Neural Information Processing Systems , volume 36, 2023.

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.

D. Z. Huang, J. Huang, and Z. Lin. Convergence analysis of probability flow ODE for scorebased generative models. arXiv preprint arXiv:2404.09730 , 2024.

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.

A. Kazerouni, E. K. Aghdam, M. Heidari, R. Azad, M. Fayyaz, I. Hacihaliloglu, and D. Merhof. Diffusion models in medical imaging: A comprehensive survey. Medical Image Analysis , page 102846, 2023.

Y. Kim, B. Na, M. Park, J. Jang, D. Kim, W. Kang, and I.-C. Moon. Training unbiased diffusion models from biased dataset. In Proceedings of the International Conference on Learning Representations , 2024.

D. Kingma and R. Gao. Understanding diffusion objectives as the ELBO with simple data augmentation. In Proceedings of the Advances in Neural Information Processing Systems , volume 36, 2023.

Z. Kong, W. Ping, J. Huang, K. Zhao, and B. Catanzaro. DiffWave: A versatile diffusion model for audio synthesis. In Proceedings of the International Conference on Learning Representations , 2020.

Y. LeCun, C. Cortes, and C. J. Burges. The MNIST database of handwritten digits. http: //yann.lecun.com/exdb/mnist/ .

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 arXiv:2302.12192 , 2023.

G. Li, Y. Huang, T. Efimov, Y. Wei, Y. Chi, and Y. Chen. Accelerating convergence of scorebased diffusion models, provably. arXiv preprint arXiv:2403.03852 , 2024.

G. Li, Y. Wei, Y. Chen, and Y. Chi. Towards faster non-asymptotic convergence for diffusionbased generative models. In Proceedings of the International Conference on Learning Representations , 2024.

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.

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 on Computer Vision , pages 423-439, 2022.

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.

Z. Liu, P. Luo, X. Wang, and X. Tang. Large-scale celebfaces attributes (celeba) dataset. https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html .

I. Loshchilov and F. Hutter. Decoupled weight decay regularization, 2019.

A. Lou and S. Ermon. Reflected diffusion models. In Proceedings of the International Conference on Machine Learning , pages 22675-22701, 2023.

A. S. Luccioni, C. Akiki, M. Mitchell, and Y. Jernite. Stable bias: Analyzing societal representations in diffusion models. arXiv preprint arXiv:2303.11408 , 2023.

C. Luo. Understanding diffusion models: A unified perspective. arXiv preprint arXiv:2208.11970 , 2022.

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.11420 , 2023.

R. Naik and B. Nushi. Social biases through the text-to-image generation lens. In Proceedings of the AAAI/ACM Conference on AI, Ethics, and Society , pages 786-808, 2023.

G. Parmar, R. Zhang, and J.-Y. Zhu. On aliased resizing and surprising subtleties in gan evaluation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11410-11420, 2022.

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. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. PyTorch: An imperative style, high-performance deep learning library, 2019.

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.

D. J. Rezende and F. Viola. Generalized ELBO with constrained optimization, GECO. In Workshop on Bayesian Deep Learning, NeurIPS , 2018.

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 on computer vision and pattern recognition , pages 10684-10695, 2022.

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 (IJCV) , 115(3):211-252, 2015.

C. Saharia, W. Chan, S. Saxena, L. Li, J. Whang, E. L. Denton, K. Ghasemipour, R. Gontijo Lopes, B. Karagol Ayan, T. Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. In Proceedings of the Advances in Neural Information Processing Systems , volume 35, pages 36479-36494, 2022.

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.

J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Proceedings of the International Conference on Machine Learning , pages 2256-2265, 2015.

V. Solo and X. Kong. Adaptive signal processing algorithms: stability and performance . Prentice-Hall, Inc., 1994.

J. Song, C. Meng, and S. Ermon. Denoising diffusion implicit models. In Proceedings of the International Conference on Learning Representations , 2021.

C. Szegedy, W. Liu, Y. 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.

W. Tang. Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond. arXiv preprint arXiv:2403.06279 , 2024.

M. Uehara, Y. Zhao, K. Black, E. Hajiramezanali, G. Scalia, N. L. Diamant, A. M. Tseng, T. Biancalani, and S. Levine. Fine-tuning of continuous-time diffusion models as entropyregularized control. arXiv preprint arXiv:2402.15194 , 2024.

M. Uehara, Y. Zhao, E. Hajiramezanali, G. Scalia, G. Eraslan, A. Lal, S. Levine, and T. Biancalani. Bridging model-based optimization and generative modeling via conservative fine-tuning of diffusion models. arXiv preprint arXiv:2405.19673 , 2024.

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://github.com/huggingface/diffusers , 2022.

J. L. Watson, D. Juergens, N. R. Bennett, B. L. Trippe, J. Yim, H. E. 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.

T. Weiss, E. Mayo Yanes, S. Chakraborty, L. Cosmo, A. M. Bronstein, and R. Gershoni-Poranne. Guided diffusion for inverse molecular design. Nature Computational Science , 3(10):873-882, 2023.

X. Wu, K. Sun, F. Zhu, R. Zhao, and H. Li. Human preference score: Better aligning textto-image models with human preference. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 2096-2105, 2023.

J. Xu, X. 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.

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 Comput. Surv. , 56(4), nov 2023.

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.

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.

C. Zhang, C. Zhang, S. Zheng, M. Zhang, M. Qamar, S.-H. Bae, and I. S. Kweon. Audio diffusion model for speech synthesis: A survey on text to speech and speech enhancement in generative AI. arXiv preprint arXiv:2303.13336 , 2023.

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.

Supplementary Materials for 'Constrained Diffusion Models via Dual Training'

A Details on ELBO

Recall the evidence lower bound (ELBO),

we can utilize conditionals to expand it into

︷︷ denoising matching term where the first term is the reconstruction likelihood of the original data given the first latent x 1 , 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.1, it is known that the reconstruction likelihood and final latent mismatch are negligible, and thus the approximation in (4) is almost exact, which is our focal setting of this paper.

We next focus on one summand of the denoising matching term,

Application of the reparametrization trick leads to x t = √ α t x t -1 + √ 1 -α t ϵ t -1 , where ϵ t -1 ∼ N (0 , I ) is a white noise sample. Using Bayes rule, we can express q ( x t -1 | x t , x 0 ) as a Guassian distribution

where µ q ( x t ) = 1 √ α t x t + 1 -α t √ α t ∇ log q ( x t ) is the mean and σ 2 q ( t ) = (1 -α t )(1 -¯ α t -1 ) 1 -¯ α t is the variance.

To stay close to the ground-truth backward conditional q ( x t -1 | x t , x 0 ) as much as possible, we take p ( x t -1 | x t ) to be the same as q ( x t -1 | x t , x 0 ) except replacing ∇ log p ( x t ) by ̂ s ( x t , t ) and σ 2 q ( t ) by σ 2 p ( t ) ,

where the second equality is due to the KL Divergence between two Gaussians. Since σ 2 p ( t ) and σ 2 q ( t ) are constants, the variance mismatch term is irrelevant to optimization. Denote ω t := (1 -α t ) 2 2 σ 2 q ( t ) α t and ¯ ω := ∑ T t =2 ω t . We can define a discrete distribution over the set { 2 , . . . , T } as

p ω ( t ) := ω t ¯ ω . Also denote v := ∑ T t =2 1 2 ( d log σ 2 p ( t ) σ 2 q ( t ) -d + d σ 2 p ( t ) σ 2 q ( t ) ) . Hence, the ELBO maximization: maximize p E ( p ; q ) , is equivalent to the quadratic loss minimization,

where E x 0 , t, x t is an expectation over the data distribution q ( x 0 ) , the discrete distribution p ω ( t ) from 2 to T , and a forward process q ( x t | x 0 ) given the data sample x 0 . Since shifting an objective function by a constant and scaling an objective function by a constant don't change the solution of an optimization problem, we omit constants v and ¯ ω for brevity, and only emphasize them whenever it is necessary. Hence, the ELBO maximization equals the quadratic loss minimization,

̂

up to some scaling and shifting constants, where E x 0 ,t,x t is an expectation over the data distribution q ( x 0 ) , the discrete distribution p ω ( t ) from 2 to T , and a forward process q ( x t | x 0 ) given the data sample x 0 . In practice, however, we have to parametrize the estimator ̂ s ( x t , t ) as ̂ s θ ( x t , t ) with parameter θ ∈ Θ ,

which is our focal objective of generative modeling. A parametrized representation of p ( x t -1 | x t ) associated with ̂ s θ ( x t , t ) is denoted by p θ ( x t -1 | x t ) and the backward process has a parametrized joint distribution p θ ( x 0: T ) . We remark that the above prediction problem can be reformulated as data or noise predictions [55], with our results directly transferrable to these formulations.

B Proofs

We provide proofs of all lemmas and theorems in the main paper.

B.1 Proof of Lemma 1

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 are the same. Since q ( x 0: T ) ∈ 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:

It is clear that the likelihood maximization problem maximize p E q ( x 0 ) [ log p ( x 0 ) ] is equivalent to minimize p D KL ( q ( x 0 ) ∥ p ( x 0 )) . Therefore, any distribution p ⋆ ( x 0: T ) whose marginal satisfies p ⋆ ( x 0 ) = q ( x 0 ) , will be a solution of the likelihood maximization problem. This includes the solution of the KL divergence minimization and ELBO maximization which is p ⋆ = q . Therefore,

which concludes the proof.

B.2 Proof of Lemma 2

Proof. It is straightforward to check the zero duality gap in convex optimization; see e.g., [4, Proposition 5.3.2]. Furthermore, for a convex optimization problem, an optimal dual variable λ ⋆ that maximizes the dual function is a geometric multiplier. Hence, ( p ⋆ , λ ⋆ ) is an optimal primal-dual pair of the convex optimization problem.

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 a geometric multiplier. Thus, Problem (U-KL) reduces to an unconstrained problem,

where the objective function results from plugging an optimal dual variable λ ⋆ into Lagrangian L ( p, λ ) .

By the definition of Lagrangian,

By taking λ = λ ⋆ , Problem (12) is equivalent to

From the definition of ELBO, we know that

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 q ( λ ⋆ ) mix . Thus, Problem (13) is equivalent to

which, together with Lemma 1, completes the proof.

B.4 Proof of Theorem 2 and Feasibility Criterion

We start with the proof of Theorem 2.

Proof. Similar to the proof of Theorem 1, we begin with the Lagrangian of Problem (U-ELBO),

where from (14) onwards we use notation: λ 0 = 1 , ¯ b 0 = 0 , λ = [ λ 0 , . . . , λ m ] T , ¯ b = [ ¯ b 0 , . . . , ¯ b m ] T , and use q 0 to represent q , which will be used in the rest of proof. To formulate the dual problem, we check the minimum of the Lagrangian,

where the only term that depends on p is the ELBO. Recall that:

Since the minimum value of the KL divergence is zero (attained when p = q ), the minimum of -E ( p ; q ) is likewise attained when p = q . Thus, it is straightforward that the minimum is equal to the entropy of the distribution q , denoted by h ( q ) := -E q ( x 0 ) [ log q ( x 0 ) ] . With this in mind, from (15) we have

Thus, the dual problem reads

We first reformulate the entropy of the mixture distribution q ( λ ) ,

where going from (16) to (17) we utilize the assumption that the distributions { q i } i m =0 have disjoint supports; see Remark 1 on when this is the case.

Now, we can compute the gradient of the dual function over λ i , i = 1 , . . . , m ,

Setting the gradient be zeros allows us to find the optimal dual variables λ ⋆ ,

Hence,

We clarify that in (18), λ ⋆ = [ λ ⋆ 0 , . . . , λ ⋆ m ] T with its first element being λ ⋆ 0 = 1 . Finally, if we return back to notation λ ⋆ = [ λ ⋆ 1 , . . . , λ ⋆ m ] T , then,

which completes the proof.

Remark 1. We remark on the assumption of the distributions { q i } i m =0 having disjoint supports. In the setting of adapting model to new data in Section 3.2, this is a reasonable assumption. Since we often finetune a pre-trained diffusion model on new data not seen in the original pre-training 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 i } i m =1 and the objective distribution q 0 usually are not disjoint. However, since the distributions { q i } i m =1 often are often restrictions of q 0 to subsets of the support of q 0 , 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 to cases where the distributions are neither disjoint nor restrictions of the objective distributions, is challenging and has been left to future work.

To prove a feasibility criterion, we first show that the dual function is finite in Lemma 5.

Lemma 5 (Boundedness of the optimal dual function) . Let the differential entropy h i of each distribution q i be finite. Then, the optimal value of the dual function D ⋆ := max λ ≥ 0 g ( λ ) is finite if and only if

Proof. ( ⇐ ) From ∑ m i =1 e h i -¯ b i < 1 , the otpimal dual variable λ ⋆ given by (18) is finite. Thus, the optimal value of the dual function g ( λ ⋆ ) becomes

and λ ⋆ i = e h i -¯ b i > 0 , and also { h i } i m =1 are all finite. Therefore, D ⋆ is finite.

( ⇒ ) We prove it by contradiction. Assume ∑ m i =1 e h i -¯ b i = e δ ≥ 1 for some δ ≥ 0 . For any λ ≥ 0 , there exists a direction in which g ( λ ) increases, i.e.,

To see (19) by contradiction, we check that

which contradicts our assumption that ∑ m i =1 e h i -¯ b i = e δ . By contradiction, we have (19). Furthermore, (19) implies that g ( λ ) is unbounded above, which contradicts the finiteness of D ⋆ . Therefore, we must have that ∑ m i =1 e h i -¯ b i < 1 .

Lemma 6 (Feasibility criterion) . Let the differential entropy h i of each distribution q i be finite. Suppose that there exists a feasible solution to Problem (U-ELBO) such that its objective function is bounded from below. Then, Problem (U-ELBO) is feasible if and only if

Proof. ( ⇒ ) Since the primal problem is feasible, the optimal objective function F ⋆ is bounded from below and it is attained at a feasible point. For the sake of contradiction, we assume ∑ m i =1 e h i -¯ b i ≥ 1 , which is equivalent to D ⋆ = ∞ according to Lemma 5. However, this violates weak duality, i.e., D ⋆ ≤ F ⋆ . By contradiction, we must have ∑ m i =1 e h i -¯ b i < 1 .

(

⇐

) We consider a set

A

,

The set A is convex since it is the intersection of m +1 epigraphs of convex functions. We also introduce another convex set B ,

We utilise proof by contradiction. Assume that the primal problem is infeasible. Then there doesn't exist any p ∈ P such that -E ( p, q i ) -¯ b i ≤ 0 for all i = 1 , . . . , m . Hence, A and B are two disjoint convex sets. From the separating hyperplane theorem, there exists a hyperplane that separates them, i.e., ∃ v ∈ R m +1 and c ∈ R ,

Let v = [ λ 1 , . . . λ m , γ ] T . Then, (21) reduces to

This is because γt ≤ c for any t ∈ R . Note that γ = 0 means the separating hyperplane is vertical, i.e., being parallel to the t -axis. Furthermore, from (20) we can write

The above is true because the set of values that each u i can take in A is unbounded above. Thus, since λ T u ≥ c , necessarily every λ i has to be non-negative. Now, we consider

which shows that D ∗ = ∞ . This contradicts D ⋆ being finite ( D ⋆ is finite due to ∑ m i =1 e h i -¯ b i < 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 of the convergence theory of DDPM [47, Theorem 3]. We next check all assumptions of [47, Theorem 3]. It is easy to see that we can cast q ( λ ) mix as a target distribution

of a diffusion model. By the definition,

Thus, the partial minimization of ¯ L ( θ, λ ) over θ is equivalent to a weighted EBLO minimization,

where we normalize the weighted ELBO objective by introducing a mixed data distribution q ( λ ) mix . We note that ¯ p ⋆ ( λ ) is also a minimizer of Problem (22).

On the other hand, using Problem (P-LOSS), we can rewrite Problem (22) as

which is equivalent to the score matching objective in DDPM. Therefore, the score matching assumption in [47, Assumption 1] is satisfied with the error bound ε 2 score from Assumption 3. Viewing Assumption 2, and using appropriate stepsize and variance, all assumptions in [47, Theorem 3] are satisfied. Therefore, application of [47, Theorem 3] completes the proof.

B.6 Proof of Lemma 4

Lemma 7. The TV distance between two mixture data distributions q ( λ ) mix , q ( λ ⋆ ) mix is bounded by

or, equivalently,

Proof. By the definition,

where the first inequality is due to (1 + λ ⊤ 1)(1 + ( λ ⋆ ) ⊤ 1) ≥ 1 , and we use triangle inequality in the second inequality.

We recall the Lagrangians for Problems (U-LOSS) and (P-LOSS),

and their associated dual functions,

For brevity, we use shorthand E q and E q i for E q ( x 0 ) , t, x t and E q i ( x 0 ) , t, x t , respectively.

Lemma 8 (Parametrization gap) . Let Assumption 4 hold. Then, 0 ≤ ¯ g s ( λ ) -g s ( λ ) ≤ 4 R (1+ ∥ λ ∥ 1 ) ν for any λ ≥ 0 .

Proof. Let the partial minimizer of L s ( ̂ s, λ ) over ̂ s be ̂ s ⋆ ( λ ) := argmin ̂ s L s ( ̂ s, λ ) for any λ ≥ 0 . For any λ ≥ 0 , there exists ˜ θ ∈ Θ such that ∥ ∥ ̂ s ⋆ ( λ ) -̂ s ˜ θ ∥ ∥ L 2 ≤ ν for any λ ≥ 0 , according to Assumption 4. Thus,

where the first inequality is due to that the quadratic function is locally Lipschitz continuous with parameter 4 R , and the second inequality is because that there exists ˜ θ ∈ Θ such that ∥ ∥ ̂ x ⋆ ( λ ) -̂ x ˜ θ ∥ ∥ L 2 ≤ ν for any λ ≥ 0 , according to Assumption 4.

By the definition ̂ s ⋆ θ ( λ ) ∈ argmin θ ∈ Θ ¯ L s ( ̂ s θ , λ ) ,

Therefore,

Lemma 9 (Differentiability) . The dual function g s ( λ ) is differentiable with gradient ∇ λ L s ( ̂ s ⋆ ( λ ) , λ ) .

Proof. For any λ ≥ 0 , the Lagrangian L s ( ̂ s, λ ) is strongly convex in function ̂ s ∈ S . Since S is convex and compact, any partial minimizer ̂ s ⋆ ( λ ) is unique. By Danskin's theorem [4], g s ( λ ) is differentiable and its gradient is the gradient of L s ( ̂ s, λ ) over λ at ̂ s = ̂ s ⋆ ( λ ) .

Lemma 10 (Convexity) . The dual function g s ( λ ) is µ -strongly concave in λ ∈ H , where

Proof. For any λ 1 , λ 2 ∈ H , we denote ̂ s ⋆ 1 := ̂ s ⋆ ( λ 1 ) and ̂ s ⋆ 2 := ̂ s ⋆ ( λ 2 ) , which are unique partial minimizers of the Lagrangians L s ( ̂ s, λ 1 ) and L s ( ̂ s, λ 2 ) . Denote ℓ 0 ( ̂ s ) := E q [ ∥ ̂ s ( x t , t ) -∇ log q ( x t ) ∥ 2 ] , ℓ i ( ̂ s ) := E q i [ ∥ ̂ s ( x t , t ) -∇ log q ( x t ) ∥ 2 ] for i = 1 , . . . , m , and ℓ ( ̂ s ) := [ ℓ 1 ( ̂ s ) , . . . , ℓ m ( ̂ s ) ] ⊤ . By the convexity of ℓ i ,

If we multiply the first inequality above by λ 2 ,i ≥ 0 and the second inequality above by λ 1 ,i ≥ 0 , and add them up from both sides, then,

We apply Lemma 9 to the LHS of the inequality above,

On the other hand, by the optimality of ̂ s ⋆ 1 and ̂ s ⋆ 2 ,

which allows us to simplify the right-hand side of (23) and obtain

where the last inequality results from the strong convexity of quadratic functionals.

By the smoothness of quadratic functionals with parameter 1 ,

where the equality is due to the optimality condition (24) and the last inequality is due to triangle inequality. By Assumption 5,

We also notice that

where the first inequality is due to triangle inequality and the second inequality is due to the smoothness of quadratic functionals. Hence,

or, equivalently,

Therefore, (25) becomes

which completes the proof by choosing the smallest modulus over λ 1 ∈ H .

Proof. By Lemmas 9 and 10, for any λ ∈ H ,

Thus, if we choose λ = ¯ λ ⋆ , then

Optimality of ( ̂ s ⋆ ( λ ⋆ ) , λ ⋆ ) leads to the complementary slackness,

and the feasibility,

Therefore,

According to Lemma 8, ¯ g s ( ¯ λ ⋆ ) -4 R ( 1 + ∥ ∥ ¯ λ ⋆ ∥ ∥ 1 ) ν ≤ g s ( ¯ λ ⋆ ) . Hence,

Thus,

where the last inequality is due to that g s ( λ ) ≤ ¯ g s ( λ ) for any λ ≥ 0 , and the optimality of ¯ λ ⋆ ,

B.7 Proof of Theorem 3

Proof. By the triangle inequality for TV distance,

where the second inequality is due to Lemma 7 and Lemma 3, and the last inequality is due to ∥ λ ∥ 1 ≤ √ m ∥ λ ∥ and Lemma 4.

B.8 Proof of Theorem 4

Lemma 11. For a stochastic variant of Algorithm 1 in Section 4.3, we have

Proof. For brevity, we let the stochastic gradient be ̂ f ( h ) = [ ̂ f 1 ( h ) , . . . , ̂ f m ( h )] with

By the definition of λ ( h +1) ,

where the inequality is due to the non-expansiveness of projection. Application of the conditional expectation over both sides of the inequality above yields,

which gives our desired result when we use the fact that E [ ̂ f ( h ) | λ ( h )] is an approximate descent direction of the dual function ¯ g ,

Lemma 12. In the stochastic variant of Algorithm 1 in Section 4.3, the maximum prarametrized dual function in history up to step h satisfies

Proof. The proof is based on the supermartingale convergence theorem [67, Theorem E7.4]. We introduce two processes,

where α ( h ) measures the gap between λ ( h ) and ¯ λ ⋆ when the optimality gap ¯ D ⋆ -¯ g best ( h ) is below the threshold, and β ( h ) measures the gap bewteen ¯ D ⋆ and ¯ g ( λ ( h )) (up to some optimization errors) when the optimality gap ¯ D ⋆ -¯ g best ( h ) is below the threshold. Clearly, α ( h ) is non-negative. Also, β ( h ) is non-negative due to

Let F h be the σ -algebra generated by sequences: α ( h ′ ) , β ( h ′ ) , and λ ( h ′ ) for h ′ ≤ h . Thus, {F h } h ≥ 1 is a natural filtration. We notice that α ( h +1) and β ( h +1) are determined by λ ( h ) in each step. Hence,

We first show that

A simple case is when α ( h ) = 0 ,

There are two situations for α ( h ) = 0 . First, if ¯ D ⋆ -¯ g best ( h ) ≤ ηS 2 2 + ε 2 approx , then α ( h ) = β ( h ) = 0 . Due to ¯ g best ( h +1) ≥ ¯ g best ( h ) , we have β ( h ) = 0 and ¯ D ⋆ -¯ g best ( h +1) ≤ ηS 2 2 + ε 2 approx . Thus, α ( h + 1) = 0 , and (26) holds. Second, if λ ( h ) = ¯ λ ⋆ , but ¯ D ⋆ -¯ g best ( h ) > ηS 2 2 + ε 2 approx , then ¯ D ⋆ = ¯ g ( λ ( h )) . Hence, β ( h ) < 0 , which contradicts the non-negativeness of β ( h ) . Therefore, ¯ D ⋆ -¯ g best ( h ) ≤ ηS 2 2 + ε 2 approx has to hold, which is the first situation.

We next show (26) when α ( h ) > 0 .

where the last inequality is due to Lemma 11, and the last equality is due to ¯ D ⋆ -¯ g best ( h ) > ηS 2 2 + ε 2 approx . Therefore, (26) holds.

Finally, application of the supermartingale convergence theorem [67, Theorem E7.4] to the processes α ( h ) and β ( h ) for h ≥ 1 concludes that β ( t ) is almost surely summable,

Denote the step when ¯ g best ( h ) achieves ¯ D ⋆ -( ηS 2 2 + ε 2 approx ) by h best, and the associated dual variable be ¯ λ best = λ ( h best ) .

Lemma 13. For a stochastic variant of Algorithm 1 in Section 4.3, we have

Proof. We denote the segment between ¯ λ best and λ ⋆ by B . By Lemma 10, the dual function g ( λ ) is strongly concave on B with parameter µ . Thus,

where the second inequality is due to Lemma 8. We note that ¯ g ( ¯ λ best ) = ¯ g ( λ ( h best )) and ¯ D ⋆ ≥ ¯ g ( ¯ λ best ) , and the third inequality is due to Lemma 12.

Proof. By the triangle inequality for TV distance,

where the second inequality is due to Lemma 7 and Lemma 3, and the last inequality is due to ∥ λ ∥ 1 ≤ √ m ∥ λ ∥ and Lemma 13.

This means that either

C Experimental details

We provide implementation details of our computational experiments in Section 5. The source code is available here. 2

Algorithm details: We train our constrained diffusion models by replacing the exact primal minimization step in Algorithm 1 with N steps of gradient descent with the Lagrangian as a 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 no bearing on our main results. Algorithm 2 depicts our practical implementation of Algorithm 1 .

Algorithm 2 Practical Implementation of Algorithm 1

1: Input : total diffusion steps T , diffusion parameter α t , total dual iterations H , number of primal descent steps per dual update N , dual step size η d , primal step size η p , initial model parameters θ (0) .

2: Initialize : λ (1) = 0 .

3: for h = 1 , · · · , H do

4: for n = 1 , · · · , N do

5: θ 1 = θ ( h -1) .

7: θ ( h ) = θ N +1 .

8: end for

9: Update the dual variable

10: end for

In Algorithm 2, the unbiased estimate of the noise prediction loss is evaluated via

where { x ( i ) 0 } B i =1 is a randomly chosen batch of samples from q , { t ( i ) } B i =1 are time steps randomly sampled from the interval [2 , T ] , and x t ( i ) is the noisy version of x ( i ) 0 at time step t ( i ) . Then, the noise ϵ 0 is sampled from the standard Gaussian, and x t is derived from x t = √ ¯ α t x t + √ 1 -¯ α t ϵ 0 .

We remark an important implementation detail in the fine-tuning experiment. In the fine-tuning constraints, we have to evaluate the KL divergence

where p θ pre ( x 0: T ) is the joint distribution of the samples and latents generated by the backward process of the pre-trained model. The KL divergence constraint in (27) further reduces to

To estimate the expectation in (28), we need to sample latents x t from the backward distribution p θ pre ( x 0: T ) . In practice, this is computationally inefficient, since it requires running inference each time one wants to sample a latent. This is why we implement this with sampling x t as random Gaussian noise instead. This still ensures that the predictions of the new model p θ don't differ too much from the pre-trained distribution p θ pre while making sampling batches much faster.

Resilient constrained learning. The choice of the constraint thresholds { b i } i m =1 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

adjust the thresholds { b i } i m =1 during training. In essence, the resilient constrained learning adds a constraint relaxation cost to the loss and relaxes the thresholds by updating them through gradient descent each time we update the dual variable. It can further be shown theoretically that an equivalent formulation of resilience is achieved by adding a quadratic regularizer of the dual variable into the loss objective and setting the constraint thresholds to be zero, i.e., ˜ b i = 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 9 of Algorithm 3).

Algorithm 3 Resilient Constrained Diffusion Models via Dual Training

10:

1: Input : total diffusion steps T , diffusion parameter α t , total dual iterations H , number of primal descent steps per dual update N , dual step size η d , primal step size η p , constraint step size η c , initial model parameters θ (0) , constraint relaxation cost γ . 2: Initialize : λ (1) = 0 . 3: for h = 1 , · · · , H do 4: for n = 1 , · · · , N do 5: θ 1 = θ ( h -1) . 6: θ n +1 = θ n -η p ∇ θ ( ̂ E x 0 ∼ q, t, x t [ ∥ ̂ ϵ θ n ( x t , t ) -ϵ 0 ∥ 2 ] + m ∑ i =1 λ i ̂ E x 0 ∼ q i , t, x t [ ∥ ̂ ϵ θ n ( x t , t ) -ϵ 0 ∥ 2 ] ) . 7: θ ( h ) = θ N +1 . 8: end for 9: Update the dual variable λ i ( h +1) = [ λ i ( h ) + η d ( ̂ E x 0 ∼ q i , t, x t [ ∥ ∥ ̂ ϵ θ ( h ) ( x t , t ) -ϵ 0 ∥ ∥ 2 ] -˜ b i ( h ) -2 γλ i ( h ) )] + 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 a positional embedding of the time to the input image. The parameters of the model are summarized in Table 1. 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 to 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 to 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 as 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 as we update model parameters doesn'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 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). 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 η , we refer to Theorem 8 in the paper, showing a smaller error bound for smaller η while slowing convergence. In practice, as long as η ≤ 1 , we observed that the model converges to similar results reliably.

Efficiency of constrained diffusion: 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.

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, 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.

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.

The 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 (see discussion on batch sizes in hyperparameter sensitivity section).

FID scores. As a quantitative means of evaluating our constrained diffusion models, we use the FID (Frechet Inception Distance) as a metric to gauge the quality of the samples generated by diffusion models. The FID score was first introduced in [33] in form of

where m 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, m w and C w 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 is an imbalance in the original dataset, we consider a 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 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.

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

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 libraries for training our constrained diffusion models, and Adam with decoupled weight decay [52] as an optimizer. 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 models, 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: //huggingface.co/farleyknight-org-username/vit-base-mnist ). Aclassifier for images of male/female faces with %98 . 6 validation accuracy ( https://huggingface.co/cledoux42/ GenderNew_v002 ). For classifying the image-net data, a zero-shot classifier based on a CLIP model ( https://huggingface.co/openai/clip-vit-base-patch32 ) was used. The Autoencoder for the latent diffusion model was the stable diffusion V AE with KL regularization found on ( https://huggingface.co/stabilityai/sd-vae-ft-mse ).

NeurIPS Paper Checklist

1. 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 guarantees 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 5.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We mention two potential limitations of this work as 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 a complete (and correct) proof?

Answer: [Yes]

Justification: We provide all assumptions, lemmas, and theorems in Sections 3 and 4, 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We summarize our experimental results in Section 5, and provide implementation 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 instructions to faithfully reproduce the main experimental results, as 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 are open source machine learning datasets that are accessible online.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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]

Justification: While the histogram plots that showcase our main results do not include error bars, we do provide Frechet Distance metrics (that are a measure of statistical distance between sample distributions) which emphasize our results.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) 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 NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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

Question: Does the paper describe 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, or scraped datasets)?

Answer: [NA]

Justification: We release no models and the supplemental code for training image generation model, trains the model on MNIST and Celeb-A datasets neither of which have potential for misuse.

12. Licenses for existing assets

Question: 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?

Answer: [Yes]

Justification: The libraries and assets used have been noted and credited in Appendix C.

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 as details about compensation (if any)?

Answer: [NA]

Justification: [NA]

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

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]

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n1. **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.\n\nWEAKNESSES:\n**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.\n\nQUESTIONS:\nPlease see weaknesses.\n\nLIMITATIONS:\n1. **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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nThe 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", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThere 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.\n\nQUESTIONS:\nSee weakness\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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)\n\nQUESTIONS:\n- 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].\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}