Composition and Alignment of Diffusion Models using Constrained Learning

Re: Weakness 1

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

When 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

$$P_\theta^\star - D_\theta^\star \leq O(\nu),$$

where $\nu$ is a parametrization gap that satisfies

$\|\| p(\cdot; s) - p(\cdot; s_\theta) \|\|_1 \leq \nu \quad \text{for any } s \in \mathcal{S}.$

Thank you for pointing out this important point. We will include all these points formally in the final version.

Reference

[1] (Chamon et al.) Constrained Learning With Non-Convex Losses

---

Re: Weakness 2

In 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:

$$D_{\text{KL}} \left( p_{0:T}(\cdot ; s_p) \Vert p_{0:T}(\cdot; s_q) \right)$$ $$= \sum_{t=1}^T E_{x_{0:T} \sim p} \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]$$

Importantly, $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

$$p(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),$$

and

$$p(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),$$

is still tractable and given by:

$$\begin{aligned} D_{\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]. \end{aligned}$$

If variance and noise schedules are the same, i.e., $\alpha_t = \alpha'_t$ and $\sigma_t = \sigma'_t$, the above reduces to:

$$\frac{\beta_t^2}{2 \sigma_t^2} \|\| s_p(x_t, t) - s_q(x_t, t) \|\|^2.$$

The derivation in Lemma 2 becomes much more cumbersome in this relaxed setting; however, the point-wise KL is still tractable.

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

---

Re: Weakness 3

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

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

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

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

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

References

[1] (Biggs et al.) Diffusion Soup: Model Merging for Text-to-Image Diffusion Models
[2] (Skreta et al.) The Superposition of Diffusion Models Using the Itô Density Estimator
[3] (Christopher et al.) Constrained Synthesis with Projected Diffusion Models
[4] (Liang et al.) Multi-Agent Path Finding in Continuous Spaces with Projected Diffusion Models
[5] (Narasimhan et al.) Constrained Posterior Sampling: Time Series Generation with Hard Constraints
[6] (Zampini et al.) Training-Free Constrained Generation with Stable Diffusion Models
[7] (Fan et al.) DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models
[8] (Giannone et al.) Aligning Optimization Trajectories with Diffusion Models for Constrained Design Generation
[9] (Chen et al.) Towards Aligned Layout Generation via Diffusion Model with Aesthetic Constraints.

---

Re: Weakness 4

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

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

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

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

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

References

[1] Du et al. Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC
[2] Du et al. Compositional Visual Generation with Energy Based Models

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.

---

Re: Weakness 5

All experiments were run on a single Nvidia A6000 GPU.

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

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

We thank the reviewer for highlighting this and we will include the above and additional runtime details in the final version.

Re: Weakness 1
An 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.

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

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

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

Proof Summary of Lemma 2:
In 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.

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

Reference

[1] (S. Lyu) Interpretation and generalization of score matching

---

Re: Weakness 2
By 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:

$$\begin{aligned} \min_{u \geq 0, p} \quad & u \\ \text{subject to} \quad & D_{\mathrm{KL}}(p \Vert q_i) \leq u \quad \text{for } i = 1, \ldots, m, \\ & \mathbb{E}_{x \sim p} [r_j(x)] \geq b_j \quad \text{for } j = 1, \ldots, n. \end{aligned}$$

where the goal is to find a model that is close to multiple pre-trained models and also satisfies one or more reward constraints.

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

---

Re: Weakness 3
The 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:

As mentioned in Lemma 2, computing point-wise KL requires 'valid' score functions, i.e., gradients of log-likelihood of some forward process.

We 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)$.

Crucially, 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:

$$q_{\mathrm{rw}}^{(\lambda)} \propto q e^{\lambda^\top r}.$$

Luckily, unlike in composition, in alignment, since there is only a single KL objective, point-wise KL and path-wise KL lead to similar solutions:

$$p_0^\star(x_0) = \int q_{0:T}(x_{0:T}) e^{\lambda^\top r(x_0)} dx_{1:T} \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} = q_0(x_0) e^{\lambda^\top r(x_0)}$$

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

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

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

Re: Weakness 1

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

To 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:

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

Minimum CLAP scores across prompts for each method:

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

Reference

[1] CLAP: Learning Audio Concepts From Natural Language Supervision (Elizalde et al.)

---

Re: Weakness 2

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

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

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

---

Re: Weakness 3

We 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:

• KL divergence constraints in the constrained composition problem:
  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.

• Reward constraints:
  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.

---

Re: Question 1

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

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

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

---

Re: Question 2

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

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

Re: Weakness 1

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

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.

We thank the reviewer for the suggestion, and will include ablations on the constraint levels that illustrate the above points (see tables below).

Table 1. MPS reward alignment with saturation and contrast constraints, for varying thresholds.

Table 2. Pickscore reward alignment with saturation and contrast constraints, for varying thresholds.

---

Re: Weakness 2
As 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:

We thank the reviewer for bringing up this important setting and will include the above example in the final version.

---

Re: Question 1
The first letter indicates parameterization: U and S in UR-A and SR-A denote Unparameterized and Score function parameterized problems, respectively.
The second letter, R or F (as in UR-C or UF-C), stands for Reverse or Forward KL divergence constraints.
The final letter, A or C, signifies Alignment or Composition. We will clarify these abbreviations in the final version.

---

Re: Question 2
Based 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

$$p^\star \propto p_{\text{pre}} e^{\sum \lambda_i r_i}$$

but with potentially different $\lambda$'s.

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

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

We thank the reviewer for raising this interesting question, and we will comment on this in the final version.