Direct Distributional Optimization for Provable Alignment of Diffusion Models

Clarification of the presentation

We appreciate your feedback on the presentation. We have made the modifications as suggested:

• We have clarified the differences between Opt.1 and Opt.2 in the theoretical analysis of Dual Averaging. In particular, when the objective $F$ is convex, we have $\mathcal{O}(1/K)$ convergence of the weighted sum of the losses in Opt.1 with $K$ iterations. On the other hand, when $F$ is not necessarily convex (but smooth), the functional derivative of the regularized loss $\delta L/ \delta q (\hat{q}^{(K)})$ converges to zero, up to a constant w.r.t $x$, in Opt.2.

• We have also updated the figures and the implementation details from the previously conducted experiments to make them more comprehensible. Just to clarify, the terms linear and nonlinear refer to the fact that when the target functional is linear, a single iteration of Dual Averaging is sufficient, whereas in the nonlinear case, multiple iterations are required.

• We have refined the expressions related to the LSI to make them more precise.

If you feel that all of your concerns have been properly addressed, we would deeply appreciate your consideration of a higher rating score, if you believe it is warranted.

[R1] 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).

[R2] 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.

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

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

We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.

Comparison against existing methods

We absolutely agree that we need to compare our methods against existing methods, so we have done experiments with an existing DPO optimization method, called Diffusion-DPO [R1].

• We empirically compared the performance of the methods with Gaussian Mixture Model, setting the same DPO objective.
• The Dual Averaging was done with $\beta'\geq\beta=0.04$. We have illustrated results for the case where $\beta'>\beta$ in the revised manuscript.
• As a counter method, following Diffusion-DPO [R1], we have optimized the score networks by minimizing the approximate upperbound, which can be calculated with the score, of the (not regularized / $\beta$-regularized) true DPO loss. In practice, it is hardly realistic to compute the true loss during Diffusion-DPO, but in this case, we forcefully carried out the computation. The densities $q$ and $p_\mathrm{ref}$ were estimated by repeatedly computing the denoising path and empirically obtaining their marginal densities.
• We evaluated how the true DPO loss of Diffusion-DPO behaves while minimizing the approximated upperbound.
• In addition, we evaluated the averaged squared Euclidean distance from the target mean of Gaussian as the Metric Loss.

We have compared the performance and the computational costs in the optimization phase:

From this result, we observe the following:

• We got a smaller value of the true DPO loss and the metric loss for our method.
• It is crucial to note that the prior approach makes it difficult to control the true loss. That is, minimizing their approximated upperbound failed to decrease the true loss indicating that the upperbound is too lose to be used as a surrogate of the true loss (please refer to the loss curves in the revised manuscript).
• In our approach, we directly optimize the (regularized) true loss, ensuring that the iteration can continue until it reaches a small value albeit at a higher computational cost. In phase 1 (the optimization phase), regarding time and space complexity, our method was comparable to the existing method to achieve roughly the same loss. As you mentioned, while we only need to store one previous density ratio model, our method does require more space to be utilized.

In phase 2 (the sampling phase), 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. The sampling error for each correction term would be $\mathcal{O}(1/\sqrt{N})$. This leads to severe computational effort in phase 2.

However, we emphasize that Doob's $h$-transform techniques used in our phase 2 have been empirically applied in image generation tasks [R2], Bayesian computation [R3], and online filtering [R4] with efficient sampling schemes. In particular, as a more practical alternative of our phase 2, the idea of approximating the correction term using neural ODE solvers for faster test-time implementation has also been proposed in [R2]. This would enhance the practicality of our phase 2.

Thank you very much for your work on reviewing this submission.

Could you read the author's reply, check if they have addressed your comments, and acknowledge it?

Further, please make sure, if necessary, to explain what your new stance is based on this information.

Thank you for replying our rebuttal. Below, we address your questions in detail:

Numerical Experiments

• We note that “Upperbound” refers to the upper bound on the approximation to the DPO loss, not on the true DPO loss. We refer you to paper [R1], which approximates the true DPO loss by replacing the current reverse process with the ROU process and derives the upper bound based on this approximation. Hence, the upper bound could lie below the true DPO loss since the aligned reverse process will deviate from ROU.

