Composition and Alignment of Diffusion Models using Constrained Learning
Authors: Shervin Khalafi
Ignacio Hounie
Dongsheng Ding
Alejandro Ribeiro
Venue: NeurIPS 2025 poster ()
Keywords: Diffusion Models
Constrained Optimization
Alignment
Composition
Generative Models
Primary area: deep_learning

Abstract: Diffusion models have become prevalent in generative modeling due to their ability to sample from complex distributions. To improve the quality of generated samples and their compliance with user requirements, two commonly used methods are: (i) Alignment, which involves finetuning a diffusion model to align it with a reward; and (ii) Composition, which combines several pretrained diffusion models together, each emphasizing a desirable attribute in the generated outputs. However, trade-offs often arise when optimizing for multiple rewards or combining multiple models, as they can often represent competing properties. Existing methods cannot guarantee that the resulting model faithfully generates samples with all the desired properties. To address this gap, we propose a constrained optimization framework that unifies alignment and composition of diffusion models by enforcing that the aligned model satisfies reward constraints and/or remains close to each pretrained model. We provide a theoretical characterization of the solutions to the constrained alignment and composition problems and develop a Lagrangian-based primal-dual training algorithm to approximate these solutions. Empirically, we demonstrate our proposed approach in image generation, applying it to alignment and composition, and show that our aligned or composed model satisfies constraints effectively. Our implementation can be found at: https://github.com/shervinkhalafi/constrained_comp_align.

Composition and Alignment of Diffusion Models using Constrained Learning
Shervin Khalafi University of Pennsylvania shervinkèseas upenn. edu
Ignacio Hounie University of Pennsylvania ihounie@seas: upenn. edu
Dongsheng Ding University of Tennessce, Knoxville dongshed@utk.. edu
Alejandro Ribeiro University of Pennsylvania aribeirodseas. upenn. edu
Abstract
Diffusion models have become prevalent generative modeling due to their ability sample from complex distributions. To improve the quality of generated samples and their compliance with user requirements, two commonly used methods are: (i) Alignment, which involves finetuning diffusion model align it with reward; and (ii) Composition, which combines several pretrained diffusion models together, each emphasizing desirable attribute in the generated outputs. However, trade-offs often arise when optimizing for multiple rewards or combining multiple models, as they often represent competing properties. Existing methods cannot guarantee that the resulting model faithfully generates samples with all the desired properties. To address this gap, we propose constrained optimization framework that unifies alignment and composition of diffusion models by enforcing that the aligned model satisfies reward constraints and/or remains close each pretrained model. We provide a theoretical characterization of the solutions to the constrained alignment and composition problems and develop a Lagrangian- -based primal-dual training al- gorithm to approximate these solutions. Empirically, we demonstrate our proposed approach image generation, applying it to alignment and composition, and show that our aligned or composed model satisfies constraints effectively. Our implemen- wttpr/ootttrreeeeeccooepeee_oeee cacpoodooooppppppaoeee aliggn
tation can be found at:
Introduction
Diffusion models have emerged as the tool of choice for generative modeling in a variety of settings [38, 50, 9], image generation being most prominent among them [37]. Users of these diffusion models would like to adapt them to their specific preferences, but this aspiration hindered by the often enormous cost and complexity of their training [48, 56]. For this reason, aligument and composition of what, in this context, become pretrained models, has become popular [29, 31].
Regardless of whether the goal is alignment composition, likely conflicting requirements, In alignment, we want to stay close to the pretrained model while deviating sufficiently pretrained models, our goal is to sample from the union or intersection of their distributions [14, 1]. The standard approach to balance these requirements involves the use of weighted averages. This can be a linear combination of score functions in composition [14, 1] or may involve loss given by linear combination of Kullback-Leibler (KL) divergence and reward [17] in the of alignment.
to balance what are most
affect some rewards of interest [17, 13]. In composition, given several
Conresponding authors.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

In practice, weight- based methods are often designed in an ad hoc manner, with the weights treated as tunable hyperparameters, which makes the approach notoriously difficult optimize and generalize.
In this work, we propose a unified view of alignment and composition via the lens of constrained learning [7, 6]. As their names indicate, constrained alignment and constrained composition problems balance conflicting requirements using constraints instead of weights. Learning with constraints and learning with weights are related problems indeed, we will train constrained diffusion models in their Lagrangian forms. Yet, they are also fundamentally different. In the constrained formulation, the hyperparameter tuning spaces are more interpretable (see Section 3), and in some cases such as the constrained composition formulation hyperparameter tuning can even be avoided entirely (see Section 4). These advantages are particularly evident in constrained problems. as discussed in Sections 3 and 4. We summarize our key contributions in three aspects below.
ii) Problem Formulation
For alignment, we formulate a reverse KL divergence--constrained distribution opti- mization problem that minimizes the reverse KL divergence to pretrained model, subject to expected reward constraints with user-specified thresholds.
For composition, we propose using KI. divergence constraints to ensure the closeness to each pretrained model. It is important to distinguish composition with reverse KL and forward KL constraints as they lead a weighted product or weighted mixture [22] of the individual distributions, respectively. In this work, we focus composition with reverse KL constraints, and discuss forward KL constraints in Appendix E.
Theoretical Analysis
In Section 3, we characterize the solution of the aligmment problem as the pretrained model distribution scaled by an exponential function of weighted sum of reward functions. In Section 4, we characterize the solution of the constrained optimization problem with reverse KI. divergence constraints as a tilted product of the individual distributions. We establish strong duality for both problems, which enables us to use a dual-based approach to develop primal-dual training algorithms for solving them.
We illustrate the distinction between the KL divergence between diffusion trajectories (path-wise), and the KL divergence between the final distributions (point-wise) in Section 2.2. We also propose new method to evaluate the point-wise KL divergence.
(iii) Empirical Results
For alignment, we demonstrate the difference between constrained and weighted align- ment through experiments in Section 5.1. The constrained approach scales naturally to finetuning with multiple rewards, eliminating the need for extensive hyperparameter searches to determine suitable weights. Moreover, specifying reward thresholds is often more intuitive than choosing regularization weights. Without constraints, however, the model can easily overfit to one or several rewards and diverge substantially from the pretrained model. In contrast, our method identifies the model closest to the pretrained one that still satisfies the desired reward constraints (see Figure 4).
For composition, show the properties of constrained composition of diffusion mod- els through experiments in Section 5.2. We see that when the composition weights are not chosen properly, the resulting model can become biased towards certain individ- ual models while neglecting others. Constrained composition addresses this issue by finding optimal weights that preserve closeness to each individual model. Particularly, when composing multiple text-to-image models each finetuned on a different reward, imposing constraints yields weights that enable the composed model to achieve higher performance across all rewards, compared to composition with equal weights.
2 Composition and Aligmment of Diffusion Models
We introduce constrained distribution problems for alignment and composition in Section 2, and characterize the reverse and forward KL divergences for diffusion models in Section 2.2.

---

