Direct Distributional Optimization for Provable Alignment of Diffusion Models
Authors: Ryotaro Kawata
Kazusato Oko
Atsushi Nitanda
Taiji Suzuki
Venue: ICLR 2025 Poster (2025)
Keywords: Diffusion models
Optimization
Primary area: 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.

Published as a conference paper at ICLR 2025
DIRECT DISTRIBUTIONAL OPTIMIZATION FOR PROV- ABLE ALIGNMENT OF DIFFUSION MODELS
Ryotaro Kawatal 12, Kazusato Oko Atsushi Nitanda+.5.1 Taiji Suzuki!
1 Department of Mathematical Informatics, University of Tokyo, Japan
2 Center for Advanced Intelligence Project, RIKEN, Japan
& Deeaarmeent of ECCS, UC Berkeley
ACFAR and IHPC, Agency for Science, Technology and Research (A*STAR), Singapore
"College of Computing and Data Science, Nanyang Technological University, Singapore
loko@berkeeyy.ee
dttsushi.ittanddeffrrrccom edutaarreduu.e edu.sq
taijf @mi st. u-
ABSTRACT
We introduce novel alignment method for diffusion models from distribution optimization perspectives while providing rigorous convergence guarantees. We first formulate the problem generic regularized loss minimization over proba- bility distributions and directly optimize the distribution using the Dual Averaging method. Next, we enable sampling from the learned distribution by approximat- ing its score function via Doob's he-rransfrrm 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 general distri- bution optimization problems, including aligmment tasks in Reinforcement Learn- ing with Human Feedback (RLHF), Direct Preference Optimization (DPO), and Kahneman-Tvertt Tvessky Optimization (KTO). We empirically validate its performance on synthetic and image datasets using the DPO objective.
INTRODUCTION
Diffusion models (Sohl-Dickstein et al. 2015 Ho al. 2020, Song al., 2021) have recently emerged as powerful tools for learning complex distributions and performing efficient sampling. large-
Within the framework of foundation models, common approach involves pre-training scale datasets, followed by adapting the model to downstream tasks or aligning it with human preferences (Ouyang al. (0022). This alignment is typically formalized tion optimization, with a regularization term that encourages proximity to the pre-trained distribu- tion. Examples of such alignment methods include Reinforcement Learning with Human Feedback (RLHF) ZZiegler et al. 2200)), Direct Preference Optimization (DPO) (Rafailov et al. 2023
nonlinear distribu-
Wallace
al. 2024), and Kahneman-Tverske Optimization (KTO) (Ethayarajh et al. 2024
2024a)
Specifically, there methods solve the minimization problems of regularized functional F(g) + (D (gllPet) over q in the probability space P, where Pref the probability density corresponding to the pretrained diffusion model. However, this type of distributional optimization problem over the density q is challenging because the output density governed by the reference model cannot be evaluated and neither is the aligned model g. They are accessible only through samples generated from their corresponding generative models. Existing distributional optimization methods such as mean-field Langevin dynamics (Mei et al. 2018) and particle dual averaging (Nitanda et al.) 2021) resolved this problem by adapting a Langevin type sampling procedure to calculate a functional derivative of the objective without explicitly evaluating the densities. However, the distributions and Prof 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 be mathemati-

---

Published as a conference paper at ICLR 2025
cally characterized by isoperimetric conditions, such logarithmic Sobolev inequality (LSI) (Bakry jet al. 2014), 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 L.II constant, so they suffer from severely slow convergence, failing to align diffusion models.
That to say, alignment diffusion models has two challenges: (i) inaccessibility of the output densities and (ii) muitimodality of the densities. This naturally leads fundamental question:
Can we develop an alignment algorithm for diffusion models from distribution optimization perspectives, while ensuring rigorous convergence guarantees without isoperimetric conditions?
We address this question by developing a diffusion- based distribution optimization method and providing rigorous convergence and sampling error guarantees, and demonstrate its applicability to several tasks involved with diffusion model alignment. Our method represents the aligned model by diffusion model that can be described by merely adding a correction term to the score function of the original reference model. During optimizing the model, we don't 't rely on any MCMC sampler but only use samples generated by the original reference model (and the aligned diffusion model). This characteristics is helpful to resolve the issue of isoperimetric condition.
h-transform
Aligned Diffcsion
Patt
I...K
Averaging
600.0 Viognto = VlogEle [
the density ratio
to minimize linearized i':,
Referanca dersity
1=1
==2
k=3
Pes Ja
Figure 1: Overview of the proposed method integrating Dual Averaging and Doob's h-transform.
1.1 OUR CONTRIBUTIONS
To tackle the two challenges mentioned above: absence of (i) sampling guarantee and (ii) isoperime- try, propose general framework that integrates dual averaging (DA) method (Nesterov) 2009) and diffusion model (Sohl-Dickstein et al. 2013 Ho et al. 2020 Song et al., (021). The DA method is iterative algorithm that constructs the Gibbs distribution converging to the optimal distribution (i.e., alignment). A key advantage of this DA scheme is its ability to bypass isoperimet- ric conditions with the help of high sampling efficiency of the reference diffusion model, enabling isoperimetry-free sampling from an aligned distribution. Specifically, an aligned diffusion process that approximately generates the Gibbs distribution obtained by DA method can be constructed through the Doob's ' h-traasform (Rogers & Williams) 2000) and density ratio estimation with re- spect to a reference distribution using neural networks (see Figure for illustration).
We summarize our contribution below.
We establish a model alignment method to align diffusion models, with convergence guaran- tees for both convex and nonconvex objectives, based on distribution optimization theory and Doob's hi-transform. Notable distinctions from other mean-field optimization methods are that our method works without isoperimetry conditions such as LSI and allows for sampling from multimodal distributions.
We also analyze the sampling error due to the approximation of the drift estimators with neural networks and the discretization of the process, and evaluate how these errors affect the final sampling accuracy.
Our general framework encompasses several major alignment problems such RLHF, DPO, and KTO, establishing a provable alignment method for these scenarios. We demonstrate the applicability of our framework to these settings and empirically validate its performance
2

---

