Direct Distributional Optimization for Provable Alignment of Diffusion Models

Authors: Ryotaro Kawata, Kazusato Oko, Atsushi Nitanda, Taiji Suzuki
Primary area: optimization
Keywords: diffusion models, optimization

Abstract:
We introduce a novel alignment method for diffusion models from distribution optimization perspectives while providing rigorous convergence guarantees.
We first formulate the problem as a generic regularized loss minimization over probability distributions and directly optimize the distribution using the Dual Averaging method.
Next, we enable sampling from the learned distribution by approximating its score function via Doob's $h$-transform technique.
The proposed framework is supported by rigorous convergence guarantees and an end-to-end bound on the sampling error, which imply that when the original distribution's score is known accurately, the complexity of sampling from shifted distributions is independent of isoperimetric conditions.
This framework is broadly applicable to general distribution optimization problems, including alignment tasks in Reinforcement Learning with Human Feedback (RLHF), Direct Preference Optimization (DPO), and Kahneman-Tversky Optimization (KTO). We empirically validate its performance on synthetic and image datasets using the DPO objective.

Decision: accept (poster)
Reviewer scores: 8, 8, 5, 8, 6

Reviews and discussion:
<thread>
creation date:
1730142347791

thread_type:
official_review

<thread_object>
creator: reviewer_tyit

<content>
summary:
This paper proposes a new optimization procedure to minimize entropy-regularized versions of finetuning losses for diffusion models such as DPO or KTO losses. This optimization method is an approximation of Dual Averaging (DA), where, instead of tracking a set of particles to approximate the distribution as done in prior work, an estimate of the unnormalized density ratio between the iterate and the original reference model is tracked. These changes make it possible to sample from finetuned model (e.g. the final iterate) by simulating the backward SDE of a OU forward process (à la diffusion model), where the score of the forward process is estimated by a Monte Carlo approximation its expression obtained using Doob's h-transform formula, whose ``correction'' terms can be estimated from the density ratio-estimate.

On the theoretical side, global convergence (in TV) of the final DA iterate to the true minimizer is provided when the loss is Convex, its first variation is bounded and smooth, and the density ratio estimates are accurate enough, while non global convergence (convergence of the gradient) to some limit is established if the loss in nonconvex. Finally, the TV between the distribution of the final finetuned model samples and true finetuning loss minimizer is controlled, assuming a bound on the error of score of the reference OU process, and a bound in the error of the Doob correction term.

The method is implemented on a synthetic 2d mixture of Gaussians reference model with mode alignment, an image dataset with color alignment, and Head CT scans from Medial MNIST with orientation alignment, where it is show that the loss trends downwards, and that the finetuned samples are more aligned than the ones from the reference models.=

strengths:
The present submission tackles an important problem (diffusion model alignment). The method is an interesting twist on existing dual averaging approaches, whose setting does not involve estimating density ratios.

A significant effort is spent providing end-to-end guarantees for both the optimization procedure and the sampling procedure. These guarantees do not directly require a Log-Sobolev inequality on the target. 

Finally, on the experimental side, the method is tested on high-dimensional dataset, and display reasonable performance (although there is no relative point of comparison due to the lack of results reporting the performance of existing methods on the datasets and tasks explored by the paper).

weaknesses:
My main concern is the lack of comparison against existing methods to finetune diffusion models, both in terms of performance and in terms of space and memory complexity. I personally like the paper, however I believe it needs comparison against the applicable baselines before acceptance.
An investigation of the impact of $\beta$ would also be interesting.

The method involves some computationally expensive operations such as Monte Carlo estimation at every iteration of the DA loop. It would be nice to have the total time and space complexity of the algorithm described somwehere, complementing a Computational time comparison with other methods. The fact that only one previous density ratio estimate (and not k) are needed to approximate the next one (see l 1953, algorithm D1) seems important to cap the total time and space complexity of the method and not directly obvious from equation 2 and 3, should be mentioned in the main.