Figure 1: Product composition (AND). Three Gaussian distributions being composed (Left). Compo- sition using equal weights (Middle), and with constraints (Right). The constrained model samples from the intersection of the three models.
2.1 Composition and alignment in distribution space
We formulate Unparameterized constrained distribution optimization problems using Reverse For- ward KL divergence, for Alignment and Composition as illustrated in (UR-A), (UR-C), and (UF-C).
Reward alignment: Given pretrained model and a set of rewards (r.(z)) = that can be evalu- ated on a sample æ, we consider the reverse KL divergence (KL || q) := p(x) log(p(z) (q)))) that measures the difference between distribution and the pretrained model Additionally, for each reward 7ú, we define a constant b, standing for requirement for reward ry. We formulate a constrained alignment problem that minimizes reverse KL divergence subject
constraints:
p² = argmin DKL (p
subject to Ex [r:(z)] 2 b; for i 1:::=1... 1,..., m (UR-A)
P
As (UR-A), the constrained alignment problem is solved by the distribution p¹ that is closest to the pretrained one q as measured by the reverse KI. divergence (p ()) among those whose expected rewards Ez (ri((x)) accumulate to at least b; By 'pretrained model' we refer to a sampling process that produces samples, not the underlying distribution. Let the primal value be PAL == DxL (p" g)
Product composition (AND): Given set of pretrained models {q');= formulate strained composition problem that solves a reverse cL-constrained optimization problem:
(p' u²) argmin u
subject DKL DKL(pll (pllg') < u for I...,, m.
(UR-C)
In (UR-C), the decision variable serves as an upper bound on the m KL divergences between distribution and m pretrained models {q ] 1. Partial minimization over u allows us to search for a distribution p that minimizes this common upper bound. Hence, the optimal solution p' minimizes the maximum KL divergence among m terms, each computed between and pretrained model gi- Let the primal value be PAND := u*. The epigraph formulation (UR-C) is practical, as the constraint threshold can be updated dynamically during training. In contrast, Figure shows that the model composed with equal weights biased toward the two most similar distributions.
Mixture composition (OR): A different composition modality that also fits within our constrained framework is the forward KL-constrained composition problem. We obtain this formulation by replacing the reverse divergence (KL(p 4s) in (UR-C) with the forward KL divergence DKL p):
up)
uu) = argmin u P. a
subject to DKL ()) < for =1==1.... l.....m.
(UF-C)
Mixture composition was studied in related but slightly different constrained setting [22]. In fact, the solution of the constrained problem (UF-C) learns to sample from each distribution in proportion to its entropy; see [22, Theorem 2). As shown in Figure 2, the constrained model samples more frequently from the higher-entropy distribution with two models, whereas the equally weighted composition samples equally from both distributions, leading to unbalanced sampling. Since the algorithmic design and analysis for (UF-C) follow those [22], mixture composition not the main focus of this work. For completeness, compare it with product composition in Appendix E.

---

Distributionn Unronstrainee Consireined
Mixture composition (OR). Two of Gaussian mixtures being composed (Left). One has two modes and the other has only a single mode. Composition using equal weights (Middle), and with constraints (Right).
Figure 2:
The reverse KL-based composition (UR-C) tends sample from the intersection of the pretrained models {q V 1. whereas the forward KL-based composition (UF-C) tends to sample from their union. Thus, product composition enforces a conjunction (logical AND) across pretrained models, while mixture composition corresponds to disjunction (logical OR). We emphasize that Problems (UR-A), (UR-C), and (UF-C) should serve as canonical formulations; the proposed constrained methods can be readily adapted to their variants, e.g., mixture composition with reward constraints.
2.2 KL divergence for diffusion models
generative diffusion model consists of forward and backward processes. In the forward process, add Gaussian noise e, to clean sample Xo over T time steps follows
X,
an X,-1 + V at Cr for 1,... ar11
a-1
where Et (O, 1) is the standard Gaussian noise, and {ar 12' is a decreasing sequence of coefficients called the noise schedule. We denote the marginal density of X₁ at time t as (-).). Given d-dimensional score predictor function s(x, f): Rd x (1....,777 - Rd, we introduce backward denoising diffusion implicit model (DDIM) process [42] as follows
04 Xx + s(x,,, + occ. ae
X,-1
(2)
where (O, 1) is the standard Gaussian noise, and {o?) is the variance schedule that determines the level of randomness in the backward process (e.g., 0, reduces deterministic trajectories), and Br = 0--1 ((--a)) a.) (a)) (11 o1 - ā+) is determined by the variance schedule and the noise schedule aç. Here, we use the equivalence between the score-matching and denoising formulations of diffusion model replace the denoising predictor in [42] by the score function. Given a score function s(x, t), we denote the marginal density of X, as (+)s) and the joint distribution over the entire process Po:r (Nuri s).
or
In the score-matching formulation [43], denoising score- motee-mtthhng objective is minimized to obtain function s* that approximates the true score function of the forward process, i.e., (xx se V log Pe (z). Then, the marginal densities of the backward process (2) match those of the forward process (1), i.e., pe(: s*) (+(.) (-) for all f.. Thus
the backward process to generate samples Zo ~ Pa that resemble samples from the original data distribution Hp Po-
We denote the KL divergence between two joint distributions p. q over the backward process by DKL (Po.r(-) (-)r())) which is known as path-wise KL [17, 19J. The path-wise KL divergence is often used in alignment to measure the difference between finetuned and pretrained models.

---

Lemma (Path-wise KL divergence). If two backward processes Po:r (-) and qu:r have the same variance schedule of and noise schedule at, then the reverse KL divergence between them is given by
(KI.(Po:(((55))) (1)))
Σ Re, ~ (110p)
[1)( 20
8 S2 [01
(3)
See Appendix C.I for the proof. When the two backward processes differ in their variance and noise schedules, the path-wise KL divergence remains tractable, and we omit it for simplicity. While the path-wise KL divergence
useful regularizer for alignment, when composing multiple models, the point-wise KL divergence Dri (po(-) || go(-)) is a more natural measure of the closeness between two diffusion models. This is because we mainly care about the closeness of the final sampling distributions: P-).), (-). and not the underlying processes: Po:r (-). (o:r(-). However, since our proposed approach to compute the wise KL is intractable for alignment, we adopt the path-wise KL for alignment and retain the point- wise KL for composition; see more discussion in Section 4.2.
However, it is not obvious how to compute the point-wise KI. divergence, as evaluating the marginal densities is intractable. We next establish a similar formula as (3) by limiting the score function class.
Lemma (Point-wise KL divergence). Assume two score functions sp (z.t) = V log p. (z), (2 t) V log (z). where Pt. qe are two marginal densities induced by two forward diffu- sion processes, with the same noise schedule, starting from initial distributions pu and qu. respectively. Then, the point-wise KL divergence between two distributions of the samples generated by running DDIM with Sp and Sy is given by
T
(pols :8p) go :1))
= Σ
We Ee Be-p(((ten)
[1ss(cc,1) (15p(,11-8s 80(2, 111]]
+ eT
where We a time-dependent constant, and er a discretization error that depends on the total number of diffusion time steps T.
See Appendix C.2 for the proof. The key intuition behind Lemma 2 is that if two diffusion processes are close, and their starting distributions are the same (c.g.. -V( I) at time t 7"), then the end points (i.e., the distributions t = 0) must also be close. The sum on the right-hand side of (4) can be viewed as the difference of the two processes over time steps, up to a discretization error.
3 Aligning Pretrained Model with Multiple Reward Constraints
We provide characterization of the solution Problem (UR-A) in Section 3.1, and establish strong duality for diffusion models in Section 3.2, together with a dual-based training algorithm.
3.1 Reward alignment in distribution space
To apply Problem (UR-A) to diffusion models, first employ Lagrangian duality to derive its solution in distribution space. Alignment with constraints is related but fundamentally different from the standard approach of minimizing a weighted average of the KI. divergence and rewards [17]. They are related because the Lagrangian for Problem (UR-A) is precisely the weighted average:
LAII(P, = (xllg)) - AT (Ex=n[((x)) b).
(5)
where we use shorthand b := [b,...., bm IT. r r= /ri,...,,mi Fm JT, and À := Da,.... dmd is the Lagrangian multiplier dual variable. Let the dual function be DA == minimize, (A A) and an optimal dual variable be X" E argmax, 20 DAu(A). Denote DA == DA For any
0, we define the reward weighted distribution g/2)) (subscript rw for reward weighted):
(A) (-) == 7Tm (A) (()
where (A) = ((x)ee (r(e) is the normalizing constant.
In the distribution space, Problem (UR-A) a convex optimization problem, since the KI. divergence is strongly convex and the reward constraints are linear in p. Thus, apply strong duality in
5

---

convex optimization [4] to characterize the solution to Problem (UR-A) in Theorem 1. Moreover, it is ready to formulate the constrained alignment problem (UR-A) as an unconstrained problem by specializing the dual variable to a solution to the dual problem.
Assumption 1 (Feasibility). There exists a model such that Eg [x))(r(<))) > by all = 1,..., ......... n.
Theorem (Reward alignment). Let Assumption I hold. Then, Problem (UR-A) is strongly dual, i.e, PuL = DAL- Moreover, Problem (UR-A) is equivalent to
minimize DKL (p (pl] ((x
where X is an optimal dual variable, and the dual function has the explicit form: DALI(A) log Zow (A). Furthermore, an optimal solution of (UR-A) is given by
p'
(2*) Trw
(8)
See Appendix C.3 for the proof. Theorem provides closed-form solution to the constrained
alignment problem (UR-A), ie., gre (A") This solution generalizes the reward-tilted distribution [13], which corresponds to finetuning model with expected reward regularizer. In Problem (UR-A),
the optimal dual variable X assigns weights to the rewards such that all the constraints are satisfied optimally, while remaining close possible to the pretrained model.
3.2 Reward alignment of diffusion models
We introduce diffusion models into Problem (UR-A) by representing p and q as two diffusion models: Po: (pp)) and ("; 54), with score functions sp and sg. respectively. The path-wise KL divergence has been widely used in diffusion model alignment to capture the difference between two diffusion models [44]. Hence, we instantiate Problem (UR-A) in space of score functions as follows
minimize *p ES DKL. (Po:rPo 8p) (o:r(15 8g)
(SR-A)
subject to
Bzo-pol::9n [xo)] 2 b; for I,..., m.
We define the Lagrangian for Problem (SR-A) as LALI(Sp, A) LALI(POTT sp), A). Similarly, we introduce the primal and dual values: PA and D In general, Problem (SR-A) is not guaranteed to be convex, since the path-wise KL divergence (3) involves an expectation taken over the backward process (.:r()). Nevertheless, the path-wise KL divergence is convex in the entire path space {po:r(-)); and constraints are linear. Hence, when the score function class S is expressive enough to induce any path distribution, we establish strong duality for Problem (SR-A) in Theorem 2.
Theorem 2 (Strong duality). Let Assumption I hold for some & S. any path distribution po-r(-) can be induced by a score function 8y E S, then Problem (SR-A) is strongly dual, i.e., PAL DAL
See Appendix C.4 for the proof. It is mild to assume the score function class is expressive, as diffusion models typically employ overparameterized networks (e.g., U-Nets or transformers) in practice. Motivated by strong duality, we propose dual-based method for solving Problem (SR-A), alternating between minimizing the Lagrangian via gradient descent and maximizing the dual function via dual sub- ara-gradient ascent below.
Primal minimization: At iteration n, we obtain new model g (n+1) via a Lagrangian maximization:
(n+1) E argmin TALL((p, (m)).
S
Dual maximization: Then, the model ((++1) estimate the constraint violation Ezg [r(aro)] b, denoted as r r(s(++1)) b, and perform dual sub-gradien ascent step:
(m+1) [1m)) + -
4 Constrained Composition of Multiple Pretrained Models
We provide characterization of the solution to Problem (UR-C) in Section 4.1, and establish strong duality for diffusion models in Section 4.2, together with dual-based training algorithm.

---

4.1 Composition in distribution space
To apply Problem (UR-C) diffusion models, we first employ Lagrangian duality to derive its solution in distribution space. Let the Lagrangian for Problem (UR-C)
LAND(P: u. A) = u + Σ X (P u)
(9)
and the associated dual function (A), which is always concave, defined as DAND(A) = max LAND u, A).
(10)
Let a solution to Problem (UR-C) be (p", u²), and let the optimal value the objective function be PAdD u², Let optimal dual variable be X E argmax, 2 (A) and the optimal value of the dual function be AND := DAND(A*). For any À > O, we define the tilted product distribution
AAND as a product m tilted distributions {q') {r')==..
(A)
II (A)) TX
AAND (-) = ZAND (A)
(11)
where (AND(A) = f (x)) 1Tx de the normalizing constant.
In the distribution space, Problem (UR-C) a convex optimization problem, since the sub-level set of the KL divergence is convex. Again, apply strong duality in convex optimization [4] to characterize the solution to Problem (UR-C) in Theorem 3. Moreover, is ready to formulate the unconstrained problem by specializing the dual
constrained composition problem (UR-C) variable to solution to the dual problem.
Assumption 2 (Feasibility). There exists pair (p, such that (r < u for all 1,...,P.
Theorem 3 (Product composition). Let Assumption 2 hold. Then, Problem (UR-C) is strongly dual, i.e,, PADD DAND Moreover, Problem (UR-C) is equivalent
minimize DRI. (v) [IID) (X") )
(12)
where X is the optimal dual variable, and the dual function has the explicit form, D(A) log ZAND (A). Furthermore, the optimal solution of (12) is given by
p² AANDD
(A*)
(13)
See Appendix C.5 for proof. The distribution LAND (Z(() Atx allows sampling from weighted product of m distributions {q').² - where the parameters {>,/1T A weight the importance of each distribution. The geometric mean is special case when all X are equal [1].
proposed constrained optimization problem (UR-C) to the well-known problem of sampling from a product of multiple distributions [1, 14]. Furthermore, our constraints enforce that the resulting product is properly weighted to ensure the solution diverges as little as possible from each of the individual distributions (see Figure 1 for illustration).
Remark 1. Theorem 3 connects
4.2 Product composition of diffusion models
We introduce diffusion models into Problem (UR-C) by representing p and q² as two diffusion models: p( sp) and qº (zo; sy), with score functions 8p and s respectively. The point-wise KL divergence naturally measures the closeness of the final sampling distributions we care about. Hence, we instantiate Problem (UR-C) a space of score functions as follows
minimize &p E subject to
(SR-C)
(pl sp) g(=o; st)) S a for 1,...,m m.
We define the Lagrangian for Problem (SR-C) as LAND (Sp, A) == LAND( sp), u, A). Similarly, we introduce the primal and dual values PAND and DAND- Although Problem (SR-C) is non-convex.
7

---

since the point-wise KL divergence (4) involves an expectation taken over the backward process (o: Nevertheless, the point-wise KL divergence is convex in the final distribution space. Hence, when the score function class S is expressive enough to induce any path distribution (hence any final distribution), we establish strong duality for Problem (SR-C) in Theorem 4.
Theorem 4 (Strong duality). Let Assumption 2 hold for some $ E S. If any path distribution py:r can be induced by score function Sp S. then Problem (SR-C) is strongly dual, i.e., PAND ANDD
See Appendix C.6 for proof. It is mild to assume the score function class is expressive, diffusion models typically employ overparameterized networks (e.g., U-Nets or transformers) in practice. To solve Problem (SR-C), similar to the one in Section 3.2, we apply dual-based approach below.
Primal minimization: At iteration n, we obtain a new model g(++1) via a Lagrangian maximization:
8 ((+1))
E arggmin LAND(Ep, (x)).
Sy ES
Dual maximization:
Then, we use the model g(x+1) to estimate the constraint violation and perform a dual sub- Siih--radient ascent step:
(n+1)
((()) + (DKL (p(x0: s(n+1)) ( s()) - [, for i = f 1....,m.
It is nontrivial to compute the point-wise KL. divergence in the Lagrangian LAND (8p: A(m) and the constraint violations above. Recall that Lemma 2 gives us a way to compute the posnt-wise KL: DxL (p(r0; s) (xx; sk)). However, requires the functions s and sy each to be a valid score function for some forward process. Indeed, this is the case for sy since it a pretrained model where it would have been trained to approximate the true score of a forward process. Yet, regarding the function 8 that we are optimizing over, there is no guarantee that any given 8 E S is a valid score function. To address this issue, we introduce Lemma 3 that allows us to minimize the Lagrangian. Lemma 3. The Lagrangian for Problem (SR-C) is equivalently written as
LAND(s, A) = DKL p(200 s) ll (To) ) - log (AN) Furthermore, a Lagrangian minimizer g(A) E argmin, LAND(S, >) is given by
(14)
s(A) argmin SES
Σ we E ti Bay q(xe(=o) ()) - V logg((x =o))l
(15)
(2) where g( (r)) ~ ((Varo, (1 L. and ((1)= v log AAND,
See Appendix C for proof. With Lemma 3, as long as we can obtain samples from the distribution (X)
LAND. we can approximate the expectation in (15) and use gradient- optimization methods to find a Lagrangian minimizer s( To do so, we use annealed Markov Chain Monte Carlo (MCMC) sampling [14], which requires having access to the scores of sequence of distributions that interpolate (A)
(X) smoothly between LAND (IT) and SAND (am):: V log AND (ar.) Σn = > V log g (Iz). In alignment, since we don't 't have these 'intermediate" scores, we cannot employ the approach in Lemma 3. See
(X)
Appendix B for sampling details.
For the dual update, we evaluate the KL divergence Dra (pa(:s())) |I so(:s")) between the marginal densities induced by the Lagrangian minimizer s(A) and the individual score functions s² using Lemma 2, since both are valid score functions.
Remark 2. In practice, the primal step only yields an approximate Lagrangian minimizer s(8) (xx x V log SAND, ((r). This results in two sources of error in evaluating the expectations on the RHS of (4):
T
[[LON(E. - s stx,,)]]] + ET (16)
DKL (p)
Σ Th Th E 100)) B t=o
The first error caused by not using the exact s(a) in (1, 1)) The second error introduced by not evaluating the expectation on correct trajectories given by x s(A)) (X)
However, the second error reduces, we have a way of sampling from the true product We ~ LAND 0 because can get samples from s(1) just by adding Gaussian noise to To-
See Appendix F for the detailed algorithm of product composition.
8

---

Computational Experiments
We demonstrate the effectiveness and merits of our constrained alignment and composition in series of computational experiments in Section 5.1 and Section 5.2, respectively.
5.1 Alignment of diffusion models with multiple rewards
We extend the AlignProp framework [35] to handle multiple rewards as constraints. We finetune Stable Diffusion v1.5² using several widely-used differentiable image quality and aesthetic rewards: aesthetic [39], hps [52], pickscore [23], imagereward [54] and MPS [60]. Since these rewards vary substantially in scale, making it difficult to set constraint levels, normalize each by computing the average and standard deviation over a number batches. all experiments, models are finetumed using LoRA [21]. Experimental settings and hyperparameters provided in Appendix G.
MPS, local contrast, and saturation constraints. common shortcoming of several off-the-shelf aesthetics, image preference, and quality reward models is their tendency to overfit to specific image characteristics such as saturation and sharp, high-contree textures; see, for example, images in the first column in Figure 3 (Right). To mitigate this issue, we add regularizers to the reward function to explicitly penalize these characteristics. However, the regularization weight is not carefully tuned, models may overfit to the regularizers instead of optimizing for the intended reward. As shown in Figure 3, when using equal weights, the MPS reward decreases (Left). In contrast, our constrained approach effectively controls multiple undesired artifacts while ensuring none of the rewards neglected, achieving a near feasible solution at the specified constraint level: a 50% improvement.
MPS Only
Unconstrained
Constrained
Constraint level
MP5
Saturation
Local Contrast
aled Reward
150
Normalized
50
Equal Weights
Constrained
Figure 3: Reward alignment. Stable diffusion is finetuned using one reward that emphasizes aesthetic quality (MPS), and Saturation and Local Contrast regularizers. Reward values for the equal weights method and our constrained alignment (Left). Images are sampled from the aligned models (Right), and the model trained solely with MPS reward is used for comparison.
II. Multiple aesthetic constraints. When finetuning with multiple rewards, arbitrarily assigning fixed weights can lead uneven performance across rewards. As shown in Figure 4 (Left), the model tends to overfit one reward while neglecting more challenging ones (e.g.. hps). In contrast, constraining all rewards enables the model to improve each reward up its specified level, including the challenging ones. From Figure 4 (Middle), minimizing the KL divergence subject to these constraints also yields smaller KL divergence to the pretrained model. Without constraints, overfitting a subset of rewards causes the model to deviate excessively from the pretrained one, which is undesirable (Right).
5.2 Product composition of diffusion models
In high-dimensional settings such as image generation, obtaining samples from the true product distribution via MCMC and then minimizing the Lagrangian in (15) to estimate the true product score function is prohibitively expensive. To address this, we employ a surrogate for the true score both for sampling and for computing the KI. divergence, as detailed in Appendix G.
Composing models finetuned on different rewards. We investigate the composition of several finetuned variants of the same base model, where each model is trained with LoRA a different reward
hatps://uuggiggiaee.c//saaabb/b wotfic///ooobbliiiifiiiiooooo 5/stable-diffusion

---

Pre-trained
Unconstrained
Constrained
Unconstrained Corstrained
saathetic
palacore
unconsianeed
Corstroned
1.5
Normelized Reward
0.5
(LDivrreenos
0.20
6.0
0.33
-1.0
6.05
0.00
2.5
5.0
7.5
10.0
12.5
15.0
2.5
7.5
10.0
125
15.0
Epoch
Epoch
Figure 4: Reward alignment. Stable diffusion is finetuned using multiple image quality/aesthetic rewards. Reward trajectories for the regularization-base method and our constrained alignment during training (Left). KL divergences to the pretrained model (Middle). Images are sampled from the aligned models (Right), and the pretrained model is used for comparison.
function. A key challenge is determining appropriate combination weights: arbitrary choices can lead to undesirable trade-offs and underrepresentation certain models in the mixture, as evidenced in Figure 5 by drops in up to 30% in some rewards. Our constrained composition provides a principled way to select weights that maintain proximity to each model, improving rewards across all models.
A
Mpp
Piducors
Im Imageroward
hpp
Nomalized Reward 1%:
100
60
Min. CLIP (f)
Min. BLIP (t)
40
Combined Prompting Equal Weights Constrained (Ours)
22.1
0.204
20
22.7
0.252
22.9
0,268
Equal Weigts
Constrained
Figure 5: Product composition. Stable diffusion with LoRA is finetuned using different rewards, for equal weighted and product mixtures. 100% represents the reward levels attained by models aligned solely with the individual reward. Higher better.
Table 1: Product composition. We compare our constrained composition with two baselines us- ing minimum CLIP and BLIP scores. The score is averaged over 50 different prompt pairs that are sampled from a list of simple prompts.
II. Concept composition with stable diffusion. Following the setting in [40], we compose two text- to-image diffusion models, each conditioned different input prompt. We apply the constrained composition (SR-C) to determine the optimal weights for composing two models, and compare against the baseline that uses equal weights. Closeness each model encourages faithful representation of both concepts in the images generated by the composed model, as reflected by improved text-to-image similarity metrics: CLIP [20] and BLIP [25], which are reported Table 1. We compute similarity scores between the generated images and each of the two prompts and compare their minimum values. We also include baseline where images are generated from a single combined prompt containing both inputs. Images from all approaches, along with implementation details and additional experimental results, are provided in Appendix G.
Conclusion
We have developed constrained optimization framework that unifies alignment and composition of diffusion models by enforcing that the aligned model satisfies reward constraints and/or remains close to each pretrained model. Theoretically, we characterize the solutions to the constrained alignment and composition problems and design dual-based training algorithms to approximate these solutions. Empirically, we demonstrate constrained approach image generation tasks, showing that the aligned or composed models effectively satisfy the specified constraints.
10

---

References
[1] B. Biggs, A. Seshadri, Y. Zou, A. Jain, A. Golatkar, Y. Xie, A. Achille, A. Swaminathan, and
S. Soatto. Diffusion soup: Model merging for text-to-image diffusion models. arXiv preprint arXiv:2406.08431, 2024.
[2] K. Black, M. Janner, Y. Du, I. Kostrikov, and S. Levine. Training diffusion models with reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024.
[3] A. Blattmann, R. Rombach, H. Ling, Dockhorn, 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 22575,, 2023.
[4] P. Boyd and Vandenberghe. Corvex optimization. Cambridge university press, 2004.
[5] A. Bradley and P. Nakkiran. Classifier-free guidance is predictor-corrector arXiv preprint arXiv:2408.. 09000, 2024.
[6] 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-176) 1760,, 2022.
[7] L. F. O. Chamon and A. Ribeiro. Probably approximately correct constrained learning, 2021.
[8] J. Chen, R. Zhang, Y. Zhou, and C. Chen. Towards aligned layout generation via diffusion model with aesthetic constraints. In The Twelfth International Conference on Learning Representations, 2024.
[9] C. Chi, Z. Xu, S. Feng, E. Cousineau, Y. Du, B. Burchfiel, R. Tedrake, and S. Song. Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal of Roboties Research, page 02783649241273668, 2023.
[10] M. Chidambaram, K. Gatmiry, S. Chen, H. Lee, and J. Lu. What does guidance do? fine-grained analysis in a simple setting. arXiv preprint wrXiv2249..1077 2024.
[11] J. K. Christopher, S. Baek, and N. Fioretto. Constrained synthesis with projected diffusion models. Advances in Neural Information Processing Systems, 37:89307-89333, 2024.
[12] K. Clark, P. Vicol, K. Swersky, and D. J. Fleet. Directly fine-tuning diffusion models on differentiable rewards. In The Twelfth International Conference on Learning Representations, 2024.
[13] C. Nomingo-Enrichh M. Drozdzal, B. Karrer, and R. T. Q. Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control, 2025.
[14] Y. Du, C. Durkan, R. Strudel, J. B. Tenenbaum, S. Dieleman, R. Fergus, J. Sohl-Dickstein, A. Doucet, and W. Grathwohl. Reduce, reuse, recycle: Compositional generation with energy- based diffusion models and meme, 2024.
[15] B. Flizalde, S. Deshmukh, M. A. Ismail, and H. Wang. Clap: Learning audio concepts from natural language supervision, 2022.
[16] Y. Fan and K. Lee. Optimizing DDPM sampling with shortcut fine-tuni aie-tuning. In International Conference Machine Learning, pages 9623-9639, PMLR, 2023.
[17] Y. Fan, O. Watkins, Y. Du, H. Liu, M. Ryu, C. Boutilier, P. Abbeel, Ghavamzadeh, K. Lee, and K. Lec. DPOK: Reinforcement learning for fine-tuning rext-to-image diffusion models, 2023.
[18] G. Giannone, A. Srivastava, O. Winther, and F. Ahmed. Aligning optimization trajectories with diffusion models for constrained design generation. Advances Neural Information Processing Systems, 36:51830- 1180-5131111 2023.
11

---

[19] Y. Han, M. Razaviyayn, and R. Xu. Stochastic control for fine-tuning diffusion models: Optimality, regularity, and convergence. arXiv preprint arXiv:2412.18164, 2024.
[20] Hessel, A. Holtzman, M. Forbes, R. L. Bras, and Y. Choi. Clipscore: A reference-free evaluation metric for image captioning, 2022.
[21] E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang. L. Wang, W. Chen, et al. Lora: Low-rank adaptation large language models. ICLR, 1(2):, 2022.
[22] S. Khalafi, D. Ding, and A. Ribeiro. Constrained diffusion models via dual training. In The eiihth- Annual Conference on Neural Information Processing Systems, 2024.
[23] Y. Kirstain, A. Polyak, U. Singer, S. Matiana, J. Penna, and O. Levy. Pick-a-pic: An open dataset of user preferences for text-to-image generation. Advances in Neural Information Processing Systems, 36:36652-36665 -36623,6633...2 2023.
[24] K. Lee, H. Liu, M. Ryu, O. Watkins, Y. Du, C. Boutilier, P. Abbeel, M. Ghavamzadeh, and S. S. Gu. Aligning rext-to-image models using human feedback. arXiv preprint arX1122302.12192, 2023.
[25] J. Li, D. Li, C. Xiong, and S. Hoi. Blip: Bootstrapping language-image pre-training for unified vision-language understanding and generation, 2022.
[26] S. Li, K. Kallidromitis, A. Gokul, Y. Kato, and K. Kozuka. Aligning diffusion models by optimizing human utility. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[27] J. Liang, J. K. Christopher, S. Koenig, and F. Fioretto. Multi-agent path finding in continuous spaces with projected diffusion models. arXiv preprint arXiv:2412.1799,, 2024.
[28] Liang, K. Christopher, S. Koenig, and F. Fioretto. Simultaneous multi-robot motion planning with projected diffusion models. arXiv preprint arXiv:2502.03607 2025.
[29] B. Liu, S. Shao, B. Li, L. Bai, Z. Xu, H. Xiong, J. Kwok, S. Helal, and Z. Xie. Alignment of diffusion models: Fundamentals, challenges, and future. arXiv preprint arXiv:2409.07253, 2024.
[30] H. Liu, Z. Chen, Y. Yuan, X. Mci, X. Liu, D. Mandic, W. Wang, and M. D. Plumbley. Audioldm: Text-to-audio generation with latent diffusion models, 2023.
[31] N. Liu, S. Li, Y. Du, A. Torralba, and J. B. Tenenbaum. Compositional visual generation with composable diffusion models. In European Conference Computer Vision, pages 423-439 Springer, 2022.
[32] S. Lyu. Interpretation and generalization of score matching, 2012.
[33] W. Mou, N. Flammarion, J. Wainwright, and P. L. Bartlett. Improved bounds for discretiza- tion of langevin diffusions: Near-optimal rates without convexity, 2019.
[34] S. Narasimhan, S. Agarwal, L. Rout, S. Shakkottai, and S. P. Chinchali. Constrained posterior sampling: Time series generation with hard constraints. arXiv preprint arXiv:2410. 12652, 2024.
[35] M. Prabhudesai, A. Goyal, D. Pathak, and K. Fragkiadaki. Aligning text-to-image diffusion models with reward backpropagation, 2024.
[36] M. Prabhudesai, R. Mendonca, Z. Qin, K. Fragkiadaki, and D. Pathak. Video diffusion alignment via reward gradients. arXiv preprint arXiv:2407.08737, 2024.
[37] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models, 2022.
[38] C. Saharia, W. Chan, S. Saxena, L. Li, Whang, L. Denton, K. Ghasemipour, R. Gon- tijo Lopes, B. Karagol Ayan, T. Salimans, et al. Photorealistic text-to-image diffusion mod- els with deep language understanding. Advances in neural information processing systems, 35:36479- 36494, 2022.
12

---

[39] C. Schuhmann, R. Beaumont, R. Vencu, C. Gordon, R. Wightman. M. Cherti, T. Coombes, A. Katta, C. Mullis, M. Wortsman, et al. Laion-5b: An open large-scale dataset for training next generation image-text models. Advances in neural information processing systems, 35:25278- 25294, 2022.
[40] M. Skreta, L. Atanackovic, J. Bose, A. Tong, and K. Neklyudov, The superposition of diffusion models using the itô density estimator, In The Thirleenth International Conference Learning Representations, 2025.
[41] M. Sohrabi, J. Ramirez, T. H. Zhang, S. Lacosst-Julien, and J. Gallego-Posada. On pi controllers for updating lagrange multipliers in constrained optimization. In International Conference on Machine Learning, pages 45922-45954 PMLR, 2024.
[42] Song, C. Meng, and S. Ermon. Denoising diffusion implicit models. 2022.
[43] Y. Song, Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations, 2021.
[44] M. Uehara, Y. Zhao, T. Biancalani, and S. Levine. Understanding reinforcement learning- -baed fine-tuning of diffusion models: A tutorial and review. arXiv preprint (xXXx:2407.13734 2024.
[45] M. Uehara, Y. Zhao, Black, E. Hajiramezanali, G. Scalia, N. L. Diamant, A. M. Tseng, T. Biancalani, and S. Levine. Fine-tuning of continuous-time diffusion models as entropy- regularized control, arXiv preprint arXiv:2402.15199 2024.
[46] M. Uchara, Y. Zhao, K. Black, E. Hajiramezanali, G. Scalia, N. L. Diamant, A. M. Tseng, S. Levine, and T. Biancalani. Feedback efficient online fine-tuning of diffusion models. In Forty-first International Conference on Machine Learning, 2024.
[47] M. Uehara, Y. Zhao, E. Hajiramezanali, G. Scalia, G. Eraslan, A. Lal, S. Levine, and T. Bian- calani. Bridging model-based optimization and generative modeling via conservative fine-tuning of diffusion models. Advances in Neural Information Processing Systems, 37:12751 11175555517755 2024.
[48] A. Ulhaq and N. Akhtar. Efficient diffusion models for vision: A survey. arXiv preprint arXiv:2210009292. 2022.
[49] B. Wallace, M. Dang, R. Rafailov, L. Zhou, A. Lou, S. Purushwalkam, S. Ermon, C. Xiong, S. Joty, and N. Naik. Diffusion model alignment using direct preference optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8228 8238, 2024.
[50] L. Wang, C. Song, Z. Liu, Rong, Q. Liu, and S. Wu. Diffusion models for molecules: survey of methods and tasks. arXiv preprint arXiv:2502.095// 2025.
[51] X. Wu, Y. Hao, M. Zhang, K. Sun, Z. Huang. G. Song, Y. Liu, and H. Li. Deep reward supervisions for tuning text-to-image diffusion models. In European Conference on Computer Vision, pages 104-124, 2024.
[52] X. Wu, K. Sun, F. Zhu, R. Zhao, and H. Li. Better aligning cextto--mmaee models with human preference. arXiv preprint 2303.
1(3), 2023.
[53] X. Wu, K. Sun, F. Zhu, R. Zhao, and Li. Human preference score: Better aligning text- to-image models with human preference. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 2096-2105, 2023.
[54] J. Xu, X. Liu, Y. WU, Y. Tong, Q. Li, M. Ding, Tang, and Y. Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. Advances in Neural Information Processing Systems, 36: 5903-155355 2023.
[55] J. Xu, X. Liu, Y. Wa, Y. Tong, Q. Li, M. Ding, J. Tang, and Y. Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. Advances in Neural Information Processing Systems, 36, 2023.
13

---

[56] J. N. Yan, J. Gu, and A. M. Rush. Diffusion models without attention. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8239-8249, 2024.
[57] K. Yang. J. Tao, J. Lyu, Ge, J. Chen, W. Shen, X. Zhu, and X. Li. Using human feedback to fine-tune diffusion models without any reward model. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1941-9951,2 2024.
[58] S. Zampini, J. Christopher, L. Oneto, D. Anguita, and F. Fioretto. Training-free free constrained generation with stable diffusion models. arXiv preprint arXiv:2502.0562) 2025.
[59] H. Zhang and T. Xu. Towards controllable diffusion models via reward-guided exploration, 2023.
[60] S. Zhang, B. Wang, J. Wu, Y. Li, T. Gao, D. Zhang, and Z. Wang. Learning multi-dimentn dimensional human preference for text-to-image generation. Proceedings of the IEEF/CVE Conference on Computer Vision and Pattern Recognition, pages 8018-8027 2024.
[61] Z. Zhang, L. Shen, S. Zhang, D. Ye, Y. Luo, M. Shi, B. Du, and D. Tao. Aligning few-step diffusion models with dense reward difference learning. arXiv preprint arXiv:2411.11777, 2024.
[62] H. Zhao, H. Chen, Zhang, D. D. Yao, and W. Tang. Scores actions: framework of fine-tuning diffusion models by continuous-time reinforcement learning. arXiv preprint arXiv:2409.08400, 2024.
14

---

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: There are references in the introduction sections of the paper were we dicuss in depth the claims made.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the limitations briefly in the conclusion, and more thoroughly in Appendix A.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiitiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
15

---

Answer:
[Yes]
Justification:
Yes, the full proofs for the theoretical results are provided in Appendix C.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We provide details for reproducing the experiments in the paper in Appendix G. We will also provide the code used for all the experiments upon the paper's 's publication.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path reproducing verifying the results. 5. Open access to data and code
16

---

reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes)
Justification: We will make public the repository with our code and implementations used for all of the experiments upon the paper's publication and include link to it in the paper. Instructions and implementation details are provided in Appendix G. Anonymized code for implementing some of the experiments will also be provided with the supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
public/ /eaiiiiggge/////eeeeeeeiiiiissee wetgggg//cdddubbnneeeee http///ooeuuissiooooo.. ceiioonoooiioo Pooiollyy))).. for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: Coo///ooddeeseseeeeeeeeeee / httpscccccbbcccccc hediii////ddess/eeeuuuiiiiii wwwisssssnnpoooooo onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: These details are provided in Appendix G.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Most of the plots in the main paper include error bars. Some don't t for visual clarity. More details statistical significance of the results provided in Appendix G.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
17

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We ran the experiments on system with 2 NVIDIA RTX A6000 GPUs with 48 GB of GPU memory each. More details be found in Appendix G.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https atttps//nneus http///eerrpp..ee
Answer: [Yes]
Justification: Yes.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. if there special consid- eration due to laws or regulations in their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: These impacts discussed in Appendix A.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
18

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We use publicly available pretrained models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
The models and code used have been properly cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URI.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
19

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper doess not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
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 equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
20

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: /ttp///errpp/////ee cc/Confferences/2025 /LLM) for what should or should be described.
21

---

Supplementary Materials for "Composition and Alignment of Diffusion Models using Constrained Learning"
A Limitations and Broader Impact
Despite offering a unified constrained learning framework and demonstrating strong empirical results, further experiments are needed assess our method's effectiveness on alignment and composition tasks beyond image generation, under mixed alignment and composition constraints, and in combination with inference- time techniques. Additionally, further theoretical work needed to understand optimality of non-convex constrained optimization, convergence and sample complexity of primal-dual training algorithms.
Limitations:
Our method can enhance diffusion models' compliance with diverse requirements, such as realism, safety, fairness, and transparency. By introducing a unified constrained learning framework, our work offers practical guidance for developing more reliable and responsible diffusion model training algorithms, with potential impact across applications such as content generation, robotic control, and scientific discovery.
Broader impact:
B
Related Work
Alignment of diffusion models. Our constrained alignment is related to a line of work on finetuning diffusion models. Standard finetuning typically involves optimizing either a task-specific reward that encodes desired properties, or a weighted sum of this reward and regularization term that encourages closeness to the pretrained model; see [16, 55, 24, 53, 59, 51, 2. 12, 61] for studies using the single reward objective and [45, 62, 47, 46, 36, 17, 19] for those using the weighted sum objective. The former class of single reward-based studies focus exclusively on generating samples with higher rewards, often the cost of generalization beyond the training data. The latter class introduces a regularization term that regulates the model to be close to the pretrained one, while leaving the trade-off between reward and closeness unspecified; see [44] for their typical pros and cons in practice. There are three key drawbacks to using either the single reward or weighted sum objective: (i) the trade-off between reward maximization and leveraging the utility of the pretrained model is often chosen heuristically; (ii) it is unclear whether the reward satisfies the intended constraints; and (iii) multiple constraints are not naturally encoded within single reward function. In contrast, we formulate alignment as a constrained lcarning problem: minimizing deviation from the pretrained model subject to reward constraints. This offers a more principled alternative to existing ad hoc approaches [8, 18]. Our new alignment formulation (i) offers theoretical guarantee of an optimal trade-off between reward satisfaction and proximity to the pretrained model, and (ii) allows for the direct imposition of multiple reward constraints. We also remark that our constrained learning approach generalizes finetuning of diffusion models with preference [49, 57, 26].
Composition of diffusion models. Our constrained composition approach is related to prior work on compositional generation with diffusion models. When composing pretrained diffusion models, two widely used approaches are (i) product composition (or conjunction) and (ii) mixture composition (or disjunction). In product composition, it has been observed that the diffusion process is not compositional, e.g., a weighted sum of diffusion models does not generate samples from the product of the individual target distributions [14, 10]. To address this issue, the weighted sum approach has been shown to be effective when combined with additional assumptions or techniques, such as energy-based models [31, 14], MCMC sampling [14], diffusion soup [1], and superposition [40]. However, how to determine optimal weights for the individual models is not yet fully understood. In contrast, we propose a constrained optimization framework for composing diffusion models that explicitly determines the optimal composition weights. Hence, this formulation enables an optimal trade-off among the pretrained diffusion models. Moreover, our constrained composition approach also generalizes to mixture composition, offering advantages over prior work [31, 14, 1, 40).
22

---

Diffusion models under constraints. Our work is pertinent to a line of research that incorporates constraints diffusion models. To ensure that generated samples satisfy given constraints, several ad hoc approaches have proposed that train diffusion models under hard constraints, e.g.. projected diffusion models (27, 11, 28], constrained posterior sampling [34], and proximal Langevin dynam- ies [58]. In contrast, our constrained alignment approach focuses expected constraints defined via reward functions and provides optimality guarantees through duality theory. A more closely related work considers constrained diffusion models with expected constraints, focusing on mixture composition [22]. In comparison, we develop new constrained diffusion models for reward alignment and product composition.
C Proofs
For conciseness, wherever it is clear from the context we omit the time subscript:
sp)) = DKL (p(Zu:r;
(17)
C.1 Proof of Lemma
Proof. The DDIM process is Markovian in reverse time with the conditional likelihoods given by
(v | as Br s(2,,)), of
p x(is)
~N((
N
(18)
at
Using (18) we expand the path-wise KL:
(Po.To 8p) ())T)::8)))))))
Bzar ~ sp) log q( sy)
(a)
Σ log
p 1
Zz; 8p)
g
20, 8.)
(b)
1ig 8p) p
log
g
Bzor [DK 525 / 11 ar: 89))
g [| - sq (Tt, 1)]12
Biz pe+1
20
1
Σ
(p.)
83 (1p((1n) 20
sq sz (x(,)))))
where (a) is due to the diffusion process, (b) is due to the exchangeable sum and integration, (c) is the definition of reverse KL divergence at time (d) is due to the reverse KL divergence between two Guassians with the same covariance and means differing (e sy (ara, ()), and in (e)
we abbreviate E., ~ p as E(px) that is taken over the randomness of Markov process.
C.2 Proof of Lemma
The proof for Lemma 2 is quite involved, thus we have divided it into multiple parts for readability, In Section C.2.1, I, we give a few definitions for continuous time diffusion processes. Section C.2.2, prove analogue of Lemma 2 in continuous time. In Section C.2.3, we bound the discretization error éT incurred when going from continuous time processes to corresponding discretized processes and thus complete the proof. The proofs for all lemmas presented here can be found in Appendix D.
C.2.1 Continuous time preliminaries
Throughout the proof, we will dealing with continuous time forward and reverse diffusion processes and their discretized counterparts.
Notation Guide:
23

---

We denote the continuous time variable т E [0, 1] to differentiate it from the discrete time indices E {0.... J. corresponds to 7 1 and I corresponds
We denote I, the continuous time reverse DDIM process and X the corresponding discrete time process.
The forward processes we denote with an additional bar e.g. I,, X, denote the continuous time and discrete time forward processes respectively.
Marginal density of continuos time DDIM process with score predictor s(x, r) at time 7 we denote as: Pr ().
Given a function s(x, 7) : F2 [0, Rd, and a noise schedule a, increasing from āg 0 to ā1 = 1, we define continuous time reverse DDIM process
âr + o² ) T))dt + o,ddB Ho~ NN0,10 3r + f 2ā, 2 2a,
âr
d3,
(19)
The variance schedule o, is arbitrary and determines the randomness of the trajectories (e.g. O, 0 for all T, then the trajectories will be deterministic). The DDIM generative process (19) induces marginal densities (x, 8) for T E [0, 1].
For reference the Discrete time DDIM process defined in the main paper is
X
Ot-1 0 X,
Be (xs(X + oeet.
(20)
Up first order approximation, the discrete time process (20) is the Fuler-Maruyama discretization of the continuous time process (19). A uniform discretization of time is assumed, i.e., T (See [13, Appendix B.1] for the full derivation).
Given random variables Xo ~ Po = NO,, I) and I, ~ PT, where i1 is some probability distribution (e.g., the data distribution), we define a reference flow I, for r E [0, 1] as
Σ, = No + G1
(21)
Note that there is no specific process implied by the definition above, since different processes can have the same marginal densities as the reference flow at all times T. We denote by p-(-) the density of K,. As a, decreases from ae to a1 0, and &r increases from Ço 0 Ç1 the reference flow gives an interpolation between Pu = NO I) and P1.
If the score predictor s(r, 7) V, log (x), then the DDIM process (19) has the same marginals as the reference flow (21), i.e., P. (r, = (x) for 7 E [0, 1]. This assuming proper choice of a,, Gr, i.e., 0,, VI a,, vā,-
C.2.2 Proof for continuous time
We generalize [32, Theorem 1] to characterize how the KL. divergence between the marginals two continuous time forward processes changes with time.
Lemma 4. Consider reference flows defined as X, a, K1. for [0, 1] where To
N(,, 1). Denote by P. (.), the marginal density of X, when X, ~ P1 and similarly (:) the marginal density of X,, when Xn q1. The following then holds:
DKa qz(-) = - ( dr
(22)
where = -r/aa, and Dr(p) q) denotes the Fisher divergence.
By integrating the derivative of the KL. divergence as given by Lemma 4, we obtain the following continuous time analogue of Lemma 2, which characterizes the point-wise wise KL divergence of two continuous time diffusion processes.
For consistency with other works from whom we will utilize some results in our proofs, namely [13, 32). the direction of time we consider in continuous time is reversed compared to discrete time. This does not affect derivations and results beyond notation change.
24

---

Lemma 5. Consider two score predictors sp(2,T) Ve px(),, sq(z.T) Va log gr (α). where P-. 9, are marginal densities of two reference flows, with the same noise schedule, starting from initial distributions Po and qo. respectively. Then, the point-wise KI. divergence between two distributions of the samples generated by running continuous time DDIM (19) with sp and Bq is given by
Dra(pol:35p) (p) (990((
wi Ex (xwp)()))) [1sp(1, (2, 1)
(23)
where Wy is time- time-dependent constant
C.2.3 Bounding the discretization error
now turn to bridging the gap between continuous and discrete times. [33], they bound this gap which arises from the discretization of the continuous time diffusion process. We will utilize the main result from this paper with minor modification in that we consider a time-dependent drift term. This is formalized in Lemma 6 which allows bound the KL divergence between the marginals pa (-) of the discrete time backward DDIM process and the corresponding marginal pe,r(-) the continuous time backward process.
Lemma 6. (Modification of Theorem I from (33].) Under mild assumplions on the score function (outlined in the proof). the KL divergence between the marginals of the discrete time backward process pc(-) and continuous time backward process P-/r(-) can be bounded as follows:
(xe(pe('s 8p) (1)))) S 7
(24)
where e is a constant depending on the assumptions.
Next we need to characterize the sensitivity of the KI. divergence to perturbations the first and second arguments that apply Lemma 6.
Lemma 7. Assume M max, log( Po(sse)) is bounded. Then, the point-wise KL between the continuous time processes approximates the point-wise KL. between the discrete time processes up to a discretization error (T)::
[Dx 90(:ss)) (pa(pp(:s))) (o)) S (T),,
(25)
where e1 (T) O(1TT)
And lastly, we need to characterize the discretization error going from integral over continuous time to a sum over discrete time steps.
Lemma 8. Assume B1, B2 as defined below are finite:
B1 sup || r) 52(21, T)
(26)
x,7
B2 sup (s)(2,T) s -o
d
T))|
(27)
2.
dr
Then the integral from Lemma giving the point-wise KI. in continuous time can be approximated with discrete time sum as follows:
["
1 T Wr/T En ~ pel [|sp(2, s2(2,, 1)11
0= 3 E, Egppp(
Σ ==
< (2())
(28)
where the discretization error e2(T) O(1TT).
remains to combine Lemmas 7 and 8 to complete the proof of Lemma 2:
DRI(ppl: ; 8p) [| (100)=;8)))))))))))))))))))))))
Σ We (e
[|ssp(r, s,(x,, 1)11 + eT (29)
where lerl S e1 (T) + (2 = O1/T). (We abuse notation to denote # W in (29) and in the main paper.)
25

---

C.3 Proof of Theorem
Proof. For any 2 O, the optimal solution p" p²(:A) is uniquely determined by solving a partial minimization problem,
minimize (p A).
PEP
Application of Donsker and Varadhan's 's variational formula yields the optimal solution
p p=(;s) X) x q(-)e)
Since the strong duality holds for Problem (UR-A), its optimal solution is given by p² (=:A) evaluated at X It is straightforward to evaluate the dual function by the definition D(A) (::p(;)),)
C.4 Proof of Theorem 2
Proof. We first consider the constrained alignment (SR-A) in the entire path space (pu:r()) Since the path-wise KI. divergence is convex in the path space and the constraints are linear, the strong duality holds in the path space, i.e., there exists a pair (-), A²) such that
PÂLI = (xL(P ll g-:r(:;;,)) = DALI(A") = DA Equivalently, (Po:r ())x*) a saddle point of the Lagrangian (t)
LAuI(Pōrr()) A) s TAun(Po:(().'') < [Au((o:r(())**)) for all pa:r()) and X 2 0.
Since the score function class S is expressive enough, any path distribution P.:r(.) can be represented as po:r(:i 8p) with some 5p E S; and vice versa. Thus, we can express (o:r(-) as po.r(is)) with some 2 E S. We also note that the dual functions DALI(A) (A) in the path and score function spaces are the same. Hence, the dual value for Problem (SR-A) remains to be DAu(A*). Thus, (s), X*) is saddle point of the Lagrangian LALI(Sp, A) == LALI(PO sp),
LAII(S;; X*) < LAii(sp; for all 8p E and X 2 0. Therefore, the strong duality holds for Problem (SR-A) in the score function space S.
LALI(S); A)
C.5 Proof of Theorem 3
Proof. By the definition,
LAND (P, u; A) =
de (Dka(ell u)
m
ux
(N,E [log(p))] A,E anp log q' (z)
1 u-uAt
E, [log(()]] E Exap
log II (g())))
i=1
i=1
ut 1
Eznp [logp(x) - Exno log II (d'(z)) At 1)
+
A
A (DaL (2) u) 1 A (A).
By taking À = A obtain a primal problem: maximize, e P.u o LAND (p. u; X*), which solves the constrained alignment problem (UR-A) because of the strong duality. By the varational optimality, maximization of (AND(p, u; X+) over p and u is at unique maximizer,
p (=(;2t)
9AND (-)
(2*)
26

---

and u² = 1 1 2 0 and u² oo otherwise. This gives the optimal model p² (=
Meanwhile, for any X 2 0, the primal problem: maximize, P., 2 0 LAND (p, u; X) defines the dual function (A). By the varational optimality, maximization of LAND (p, u; X) over p and u is at unique maximizer,
p¹ p²(;; A. & (AND())
and u² (A) = 00 1 and u² v*(A)= oo otherwise. This defines the dual function, DAND(A) LAND(V" (AND()"(s)); (=A)) u² (A);
(A(A) + Σ A (D (:; (:: (N)D (100)))w(((() v²((a)) - 1 log (A)
(1-1"A))²(A) 11 (AND
which completes the proof by following the definition of the dual problem and the dual constraint
1
C.6 Proof of Theorem
Proof. Similar to the proof of Theorem 2, we can establish saddle point condition for the Lagrangian (sp: u, A) by leveraging the expressiveness of the function class S which represents the path space {.)). As the proof follows similar steps, omit the detail.
C.7 Proof of Lemma
Proof. From section C.5, recall::
m
LP u; A) u
d (DaL (p) (NN - u) -1 A Alog (AND(A).
(30)
Since in the diffusion formulation of the problem (SR-A) have p Po( o² po( we can derive similarly to (30) that:
(PNN(P)(:)). u: +
Σ A (DKL(PO(:ss) (I u) A Alog ZAND(A). (31)
1=1
Since minimizing over u would trivially give min, LAND(P, u; X) = -o unless A = 1, we
consider the Lagrangian the non-trivial case where A 1. Then we have:
LAND(P(:3s):)) LAND(S, A) = DRL (po (AND,)) - log (A).
(32)
The second term log ZAND(A) does not depend on 8, thus it suffices to minimize DRL (Po(:s) (A)
(AND,)) to find the Lagrangian minimizer which we call sAA), The KL is minimized
when Po(:; (A)) lADDoD If we have access samples from lAND,0
(A)
(A)
fit LAND, by
(X)
optimizing the Denoising score matching objective similar to Equation (1) in [43]:
T
Le (8, A) =
Σ
wi 30
(2)
Ber [|s(cc t) V log (x, |
(33)
t=0
From [43] know that given sufficient data and predictor capacity of we have argmin, L A) AANDo which concludes the proof.
(A)
D Additional Proofs
We provide detailed proofs for all lemmas in Section C.2.
27

---

D.1 Proof of Lemma 4
Proof. We start by defining 2, as a time-dependent scaling of I,:
D- = ar = +
(34)
where 7/ := 5 Denote by p- (2))) the marginal density of 2) when &1 ~ P1 and similarly Te (2),), the marginal density of 2), when X, 91- Now generalize Theorem from [32] to
show that (22) holds for pr, 9r- Their Theorem is for the specific case of /r v1 - t.
We now present Lemmas and 10 which we will need the remainder of the proof.
Lemma 9. For density p- (2+) defined in Theorem 1, the following identity holds:
= ".. (.
(35)
dt
Proof. Proof of Lemma 9. We start with p, (2), with pı (3t):
(2)))
which the convolution of a Gaussian distribution
/ xx. (2m72)d/22 exp
[0) 31
p. (2))
(381)
(36)
27?
Taking the derivative we have:
dt (2rr) = / 7/4 |0) -x1l1²
d
||]]
Ezll²)
pi pi())) (Zi)
7
(2m7?)d/22
exp
27?
(37)
7t
|
3t $11
P1 (1()))
81
(2m72)d/22
exp
27?
On the other hand, taking the gradient of (2))) with respect to 2) we get:
2, x1 exp
| 1111
V (-)))
(i()))
(38)
8x, 2 (2m72)a/22
Taking the divergence of the gradient, we have:
(2) E1
112
4g. Pr (2),) =
(2m72)d/22 exp
p((
R1
7
277
(39)
||
xt,
pi (2t).
R1 fo (2mr)2)a/2 exp
27
Comparing Equations (37) and (39) proves the result.
Lemma 10. For any positive valued function fx) : IRd - whose gradient Vzf and Laplacian A are well defined, we have the identity
A.f(x) f(x)
A2 log ((x) + | log f(x)]
(40)
*Just to avoid any confusion, in [32]. at f 0 we have the data distribution and as t increases the distributions converge to Gaussians. However in the current paper, the direction of time is the opposite, meaning t = 0 corresponds to the pure Gaussians and t have the data distributions.
28