• In fact, we empirically observed that the upper bound successfully constrained the true objective during the early stages (up to approximately 25 steps). However, the true loss began to increase and diverge without additional regularization, implying "catastorophic forgetting" or "overoptimization" to the upper bound, and the true loss ceased to decrease with additional regularization.

• As far as we know, Diffusion-DPO [R1] is the first paper that applied Direct Preference Optimization to diffusion models. Diffusion-DPO was presented as a computationally efficient alternative of Reinforcement Learning approach. This method is empirically powerful: for example, utilizing Diffusion-DPO, [R2] aligned diffusion models for text-to-audio generations. However, their method approximates the true DPO loss by replacing the current (aligned) reverse process with the ROU process and derives the upper bound based on this approximation, which can be crude. Therefore, we selected Diffusion-DPO as an alternative because this issue has not been thoroughly studied in prior work.

Citation

• As you have suggested, we will replace the citation in p.181 to [R3] in the final version.

We will clarify and fix these points you have mentioned in the final version.

[R1] Bram Wallace et al. (2024). Diffusion Model Alignment Using Direct Preference Optimization. ArXiv preprint. https://arxiv.org/abs/2311.12908

[R2] Navonil Majumder, Chia-Yu Hung, Deepanway Ghosal, Wei-Ning Hsu, Rada Mihalcea, and Soujanya Poria. 2024. Tango 2: Aligning Diffusion-based Text-to-Audio Generations through Direct Preference Optimization. In Proceedings of the 32nd ACM International Conference on Multimedia (MM '24). Association for Computing Machinery, New York, NY, USA, 564–572. https://doi.org/10.1145/3664647.3681688

[R3] Chizat, L. (2022). Mean-field Langevin dynamics: Exponential convergence and annealing. Transactions on Machine Learning Research. https://openreview.net/forum?id=BDqzLH1gEm

We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.

We are grateful to the reviewer ei8r for highly valuing our contributions and for suggesting improvements to the readability of the paper.

Following your suggestions, we have revised the manuscript:

• We have revised the introduction to make it more accessible to readers who are not familiar with distribution optimization.
• We have also revised the convergence analysis of Dual Averaging, including a proof sketch, to ensure it is fully understandable.

We would appreciate it if you could review the revised manuscript.

Below, we will provide responses to the valuable questions we received:

Challenges in our paper

Alignment of diffusion models as a minimization problem over the probability space has two challenges:

• (i) inaccessibility of the output densities.

As you summarized in the summary, the issue lies in the fact that the output distribution of diffusion models is parametrized by the score network.

• (ii) muitimodality of the densities.

Existing distributional optimization methods such as mean-field Langevin dynamics [Mei et al., 2018] and particle dual averaging [Nitanda et al., 2021] resolved this minimization problem by adapting a Langevin type sampling procedure into its optimization procedure so that we can calculate a functional derivative of the objective. Concretely, we find the functional derivative $\delta F/ \delta q$ in mean-field Langevin dynamics that is described as

where $X_t \sim q_t$ and $(B_t)_t$ is Brownian motion. This density $q_t$ approaches the minimizer of the optimization problem.

However, the densities in our problem setting $q$ and $q_\mathrm{ref}$ are highly complex and multimodal from which it is extremely hard to generate data by a standard MCMC type methods including the Langevin dynamics. This difficulty can also be mathematically characterized by isoperimetric conditioins, such as logarithmic Sobolev inequality (LSI), that has usually exponential dependency on the data dimension $d$ for multimodal data yielding the curse of dimensionality. Unfortunately, the existing distribution optimization methods mentioned above are sensitive to the LSI constant, so they suffer from severely slow convergence, failing to align diffusion models.

Necessity of assumptions for Dual Averaging

The first assumption (i) is quite standard assumption known as a Lipschitz smoothness condition. This is required to obtain $O(1/k)$ convergence, otherwise, we would only have $O(1/\sqrt{k})$ for convex settings and no guarantee for non-convex settings.
The condition (ii) is required to take the approximation error of $q^{(k)}$ from $\hat{q}^{(k)}$ into account. If we have the exact solution (i.e., $q^{(k)} = \hat{q}^{(k)}$), then this is not required.

Is it possible to deduce (i) from (iii)?

Thank you very much for your insightful comment. Yes, that is true. We can derive the condition (i) from the third one (iii). Acutally, $q,q'\in\mathcal{P}$, we define $q_t = q' + t (q-q')$ that is mixture of $q$ and $q'$. Then,

$F(q) - F(q') - \int \frac{\delta F}{\delta q}(q')\mathrm{d}(q-q') = \int_{t=0}^{t=1} \int \frac{\delta F}{\delta q}(q_t) \mathrm{d}(q-q') \mathrm{d}t - \int \frac{\delta F}{\delta q}(q')\mathrm{d}(q-q')$

by the fundamental theorem of calculus. Then, the RHS is

$\int_{t=0}^{t=1} \int \left[ \frac{\delta F}{\delta q}(q_t) - \frac{\delta F}{\delta q}(q') \right] \mathrm{d}(q-q') \mathrm{d}t.$

Using the $L_\mathrm{TV}$-Lipschitz continuity of $\frac{\delta F}{\delta q}$, it is bounded by

$\int_{t=0}^{t=1} L_\mathrm{TV} \mathrm{TV}(q,q')^2 \mathrm{d}t = \frac{L_\mathrm{TV}}{2}\mathrm{TV}(q,q')^2,$

which implies the smoothness

$F(q) \leq F(q') + \int \frac{\delta F}{\delta q}(q')\mathrm{d}(q-q') + \frac{L_\mathrm{TV}}{2}\mathrm{TV}(q,q')^2.$

From Pinsker's inequality, we also obtain

$F(q) \leq F(q') + \int \frac{\delta F}{\delta q}(q')\mathrm{d}(q-q') + L_\mathrm{TV}D_\mathrm{KL}(q\|q').$

Here, we emphasize that when the inner-loop error is ignored, it is possible to prove convergence using only the smoothness w.r.t. $D_\mathrm{KL}$ instead of the Lipschitz continuity of $\frac{\delta F}{\delta q}$. Your excellent suggestions have been incorporated into the manuscript.

Characterization of convergence in the nonconvex setting

Thanks for your insightful question. It is indeed shown that the solution converges to a (near) stationary point as we presented in Corollary 1. What Corollary 1 claims is that the functional derivative of the regularized loss also converges to zero: with sufficiently large $k$, we can approximately evaluate

$\min_{k=1,...,K} \mathbb{E}\_{x\sim q^{(k)}}\left[ \left( \frac{\delta L}{\delta q}(q^{(k)},x) - \mathbb{E}_{x'\sim q^{(k)}}\left[\frac{\delta L}{\delta q}(q^{(k)},x')\right]\right)^2\right] = \mathcal{O}\left(\frac{1}{K}\right),$

because $D_\mathrm{KL}(\hat{p}^{(k)}\|\hat{p}^{(k+1)})$ will be approximately the moment generating function $\psi_{q^{(k)}}$ of $\frac{\delta L}{\delta q}(q^{(k)})$ and asymptotically approaches the variance of it. Therefore, we interpret the above equation as

\frac{\delta L}{\delta q}(q^{(k)},x) \to 0~~~(\text{up to constant w.r.t. x}),

which implies that the solution $q^{(k)}$ converges to a stationary point as in the standard finite dimensional optimization.

In the event that all your concerns have been fully resolved, we would be grateful if you could think about adjusting the rating score positively, if it aligns with your judgment.

We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.

Implementation details for applying Dual Averaging to diffusion models

This is a crucial question that addresses how we designed our algorithm. First of all, we would like to emphasize that our algorithm is designed so that we can avoid direct computation of the quantities $D_\mathrm{KL}(q\|p_\mathrm{ref})$ and $q \propto \exp(-f)p_\mathrm{ref}$ because they are hardly obtained as you pointed out.

We addressed these computational challenges by focusing on estimation of the density ratio of samples at the final step $q(x_T) / p_\mathrm{ref}(x_T) \propto \exp(-f(x_T))$ that is tractable with Dual Averaging. Actually, we highlight that $f$ is the weighted sum of the functional derivatives of $F$, which can be calculated with the estimated $f$ in the previous iterations within the Dual Averaging framework.

For instance, when $F$ is the DPO objective function, the functional derivative of $F$ can be calculated by samples generated from $p_{\mathrm{ref}}$ (which is already given by a diffusion model) and the last iterate $f_k$. More specifically, the DPO objective $F$ can be represented as the expectation of $\log \sigma(-f(x_1) + f(x_2))$, where $x_1, \; x_2$ are the win and lose data points respectively which can be easily generated through the diffusion model corresponding to $p_\mathrm{ref}$ (and $\sigma$ is the sigmoid function). Additionally, the computation of KTO objective and its functional derivatives are also feasible in a similar manner.

Just for clarification, we present an implementation of the Dual Averaging method as follows:

1. Initialize $\hat{f}_0$ by neural networks;
2. Get i.i.d. data points $x_1,\dots,x_n$ from the diffusion model corresponding to the reference model $p_\mathrm{ref}$;
3. Iteratively estimate $f$;
   1. Empirically calculate $\frac{\delta F}{\delta q_k}$ for each $x_i$ with the density ratio $\exp(-\hat{f}_k)$ given in the previous steps;
   2. Construct $\hat{f}\_{k+1}$ by estimating the average of $\frac{\delta F}{\delta q_k}$ by minimizing the mean squared error on the tuples $(x_i, \frac{\delta F}{\delta q_k}(x_i))_{i=1}^n$.

Note again that, during the algorithm, we don't need to calculate $D_\mathrm{KL}(q\|p_\mathrm{ref})$ and $q \propto \exp(-f)p_\mathrm{ref}$.

Role and utility of phase 2 (sampling with correction term)

I appreciate your input, but Dual Averaging optimizes obtains only the output density ratio of samples at the final denoising step. The role of Phase 2 is to construct the denoising path with the correction terms, looking ahead the optimized final output density. The correction term at each denoising step is uniquely determined by the final output density ratio with Doob's $h$-transform technique.

In addition, we agree that the proposed correction term requires multiple reverse process, however, this kind of techniques like Phase 2 has been empirically used in image generation tasks [R1], Bayesian computation [R2] and filtering [R3].

Particularly, as a more practical option, the idea of approximating the correction term using neural ODE solvers for faster implementation has also been proposed in [R3].

Why do we need Phase 2?

It should be noticed that knowing the density does not imply that we can sample data from the density.
You may consider that we can apply the Langevin dynamics to generate data from $q_K$, however, it is almost intractable in real complicated data sets such as images due to their high dimensionality and multimodality. Instead, our strategy is to make the most use of the sampling efficiency of the diffusion model. It has been empirically and theoretically known that the sampling efficiency of the diffusion model is independent of the isoperimetric inequality (such as log-Sobolev inequality) of the target distribution, which is especially favorable for high dimensional multi-modal data.
Our algorithm aims to construct an aligned diffusion model just by adding a correction term in Phase 2 in which we use the $h$-transform technique. In other words, Phase 2 is required to convert the density ratio to an (efficient) sampling algorithm.

$L_p$-smoothness of the score $\nabla \log p_t$ and isoperimetric conditions

Your point is important; indeed, it is the case that $L_p$-smoothness of the score $\nabla \log p_t$ implies log-Sobolev in equality with some constant $C > 0$, assuming that $p_\mathrm{ref}$ satisfies log-Sobolev inequality. However, the log-Sobolev inequality becomes very loose when the dimension $d$ goes extremely high. The constant $C$ becomes exponentially small with $d$, as we mentioned in the line 178 to 180 in the first submission; which is highly complex and has multi-modality, leads to failure or weak satisfaction of LSI. As a result, the convergence of MFLD is significantly slowed down. Therefore, $L_p$-smoothness of the score is weak enough to add it.

Since your suggestions were very valuable, we have clarified these parts you pointed out.

If all the issues you raised have been adequately resolved, we would humbly request you to consider revising the rating score upwards, should you find it reasonable.

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

Reviewer 141C,

Thank you very much for your work on reviewing this submission.

Could you read the author's reply, check if they have addressed your comments, and acknowledge it?

Further, please make sure, if necessary, to explain what your new stance is based on this information.

Thank you for your thoughtful and insightful questions. Below, we address them in detail:

• Assumption on the Correction Term

  • You are correct that we assumed the correction term at each step of the denoising process can be estimated with the error $\epsilon_{\rho,l}^2$ as stated in Theorem 3. However, we also demonstrated in Theorem 4 that $\epsilon_{\rho,l}^2$ grows linearly with $d$, which is relatively acceptable in high-dimensional settings.

• Catastrophic forgetting

  • First, our approach leverages entropy-regularized fine-tuning, which mitigates catastrophic forgetting in the aligned model by incorporating KL-regularization into the loss function, as detailed in [R1, R2, R3, R4]. This KL-regularization enables us to preserve the "naturalness" of the pre-trained model by forcing the aligned model to have the same support as te pre-trained diffusion model, which prevents reward collapse and catastrophic forgetting.
  • Second, the correction terms are produced by a separate model, ensuring that their inclusion helps safeguard the pre-trained model and reduces the risk of catastrophic forgetting compared with such methods as classifier-free guidance and RL-finetuning which directly fixes the pre-trained model, as demonstrated in [R2].

• Purpose of Phase 2

• • We would like to emphasize that in our setting, the densities $p_\mathrm{ref}$ and $q$ themselves remain unknown, instead, only the estimate of the density ratio $q / p_\mathrm{ref}$ is available. Hence, we cannot exploit the density information directly.
  • Additionally, the sampling from the Gibbs distribution $\exp(-f)$ is quite challenging when it does not or weakly satisfy isoperimetric condition even when $f$ is accessible. In fact, this motivates the development of new sampling methods (e.g., [R5]) that utilize diffusion models without relying on ispperimetry. We note that our problem also falls in the similar situation.

[R1] Masatoshi Uehara et al. (2024). Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control. ArXiv preprint. https://arxiv.org/abs/2402.15194

[R2] Xiner Li et al. (2024). Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding. ArXiv preprint. https://arxiv.org/abs/2408.08252

[R3] Wenpin Tang. (2024). Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond. ArXiv preprint. https://arxiv.org/abs/2403.06279

[R4] Hanyang Zhao et al. (2024). Scores as Actions: a framework of fine-tuning diffusion models by continuous-time reinforcement learning. ArXiv preprint. https://arxiv.org/abs/2409.08400

[R5] Huang, X., Zou, D., Dong, H., Ma, Y. & Zhang, T.. (2024). Faster Sampling without Isoperimetry via Diffusion-based Monte Carlo. Proceedings of Thirty Seventh Conference on Learning Theory, in Proceedings of Machine Learning Research 247:2438-2493 Available from https://proceedings.mlr.press/v247/huang24a.html.

We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.

The role of hyperparameters

Thanks for paying attention to an important point.
The hyperparameter $\beta'$ plays a similar role to an inverse learning rate and should be set so that the solution does not diverge in the theoretical analysis. More concretely, when the objective $F$ is smooth, we have shown that the objective function with additional regularization imposed by the hyperparameter $\beta'$

$\tilde{L}\_k(q) = F(q) + \beta D_\mathrm{KL}(q\|p_\mathrm{ref}) + \frac{\beta'}{k} D_\mathrm{KL}(q\|p_\mathrm{ref})$

monotonically decreases during our DA algorithm under the constraint of $\beta'$. This result also implies that the hyperparameter $\beta'$ controls the speed of the convergence of the regularized objective $F(q) + \beta D_\mathrm{KL}(q\|p_\mathrm{ref})$. In other word, the smaller $\beta'$ yields a more aggressive update while there would not be a theoretical guarantee when $\beta'$ violates the constraint.

In our experiment, we chose the hyperparameter settings in a naive manner to simplify the experiments. However, we have observed how the DPO loss during the Dual Averaging behaves when $\beta'\, (>\beta)$ varies with Gaussian Mixture Model. The additional experimental results with $(\beta,\beta')=(0.04,0.04\times i)~(i=1,\dots,5)$ showed that the smaller $\beta'$, the faster the convergence is, confirming that it indeed plays a role similar to that of an inverse learning rate.

The score estimation error

Thanks for pointing out an important point. Indeed, the $L^\infty$-error constraint is not necessary, and it can be replaced by the $L^2$-error constraint because we are using only the $L^2$-error in the proof. We have fixed this point in the revised manuscript.

Should all of your concerns be satisfactorily addressed, we would greatly appreciate it if you might consider improving the rating score, if you deem it appropriate.

We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.

The motivation of Dual Averaging

Indeed, we can derive a Gradient Descent (GD) version of DDO, however, we also see that the GD approach does not directly give the density ration $q^{(K)}/p_{\mathrm{ref}}$ required for the $h$-transform calculation. Indeed, the update of the GD method can be formulated as the following type of an optimization problem:

$q^{(k)} \leftarrow \arg\min_q \mathbb{E}_q[g] + \beta'^{-1} \mathrm{KL}(q || q^{(k-1)}).$

This update gives only the density ratio between the current solution $q^{(k)}$ and the last update $q^{(k-1)}$ (i.e., $q^{(k)}/q^{(k-1)}$) instead of the direct ratio $q^{(k)}/p_{\mathrm{ref}}$ that can be obtained by the DA approach. We employed the later one because the GD approach should take the product $\prod_{k=1}^K q^{(k)}/q^{(k-1)}$ to obtain the final outcome $q^{(K)}/p_{\mathrm{ref}}$ which is complicated, computationally demanding and could cause an accumulated error.

Estimation of the correction term

Note that we are essentially estimating the Radon-Nikodym derivative between $q^{(K)}$ and $p_{\mathrm{ref}}$ which can exist even when the density is degenerated. Actually, under some boundedness condition on the first order variation, the density ratio can be bounded and hence we can estimate the $h$-transform even when the density is ill-posed.

Moreover, the time-discretization error in the reverse process can be well regulated as in our Theorems 3 and 4 by paying attention to the short-time Gaussian approximation at each time step.
This discretization error vanishes as the time step size goes to zero and the score estimation becomes accurate.

Presentation of proof

Thanks for your suggestion. Please notice that there is a technical overview of the argument at the beginning of the proof section (C.1). However, following your comment, we have added a concise explanation of each lemma before each lemma is presented. Especially, we considerably enlarged Section A to give the proof overview for the convergence of our DA method.

Notation

The definition of the functional derivative (also known as first order variation) is quite standard. We have added a sentence to make the definition clearer. We also added a table of notations at the beginning of the appendix, which we believe helps the readers a lot.

To address your other questions regarding the notation, $\nabla_x \nabla_x^\top$ refers to Hessian matrix, so the $i,j$-th component of the $d\times d$ matrix $\nabla_x \nabla_x^\top f(x)$ is $\frac{\partial}{\partial x_i}\frac{\partial}{\partial x_j} f(x)$. $\nabla_x \nabla_x^\top$ is also abbreviated as $\nabla_x^2$, so we have standardized them to the latter. In addition, the sub-indices $\mathbb{E}[\dots | \bar{X}\_t = x]|_{x=\bar{X}_t}$ refers to substituting the random variable $\bar{X}_t$ into $x$. It might have been simpler to write $\mathbb{E}[\dots|\bar{X}_t]$, so we have changed this notation to this.

Proof citations for the Dual Averaging method

The convergence rate of the DA method for non-convex objective has not been so much interested in the literature, and thus as far as we know, there is not a direct reference for that (although an accelerated DA method with a line searched step-size is considered in [R1]). For example, the original paper by Nesterov (2009) dealt only the convex setting. Thus, we cited Liu et al. (2023) for the most relevant reference.

[R1] Yurii Nesterov, Alexander Gasnikov, Sergey Guminov, Pavel Dvurechensky: Primal-dual accelerated gradient methods with small-dimensional relaxation oracle. 2018.

Minimizer of the DPO objective

The minimizer of the DPO objective would not be a strictly single Gaussian, because we also have the regularization coefficient $\gamma$ in the DPO objective. However, if the preference label is generated exactly from the Bradley-Terry model, then the minimizer of DPO with infinitely many preference labels (pairs of winning and losing data) and zero regularization parameter will exactly become the single Gaussian.

Guarantees for learning the sequence $q^{(k)}$

For a non-zero $\beta$, the optimal density ratio is bounded and smooth for a non-zero KL-regularization term. In that situation, our theory implies that our method with a sufficiently large neural network can estimate the optimal density ratio with small error.
As for an estimation error, we need to go through a finite sample generalization error analysis but it is not in the scope of this paper. We would like to defer it to the future work.

If all of your concerns have been fully addressed, we would sincerely appreciate it if you could kindly consider raising the rating score, should you find it appropriate.

Thank you very much for all your replies, and for the added work that you have done on this submission. I will take of all these into account for my meta-review.

Would it be possible for you to comment on the relationship of your work with (Marion et al., 2024)? It seems related to your project.

Marion et al: Implicit Diffusion: Efficient Optimization through Stochastic Sampling - https://arxiv.org/abs/2402.05468