The presentation needs to be clarified. Currently,
  - How is the reader supposed to know that the equality on equation (2) holds? It seems hidden as specific instances of Lemma 1, but this should be highlghted more.
  - The recipe for estimating the density ratios during DA is not clearly described in the main body. 
  - The presentation of the image generation results could be improved: what is the color (0.9/0.9/0.9)? this deserves to be on one of the figures IMO. Figure 7 definitely helps to understand how much the model has learned, and should be at least referenced in the main text.
  - I did not understand the meaning of the ("Linear/nonlinear", "k=1/2/3 DA Loops") terms. Do you mean number of DA iterations? Why linear and nonlinear?
  - The reason "our method works without isoperimetry conditions such as Log-Sobolev inequality" could be explained more clearly: my understanding is that the original DA update for MLFD required running an inner Langevin sampler, which is sensitive to LSI. On the other hand, the current method estimates the density ratio by using the backward SDE to sample from the reference model, which does not rely on an LSI to converge.

Additional minor style/typos improvements are suggested in minor.




### **Minor**
The style could be substantially improved. The english grammar mistakes and imprecision should be corrected.

- the differences between the two DA techniques are not explained. Turns out some results hold for Opt 1 or Opt 2.
- Typo l 1235, 1236, 1249: $\nabla p_t$ -> $\nabla \log p_t$
- "Particle-based optimization methods are not designed for resampling from the distribution where obtained particles follow" -> "...that the obtained particle follow"
- "pre-trained distribution is typically highly complex multimodal distributions, failing to satisfy isoperimetric conditions". Not super accurate. Isoperimetric conditions are verified for Mixture of Gaussians for instance. But the isoperimetric constant is very large.
- "The convergence rate of MFLD (Nitanda et al., 2022; Chizat & Bach, 2018) has been established under the condition where the proximal"... I don't think (Chizat & Bach, 2018) tackle MFLD.
- Consider using ``test'' for correct quotes formatting.

questions:
.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730589417906

thread_type:
official_review

<thread_object>
creator: reviewer_ei8r

<content>
summary:
From my understanding, the paper presents a novel distribution alignment technique for score-based diffusion models based on optimization techniques over probability measures. Standard measure optimization techniques face challenges in the diffusion model setting because both the density governed by the model and the density of the true reference distribution cannot be evaluated. The model parametrizes its distribution using the score function/networks, whereas the true reference distribution is accessible only through samples. The work reformulates the measure optimization as an optimization problem over potentials. Then, sampling using the learned potential proceeds by the Doob's $h$-transform technique.

strengths:
- The proposed method is usable on relevant examples such as reward-based alignment, DPO, and KTO. There are contributions even in the small details of this derivation, including rewriting the DPO objective in a way that is more amenable to statistical analysis.
- The sampling error is bounded while taking into account all sources of approximation: the discretization of the SDE, approximation of the score, and approximation of the Gibbs distribution resulting from the potential.
- While the theoretical derivations in the paper are already a strong contribution, the experiments are helpful in confirming that the method works on simple tasks.

weaknesses:
Broadly speaking, the significant portions of the paper are unreadable. While I applaud the authors effort in producing this work, I request an equal effort in presenting it well. Some specific points are outlined below.
- The "two main challenges" addressed by paper are stated at least three times (045-048, 162-168, 185-194). While I believe this was done for clarity sake, the "challenge" is never described in an accessible way. While I interpret them as what I wrote in the summary, it is not explicitly clear to me. To say "particle-based optimization methods are not designed for resampling from the distribution where obtained particles follow" is meaningless without having described what particle-based methods are. The pre-trained distribution "failing to satisfy isoperimetric conditions" is meaningless without clearly explaining the implication of this (for example, are there no distributional optimization guarantees for distributions that do not satisfy isoperimetric conditions)? 
- The appendix is a literally a list of theorems and proofs. Whether for optimization or sampling-style analysis, it is important to give context to all results. Concretely, this means: 1) providing a proof outline, 2) describing which technical steps are identical to/adapted from previous work (with references), and what proof techniques are specific to this paper. In its current state, it is hard to even verify the technical correctness of the paper.
- The notation is excessive, and not always introduced. Unless I am misunderstanding, $\hat{q}$ is presented as the minimizer of the objective, but seems to be changed to $q^*$ towards the end. I am also not sure that $\tilde{X}_{kh}^{\leftarrow}$ is introduced in 416. In the final version, please provide a global notation table in the first appendix containing all symbols and a verbal description. 