---

We now continue with the proof of Lemma 4. We start with the definition of Fisher divergence for generic distributions p. q:
J p(z) fp(x) |V ((x) - - Flogg
Dr(pl1g)
dr
J. p(z)
Vp(a)
Vq((+) 2
dr
(41)
p(a)
g(x)
J. p(a))
()
Vg(z)
Vp(a) T VV((x) p(x))((x)
+
de
p(a))
We apply integration by parts to the third term. For any open bounded subset f of Rd with piecewise smooth boundary 02::
/ JEER Vp(x)]
Vg(z) (xx)
dr [ [p(x)
(Vlog
JZER
(42)
log (())de + J p(z)(V log ((x) Rt)ar
ZEQ
Assuming that both p(z) and q(xx) are smooth and fast-decaying, the boundary term in (42) vanishes. Then we can combine (41) and (42) to write:
DF(p p(xx) (I [[Vlol()>)] | [llog(>)) 2A, log ((g)(z) dr
(43)
Returning to our distributions 5. (2))) and q- (2)) we can rewrite (43) as:
(2).) log (2), | | log q, (2+) + 24g, log 9, (2),) dw) / 2)
(p(P (()
(44)
For conciseness in notation, we drop references variables 2, and ž, in the integration, the density functions, and the operators whenever this does not lead to ambiguity. We start by applying Lemma 10 Equation (43):
/ pr (IV log
Dr(p qr))
V log 12 2A log 7r)
(45)
Aqr + A log बा
pr log P-1
Next, we expand the derivative the KL divergence:
dr DKL (+)) = =f d
dr P- log PT + f pr dr log - f pr- dr
d log qr-
We climinate the second term by exchanging integration and differentiation of
/ pr d dr
/
dp- d
p 0. / dr
logp,
dr
As result, there are three remaining terms in computing substitute using Lemma 9, as:
d. DDL which we can further
d DKL (
d 1 - f d - log - -fB pr dt
p.
logğr. नM,
dr
(46)
= 7rir (f logø, - Ap, log log qr - -f |p, qr
Aq,
29

---

Using integration by parts, the first term in (46) is changed to:
f Ap, lo
d Opr
Σ dy
log
(v)
Vp- V
log
The limits in the first term become zero given the smoothness and fast decay properties of p. (g)) The remaining term can be further simplified as:
/ Vog..
/ V log / p V. Vog Vo.g...
The second term in (46) can be manipulated similarly, by first using integration by parts to get:
/ Ap, dr
Σ
0p.
log 9r
Vp, 7 log dr.
py:
==
Applying integration by parts again to Vp
V log dr, we have:
/ Vp lo
alog qr
f Pog.og
Pr
Oy:
The limits at the boundary values are all zero due the smoothness and fast decay properties of (v)) Now collecting all terms, we have f Pr log Pr = - PrIV log pr|2 and | p log 4r
J P log 4r. Thus (46) becomes:
d (p. (Tă+) -7tr / p. (IVlogp,1* Alog बr
Adr 9r
dr
Combining with (45), this leads to the following:
d DrL (5r (1âr) -FG. (.)))
d
(47)
dr
Recall that (-) and (-) were the densities of the sealed random variable 2) R, This leads to (X )dK, pr (D))d)), Thus, it is straightforward to show that both KL divergence and Fisher divergence are invariant to the scaling the underlying random variables. For KL divergence have
DEL (P))(())))))) Pr (2+) log (2)) d2), = / (tt) log (x (2),)
d3, = () q+(x+
(48)
And for Fisher Divergence we can write
()(p /b((0.)
Vp. (2),))
(2)(
dD)
(2))
(20))
/ (x.)
Vpr(E,)
(49)
d3
(3E+)
(x{,)
= D+(p+(-) () (+)))
Thus we can replace the divergences in (47) with those of the non-scaled distribution, which concludes the proof.
30