Published as a conference paper at ICLR 2025
both synthetic and image datasets, aiming at data augmentation for specific mode of distribution, using DPO objective.
We also emphasize that method has the potential to be applied to general distribution opti- mization problems beyond alignment tasks such as density ratio estimation under the covariate shift setting (Sugiyama et al. 2008 T'sbooi al. 2009) and climate change tracking (Ling et al. 2024).
1.2 RELATED WORK.
Mean-field optimization. PDA method (Nitanda et al. 202 Nishikawa et al. 2022 an exten- sion of DA method (Nesterov] 2009) to the distribution optimization setting, was the first method that proves the quantitative convergence for minimizing entropy regularized convex functional. Sub- sequently, P-SDCA method (Oko et al.) 2022), inspired by the SDCA method (Shalev-Shwartz & Zhang 2013), achieved the linear convergence rate. Mean-field Langevin dynamics (Mei et al.) 2018) is the most standard particle-based distribution optimization method, derived as the mean- field limit of the noisy gradient descent, and its convergence rate has been well studied by Mei ic al. 20188; Hu et al. (202T); Nitanda et al. (2022); Chizat (2022); Suzuki et al. (2023); Nitanda (2024). Additionally, several mhean-field optimization methods such meen--field Fisher-Red Rao gradient flow (Liu et al., (023b) and entropic fictitious play (Chen et al.) [023a Nitanda et al. (023) have
been proposed with provable convergence guarantees. We note that the convergence rates these
methods were established under isoperimetric conditions such as log-Sobolev and Poincaré inequal- ities, which ensure concentration the probability mass. IKLPD method (Yao et al.) (024) shares similarities with our method, as it employs the normalizing flow to solve intermediate subproblems in the distributional optimization procedure, and its convergence does not depend on isoperimet- ric conditions. However, the applicability of IKLPD to alignment tasks remains uncertain since handling the proximity to the reference distribution is non-trivial.
Fine-tuning of difflusinn models. Recently, alignment of diffusion models have been investigated. inspired by LLM fine-tune methods, such as RLHF (Ziegler et al. 2020). DPO (Rafailov et al., 2023), and KTO (Ethayarajh et al. 2024). Applying them diflusion models entails additional difficulty since the output density Pref of the diffusion model is not available, and hence several techniques have been developed to circumvent the explicit calculation of Pref. For instance, |aan et al. (2023); ; Black et al. (2024); Clark et al. (2024) invented the maximization algorithm of the reward each diffusion time step. Uehara et al. (2024b) used Doob's l-ttranffrrm to compute the correction term that can be automatically derived from the density ratio between the generated and the reference distributions. Instead of optimizing original DPO and KTO objectives, Wallace et al. (024) considered the evidence lower bound (ELBO) and Li et al. (204a)) defined a new objective function to replicate KTT. Marion et al. (2024) also studied fine-tuning of diffusion models as distributional optimization within the RLHF framework and conducted convergence analysis for the one-dimensional Gaussian distribution.
2 PROBLEM SETTING
Distributional Optimization. Let P be the space of probability density functions with respect to the Lebesgue measure on (Rd, B()).. Let F : P > R be functional and pref E P be the reference density. In this work, we consider the regularized loss minimization problem over P:
min (T(g) = F(q) + PKL (allPret)) gEp
where DKL (gl(gliPeet) = Eyllog 2 is the Balllack--eibbee divergence, and 3 > is regularization Prot coefficient. We assume that F is differentiable. That is, the functional F has first order variation bg P x Rd (g. + bq oF (q, x) E R such that for all q, q' E P.
d dF + c(d' 9))
=
8F (q.z)(d (x)dr. 89
de
==
In the following we assume that there exists unique minimizer Ωopt := arg min min yEP L(g)-
Diffusion Models. Pref is the output density of a pre-trained diffusion model
et al.
2013 [Song& & Ermon|| 2 Ho et al.] 20200 Song et al.) [202T] Vahdat et al., 20.1), while p. E P the target distribution of pre-training. A "noising" process [X, 1220 denotes the Constein-Uhlenbae
3

---

Published as a conference paper at ICLR 2025
(OU) process from p. (z). The law of X, can be written (x)) [m, 01 dd) (Xo) with my e of 1 E Then, the reverse process 1x: (T' 2 0) can be defined as
X Pe, dx; f + 2V log log ) v2dB.
Then it holds that Law (o,((x)==> Law( (x which enables us to sample from p.. In practice, we approximate the score V PlggPT (x: function by a score network : IR - Rd (a, r> s(x, t). In addition, we initialize X (0, La) = pr and the process is time-discretized. The random variable generated by the following dynamics with step size h is denoted by {X; JOSIST.
XO ~N, Ia), dx;= = (X +2s(( (dd E [th, (1+1))), 1= 1...., I. T//h
In the same way. we define q: as the density of the diffusion process corresponding to qu, and {X.)e20 and {Xf JoSeST as the corresponding noising and the backward process.
Now we have two challenges in this problem
(A). Multimodality of Prof and Tops. Our goal to obtain samples from Jopt. Now we tackle with the case that Pref and Copt have high multimodality; Pref and Jopt have multiple modes or maxima, in other words, the potentials - log Pref and log 4opt are extremely far from concave. Probability distributions of real-world data like images often have such complex structures, which implies that LSI is significantly weak.
Pret
(B). Inaccesibility to the density Pref- We cannot directly calculate the density Pref because diffusion models only have the score network. We only have information of Pref as samples from Pref-
Inapplicability of Mean-Field Langevin Dynamics. Mean-field Langevin dynamics (MFLD) (Mei et al.) 2018) is the most standard particle-based optimization method tailored to solve the special case of the problem ( where F convex functional and Pref is a strongly log-concave distribution such Gaussian distribution. The convergence rate of MFLD (Nitanda
ct al.
2022
Chizat
2022) has been established under the condition where the proximal distribution:
3 16 6F by (a.-)) associated with MFLD iteration q satisfies logarithmic Sobolev inequality (LSI), which says sufficient concentration of Pq- Typically, LSI for Py should rely on the isoperimetry of the reference distribution Pref with Holley-Strook argument (Holley & bg (g. is not expected to encourage it in general. Therefore, the replacement of Prof to pre-trained distribution, 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 that finite particles approximately follow
Py X Pref exp
Stroock 1987) since of
down. Additionally, MFLD is practically implemented the ideal dynamics of MFLD, and it is not intended for resampling from the final distribution represented by these particles. In short, the lack of (i) sampling guarantee and (ii) isoperimetric condition limits the applicability of MFLD in our problem setting.
To address these challenges, we employ the following strategy, which will be detailed in Section 4
(A). Modifications of diffusion models for sampling from multimodal distribution. First, we obtain samples from Pref using diffusion model which works for a broad range of probability distributions that have complex structures such as multimodality. Then, we reconstruct the diffusion model to sample from Ωope by simply adding a correction term.
(B). A distributional optimization algorithm that does not require the density Pref- Sec- ond, following the dual averaging (DA) method in the distribution optimization setting, constract sequence of distributions ((t)) converging to the optimum, which guides an aligned diffusion process in combination with the density ratio estimation using neural net- works with Doob's li-transform. We only have to calculate the density ratio between g(() and Pref and the samples from Pref to run our algorithm.
APPLICATIONS
Our framework of the distributional optimization for the pre- pre-trained diffusion models includes im- portant fine-tuning methods for diffusion models. The first term F(g) in ( represents the human preference, acting as feedback from the outputs of Pref. Note that F(q) can be dependent of Pref- The second term 8D (gllPeat) in I provides regularization to prevent q from collapsing.
Example 1 (Reinforcement Learning.. Our study includes the case F is limited to be a linear functional:
4

---

Published as a conference paper at ICLR 2025
min (Eg]-r(x)] + 3DKL (gllpzet)) qEp
where r(z) is a reward function. In this case, the optimal distribution is obtained as qa (x.
exp
Pret(x). This type of the problems has been studied as the Reinforcement Learning Fan
et al. 2023 Black et al. 2024 Clark al. 2024 Uehara al. 2024b Marion et al. 2024
The following two examples have not been directly solved via diffusion models:
Example 2 (DPO). Direct Preference Optimization (DPO) (Rafailov et al.) 2023) is an effective ap- proach for learning from human preference for not only language models but also diffusion models. Our algorithm can directly minimize the DPO objective, while Wallace al. al.442024) tried applying DPO diffusion models via minimization of upper bound of the original objective. In DPO, humans decide which sample more preferred given two samples from Pret. Let Iw and zı be "win- ning" and "losing" samples from Pret- Fww i In denote the event that Fw is preferred to Et. The DPO objective can be written as
(ạ) -Bearpied B Barp-Postt [oo (7 (7log
g(aw)
-olog
1 lourn (Iro, mi)
Pret Pret((w)
Peet (ies(
where Ezap denotes expectation with respect to x whose probability density is E P, a is a sigmoid function, Izvy y) is one if x and is zero otherwise. Precisely, the functional derivative of LDP is calculated as
Ólpeo (1,z) Beeprrper 69
a (->/(x) +J
[e-1 e-1(3) x- (1r mi)
+ yTBr Proo (1-o(-7s(x + +f())))
[c-Idp
I (Tur Ta)
(2)
where q == = e See Appendix B for the derivation. Therefore, we only need samples from Pref and the log- log-density ratio or the potential f to calculate the functional derivatives. Example (KTO). Our algorithm can also minimize LuTO directly. Assume that the whole data space Rd is split into desirable domain DD and undesirable domain Do. The objective of the original KTO (Ethayarajh et al. (222) is formulated as
LKTO(g) =law-peat
(D (1 o (Allog 2
D DKL (KI(glprret)) LizEDot
Prot
(1- (DKL(glIPeet) - log -. 1IEDDU) +TU
where NDD NUs K. are hyper parameters, and a sigmoid function. Li et al. (2024a) defined objectives compatible with diffusion models based on KTO, but algorithm can directly minimize LKTO. Like the DPO objective, we only have to calculate samples from peet and the potential f of the density ratio (f is defined as q- = e f dpret) ) to calculate the functional derivatives. Please refer to Appendix B for the concrete formulation of 8Lgro
0g
THE NONLINEAR DISTRIBUTIONAL OPTIMIZATION ALGORITHM
Now we make a concrete introduction of our proposed approach. Our goal of the distribution opti-
mization 11 is to train neural network that approximates fopt = llog 4opt + (const.). To achieve
this, we utilize the Dual Averaging (DA) algorithm (Nesterov] 2009 Nitanda et al. 2021 Nishikawal et al. 2022). and we iteratively construct tentative local potential Ji. by approximating the update the DA algorithm. After we obtain fK. we estimate the diffusion model that generates the desired output (approximately) following q.. through Doob's 's h-transform technique.
Dual Averaging. Let fi be randomly initialized potential. First, initialize q²)) oc exp(-fi) where fi is randomly initialized neural network. Then, the distribution (8) is updated recursively by pulling back the weighted sum of gradients from the dual space to the primal space. There are two options of DA methods. For a given hyper-parameter 3' > 0, the update of Option 1 is given as
Phase 1:
(Opt. 1) g(++1) arg min PEP
Σ 1=1
(EE
5F
(q(i)) BDKL (glppor))) + 28" ( og
R(R+1)
exp (-g Proff
(3)
5

---

Published as a conference paper at ICLR 2025
..=,,...,AA) By Lemma in Appendix A.2 g( can be explicitly determined. We train neural network fa+t to approximate Similarly, the update of Option 2 is g and define the next step glit1) & exp given as
where g(k) Ej=1 wj (R) (1) w) By
3M(+01/240(44)
k Σ
(Opt. 2)
g min PEP
2 R(K+1]
Eg
OF 0g
(g - -3g5o)) 2 (illeettr)))
(4)
1=1
Here, we again express as 1( (xx 6F bg
exp(-g (x) Pret( where g(
Σj=1 Wj manner as Option This phase DA update summarized in Algorithm
(66 (q x) Bz))())) with w w/6
3(++1J' Then, q (k++1) is obtained in the same
For the more detailed
algorithm in Option 1, please refer to Algorithm D.1
Algorithm 4.1 Dual Averaging (DA)
Input: s: pre-trained score. fi: initialized neural networks
Output: fx: a trained potential.
Set (()) Pref and aL о exp(-fi)Prof
for == 1,
K- do
Obtain g(k) via the DA algorithm with Option (Eq. (3) or Option 2 (Eq. (4) recurrence formula (47). where g(K+1) x exp(-gi (g)()) Pror the ideal update.
using the
Train a neural network fx+1 to approximate g()), and set ((+++) xx exp (- (- f-fk1))
end for
Phase Sampling with Doob's h-transform. After we obtain the solution [K. we want to sam- ple from 9K xxp(-fx)a Pref. which approximates the optimal solution of (I). When sampling from 9K. it is necessary to obtain the score function related to this distribution. However, con- structing the score function of q, only from the score function of Pref and fx requires particular technique. Specifically, we apply Doob's h-transform (Rogers & Williams 20000) Chopin et al. 2023 Uehara et al. 2024b Heng al., 2024). By introducing the correction term u, RR+1 1 defined by
u. (. t) V log (lcxp((f.(Xx)) X =
the score function of q+ (x, f) written as V log (y) V log Po -) u. (y, t). where is the law of the backward process at time t whose output distribution is the optimal solution We provide the derivation in Lemma 12 and refer readers Rogers & Williams (2000); Chopin et al. (2023) for more details and a formal treatment of Doob's h-transtorm. By approximating log P. (X, T t) by the score network s(x, and the correction term u, (I, t) by u(z, t) and discretizing the dynamics, we obtain the following update
Y ~ ~V Ia), dY," (Y, - 2(s 1h) + u(Y Ih)) V te [1h, (I 1)A).
where (xr, can be computed u(x; () V (xxxp(xx(xx)))x
zz), which can be timated by running the reference diffusion model (X;-). The practical treatment for this is discussed in Appendix C For experimental information, please have look at Section and Algorithm D.2 in
Appendix D.2
5 THEORETICAL ANALYSIS
In this section, give theoretical justification of our proposed algorithm. More concretely, show the rate of convergence of the (inexact) DA method and give an approximation error bound on the diffusion model based on the h-transform.
5.1 CONVERGENCE RATE OF THE DA METHOD
We give the convergence rate of the DA algorithm in the two settings: when F'ss is convex and (II) non-convex, respectively.
For DPO and KTO, it suffices obtain the neural network (*) minimize Ep... ( where the expectation with respect Pref is simulated by generating data from Pret. while obtaining g 6) for general settings requires Doob's h-transform similar Phase 2. Please also refer Section

---

Published as a conference paper at ICLR 2025
(I): Convex objective F. First, we show the rate when F is convex. We basically follow the proof technique of Nitanda et al. (022): Nishikawa et al. (002)). the analysis, we put the following assumption on F.
Assumption 1. The loss Sunction F satisfies the following conditions:
(i) 8F is bounded: There exists By such that | (1)lls> S Bp for any q E P, bq
(ii) &F is Lipshitz continuous with respect the TV distance: There exists Liv such that 11 (q) 87 (d))ll I LTv'VV for any q. 6 E.P
6q
(iii) F is convex: F(g) F(a') + 6F (4)d(q for any q. EPP
Then, we can show that the (inexact) DA algorithm achieves the following convergence rate that yields O(K convergence of the objective.
Theorem 1 (Convergence of the objective in Option 1 Suppose that 3' 2 and train the potential 1 that it is sufficiently close to glk) as (V(g((), ((()()),(())) I ETV for all k. Then, under Assumption Option I satisfies the following convergence guarantee:
K(K + 1) k [() I(aopt) < 2TTTVTTV+
K
2BF K(K+1)
2/' DRL(q (xllett) + 2BB 3-1 +
K
See Appendix A. for the proof. From this theorem, we see that the DA algorithm with Option I achieves O11KK convergence. The assumption (ii) in Assumption [ is required to bound an expectation of ( in the bound. The assumption (iii) is required to bound the difference between the exact update (A)) and the inexact update (*) If the update is exact, we don't 't need this assumption.
(II): Non-convex objective F. We also give a convergence for non-convex loss F. Here, we put the following assumption.
Assumption 2. The loss function F satisfies the following conditions:
(i) 8F is bounded: There exists Br such that | (q) oc < Br for any q E P.
5q
(ii) 8F is Lipshitz continuous with respect to the TV distance: 6F (q)
bn
8F
bq
(1) S Lr TTVTV((,
for any 4,4' EPP,
(iii) F is lower bounded.
As Assumptions (i) and (ii) are the same as the convex case (Assumption]] []. and lower boundedness (iii) is weaker than convexity in [ssumption[] Assumption (ii) induces the following type of smoothness commonly observed in standard optimization: (ii)' There exists SF 2 such that F(q)) S F(g') +
J &F (g')dlg - d') + Se DKL (4llqllg'). When the inner-loop error is ignored, it is possible to prove convergence using only the smoothness assumption (ii)' instead of assumption (ii). For details, please refer to Appendix AAIII]
2
Under this assumption, we can show the following convergence guarantee with respect to the se- quence ( even non-convex setting.
Theorem 2. Suppose that B' > L+L and (M) (**)) both Option and 2 yield the following convergence:
s crv for all k then under Assumption
(i) Tim DKL( (1(()++)11()))))))))))))))1 (1)((()) = O.
(ii) For all K, it holds that
min {e Dxt (Q ( (L)()(1) L(lopt)) + (LTV + BF) =: 1 K Am Amly.., K.
K8
where Cj 38425' for Option and Ck for Option
See Appendix A.I.I for the proof. The proof utilizes analogous argument with Liu al. (202aa) that analyzed convergence of a DA method in a finite dimensional non-convex optimization problem. However, need to reconstruct proof because should work on the probability measure space,

---

Published as a conference paper at ICLR 2025
which is not a vector space, and carefully make use of the property of the KL-divergence. We
see that Vxx K = ( if érv is sufficiently small as cTV = () and thus the discrepancy between ( and g(k) converges. Especially, the convergence of (g
yields the
convergence of the "dual variable" for Option in the following corollary. For that purpose, we define (g) == [(g)+ 8' DrI.(g (geeppet)) (see its similarity to the inner objective of the DA update 3 and 4 and define 1/g (g) = log (Ey(exp((1) + (o))) which is "moment generating function"
of a dual variable g.
Corollary (Convergence in Option 2 Under the same condition as Theorem
we have that
min Wyua) ( Igkgk
8L4 (q(K))
K-
oq
Roughly speaking, this corollary indicates that the variance of the gradient sL4 (Q)*) converges as og
Varguo bq (q(*)) (1/K) because we may approximate (by(g) = 2 Var,(g) when [g Eg [g] is sufficiently small. Therefore, we have convergence guarantee of the dual variable (gradient) aLe (z) + 0 (up to a constant w.r.t. 2) even non-convex setting, which justifies usage of by
our method for general objective functions.
5.2 DISCRETIZATION ERROR OF DOOB'S H-TRANSFORM
We provide the sampling error analysis after obtaining the approximate solution fK- In the analysis of Dual Averaging, Pref was assumed to be known, and the goal was to obtain the optimal solution of ( that corresponds to 9opt = exp(-fopt)Pree
From here, considering that Pref esti- mates P.. we shift our focus to sample from the optimal alignment q. о p.Pr. while the dynamics we implement involves several approximations: (i) time discretization approximation, (ii) approx- imation of the score V. log Pt by s(x, T t), (iii) approximation of p" by p exp(- (f(),.)) (iv) approximation of u. by u.
To evaluate how such approximation affects the final quality of our generative model, we will give a bound of the TV-distance between q. and q PPret by putting all those errors together. To do so, we put the following assumption.
Assumption 3. 1. V Pog is Lp-smooth at every time and it has finite second moment
B[| [ S m < co for all t E R. and x E TR...
2. V log p. is -m-sohth and bounded as Cp P. S SCp constant Cp-
3. The score estimation error is bounded by Ex: (0 (1) (VVV)))ogPrr PT (X (XE)) Se e at each time l.
4.
(X,
u (1 [h)11 S 20.1 for any 1 < S T/h and t E [1h, (1 + 1)h).
This assumption is rather standard, for example, Chen al. (2023c) employed these conditions except the last condition on u and u.. The fourth assumption in Assumption Bis imposed to mathe- matically formulate the situation: qopt/Pret = q-/p.. Then, we obtain the following error bound: satisfied. Then, we have the following bound on the
Theorem 3. Suppose that Assumption
distribution @ of Y
TV 2) < Te + EThh he? pd
+ TI Lp) mh²)
I)).
The proof is given in Appendix C.I It basically follows Chen et al. (202cc b). but we have derived the smoothness V log(4e.:) from that of V log(pt). In this bound, we did not give any evaluation on 2. however, this error term can be bounded follows with additional conditions.
Assumption 4. (i) (. -) is H Lipschitz continuous in a sense that || [V,s(I,)) Vys(y,t) I
Hallæ - yll for any x, y E Rd and 0 < t < T and [[[|s (X hn k < for any k. (ii) There
exists R > 0 such that supt, ( ( log Pe (1) V| log s(x, 1) lopp < R.
Theorem 4. Suppose that Assumptions B Fand (anddja hoold and [|p, pllso S s' and || I SCp. We also
assume pp is bounded and Lipschitz continuous. Then, forany
any choice of0 < h < $< 11({12R))
we have that
22 ZCB {=o.c + R$ (e² + +230(6+ + [T. m + Q2 + dh) + R3 (1 + 2R)
+ min(T 1h, 1/22
8

---

Published as a conference paper at ICLR 2025
where Es, := CL + 2R)28 8 + 02 1+ |=² + Ljd(h + +m and Rip and Lip are constants
introduced in Lemmal T0
The proof given in Appendix C.3
We applied the so-called BismutEElworthy-ii integration-by- parts formula (Bismut 1984 Elworthy & Li 1994) to obtain the discretization error. By substituting
5 vh to balance the terms related to h and 6, we obtain simplified upper bound of Σi T/h h
as ΣT/h her 2 (1+ + 2T + TVh + (T' + log(1/h))s=" These results can be seen as h- Transform extension of the approximation error analysis given in Chen et al. (023c). However, the approximation error corresponding to u has worse dependency on h. This is because the computation of u uses the discretized process Prof and is affected by its error while the ordinary diffusion model does not require sampling to obtain the score.
NUMERICAL EXPERIMENTS
We conducted numerical experiments to confirm the effectiveness of minimizing nonlinear objec- tives. We used Option 1 and / was set to be 0.04. We also compared the DPO objective (Rafailov ett al. 2023) we minimized with the upper-bound (Wallace et al. 2024) using a toy case, Gaussian Mixture Model (GMM). For detailed experimental setting, please refer Appendix [
Alignment for Gaussian Mixture Model. We aligned a score-based diffusion model to sam- ple from 2-dimensional GMM using the DPO objective. The reference density was defined as 2 (W Σ) - N V(22E)) [41 [-2.5, OJ, μ2 = [2.5, 0j, Σ = [[1, OJ, 5]]. The target point was Nw == M22 The preference of points Tw and xy was determined by the Euclidean distance d(,, Pw)
from Mw. Zw > I4 if and only if d(I (ww) < d(x),Mw). We describe the implementation details below:
Dual Averaging. We set the hyperparameter B in [0.04, 0.2), which controls the learning speed illustrated the middle in Figure _
In ktth DA iteration (k: 1.... K), using empirical samples
{2r (1i, 22,)) 1000 from Pref and the previous potential fk-1 (q x e -- Pret) imple- (g(h-1), xx4)) {)))) |,000 to construct fa g(h-1) with mented by neural networks, we prepare ((x; 5F og
the recurrence formula (47) in Appendix D Please note that, to calculate 5F for DPO, only need
8g
the fA and empirical samples from Pret to calculate the expectation Egep..rl as described in Eq. ( The convergence of the true DPO loss is depicted in Figure [ compared to Diffusion- DPO (Wallace) jet al., 2024). The true DPO loss, upperbound, and metric loss (mean of the Euclidean distance of the particles from ).,)) attained by Diffusion-DPO and our method, with optimization time and GPU memory consumption (in Phase for ours), are summarized in Table [ Diffusion-DRP optimizes the approximate upperbound instead of the true loss, and therefore it failed to completely control the true loss. Compared to Diffusion-DPO, our algorithm allowed us to optimize the loss to a smaller value within acceptable computational budget.
Doob's 's h-transform. We constructed the aligned model by calculating the correction terms. The conditional expectations were calculated using naive Monte Carlo method with 30000 samples. The number of the diffusion steps was 50. The histograms of aligned samples are shown at the right of Figure 2 In the simplest phase 2 that present (Algorithm D.2), we used significant computa- tional resources, with O(L2) time complexity. To estimate the correction term u for each sample simultaneously by Monte Carlo with N samples, O(N) memory space is required. However, this Doob's h-transform technique itself has been used in image generation (Uehara et al. 2024a b), Bayesian samping (Heng et al. (2024), and filtering (Chopin et al. 2023). In particular, computa- tional resources can be saved by approximating the correction term using neural ODE (Uehara ct al. 2024b).
Table 1: The quantitative comparison between Diffusion-DDP and our method
Diffusion-DPO (50 iter., w/o reg.)
Diffusion-DPO (200 iter., w/ reg.)
Ours
Ref.
(8' 0.04)
True DPO loss
0.346
0.340
0.343
0.328
(Approx.) Upperbound
0.311
0.337
Metric loss
3.256
3.011
3.303
2.022
Opt. time (s)
585
2280
1320 (Phase D)
GPU memory (%)
6.54
6.54
8.75 (Phase I)
9

---

Published as a conference paper at ICLR 2025
Diffusion-DPO
Ours
Sampling (Ours)
reference
Oijjective (8- 0.00)
0.3475
True Objective ( 0.00
03450
0.300
True Object (8 0.12)
0.423
0.3423
Objective
True Objective (5 0.2)
Sarget
iter
0.1500
0.130
03525
Tue Objective (DHB-PPP)
Upperbound (Def BPO)
0.3225
Thue Objective 5P0 reg)
Uppertouudd (Def-D20 Peli--PD0-wa wi reg)
50
100 Iteration
150
175
100
Iteration
Figure 2: Left and Middle. The smoothed loss during optimization for Gaussian Mixture Model in Diffusion-DPP with/without regularization (left) and ours (middle). "True Objective": the true DPO loss (Rafailov et al. 2023) whose target point was Mer [2.5, 0). "Upperbound": An approximate upperbound "Objective" optimized by Diffusion-DPO (Wallace et al., 2024). Right. Aligned samples by Doob's 's h-transform. "Ref." represents the empirical density of Pref-
Target Color. We aligned the image generation of the basic pre-trained model available Diffusion Models Course (source: HuggingFace (2022) to simulate specific generations for data augmentation. The pre- pre-trained model generates RGB-colored 32 x 32 pixel images of butterflies. In this experiment, the model was aligned based on a target color, which is (R, G. (0.9, 0.9, 0.9) while (1,1.1) corresponds to white, using the DPO objective. Please refer to Figure 7 for an illustration of the preference indicated by DPO. We leveraged 1024 images generated by the pre-trained model train f& in each DA iteration. We see that the (regularized) objective went down during DA and the generated images reflected the target density more, on the right side of Figure Please refer Figure in Appendix for the convergence results.
Image Generation Alignment based
Tilt Correction for Image Generation. The goal of this experiment is data augmentation of im- ages facing the same direction, given only dataset of rotated images. We used 10000 images of Head CT in Medical MNIST (Lozano (017)) and augmented it by rotating images up 40000. Data augmentation of brain images useful for tumor analysis and anomaly detection (Shen al. 2024 Fontanella et al. 2024). In each DA iteration, 6400 images were generated to train Js. The results are in Figure 3] For the convergence results, please refer to Figure 12 in Appendix. From the results, see that the (regularized) objective decreased during DA and the generated images became more preferred.
Ours
Ours
Ref.
iter
iter 2
Ref.
iter
iter 3
Target: light-colored
Target: tilt-corrected
Figure 3: Left. Examples of aligned image generation. Our goal was to generate light-colored butterflies. "iter=2": ours with DA iterations, "iter=2": ours with k = DA iteration. "Ref." samples from Pret- Right. Tilt-correctdd Head CT image generation. "iter=3": ours with & = 3 DA iterations, "iter=1": ours with DA iteration. "Ref,": samples from Pref-
CONCLUSION
We proved that the distributional optimization can be solved even if Pref and dopt are mutilmodal and we only have access to the score, not the density. This setting includes important fine-tuning methodologies for diffusion models: Reinforcement Learning, DPO, and KTO. Our algorithms are based the DA algorithm and Doob's h-transform technique and it can solve them more directly than previous works. They are guaranteed to useful for general objective functions, the dual variable converges even in a non-convex setting. While framework has potential applications in general distributional optimization problems, such as density ratio estimation under covariate shift and climate change tracking, further exploration of these applications is left for future work.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
RK was partially supported by JST CREST (JPMICR2015), KO is partially supported by JST ACT- X (JP-MJAX23C4), TS was partially supported by JSPS KAKENHI (24K02905, 20H00576) and IST CREST (JPMJCR2115. This research is supported by the National Research Foundation, Sin- gapore, Infocomm Media Development Authority under its Trust Tech Funding Initiative, and the Ministry of Digital Development and Information under the AI Visiting Professorship Programme (award number AIVP-2024-004). Any opinions, findings and conclusions recommendations ex- pressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore, Infocomm Media Development Authority, and the Ministry of Digital De- velopment and Information.
REFERENCES
Dominique Bakry, Ivan Gentil, Michel Ledoux, et al. Analysis and geometry of Markov diffusion operators, volume 348 of Grundlehren der mathematischen Wissenschaften. Springer, 2014.
Jean-Michel Bismut. Large deviations and the Malliavin calculus, volume 45 of Progress in Math- ematics. Birkhäuser Boston, Inc., Boston, MA, 1984.
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion wttpa//ooor//////0001.1.333
models with reinforcement learning, 2024.
Fan Chen, Zhenjie Ren, and Songbo Wang. Entropic fictitious play for mean field optimization problem. Journal of Machine Learning Research, 24(211):1-36, 2023a.
Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of score-based generative modeling: User-friendly bounds under minimal smoothness assumptions. In International Conference on Machine Learning, pp. 4735-4763, PMLR, 2023b.
Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In The Eleventh International Conference on Learning Representations, 2023c.
net/ hat//oorrmmiii ZyLVMgs20U
Lénaïc Chizat. Mean-field langevin dynamics: Exponential convergence and annealing. Transac- fions on Machine Learning Research, 2022. ISSN 2835-8856. URL htttthsss///peeeeiiiiiie
het/forummicie wtt.ormmmmddaaq..... gf.
Nicolas Chopin, Andras Fulop, Jeremy Heng, and Alexandre H. Thiery. Computational doob h- transforms for online filtering of discretely observed diffusions. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Pro- ceedings of the 40th International Conference on Machine Learning. volume 202 of Proceed- ings of Machine Learning Research, pp. 19904-922 PMLR, 23-29 Jul 2023. URL https: mtr.cceedingsmmrr.ppop 2
/v02////22222///// chop wataooopn2aaaha.mm
Kevin Clark, Paul Vicol, Kevin Swersky, and David Fleet. Directly fine-tuning diffusion models on differentiable rewards, 2024. Ut/pttss://aaxxv wrw./orgggbbb..o aa////b//223/
K.D. Elworthy and X.M. Li. Formulae for the derivatives of heat semigroups. Journal of Functional Analysis, (25()):25--286 1994.
Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela. Kto: Model alignment as prospect theoretic optimization, 2024. htthtttpss//axxiii
2024.
abs/
2402. 306
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Dpok: Rein- forcement learning for fine-tuning text-to-image diffusion models. In A. Oh, T. Nau- mann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36. pp. 79858-79885. Curran Associates, Inc.,
2023. URL utt.ttttsss/ppppeeeee.eeeeoooeeo
(tcpsaeerr_ieeeeeeeeee wwwes//pppe///2////// fi F1144555aaanaa84000cceeeeeeeeee-

---

Published as a conference paper at ICLR 2025
Alessandro Fontanella, Grant Mair, Joanna Wardlaw, Emanuele Trucco, and Amos Storkey. Diffu- sion models for counterfactual generation and anomaly detection in brain images. IEEE Transac- tions Medical Imaging, pp. 1-1, 2024. doi: 10.110//MM.22244460999
Jeremy Heng, Valentin De Bortoli, and Arnaud Doucet. Diffusion Schrödinger Bridges for Bayesian Computation. Statistical Science, 39(1):90 99, 2024. doi: 10.1214/23-STSS9 URL https:
//doi.org/li 1214/23- STS908
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and Lin (eds.), Advances Neu- ral Information Processing Systems, volume 33, pp. 6840-6851. Curran Associates, Inc., 2020. URL File htt.:/ccbbcccc555aaa00511111.1.0
http://pprooeddnnn...e.e......es Paper. ppp
files/paper/202c/
Richard Holley and Daniel Stroock. Logarithmic sobolev inequalities and stochastic ising models.
Journal of statistical physics, 46(5-6): (6655-6:1159-111 159-1199, 1987.
Kaitong Hu, Zhenjie Ren, David Šiška, and Lukasz Szpruch. Mean-field langevin dynamics and energy landscape of neural networks. In Annales de l'Instiuu Henri Poincare (B) Probabilites statistiques, volume 57, pp. 2043-2065. Institut Henri Poincaré, 2021.
HuggingFace. The hugging face diffusion models course, 2022, htttps:/huuginggacc.cco course 2022.
Shufan Li. Konstantinos Kallidromitis, Akash Gokul, Yusuke Kato, and Kazuki Kozuka. Align- ing diffusion models by optimizing human utility, 2024a. URL BRR/htttpp::/axii.ooro/
Xiner Li, Yulai Zhao, Chenyu Wang, Gabriele Scalia, Gokcen Eraslan, Surag Nair, Tommaso Bian- calani, Shuiwang Ji, Aviv Regev, Sergey Levine, and Masatoshi Uchara. Derivative-free guid- ance in continuous and discrete diffusion models with soft value-based decoding, 2024b. URL hitpps:/aaran
org/ abs 2408 08252
Fenghua Ling, Zeyu Lu, Jing-Jia Luo, Lei Bai, Swadhin K Behera, Dachao Jin, Baoxiang Pan, Huidong Jiang, and Toshio Yamagata. Diffusion model- based probabilistic downscaling for 180- year east asian climate reconstruction. npj Climate and Atmospheric Science, 7(1):131, 2024.
Changxin Liu, Xuyang Wu, Xinlei Yi, Yang Shi, and Karl H. Johansson. Rate analysis of dual aver-
aging for nonconvex distributed optimization. IFAC- Paper-napesnnLLine, 56(2):5209 5214, 2023a. ISSN 2405-8963. doi: http:///o.orr/oo.0.00iiaaao.o0.0
10.11...........22222....11....1 URL https: //oww. www.
scciencedi rect. wttp:/cc..................eeaaaa World Congress.
22nd IFAC
Linshan Liu, Mateusz B Majka, and Lukasz Szpruch. of measures and convergence of mean-fiell birth-death processes. Applied Mathematics & Opti- mization, 87(3):48, 2023b.
inequality on the space
Arturo Polanco Lozano. Medical mnist classification, 2017. URL https:// com/
apolanco httpa///222.eeaaaaaaaaa-a-aaaaa-
Pierre Marion, Anna Korba, Peter Bartlett, Mathieu Blondel, Valentin De Bortoli, Arnaud Doucet, Felipe Llinares-López. Courtney Paquette, and Quentin Berthet. Implicit diffusion: Efficient optimization through stochastic sampling, 2024. URL http::://aaaii..rrrrts/// wwwp/////////..../////////
05468
Song Mei, Andrea Montanari, and PhanMinh Nguyen. mean field view of the landscape of two- layer neural networks. In Proceedings of the National Academy of Sciences of the United States of America, volume 33, pp. E7665-E7771, 2018
Stéphane Mischler. An introduction to evolution PDEs, Chapter 0: On the Gron- wall lemma. 2019. URL https:/ /www ceremade dauphine ft/-miichher/ Ensei /ttp//oee120000 chapp....
12

---

Published as a conference paper at ICLR 2025
Yurii Nesterov. Primal-dual subgradient methods for convex problems. Mathematical Programming, 120(1):221-259, Aug 2009.
Naoki Nishikawa, Taiji Suzuki, Atsushi Nitanda, and Denny Wu. Two-layer neural network on infinite dimensional data: global optimization guarantee in the mean-field regime. In Al- ice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neu- ral Information Processing Systems, 2022. URL https: /wtp//ppeeeeveeeeeeee net/ tenee/forume forum?id Hr 87755GGGK
hnttps://openreviieet
Atsushi Nitanda. Improved particle approximation error for mean field neural networks. arXiv preprint 2405. 15767, 2024.
Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Particle dual averaging: Optimization of mean field neural network with global convergence rate analysis. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neu- ral Information Processing Systems, volume 34, pp. 19608-19221. Curran Associates, Inc., 2021. URL https://pproceeeinggs.oo nwwr:ps...eeuuusse....co.. cc/paper files/paper/20211 Htt.://3....bbb02311100022222223
Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field langevin dynamics. In Gustau Camps-Vall Francisco J. Ruiz, and Isabel Valera (eds.). Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pp. 9741 97577. PMLR, 28-30 Mar 2022. URL https:// http:///w.........ss///////////s
html
Atsushi Nitanda, Kazusato Oko. Denny Wu, Nobuhito Takenouchi, and Taiji Suzuki. Primal and dual analysis of entropic fictitious play for finite-sum problems. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Pro- ceedings of Machine Learning Research, pp. 26266-26282. PMLR, 23-29 Jul 2023. URL mlr ptt.eeeddiggsmm..p.res//s/eeeeee wwtp://222/mmmm/ httm:///m//m/aaaaaaaaaatt
Kazusato Oko, Taiji Suzuki, Atsushi Nitanda, and Denny Wu. Particle stochastic dual coordinate
ascent: Exponential convergent algorithm for mean field neural network optimization. In Interna- tional Conference Learning Representations, 2022. URL https::////
hntp::://ppenneeieeww... http///ppneeeeeeeeeeeeee net/
Forum??d=POgs7AJJea
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal. Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kel- ton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neurul Information Processing Systems, volume 35, pp. 27730-27744. Curran Associates, Inc., 2022. URL
fillesppaper20222 File/ (a2.baaeee2333333333000001011112
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model secretly a reward model. In A. Oh. T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (cds.), Advances in Neural Information Processing Systems, volume 36, pp. 35772-53771.. Curran Associates, Inc., 2023. URL /llee/88obaacoo6//aat3225566ee
http://ppooceeeiie...eeeeer..... wwwrips.ec..c.... wwtppppr.ppppppppppppp filess/paerr222257
Ceteeee-cnffeeeee..c pdf
L. C. G. Rogers and David Williams. Diffussions, Markov Processes and Martingales. Cambridge Mathematical Library. Cambridge University Press, 2 edition, 2000.
Shai Shalev-Shwartz and Tong Zhang. Stochastic dual coordinate ascent methods for regularized loss minimization. Journal of Machine Learning Reseurch, 14(1), 2013.
Yiqing Shen, Guannan He, and Mathias Unberath. Promptable counterfactual diffusion model for unified brain tumor segmentation and generation with mris, 2024. URL https: htttp::/aaxxi; //arxiiv brg//bs//44u 2407. 12678
13

---

Published as a conference paper at ICLR 2025
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrimm thermodynamics. In Francis Bach and David Blei (eds.). Pro- ceedings of the 32nd International Conference Machine Learning. volume 37 of Proceedings of Machine Learning Research, pp. htttps://proceest ngs ml wre.s...ss./...
Lille, France, 07-09 Jul 2015. PMLR. URL html
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data dis- tribution. In H. Wallach, H. Larochelle, A. Beygelzimer, F. 7TAché-BBuc, E. Fox, and R. Gameet (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https::/proceeddngg www.:ngg....urir..... wwwpppppppppppppccpecpce.eee...e files/
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- fional Conference Learning Representations, 2021. URL https: https:://openreview: htttanrreiieeee////e//eee
Forum21d-Px11G120EE
Masashi Sugiyama, Taiji Suzuki, Shinichi Nakajima, Hisashi Kashima, Paul Von Búnau, and Mo- toaki Kawanabe. Direct importance estimation for covariate shift adaptation. Annals of the Insti- tute of Statistical Mathematics, 60:699-746, 2008.
Taiji Suzuki, Denny Wu, and Atsushi Nitanda. Convergence of mean-field langevin dy- namics: time-space discretization, stochastic gradient, and variance reduction. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neu- ral Information Processing Systems, volume 36, pp. 15545-15577. Curran Associates, Inc., www.rnng.nnnuuugooo http:///pppeeti.eeeppppr///...3 pdf
2023. URL
File/
Wenpin Tang. Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond, 2024. JRJ/Atts../arxx...orrr/aaa/////
Yuta Tsuboi, Hisashi Kashima, Shohei Hido, Steffen Bickel, and Masashi Sugiyama. Direct density ratio estimation for large-scale covariate shift adaptation. Journal of Information Processing, 17: 138-155, 2009.
Masatoshi Uchara, Yulai Zhao, Tommaso Biancalani, and Sergey Levine. Understanding rein- forcement learning-based fine-tuning of diffusion models: A tutorial and review, 2024a. URL
htttps:/aaxiv..rgg/// awtp://b///// 13734
Masatoshi Uchara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathanicl Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine-tuning of continuous- time diffusion models as entropy-regularized control, 2024b. URL https::/ /aapps://arxi...o org/
hbs/2402.1510
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling latent space. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 11287-11302. Curran Associates, Inc., 2021. URL
paper files/ pdd.
Patrick von Platen, Suraj Patil, Anton Lozhkov. Pedro Cuenca, Nathan Lambert, Kashif Ra- sul, Mishig Davaadorj, Dhruv Nair, Sayak Paul, William Berman, Yiyi Xu, Steven Liu, and Thomas Wolf. Diffusers: State-of-the-arr diffusion models. https:///itthu com/ 2022.
Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 8228-8238, 8238, June 2024.
Rentian Yao, Linjun Huang, and Yun Yang. Minimizing convex functionals over space of probability measures via KI. divergence gradient flow. In Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li
14

---

Published as a conference paper at ICLR 2025
(eds.), Proceedings The 27th International Conference Artificial Intelligence and Statisties, volume 238 of Proceedings of Machine Learning Research, pp. 2530-25355 PMLR, 02-04 May 2024. BRLLLttps::pp/rceeiiiiee
mlr. wwtppppwwseeeppeppppppppppt.ete
v2
html
Hanyang Zhao, Haoxian Chen, Ji Zhang, David D. Yao, and Wenpin Tang. Scores actions: framework of fine-tuning diffusion models by continuous-time reinforcement learning, 2024.
URL http::/p///i....o////aaaa/a abs/249../244...../// 08400
Daniel M. Ziegler, Nisan Stiennon, Jeffrey Wu, Tom Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences, 2020.
URL 0ttxxi/..rr/bbb/190.008855 08599
15

---

Published as a conference paper at ICLR 2025
NOTATIONS
Table 2: List of commonly used symbols
Symbol
Description
dimension the sample space R2
B(Rd)
Borel set on IR
P
The space of probability density functions to the Lebesgue measure on (IR", B(R))).
F Rd P + R
A functional.
0q ())
The functional derivative of F. We often abbreviate this as 8F bq
(4)
(gl())
Kullback-Leibler divergence between g and p.
TV(g,))
Total Variation between q and p.
L(g)
The regularized loss F(g) + 3DI (gllpro).
B
The regularization coefficient of L(g).
P. E dP
The target distribution of pre-training.
Pref E P
The output distribution at the final denoising step of the pre-trained model.
9. Peps & oxp(-S.)y- EP
The target distribution of alignment.
4op: PoptPret exp(- Pret E P
The unique minimizer of L(g), that corresponds to q. when Pref = P.. We assume that pa x exp(-f<) can be estimated by minimizing L(g), i.c. Popt ~ p..
g(k) x
The optimal solution of the subproblem solved iteratively during Dual Averaging.
exp(-g (()))
(a) x exp(- fw Pret
The implemented density that estimates ((k)
gr
The hyperparameter of Dual Averaging that controls the
optimization speed.
The optimal correction term sample from q.-
The implemented correction term using the estimated output density ratio Pupt.
N(1,o)
Gaussian with mean and variance 2
T>0
The diffusion time.
(Biheeoom
Brownian motion on iRd
(E.hrejo.T)
The Cinstenn-Uhlenben (OU) process starts from P..
16

---

Published as a conference paper at ICLR 2025
The reversed Ornstein- Omstein-Uhlenbeck (ROU) process of {Xi hejo7) initialized at pr.
(XF he(o.T]
The time-discretizen and score-estimated ROU process of (Xrhe initialized as X NN La).
(Yihejo
The Ormstein-Uhlenben (OU) process starts from 9x.
The reversed Ornstein-Uhlenbee (ROU) process starts from (r-
{xr hej..
The time-discretized dynamics using s and u.
The law of Xt, and the law of X The law of the OU process Y, corresponding qx P.P-- The discretized time step of { Jæe[o,
41
The number of steps of {XF Jee
R: R - Rddd
A score network to implement (x; )
BF
A constant that bounds 5F
6g*
LTv
Lipschitz constant of bg
SF
Smoothness constant of the weaker smoothness of F.
Vaf(x)
Gradient of with respect to ar
Vaf(x)
Hessian matrix of
-|
Euclid norm
10
_ norm
llop
The operator norm
E||
the conditional expectation given a randam variable X.
E X = x.]
the conditional expectation given a randam variable X, evaluated at X
Big-O notation
22
inequalities ignoring constants.
17

---

Published as a conference paper at ICLR 2025
A CONVERGENCE ANALYSIS OF DUAL AVERAGING
Overview for convex loss
In Section A. we interested in the convergence of Option when the objective functional P is
convex in distributional sense:
8F (q')d(q d'), for any 9.9' P. δq
F(4) F(d) +
Then, the regularized objective L(q) F(q) 3KL (gllptet) becomes strongly convex with B >> 0. Formally, we assume that
(i SF is bounded: There exists BF > 0 such that 1| (g) S BF for any g E P,
Ag
(ii) is Lipshitz continuous with respect to the TV distance: There exists Lrv > such that
6g
( (l S LTVTV(q q') for any q, E.P 3g (iii) F' is convex: F()) 2 F(d') + AF by (g'ddlq d') for any 9.9' E.
ag (q)
In the update of Option 1, we iteratively use the distributional proximal operator. We minimize
2j k(k + g6ا
+ 8 2 DKa(gllPret)
En
+
(5)
by the minimizer (()), where (k denotes by (x)) with (k)) = q()) (a( constructed by neural networks and samples from Pret. while do not directly compute ( Intuitively, with k oc,
we hope that is almost converged around the minimizer Πopt. i.e. = Лopt. Then, from the differentiability of F, each [g in the first term would be linear functional that well approxi- mates ()) with suficiently large j. In addition, 3 DKL((llPedd vanishes, so the equation (5) roughly written as
k
2j
Ey [00)] + (gllpe))) + 8 (Kl(ellree))
k(K + 1)
1=1
k
2j
(F( + 3DKLL (ppe))))))))))))))))))))())))( + 3 DDLLellPPed)
k(k + 1)
k
2j
(F(G()) (I
k(k + 1)
This is the weighted average of the regularized losses L(q) F(q) + (ellpell))t... From this concept, we will show that the convergence represented as
Weighted average of L()(t)) = [Weighted average of L(gone)) + +0
Overview for nonconvex loss
In Section A.2
interested in the case that F is not necessarily convex (mainly) in Option 2.
Alternatively, we use the assumption that smooth in terms of Ki--divergence:
(ii) (The weaker smoothness derived from (ii).. There exists SF 2 such that
/
5F (d')dla- d) + Dkl(glla') for any g. (6) ōg
SF
F(g) F(())
This is induced by Lipschitz continuity of 6F
by ((i) of Assumption and 2 and Pinsker's inequality.
When the inner-loop error is ignored, is possible to prove convergence using only the smoothness T instead of Lipschitz continuity of &F
87 Please note that KI KL-divergence would be the "quadratic
18

---

Published as a conference paper at ICLR 2025
term" in the ordinary definition of smoothness. Following the theoretical analysis of standard non- convex optimization, our goal is to show that the "derivative" of the (regularized) objective L(g)
goes to zero in the form of functional derivative:
(g(h) .))
(up to constant w.r.t. 22.)
óq
Please note that we can ignore derivatives.
8L (g x)) is a constant because of the definition of the functional bq
At the first step, both in Option and 2. the regularized objective L(g) (roughly) monotonically decrease during the Dual Averaging. Ignoring some terms, we want show that
L(q L(q - DKL (q (1 < O.
Next, by taking a telescoping sum, we obtain that
K (
K
From the above equation, it is immediately shown that
(
min ( k=I....
where (weight) à: & in the Option 1, while (weight) ~ in the Option Therefore, when we choose the Option 2, it holds that
min DKL (1 (((h+)) k=1....K K
K
(L) 0 (up to constant w.r.t. X) because DKL (KL(g()) (g(k+))) is ag
Finally, we get the conclusion
the approximate "moment generation function" of (L)()))),, when t: co,
5y
min "Variiacee" of (L(g())) k=I...K
min k=1....K
DKL (t ((((1))) =0
1 K
5q
This result aligns with the result for standard, non non-distributinal Dual Averaging (Liu et al. 2023a).
mink=l.....K (gradieett of the objective) (1//K) because the variance is the second moment.
A.1 WHEN F IS CONVEX
We will show that, when is convex in the distributional sense, the convergence of the Option can be written
(
1 K
Weighted average of (g)) Weighted average of L(
We require 3 to be positive to make the regularized objective F(q) +B)). (gllpret) strongly convex. In addition, one of the necessary conditions 3' implies that the "learning rate" should be controlled to converge.
Theorem (restated Theorem Assume that F is convex, 0g respect in TV distance, TV (G(h), (())) eTV, and 3' 2 B. Then,
S Br,
dg Terv-Lipshitz with
2
k F(g(t)) (g()()(A))pppall - F(Gopt) 3D (Sops [Pwett)
K(K +
2
<
Σ
22TTVTTV + Eg00 Eggyy (x) + ( DICT. (Gopt (lPret))
K(K +1
2
BF + 3(K + DKL (Rops llPeet) + Eggoss (1 B..K
<2TTVERD +
K(K +
19

---

Published as a conference paper at ICLR 2025
First, want show that the DKL regularization term plays role of the "quadratic" regularization, which makes the regularized objective F(q) + 3D (q Pree) strongly convex. We observe that the sum of the linear (convex and concave) objective and the DxL. regularization term
(4)) = [x))r())] + (Dxr.
is strongly convex in the distributional sense. In particular, if q is the minimizer of F, then
B " F(4) for all g' P.
F(4) +
This result is useful to prove Lemma Bgl 68F [ + 3DL (ellpttt)) 6q
which controls the regularized and linearized objective
Lemma (Nitanda et al.
2021 ) F: + F(g) = = Assume that is
(gll()
bounded. We putri r1 dp 222
dg'
dp
F(g) F(y') + Ey((r1 (2)((()) + Blg B In2 r (2) Especially, ifg x exp(-5)p.
B ||1 rr2l r21)2.00)
F(g) S F(d')
2
Proof. Omitted. See Nitanda et al. (2021). We use Pinsker's inequality TV(g,4)3 & DKl(gll').
Given that F(a) + DDc. is strongly convex, let us prove Theorem 1
Proof. The main parts of the proof follow existing papers NNtanda et al. 2021 Nishikawa al., 2022). The key difference lies in that we are not using Langevin sampler in the inner loop, which slightly changes how the eтror of (8))) is handled.
6F (1) and ÿopt denotes the We begin by linearizing the weighted sum the losses. (a) denotes by minimizer of F. From the convexity F,
K
2
(( + (DLL( (Pret) F(oopt) BDDL (Gope (Pree)))))t)
K(K +
1)
K
6F
ki
(q
(d((t) - (opt) + B(DKT (f (g(h))lprett)) DKT.(Gope (Ipror))
K(K +1)
5q
k=1
2
K
k
8F
(k)))
5F
(a (d)((())))) @ops) + Bgres (0 - Eg (b)]
K(K
+1)
ōg
oq
+3( (R)
([pret)
DKL. (aopa [1prrot)
Here, bound the estimation error of (k)) respect to q:
8F
(k)
using the Lipschitz continuity of 8F with
dq
bg
8F | (gth)
dl) pope)
6F (x)),, xx) g (x)
(H) aopt)
sup
<2Lgerv.
20

---

Published as a conference paper at ICLR 2025
Then obtain
F(((()) 3B (P F(oopt) - (oope (Ppeet))))
K(K + 1)
k 22 + Egess [k)]
K(K 1)
+B(L (1 (pprot) DRL (Sops [lPrett)]
This implies that it is sufficient to bound the weighted sum of the (regularized) linearized objectives Eges (x) + 3{ DxL ( (Ipoot) k = 1,..., K.
In each update of Option 1, ((++1) is obtained by maximizing
Bk(k 1)
V = -En
Σ
ig
DKL (gllpeet)) - Bk k + ()llx(el)
We also define
dg(k+l)
(k+1)
V V=+1(1(H=1))
74
dpret
Then, we can show that V/2 has the following recursive relation. Lemma approximately implies
ge [g [*) k DKI. (g [pror)) S VE Vi+1 O(1,, then, summing from Ae to K, roughly obtain,
2
S 0 AK
G [g + 3DKL (g() lIrret) VK+1
K(K +1
So, as long as we prove Lemma we can connect this inequality to the inequality
because
V VK+1 2 ΣK Ek=1 k(Eg.p p (g (k) + 3DkL (@opellPrct)):
Lemma 2. For ls 2 1.
B.K
V2+1 S V2 Lg (3k B¹)
(KL( ((pee)) +
3(k 1) + 2,3"
For the proof of Lemma
2 please refer to Section A.1.I in which we use Lemma 1
From here, we will rigorously explain the result derived from Lemma 2 When h 1, because
(l) < Br,
V Br 3DKI (1 (|peet).
By taking a telescoping sum of 9
K
V
<BF
1:
ERuNs 1g + BDI (f - [[((K + 1) + DKt (2+) (
k=1
B2.k
+
3(k 1) 2,3'
k=1
K
B&K 1: (Egoo lg (0)) + DKL DDKII.9' (Ipett)))) +
<Bp
(10)
B
k=1
where we used 8¹ 2 in the last inequality. On the other hand, we see that
k
BK(K + 1) 23' (K + 1) DKL DKL(Gopelll Ilpee)
V 2 - -Egors
Σ (B)
(k)
2
k
=- =-Easp
(k) + B.D. (Gope ((ror)))
BB'(K + D)T. (dopt Plett)
(11)
21

---

Published as a conference paper at ICLR 2025
because (Eq.)
V is the maximal value. Combining the upper bound (Eq. 10 and the lower bound
K h (Bgus) (0 Eggr [0 (Q PPett) DKL (Pope
(loeeelllPPe)t))) IPret)))
BF + (K( + 1)DRL (aopellll (Pret) + B3K K
8
Finally we get the convergence rate:
2
[( + (T - F(Copt) - 3DL (oops (leret)
K(K + 1)
2
Σ
le
22TTveVV + Bgcwlll Rgops (x)))] +((DLL (G DKL (lPrett)
K(K + 1)
k=1
2
BI. K
<2LTVEED
Br (KK - D)D (Pret) +
K(K +1
8
This concludes the assertion.
A.1.1 PROOF OF LEMMA
Proof. We can calculate the relation between Vi+1 V++1(G(h+1)) and Va (g(ght+1) as
h
( 1)(kk + 233)
V++
Ex((41)
DKI. FKKT(Qthe (g (|Pret)
2
6-1
Σ أ96و) 1=1
Bhik -1 1)
DKL(a (Q - B8 3DK (K (Q(h+11)
--Rg(e+)
-kEg(t)) (x) - (Bk + D'T, (K))D((( (G(h+2) (|Pret)
()) k:ggu [g [x) ((Eg(s)) (k( 1g (3k 8') D') (+(t1))))))
Prct) (12)
where we used the definitions of V& and Ve+1- Next, we upper bound the RHS of Eq, 12 using Lemma abbout the convexity of DKL:
(RHS of Eq. (12)
<V
Bk(k- 1) 2
1|r (++(tt --(()) (1) (pea) kBg (x)
1++ + k((Egeo) (x( Egu+ 1) [g (x)
(9k- + B¹) Dx) (k+ (t. the second equation in Lemma I and the optimality of V%²
<VE
Bk(kk 1) + 2
11r
(k+1)
- -rt (k) 1 (1( - kBggt (x) + k Eges (4) Egee+1) +1 (K))
4
3k 1)DKL (((())))|ppedt))
<VZ Bk(k 1) 2.8'
|| (++t)))) , 11)) (rea)) - kBgges (x) + Brokllr - ( (p(rt)
3(k DKL((p (2(2+1)1199
(:: (1) S Br)
BZ-K
<VZ - -EEgIi (1) B(k (g++ (Iptet)
3(k - 1) 2
where we used the arithmetic geometri mean inequality the last inequality. This concludes the proof.
n
22

---

Published as a conference paper at ICLR 2025
A.2 CONVERGENCE PROOF FOR NON-CONVEX LOSS F
Here, we give the proof of Theorem 2 and Corollary [ In this section, we always assume Assump- tion hools.... We are mainly interested the property that F smooth with respect to KL-divergence instead of convexity:
(ii' (A weaker version of (ii)). There exists SF such that
8g /o ( )d(a - d') (xl(glla') for any q. P. 2
8F
SF
F(g) S F(d'++
(13)
This property can be derived from (ii) in Assumption 2 that is Lipschitz continuity of oF. bq
Lemma 3. Assume that of is Lipshitz continuous with respect to the TV distance: There exists AF 8g
LTv > O such that
(q)
&F
(o) I LTvTV(g.d)) for any q, E E.P. Then, 8q
6F
(d')d(a- - (0' LTV rV,,''/ for any q. EP. by
F(4) S F(d')
2
From Pinsker's inequality,
oF
F(1) S F(d'+
(d))d(a - d') + Lov (allalla') for Jny q.d EPP 0g
Proof. The proof is almost identical the proof of Lyy-smoothness commonly discussed in the context of standard optimization. For o' E P. we define q: d' + ((g - g') E P. Then,
8F (1')dla -d)
F(1) F(d)
89
1=1
5F
(9 )qq d')dt -f
5F /
(4)) 9)) 5g
ōq
18
oF (d) d(q dt oq
(qt)
Eq
SF 5q
''' (LTv-Lipshitz continuity of
LTV TV(..g')²
2
while we used the fundamental theorem of calculus in the first equation and the assumption in the inequality.
We emphasize that when the inner-loop error ignored, it is possible to prove convergence using only the smoothness (13) instead of (ii) in Assumption Therefore, we will use the notation Sp for the parts that can be derived using Assumption (ii)' instead of Assumption (ii).
Under the assumptions, our goal is to show that the functional derivative of L(g) = F(g) +
(gll goes to a constant:
8T.
-
(up to constant w.r.t. x)
5q
We prepare the following Lemma about the convexity of DKL:
Lemma 4. The following equations hold:
(i) DRL(allP) = DKI(g) (g'llp) + á Dxa (lp)d(g (g g¹) + DKL
(ii) Letr Rd R be smooth fmooth, function. For F(g) Ey(r(x)] Dku(allp) o' = =exp(-r)g, it holds that
F(g) F(g') + Dka.(glla').
F(g) is uniquely minimized 1==-/ q'
23

---

Published as a conference paper at ICLR 2025
We omit its proof because it is straight forward. Lemma implies that DKL also plays important role as "quadratic" penalty term of the proximal operator whose output is
argmin {H + DKL (gllee))))))) for any functional H. qEP
This property allows for the use of standard nonconvex convergence analysis of Dual Averaging based on the proximal operator (Liu et al. 2023a).
Roughly speaking, the intermediate goal is to show that L(g) monotonically decreases in each kth iteration:
L(q L(g S 0.
By the weaker smoothness of F (Eq.433)
13),
the left hand side is approximately bounded
8F
L(q(k+1)) L()()))
- (()()))
(g
óq
/ 5q
[6
(d))))((((t)))) - -(()) + DKL (
(14)
ignoring various minor terms and constants. To bound the right hand side in 14 we show the following inequality regarding the Option using Lemma
Lemma 5. Assume that (V (x())) S crv/2 2 for all k. Then, the Option achieves
8L (( g(k1)
Eq
<BpeTV (DKL(G"))Prat) (KL( ((ret) &DKL
8'
(I + 1) (Q(h) ( (Q(K+D))
Proof. In the Option 2, (k++ minimizes
/
8L j (10))d(a 5q
Tk+1 (4) = Σ
(00 B': + 1)DRL (allProt).
We interpret that Option 2 computes the proximal operator of the weighted sum of 5L concept is justified by Lemma By the definition of rı and r1+11 it holds that by and this
8T. (g )dgg g(h)) + 8 DKL (gl(o))
Tk+1 (4) rx(
(15)
5g
From Lemma we also have
rx(g) - nu (k) = - (a) for all E
(16)
because (4)(4) is just the sum of the linear functional of and the KI-divergence ignoring the constant. Letting = gl++1) and we obtain
0 (g - (k) - 3'K (f
(17)
=T - hi39
8L
(1 ( - -(())) - 3' DKI DKI (k(t)) I|prot)
())( &'k (KLK(((h+t))) (1
In the last equality, we decomposed the sum using the equation 15 Thus,
8L (d) (()))
l 5q
(rk+i(g(h+1)) Tx ())))) - 3 DKL DKL(K (kL(@"h (g (ppeet) - B'DD KL ( (g(h+))11)())))
By using the same argument Eq. (66] for & K + 1, we have
Tk+ ( (g(++1)) 38 (k: + DKL (1) (1
(1)))
- B'D. (pPoot))
Substituting this relation to Eq. 117 and noticing
6F 80
Br, obtain the assertion.
24

---

Published as a conference paper at ICLR 2025
We also show the result for Option that is similar to Lemma 5 Lemma 6, Assume that (V ( ((()) S 2Tv/2f0 for all k. Then, Algorithm I achieves
oF (a()))((((kt)) -(th) g(k))
< <Brrrv + 5k +
(DKL( - DKL (g (g(c+t1)/PPret)
k
B(k 1) - 28'
) DRL (KL(t(HDD11((t)))
B(k + +) ((
DKL. ((th))
Proof. Please refer to Appendix A.2.1
The proof is almost identical to the proof of Lamma 5
We mentioned that our rough intermediate goal was show that L(g) monotonically decreases in each kth iteration:
L(q(k-1)) L(q(k)) So.
Rigorously, we will show that
(g) = I(g) + 3'
DKL DKL (gllptet)
decreases. This is an objective function with additional regularization imposed by the hyperparam- eter B¹. From Lemma 5 we prove that Lx (q)))) decreases as in the following lemma.
Lemma 7. Assume that (V (*)) S erv/2 for all & and 28 + SF S 23'. Then, Option satisfies
LK+1 - (n( S(LTV BFJerv 3(k+ 1)
DKL (0 (1((+11))
(18)
and Option I satisfies
( - (n( S(LTV + BF)erV Bk(k 1) + 2B +
DKL (q (1(1)1(((h+))
2k
(19)
By this lemma, it can be shown that the sum the KL divergences DKL (g(0))11)((+1)) converges to O at rate of O(1/K) through a telescoping sum: we obtain that
(wigght); DEL (
((th)
K
K
From the above equation, it is immediately shown that
(
min I=I..... K
(
0
where (weight), ~ k in Option 1. while (weight) ~ 1 in Option 2. In Option 2, will be able to show that min; DKL (4 [|ğ ([)))=))) 0, which implies that g(a)) converges to some point (in fact, this is the stationary point).
Now let us prove Lemma
Proof. We only prove the inequality for Option 2. By the smoothness of F (the weaker smoothness (i)' Eq.13 the Lipschitz continuity of éF dg and the property of the KL-divergence (Lemma
we have
5F ōq (a (((gktt)))) - (()) + Sf DKL (gk(g"tt (1(
(( ((g <
2
+ 3DKL (Q (1())))1(() + / (
g(())
1
8F
(1)))) - d(glkt)) ((t)
8y
SF
D DK.. (1 + 3Dr.. (f + LTVeTV. (20)
+ 2
25

---

Published as a conference paper at ICLR 2025
We have used the induced smoothness of F (Eq. (3) in the first inequality, and used the Lipschitz continuity in the second inequality. By Lemma [ we can bound
of (a(*)) - ()( d(q(k+))) f ogg (d())) Q())) Thus the RHS of Eq. (20) bg can be bounded as 3'
(RHS) - (DKLG (Pret) DKL (gl DDLL (G (1((t)) (k 1) (g (1(1)1((((+1
23 SF PKL(G(+DD/1G(C)) + (LTV BF)ETV
8¹ (KL(G (9 (pret) DKL (Qth(t) (k+ + 1) (G
(23 SF -28") DKI. (k (art((a-11000 (1a + (LTV + BF)eTV,
2
which gives the assertion in Option by noting E DKL (P(tt//)))))) 2 E+I DEL (ppet)) As for Option 1, we repeat the same argument to show the desired result.
Combining Lemma 5 and Lemma
we will prove that
K
DkI. (1) =
1 K
K
k=1
where (weight) ~ & in Option and (weight)s in Option 2. When take Option 2, it holds that ming DRL (q (g(k+)))) O(1KK). It is important that DRL (g(h) (|a()+))) is also the approximate "moment generation function" tig (g) = log (e,p(( + Ey]gl]) of (g when 87
k + X0,
min K=I..K
"Variance" of 5q
8L.x
(Q(h), min
DKL (g (6 = 0 (K) K
k=
This can be interpreted as
éLc (g 22 + 0 (up constant w.r.t. x).
5g
This result is consistent with the findings for standard, non-distributional Dual Averaging (Liu al., 2023a), where mi.... (graadiet of the objective),.l' O(1/K) as the variance corresponds the second moment..
We rigorously formulate the above discussion as the following proposition. Proposition 1. Let
1 (Li( - -L)++ (LTv + BF)
K
Σ eTV.
K
Kg
KB'
Then, Algorithm I satisfies
K
8k + 2,3'
DKL(G DkL(g(W)lg(tt11) < Vx
K k=
and Algorithm 2 satisfies
K DKL( S Mx
1=
This also yields that the following bound holds for Algorithm 2:
3 + 1) dq (a(t))
k (gges 3'kks + 1) dq
oi (g(h))
min ISKSK
< 1 K,
where the expectation the middle term is taken over a random index lo uniformly chosen from [K] (1,....K K).
26

---

Published as a conference paper at ICLR 2025
Proof. Summing up 19 for lo 1..... K, we have that
K
K
L - (2 S(LrV + BF) Σ crV - 5
k
+
(KL(G/1)
h
This yields that
B
k ( + K ((xx(((((t)))
K k=1
$
1 (L(( - LD)
(ITV + Br)
Σ
K
erV.
(21)
k
K
k=1
Please note that
DKL (g
8'
k +
1
DKL
(1 + 1 (kk - L7)
k
by the optimality of L and that the RHS in 21 is O(1/K). Here, we see that
(KL(( (1)(())((((+11)
=Bguo| 1- + g [h0]] loggBBner (xpppgghee 1))) + logg (19g(11pna (exp(-g(*)
-Egrs
+ ( log( (exp( (b()) + 1(
=log [ggcs cxp (g()) + ((h Egges [-g + +)(
Now, the term g(k) -ggh-1) can be evaluated as
- (kt-1)
k
8L (g((0)
1 xb--1)
k
(1))
30 1) ōg
1:
3 + 1) 8q
which yields the assertion because
DKL (gl(g(A) (((httt)
k
8Lx (b)) BBua)
k
8L, (a (3(00)
1)}
=log
E Egus
exp
Bl(k 1) dq
Bikk + 1)
8g
which is the "moment generating function" wha of ((++1) oq
1 (Q())).
A.2.1 PROOF OF LEMMA 6
Proof. The proof is almost identical to Lemma 5 In Algorithm 1, (((++1) is defined the minimizer of the following quantity:
8F ( ōq
Tk+ (g)
Ei j=1
(d g(0) (gl((alrroo))
1) (allllllooo
Hence, we have
Tk+1(g)= = =x(()) + +/ /
8F - ( (8k + B¹) DKI(allPeer)
(22)
by the definition of rg- Moreover, Lemma 4 and the optimality (()) gives that
3k(k -1) + B8kk ()ll for all g E P.
(g) - (x))
(23)
2
27

---

Published as a conference paper at ICLR 2025
Substituting -- 9 (( obtain
0 (r=(q(h+1) rg(() Bki(k: ( (
+ (38%)
DKL (It(ghtDll
2
8F
-Tk-1(((kt2))
(1 - -Q(h)) (3k + BB DKI. (q. (1 Prot) (used 22
(1
8q
- -ra((((0)-
1 18%)
[3k(k -1)
DEC P
(24)
2
This is equivalent
I
(g (d(((e+)) ((A))
5F
bq
Sri+1
((())) (g
(3k + DKL
[5k(k-1] -1)
[gk(k:
3k) DKL DKL (
2
By using the same argument as Eq.
23 for l t k + 1, we have
Tk+1 Tk++( (Q(h+ +
Bk k + +1)
+ 8' (k+++ 11) 1))
(KL(
2
(g(+)(
(*)) (3k D (1
Substituting this relation to Eq. (2 and noticing
6F by
S BF, we obtain the assertion.
B DERIVATIONS OF FUNCTIONAL. DERIVATIVES
B. DIRECT PREFERENCE OPTIMIZATION
We can apply the proposed algorithm the (populational) objective of Direct Preference Optimiza- tion (DPO) (Rafailov et al. 2023). DPO is effective approach for learning from human preference for not only language models but also diffusion models.
Original DPO objective Let Two, be "winning" and "losing" outputs independently sampled from the reference model Pref- The event {ae > at} determined by the human preference. The original DPO objective is formulated as
LDPO.original(9 -EE log (7log ((F.rc) y log g(x1)
Prof(Iw.
Pref pref(zi)
where ? > 0 is a hyperparameter. The expectation taken by zw: 2y. that are "winning" and "losing" samples from Prof- Wallace et al. (2024) derived new objective that is the upper bound of LDPO, (q), but is specialized derivation of the optimization method for DPO.
Reformulating of the DPO Objective The goal is to directly minimize LDPO, original (g), however, in the above expression, we cannot apply DPO to diffusion models directly because the expectaion with tuple (Iw, 22) is not formulated well. We start with another expression of the objective of DPO:
LDPO(g) = -EE -Eewrpest E Exyt-peat log
(7
(/log
Pret y log (21)
((w))
q(x1)
1
(Iw 22))
Layy ( w) is if > otherwise. This LDPO is in the regime of our algorithm and the functional derivative can be derived:
Proposition 2. The functional derivative of LDPo(g) is calculated as ÓLDPO (4,x:
5g
e dpret
=-7Bmmpie
(1 (1-o (->f(x) +zf(z))))
1 aara mn (x, (i)
(((())
(1 (1-o(-a((rm)) +x)))) J 5
+ +BBuuarpid
e-f(x) Jaura ( (Rins
where q = e pros/ a @-- This functional derivative traciable in seftings.
28

---

Published as a conference paper at ICLR 2025
Proof. In this proof, use the following notations:
Prof the output distribution of a pre-trained model,
q == e= f e -1 dprot: the output distribution of aligned model,
LSL(g), Q2) = log (7 log Q1 /Peer log 92 Pret))
L,LSL( 92) - 7(1 (7 log 91 /Peet - log 92 Prof))
0,LSL(91 -(1 o(7 log 91 Pret - 7 log 42 Pret 1 42'
((1,,22)= = 7llog - log r2),
ddier
Inv(f,x) e-1(0)
The objective of DPO is written
(D Ex Empxpeer [LSL(g(Zw), (p( (Eru 21)l
We obtain the first variation of the objective as follows:
LNpo(a elğ--0 ())
-Lpoo(g) - e Barrner ((0LSL((((w) (=w)
+8.LSL(g(Ew)) ((=)))((-))))))))))))))))))))))) 22(2)) 1 Lourn (Fw, + O
=Lpro(p)
Exprepint [f (115b
olvi g(sww)) ((x)) )) (Tw) (f. Ire d zn 1
(Del (21)
( 7(1 o ((sm)) a(x1) ())))))))))))))))))))))) (ar)) m1)de (odddllluuaa 1nwras
(rot(Fw)) (1))(
+ +
Then, the first derivative of
őLDPO (p.x) Exieprnt
- vi ((x) g(xi))
Prer(x) pror(11) 1) 1 (z, mn) Inv(f,x)
öp
+ Eee Pppo (111 o vi( g(am)) g(x)
Pret(Ew) paatt))) dewro(dw::)) 2)) Inv(f,x).
B.2 KAHNEMAN-TYERE OPTIMIZATION
Assume that the whole data space Rd is split into desirable domain Dp and an undesirable domain Du. The objective of original KTT (Ethayarajh et al. 2024) is formulated as
LKTO(g) =Ex~prn yD (1-o [:log g.
DKL(ellpret) ((NED))
Pref
+ru 1 11-0 (DLL(llPeat) klog 2
) 1 ((neRu))
where 7D, yU, are hyper parameters, and is sigmoid function.
29

---

Published as a conference paper at ICLR 2025
Proposition 3.
The functional derivative of LKTO is calculated as åLTTO (q, x)
e dpref f(m)
(e(pdderv(((x)))
(wEDD)
+ (-f(x)- log
f(a) dprct) P [nDoderiv (((v)) 1 (vEDa))
+ NNTOderiv (- /
e '
e-1(x) 1l (mEDD))
- (-f((x) - log / e ) [dooderiv (- 1 (vEDe)]
where deriv ol)) o.))) x(xx) = log a(w)
DKL Preet), Je- ddeect
o pppp
P=a(()
The functional derivative of KTO can be caleulated if you have f(x) and the samples from Pref- Note that log ((z)/Prer(2)) = - f(x) log f -- dpret.
Proof. The first variation of LKTO is
LKTOla + elğ - q)) (ạ)
myDoderiv ((x) ((x)- g(x)
((x) + If log
A(v) d(ğ ))()) [
Prel(y)
+ eBxuppoe Nyvoderv ((x)) g(z) - f log g(v) d(ā (00) LivEDo
q(x)
(ret(y)
/
+ log g(x) E (NTDaddorv( ( } dd(7-))0
d(j
4)())
Prot(w))
+e / (-
llog g(x) E. (rrpDoderrvv-<(y)) Pref(2))
(oEPn)] ( } a(i - a())
dpee [00)
where Inv(f
Now the desired result immediately follows.
C ERROR ANALYSIS OF DIFFUSION MODEL
C.I OVERVIEW
Here, analyze the sampling error caused by the diffusion model. Let organize the settings and notations used in this section.
Target distribution. The target distribution is g. go- which is decomposed q- (x) pa (x)p. (x) with p. and p.. Here p. and represent the distribution of the original model and the density ratio obtained by fine-tuning, respectively.
Sampling with score-based diffusion model. In the score-based diffusion model, we start with the forward process, which is written as stochastic differential equation (SDE). Choosing the Ornstein-Uhlenbae (OU) process, {Xt follows the following SDE:
Xo ~ Pe, dX, X,dl + V2dBr.
(25)
where (B.h20 is the standard Brownian motion. At each time I, the law of X, is written as
p=(x)= / p. exp (
1
20 |m x1|3 dr,
30

---

Published as a conference paper at ICLR 2025
where my c- and of e 2t, In the same way, define (Yi)20 by replacing p. by q, in 25 and let dj be the law of Y,
Then, for some 2 0, we can define the reverse process for {X; JUSegT. (Let Bg below be distinct from the one in 25
X
dx; = 1x;- + 2V 2V )dd V2dB.
(26)
Then, the law of X equals pr-t. which is why we call 26 as the reverse process. the same way, we define (Y,- logier the reverse process of (Y,
Doob's h-transform. By applying Doob's h-transform to Vlog 4r- it can be decomposed into the original score V VlogPr- and correction term.
(V, (logpr-((V;)) + V log(B 8x = x)) x))l-
(27)
See Lemma 12 for derivation.
Approximation of the score and correction term. Because we do not have access to the exact value of pt and p. and therefore cannot implement (27 exactly, we consider approximating them by, e.g..
neural networks. We approximate V log (x) by s(x, (): IRd+1 + Rd. Also, we approximate
(z, V (Xo)) = z|) by u(x,t) t): Id+ + =².
Discretization. Also, we need to discretize the stochastic differential equation. We finally obtain the approximation of (26), denoted by {y,+ JoSEST, as
Yo NO Ta), dy," (Y; + 2nh kh) 2u(Vin kh)]dt + V2dB, t E [kh,, (k + 1)h).
Obtaining the correction term (approximately). Given the score network s(x, t) approximating V logg pogge(e (x) and the function h that approximates h., we can approximate the correction term u(æ, 4).
Remind that, for fixed x E Rd 1 = kh, 8 k(h + +1 E R, the correction term is calculated as
EEp (Xo) Nr I'I P[X XT_ 2' = x)dr'
V (p (2.) xi])
(Xp) NT
(28)
One way to approximate diffusion process
28 starts from approximating (X,) Xx- = x². we run the reverse
X- = dx; (X; 2V 1X, )dr + V2dB,,
obtain that the law of XF is equal that of Xolx v. Therefore, by running the approxi- mated reverse process (with a slight abuse of notation)
X, = v, dX, = (&, + 23( (i,, dd + VddB, E [1h, (1 1)h1,
multiple times, the sample of 7r, denoted by {ir, 1 can approximate Elp, (Xo) X = x]
(xo)
(TT,))
(29)
n
where p' is the approximation of p...
On the other hand, we approximate 3- P = x11T-- x) by approximating P P x'I X æ] with Gaussian distribution. Specifically, because Xx-. 2 = x obtained by the following reverse diffusion process
X = dx= = (X= + 2V log log (X: VdB
approximate V logpr-- by s(Ithr kh) to obtain
x Ikhi dxx= (X= 2s(xun, kh))drr + vddB
The distribution of X is denoted by Ne + 2(e 1)s(Iun,
(30)
31

---

Published as a conference paper at ICLR 2025
Using this, approximation is
0 P|Xr-= x xx = x²)
dr
(x' - (e x + 2((e - 1)s(ZRM, kh)])2
2
Or (2T (e2h - 11)
exp
2(c2h -1)
c² ( + 2(c²t ()s(IKA, kh)))
(x' (e 2(e 1)s((1 kh)))² 2(E2h - 1)
exp
(e2h - 1)( (e2h - 1))
This implies that, if we sample { from (30
/ (Xo] = = x'1
0
P(XT_s = xT = x|da'
0x
1
Blpa [x - (c"x: + 2(c" 1)s(21, t))),
= =2]((ar
(e2h-1) m
and approximate each Elp. (Xo) x')l in the same way as 29 we can approximate the correction term (28).
Now, our goal is bound the error of the whole pipeline under the following assumptions of p. and P* and the score approximation eπror.
Assumption 5 (Assumption restated). 1. V Po is Ly-smooth at every time t and it has finite X S m < oo for all E RR and x E Rdd.
second moment E
2. V log Pa is Lm-smooth and bounded as Cp Ss S for constant Cp-
logpr-((85 S at each 3. The score estimation error is bounded by Ex- [|(=(X[ time t.
4. Ep. [ll,, 1h) - u(x, [h)]l?] S Ep., for any 1 S 1 T/h.
Theorem (Theorem 3 restated). Suppose that Assumption is satisfied. Then, we have the fol- lowing bound the distribution q of Y
T/h
TV 2) S T
hea TLL, + 2 + m²h²) + exp(-2T) DKL (0, I)).
Proof. Suppose that [[s(,,) + uu t) log fl--1 v S and log - LygLsyLsooon smooth every
time t. According to Chen et al. (023c) and Pinsker's 's inequality, the distribution of Yr satisfies
T/h
(L(g)l)=²² S T +
he a. TDI m²h²)= +
(0,
According to Lemma
Ly is bounded by
Leg S LOCO + Lor
which yields the assertion.
In the bound, we assumed that the term e"_, can be bounded, however this approximation error can be derived as in the following theorem with additional technical conditions.
Assumption 6 (Assumption restated).. (i) (.s(,-) -) is Hg-Lipschitz continuous in a sense that
V 1) - V ,,s(y, l)pp < H. wll for any x; y E IRd and 0 < L < T and [[][= (8 kh) [|2] < Q² for any k. (ti) There exists R > 0 such that sup,... V log P. (z) llop: V log s(x, t)l } S R.
Theorem (Theorem (restatedd, Suppose that Assumptions and 4 hold. Assume that ollo0 z e', lle]lco < Cp, and sup, /l (x)]| < Rp. [V V (v)l -
L vll (Vz, v). Then, for
any choice of < h < 8 < 1/ + 2R), we have that
Eg <C3 (5a + R$. (e² + 12 m²²)) [T2 (m m + Q² + dh) + R2 (1 + 22)
+ min(T 1h, 1/(2 2R) -22
where Exy- := 01 introduced in Lemma 10
C2 1+6R IE2 + Ljd(h + m/ and RG and Ly are constants
32

---

Published as a conference paper at ICLR 2025
C.2 BOUNDING THE SMOOTHNESS
The proof of Theorem 3 (i.e., Theorem 5) utilizes the smoothness of the density corresponding to the aligned model. The following lemma gives its bound.
Lemma 8. Suppose that V log Pt is Lpp-Lipschit, for all t and V log pa is L--Lipschitz, and Cp" - pa S Cp. Then, V log qt is Lg-Lipschitz,fe all where Ly bounded by
4(Lp + Lp)² 2 + pp) - (4 (4+ Cp)lpp + 4hp S Lp Lop
big S min
Proof. We divide the proof into two parts, with of = Tbpt as the boundary.
First consider the case when of S x(pp Remind that
(c(x) = f 9- (v(()) (2TO
exp 207
m.y x dy.
with of -
-2t and m7 =e From this, V (xx( and V. are computed as
1
V (xlo(x)= v go (v) (2TO?))
exp (
|m
==13) 12 dy
20
m / (V)))(())
1
exp (
Immy xe| 2) Pdy
(2TT?)
20
and
(z((z) = m, (V (2TO?)
exp (
1
m.y - x xdy
20?
Thus, can compute V (x) as
2 x
(x)
Vel(x)
(e(z)
(q:(T))2
m [(
(z9((v)
VVV. (v) (1) exp
Ima x1l
(-(
207
(31)
(x) exp 2 [ma - -x|1=) x1|²)dy
mi² 2 ʃ V. (v) (-(y) exp (
1 |m - x1l
(1)²
20?
(32)
(v) exp
|ma xx12)dy
20
mi 2 f Fue-ly) ao (y) exp
1 mmy -x|1²) =zl²)dy l Vvg.(v)
q- (1) exp ( -/77 |ma x dy
2- (x)
207
9-(1)
(ʃ qx (v) exp (-27 प |mm - x || dy)2
(33)
Eq. 31 is expectation of V log g. (v) vzg.(z) 9-(3)
V.g. (v)V.gg (z) (b))
with respect to distribution
A(y|x) & (v) exp ( -- 20?
m.y x=13).
is bounded by m; (Lp + Lp) from the assumption. On the other hand, the other two terms are regarded as the covariance of V log qu (v(v) with respect to that distribution. Because of < L,, A x) Lp))strongly concave, and therefore satisfies the Poincaré inequality with
Therefore, 31
constant I,+ Therefore, for any E Rd we have
(3
= m, 2 (E (E log q. (V) log q. (1)) EAcy(a) V log 9u 9 (v)] log q. (v) )a
< m,
m,
(Lp + Lo)² m, Lp + Lp
Lp + Lp
Lp Lp
33

---

Published as a conference paper at ICLR 2025
This implies that 32 133 is m, (Lp Lp) smooth and V log q((z) is 2m, 2 + Lp) smooth.
or we have m, V 2(,+ By applying this, we have 2m, -2 (Lp +
Because m,
Ipp) 3(4p+b,)-1
Next, let us consider the case when of 2 (Lp+hp)
2 Note that (=)((x) and V2 are also written
as
Va(((x)== = a-(1))
1
exp (
may x1(2) I may dy
(2T07)
20
of
and
2 (x)
= =-f -f -faa()) ga((1)
1
I
exp (
may x1(2)
dy
(2TT7)
20,
o;
S (/(())) (2T07)
1
(x
m.y)( mma) dy.
mma 21(2)
exp
20
Thus, we can compute V log 4t (x) as
2 log (x²)
Vi((x)
(V (9)(x))2
q((x)
I (v) exp (- 2077
|la - x: x=|²2) ddy
(34)
(=())) exp
|,a x 11)cdl
20
I (v) exp (
1
mm - x=| (a m, m. dy
207
(35)
S (q (y) exp
may x x112)dy
(U) (
1 mm xz|(2) (a (-- ( (ʃ qo (x) exp (-
1
|ma
xx13)
(a (r- m.y) dy)
20?
of
20
of
(J 9-(1) exp (- 20? |m. x (2)
(36)
Eq. 34
is bounded by oe 2(,, + Lp). On the other hand, 33 + +5) are transformed into
35
36
m J (v) exp ( 1
207 |ma x=|(2 () yy dy
of f q. (v) exp
mma - x1| dy
207
m (ʃ (v) exp
1 may x xdy)) (ʃ (-())) exp ( (- 207
|m x ydy))dyy)
207
of
(J q- (y) exp ( 207 L
|ma - =112)dy)²
mi Var aiy (2))
of
where Var m) (y) means the variance of X, conditioned on X, x, with respect to gr-
Thus, bounding Var (x)())) yields the conclusion. Lemma implies that
(ul))) < C² Car (v).
(37)
Similarly to 2 log (x). 2 lzg satisfies
m
2 (x) -0, 2 2
Var
po (ll (V).
(38)
o
34

---

Published as a conference paper at ICLR 2025
By combining (37 and 38
we have
2
Var (ula) (v) - I Vog
of
Therefore, from the assumption that V log Lp Lipschihi and of 2 2(Ipp + Ip),, obtain
that || mi Var (u|z) (x)l)) (v) - (2 Cp) p
By putting it all together, V log q: is ((4 + Lg + 4I, -Lipshhisshiizz
Lemma 9, When Cp I (x) S Cp. have
Poje(x|(') S 0
for all x, x and t.
Proof. We can write Zotz (z|2²)
g0,c(x,
Ps (2)))=())))))))))))))))
Zoje (x](==))))) I go (x" Jaqo (a " J p. (x") )r. (a²)) (ee"))oo (oll )
Because 7 S px(x) S Cp and (0)((11)))) (=jo(>"((xx) we have
pa (x)
(oje(x)(a)) S ch pa peoo C2po (zr22')
which concludes the proof.
C.3 ESTIMATION ERROR OF THE CORRECTION TERM
Since have shown the time discretization error in the previous section, what remain show is just an upper bound of e².1- For that purpose, we put an additional assumption which is almost same as Assumption 4 except the condition (iii). A bound of EÊy in the third condition (iii) will be given as 2 (e 3 h) by Chen et al. (202cc).
Assumption 7. (i) (..(::))
H,-Lipschitz continuous in sense that [ V Vys(y, O)pppp S H.Ilæ for any x, y E Rdd and 0 S LSS T and B[]ls(X kh) 2] S Q² for any k.
(ii) There exists R > 0 such that sup,,, 11V log Pt (1)ooop V log s(x, t)pp < R
(iti) Ex, NIV(XFXXFIXX
X = 8,) < stv for any E [0, T7.
Theorem 7. Suppose that 0 h 5 1/(1+ 2R) and Assumptions 7 and 5
hold. Let Lip and Rip
be as given in Lemmal 10 Then, it holds that
Æx- VI pp. [x.] V.Elo- (X- [
SR
(s² + 1 12 + m²²1) + E + [1 (m + 4Q² + dh) + 12 +2
2
=0 + 8 +
FTV
8
where
E C1 5 22xp(2)
c2 (2 .π. + CRR 12d(6+m66)1)
and Cy > 0 is universal constant.
By Chen et al. (023c). stv Os + h), and thus by substituting 5 +- Vh, finally obtain an emor estimate as
(1 (1+ Vh
2 + Vh.
35

---

Published as a conference paper at ICLR 2025
Proof. Let [kh, (k: + 1) and t* kh + where å is larger than equal h: 8 2 h. We only consider a situation where T - t 2 The situation where < - can be treated in the same manner by noticing trivial relation (xF [X x (p.(x). Then, for given initial state E R4, define the stochastic processes
X
= + 2V 2V Logpr- )d v2dB,, i
and its numerical approximation as
= =x; + 2s(X 2 k,))ddr+ + V2dB, TT)
where H is the integer such that T E 1 (kr + 1)h).
Note that
E P. (XF) V
V Ep.
(6)
(39)
We first evaluate the term (a):
V=EI) = - (XF) = =]]
As we have seen above, the derivative can be expressed by the following recursive formula of the conditional expectation:
V.E
X; z].
For notation simplicity, we let (x)
p-
XX and py (Xp
zt). Then, the Aysmut-Elworthy formula [Bismut v E RRd.
1984
Elworthy & Li 1994) yields that, for any
1
V Ex X = x] =
(nx (B.)ox(XF) x; = x
8
where 78,r is the solution of
dnx,- (I +
1x,0
Similarly, we define 7x,r for the process X as
dnx,= (1x,r + 2V 8 (X ky h
1x,0
Then,
V V. (XF) = - - (XF
1
(nx,- - -xx,, 8 Xτ
1
+ E
(Bx) py (x:)) X X
(40)
36

---

Published as a conference paper at ICLR 2025
By the Ito isometry, the first term of the right hand side can be bounded
1x = X; = 3t {
E 5 (nx,- - 7x,, -
2
/ (nx.- 11x, (B,))BBy
<B
|xt = Xx = a
3
I (nx 1x :dxr
<CAE
1xt = Xx = x
h || 1x.=112dT Xx Xt=
=CPR
=l( X X; x 1 (1118,7 -1lp + ||x,,
Hence, we just need to bound differential equation:
|7xr- -vll v/12
the right hand side. We note that obeys the following
d|l,x,r 011
dr
-2(1x,= 8) drxx,
dr
=2(18,= 0) TT + 2v log
-2(18, v) TI 2 (8 v) +
52(1 + 21)) |1x,, - -112 2(1 + 211 -lll
which also yields that
2 + 2R)11x,, dllu ull
2 2|,,,, - w/|l
- w/12 + 2(1 + 2R]][||llxx bll
dr
d|nt,r - w// 2 2R)(1188, w|l + (o/l |178,r 1/| (2xp(2(1 + 2R)T) 1|]|a||
dr
||88, u/|2 < [exp(2(1 + 2R)7) ||]1
Therefore, if is sufficiently small (such as 5 < 1/(1 + 2R)), then we arrive at
Inx,- u/1 < (n(1 /
with a universal constant Cy,. for any S < 8. the same vein, we also have
|/xx,r - v/2 < (1 + 2/)
1| exp(4(1 + 22)r) 1| z exp(4(1 + 2
for 0 S S 6. These bounds yield that
2C R
(178,- - v/12 1| wll3)dr Xt = X; =
J0
fo 72d
1c2,0/11 + 2R)2
C + 21)²
3 8.
82
37

---

Published as a conference paper at ICLR 2025
Next, we bound the second term of the right hand side in Eq. (40):
2
E
J
(xx.:
(B.)(xxxx)))
(or X = Xt =
6
S
|| 1 Xt Xt- E' - (xx(x )
Xx= Xt=
82
exp(2(1 + 22) [xxx(xx;)) - (X; )) X = X = a]
(41)
Otherwise, we also have the following inequality:
(nx, -xr X = =x;-- =
-
1
8
2
<2E
8 1, dB. py X X, x
1 ( (x FO x = xx-= = x 2
+ 2EE
exp(2(1 + 22]6][[u1
E [1p
py (x+))) ) X =x =
X, = xx)2 2E [1 (VVxx=(x)xXxxxxxx))))x))) (xx X
(42)
For bounding these quantities, need to bound the discrepancy || - X 1. quantity follows the following ODE:
Note that this
x X 1|
dr -2(X;- -
2V, - (X+ - 22 kh))] pr Pr (8,,) s(sr, kh))
- - -AX;=-((X;-- - x;) (V, loggr-
<4I18;- -x, - 2 log(pr
s2z, kh)l1
Therefore, it satisfies that
S L - X 1s +2 II (ox s(x, kh) dss
- Taking its expectation, we see that
/. - 1s + I' (6² + L + 1.rms²) )ds,
- 1 -
where we used Theorem 10 (and its proof) of Chen et al. 2023b) for obtaining {(T). Then, Gronwall inequality yields
E[|| - x s + f 4E(s)c S Ljd(7²² + mr²),
(see Mischler 2019) for example). Then, the Lipschitz continuity of by (Lemma 10 yields that E (upy (Xt=) - (r X = X = x] S LglET + Lid(1² + m)))
Bound for t kh: First, we show bound for t kh. The right hand side of Eq. 41 with ò h can be bounded by exp(2(1 + 22) x [Exxx - (r 8 = X]
8
exp(2(1+ +2)))
2Ex; [[((xxXX - py (x ))² + (pr - (Xx+)) X x
8
()p(2(1+2)))]ll)))
CRetiv + RLExx
LE [X x; XF = x1
8
38

---

Published as a conference paper at ICLR 2025
where we used
E [(px(X ox = <C [TV I xf
Xt= Xx-+))
SC
Here, by using Theorem 10 of Chen et al.
2023c again, the right hand side can be bounded as
2 exp(2 + ||
(C + CR + 12 + ml
5
C2 c c + CRRIEE + L,d(6 + mm mo8)]
=2 expp(2(1 + 2R))) |0]2
Then, with the contraint h S 1/(11 + 2R), it can be further simplified as
ty CRİle + 15d(0 mo²)]
2 exp(())|l||
Therefore, by taking maximum with respect to E He with constraint ||0| = 1,
Ex Ex- (1) V [p [Xx))xxx,,>> - (x |X; = ]1
4c C2 2R)2
+ 2 22xp(2) { c3 t
xo + +RR + Lga(6 + mo²)] = Es,e.
(43)
3
We see that E6,5 0(8 + =² + /0)
Bound for general t E (kh, (k + 1)h): In this setting, we utilize the inequality 422 constraint 5 < 1/(1 + 2)) and ||||| 1, the right hand side of 42 can be bounded by
Using the
exp(2)
E [(py (X (x x = X
= x]
8
+ 2E (exp(2) 2) V (x2x (X;))) - py (X,- M| X = X = =x]²
exp(2)
R² [X - X )² X = x = 2| + 2 eexp(2)E6, (:: Eq. (43
8
By taking the expectation with respect to x = X we arrive at Ex, 11V (Xx X V.Eip. | X |
<C exp(2) R2 (528 + 12 m) + exp (2)
<Cexp(2) R (s" 17 6 + m + 2 exp cxp (2)
(44)
This gives an upper bound of the term (a) in Eq. term (b) in Eq. 39
39
Then, we just need to bound the remaining
For that purpose, we define (v. = formula again,
= zc. Then, using the Bismut-Elworthy-Li
(VVrrx( ()yrRa(
(v))(())) - E - X x]
=Ele = (vV Vipr,( + (1) (X,,, - v) (X =x]
<E|Lp|Io - X; || + Ro - w/l X = x]
<EELLpI(((K 1)h- 1 (x - 25(2, kh)) + / (h BR R&p R.pl - w/l X = x).
which yields that
Ex- Ex-1IV Vpr.an (X 11
(kk - 112 + 4 (X)) + d((k= + 1)/ - t)
+ +R,c(1+2R))((( 1)h- - t)²
$ 2213 + dh) + (1
-121.g + +42 +h) RR,, + 112
(45)
Combining 44 and 445) gives the assertion.
n
39

---

Published as a conference paper at ICLR 2025
Lemma 10. Suppose that sup, [V < Rp. (V) Vp-(v)]l - Lp yll (Vx, y), and
V.s(:.) is H2 H--iischhitz continuous with respect to z. Let (y, = (X- = =x]. Then,
V (2) is bounded by Ri and Lipsschhi continuous for any 0 S t T, where
R max{C,2 2R),,
e22 21
222
1/2
L. max
2R + + 2R) exp(6)
+
1+2R
for a universal constant Cy > 0.
Proof. We show it only when t = kh for a positive integer & just for simplicity. The proof for general can obtained in the same manner.
(i) First, we assume that --t t 2 1/4(1 + 2R). In the following, we let E Rd be an arbitrary vector with ||e|| = 1. We again utilize the Bismut- Lsnutt-lworthy-Li formula:
e (x))(x
V VxE(p.(XF)) x x]
5 / dB,) x =
Hence,
(v
(rr(())))
2 2 10 (wr.
SCRR
1 = x
1
CCR
|| Xx x
52
Here, we know that 11 < exp(4(1 + 22) || and thus
(u vx)x))x))))² < C 2
Hence, by taking S = 4(1+2R) we have that
1
(rr) C2
This shows that || (x)]l is bounded by R. 2 + 2R)e.
Next, we show its Lipschitz continuity. For that purpose, we define two stochastic processes
X = x, dx = (X; + 28 (5 kh) + V2dB, (r E [kh, A(h 1))), z y džt =Z,- + 2s kh) + V2dB, (T [kh, 1)]),
where x, y E IRd with || y/1 < E. Accordingly, we also define
1/.00 v, dnx,F dr
(I+2V x X kk))xx...
12.0 = dnz,r
(1 + 2V s 2 kh))
dr
Thus,
Xi+1IM which yields
Xthh 2th + N| 2t. 2(s kh) s(Zth, khi)),
|| ||xt+t1jA
2 T
5(1 + h11 + R 22a1l
<(1 h(1 + R v/-
40

---

Published as a conference paper at ICLR 2025
Now, assume & Σ S/h that have (K+1) Z ik+1jall & exp(S(1 R))][= yll for k 1..... S/h. Hence,
d(nx,r 1]Z,r)
dr
- 117,-) + 2V 6( kI)) - 2V s( 2Z kt)))zzs
- T12,,) + 2(V 8 (Khh) kh) - V 8 2khh MX)) - 2V a (2 kh) - nxx,);
which also yields that
d 7x.F 112.-
dr
-2 11, + 4(1xxx 112,-) [VV ) s(X kh) V s s(Z kh))
- 4(1xxr - -12,7) s( kh) (TZ,r - '7x,-)
<2 - 12 + 41|11x.7 - 12 H1 exp(S(( + ex))2x + 2R)S)
+ 4R|| 112,,11"
Therefore,
d 1x,r 72.7
dr
<1 + 2R) [|]xxr- - 72, + 2H, exp(S(1 + R)) exp(2(1+ +
and thus by noticing ||11,0 - 1/2,0|| = 0, we have
||x,- 72.- & (exp(S(11+ 015511+2000 (2()) -1] 2H, exp(S(1 + R)) exp + 2R)5)
E,
1 + 2R
for any + & S. Then, by setting 5 = 1/((1 + 2R), the right hand side can be rewritten as
H,
||x, 7/,,~1 Snn 1 2R
for an absolute constant Cyp Therefore, we arrive at
(u (VPr. Ver:)(u)))²
(B,Yr. - /' (nz,,, dB,
S
1
<2E
(1xx,r - dr [x)
S2
1
E 3((xrs(Xx3) sr (r ))²
2E
112 dr
$2
2
H2
<
02
C2 + exp(4(( + 2 R2 exp((2(1 + R))s
5
(1 + 2R)
Then, for the choice of S 1/(1 + +2). the right hand side can be bounded by
HC2
+ 2(1 + 2R) 2x) exp(6) R3 e²
1+ 2R
This
implies
that V.y.(())
is Lipschitz
continuous
with
constant
L
#2 2(1 + 2R) exp(6) R??)
1/2
1428
(ii) Next, assume that T 1. 5 1/4(1 + 2R). In this situation, may use the following relation:
u. TV - - (Vp. X;- x].
41

---

Published as a conference paper at ICLR 2025
And, tracing an analogous argument by replacing py,s with h, obtain the assertion with
R2 = Rp: lup = (2 + +,2),)) 1/2
Lemma 11. Ifll.. plloc S E then
e (XF = x]] - [x = = - /in -,, 1/(2 + 2R)
Proof. It can be proved by the Bismut-Elworthy-Li formula again. We omit the details.
Combining all inequalities, we arrive at (the formal version of) Theorem
Theorem 8 (Formal statement of Theorem 4. Assume that Assumptions 7 and 5 hold and the conditions in Lemma 10 are satisfied. and Ilp, - plloc S and llpllso S Cp- Let lug and Reg be as given in Lemma 10 Then, for 0 S h S 8 S 1/(1 + 2R). we have that
Bp. [IIv,(Y; (Y, .() u(Vkin 1)13
<C, { RR xp (=² + d5 + m)²))
+ Es,e
[RRg + ( + 4Q² + dh) + R5 22)
(m+
6²
12 + Cp(1 + +2R) nh
min (T - t, 1/ (2+ 2R))
where
4e (21
{ 03 (03(12g
R +x²)) 11 =πv + CRRIE* L2 + +
Es,e
8 + e2x((2)
3 and cy > 0 is universal constant.
Proof. Define
Pe, Elp (X7 [X; = =]],
pc(x)= = (lp((E]() =x].
First, note that
(p),(()) - Vp((s)
+ (0) -
||u" (z,, - u(x, t)]
Pv)((2)
(x)(
Vp,,(1)) VO((T))
+2 (p,())(p (x) 2
<
pr,)(z)
(re
<20 (z) (x)(2) +2 / (x) (1) p((())
(x)))
220 VV,, (x) + 2R (x) - p
Therefore, the expectation of the right hand side with respect to X,- can be bounded by Ex: [Xl .t) - =(x;; {)
< pp). (=:)))
+ 2R 2R, C&R [10)(())- -x (x:)
220 (x)
+ 2R 2R. -V>(c()))VV<(()??
{p((=))
[V(XF
X = =x)
The first term of the right hand side can be bounded by Theorem and Lemma 11 The second term can be bounded by eyv by Assumption 7
In the same vein, we can bound the difference t| k2
<20 [//)(>) - (2)11 +2R,CEEx, [l(e(x) (x))())
2/2 (m + 1Q² + dh) + &R (1 +2)22162
+ 2R 222 (m + 1Q² +dh)
<(2(133 + (m + 4Q² + dh) + R R1 + 1??
42

---

Published as a conference paper at ICLR 2025
where used (45) and the same argument as (45) in the last inequality with Rgg Lipschitz continuity of Pr.
Finally, convert the expectation w.r.t X; to that w.r.t X; However, the density ratio between PN and qi is bounded by Cp. which yields the assertion.
C.4 DOOB'S H-TRANSFORM
Lemma 12. For all E (0, T the following relationship holds.
V log qe (a) v log (x)) Vz log(Elp æI).
Proof. Let us denote the joint distribution of Xn and X, as Pa,(X X,). the conditional dis- tributions of Xo given X, as (Xoo and the conditional distributions X, given Xo as (X, Xo). Define Or in the same way. is straightforward to see that
(o)) + log(Eja (Xo)) x|) (xg(p(())
= x))
(=())) (x)d = logpt(x) /.
log
Pa (a))o,eea' x)dd'
log (1()))o(x (x)l(
log
(46)
Note that (=)(((/=') and (=)((=)>²) are the same in (46). Therefore,
+ (og(l)) = æ||) = log / 20(x')4no(xx)e (ejo(<laa'
= log (2),,
which concludes the proof.
D DETAILS OF NUMERICAL EXPERIMENTS
In this Appendix, we compare our method with existing approaches in a very simple numerical experiment and confirm that our method is also capable of performing image generation alignment. more practical and clearer comparison using realistic benchmarks is left for future work.
D.1 ALIGNMENT FOR GAUSSIAN MIXTURE MODELS
We explain how to align the pre-trained diffusion model for Gaussian Mixture Models and the tech- nical details of our algorithm. Almost the same algorithm was used for the other experiments.
Pre-Trained Score-Basee Diffusion Model We pre-trained score-based diffusion model to sample from the 2 dimensional mixture of Gaussian Mixture Models. The target density was E) (N2,, Σ)). 11 2-2.5, 0), H2 [2.5, 0).
[[1. 0]. [o. 5]]. The score model was implemented as simple 4 layer neural networks. The learning rate of pre-train was 0.0005, the batch size was 100, the number of epochs was 1000. The pre-train MSE loss shown the left side in Figure The minimum losses until the current epoch were plotted. The histogram of 20000 sam- ples from the pre-trained score-based diffusion model is the right figure in Figure) For sampling, T was set to be 10 and the number of sampling steps was 100.
The objective The target was the mean of the Gaussian in the right side, Mue := μ2 = [2.5, 0]. The preference of point Fn and zi determined by the Euclidean distance d(- (w)) d(., from p [2.5, 0J. F.. > 2j if and only if d( (w) < d(z), (.a). The DPO objective was used this setting. 2000 points from Prof were sampled calculate the expectation of the functional derivative. Note
43

---

Published as a conference paper at ICLR 2025
1055
10
00
107
x1
epoch
Figure 4: Left. Pre-train MSE loss of denoising score matching. The minimum losses until the current epoch were plotted. Right. The histogram of 20000 samples from the pre-trained DDPM.
that the reference o can be (and were) analytically calculated because q/Pref The regularization terms 3 and 7 were 0.04 and 0.1. The metric loss was calculated with the samples s.t. d(z,pa) < 10 random 1000 samples.
As counter method, Diffusion-DPO (Wallace al. 2024) was implemented with the learning rate 0.0005 and 0.0001 for w/ reg and w/o reg respectively, and batch size 5000. We remark that, in practice, it is hardly realistic compute the true DPO objective during Diffusion-DPO, but in this case, we forcefully carried out the computation by estimating the density ratio ( (r) with 100 2 samples for each x:: The densities q and Pref were estimated by repeatedly computing the denoising path and empirically obtaining their marginal densities, which implies some degree of numerical instability in Figure and additional computational cos in Table 1
Please note that the reason why the values of "Upper bound" and "True Objective" in the existing method shown in Figure are reversed that their "Upper bound" approximates the true DPO loss by replacing the current reverse process with the ROU process (Wallace et al. 22024. This implies that the upperbound may be too loose to be used as a surrogate the true loss.
The phenomenon that the "True Objective w/o reg." of the existing method diverges upward in Figure 2 likely due the absence of additional entropy regularization, leading catastrophic forgetting Uehara et al. (2024a); Tang (2024); Li et al. (2024b); /hhao et al. (2024). On the other hand, when applying our method, we incorporate entropy regularization.
Additionally, some degree of numerical instability is unavoidable due to the experimental method- ology because we empirically estimated the density ratio to compute the true DPO loss in using Diffusion-DPO.
Dual Averaging in Option We implemented Option for the experiment. We used 8 NVIDIA V100 GPUs with 32GB memory. The plotted losses (smoothed with EMA) in Figure demonstrated that DA algorithm decreased the losses. Please refer to Figure 5 for losses without smoothing. We estimated the log-density ratio (I In each loop, the potential was trained by 1000 points in æ1, 222 [-5, 5]. the learning rate was 0.0005, the number of epochs was 1000. We kept (lk-1) neural networks and generate the dataset
log (k) /Pref + const. with neural networks fx = (kk-
2the DPO loss of Ref. Table 1 our algorithm at & = hefore the random initialization of the potential, and the initial (true and approx.) objectives in Diffusion-DPO in Figure 3nd
3 "Opt.time" and "GPU memory "also include the time and space required compute the "True Objective" WW took (Xe) in (2) in Wallace et al. (222)) to be 0.5 (constant, recommended in the paper). We will leave the investigation of the phenomena that arise when w is turned into time-dependent function for future work.
44

---

Published as a conference paper at ICLR 2025
0.360
0.300
D.CIO
D-010 wn rag
0.345
True Obpective 10 +0.00
Tiat Capectiivn -0.60
Oleective 12 -0.20
0.355
Olgertiun of -0.10
0.30
0320
25
50
300 Rerettoo
175
200
Figure 5: (For Reference) Losses without smoothing corresponding to Figure Diffusion-DPO. "Ours": Proposed DA.
"D-DPO":
for g(k) from the equation (3 = for simplicity)
g(k)
2
[k(k + 1)
((h-1)
k
5F
(g
B(k + 1) + 2)
2
8q
k 2 k g((h())
+
2h
8F
(q
(R))
(47)
Bk(k + 2)
ōg
So. only need one model to store for DA loops (we don't need g(k g(), The psudocode for this phase described as Algorithm D.I
Algorithm D.1 Dual Averaging (Option 1)
Input: F: objective, pre-trainn tre-trained score, 3: Regularization scale, K; number of loops.
Output: fx: a trained potential.
Define q! = Pref as the reference density obtained by the pre-trained score s Initialize NN fi randomly
Collect samples (I:)); (e.g. collect data at the final denoising step from Pref by score function 8)
for k = K do
((k)
x exp(- f01-fxxee Pref (no actual computation)
Construct dataset ((Ti, Bk(k+2) 8g
2
AF
(g
,x:)]e with (xi)): and Sa using equation 2
Train fa+1 to approximate 242fk + pk(k++2) bg
2k 6F (x) by minimizing MSE.
end for End
The heat map of the trained potential is shown in Figure
Doob's h-transform We sampled the aligned images with 50 diffusion steps. The guidance term of Doob's h-transform (x( æ]] was calculated by Monte Carlo with 30000 samples. The psudocode for this phase described as Algorithm D.2] In phase (the sampling phase), our simplest solution (i.e., estimat- ing the correction term at each time step using Monte Carlo) has a time complexity of O(L²), where 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 si- multaneously, O(N) memory space required. The sampling error for each correction term would be (I/V/N). This leads to severer additional computational effort in phase than the time and space consumed as in Table for Phase | However, this Doob's h-transform technique itself has been used in image generation (Uehara et al. 2024a b), Bayesian samping (Heng et al. 2024), and filtering (Chopin al., 2023). As more practical alternative of phase 2, the idea of approx- imating the correction term using neural ODE solvers for faster test-time implementation has also been proposed in (Uehara al. 2024b a)
calculated in every diffusion steps. The conditional expectation
In Table we showed the computational cost for Phase I, not for Phase Il because we only require Phase I compute the true loss. Note that we used Phase II to compute the metric loss, however, the true loss is much important.
45

---

Published as a conference paper at ICLR 2025
k=1
k=2
k=3
K=
1320
144
0d8
076
0345
1183
1230
129
114
0.6
0825
9765
6.80
034
Target
0.66
0:44
0385
0465
6310
6643
0.5
0.99
028
6373
0.39
0.155
6.105
0.9
k=5
k=7
1.500
1.58
1466
133
145
130
134
0.66
0735
0.78
0.70
000
660
035
0445
044
542
040
0240
E24
022
0075
0006
004
k=9
k=10
k=11
k= 12
1.02
1:66
147
144
1.0
142
1.1
0.44
0.16
10.6
042
0.2
004
004
D02
k=13
1=14
k=15
k=16
1.43
144
1.22
1224
132
1.02
1.04
1.02
082
0.86
882
0.3
842
662
062
6/47
8/0
Figure 6: The heatmap of the potential in the kth loop in Dual Averaging for Gaussian Mixture Model. The target point was [2.5, 0]. Note that Sa is the negated log-density ratio: the aligned density is exp(-fx))Pre
Algorithm D.2 Doob's h-transform (A simplest implementation)
Input:
an objective, fK: a trained potential, st pre-trained score, output approximately from exp (-fx proff. Initialize wo as white noise
Output:
Set number of steps L and the time T. Set the step size h = T/L for T, do Initialize samples (21,))sg as dn. for 1' l', (denoising step of (In,)ss)) end for
do
Store N samples of X as (2I,iiiS..
Approximate u( 1h) V Voog Elexp (-fx(XF) X 21] by Monte Carlo with (2I.a)ieN.
Sample white noise Enoise
T1+11 == कi + 5(xi + 25(2), T 1h) + 2u(x1, Uhi) v 2hinoss
end for
End
46

---

Published as a conference paper at ICLR 2025
D.2 IMAGE GENERATION ALIGNMENT
We aligned the image generation of the basic pre-trained model in Diffusion Mod- els Course (source: HuggingFacce (2022)) The pipeline path utilized was @ohnowhitaker/dddmm outterflies 32px". The summarized results are show in Figure 8
The pre-trained Model The model samples the images of butterflies of 32 x 32 pixels. Number of the sampling step was 1000.
The objective The target color was [0.9, 0.9, 0.9] in RGB. The reward is visualized in Figure 7 6400 samples from Pref to calculate the expectation F and SF. / and y were 0.05 and 1. The dq output of the functional derivative was clipped in +00 to stabilize the training step. In calculation of the DPO objective, the sample that has a higher reward is the "winning" sample.
Dual Averaging In each loop, f& was trained by 1024 images from pooled 6400 images, the learning rate was 0.0001, the batch size was 64, and the number of epochs was 5. f&S were implemented by model in Diffusers library von Platen et al. (022). How the potential fs learned the reward, The DA algorithm
the distance from the target color, was shown in the right side of Figure 7 succeeded to extract the target images from the true reward.
reward
DA from reward
Figure 7: Left. Output images of Pret sorted by the distance from the target color (= rrwaadd). Right. Images sorted by the learned potential in k 2.
Doob's h-transform We sampled the aligned images with 1000 diffusion steps. The guidance term of Doob's h-transform calculated in every 10 diffusion steps for faster sampling. We calculated the drift term of Doob's b-transform once
addition, diffusion steps for faster sampling.
We defined decay rate ra 0.95 and strongness s for V log Pt. and finally added ra log (fexp(-sfKi (XX;- 2 210m5] as drift term in 10m m''h diffusion step balance the computational cost and stability. The conditional expectation (Xp x]] was calculated by Monte Carlo with 128 samples.
D.3 TILT CORRECTION FOR GENERATION OF MEDICAL IMAGE DATA
Our goal of this experiment generating images with no rotation with unconditional pre- trained model that generates rotated images. The objective was based on DPO. The rotation angle of each image was predicted with CNN. The summarized results are in Figure 12
Dataset 10000 images of Head CT (64x 64 pixel) in Medical Lozano 2017 were leveraged.
They include some rotated images, the angles are up 90". To aggravate the existing situation, augmented the data by rotating images (up to 45") from this dataset to be 40000. A predictor of angles were trained with the augmented dataset define the reward in place of human preference due to preparation difficulties. The number of epochs was 10, 95% and 5% of the dataset were used for training and validation, the training MSE loss was 3.79, and the validation MSE loss was 3.25.
47

---

Published as a conference paper at ICLR 2025
Ours
Target color: White
iter
iter
Ref.
Objective
Regularized Objective
0.3-
02-
10
20 Iteration
30 40
Figure 8: Left. The smoothed loss during DA for image generation alignment from k 1. "Ob- jective": DPO objective. "Regularized Objective": "Objective" + [D (gllPret). = 0.05 Right. Examples of aligned image generation. "iter=2": ours with k DA iterations, "iter=2": ours with DA iteration. "Reference"; samples from Pref
Note that there were rotated images in the original dataset, so the labels of the rotation angles made in the augmentation were noisy.
Pre-trained Autoencoder We pre-trained autoencoder from scratch. encodes gray-sealed 64 64 pixels into latent 32 32 pixels. It only has convolution layers that geometrical features were preserved for simplicity. The training data was 95% augmented 40000 images and the validation data was 5% of them. This autoencoder was fed white-noised data make the model robust with noise. The pretraining MSE loss is shown in the left side of Figure 9
We also pre-trained (latent) diffusion model based on Unet2DModel in Diffusers von Platen et al. (2022). Number of sampling steps was 1000. The beta scheduler was pre-training, we leveraged the augmented dataset up to 20000 images, the number of epochs was 15, the batch size was 8, the learning rate was 0.0001. The pretraining MSE loss is shown right side Figure 9 and the generated images are displayed in Figure 10
Latent Diffusion Model
set to be squaredduao coobaredoos.com _cppssccc_2....
01920
0.000
Validation Loss
2.2755
MSC
8.570
Train
30002
E.50
E3250
Foachs
Figure 9: Left. The MSE loss in pretraining the Autoencoder. The solid blue line and the dashed red represent the validation loss and the train loss. Right. The MSE loss in pretraining the diffusion model.
The objective The reward was defined to be -predicted anglel of the pre-trained predictor (in place of humans), described in the left side of Figure [T] 6400 samples from Pref to caleulate the expectation F' and &F 33 and were 0.01 and 0.1. The output of the functional derivative was clipped in +5 to stabilize the training step. In calculation of the DPO objective, the sample that has higher reward (in the more vertical direction) is the "winning" sample.
8q
Dual Averaging In each loop, fs was trained by 6400 images from Prof All the generated images reused during DA. The learning rate 0.0001, the batch size 64, and the number of
48

---

Published as a conference paper at ICLR 2025
32 32 (latent)
64 x 64
Figure 10: Top. Outputs of pre-trained latent diffusion. Bottom. Decoded images of the outputs.
epochs was 5. We compared the learned potential with the reference reward in Figure that DA iterations worked well to replicate the reference reward.
We see
reward
DA from reward
Figure 11: Left. Output CT images of the pre-trained model, sorted by the absolute values of estimated angles (= reward). Right. Output CT Images sorted by the trained potential in the 3rd loop in DA.
Doob's 's h-transform We sampled the aligned images with 1000 diffusion steps. The guidance term of Doob's h-transform was calculated once in 10 diffusion steps. Technical settings were the same with. The conditional expectation was calculated by Monte Carlo with 128 samples.
Ours
iter
iter
Ref.
0.40
Objective
0.38
Regularized Objective
0.0
0.2
Figure 12: Left. The smoothed loss during DA for tilt correction from & = 1. "Objective": DPO objective. The target point was [2.5, OJ. "Regularized Objective": "Objective" 3D (gllptt))) 3 = 0.01 Right. Tilt-corrected Head CT image generation. "iter=3": ours with & = 3 DA iterations, "iter=1": ours with k DA iteration. "Reference": samples from Pref-
49


Reviews:
[{"review_id": "cgSCgECpAb", "reviewer": "Reviewer_hXn1", "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.\n \nAuthors 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.", "questions": "What is $\\nabla_x \\nabla_x^\\top$ in line 1471\nWhat is the meaning of subindices $|_{x=\\bar X_t}$ ? The notation is very confusing\nWhat 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?\n\nCan 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?\n\nIn 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?\n\nIn 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?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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?\n\nOther concerns:\nOne 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?\n\nThe 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. \n\nThe derivative notation in line 143 is not clear.\n224: one and zero should be written in English rather than numbers", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SShyvrII0S", "reviewer": "Reviewer_aaJd", "summary": "The authors propose a distributional optimization algorithm that builds on the dual averaging (DA) algorithm combined with Doob's \nℎ-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.", "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?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\n2. The authors present a convergence rate analysis based on suitable assumptions, enhancing the theoretical robustness of the approach.\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pKabPzzrtL", "reviewer": "Reviewer_141C", "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.", "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?\n\n2. 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?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "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.\n2. 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.\n\n2. 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.\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g46wLAj5ye", "reviewer": "Reviewer_ei8r", "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.", "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?\n- 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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "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.\n- 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.\n- 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.\n- 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)? \n- 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.\n- 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. \n\nI am recommending 5 as a score. \n\nI 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).\n\nI will **raise my score from 6 to 8** if the authors provide **all missing text** from Appendix A. This includes \n- text between A and A.1\n- text between A.1 and Theorem 5 (which should be renamed to Theorem 1).\n- description of the significance of Lemma 1 (which seems to be a strong convexity-like result).\n- description of the significance of Lemma 2\n- any in-line explanations.\n- similar for all of A.2.\nPlease 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SOaXhdKQZp", "reviewer": "Reviewer_tyit", "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.\n\nOn 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.\n\nThe 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.=", "questions": ".", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n\nA 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. \n\nFinally, 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.\nAn investigation of the impact of $\\beta$ would also be interesting.\n\nThe 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.\n\nThe presentation needs to be clarified. Currently,\n  - 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.\n  - The recipe for estimating the density ratios during DA is not clearly described in the main body. \n  - 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.\n  - 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?\n  - 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.\n\nAdditional minor style/typos improvements are suggested in minor.\n\n\n\n\n### **Minor**\nThe style could be substantially improved. The english grammar mistakes and imprecision should be corrected.\n\n- the differences between the two DA techniques are not explained. Turns out some results hold for Opt 1 or Opt 2.\n- Typo l 1235, 1236, 1249: $\\nabla p_t$ -> $\\nabla \\log p_t$\n- \"Particle-based optimization methods are not designed for resampling from the distribution where obtained particles follow\" -> \"...that the obtained particle follow\"\n- \"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.\n- \"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.\n- Consider using ``test'' for correct quotes formatting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)