I am recommending 5 as a score. 

I will **raise my score from 5 to 6** if the authors can provide in their rebuttal **a completely rewritten version of 036 - 048 in an accessible form**, including equations, for readers who are unfamiliar with optimization over probability measures, as this is a more niche field than traditional optimization theory or diffusion models at this point (feel free to ignore references).

I will **raise my score from 6 to 8** if the authors provide **all missing text** from Appendix A. This includes 
- text between A and A.1
- text between A.1 and Theorem 5 (which should be renamed to Theorem 1).
- description of the significance of Lemma 1 (which seems to be a strong convexity-like result).
- description of the significance of Lemma 2
- any in-line explanations.
- similar for all of A.2.
Please look to the outline in the weaknesses for this. For examples of papers with well-presented optimization/sampling proofs, see [Schmidt et al. (2016)](https://arxiv.org/pdf/1309.2388), [Diakonikolas et al. (2021)](https://arxiv.org/pdf/2011.00364), [Balasubramanian et al. (2022)](https://proceedings.mlr.press/v178/balasubramanian22a/balasubramanian22a.pdf).

questions:
- In Assumption 2, are both (i) and (iii) necessary? Can we not assume (iii), the derive (i) with TV^2 instead of KL, and then apply Pinkser's inequality to achieve the smoothness relative to KL?
- The nonconvex convergence analysis results only in convergence of the successive gaps of iterates. What would it take to achieve convergence to an appropriately-defined stationary point?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730695788279

thread_type:
official_review

<thread_object>
creator: reviewer_141c

<content>
summary:
The paper formulates the diffusion model alignment problem as a general distribution optimization problem. The proposed general objective, which involves regularized loss minimization over probability distributions, encompasses important tasks in language model alignment, including RLHF, DPO, and KTO. The authors propose an algorithm that first employs Dual Averaging to approximate the optimization objective function over a distribution into a potential function (thereby transforming the variable from a distribution to a point), and then uses this learned potential function to fine-tune diffusion models. The approach is supported by theoretical guarantees.

strengths:
1. The paper's formulation of a general distribution optimization problem is significant as it unifies several widely-studied alignment tasks under a single framework, making it relevant to current research priorities in the field.
2. The paper is well-written and easy to follow.

weaknesses:
1. The paper lacks crucial implementation details for applying Dual Averaging to diffusion models. Specifically, while the objective involves  $KL(q|p_{ref})$ and $exp(-f)p_{ref}$, it doesn't address how to compute these terms when only time-dependent score functions are available by diffusion models. The potential Monte Carlo alternatives would be computationally intensive.

2. In phase 2,  the proposed correction term u(x,t) in equation 17 is computationally impractical, requiring multiple reverse process simulations at each inference timestep.

3. Based on 2, the paper's key Assumption 3.4 regarding the small approximation error of u* is questionable, particularly given the gradient terms and complexity of the approximation.

questions:
1. According to Alg. 4.1, we've already got $q_K \propto exp(-f_K)p_{ref}$, which is exactly the optimal solution of (1) (also the target of phase 2), why do we need to conduct phase 2?

2. The paper claims to provide convergence guarantees for diffusion alignment problems without relying on isoperimetric conditions. However, Assumption 3 requires that $\nabla \log p_t$ is $L_p$-smooth at every time $t$. Could this requirement serve as an alternative to isoperimetric conditions?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730709176956

thread_type:
official_review

<thread_object>
creator: reviewer_aajd

<content>
summary:
The authors propose a distributional optimization algorithm that builds on the dual averaging (DA) algorithm combined with Doob's 
ℎ-transform. This framework encompasses several methods, including Reinforcement Learning(RL), Direct Preference Optimization (DPO), and Kahneman-Tversky Optimization (KTO). The paper includes an error analysis and provides numerical implementations of Gaussian mixtures and image generation tasks.

strengths:
1. The paper provides a general framework for distributional optimization, integrating a sophisticated optimization algorithm (dual averaging) with Doob's ℎ-transform to capture the score function effectively.

2. The authors present a convergence rate analysis based on suitable assumptions, enhancing the theoretical robustness of the approach.

3. Experiments are conducted on both synthetic data (Gaussian mixtures) and real-life CT images, demonstrating the practical applicability of the proposed method.

weaknesses:
The method introduces an additional entropic regularization term involving a hyperparameter, $\beta'$. In the theoretical analysis, beta prime is assumed to be larger than $\beta$, yet in the experiments, $\beta'$ is set equal to $\beta$ in all cases. There is a lack of discussion on the role of these parameters and their impact on the results.

questions:
Typically, the score estimation error is assumed in the $L^2$ sense, but Assumption 3.3 assumes it in the $L^\infty$ sense. Could this assumption be improved?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731180511876

thread_type:
official_review

<thread_object>
creator: reviewer_hxn1

<content>
summary:
Authors study the problem of sampling from distribution $\hat q$ given as the minimizer of a general optimization problem in the space of distributions given by $F(q) + KL(q||p_{ref})$, where $p_{ref}$ is the base distribution (e.g. obtained by pretraining) and $F$ is a functional in the space of distributions. $F$ can be defined as the DPO objective given user preferences of the samples from the distribution. A simple toy example of the DPO based F, which authors use for sampling from a single gaussian when the reference measure $p_{ref}$ is mixture of gaussians,  is when the DPO preference function between two samples picks the one that is closer to the correct gaussian mean.
 
Authors introduce an algorithm for this task based on applying the Dual averaging method in the space of distributions, and learning the gradient of $F$ at each iteration by neural networks. At the end of DA, they get access to a neural net that approximates the log density and the score of $\hat q$ with respect to $p_{ref}$. Then, they use a trick based on Doob’s $h$-transform relation to translate this information to the score function of  $\hat q$, which they use to obtain fresh samples from $\hat q$ using the backward process of Ornstein-Uhlenbeck.

strengths:
The overal approach linking optimization in the space of distributions to diffusion models is new to me and looking interesting. Authors seem to derive bound for smoothness coefficients of score function of $\hat q$ based on smoothness of the diffusion for the reference measure. Their approach in using Doob's h-transform for estimating the score functions of the corrected distribution is also very interesting. The presentation of the beginning part of the paper looks clean and authors clearly state their high level goal.

weaknesses:
Overall choice of the Dual averaging for this task seems unmotivated, I don’t see why we can’t use simpler algorithms like GD in the space of distributions, and learning the gradients with neural nets?

Other concerns:
One concern that I have is where authors use Doob’s h-transform to estimate the score of $\hat q$ as \nabla \log\E[e^{-f(X_t)} | X_T], by estimating the derivative of probability density instead of log density, in line 1252, with samples. Given that estimating the density is ill-conditioned in general, can authors elaborate on the exact guarantee they get for this approach?

The presentation of the proofs can be more clear, right now the logical order between Lemmas, etc is not clear. I would suggest authors add more intuition before stating the mathematical formulations in each Lemma, in how that relates to the bigger picture. 

The derivative notation in line 143 is not clear.
224: one and zero should be written in English rather than numbers

questions:
What is $\nabla_x \nabla_x^\top$ in line 1471
What is the meaning of subindices $|_{x=\bar X_t}$ ? The notation is very confusing
What is the benefit of using the Dual averaging method compared to gradient descent in the Wasserstein space? I assume you can also implement that by learning the gradient of the functional by neural nets?

Can authors clarify why they use the proof of Liu et al for their proof of convergence of the Dual averaging method in the non-convex case, instead of the more classical older papers on Dual averaging?

In the case of gaussian mixture where the authors define the objective of DPO  with the goal of sampling from one of the single gaussians, is it the case that it can mathematically be shown that the minimizer of the DPO objective is exactly a single gaussian?

In the case of gaussian mixture, Is it possible to prove guarantees for learning the sequence $\bar q_k$ rigorously with neural nets or some polynomial regression?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732150566698

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Comment

comment:
First of all, we would like to express our sincere gratitude for your thoughtful and constructive feedback on our work. Your comments and suggestions have been invaluable in helping us identify areas for improvement and clarify our contributions.  

We have carefully reviewed all your comments and would like to address a few key points here:  

1. **Motivation of our algorithm**:  
    - The reviewers expressed concerns about why Dual Averaging is effective in the current distribution optimization task.
    - By employing Dual Averaging, we only need the neural network to directly learn the density ratio $q^{(K)}/p_\mathrm{ref}$ of the output distributions at the final step of the denoising process in the diffusion models, while a naive gradient descent type method does not. This approach enables us to achieve an aligned output density. 
    - Using this density ratio, the Doob's $h$-transform uniquely determines the correction terms for each time step in the aligned diffusion model.  

2. **Computational efficiency of our algorithm**

    - The reviewers raised questions about the computational cost associated with the Doob's $h$-transform technique.  
    - Our simplest solution (i.e., estimating the correction term at each time step using Monte Carlo) has a time complexity of $\mathcal{O}(L^2)$, where $L$ represents the number of time steps in the denoising process. When we set $N$ as the number of particles needed to estimate one correction term, to compute the correction term for each sample simultaneously, $\mathcal{O}(N)$ memory space is required.
    - We emphasize that Doob's $h$-transform techniques used in our phase 2 have been empirically applied in image generation tasks [R1], Bayesian computation [R2], and online filtering [R3]  with efficient sampling schemes. 
 

3. **Comparision between the prior work**:  
   - The reviewers pointed out the need for a comparison with existing methods [R4]. In response, we conducted additional numerical experiments and confirmed that our algorithm achieves a smaller loss while remaining within acceptable computational resource limits during Dual Averaging phase.


4. **The role of the hyperparameter**:  
   - The reviewers raised questions about the role of the hyperparameter $\beta'$ in our algorithm. In our approach, $\beta'$ control the speed of optimization: The objective function with additional regularization imposed by the hyperparameter $\beta'$, which is $L(\hat{q}^{(k)}) + \beta'/k D_\mathrm{KL}(\hat{q}^{(k)}\|p_\mathrm{ref})$, monotonically decreases during Dual Averaging. Through experiments, we observed that smaller values lead to slower optimization, and we have included these results in the manuscript.

5. **Presentation of our contributions**
    - The reviewers suggested that the introduction and theorem proofs should be made more accessible to readers who may not be familiar with the relatively niche field of distribution optimization. We have incorporated this advice into the manuscript and then the manucript was substantially improved.
    - Additionally, we received numerous suggestions to improve the clarity of notation and enhance the precision of expressions.

We deeply appreciate the time and effort you have invested in reviewing our submission. Your feedback has helped us significantly refine and strengthen our work. Thank you again for your valuable insights and for contributing to the quality of our research.  


Sincerely,  
Authors of Submission4323.

---

[R1] Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, Sergey Levine: Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control. arXiv, 2024.

[R2] Jeremy Heng and Valentin De Bortoli and Arnaud Doucet: Diffusion Schrödinger Bridges for Bayesian Computation. Statistical Science, 2024.

[R3] Nicolas Chopin, Andras Fulop, Jeremy Heng, Alexandre H. Thiery: Computational Doob $h$-transforms for Online Filtering of Discretely Observed Diffusions. ICML, 2023.

[R4] Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, Nikhil Naik (2024). Diffusion model alignment using direct preference optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 8228-8238).
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a new algorithm for distribution optimization, specifically for aligning diffusion models, based on dual averaging and Doob's h-transform. The approach offers a general framework applicable to various alignment tasks and demonstrates promising empirical results. Reviewers found this paper interesting. 

While the paper presents a novel and potentially valuable contribution to distribution optimization, it could benefit from clearer explanations for a broader audience, more detailed comparisons with existing methods, and further investigation of the computational costs and the role of hyperparameters. In particular, for the camera-ready, some of the citations suggested by the reviewers on RL for distribution optimization should be mentionned in Example 1 (e.g. Marion et al. (2024))