---

D.2
Proof of Lemma 5
Proof. We start with direct application of Lemma
(po((po() (190 - f"' (1)
(
|| VV v Voog 9r (z) dr
r=0
(50)
7 7T
[1sp(1,T))
s,(2, m) dr
Qt E Pa [|sp(E,7) Tx)
[) dr
In the second line we used the fact that Po(-) = go(+) = (0, 1). therefore DRI. (xx.(p()))ll(())) The third line follows from our definition of the score functions. Finally, in the last line we used the fact that 7=7= or which follows from = ST/a, and a + c 1:
Pr7
dr ar ar Sr à
(51)
- à,a
à a??)
at
dr
by denoting W, = we conclude the proof.
D.3 Proof of Lemma
Proof. In [33] they prove this result assuming drift term that only depends For modifi- cation, we begin by defining the time-dependent drift b, : IRd + Rd of the diffusion process (19)
(ăr 2
(árr
dr of 20xT
s(a,)))
(x)
Assumption 3. The drift (-) satisfies the following properties for all times [0, 1].
1. Lipschitz drift term.There is a finite constant L1 such that
[bx())) b-(u)l22 & Ia || ul12 for all x. y E TRd.
(52)
2. Smooth drift term. There is a finite constant Le such that
[br(x) - (brly)loo s T2 || - -ll2 for all x. E TRd
(53)
3. Distant dissipativity. There exist strictly positive constants p. B such that (z). < - | + B for all x Rd.
(54)
4. Time-continuous drift term. There is a finite constant La such that
0b, 0b,(z) (x)
< La for all Rd.
(55)
θT
There is an additional assumption in [33] on the smoothness of the initial densities of the continuous and discrete processes. our case both are the standard Gaussian which satisfies the assumption. We do not provide the whole proof here would consist of almost the entirety of [33]. We focus
31

---

small part of the proof, that is the only part that changes when we use time-dependen drift b, (z) as opposed to [33] where they assume -ndepenndend drift b(z).
Consistent with their notation, we define a continuous time diffusion process:
dx, b, (X,)dT dB,;:
(56)
and its Euler- Euler-Maruyee Maruyama discretization parameterized by step size > 0 (In our case = +):
((++)) Rnn + nbbnn (Xxx) + VTka.
&x (N(,
(57)
Furthermore they construct a continuous time stochastic process over the interval 7 E [1,, (I + 1)n) that interpolates (57):
! kon
f dB,
X, = Xkn +
bky (Knn) +
(58)
kn
Then they prove that the densities of the two continuous time processes given by (56).(58) denoted as Ty and म,, respectively satisfy the following (Lemma from [33]):
d
1 / (x) (x) - b, (=0) dx.
DEL a+ || TT))
(59)
dr
2 Rd
where (xx) = E bkyy kn)) x where the expectation is over the process (58). Then they proceed bound the norm inside the integral. The next equation based on equation (18) from [33] is where the time dependence of the drift term in our case enters the picture.
Br (x)) (x) =
bee (Xxy)
1 8r, = = - (x)
E bknn X 18, = 2) (bm)(()) +
(b,(z)
(T kn)) + +0( -kn)²)
0b, (x)(
bun (Nmx) baon (8+) = 2
(T kn) + +0((T - kn))
=
r=
(60)
They prove that the first term in (60) is (n) and from our additional time-continuity requirement for the drift Assumption 3, the second term is also (n)) (Note that 7 E [K77, (k 1)y) thus (r kn) can be at most 7)). With this, the rest of the proof from [33] goes through.
D.4 Proof of Lemma
Proof. We first prove a similar relation for generic distributions x(z), p(z) and their perturbations (x), p(z);
Where it clear from the context, we omit the integration variables. Perturbing the first argument gives us:
( (1) DEL I)]
म log
7 - -f क
/loo log
/ @ log
7 log
p
D + / TT) log
p
(61)
< (f [| + max (1os ( (m)) | -m
= ( [| + 2 log (Tv(
32

---

where log M max, log( [(x) p(a)
and drv denotes the total variation distance between distributions.
Next, perturbing the second argument we get:
/ π log
π -
T (
(f (1 DRL || =
$ Flog
p
- ft log
L
/
p P-P
Ft log
1
p p
π
p)
(62)
π p
I max (m)
/
10 pl
= 2M drv PP.
Using (61), (62) we get:
DKL ( p) (x [l
z [x [ DxL + | (x ll))
I D
x)) + 2M drv (p, p) + 2 M dIv(
S (
||
π) + 2M
DKu(0llP) 2 M V DKL( m)
2
(63)
where in the last line we utilized Pinsker's inequality bound the TV distance with the square root of the KL divergence. Now we apply (63) to diffusion models:
|Da (po(:: sp) ll (990(18)) - (pI sp) ll [o(0 sg)) S (Po(s88()8p)))) |l 90((;8p))
2M
(po(:ss))) (19((=;.)))
log M / DKi PkI (1
(64)
Furthermore from Lemma 6 know:
DKL (po(p sp) || 90(-ss)) I c/T?, DKL 54) (90(;ss)) c/T²
(65)
Putting together (64) and (65) we get:
|Di (Po sp) (9(()8g) - DKI (po(<)p)))) 40(55 s,))1 S (T)
(66)
where (T) c/T2 (2M 2og M) /c/T2, The second term dominates therefore 61 (T)
O(1/T) which concludes the proof.
D.5 Proof of Lemma
Proof. There are two sources of error need to consider. First we bound the error in approximating an integral with a sum:
5"' wr E. p. (1180) 1=
(((xx,7)] - Σ L'' f(r)dr 1(/T)) 1
-
T wt/T E. p))((
[|sp(x, (x
==
t=0
df
sup T (0.1)
dr
33

---

where have defined f(T) Bx (138p)) [|sp((,,)) (,(x,))11
We now upper bound the
supremum to show that it is finite:
dr dr ((P (1sp(1) - sy(2, ()1,,))122 dr
df d
/ d (pr(s:s))) (1sp(1,7) sq(21 1)2 dx
(67)
d / dr (1sp((,,)) sz(2, 7)1
We bound each term in (67) separately. Then the first term in (67) is bounded because dr (x:; finite as characterized Lemma The second term in (67) we expand further:
d
d / pr (pp)) dr (11sp(2, sy(a, T) 11
/ (p=(1;;p)) (32(1,7) s,(u, T),
dsp(a,t) dr
(sp(xzTT) cle
dr
2sup [|sp(I,)) ( ()12 dr (sp(2, -s(((s))))
2,r
2B B2.
The second source of error is replacing expectation the continuous time marginal Pt/r(1 8p) with expectation over the discrete time marginal (=(3;8p) which we can bound by using the fact that the two aforementioned marginals are close to each other.
T T W E. payr(( 19 [1ss((r, t) - $2 1)]
T - E t=
Wr/T Ey mel-sse)
[1sp(())) () - (2 1)
T
sup [1 7) so(2ET)112
T Wt/T T BB
T
T T2 B?
OF
where we used Lemma to get the last line which concludes the proof.
34

---

E Composition with Forward KL Divergences
We start with the constrained problem formulation using forward KL divergence (UF-C) which we rewrite here:
minimize U
(68)
subject to (g' (1 S for i 1,.... m.
In the case of diffusion models, the KI. divergence in (68) becomes the forward path-wise KI. between the processes:
minimize u.
(69)
subject to DxL (1o ((por(ppp())))) S 2 for i ,,...,77. m..
is important to note here that using the forward KL as a constraint makes sense when q' represent forward diffusion processes obtained by adding noise to samples from some dataset. We can also solve this forward KI. constrained problem to compose multiple models; In that case we treat samples generated by each model separate dataset with underlying distribution % (3r0).
In summary, the two key differences of Problem (69) to Problem (UR-A) are: (i) The closeness of model pretrained model q¹ is measured by the forward KL divergence (x p). instead of the reverse KL divergence DRL. (p || q²); (ii) The distributions {q') Im can the distributions underlying datasets, not necessarily m pretrained models.
Regardless of whether the q¹ represent pretrained models datasets, evaluating (-):(()) (-) (+)r(:;)) is intractable since it requires knowing (.r(()) which in turn re- DKL DKL(gor(
quires knowing (-(-) exactly. To get around this issue we formulate a closely related problem to (69) by replacing the KL with the Evidence Lower Bound (Elbo):
minimize subject to Elbo(uh.r::oor) S for i 1,...,71 m
u
(70)
where the Elbo is defined as
Elbo(qu::: Po:T) Exgrggo (x( log (o:T(
(71)
g( xx)
We note that the typical approach to train a diffusion model is minimizing the Elbo. Fur- thermore, minimizing Elbo(go:T; Po.r) over p is equivalent to minimizing the KL divergence
DKL. ( || ()) since they only differ by constant that does not depend on (see [22] for more details on this)
For a given X, we define a weighted mixture distributions as
A ATT d'()
(2)
Σ
(72)
and we denote by (g) the differential entropy of given distribution g.
H(q) = -Es-ollog q(x)].
(73)
Theorem 5. Problem (70) equivalent to the following unconstrained problem:
minimize D (x*) I| p)
(74a)
P
where the optimal dual variable given by X* argmax, > o D()). The dual function has the explicit form, D(A) H (g Furthermore, the optimal solution of (7) is given by
(X*) 0mix
p
(74b)
Unlike the reverse KL case, here we can characterize the optimal dual multipliers, and the optimal solution further; Note that the optimal dual multiplier X argmax,> o D(A) D(A) argmaxx2 o H (gmix (; (;)+)) is one that maximizes the differential entropy H() of the distribution of
35

---

the corresponding mixture. This implies that the optimal solution is the most diverse mixture of the individual distributions.
There are many potential use cases where we may want to compose distributions that don't overlap in their supports; For example when combining distributions of multiple dissimilar classes of dataset. The following characterizes the optimal solution in such settings.
Corollary 1. For the special case where the distributions q' all have disjoint supports, the optimal dual multiplier X² of Problem (70) can be characterized explicitly
eB(y')
X;
E) e(((s')
Algorithm Details
F.1
Alignment
Recall from Section 3.2 that the algorithm consists of two alternating steps:
Primal minimization:
At iteration n, we obtain new model g(++1) via a Lagrangian maximization:
(x++1) argmin LALI(Sps (m)).
SES
Dual maximization: Then, use the model 1(++1) to estimate the constraint violation Exg [r(o))] b, denoted as rr r(s(++1))- b, and perform dual sub-gradient ascent step:
(m+1)
(()) + (r(s(n-1))bb)) =)
In practice we replace minimization over S with minimization over a parametrized family of functions Sp. The full algorithm is detailed in Algorithm 1.
Algorithm 1 Primal- Dual Algorithm for Reward Alignment of Diffusion Models
1: Input: total diffusion steps T, diffuusin parameter a,, total dual iterations H, number of primal steps per dual update N, dual step size na. primal step size 7ps initial model parameters ((0).
2: Initialize: ((1) = 1/m.
3: for h 1, H do
4: Initialize 8 = ((h - 1)
5: for 1, 1.... N do
6: Take a primal gradient descent step:
Op+1 On Tpp VOLALI(O, (m)).
(75)
7:
end for
8:
Set the value the parameters be used for the next dual update: 0(h) O
9:
Update dual multipliers for i = 1, 1................ m.
Aahh 1 = (h) + (o [ri(xo)] [r( b:)]
(76)
end for
We now discuss the practicality of the primal gradient descent step (75) regarding the Lagrangian function,
L)) DKI. (Pae(Porr(s so) go:r (::r(:))))) (::s))
Σ AE po(:s0)) [n(sss) 6z). (77)
[ra(2)]]
To derive the gradient of (O 0, A), we first take the derivative of the expected reward terms by noting that the expectation is taken over distribution that depends on the optimization variable We can use the following result (Lemma 4.1 from [17]) to take the gradient inside the expectation.
36

---

Lemma 11. If po (Zo.r (r())) and (r are continuous functions of B. then we can write the gradient of the reward function
V Vol.ng po( r(ao)) Bear (or(11))))))))
r(zo) tm
Vo log p x4-
I4 s0)
For the gradient of the KL divergence, we have
VoDKL (::r((;8)) 80) (o:r(:;88))
(2e,)) - (g t)] 1)
V
Ex (p(
20
Vo
Σ
E (1600)
(KL ar 80) || q mx 8q))]
Bx - (((00)
[oDKL(p(2)111 11 1
T
T
Egg ve(
Σ Lest
Vo log zez (xL -2; 80) [ m2; 141))
For simplicity, we omit the second term in practice, has negligible effect on performance. See [17, Appendix A.3] for the derivation.
F.2
Composition
For composition, we take similar approach to Algorithm 1. Recall from Lemma 3 that the Lagrangian minimizer for the constrained composition problem can found by minimizing
T
LAND (0, A)
Σ
E 2o
Bux Be.
g(as mo) [1so(1e,) t) V log ((2))(20)))
U
Thus, we detail the algorithm for composition Algorithm 2.
Algorithm Primal-Dual Algorithm for Product Composition (AND) of Diffusion Models
Input: total diffusion steps T. diffusion parameter at, total dual iterations H, number of primal steps per dual update N, dual step size Vas primal step size 7pp initial model parameters 0(0).
2: Initialize: X(1) = 1/m.
3: for h h 1,
H do
4: Initialize 0, = 0(h - 1) 5: for 1.... N do
6: Take a primal gradient descent step:
0 0, Tp (O
(78)
7:
end for
8:
Set the value of the parameters be used for the next dual update: 0(h) ON+1.
9:
Update dual multipliers for i = 1, ..... m.
(h+1) = (h) + (a (po((: ll +=)=)
(79)
X(h + 1) proj (x) where proj(y) projects its input onto the simplex X 1.
11: end for
The projection the dual multiplier vector (line 10) ensures that À AT 1, as required when maximizing the dual function (see the proof of Theorem 3).
37

---

Note that Algorithm 2 implicitly requires samples from the weighted product distribution LAND (-) in order to minimize the Lagrangian LAND(e, X). We obtain these samples using the Annealed MCMC sampling algorithm proposed in [14].
Skipping the primal. As discussed in Section both Annealed MCMC sampling and the mini- (A) mization of the Lagrangian LAND( (0, X) at each primal step-to to match the true score V log 9AND -rre challenging and computationally expensive. Therefore, for all settings except the low-dimensional case described in Appendix G.2, we employ Algorithm 3, which skips the primal step entirely.
In Algorithm& 3 we bypass the primal steps by using the surrogate product score, rather than the true score, to compute the point-wise KL. used in the dual updates. The distinction between the true and surrogate scores is discussed in detail in [14].
true score:
log LAND, (ft) V 'log
(co(x
(80)
log CAND, (xx) = Σ A log (f (oo
surrogate score:
(81)
Algorithm 3 Dual-Only Algorithm for Product Composition (AND) of Diffusion Models
1: Input: total diffusion steps T, diffusion parameter au, total dual iterations H, dual step size V)-- 2: Initialize: A(1) = 1/m.
3: for h 1,. .... H do
4: Update dual multipliers for i = 1,. m.
2,(h + = (1) + Na TAND (.) (
(82)
5: X(h 1) proj (7hh+ + 1)), where proj(y) projects its input onto the simplex X X 1.
6: end for
For given X, the surrogate score can be easily computed:
v log TAND, (x+)
Σ A log log
Σ A.V.oog
and thus we can use Lemma 2 to compute the point-wise KLs needed for the dual update. As for the samples needed from the true product distribution, we also replace them with samples obtained by running DDIM using the surrogate score.
38

---

G Additional Experiments and Experimental Details
G.1
Related work
Here we review related work and explain why these approaches are not directly applicable as baselines for our experiments.
In [40] they propose a superposition method to sample from the mixture of diffusion models with arbitrary weights. However, they only use equal weight mixtures and don't ' discuss different weights. They also devise a method to sample points that have equal likelihood under different models which fundamentally different to the product composition that discuss this work.
Existing works including [11, 27, 34, 58] discuss constrained sampling from diffusion models, but the nature of their constraints is completely different from work it mainly involves sampling from a constrained set and they propose to do this through projection onto feasible set at each diffusion time step. It is not clear how apply these methods to reward constraints or how to them to preserve distance to a model.
Other works like [18, 8] enforce very specific constraints by adding additional losses with fixed weights to the objective which implicitly enforces the constraint. These methods are very specific to the constraints they are designed for and do not generalize to arbitrary reward functions and don't t give us a way to constrain closeness to a model.
G.2 Low-dimensional synthetic experiments
To visually illustrate the difference between the constrained and unconstrained approaches, conduct experiments where the generated samples lie R.?. For the score predictor we used the same ResNet architecture as used in [14].
Product composition (AND). Unlike the image experiments, in this low-dimensional setting we use Algorithm 2 for product composition. See Figure for visualization of the resulting distributions.
Mixture composition (OR). For this experiment used the same Algorithm the used in [22] for mixture of distributions. The only modification is doing an additional dual multiplier projection step similar the last step of the product composition Algorithm 2. See Figure for visualization of the resulting distributions.
G.3 Reward product composition (Section 5.2 (I)
Reward
Acsthecc
Imagermard
0.14
0.00
0.06
0.04
0.52
0.00
Equal Whights
Consbialhod
Figure 6: KL divergence for the product composition 5 adapters pretrained with different rewards. Error bars denote the standard deviation computed across 8 text prompts each with four samples.
Implementation details and hyperparameters. We finetumed the model using the Alignprop [35] official implementation for each individual reward using the hyperparameters reported in Table 2. We then composed the trained adapters rumning dual ascent using the surrogate score as described in section F.2. We the average of scores (denoted as "Equal weights") baseline. Hyperparameters are described in Table 3. The reward values reported Figure were normalised so that 0%
39

---

corresponds to the reward obtained by the pretrained model, and 100% the reward obtained by the model finetuned solely on the corresponding reward.
Additional results. As shown in Figure 6, equal weighting leads to disparate KI. divergences across adapters in particular high KL. with respect the adapter trained with the "aesthetic" reward while our constrained approach effectively reduces the worst case KL, equalizing divergences across adapters. Figure 13 shows images sampled from these two compositions exhibit different characteristics, with our constrained approach producing smoother backgrounds, shallower depth of field and more painting-lit like images.
Hyperparameter
Value
Batch size Samples per epoch
64
128
Epochs
10
Sampling steps
50
Backpropagation sampling
Gaussian
KL penalty
0.1
Learning rate
1 x 10-3
LoRA rank
4
Table Hyperparameters used to finetune models using individual rewards.
Hyperparameter
Value
Base model
{"cheetah", "snail", "hippopotamus",
Prompts
"crocodile", "lobster", "octopus")
Resolution
512
Batch size
4
Dual steps
5
Dual learning rate
Sampling steps
25
Guidance scale
5.0
Rewards
aesthetic, hps, pickscore, imagereward, mps
Table Hyperparameters for product composition of models finetuned with different rewards.
G.4 Concept composition (Section 5.2 (II))
present additional results for concept composition using three different concepts (as opposed to just in the main paper and in [40]) As seen in table 4, our approach retains clear advantage in both CLIP and BLIP scores. See Table 5 for examples of images generated using cach method. Images with the constrained method typically do better job of representing all concepts.
Min. CLIP (f)
Min. BLIP (1)
Combined Prompting 21.52
0.206
Equal Weights
22.18
0.203
Constrained (Ours)
22.45
0.221
Table 4: Comparing constrained approach to baselines minimum CLIP and BLIP scores. The scores are averaged over 50 different prompt triplets sampled from a list of simple prompts.
40

---

Combined Prompting
Equal Weights
Constrained
Table 5: Concept composition examples for each method. Prompts used for each row:
Row 1: "a pineapple", "a volcano". Row 2: "a donut", "a turtle". Row "a lemon", "a dandelion".
Row 4: a dandelion". "a spider web". "a cinammon roll".
G.5 Concept composition for text-to-audio diffusion models
We note that our proposed framework and theoretical analysis do not depend on any specific modality task types. From theoretical guarantees, would expect experiments in other modalities to provide results similar to those presented for images. To validate this, we conduct concept composition experiments with text-to-audio diffusion model example of another modality. We treat a text-to-audio diffusion model (in this case AudioLDM [30]) conditioned on different inputs, each representing a concept, the models be composed. We apply our constrained learning to find the optimal weights to compose these two models, and use the CLAP score [15] to measure the similarity between the generated audio samples and the text prompts representing each model.
Min. CLAP Score(t)
Combined Prompting Equal Weights Constrained (Ours)
0.816
1.57
1.92
Table Minimum CLAP scores across prompts for each method
Similar to concept composition for images, we observe in Table 6 that using constrained approach, the minimum CLAP score across prompts increases compared to the two baselines. The constraints ensure closeness to each model, which in turn results in more equal representation of the concepts.
41

---

G.6
Alignment experiments
Reward normalization. In practice, setting constraint levels for multiple rewards that both feasible and sufficiently strict to enforce the desired behavior is challenging. Different rewards exhibit widely varying scales. This is illustrated in Table 7, which shows the mean and standard deviation of reward values for the pretrained model. This issue can be exacerbated by the unknown interdependencies among constraints and the lack of prior knowledge about their relative difficulty sensitivity.
In order to tackle this, propose normalizing rewards using the pretrained model statistics simple yet effective heuristic. This normalization facilitates the setting of constraint levels, enables direct comparisons across rewards and enhances interpretability. In all of our experiments, apply this normalization before enforcing constraints, Explicitly, we set
Ppre
(83)
δ pre
where r denotes the original reward and Prres Pre the sample mean and standard deviation of the reward for the pretrained model. We find that, with this simple transformation, setting equal constraint levels can yield satisfactory results while forgoing extensive hyperparameter tuning.
Reward
Mean
Std
Aesthetic
5.1488
0.4390
HPS
0.2669
0.0057
MPS
5.2365
3.5449
PickScore
21.1547
0.6551
Local Contrast
0,0086
0,0032
Saturation 0.1060
0.0706
Table 7: Mean and standard deviation reward values for the pretrained model.
The effects varying the constraint thresholds.
What we observed by varying the reward constraint thresholds in our experiments was that for thresholds up to (i.e. Tyr + 1.0 ruee for each reward) the model was typically able satisfy the constraints with minimal violation. Another trend that we observed was that increasing thresholds usually leads to constraints that are harder to satisfy leading to higher Lagrange multipliers and resulting in higher KL to the pretrained model. See Tables 8, 9 below.
An advantage of our constrained approach is that Lagrange multipliers give information about the sensitivity of the objective with respect to relaxing the constraints i.e. if the multiplier for a certain reward ends up being much higher than the rest it means that constraint is particularly harder to satisfy. Consequently, even slightly relaxing the threshold for the corresponding reward can lead to much smaller KL objective.
Constraint
Threshold
Slack
Dual Variable
DKL
contrast
0.250
-0.245
0.282
0.177
contrast
0.500
-0.985
0.000
0.296
contrast
1.000
-0.381
0.000
0.332
saturation
0.250
-0.126
0.081
0.177
saturation
0.500
0.060
0.006
0.296
saturation
1.000
0.052
1.195
0.332
Table 8: MPS reward alignment with saturation and contrast constraints, for varying thresholds.
42

---

Constraint
Threshold
Slack
Dual Variable
DEL
contrast
0.250
-0.684
0.000
0.136
contrast
0.500
-1.011
0.000
0.109
contrast
1.000
0.661
0.192
0.293
saturation
0.250
-0.025
0.014
0.136
saturation
0.500
-0.060
0.000
0.109
saturation
1.000
0.062
1.020
0.293
Table 9: Pickscore reward alignment with saturation and contrast constraints, for varying thresholds.
I. MPS + local contrast, saturation.
In this experiment, augment standard alignment poss-trained on user preferences- -with two differentiable rewards that control specific image characteristics: local contrast and saturation. These rewards are computationally inexpensive evaluate and offer direct interpretability in terms of their visual effect on the generated images. In addition, the unconstrained maximization of these features would lead undesirable generations. other potentially useful rewards not explored in this work brightness, chroma energy, edge strength, white balancing and histogram matching.
Local contrast reward. In order to prevent images with excessive sharpness, we minimize the "local contrast", which we define as the mean absolute difference between the luminance of the image and low-pass filtered version. Explicitly, let Y denote the luminance, computed as Y 0.2126R+ 0..1122G + 0.0722 and G.+Y +Y the luminance blurred with a Gaussian kernel of standard deviation 1.0. We minimize the average per pixel difference by maximizing the reward
HW Σ Yo - (Go Mis
TC
where H, W denote image dimensions.
Saturation reward. To discourage overly saturated images, we simply penalize saturation, which we compute from R, G, B pixel values as
(c) ma F. mR.G. Fi.
rs
HW
(R..E( Zij +
4
where & = x 10-8 80-8:0 is small constant added for numerical stability.
Implementation details and hyperparameters. We implemented primal-dual alignment ap- proach (Algorithm 1) in the Aliguprop framework. Following their experimental setting, we use different animal prompts for training and evaluation. Hyperparameters are detailed in Table 10.
Hyperparameter
Value
runvayml/stable-
Base model
diffusion-vi-5
Sampling steps
15
Dual learning rate
0.05
Batch size (effective)
4 x 16 = 64
Samples per epoch
128
Epochs
20
KL penalty
0.1
LoRA rank
4
MPS: 0.5
Constraint level
Saturation: 0.5
Local contrast: 0.25
Equal weights
0.2
Table 10: Hyperparameters for reward alignment with contrast and saturation constraints. Constraint levels correspond to normalized rewards.
43

---

Additional results. We include images sampled from the constrained model in Figure 14 for hps and aesthetic reward functions. Samples from model trained with an equally weighted model are included for comparison. Constraints prevent overfitting to the saturation and smoothness penalties.
II. Multiple aesthetic constraints
Implementation details and hyperparameters. We modified the Alignprop framework to acco- modate Algorithm 1. Following their setup, use text conditioning prompts of simple animals, using separate sets for training and evaluation. this setting, due to the high variability of rewards throughout training, utilized exponential moving average reduce the variance in slack estimates (and hence dual subgradients) [41]. Hyperparameters are detailed in Table 11.
Hyperparameter
Value
runwayml/stable
Base model
diffusion-vi-5
Sampling steps
15
Dual learning rate
0.05
Batch size (effective)
4 x 16 64
Samples per epoch
128
Epochs
25
KL penalty
0.1
LoRA rank
4
MPS: 0.5
HPS: 0.5
Constraint level
Aesthetic: 0.5
Pickscore 0.5
Equal weights
0.2
("cat", "dog", "horse", "monkey", "rabbit", "zebra" "spider", "bird", "sheep", "deer" "cow", "goat" "lion", "tiger", "bear", "raccoon", "fox", "wolf" "lizard", "beetle", "ant", "butterfly", "fish", "shark" "whale", "dolphin", "squirrel", "mouse", "rat", "snake" "turtle", "frog", "chicken", "duck", "goose", "bee" "pig", "turkey", "fly", "llama", "camel", "bat" "gorilla", "hedgehog". "kangaroo") ("cheetah", "snail", "hippopotamus", "crocodile", "lobster", "octopus")
Training Prompts
Evaluation Prompts
Table 1I: Hyperparameters for reward alignment with multiple rewards. Constraint levels correspond normalised rewards.
Additional results. We include two images per method and prompt in Figure 15. These are sampled from the same latents for both models.
G.7 Combining constrained alignment and composition
As mentioned in Section The constrained alignment and composition problem formulations can be combined. An example of this is composing reward-specialle models while enforcing minimum aggregate reward level. To demonstrate the viability of this approach, we conducted simple experiment in which we finetune a pretrained model using two KL constraints: one for pretrained stable diffusion and one for another model finetuned on the aesthetics reward, respectively, along with constraint on the saturation reward. The finetuned model achieves less than 10% reward constraint violation and similar KI. divergences with respect to both pretrained models, as seen Table We leave in- depth exploration of the combined Alignment and Composition problem to future work.
44

---

Constraint
Dual Variable
Initial Value
Final Value
Slack
Saturation
1.080
0.100
0.533
0.033
KL (pretrained)
0.493
0.000
0.101
0.004
KL (aesthetics)
0.507
0.260
0.097
0.000
Table 12: Results for finetuning a model with both expected reward and KL divergence constraints.
G.8 Computational details
experiments were run on single Nvidia A6000 GPU.
For alignment, there is little additional time overhead compared to baselines like AlignProp. For example, for the experiment Figure runtime is 33 minutes for both constrained and unconstrained methods, and for the experiments in Figure 4, constrained runtime is 64 minutes, unconstrained 60 minutes. Existing approaches already estimate the KL and sample batches evaluate and back-propp propaaate through the reward. The only additional computation for our method is the dual updates which is negligible terms of added time.
For composition, there no meaningful comparison to the equal weights baseline since the weights not learned in the equal weights baseline. For constrained composition, it takes around 5-10 dual updates for dual variables to converge which for composing the 5 finetuned stable diffusion models takes 9 minutes total and for concept composition it takes 2 minutes.
45

---

Equal Weights
Constrained
Table 13: Images sampled from the same latents for the product of adapters using the equal weights and when using the proposed KlL-constrained reweighting scheme using 5 dual steps.
46

---

Equal Weights
Constrained
Table 14: Images sampled from models finetuned to maximize MPS [60], along with sharpness and saturation penalizations. We compare optimizing equally weighted objective against constrained approach.
47

---

Equal Weights
Constrained
Table 15: Samples from models finetuned using multiple rewards with equal weights and with our constrained alignment method.
48


Reviews:
[{"review_id": "MVjhcYYwvF", "reviewer": "Reviewer_oCmM", "summary": "The paper proposes a unified constrained optimization framework for alignment and composition of diffusion models. Instead of weighting competing objectives, the proposed framework impose explicit constraints (e.g., minimum reward levels or maximum KL distance to each source model). The paper theoretically proves strong duality for both constrained optimization problems, characterize the solution, and implement a primal–dual training loop for score-based diffusion models. Experiments on Stable-Diffusion-v1.5 show improved balance across rewards and fairer representation of composed models relative to equal-weight baselines.", "questions": "1. What does each letter in abbreviations like UR-A stand for?\n2. Do direct reward alignment and composition of reward-specific models yield different outcomes in practice?", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \n1. The paper is well written and easy to follow\n2. The paper introduces a principled constrained optimization formulation that unifies reward alignment and model composition satisfying strong duality—bringing theoretical clarity and practical tractability to both problems.\n3. Compared to weight-based approaches, constrained formulations offer more interpretable hyperparameters and can eliminate tuning entirely in composition tasks.\n\nWeaknesses: \n1. While the constrained alignment formulation is built around explicitly enforcing reward thresholds, the experiments do not systematically analyze how varying these thresholds affects model behavior. For example, it would be informative to report sensitivity curves showing how sample quality, reward satisfaction, or KL divergence vary as the reward constraint levels are tightened or relaxed.\n2. (suggestion) It would be valuable to discuss whether a setting exists where both alignment to external rewards and composition of multiple source models are required simultaneously—for instance, composing reward-specialized models while enforcing a minimum aggregate reward level. Such a setting could further test the expressive power and scalability of the proposed framework.", "was_revised": true, "final_justification": "I lean towards acceptance. The paper is clearly written and introduces a principled constrained optimization formulation that unifies reward alignment and model composition under strong duality, offering theoretical clarity, practical tractability, and more interpretable hyperparameters than weight-based methods. While the empirical evaluation appears limited, I think the conceptual novelty and potential impact are sufficient to justify acceptance.", "rebuttal": "**Re: Weakness 1**\n\nWhat we observed by varying the reward constraint thresholds in our experiments was that for thresholds up to 1.0 (meaning mean + 1.0*std for each reward) the model was typically able to satisfy the constraints with minimal violation. Another trend that we observed was that increasing thresholds usually leads to constraints that are harder to satisfy leading to higher Lagrange multipliers and resulting in higher KL to the pre-trained model.\n\nAn advantage of our constrained approach is that Lagrange multipliers give information about the sensitivity of the objective with respect to relaxing the constraints i.e. if the multiplier for a certain reward ends up being much higher than the rest it means that constraint is particularly harder to satisfy. Consequently, even slightly relaxing the threshold for the corresponding reward can lead to much smaller KL objective.\n\n We thank the reviewer for the suggestion, and will include ablations on the constraint levels that illustrate the above points (see tables below).\n\nTable 1. MPS reward alignment with saturation and contrast constraints, for varying thresholds.\n| Constraint | Threshold | Slack  | Dual Variable | $𝐷_{KL}$ |\n| ---------- | --------- | ------ | ------------- | --------------- |\n| contrast   | 0.250     | -0.245 | 0.282         | 0.177           |\n| contrast   | 0.500     | -0.985 | 0.000         | 0.296           |\n| contrast   | 1.000     | -0.381 | 0.000         | 0.332           |\n| saturation | 0.250     | -0.126 | 0.081         | 0.177           |\n| saturation | 0.500     | 0.060  | 0.006         | 0.296           |\n| saturation | 1.000     | 0.052  | 1.195         | 0.332           |\n\nTable 2. Pickscore reward alignment with saturation and contrast constraints, for varying thresholds.\n| Constraint | Threshold | Slack  | Dual Variable | $𝐷_{KL}$ |\n| ---------- | --------- | ------ | ------------- | --------------- |\n| contrast   | 0.250     | -0.684 | 0.000         | 0.136           |\n| contrast   | 0.500     | -1.011 | 0.000         | 0.109           |\n| contrast   | 1.000     | 0.661  | 0.192         | 0.293           |\n| saturation | 0.250     | -0.025 | 0.014         | 0.136           |\n| saturation | 0.500     | -0.060 | 0.000         | 0.109           |\n| saturation | 1.000     | 0.062  | 1.020         | 0.293           |\n\n---\n\n**Re: Weakness 2**  \nAs the reviewer suggests, composing reward-specialized models while enforcing a minimum aggregate reward level is a great example of combining constrained Alignment and Composition. To demonstrate the viability of this approach, we conducted an experiment using two KL constraints: one for pre-trained stable diffusion and one for a model fine-tuned on the aesthetics reward, respectively, along with a constraint on the saturation reward. The fine-tuned model achieves less than 10% reward constraint violation and similar KL divergences with respect to both pre-trained models, as seen in the table of results below:\n\n| Constraint      | Dual Variable | Initial Value | Final Value | Slack  |\n|-----------------|---------------|---------------|-------------|--------|\n| Saturation      | 1.080         | 0.100         | 0.533       | 0.033  |\n| KL (pretrained) | 0.493         | 0.000         | 0.101       | 0.004  |\n| KL (aesthetics) | 0.507         | 0.260         | 0.097       | 0.000  |\n\nWe thank the reviewer for bringing up this important setting and will include the above example in the final version.\n\n---\n\n**Re: Question 1**  \nThe first letter indicates parameterization: **U** and **S** in UR-A and SR-A denote Unparameterized and Score function parameterized problems, respectively.  \nThe second letter, **R** or **F** (as in UR-C or UF-C), stands for Reverse or Forward KL divergence constraints.  \nThe final letter, **A** or **C**, signifies Alignment or Composition. We will clarify these abbreviations in the final version.\n\n---\n\n**Re: Question 2**  \nBased on our theoretical analysis of the solutions for the Composition and Alignment problems, it is straightforward to show that the solution in both cases will end up being in the form\n\n$$\np^\\star \\propto p_{\\text{pre}} e^{\\sum \\lambda_i r_i}\n$$\n\nbut with potentially different $\\lambda$'s.\n\nThere is no theoretical guarantee that the optimal dual multipliers, $\\lambda$, will be the same or similar for direct reward alignment and for the composition of reward-specific models. This is reflected in our experiments, where we observe that these two approaches lead to different final results in terms of expected reward values.\n\nRegarding where each approach would be useful, we note that reward alignment requires a reference model and access to each reward function, while composition can be used when we have access to models fine-tuned for specific rewards or tasks but do not have access to the rewards or task data themselves.\n\nWe thank the reviewer for raising this interesting question, and we will comment on this in the final version."}, {"review_id": "8gzOmGVaFX", "reviewer": "Reviewer_tobx", "summary": "This paper introduces a constrained optimization framework for aligning and composing diffusion models using reverse KL divergence and reward constraints. It leverages strong duality to design a primal-dual algorithm that enables fine-tuning with multiple rewards while preserving proximity to pretrained models.", "questions": "How sensitive is the method to the choice or scaling of reward functions, and how does the heuristic for automatic threshold selection perform across different tasks or domains?\n\nCan the proposed alignment and composition framework be extended to unconditional generation tasks, or is it inherently tied to reward-conditioned settings like text-to-image generation?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The constrained alignment method is principled and interpretable, offering clear control over reward satisfaction while avoiding overfitting.\n- The approach scales well to multiple reward signals, avoiding exhaustive hyperparameter searches required by weighted objectives.\n- The use of strong duality and primal-dual optimization brings theoretical rigor and enables efficient training with provable guarantees.\n\nWeaknesses:\n- The paper lacks experiments beyond text-to-image diffusion, limiting generalizability across modalities or task types.\n- The choice and tuning of reward functions is still nontrivial, and their quality directly affects alignment; this is underexplored.\n- While the method avoids explicit weight tuning, automated threshold selection is heuristic-based and may not always generalize well.", "was_revised": true, "final_justification": "My concerns and questions have mainly been addressed, and I raise my score to 4.", "rebuttal": "**Re: Weakness 1**\n\nOur proposed framework and theoretical analysis do not depend on any specific modality or task types. From our theoretical guarantees, we would expect experiments in other modalities to provide results similar to those presented for images in the paper. The reason why we focus on text-to-image diffusion is that it is currently the most widely used case of diffusion models: standard pre-trained models and reward functions for this modality are readily accessible. Furthermore, the reviewer seems to have focused mostly on Alignment but we emphasize that constrained Alignment, Product *and* Mixture Composition, and even the combination of composition and alignment (We give an example in our response to reviewer QwoA) together cover a wide variety of different tasks that our framework can be applied to.\n\nTo further address the reviewer's concern regarding generalizability, we have conducted concept composition experiments for text-to-audio generation as an example of another modality. We briefly outline the setting and our results below:\n\nWe treat a text-to-audio diffusion model (in this case AudioLDM) conditioned on different inputs, each representing a concept, as the models to be composed. We apply our constrained learning to find the optimal weights to compose these two models, and use the CLAP score [1] to measure the similarity between the generated audio samples and the text prompts representing each model.\n\n**Minimum CLAP scores across prompts for each method:**\n\n| Method             | Minimum CLAP Score |\n|--------------------|-------------------|\n| Combined prompting | 0.816             |\n| Equal Weights      | 1.57              |\n| Constrained (Ours) | **1.92**          |\n\n\nSimilar to Section 5.2 (II), where we did concept composition for images, we observe that using our constrained approach, the minimum CLAP score across prompts increases compared to the two baselines. The constraints ensure closeness to each model, which in turn results in a more equal representation of the concepts.\n\n**Reference**\n\n[1] CLAP: Learning Audio Concepts From Natural Language Supervision (Elizalde et al.)\n\n---\n\n**Re: Weakness 2**\n\nWe provide a framework that can be used for Alignment with arbitrary reward functions. To evaluate the performance of our proposed approach, we measure whether each expected reward is higher than the chosen threshold and the closeness to the pre-trained model.\n\nWe agree with the reviewer's point that depending on the choice of reward functions, the generated sample distributions can be different; however, this is independent of the merit of our constrained framework. Hence, whether increased reward leads to generations with higher quality or not (and what quality means in this context) is not the focus of this work.\n\nWe do note that some commonly used aesthetic rewards tend to overemphasize certain qualities like saturation or contrast in images. This is another motivation for our constrained alignment approach: by constraining simple saturation and contrast rewards alongside the original aesthetic reward, we can mitigate these issues of the aesthetic reward (see Figure 3 in the paper).\n\n---\n\n**Re: Weakness 3**\n\nWe would like to clarify that our automated threshold selection is not purely heuristic, as it is informed by our theoretical analysis. There are two types of constraints that we consider in the paper, which we now explain below:\n\n- **KL divergence constraints in the constrained composition problem:**  \n  While our framework supports choosing fixed thresholds, we note that the values of the KL divergence are not as interpretable as rewards. This is why we present the epigraph formulation in the paper, which essentially finds a minimum shared threshold $u$ for all the KL constraints. This epigraph formulation could be viewed as \"automated threshold selection.\" We have theoretically analyzed the solution of this problem and discussed its properties, noting that the solution has a distribution equally close in KL divergence to each of the models, which is desirable for composition. Hence, we do not view it as a heuristic since we present a theoretical justification and analysis for it.\n\n- **Reward constraints:**  \n  We agree with the reviewer that choosing reward thresholds is more interpretable, avoiding exhaustive tuning of multiple reward weights. To ensure that the chosen thresholds lead to a feasible problem, we suggest normalizing the rewards using the mean and standard deviation of reward values estimated from the pre-trained model. This helps alleviate the issue of rewards that have different scales.\n\n---\n\n**Re: Question 1**\n\nThe choice or scaling of the reward functions mainly informs what the constraint threshold for each reward should be. As long as the chosen thresholds are feasible, the dual variables (i.e., the learned weights for every reward) converge to the values that result in satisfying all the constraints.\n\nWe note that depending on the reward functions chosen, different scales might make choosing feasible thresholds less straightforward. To alleviate this in practice, we propose to use reward normalization by estimating the mean and variance of reward values on a few batches sampled using the base pre-trained model. Then, we observe that setting the threshold for each reward to be (threshold = mean + alpha $\\cdot$ std) with $\\alpha$ between 0.5 and 1, usually works well in practice.\n\nRegarding \"automatic threshold selection\" for composition, we maintain that it is theoretically justified and is not a heuristic. We refer to our comments on **Weakness 3** for a detailed discussion. Furthermore, our experiments on constrained composition for audio diffusion models verify that the applicability of our approach (including the automated threshold selection) is not domain-dependent.\n\n---\n\n**Re: Question 2**\n\nThe proposed framework is not limited to conditional generation tasks. In fact, the illustrative examples shown in Figures 1 and 2 are from constrained unconditional models trained to generate data from Mixture of Gaussian distributions. For the main experiments, we focus on the text-to-image conditional setting, since it is a commonly used application of diffusion models.\n\nWe would appreciate clarification on what the reviewer means by \"reward-conditioned settings\" in this context. We note that the rewards we used for alignment include both rewards that measure similarity between text-image pairs (like HPS) and rewards that only take images as input (like saturation and local contrast rewards), which could be used for unconditional generation tasks."}, {"review_id": "2dcWrwAwOk", "reviewer": "Reviewer_QwoA", "summary": "This work proposes addressing the problems of reward alignment and composition in diffusion models through the lens of constrained learning via Lagrangian duality. The authors present a coherent formulation that yields both structural insights into the solution and practical algorithms for solving the optimization problem.", "questions": "Please see the weaknesses above. In summary:\n- Are there new technical contributions beyond the formulation itself?\n- Why is path-wise KL divergence used for alignment rather than divergence at the final distribution?", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strengths\n- The Lagrangian formulation offers a clean and principled approach to the problem, and the technical work appears solid.\n- The method allows users to specify thresholds instead of weights, which is arguably more intuitive and user-friendly.\n\n### Weaknesses\n- While the constrained learning perspective is appealing, the technical novelty appears limited. The analysis seems relatively straightforward, and it's unclear whether the proofs or techniques used contain significant new contributions. It would be helpful if the authors clarified this point.\n- Although the title and framing suggest a unified treatment of reward alignment and product composition, the paper handles these two problems separately. The connection between them seems limited to the fact that a similar analytical approach is applied in both cases.\n- The motivation for using the path-wise KL divergence in alignment is somewhat unclear. Similar to composition, the final distribution $p_0$ seems more important, so why is the path-wise divergence more appropriate? Furthermore, the formulation appears to rely on a discrete-time variant; can this be extended to continuous time?\n\n### Minor\n- Section 4.1, line 211: I believe it should be problem (UR-C), not (UR-A).", "was_revised": true, "final_justification": "Overall, this work on reward alignment and composition using a lagragian formulation is interesting and the work is solid, but not particularly remarkable in terms of novelty and potential impact. As such, I have recommended a borderline accept.", "rebuttal": "**Re: Weakness 1**  \nAn integral part of our approach lies in Lemma 2 which allows us to evaluate the KL divergence between marginals of a backward diffusion process (i.e., point-wise KL). As far as we know, our work appears to be the first to provide a tractable characterization of point-wise KL and highlight its difference to path-wise KL.\n\nWe argue that in composition, constraining the KL divergence between the final marginal distributions at $t = 0$ is more natural, since we care about distributions of the generated samples and not the whole diffusion process. At first glance, computing this KL seems intractable, since in general it would require marginalizing the backward diffusion processes with arbitrary score functions to derive $p_0$.\n\nWe would like to highlight that the proof is quite involved as it requires characterizing the derivative of the KL between marginals in continuous time, and then bridging the gap to discrete time and using properties of the KL along with novel technical analysis to bound the discretization error. We expand on this further below.\n\nFurthermore, we would argue that even some parts of the analysis that are relatively straightforward from a technical standpoint (e.g., strong duality), are still important and novel, since the results give new insight into the solutions of the constrained problem, connecting our work to the broader model composition and alignment literature.\n\n**Proof Summary of Lemma 2:**  \nIn the proof, we begin with the intuition that if the initial distributions at time $t = T$ are the same and the two backward processes are close, then the final distributions at time $t = 0$ will also be close. We then generalize a result from [1] to formalize this in continuous time, giving us an analogue of Lemma 2 in continuous time.\n\nThen, we use results from the literature on discretization of Langevin dynamics along with additional theoretical analysis of our own to prove that the continuous time point-wise KL computation can be extended to work for discrete time diffusion processes, where we characterize the discretization error to be of order $O(1/T)$. We refer the reviewer to Appendix C.2 in the supplementary material for more details.\n\n**Reference**\n\n[1] (S. Lyu) Interpretation and generalization of score matching\n\n---\n\n**Re: Weakness 2**  \nBy unified treatment, we mean that the same constrained learning framework can be applied to both problems. Beyond the similar analytical approach, we also use essentially the same primal-dual algorithm to tackle both, the only difference being a workaround in composition to minimize the Lagrangian. Furthermore, as we briefly mention in the last paragraph of Section 2, the problem formulations for composition and alignment should be viewed as canonical formulations, and similar approaches can be readily adapted to solve combinations of these, for example:\n\n$$\n\\begin{aligned}\n\\min_{u \\geq 0, p} \\quad & u \\\\\n\\text{subject to} \\quad & D_{\\mathrm{KL}}(p \\Vert q_i) \\leq u \\quad \\text{for } i = 1, \\ldots, m, \\\\\n& \\mathbb{E}_{x \\sim p} [r_j(x)] \\geq b_j \\quad \\text{for } j = 1, \\ldots, n.\n\\end{aligned}\n$$\n\n\nwhere the goal is to find a model that is close to multiple pre-trained models and also satisfies one or more reward constraints.\n\nIn order to emphasize this unified treatment, we present an example of the combined constrained alignment and composition formulation above. We conducted an experiment using point-wise KL constraints to pre-trained Stable Diffusion and a model fine-tuned on aesthetics reward, along with a constraint on the saturation reward. The resulting constrained model achieves less than 10% reward constraint violation and similar KLs with respect to both pretrained models.\n\n| Constraint      | Dual Variable | Initial Value | Final Value | Slack  |\n|-----------------|---------------|---------------|-------------|--------|\n| Saturation      | 1.080         | 0.100         | 0.533       | 0.033  |\n| kl (pretrained) | 0.493         | 0.000         | 0.101       | 0.004  |\n| kl (aesthetics) | 0.507         | 0.260         | 0.097       | 0.000  |\n\n---\n\n**Re: Weakness 3**  \nThe main reason is that computing point-wise KL in the alignment problem is intractable, at least with our suggested approach. However, in the alignment setting, the KL divergence objective acts as a regularizer to prevent too much deviation from the pre-trained model, and path-wise KL is commonly used to fulfill the role of a regularizer. Furthermore, under a mild assumption, it can be shown that using point-wise or path-wise KL leads to the same solution in alignment. We expand on these points below:\n\nAs mentioned in Lemma 2, computing point-wise KL requires 'valid' score functions, i.e., gradients of log-likelihood of some forward process.\n\nWe might reasonably assume that the score of a pre-trained model, $s_q(x,t)$, approximates a valid score function, since it was learned by minimizing the distance to the true score for some data distribution. However, for $s_p(x,t)$ which we are updating at each primal step of training, there is no longer a guarantee that it is a valid score function. This is why we need Lemma 3 which gives us a way to find the Lagrangian minimizer without explicitly computing the KL between $p_0(\\cdot; s_p)$ and $p_0(\\cdot; s_q)$.\n\nCrucially, Lemma 3 assumes we have access to samples from the Lagrangian minimizer distribution ($q_{\\mathrm{AND}}^{(\\lambda)}$ in the case of composition). This is okay in composition since we have ways of sampling from this Lagrangian minimizer. However, for alignment this is no longer the case, since we have no way of directly sampling from:\n\n$$\nq_{\\mathrm{rw}}^{(\\lambda)} \\propto q  e^{\\lambda^\\top r}.\n$$\n\nLuckily, unlike in composition, in alignment, since there is only a single KL objective, point-wise KL and path-wise KL lead to similar solutions:\n\n$$\np_0^\\star(x_0) = \\int q_{0:T}(x_{0:T})  e^{\\lambda^\\top r(x_0)}  dx_{1:T}\n\\approx \\int q_0(x_0)  q(x_1 | x_0) \\cdots q(x_T | x_{T-1})  e^{\\lambda^\\top r(x_0)}  dx_{1:T}\n= q_0(x_0)  e^{\\lambda^\\top r(x_0)}\n$$\n\nwhere we have used the fact that the pre-trained model $q_{0:T}(x_{0:T})$ has been trained to match a forward diffusion process; therefore, we assume it is approximately equal to the forward process which is Markovian in forward time and thus can be decomposed as above. Note that $q_0(x_0) e^{\\lambda^\\top r(x_0)}$ is the exact solution if we use point-wise KL as the objective.\n\nWe thank the reviewer for bringing up this point, since emphasizing the distinction between point-wise and path-wise KL is a central aspect of the paper.\n\nFor the second concern, yes, the formulation can easily be extended to continuous time. In fact, in the proof of Lemma 2, we begin by proving it in continuous time, and then bound the discretization error, extending the proof to the discrete time diffusion setting."}, {"review_id": "2eAZyO9K3v", "reviewer": "Reviewer_GAaS", "summary": "This paper proposes a constrained framework that unifies the alignment and composition of diffusion models, aimed at tackling the trade-offs arising when optimizing for multiple rewards or combining multiple models. The authors give theoretical results of aligning with multiple reward constraints and composition of multiple pretrained models under the canonical framework of Lagrange duality. Experiments show that aligned or composed model satisfies constraints effectively when using the proposed constrained method.", "questions": "See \"weaknesses\".", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The paper provides a comprehensive theoretical framework for constrained alignment and composition of diffusion models, including characterization of the solutions and strong duality results. The proofs are detailed and well-structured, which adds significant value to the work.\n\n\n2. The proposed method unifies alignment and composition tasks, offering a principled alternative to existing ad hoc approaches. \n\n3. The experimental results of finetuning alignment and product composition are consistent with theories, satisfying constraints effectively.\n\nWeaknesses:\n\n\n1. The strong duality claims in Theorems 2 and 4 are not rigorously justified. While the KL-regularized constrained optimization problem is convex in the space of diffusion trajectories, the authors project this result onto a non-convex parameter space of score networks without formally proving that the parameterized family is expressive enough to represent all feasible distributions. This step implicitly assumes a surjective mapping from score functions to path distributions.\n\n\n2. Lemma 1 and 2 assume that the diffusion models in comparison share the same noise and variance schedule. I think it is impractical. It would be valuable if the authors could discuss whether the results can be extended to more general settings, or at least clarify the practical implications of violating this assumption.\n\n\n3. The empirical evaluation is quite limited. The experimental section only compares the proposed method against equal-weight baselines for both alignment and composition. Although these baselines are simple and commonly used, they are far from state-of-the-art. For example, recent constrained or flow-based composition baselines are not compared, even though the related-work section cites them. \n\n\n4. For high-dimensional images the authors skip all primal updates and drive dual variables solely with a surrogate product score  to avoid annealed MCMC cost . However, the paper does not provide any theoretical or empirical analysis of the approximation error introduced by this shortcut.  As a result, it is unclear whether this practical modification substantially weakens the guarantees of the method.\n\n\n\n\n5. The method involves expensive components such as repeated KL divergence estimation between large diffusion models and, in the full algorithm, annealed MCMC sampling. The authors acknowledge that the approach is “computationally costly,” but do not provide any quantitative reporting of GPU hours, wall-clock time, or memory usage. It would be valuable if the authors could include such metrics.", "was_revised": true, "final_justification": "The theoretical contributions of the paper is solid. However, the empirical evaluation remains relatively limited, especially in comparison to other diffusion-related papers at NeurIPS. For this reason, I will keep my score (weak accept), but I will increase my confidence from 3 to 4.", "rebuttal": "**Re: Weakness 1**\n\nOur strong duality results show that there is no duality gap when the score networks are sufficiently expressive. This can be justified by the widely-used practice that utilizes over-parameterized networks (e.g., U-Nets or transformers) to train diffusion models. When Slater condition holds for the convex hull of score networks $\\overline{\\mathrm{conv}}(\\mathcal{S})$, we can apply the known strong duality in non-convex constrained learning (e.g., see Proposition 3.2 of [1]) to our constrained alignment and composition problems. Hence, mapping from score networks to path distributions is not necessarily surjective and strong duality holds, i.e., $P^\\star = D^\\star$.  \n\nWhen the score networks $s_\\theta$ are parametrized by a class of functions with limited expressiveness, we can use the duality gap analysis in non-convex constrained learning (e.g., see Proposition 3.3 of [1]) to characterize the duality gap\n\n$$\nP_\\theta^\\star - D_\\theta^\\star \\leq O(\\nu),\n$$\n\nwhere $\\nu$ is a parametrization gap that satisfies\n\n$\\|\\| p(\\cdot; s) - p(\\cdot; s_\\theta) \\|\\|_1 \\leq \\nu \\quad \\text{for any } s \\in \\mathcal{S}.$\n\nThank you for pointing out this important point. We will include all these points formally in the final version.\n\n**Reference**\n\n[1] (Chamon et al.) *Constrained Learning With Non-Convex Losses*\n\n---\n\n**Re: Weakness 2**\n\nIn both Lemmas, this assumption can be relaxed, leading to different, but still tractable sums. In Lemma 1, the path-wise KL can be decomposed into a sum of $T$ KL divergences:\n\n$$\nD_{\\text{KL}} \\left( p_{0:T}(\\cdot ; s_p) \\Vert p_{0:T}(\\cdot; s_q) \\right)\n$$\n\n$$\n= \\sum_{t=1}^T E_{x_{0:T} \\sim p}\n\\left[ D_{\\text{KL}}\\left(p(x_{t-1} \\mid x_t; s_p) \\Vert p(x_{t-1} \\mid x_t; s_q) \\right) \\right]\n$$\n\n\n\nImportantly, $p(x_{t - 1} | x_t; s_p)$ and $p(x_{t - 1} | x_t; s_q)$ are Gaussian distributions. If the two models share a variance and noise schedule, the KL between Gaussians simplifies to the squared norm of the difference between the means, which is what we present in the paper. However, if the noise and variance schedules are different, the KL between two Gaussians\n\n$$\np(x_{t - 1}| x_t; s_p) = \\mathcal{N}\\left(\\sqrt{\\frac{\\alpha_{t-1}}{{\\alpha}_t}} x_t + \\beta_t s_p(x_t, t), \\sigma_t^2 I\\right),\n$$\n\nand\n\n$$\np(x_{t - 1}| x_t; s_q) = \\mathcal{N}\\left(\\sqrt{\\frac{\\alpha'_{t-1}}{{\\alpha'}_t}} x_t + \\beta'_t s_q(x_t, t), \\sigma'_t{}^2 I\\right),\n$$\n\nis still tractable and given by:\n\n$$\n\\begin{aligned}\nD_{\\mathrm{KL}} \\big( p(x_{t - 1}| x_t; s_p) \\,\\Vert\\, p(x_{t - 1}| x_t; s_q) \\big) = \\frac{1}{2} \\Bigg[ & \\frac{1}{\\sigma_t^2}\\|\\|\\left(\\frac{\\alpha_{t - 1}}{\\alpha_t} - \\frac{\\alpha'_{t - 1}}{\\alpha'_t}\\right) x_t + \\beta_t s_p(x_t, t) - \\beta'_t s_q(x_t, t)\\|\\|^2  + d \\log \\left(\\frac{\\sigma_t^2}{\\sigma'_t{}^2}\\right) + d \\frac{\\sigma_t^2}{\\sigma'_t{}^2} - d \\Bigg].\n\\end{aligned}\n$$\n\nIf variance and noise schedules are the same, i.e., $\\alpha_t = \\alpha'_t$ and $\\sigma_t = \\sigma'_t$, the above reduces to:\n\n$$\n\\frac{\\beta_t^2}{2 \\sigma_t^2} \\|\\| s_p(x_t, t) - s_q(x_t, t) \\|\\|^2.\n$$\n\nThe derivation in Lemma 2 becomes much more cumbersome in this relaxed setting; however, the point-wise KL is still tractable.\n\nWe note that the \"same schedule\" assumption holds for all of our experimental settings. However, we agree with the reviewer that in some practical settings this assumption may be violated. We will remark on this relaxation in the final version.\n\n---\n\n**Re: Weakness 3**\n\nWe believe aside from the commonly used equal-weights baseline which we compare to, other existing methods in the literature are not applicable to the settings that we consider and are generally not comparable to our constrained learning framework. This includes papers that propose constrained methods; In all these works that we are aware of, what they mean by constraining the diffusion process is significantly different from our definition. We now discuss the most relevant papers that we have cited to explain why they are not suitable baselines in this work. If the reviewer has specific existing approaches in mind as meaningful baselines, please mention them so we can discuss and potentially include them.\n\n- In [1] which has the most compatible setting, they average model weights to sample from the geometric mean distribution. They further propose a weighted average with potentially different weights and suggest finding weights that minimize a specific objective (e.g., Image reward). They give a heuristic algorithm to do this called \"Greedy Souping\". This is not usable in our setting since there is no clear choice for the objective in both for composition and alignment with multiple rewards.\n\n- In [2] they propose a superposition method to sample from the mixture of diffusion models with arbitrary weights (but they only use equal weight mixtures and don't discuss different weights). They also devise a method to sample points that have equal likelihood under different models which is fundamentally different to product composition.\n\n- Works like [3,4,5,6] all discuss constrained sampling from diffusion models, but the nature of their constraints is completely different from our work as it mainly involves sampling from a constrained set and they propose to do this through projection onto a feasible set at each diffusion time step. It is not clear how to apply these methods to reward constraints or how to use them to preserve distance to a model.\n\n- Other papers we have cited either propose methods to improve alignment with a single reward by using a KL regularizer like [7], or they enforce very specific constraints by adding additional losses with fixed weights to the objective which implicitly enforces the constraint like [8,9]. These methods are very specific to the constraints they are designed for and do not generalize to arbitrary reward functions and don't give us a way to constrain closeness to a model.\n\n\n**References**\n\n[1] (Biggs et al.) *Diffusion Soup: Model Merging for Text-to-Image Diffusion Models*  \n[2] (Skreta et al.) *The Superposition of Diffusion Models Using the Itô Density Estimator*  \n[3] (Christopher et al.) *Constrained Synthesis with Projected Diffusion Models*  \n[4] (Liang et al.) *Multi-Agent Path Finding in Continuous Spaces with Projected Diffusion Models*  \n[5] (Narasimhan et al.) *Constrained Posterior Sampling: Time Series Generation with Hard Constraints*  \n[6] (Zampini et al.) *Training-Free Constrained Generation with Stable Diffusion Models*  \n[7] (Fan et al.) *DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models*  \n[8] (Giannone et al.) *Aligning Optimization Trajectories with Diffusion Models for Constrained Design Generation*  \n[9] (Chen et al.) *Towards Aligned Layout Generation via Diffusion Model with Aesthetic Constraints.*\n\n---\n\n**Re: Weakness 4**\n\nWe agree with the reviewer that from a theoretical standpoint, using the surrogate score is not an accurate method to sample from a product distribution. We note that it is a commonly used approach in the literature (e.g., see [2]) and more generally classifier and classifier-free guidance both use a similar approach.\n\nFurthermore, in constrained product composition for text-to-image diffusion, our results show that even with the approximation error, KL divergence constraints lead to more balanced representation of the concepts in concept composition and more balanced reward levels in composing reward fine-tuned models, which we believe is strong evidence of the effectiveness of our approach.\n\nAs we discuss in our paper, in [1] it is shown that using annealed MCMC sampling with the surrogate score can lead to samples from the correct product distribution. However, their full implementation needed to reproduce the relevant results, which we could then apply in our constrained learning framework, is not publicly available$^*$ and reproducing their results with an incomplete implementation is beyond the scope of our paper.\n\nRegarding the theory, it can be shown that in the case where the initial distributions at time $t = 0$ are Gaussians, the surrogate score and the true score coincide for all values of $\\lambda$. Anything beyond this simple setting quickly becomes intractable to analyze.\n\nWe thank the reviewer for highlighting this approximation error and agree that analyzing it will provide a stronger theoretical justification of our product composition experiments. We plan to more thoroughly study and address the effects of using the surrogate score for product composition in future work.\n\n**References**\n\n[1] Du et al. *Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC*  \n[2] Du et al. *Compositional Visual Generation with Energy Based Models*  \n\n*Note: There is a GitHub repo for the paper [1] but it does not include the implementation needed to do MCMC sampling with pre-trained text-to-image diffusion models.*\n\n---\n\n**Re: Weakness 5**\n\nAll experiments were run on a single Nvidia A6000 GPU.\n\n- For Alignment, there is little additional time overhead compared to baselines like AlignProp. For example, for the experiment in Figure 3, runtime is 33 minutes for both constrained and unconstrained methods, and for the experiments in Figure 4 constrained is 1h4m, unconstrained is 1h. Existing approaches already estimate the KL and sample batches to evaluate and backprop through the reward. The only additional computation for our method is the dual updates which is negligible in terms of added time.\n\n- For Composition, there is no meaningful comparison to the equal weights baseline since there, we don't learn the weights. For constrained composition, it takes around 5–10 dual updates for dual variables to converge which for composing the 5 fine-tuned stable diffusion models takes ~9 minutes total and for concept composition ~2 minutes.\n\nWe thank the reviewer for highlighting this and we will include the above and additional runtime details in the final version."}]

Author rebuttal:

Meta-review: This paper introduces a constrained optimization framework that unifies reward alignment and composition of diffusion models under a Lagrangian formulation. 

The reviewers all agreed that this was a good contribution. The approach was seen as technically solid and well written, with empirical results demonstrating improved balance across multiple rewards and better representation of composed models relative to simple baselines. 
At the same time, the reviewers raised several points that should be addressed in the camera-ready. The empirical evaluation remains somewhat narrow, with limited exploration beyond text-to-image diffusion and few comparisons against more advanced baselines. Clarifications are also needed regarding assumptions (e.g., shared noise schedules, convexity to parameterized score networks). The rebuttal addressed many of these concerns, and we hope to see these additions in the final version of the paper.
Decision: Accept (poster)