Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
Authors: Yinbin Han
Meisam Razaviyayn
Renyuan Xu
Venue: ICML 2025 poster (2025)
Keywords: Diffusion models
fine-tuning
stochastic control
reinforcement learning
policy iteration
optimality
regularity
convergence
Primary area: deep_learning->theory

Abstract: Diffusion models have emerged as powerful tools for generative modeling, demonstrating exceptional capability in capturing target data distributions from large datasets. However, fine-tuning these massive models for specific downstream tasks, constraints, and human preferences remains a critical challenge. While recent advances have leveraged reinforcement learning algorithms to tackle this problem, much of the progress has been empirical, with limited theoretical understanding. To bridge this gap, we propose a stochastic control framework for fine-tuning diffusion models. Building on denoising diffusion probabilistic models as the pre-trained reference dynamics, our approach integrates linear dynamics control with Kullback–Leibler regularization. We establish the well-posedness and regularity of the stochastic control problem and develop a {policy iteration algorithm (PI-FT)} for numerical solution. We show that PI-FT achieves global convergence at a linear rate. Unlike existing work that assumes regularities throughout training, we prove that the control and value sequences generated by the algorithm preserve the desired regularity. Finally, we extend our framework to parametric settings for efficient implementation and demonstrate the practical effectiveness of the proposed PI-FT algorithm through numerical experiments.

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
Yinbin Han Meisam Razaviyayn 2
Renyuan Xu
Abstract
neering fields. Traditional approaches, which fit models to offline or online data, heavily depend on data quality and availability. As a result, data limitations can signif- icantly degrade performance (Doersch, 2016; Durgadevi et al., 2021; Fetaya et al., 2020). For instance, constraints in experimental design often lead to data scarcity, hindering model effectiveness. Moreover, acquiring sufficient data for reliable experimentation is costly and time-consuming, posing scalability challenges. Most critically, data collected under specific conditions may not generalize well to new environments or downstream tasks, limiting the adaptabil- ity of learned models (Alzubaidi et al., 2023; Bansal et al., 2022; Gangwal al., 2024).
Diffusion models have emerged as powerful tools for generative modeling, demonstrating excep- tional capability in capturing target data distribu- tions from large datasets. However, fine-tuning these massive models for specific downstream tasks, constraints, and human preferences remains critical challenge. While recent advances have leveraged reinforcement learning algorithms tackle this problem, much of the progress has been empirical, with limited theoretical understand- ing. To bridge this gap, we propose a stochastic control framework for fine-tuning diffusion mod- els. Building on denoising diffusion probabilistic models as the pre-trained reference dynamics, our approach integrates linear dynamics control with Kullback-Leibler regularization. We establish the well-posedness and regularity of the stochastic control problem and develop policy iteration algorithm (PI-FT) for numerical solution. We show that PI-FT achieves global convergence at linear rate. Unlike existing work that assumes regularities throughout training, we prove that the control and value sequences generated by the al- gorithm preserve the desired regularity. Finally, we extend our framework to parametric settings for efficient implementation and demonstrate the practical effectiveness of the proposed PI-FT al- gorithm through numerical experiments.
Generative modeling offers flexible solution by creating synthetic data that maintains the properties of collected data while enhancing data diversity. Diffusion models, such as those proposed by Ho et al. (2020); Sohl-Dickstein et al. (2015); Song & Ermon (2019); Song et al. (2021), have emerged as powerful tools in this area, supporting notable advancements such as DALL-E (Betker et al., 2023; Ramesh et al., 2022), Stable Diffusion (Rombach et al., 2022), and Sora (OpenAI, 2024). These models excel by learning the score function from potentially high-dimensional and lim- ited data, extracting critical information for the data gen- eration process. To achieve cost-effective, task-specific data generation and multi-modal integration, techniques like alignment (Wallace al., 2024), including guidance (Dhari- wal & Nichol, 2021; Ho & Salimans, 2021) and fine-tuning (Black et al., 2024; Clark al., 2024; Fan & Lee, 2023; Fan et al., 2024), of pre-trained models are proposed. While dif- fusion alignment has achieved significant empirical success, the theoretical foundations remain in the early stages-a gap this paper seeks to address.
1. Introduction
The growing availability of large-scale datasets, coupled with advances in high-performance computing, has fueled the rise of data-driven methods across scientific and engi-
Aligning a pre-trained diffusion model to a specific task or dataset requires updating its parameters through additional training, typically using smaller, task-specific dataset (Dai et al., 2023; Podell et al., 2023; Rombach et al., 2022; Wal- lace et al., 2024). Among the prevalent alignment tech- niques, fine-tuning and guidance differ in how the additional criteria are handled-eithe as a soft constraint or a hard constraint. The soft constraint approach particularly ef- fective for applications that incorporate human preference
Equal contribution 'Department of Finance and Risk Engi- neering, New York University DDaniel Epstein Department of Industrial and Systems Engincering, University Southern Cali- fornia. Correspondence to: Renyuan Xu <Ix2364@ nyu.edu>.
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
or embed human values as a reward signal evaluated on the task-specific dataset (Black et al., 2024; Clark et al., 2024; Fan et al., 2024; Uchara et al., 2024c; Zhao et al., 2024b), which is the central focus of our paper.
fusion models that presents a convergence guaraniee. We address this challenge by leveraging linear dynamics and the (one-step) concavity introduced by the KL regularization term, which inherently captures the essence of fine-tuning problems. Specifically, we establish the universal regularity of both the optimal value function and the sequence of value functions generated during iterative updates. Moreoverover, the foundational techniques underlying our results extend naturally parametric algorithm design, offering insights beyond the specific context of diffusion model fine-tuning.
Our work and contributions.
We introduce a discrete-time stochastic control formulation with linear dynamics and Beulbackk-eibble (KL) regular- ization for fine-tuning diffusion models. Specifically, we establish a novel connection to denoising diffusion proba- bilistic models (DDPMs) (Ho et al., 2020) by treating the fine-tuned score as a control to be learned. In particular, the soft constraint or human preference is modeled by a reward signal evaluated on the generated data output and the KL reg- ularization term penalizes the deviation of the control from the pre-trained score function. We utilize the discrete-time formulation since fine-tuning is inherently a discrete- time problem, as it relies on a pre-trained model which is typi- cally implemented in discrete time, with DDPM being the most commonly-used example. By appropriately select- ing the regularization parameter, we demonstrate the well- posedness of the control problem and analyze the regularity of the optimal value function. A key insight arises from the concavity of the nested one-step optimization problem in the Bellman equation, a direct result of KL regularization under a properly chosen regularization coefficient. Leveraging this property, we develop a policy iteration algorithm (PI-FT) with guaranteed convergence to the globally optimal solu- tion at a linear rate. A central challenge involves preserving the regularity of the sequence of value functions and out- put controls generated during algorithm iterations. This is achieved through a novel coupled induction argument and precise estimations of the regularization parameters. Finally, we discuss the algorithm design in parametric settings and conduct thorough experiments demonstrate the practical efficiency and effectiveness of the proposed PI-FT method.
Related literature.
Our work is related to several emerging research directions.
Fine-tuning of diffusion models. Our framework is closely related to the recent studies on fine-tuning diffusion models. Motivated by advancements in the (discrete-time) RL literature, Fan & Lee (2023) was the first to introduce soward-based approach for improving pre-trained diffu- sion models. Building on this idea, two concurrent works (Black et al., 2024; Fan et al., 2024) proposed a Markov decision process (MDP) formulation for denoising diffusion processes. To mitigate reward over-optimization (Gao et al., 2023), Fan et al. (2024) examined the impact of incorpo- rating KI regularization as an implicit reward signal. Fol- lowing this idea, Uchara et al. (2024b) introduced an online fine-tuning framework where the regret is upper-bounded by the accuracy of a statistical error oracle for reward estima- tion. Just recently, Yoon et al. (2024) formulated diffusion training as an inverse RL problem and proposed a value- based algorithm. For comprehensive review of this topic, we direct interested readers to Uehara et al. (2024a). Be- yond RL-based fine-tuning methods, alternative approaches include classifier guidance (Dhariwal & Nichol, 2021) and classifier free guidance (Ho & Salimans, 2021), supervised fine-tuning (Lee al., 2023), LoRA and its variants (Hu et al., 2022; Ryu, 2023), DreamBooth (Ruiz al., 2023), and DRaFT (Clark et al., 2024) among others.
While existing methods such as DPOK (Fan & Lee, 2023) and DDPO (Black et al.. 2024) rely on generic RI. devel- opments such as PPO or REINFORCE, they do not fully leverage the fine-tuning structure. On the contrary, the spe- cific setting considered in this work, with linear dynamics is tailored towards the devel- opment of efficient fine-tuning diffusion models. For this reason, the PI-FT algorithm and its parametric extension directly computes the policy gradient of a Ke-regullrize control objective. This principled design leads to a more ef- ficient implementation in practice compared to prior works, while also offering theoretical guarantees.
All the aforementioned references focused solely on empiri- cal investigations, except for Uchara et al. (2024b), which leaves theoretical foundations largely unexplored. While the MDP formulation discussed above aligns well with the RL literature, it overlooks the structural properties inherent to diffusion models. In contrast, our approach specifically leverages DDPM as the pre-trained model, thereby fully utilizing the underlying structure of diffusion models. More- over, the stochastic control formulation is more appropriate than RL as the dynamics are known in diffusion models and exploration is not necessary. Consequently, the efficiency of our algorithm stems from taking control-based formulation.
and
Despite promising empirical evidence in the literature, ana- lyzing algorithm convergence remains challenging due to the continuous nature of the state-action space. To the best of our knowledge, this work is the first on fine-tuning dif-
Continuous- control/RL for fine-tuning, Inspired by the continuous-tin nature of diffusion processes, recent
2

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
work has explored the application of continuous-tiin con- trol/RI. for fine-tuning. For instance, Berner al. (2024)
Notations and organization.
Take d E N+. the set of all positive integers. We de- note (2,F == (Fi PP) as usual filtered probabil- ity space supporting random variable R L L' Fo), namely, R(v) E Fo has finite mean for all y E R Denote {W>) 1 a sequence of independent standard d- dimensional Gaussian random vectors. We denote by FW the natural filtration of (Wo) In addition, M(,, 2) denotes the normal distribution with mean μ E R and covariance E E Rdxdd Let f(- |1s, Σ) denote the probabil- ity density of (M, 2). For two probability distributions P and Q such that P < Q with densities p and sup- ported on Rd we denote the KL divergence by KL(p]]g) = J& p(y) log p(v) dy. Finally, Ia E R denotes the iden- s(x) tity matrix and : 2l2 denotes the Euclidean norm.
established a connection between fine-tuning and stochastic control by analyzing the Hamilton-Jacobi Bellmmn (HJB)
equation for the log density. Domingo-Earich et al. (2025) introduced an adjoint matching method for the fine-tuning of rectified flow, drawing inspiration from Pontryagin's s maximum principle. A more rigorous treatment of fine- tuning within the framework of entropy-regularized control problems is developed in Tang (2024), addressing the well- posedness and regularity of the corresponding HJB equa- tion and extending the analysis to general f-divergences, though no specific algorithm was proposed. Following this entropy-regulariiztion perspective, Zhao et al. (2024a) derived the formulas for continuous-time policy gradient method and continuous-time version of proximal policy op- timization (PPO). Furthermore, Gao et al. (2024) applied a continuous-time q-learning algorithm to simulate data that reflects human preferences directly, without relying on pre-trained models.
The paper organized as follows. Section 2 introduces the set-up of the stochastic control problem and establishes the well- well-pooednees and regularity. Section 3 proposes a policy iteration algorithm, develops the linear convergence result, and discusses parametric extension. Finally, numerical experiments are demonstrated in Section 4.
Compared with the literature that technically remain at the level of exploring HJB equations, we establish convergence rate for the proposed algorithm in discrete time. Despite the continuous-timn nature of diffusion processes, the discrete- time setting is more suitable because fine-tuning builds upon a pre-trained model, which is naturally implemented and provided in discrete time.
2. Problem set-up and regularity results
2.1. Problem set-up
To address the challenge of fine-tuning, first introduce the dynamics of denoising diffusion probabilistic models (DDPMs), a widely adopted pre-trained framework in prac- tice. Our focus is on the discrete-time formulation, which aligns with the standard implementation of diffusion models in practice. In addition, DDPMs serve as a foundational scheme for effective fine-tuning.
RL theory. Our theoretical developments are also closely connected to the literature of RL and policy gradient meth- ods in discrete time. For MDPs with finite state and action spaces, recent advancements providing global convergence guarantees for policy gradient methods and their variants can be found in Berner et al. (2024); Bhandari & Russo (2021; 2024); Cen et al. (2021); 1) Ding et al. (2020); Fatkhullin et al. (2023); Fu et al. (2021); Liu et al. (2019; 2020); Mondal & Aggarwal (2024); Wang et al. (2020b); Xiao (2022); Xu et al. (2021); Zhan et al. (2023); Zhang et al. (2020; 2021). Beyond MDPs, policy gradient methods have also been applied to Linear Quadratic Regulators (LQRs), specific class of control problems characterized by linear dynamics and quadratic cost functions (Bu et al., 2019; Fazel et al., 2018; Guo et al., 2023; Hambly et al., 2021; Han et al., 2023; Malik et al., 2019; Mohammadi et al., 2019; Szpruch et al., 2021: 2024; Wang et al., 2020a: Zhou & Lu, 2023). While the convergence analysis of policy gradient methods is well- established in the tnove-mentionedd settings, the study control problems with continuous state-action space and general cost functions has been limited. Our work ventures into this broader class of control problems, providing con- vergence guarantees for policy gradient algorithms under more general settings.
Denoising diffusion probabilistic models (DDPM). A well-trained DDPM {Y/"*)T in discrete time follows the following stochastic dynamics with state Y Rd.
1 (Yprre + (1 (1-at)s pre (Yprr) + o.W, (1) 1+1 Vai
with YpPre ~ NN(0, Here, (W T-1 are i.i.d. standard Gaussian random vectors such that W, ~ NO, Ia). The hyper-parameters {at)T=) T- with a E (0, 1) and {o.)T-) T-1 with o, > 0 represent the prescribed denoising rate sched- ules (Ho et al., 2020; Li et al., 2024). They control the vari- ance of noise in data generation'. Here, sf"ee : IRd + Rd is the score function associated with the pre-trained model. In practice, the score estimator spre is obtained by train- ing neural network to minimize the score matching loss (Hyvärinen & Dayan, 2005). With well-trained pre-trained
Whhile DDPM chooses of 1/a4- our analysis in the subsequent sections works for general at-
3

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
model, we expect the distribution of YT to be close to the target distribution, from which the pre-trained model has access to samples and seeks to generate additional samples. Given s!t", we denote Pet1je (-I16) as the conditional density of y given Y yt induced by the dynamics (1), i.e.,
Remark 2.2 (Control as score function). Our formulation is rooted in both diffusion models and control theory. Com-
paring (3) with (I), we replace the pre-trained score spre by a control policy ue. Consequently, in the context of diffusion models, the learned control sequence (u, )=0 can be viewed as the new score function of the fine-tuned model. In other words, solving the control problem (2)(3) is essentially learning a new score function in response to the reward signal R(-) for fine-tuning. Note that the linear dynamics (3) with Gaussian noise is preferred in stochastic control due to its tractability in analysis. Specifically, our theoretical analysis in the subsequent sections heavily relies on the linearity and the Gaussian smoothing effect.
Petije f (-11/ (ye+ + (=a))ss (g))) La
Remark 2.1 (Choice of the pre-trained model.). DDPMs have become the leading choice for pre-trained models across wide range of applications, serving as power- ful building block for fine-tuning in diverse tasks (Dhariwal & Nichol, 2021; Ho et al., 2020; Li et al., 2024; Watson et al., 2023). By pre-training on large datasets, DDPMs cap- ture complex data distributions, enabling relatively easier fine-tuning for specialized applications such as conditional image synthesis, text-to-image generation, and even time- series forecasting. This versatility has been demonstrated in models such as Stable Diffusion (Rombach et al., 2022) and DALL-E (Ramesh et al., 2022), where fine-tuning on taek-specific data enhances performance and customization (Fan et al., 2024).
Remark 2.3 (Rationale for the objective function and the known reward assumption). The objective function in the control formulation (2)-3) consists of two parts: a termi- nal reward function R(-) at time T and intermediate KL penalties. The reward R(-) captures human preference on the generated samples. For example, in the text-to-image generation, the reward R(-) represents how the generated data Y1 is aligned with the input prompt (Black et al., 2024; Fan et al., 2024). In addition, the penalty term ensures that the fine-tuned model not too far away from the pre-trained model, which prevents overfitting. Unlike the Shannon en- tropy of the (randomized) control policy commonly used in the RL literature to encourage exploration, the KL. regular- ization in (2) is applied between two conditional probability densities. From an optimization perspective, the KL di- vergence introduces concavity in control and consequently leads to a better landscape of the objective. Indeed, we will choose the parameter 3, sufficiently large to guarantee the existence and uniqueness of the optimal control and to satisfy certain regularity conditions; see Theorem 2.8.
Stochastic control formulation. When human prefer- ences or soft constraints are modeled through a stochas- tic reward function R(.), our objective is to maximize the following stochastic control formulation, which models fine- tuning tasks in diffusion models:
(2) Σ B.KL (p4)) || (Da)) L=
E R(YT)
with state Y E IRd and control U. E Rd following the dynamics
In practice, fine-tuning is typically performed on a dataset containing reward ratings for each data sample, which is much smaller than the pre-trained dataset. Using this fine- tuning dataset, one can estimate the expected reward func- tion. The assumption of known reward is not overly re- strictive, online learning to acquire new rewards is not generally uncommon.
1
Y441 =
(Y,+ (1++(1a))U) u.).)) + W.W, Var
(3)
where Y, ~ (0, Ia). In particular, we work with Marko- vian policies such that the objective in (2) is dell-defined. Mathematically, the control process (U+ )z= T-1 is admissi- ble if each U. can be expressed as U, = un (Y:), where ue : Rd + Ed is measurable function ensuring that the objective in (2) remains finite. In terms of the reward, we as- sume R E L ( namely, R(y) E Fo has finite mean for all y E Rdd. We assume the reward function known in this work. The KL-divergence measures the deviation of the fine-tuned model P1+ from the pre-trained model Petije where Pt+1/t denotes the conditional distribution of Y++1 given Y, induced by the control policy u, The regu- larization coefficients [ control the strength of the regularization term.
Remark 2.4 (Connection to KL divergence on path-wise measures.). Another natural idea is to impose the KI. diver- gence between the path-wise measures, i.c., (L (l where Po:r is the joint density of (Y:)² and pprr is the joint density of (Y.Pre T J==0 This choice is relevant to for- mulation (2)-(3). In particular, the Markov property of the dynamics and the chain rule of the KL divergence imply
T
KL (ITT)) = E
Σ KT. (pee (P
We have a few remarks in place.
where the expectation is taken over all random variables
4

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
(Y1)=0- Here, we define
where the supremum is taken over all admissible control policies [ The Dynamic Programming Principle
(P (IT)T) = = po:r(y))og log Po:T(y) dy, pre (y)
implies that Vi2 satisfies the Bellman equation:
(4)
with y = (Vo..... (oo)) and for given y E IR., denote
(V)
sup E
V
(y + (1 -e)u (v)) + o.W.)
Vaa
Br (1 ax)
KT. KT. (Petajel ( (o))llessjee (we))
[1ux(y) - "
(8)
20,0
= -fDa ( log Pe+ije(Be++ 112t) d P (8t+11B1))
(5)
In the next subsection, we will discuss the well-posednee of the optimal control problem (2)-(3 and the regularity of the optimal value function V
Thus, when Be for all l, the objective in (2) becomes
2.2. Regularity and well-posedness
(IVA)
E
R(YT) Σ 3.KL
In this subsection, we establish the regularity of the optimal value function and the optimal control policy. To start, we make the following assumptions on the reward and pre- trained score functions.
-E(R(Y))) - (KL(p (
KL divergence is common choice of regularization in fine- tuning of diffusion models (Fan et al., 2024; Gao et al., 2024; Uehara et al., 2024b; Zhao et al., 2024a). However, our work the first to use KL divergence over transition dynamics (on path space) to control the deviation of the fine-tuned model from the pre-trained model, whereas formulations in the literature consider KI. between the terminal stute distributions. Our formulation potentially utilizes more information from the rich pre-trained models.
Assumption 2.6 (Smoothness of the reward). Assume r④y) = E[](v)] is Lg Lipscchct and Li-grddient Lipschitz in y E TRd, i.e., the following holds for any yn, y2 E Rd:
In(3s) r(22) - L5 11yı 32ll2
(9)
Vr((n) - Vr(22)l S 41 11 92ll2-
(10)
Assumption 2.7 (Smoothness of the pre-trained score func- tion). Assume spre is Lo-Lipschitz and Li-gradient Lips- chitz in y IRd, i.e., the following holds for any V1: V2 ERd
Given two Gaussian densities Pt+1|z and pft in (2), the following lemma simplifies the KL divergence term.
Lemma 2.5. Consider the KL divergence defined in (5). For any y Ied and any admissible control policy ue, it holds that
| (v1) - >.(" S T. 1l -22122 (11)
V (7 V. (12)ll2 - Li. l1 -vell2 (12)
KL (pe42)e(-le))) (p
While Assumptions 2.6 and 2.7 guarantee the smoothness of both the expected reward r(.) and the pre-trained score (spro)?= no convexity assumptions are imposed, and thus the control problem (3) is in general non-concave with respect to {u
(1-ar) )² [(u)(ve) spre (36)
20y0?
Lemma 2.5 links the KL-divergence between two condi- tional distributions with the squared loss between the con- trol and the pre-trained score o!"e. Since the control can be interpreted as the new score of the fine-tuned model, (6) enjoys the spirit of the score matching loss in training diffusion models (Han et al., 2024: Ho et al., 2020; Song et al., 2021). The proof of Lemma 2.5 is based on direct calculations; see Appendix With this in hand, we define the optimal value function at time as
Next, we establish the existence and uniqueness of the op- timal control, as well as the regularities of the solution to problem (2)-()). For case of exposition, we define a series of constants recursively. Let de E (0, 1) for each < T. Set L Lj and L = LT For < T, define
1
0.1
(1 + (1 oz) 4
(13)
Var.
L.
(1 + (1 - an) Lo. (1 + (1 a.)
(v) = sup E
-
T-1
Σ
3r
(1 02
20u0
|lue 53 (Y:)112
at
+ 1 1 a: L -
(14)
E=ع
Van
+ R(X+)
Yt =
(7)
I.a x-1
[ + 1- X1
(15)
1 04
5

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
E [[ (1 - Ac)
Direct calculation leads to the Lipchitz condition of ou and the smoothing effect of the Gaussian random variable is utilized to obtain differentiaaility. In addition, the Lipschitz property of Vu; is a consequence of the integration by parts formula in (35); hence (iii) holds. Finally, the regularities of Vi2 follow from the Lipschitz conditions of uj and Vug, which completes the induction.
11 = 27
1i +
11 anlvo
x
(1 0x) Lo)Li. (1-
(16)
Here, E [1]W] = (2(( < 00 is constant with T(d/2)
I(.) being the Gamma function. Note that the above con- stants only depend on the system parameters Ot,, Li- Li Li and Li as well as the hyper-parameter X In the next theorem, we show that the above constants are indeed the Lipschitz coefficients for {V('),- and [u]]:-o
3. Algorithm development and convergence analysis
In this section, we propose an iterative algorithm, PI-FT, for fine-tuning and provide non-asymptotic analysis of its con- vergence. Assuming that the expected reward function r(.) is known, we develop an iterative algorithm approximate the optimal policy in high- sigh-dimensional settings, suitable for practical implementation. In practice, r(-) is approximated using a small sample set with human feedback, which then serves as the basis for fine-tuning. The unknown r(-) case is left as a topic for future investigation.
Theorem 2.8 (Regularity and well-posedness), Suppose Assumptions 2.6 and 2.7 hold and let the constants be defined according to (13)-(16). Choose Be such that 3.
1 o 4 2 A O. Then,
(i) The optimal value function Vi (y) defined in (7) is LXii-Lipschis, differentiable and L ggaddient Lips- chitz y E Rd.
3.1. Proposed algorithm
(i)) There is a unique optimal control ut RR + ad of problem (3) satisfying
Recall that the optimal control uz satisfies (17). Motivated by this observation, we propose the following iterative algo- rithm for fine-tuning and computing the optimal control ut-
(v) = -s + Varo
E VVAn (0)), (17)
3
Algorithm Policy Iteration for Fine-Tuning (PI-FT)
where y/ (1 az) (v)) + or W. Moreover, the optimal value function Vi is the unique C. solution to the Bellman equation (8).
1: Input: Expected reward function r()) pre-trained model {spre)" 1t=0- and number of iterations {me) 77-1
2: Set V ((TT) (y) = r(y) for all y E Rd.
(iii) The optimal control ut (v) is Lipsccit, differen- tiable, and L -gradient Lipschitz in E Rd.
for T 1.....00 0 do
4: Set u, 0 (y) = sp (y).
5: for m I,..., m, 1 do
When the regularization coefficient 33 is sufficiently large, Theorem 2.8 states that the optimal value function V is Lipschitz and gradient Lipschitz. The choice of 3r also ensures the right hand side of (8) is strongly concave in u.()). guaranteeing the existence and uniqueness of the optimal control uz. Furthermore, Theorem 2.8 shows the optimal control u is also Lipschitz and gradient Lipschitz. The regularities of V and ut will serve as the foundation of the algorithm design and convergence analysis in the subsequent sections.
Update the control using
ui (m+1) (v) = Vagoz
(m(mm))) (yim))
(1 or )r
+ (g).
(18)
where y 1 (v + -1-aa) (mm) (y)) + to o.W,
(mm)
vor
7:
end for
8: Compute the value function ((lme)) using
The main challenge in proving Theorem 2.8 stems from the fact that Eq. (17) is an implicit equation, in which both sides involve the optimal control uz. This prevents direct appli- cation of Lipschitz assumptions model parameters. We outline a brief proof sketch to highlight the ideas and defer the detailed proof to Appendix B, which is based on back- ward induction. Assuming V is Lipschitz and gradient Lipschitz, the first-order optimality condition implics that (17) holds, given the optimization problem is unconstrained.
V.(mme) (2) = E vilm+s) (m)))))
B (1 or
1|u (mi) (v) - -s. s. (v)
20r0?
?
(19)
9: end for
(me] T-1 Je
and (v/mea))
12
10: return
1==0
6

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
Algorithm 1 consists of two nested loops. The inner loop updates the control at each iteration m and the outer loop computes the value function at time t. Given the value function V++1 1) at time step / + 1, we update ut ((m)) using (18) and then evaluate the associated value function using (19). Note that (18) can be viewed as an approximation of the following update rule:
Let { (umme) 2-1
J and { {vimea)
be the oulput of Algo- ) tmo
rithm 1. Then it holds that
V(((
(mme) 1
(y) - ui (v) l2
(1- Am I.0.+1 + 21 Beti) var - Ar)
S (11
(1-os) Lo)
(23)
1m=+1) (s)= pre (v) + varoz
E (y (20)
(1 )33
where
F₂ = (V((m)) (v) (x) vVi
with y yml) = 1/ ( + (1 - -a)) (m)) (v)) + o. We, and (19) can be seen as an approximation of the Bellman equa- (((++)))
T-1 k-1
tion (8), where we replace the unknown Vi+1 with '+1
$
II C
C2 - (x) met1 (24)
Direct calculation shows that the fixed-point update rule
Aunt
(20) guarantees the convergence of " }
(m7)
and
to ui; see m20
C₁ = 1+ (1 (1-,a))).
(69). However, in practice we cannot access Vit1- Hence,
and
Eq. (18) can be viewed as a proxy to (20) as long as we can control the error
\ivoi
VV (m)) (y) VV; (v) for all t.
C2. = (1 + (1 (1e) (1) L0,
1-(m use (19) calculate m/(m) With uy (m²) and t+1 an estimate of the optimal value function Vi",
Vot
(1 + (1-11-a))) -1-11) Ho..)
+
Vat
3.2. Convergence analysis
To obtain simplified convergence rate from Theorem 3.1, we set d = x == minog ) and m; = M for all t. In this case, we have E. = O((1 - A) and consequently (M)) (m) - 47 (v) = O((1 A)M). In other words, T-I
In this subsection, we analyze the convergence of Algorithm 1. For fixed d E (0, 1), define a series of constants recursively as follows. Let L8,7 Lg and LY, = LT. For each & < T, set
T. =
(1+ - a.)
[=0 converges to the optimal control at a linear rate, Note that the parameter À determines the convergence rate. In particular, setting a larger results in faster convergence. However, when À is large, Theorem 3. 1 requires that the regularization parameter 3, must also be chosen to be sufficiently large. Recall that 3, controls the distance hetween the fine-tuned model and the pre-trained model, while a bigger X may result in an optimal control far away from sp Therefore, choosing an appropriate À (or equivalently (3)) is crucial to trade-off computational efficiency and closeness to the pre-trained model.
the control sequence u,
(m)
v
V
0
(1 (1+ -a))aa)) (1 -2),; (21)
+ Va
LP =
1 (1+ + (1 ax) Lo.) + (1 Lo
at
+ 1 - at
1 at (
(1-00 Li, + (m + 2)
Lÿ Li,100,41 + max
Va
m21
Vaz
E [1
x (1 + (1 (.)
1. T.0.41(1= -0,1 mm) |
(22)
The convergence analysis of the PI-FT algorithm hinges on proving that the regularity of the value function is preserved throughout the training process. Here, we present a proof sketch of Theorem 3.1, with the full proof deferred to Ap- pendix C. The key idea is to employ backward induction, breaking the problem down into a sequence of one-step analyses. Assuming V is both Lipschitz and gradient Lipschitz, Lemma C. establishes the regularities of both (!"")) and V((m) throughout the update rule (for all m 2 0). Furthermore, given the approximation error at time t + 1, Lemmas C.2 and C.3 provide the error bounds of both the control and the gradient of value function at time t. This completes the induction.
Here, the constant LY, is well-defined as the maximum must be achieved for some m < oo due to the exponential decay of (1 m Similar to the constants defined in (13)-(16), the above constants only depend on the system parameters 0g, Lg LT, Lg and Li as well as the hyper- parameter A₁. Later, we will show that the constants defined in (21) and (22) indeed serve as the Lipschitz coefficients of [v(m)
universal over all m 2 0. The main result of this subsection is stated follows.
Theorem 3.1. Suppose Assumptions 2.6 and 2.7 hold. For d each t < T, choose Be such that - 2, p 2 X > O.
7

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
LoRA (Hu al., 2022) and ImageReward (Xu et al., 2023). Following (Fan et al., 2024), we four prompts- "A green colored rabbit," "A cat and dog." "Four wolves in the park," and "A dog on the moon"-4t evaluate the model's 's ability to generate correct color, composition, counting, and location, respectively. During training, we generate 10 trajectories, each consisting of 50 transitions, to calculate the gradient with 1000 gradient steps. By default, we use the Adam W optimizer with learning rate of 3 10-4 and set the KL regularization coefficient to fixed value 3 0.01.
ImageReward Score
1.5
1.0
0.5
0.0
Pre-trained
DPOK
-0.5
PI-FT (Ours)
Evaluation. We first compare Image scores for im- ages generated by the pre-trained model, DPOK (Fan et al., 2024), and our proposed PI-FT. For a fair comparison, we configure DPOK to perform 10 gradient steps per sampling step, using a learning rate of x 10-5, Each gradient step is computed using 50 randomly sampled transitions from a replay buffer. As result, 1000 sampling steps and a total of 10,000 gradient steps in DPOK yield a computational cost comparable to that of PI-FT. As shown in Figure 1, PI- FT consistently outperforms both baselines across all four prompts. Figure 3 further shows that PI-FT more accurately captures object counts and placements (c.g., four wolves and the dog on the moon) and avoids errors like miscoloring the rabbit. It also produces more natural textures compared to the baselines.
Color
Count Composition Location
Figure 1. Comparison of the IeageReward score among three mad- els: pre-trained Stable Diffusion (red), DPOK (yellow), and PI-FT (green). The RewaeReward scores are averaged over 50 samples from each model.
3.3. Extension to parametric formulation
To enable practical and efficient update of Algorithm we adopt linear parameterizaatio in this subsection:
we (x( K,o(y). t O...., T
(25)
where (v) = (v), (1)) is a given basis func- tion, and K = (K.) an (unknown) parameter to be learned. Despite its simplicity, this parameterization is flexi- ble, expressive, and tractable, and is widely used in control and reinforcement learning (Agarwal et al., 2021; Jin et al., 2020). With an appropriate choice basis, can capture a broad class of score approximators. including random fea- tures (Rahimi & Recht, 2007), kernel methods (Steinwart & Christmann, 2008), and overparameterized neural networks in the NTK regime (Jacot et al., 2018). This choice of rameterization enables linear convergence to the optimal solution; details are deferred to Appendix A.
Effect of KL regularization. KL regularization is known to enhance fine-tuning. We study its effect in PI-FT us- ing the prompt "Four wolves in the park," varying 3 E {0.01,0.15 1.03. As shown in Figure 2a, the gradient norm decreases to zero in all cases, indicating convergence. Figure 2b shows that small 3 values improve and stabi- lize the ImageReward score, while larger / offers limited gains. This aligns with Figure 2c, where KI divergence remains high for 3 = 0.01, but stays significantly lower for 3 E {0.1, 1.0]. Figure 4 also shows that smaller / produces images with nearly four wolves, whereas larger & leads to fewer. These results underscore the importance of the KL. coefficient effective fine-tuning.
4. Numerical Experiments
In this section, we evaluate the performance of the PI-FT al- gorithm from Section 3 via numerical experiments, focusing on the following questions:
5. Conclusion
We introduce a stochastic control framework for fine-tuning diffusion models, integrating linear dynamics with KI. reg- ularization. Our approach establishes the dell-posedness and regularity of the control problem and proposes a policy iteration algorithm (PI-FT) that guarantees global linear con- vergence. Unlike prior work that assumes regularity through- out training, we prove that PI-FT inherently maintains these properties. Additionally, we extend our framework to para- metric settings, broadening its applicability. This work advances the theoretical understanding of fine-tuning dif- fusion models and provides foundation for developing more effective fine-tuning algorithms. Our algorithmic de- sign and theoretical findings are also supported by thorough numerical experiments.
In practice, how fast does the PI-FT algorithm converge to the optimal solution?
How does the choice of affect the convergence rate and the quality of the fine-tuned models?
As shown in this section, the PI-FT algorithm converges efficiently to the global optimum; increasing 3 accelerates convergence and yields a model closer to the pre-trained one, aligning with our theoretical analysis in Section 3.
Model Setup. We fine-tune the Stable Diffussin v1.5 (Rombach et al., 2022) for text-to-image generation, using
8

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
Gradient Norm
10
B m 0.01
Score T
1.6
1.7
KL Divergence
0.016
B 0.01
0
B 1.0
ImagerRewar
1.4
1.5
00.010
0.012
= 0.1
B=
10-2
1.3
0.008
0,01
0.006
1.1
8=
1.0
8 8-100
0.002
10-1
200
460
860
1000
200
400
600
860
1000
0.000
200
400
600
800
1000
Sampling Steps
Sampling Steps
Sampling Steps
(a) Gradient norm
(b) ImageReward score
(c) KI. divergence
Figure 2. (a) Gradient (logarithmic scale) PI-FT during training. The
smoothed using the exponential moving average (EMA). A linear convergence rate is observed. (b) ImageReward score of PI-FT during training. Smaller KI. regularization coefficient 3 leads to higher ImageReward score. (c) KL divergence PI-FT during training. Larger regularization coefficient leads to smaller KL divergence and faster convergence.
Color
Composition
Count
Location
Pre- tra-trained
DPOK
PL-FT (Ours)
Figure 3. Visual comparison of images generated by the original Stable Diffusion model (pre-trained), DPOK model, and PI-FT model (ours), Prompts from left right: "A green colored rabbit" (color), "A cat and dog" (composition), "Four wolves in the park" (count), and "A dog on the moon" (location).
0
β 0.1
= 0.01
B B= 1.0
Figure 4. Randomly generated samples from PI-FT model with different KL regularization coefficients. Images from single text prompt: "Four wolves the park". The model with smaller 3 > generates more accurate number wolves.

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
Impact statement
Cen, S., Cheng, C., Chen, Y., Wei, Y., and Chi, Y. Fast global convergence of natural policy gradient methods with entropy regularization. Operations Research, 2021.
This paper presents a technical contribution in the form of a fine-tuning algorithm. We do not identify any specific societal impacts that warrant mention.
Clark, K., Vicol, P., Swersky, K., and Fleet, D. Directly fine-tuning diffusion models on differentiable rewards. In The Twelfth International Conference on Learning Representations, 2024. URL ht ptttts://ee wtttps://oppeeeeee //pp///eee/.......///mme net / forum? idd vmSEVL19f.
Acknowledgements
R.X. is partially supported by the NSF CAREER award DMS-233922 and a JP Morgan Research Award. The work of M.R. is partially supported by Google Research and Meta.
Dai, X., Hou, J., Ma, C.-Y., Tsai, S., Wang, J., Wang, R., Zhang, P., Vandenhende, S., Wang, X., Dubey, A.. et al. Emu: Enhancing image generation models us- ing photogenic needles in haystack. arXiv preprint arXiv:2309. 2023.
References
Agarwal, A., Kakade, S. M., Lee, J., and Mahajan, G. On the theory of policy gradient methods: Optimality, ap- proximation, and distribution shift. Journal of Machine Learning Research, 12:98:1-98:78 2021.
Dhariwal, P. and Nichol, A. Diffusion models beat GANs image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
Alzubaidi, L.. Bai, J., Al-Sabaawi, A., Santamaria, J., Albahri, A. S., Al-dabbagh, B. S. N., Fadhel, M. A., Manoufali, M., Zhang. J., Al- Al-Timemy, A. H., et al. A survey deep learning tools dealing with data scarcity: definitions, challenges, solutions, tips, and applications. Journal of Big Data, 10(1))46, 2023.
Ding, D., Zhang, K., Başar, T., and Jovanovié, M. Nat- ural policy gradient primal-dual method for constrained Markov decision processes. In NeurlPS, 2020.
Doersch, C. Tutorial on variational autoencoders. arXiv
preprint 1606.0:6006055550 2016.
Bansal, M. A., Sharma, D. R., and Kathuria, D. M. A systematic review on data scareity problem deep learn- ing: solution and applications. ACM Computing Surveys (Csur), 54(10s):1-29 2022.
Somingo-Enriehh C., Drozdzal, M., Karrer, B., and Chen, R. T. Adjoint matching: Fine-tuning flow and diffu- sion generative models with memoryless stochastic op- timal control. In International Conference on Learning Representations, 2025. URL https: hett:///ppeeeee /opereviee. net/ / forum? iittrrumid=xxaBee d=xQaRrt.M.....
Berner, J., Richter, I.., and Ullrich, K. An optimal control perspective on diffusion-based generative mod- eling. 2024. ISSN 2835-8856 URL https:// oo
Durgadevi, M. et al. Generative Adversarial Network (GAN); A general review on different variants of GAN and applications. In 2021 6th International Conference on Communication and Electronics Systems (ICCES), pp.
Betker, J., Goh, G., Jing, L., Brooks, T., Wang, J., Li, L., Ouyang, L.,, Zhuang, J., Lee, J., Guo, Y., et al. Improving image generation with better captions. Computer Sci- ence. https://cdn. openai. com/papers/dallle-3, pdf, 2(3):
1-8. IEEE, 2021.
8, 2023.
Fan, Y. and Lee, K. Optimizing DDPM Sampling with Shorteut Fine- Tuning International Conference on Machine Learning, pp. 9623-9639, PMLR, 2023.
Bhandari, J. and Russo, D. On the linear convergence of
policy gradient methods for finite MDPs. In International Conference on Artificial Intelligence and Statistics, pp-
Fan, Y., Watkins, O.,, Du, Y., Liu, H., Ryu, M., Boutilier, C., Abbeel, P., Ghavamzadeh, M., Lee, K., and Lee, K. Rein- forcement learning for fine-tuning text-to-image diffusion models. Advances Neural Information Processing Sys- lems, 36, 2024.
2386-2394. PMLR, 2021.
Bhandari, and Russo, D. Global optimality guarantees for policy gradient methods. Operations Research, 2024.
Black, K., Janner, M., Du. Y., Kostrikov, I., and Levine, S. Training diffusion models with reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024. https: http////ppeeeeeeeee/e/eeee net/forum?? i d=YCW JhGJJFF.
Fatkhullin, I., Barakat, A., Kireeva, A., and He, N. Stochas- tie policy gradient methods: Improved sample complexity for fisher-non-de -eneeeern-deeeeeette policies. In International Con- ference on Machine Learning, pp. 9827-9869. PMLR, 2023.
Bu, J., Mesbahi, A., Fazel, M., and Mesbahi, M. LQR through the lens of first order methods: Discrete-tima case, arXiv preprint 1907.v 00921, 2019.
Fazel, M., Ge, R., Kakade, S., and Mesbahi, M. Global convergence of policy gradient methods for the linear
10

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
quadratic regulator: In International Conference on Ma- chine Learning, pp. 1467-1476. PMLR, 2018.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion prob- abilistic models. In Advances in Neural Information Processing Systems, volume 33, pp. 6840-6851, 2020.
Fetaya, E., Jacobsen, J.-H., Grathwohl, W., and Zemel, R. Understanding the limitations of conditional gener- ative models. In International Conference Learning Representations, 2020. ht.ps:/ http////eeeeeeeeeee/eeeeee//e.ee
Hu, E. J., yelong shen, Wallis, P., Allen-Zhu, Z., Li, Y.,
Wang, S., Wang. L.., and Chen, W. LoRA: Low-rank adap- tation of large language models. In International Confer- ence Learning Representations, 2022. URL https: / //openreview. net/ Torumm?i:=2ZeeeeeFYY
Folland, G. B. Real analysis: modern techniques and their applications, volume 40. John Wiley & Sons, 1999.
Hyvärinen, A. and Dayan, P. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Fu, Z., Yang, Z., and Wang. Z. Single-timescale actor- critic provably finds globally optimal policy. In In- ternational Conference on Learning Representations, httpeewwwwee/ooooo forum?
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent ker- nel: Convergence and generalization neural networks. Advances in neural information processing systems, 31, 2018.
2021. URL
_d=pq22VsxrUmKk
srUVmK.
Gangwal, A., Ansari, A., Ahmad, I,, Azad, A. K., and Sulaiman, W. M. A. Current strategies address data scarcity in artificial intelligence-based drug discovery: A comprehensive review. Computers in Biology and Medicine, 179:108733 108334, 2024.
Jin, C., Yang. Z.., Wang, Z., and Jordan, M. Provably effi- cient reinforcement learning with linear function approxi- mation. In Conference on learning theory, pp. 2137-2143. PMLR, 2020.
Lee, K., Liu, H., Ryu, M., Watkins, O., Du Y., Boutilier, C., Abbeel, P., Ghavamzadeh, M., and Gu, S. S. Aligning text- to-image models using human feedback. arXiv preprint 23Xi::230..12192 2023.
Gao, L., Schulman, J., and Hilton, J. Scaling laws for reward model overoptimization. In International Conference Machine Learning, pp. 10835-10866. PMLR, 2023.
Gao, X., Zha, J., and Zhou, X. Y. Reward-directed score- based diffusion models via q-learning, arXiv preprint arXiv:2409.04832 2024.
Li, G., Wei, Y., Chen, Y., and Chi, Y. Towards non- asymptotic convergence for diffusion-based generative models. In The Twelfth International Conference on Learning Representations, 2024. URL https: openreview net / Pooommidddaae -eteriid-VVVEEESSS
Guo, X., Li, X., and Xu, R. Fast policy learning for linear quadratic regulator with entropy regularization. arXiv preprint arXiv:2311.1416. 2023.
Liu, B., Cai, Q. Yang, Z., and Wang. Z. Neural trust region/proximal policy optimization attains globally opti- mal policy. Advances in Neural Information Processing Systems, 32, 2019.
Hambly, B., Xu, R., and Yang, H. Policy gradient methods for the noisy linear quadratic regulator over a finite hori- zon. SIAM Journal on Control and Optimization, 59(5): 3359-3391, 2021.
Liu, Y., Zhang, K., Başar, T., and Yin, W. An improved analysis of variance-reducele policy gradient and natural policy gradient methods. NeurlPS, 2020.
Han, Y., Razaviyayn, M., and Xu, R. Policy gradient con- verges the globally optimal policy for nearly linear- quadratic regulators. arXiv preprint ar3iv:2303.08431. 2023.
Malik, D.,, Pananjady, A., Bhatia, K., Khamaru, K., Bartlett, P. L., and Wainwright, M. J. Derivative-Fie methods for policy optimization: Guarantees for linear quadratic systems. Journal of Machine Learning Research, 21: 21:1-21:51, 2019.
Han, Y., Razaviyayn, M., and Xu, R. Neural network-based score estimation in diffusion models: Optimization and generalization. In International Conference on Learning Representations, 2024. https:/ oppnrevvie.. forum i d=h
Mohammadi, H., Zare, A., Soltanolkotabi, M., and Jo- vanović, M. Global exponential convergence of gradi- ent methods over the nonconvex landscape of the linear quadratic regulator. In 2019 IEEE 58th Conference on De- cision and Control (CDC), pp. 7474-7479, IEEE, 2019.
Ho, J. and Salimans, T. flassffier-free diffusion guidance. In NewrlPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021. URL https:/ // openrevi ew, met/fooumm" http://forumm..=w@sxxxxxx d=wwwAxx FYbI.
Mondal, W. U. and Aggarwal, V. Improved sample com- plexity analysis of natural policy gradient algorithm with

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
general parameterization for infinite horizon discounted reward markov decision processes. In International Con- ference Artificial Intelligence and Statistics, pp. 3097- 3105. PMLR, 2024.
Szpruch, L.., Treetanthiploet, T., and Zhang, Y. Optimal scheduling of entropy regularizer for continuous-time linear-quadratic reinforcement learning. SIAM Journal on Control and Optimization, 62(1):135-166 2024.
OpenAI. Sora: Creating video from text. https://
Tang, W. Fine-tuning of diffusion models via stochastic con- trol: entropy regularization and beyond. arXiv preprint arXiv:4403.06279 2024.
openai com/ wwm/sora, sora, 2024.
Podell, D., English, Z., Lacey, K., Blattmann, A., Dockhorn, T., Müller, J., Penna, J., and Rombach, R. SDXL: Im- proving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.011 079922 2023.
Uchara, M., Zhao, Y., Biancalani, T., and Levine, S. Un- derstanding reinforcement learning-based fine-tuning of diffusion models: A tutorial and review. arXiv preprint 1rXiv:2407.13734 2024a.
Rahimi, A. and Recht, B. Random features for large-scale kernel machines. Advances in neural information pro- cessing systems, 20, 2007.
Uehara, M., Zhao, Y., Black, K., Hajiramezanali, E., Scalia, G., Diamant, N. L., Tseng, A. M., Levine, S., and Biancalani, T. Feedback efficient online fine-tuning of diffusion models. In Forty-first International Con- ference Machine Learning, 2024b. URL https: / //openrevien ew. net. forum? d=dtVl cyybTT.
Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M. Hicrarchical conticontitional image generation with clip latents. arXiv preprint arXiv:2204. 06/25, 1(2):3, 2022.
Rombach, R., Blattmann, A., Lorenz, D., Esser, and Ommer, B. High-resolution image synthesis with latent difffusion models. In Proceedings of the IEEE/CVF con- ference computer vision and paltern recognition, pp. 10684-10695, 2022.
Uehara, M., Zhao, Y., Hajiramezanali, E., Scalia, G., Eraslan, G., Lal, A., Levine, S., and Biancalani, T. Bridg- ing model-based optimization and generative modeling via conservative fine-tuning of diffusion models. arXiv preprint arXiv:2405. 19673, 2024c.
Ruiz, N., Li, Y., Jampani, V., Pritch, Y., Rubinstein, M., and Aberman, K. Drcambooth: Fine tuming text-to-image diffusion models for subject-driven generation. Pro- ceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 22500-22510, 2023.
Wallace, B., Dang, M., Rafailov, R., Zhou, L., Lou, A., Purushwalkam, S., Ermon, S., Xiong, C., Joty, S., and Naik, N. Diffusion model alignment using direct pref- erence optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8228-88888888 8238, 2024.
Ryu, S. Low-rank adaptation for fast text-to-image diffu- sion fine-tuning, 2023. URL https: //tttbbb///ttb//g//////p com/ cloneofs imo lora.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan. N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. International conference on machine learning, pp. 2256-2265, PMLR, 2015.
Wang. H., Zariphopoulou. T., and Zhou, X. Reinforce- ment learning continuous time and space: A stochastic control approach. The Journal of Machine Learning Re- search, 21((): 8178. 2020a.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Wang. L., Cai, Q.. Yang, Z., and Wang. Z. Neural policy gradient methods: Global optimality and rates of con- vergence. In International Conference on Learning Rep- resentations, 2020b. URL https: http///pppeee..eeeeeeee/ee (dt.orrmmmmid-S0000oosoS
Song. Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score- based generative mod- cling through stochastic differential equations. In In- ternational Conference on Learning Representations, 2021. URL hattts::/oper.eeee wtwpl/ewwwee//ttttueeeeee id-PxTIG12RRUS.
Watson, L., Juergens, D., Bennett, N. R., Trippe, L., Yim, J., Eisenach, H. E., Ahern, W., Borst, A. J., Ragotte, R. J., Milles, L. F., et al. De novo design of protein struc- ture and function with rfdiffusion. Nature, 620(7976): 1089-1 100, 2023.
Steinwart, I. and Christmann, A. Support vector machines. Springer Science & Business Media, 2008
Szpruch, L., Treetanthiploet, T., and Zhang, Y. Exploration- exploitation trade-off for continuous-tin time episodic rein- convex models. arXiv preprint arXiv:2 /2. 10264, 2021.
Xiao, L. On the convergence rates of policy gradient meth- ods. Journal of Machine Learning Research, 23(282): 1-36, 2022.
forcement learning with linear-convex
12

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
Xu, J., Liu, X., Wu, Y., Tong, Y., Li, Q., Ding, M., Tang. J., and Dong, Y. Imagereward: Learning and evaluat- ing human preferences for text-to-image generation. Ad- vances in Neural Information Processing Systems, 36: 15903-15935, 2023.
Xu, T., Yang. Z., Wang, Z., and Liang, Y. Doubly robust off-policy actor-critic: Convergence and optimality. In ICML, 2021.
Yoon, S., Hwang, H., Kwon, D., Noh, Y-K.,, and Park, C. Maximum entropy inverse reinforcement learning of dif- fusion models with energy- models. In The Thirty- eighth Annual Conference on Neural Information Pro- cessing Systems, 2024. http:: / opppreviee..... net net//ooumm2 VoV..uuLYE........
Zhan, W., Cen, S., Huang, B., Chen, Y., Lee, J. D., and Chi, Y. Policy mirror descent for regularized reinforcement learning: A generalized framework with linear conver- gence. SIAM Journal on Optimization, 33(2):00111101,, 2023.
Zhang, J., Koppel, A., Bedi, A. S., Szepesvari, C., and Wang, M. Variational policy gradient method for rein- forcement learning with general utilities. Advances in Neural Information Processing Systems, 33:4572-4583, 2020.
Zhang, Y.. Yang, Z., and Wang. Z. Provably efficient actor-critic for risk-sensitive and robust adversarial RL: A quadariq case. In International Conference on Ar- tificial Intelligence and Statistics, pp. 2764-2772. PMLR, 2021.
Zhao, H., Chen, H., Zhang, J., Yao, D. D., and Tang, W. Scores as actions: a framework of fine-tuning diffusion models by continuous-time reinforcement learning. arXiv preprint +=Xxi:2409.084000 2024a.
Zhao, Y., Uehara, M., Scalia, G., Biancalani, T., Levine, S., and Hajiramezanali, E. Adding conditional control to diffusion models with reinforcement learning. arXiv preprint 4rXiv:2406. 12120, 2024b.
Zhou, M. and Lu, J. Single timescale actor-critic method to solve the linear quadratic regulator with convergence guarantees. Journal of Machine Learning Research, 24 (222):1-344 14,, 2023.
13

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
A. Extension to parametric formulation
In this section, we provide a roadmap to show how linear convergence rate can achieved with the linear parameteriza- tion (25). We make the following realizability assumption.
Assumption A.1 (Realizabilitty. Assume that at each timestamp t, the optimal control satisfies ug (v) = K o()) for some matrix K2 E _pxd
Define the value function associated with the policy K as
((K = EYOND [Vo (Yo)]
where D is the distribution of the initial state and take D (N La), which is the distribution of Y Here the value function is defined as
T-1 Σ 3r
o1)2
V (v) = B
- + Y, = w).
20
t=t
with terminal condition
(x) -r())= = E [R(V)].
Similarly, define the Q-function:
E E (1_1(0 E c+1 (va
(1- a()u) + to,W.) (1-a)) ||u. - s (x) 20,07
u) @
y.
where the expecation taken over W.. We remark that the policy K in the superscript refers to ( (Ke)e> at each t. Recall the choice of Br implies the map QK (y. u) is strongly concave uniformly over all see Theorem 2.8. If the feature mapping Φ is well-behaved, the map K, + QK (y, Kio(y)) is also strongly concave.
Next, we calculate the policy gradient with respect to K. Note that
O.J(K) = 0J(KK θ Ok,
-
-
E
Σ
t-1
Be
(1 ac)2
[Koo(Y)) - sp - (Y))) OK k(o(Y))
0K.
1=
① aK, E [Q (Y,, Kio(Y.)))
(26)
where the expectation taken over the initial state Yo D and noise T- the unique stationary point of (K). Let K be stationary point, i.c., 0.J (K)
T-1
We next show that the optimal policy K
for all In particular, we have
K-K
E
K=K =
aK
The strong concavity in Ky-1 implies that Kū KT-1. Similarly, we have
a E 2Vr (o Ky-00 K= ØKT-2 OKr-2
E
K=K
By strong concavity again, we have Kē K Repeating the argument to conclude that K* = K is the unique
stationary point of the objective function J(K). We remark that the above analysis relies on the two facts:
T 1, the map KT-1 E 1Qr_ (Yr-1, Kr- (0T-)(((--)) has unique stationary point.
(ii) For any T 1, the map K; + E [QK (Y,, (Y)()) has unique stationary point.
14

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
These conditions (i)(ii) relax Bhandari & Russo (2024, Assumption 2.A), which assumes that K₁ E [QK (Y,, [Y))] has no sub-optimal stationary points under any K. While (Bhandari & Russo, 2024) primarily focuses global convergence guarantees for infinite-horizon problems, it also introduces a condition (cf. Condition 4) for finite-horizon MDPs that aligns in spirit with our analysis.
Notably, Bhandari & Russo (2024) did not analyze the convergence rate of policy gradient methods for finite-horizon problems. Since K is the unique stationary point, the policy gradient method converges to the globally optimal solution as long as the standard smoothness condition is satisfied. Moreover, we conjecture a linear convergence rate due the presence of strong concavity. Consider the policy gradient update rule:
Ki k<(++1] k(() + (K(m)) OSIST 1,m 2 0.
For each & T, we have
K k(m+1)
Ki
k( nj I Ki
$ k(()) + ndir (K (m) K(y)) Ki + J. (m)
St -- (.((m))))]
where Kgr == (Kelé and k (Ketese- Eq. (26) implies OKJ(K) 0,J(Kz) is independent of (Keleer The (one-step) strong concavity and smoothness leads
k.(T) + nO. (m) K;,)) Ki| Sc k km(m)
K;]
for some constant 1. If ( is further L Lipschitz in K Rxxd then we have
JaJ(K K - (k((K(()) S T 1
KS F²
and consequently we prove the linear convergence rate of the policy gradient method.
B. Ommitted proofs in Section 2
B.1. Proof Lemma 2.5
Proof. Recall that for any ye E RR."
Pe- (86) = f (1((())) and p (-11)) = f (,14
114550 (v:), TTTa).
with
1 (ye + (1 (v))) and " (v+) =
1
(e(e))
(y+ + (1 of (ve))
Vau
Vae
where f(-1u, Σ) is the Gaussian density with mean μ E Rd and covariance E E dddd Thus, for any yes ye+1 E Rd we have
Pe+ (16) Pet1je (y++1 ll:))
log
11441-14 - (11)) + 207 |1 (2c))1
20?
Denote Epeci. as the expectation under the conditional density Pe+1je(-lue) of Y:41 given Y, = Vt- By definition of the KL divergence, we have
P1+1jt (Y++ (e) p (Y (1)) ]
KL (perteel (11)) = Eprraic
log
1
[[Ye+z (e(()))l
Ep [| ( (27)
E
Ep
+ 20
20
15

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
Note that
Pp--11 [1V - 14. (vi)l = Epttiie [IY+41- - (1))) + E Ep [I1ue(t)) - 14 pre (1=)( (0:)1
+ 2Ppitii [(V++1 - (e(e)) (ame))((eeee (We( - 14 pre (vr))
= Pettae | (v))))) - [Ilre(())) - - 11p" ( )1
where we use the fact that Ef+1je (V++)] = =ve(Ve). Plugging the above equality into (27), we obtain
KI. (P++ijee l1v (u))llpfijee (-114)) 207 (1 - ( 112 (1 ax [lux(ve) - "p.h (v)))1 which completes the proof.
20
B.2. Proof of Theorem 2.8
Proof. We prove Theorem 2.8 by backward induction. At time T, Assumption 2.6 implies V7(v) -r(y) is 10.7 Lipschitz and gradient Lipschitz with LØ,7 = Li and LY, Li-
Step 1: Lipschitz condition of u. Assume that V4+1 is L8,4+ Lipschitz and choice of B3 implies that the mapping
Lipschitz in y E Rd. The
E
Vi+1
(
(y + (1 aa)au) + to.W.)
-B (1 (1-or)² ox)² 2010
||u spth (ا)
(28)
Vaz
(oz -L.it)) > 0 for any E Rd. Hence, there is a unique optimal control ut is -tssogyy concave with 7 (1-aA)² satisfying
25 1t pre (1) (1 (1-o0 (1-a)) 33
vaio?
VV
(/a, (y (1- (v)) + t..
(29)
which is obtained by setting the gradient of mapping (28) as zero for each y E Rdd Here, we apply the Lipschitz condition of V and the dominated convergence theorem to interchange the operators.
We next prove the Lipschitz condition of ut in y E Rdd Note that for any y1 and y2 in Rdd, Eq. (29) implies
Varo? VVi+
(1gg (yn + (1 (n)u + oz o.W.) ]
(11) - uz (12) = 5p (311 - sp (J2) + (1 (1- (1-0x)$
varo? (1 3,
(12 (1 au)u((y2)) + oowk
E VV
Va
Utilizing the Lipschitz condition of s!re and VV we obtain
vano? E v[|vvi (1 (u1 + - +1- ac)u (11))) to,Wi)
|| (91) (12))l2 S [|sp"((()) " + (1 33)
VV 1 (1) (v2 + (1 aa)u + o.Wi)) [
/Tar
S Lū, |1 - 12ll22
Vaiož LViit1
(v + (1 az)u;((1)) - -
1
(y2 + (1- - an)
(1 a,)
v
Var
< Lo |ll|
12ll2 + (1 3x
vano?
LY
( I1 v2ll2 +
1- at
|| (11) uz (va)ll))
Voi
Equivalently, we have
(1- o7 [1it4) B 11141)
L8 + of ofLYit+1
[1((())) - (122)l2 5
|1 -v2lla-
(1 az )3
16

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
Since
B. 11.41 2 A > O, we established the Lipschitz condition of the optimal control ut:
1 1- X1 aq (11) ( 2 Lā. +
11 -02l12 LØ. ll V2l12
where we recall T5 defined in (15).
Step 2: Differentiability of u². We now argue that (v) is differentiable for y IR Let h Rd be an arbitrary non-zero vector and let y E R be fixed. Since sp" is differentiable, we have
pre (y - h) -. (1) = V (v)h ((11)))))))))))))
(30)
Moreover, by the inductive hypothesis, VV is Lipschitz and thus V=V exists for almost all y and is bounded. Folland (1999, Theorem 2.27) implies the mapping = - E [VVZ + o,W,) is differentiable everywhere and its derivative is given by
a E VV (z +2+ o. E (2 (o,
(31)
Ōz
Let
- yau (y (1 av)u and tor (h a4) (u (v h) un (v)). Eq. (31) implies
E VV + + E VV (2 + W:) = E V V (2 w.) o(1
(32)
Since ut is Lipschitz, we have k: (112) as ||1||2 + O. Combining (30) and (32), we obtain
h)) - = sp (y + h) - s (v) + Vaio
1 ( ((y (1x+1)) h) + + au + +,,))
E
VVtt1
(as
(1 a1)3
Vano?
1
(1- (w= +o, ]
E
VVi+i Vas
(w+
vaio?
(1 (h oz) (1 + - uz w(w)) + o(11b1l))
Vs! (y)h+
(1 - 31)3 Heti(u)
where (y) = /2V
1 (u+ (1 - aai) (y)) + o.W.))]
Re-arranging the terms leads to
Vae
(Ia-
ot 3r
Minz(w))
(vs)"e())
(1 - a, Merz(w)) h o((1h12))
of
u ( + h) ui (v)
Ia
which proves that ut (g) is differentiable for any y E Rd and its derivative is given by
o. 3, He -1 (u)
(w) +
of
H Het (w))
(33)
04 3
Step 3: Lipschitz condition of Vut. Next, we show that Vu, is Lipschitz. Taking the derivative of both sides of (29),
Vu) = V V.prre (v) + Vaial E [vv (var (v (1 (1)))())(()) + of W. w:
(34)
(1 3, By
For any 2 E Rdd the integration by parts formula implies that
i E VV (= 0. W.) = EE [V (z + a.W.) VV W.T
(35)
0z
Substituting
(y + (1 )u (v)) and applying the chain rule, we obtain
Var
0 E
VV
1 (
to,W.) (1- (v)) ( +
(y + +
Wi+1(v)44(0),
Oy
17

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
where (y+ = E VV4z (aa (y (1 ai)4 + t, wi_ and (4(y) = _ follows from (34) that
Voi (La + (1-1(11-p))) aa-lv(())))
Vanaz
(y) Vsp"((y
W241 W441
(36)
(1 - )
For any V11 and U2 RR, Eq. (36) implies that
(1) Vui = Vs (122 + vaio, (1 - ox)
(W++1v1))44m)) - (1+1(12)4(x2)))
(37)
Note that for any y E I2d we have
|M++(8))l2 = 1| (3)112 I L7,t,
(38)
[4(())l2
Var [1 + (1 az) L))
(39)
where (38) holds by applying the identity (35), and (39) is consequence of the Lipschitz condition of ut. Moreover, for any V/1 and y/ in IR, the Lipschitz conditions of VV and ut imply that
[11) - W1+1(12)122 -
E [llWelle]
4T, Va: (111 - 32l12 + (1 - ar) ll 82212)
E [1Wae LY,em1 (1+ (1 (1- Varoe
an) on) 118/ - y2lla-
(40)
Furthermore, the definition Ue implies
1 1- [Vn (12(11) Voa
(M((1)) 24
(41)
Combining (37)-(41) together, we have
vano?
(11) (u2(22)ll S s (01) (u2)ll + (1- a1) (1) [[ (142 ( (4
vana?
IM (14((11) - (122)122
(1 3,
Varo?
1- az | Vui(ll)ll
< Li |1 -12ll2
11,111 (1
(1- 0x
vau
1
11 + - a1) Tz).)
E
[1[Wella]
T 11y1 -vll2
a2
Since of
81 1,1+1 5 Ae, we deduce that
E Will Ae) (1+ (1 ai)(o:)) 1191 y2ll2
(2)()1) (u2)22) S A,
Li. +
(1- ac)oa
L lll 8ell2.
where we recall T. defined in (16).
Step 4: Lipschitz conditions of V and VV." Finally, we turn prove the Lipschitz and gradient Lipschitz conditions of Plugging un into the Bellman equation (8), we have
1 (u (1 (v)) o.W: - (1-wx) [1 sp. (v)
(v) E V
-Br 20,
(42)
Vai
18

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
Since 4+1 and sh are differentiable with Lipschitz gradients and of is differentiable, we know that V is differentiable and
0
ox)u o.W.) .B. (1
(1-ar)² cx1) (1u((u) (v))2
(V)
E
Vi+1
(y
dy
20
1
(I + (1 ax) (v)) T E VV
(y +1 (1 (2) + to.Wi)]
Va
Va
2²)²
3,
(Vu, (0) (t)" (ut sp. (v)).
a,o
Define Ge+1 = E VV
1 (y + (1 -n)) + or W. follows that
Vai
Van (la + a (1- (4+1 - 3c (1 - -=))²
(Vui(y) - (v)) (
varoz
VVV(()
(1- a,)3 92-1(1)
+ (1 (4+)
(43)
va:
where (29) applied to obtain the first equality. Next, we use (43) to establish the Lipschitz conditions of V,"" and VV," The Lipschitz condition of VV++1 implies (1441(1)l) & L0.41 for all y Rdd Based (43), we have
1 (1 + a.) Vaa
[V S
L
which proves the Lipschitz condition of Vr, Next, we prove the gradient Lipschitz condition of Vz. For ease of exposition, we denote Sr( = (a. (4a + (1 a, Vs We now use (43) to show the Lipschitz condition of VV,", We first note that (v)) (1) (v) is well-defined at every y E Idd and thus
||C++ (1)ll2 < ( (1))12 S
(1+ (1 L.
Voz
With the Lipschitz condition of G in hand, for any y1 and y2 in IR,, have that
[VVV - = [[S(())) (oll)))Te+1000) (e1) (12) T 9 (v2) 1|
< ((n))))) T (Ce+i(11) - 1 ||1 (Si(v)) - Si(22)) g (4+1(22)
< (((())) ( G4+1 ((S,(v)) - -S(
1 (1 -1- an) (1 ax) 4z) ll 82ll2
a1
at Li l1 Velle
vai
where the last equation uses the Lipschitz conditions of s.re, Vsr V and VV Consequently, we have [V) VV, (12)lla T Illn y2ll2>
where we recall Y defined in (14). other words, Y is indeed the Lipschitz constant of VV.". This completes the proof.
C. Omitted proofs in Section
To prove Theorem 3.1, we need few intermediate results for one-step update rule. We begin with a lemma that characterizes the regularity of control and value functions after one update. Given a function V441 Rd + R. define the following update rule:
(m++) (1) = =x) + (1 )B
ya,o
( ( (u (1 mm))((²())) + ooww))]
E
(44)
19

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
with initialization di (0) (1)) = ot 2 (v). Furthermore, let V((T) be the value function induced by ut (m))
i.e.,
v.(m) (v) E Vet1 1
Vai (V+ (1)) (v)) + o.W.) - -3 (1-(1)) -=a)²
Iu |u(m) (v)
st (v)1²
(45)
2010
Later, we will choose V V+1 follows.
((mut1) to analyze the convergence of Algorithm 1. With (44) and (45), we state the lemma
Lemma C.1 (One-step regularity and universal upper bound). Suppose Assumptions 2.6 and 2.7 hold. Let V be a function THE
that L8,44 -iipssiis and Lf.+41-radien Lipschitz. Consider {
(mm) }
and (v(m) th
defined in (44) and (45). (i) ur ((m) T. -Lipschitz and " -gradient Lipschitz, with coefficients I.g (m) and Tyy satisfying
S m=o
m=u
Fort < T, choose Be such that oi
L LY,411 de > 0. Then it holds for every m 2 that B.
I Lö, and Ly S Li
(46)
(i)) v(m)) is
" -Lipschitz and Ly, ((m)
-gradient Lipschitz, with coefficients ('m)) and Ly satisfying
18,411
0,t
v(m)
1
(1 + oz)
(1+(1- aa) (1 Ae)
(47)
Vot
Vat
Vim)
_ (1 + (1 a2 Lo, (1 + (1 at) 1,c--
1 1- Li Li
1,t
at
Var
1
1 1- at L (m 2) (1 (1 - a.)
E [[
P Lo.r+1 - A,)
(48)
at
a,or
Lemma C.1 provides regularity properties of the sequence of controls and value functions generated by (44) and (45) throughout the optimization process. Specifically, the Lipschitz and gradient Lipschitz constants of u, (m) are bounded by those of the optimal control ui; see (46). Furthermore, when V Vi+1 Eqs. (47) and (48) lead
((m) O.t LK: + O(1 >,) and LY, LY, O(m 2)(1 )))
The residual terms diminish to m - oo and thus providing the convergence of Lipschitz constants. Consequently, selecting Ve+1 as v((mez)) when feasible, directly ensures the regularity the control and value functions defined in Algorithm 1.
We outline brief proof sketch to highlight the ideas before providing the detailed proof. The Lipschitz condition of u, (m)) derived through direct calculation, while the gradient Lipschitz condition leverages the integration by parts formula (50). Additionally, establishing the regularity of (m)) requires careful analysis of the gradient expression, along with tightly controlled upper bound; see (57).
Proof of Lemma C.I. Step 1: Lipschitz conditions u, m! and Vul") We begin with the Lipschitz condition of (!!") For any V1 and in Id, the update rule (44) implies
1m+ (yn) ub (m+1) (v=))l I ui
LØ.e llv1 - 82ll2
varo? )p
LP
1,411
1
||
-22ll2
+
1 1 at
(m) ue (y1) -. _ (w2))1)
Vot
Vai
Since 1- 22 P Fe 1,441 2 A, unrolling the recursion leads to
L0,
(m+1)
of LØ, + (1
[1-a.)3
1,141
p
of
1.141
L
(m)
20

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
1 At
S A
Lü, +
LU.e-
of 11,41 X, is important it decouples the Pe
Here, we also use the fact that O. L Note that the condition
upper bound, the sense that the last line does not depend on LY,+1 directly. Next, since "N
8t differentiable, a simple inductive argument shows that Vul""! (Tr) is well-defined. Furthermore, we show that Vu, (m) is Lipschitz. Differentiate both sides of the update rule (44),
Volo? Vul (m+mm)))) (y) = (sp + (1- a,) By Vieti (vaa Tas
y + (1- (v)) + o.W..)
(49)
For any 2 E Rd, the integration by parts formula implies that
0 E
W,
VV++1 (z + o, IE [v (z + t.o = Vin (z o,
(50)
Oz
Substituting =
-
(y (1 an)ut (v)) and applying the chain rule, we have
Vas
0 Oy
Vi++
y + (1 - an)u (1)) + to.w.)
Vat
=
E
Vi+1
1
(u+ + (1 a + o.W,
oz (Vor
Ia + (1 Va. (m)))
(51)
Va+
Define
Tlml (g) = Vi++1
1
4 (1-ar)
ar) (8)) + ocWi)
V
Wiml (w) = Vi++
1
y + (1 oc)u (w)) + to.
W.
Vat
at
(mm) (v) = 1 (Ta + (1 V. (v)
Vot
It follows from (51) that
((m) (v) = Wi(m) (u))4( (m) (z). Consequently, Eq. (49) becomes By
0
(1n+1) (1) = Vopre (y) + varo W (y) (n).
Vug
(52)
(1 (1-ax) ox)
Since Vi+i and un for any and y2 in Rd, we have
(m) are both Lipschitz, we know that V (m+1) I) is Lipschitz long Vu," (mm) is Lipschitz. Specifically,
Vu (m)
(11) (21)
Vul (m( (y2) V.S.
srre (11)
Vspre
pree (Uz) + varo?
(1 a1 (Wiml (m)) (y1) whlr) (uz)(1
Note that for any y Rdd we have
Wiml (3)
|| $ L LP,
(53)
| w0)
|| Vai
(1+ o.)
(54)
where the (53) holds by applying the identity (50), and (54) is a consequence of the Lipschitz condition of u. for any y1 and y/ in Idd, the Lipschitz conditions of Vi and (m))) imply that
(m) Moreover,
Wiil (v1) - WiT (22)
E [| [2]
L
(I11
wella + (1 - -.) ((())
1lv1 3/2 l12
1
Vate
21

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
E 1
P (1 (1 an))..)) 11 - 82l12.
(55)
and the Lipschitz condition of Vu, ""( leads to
(mr)
(m (y1) 2( (v2)
1 1- -0 Vai
Vul (mm) (v1) -vu," (v2)
1
u/m) 11 32112-
(56)
Vas
Combine (52)-(56) together to have
Vul fm+11 (21) - Vu, (m+1) (12))
S Li,e I1 -22l12 + (1 (1-a4)33
Voyo? (1P - vai 111
1 1- at pulm)
1 (1 (1- - -os)
2B ([[[Wall2]]
I 11fez2) 11 92112
at
0z
ci L. F3 1,441 S 1 Ae. we deduce that
Since
E Well
" 2e)-
11,
S
1iq,e
+
+
(1 (--a) vara
(1+ (1())aa)z))))
Equivalently, we have
(m)
S 25
Li +
E [lllaalbl (1- de) (1+(1-o))ii)) = Li (1 /400
1.
Step 2: Lipschitz conditions of v.(m) and vy!mm With the Lipschitz conditions of un establish the regularity of v.(m)_ Note that the expressions in (44), (45) and (52) imply that
(m) and Vu (m)
we now turn to
vv (g)
1 [Ia + (1 (1-au ax) (1-oz)
Vu( (1m)(()) (v) - 3r (1 - (1-a (Vulm))
(Vut (()) (y)
Varož glm (w)
((m="(()))
(1-02)3 oz)
Vou
(1 -
1
s (p (v)) Giml (v) + (1 (1-ar) (Vul (mm) (v) - (v)-vop(()))) T (G(m)) (v) -1m1-1)(10) (57)
(Ia + (1 - 04) V
Vai
Vai
We will use the above expression (57) to derive the Lipschitz and gradient Lipschitz conditions of (m)) To proceed, we first get a few useful estimates. Recall the notation S,(y) Jrs (La + (1 ax) Vs. (y)). We first show that S,(y) T ( (v) is bounded and Lipschitz. Note that for any y E Rdd
< ||(()))l2
( (v)
1
(1+ + (1 - az) L8.
(58)
Va:
Also, for any y1 and y2 in Rd, we have
(((()))) T ( (ux) (12) T (w))
< ((())) T (g.m) (vr)
mlml (12)
|1 + +[((5.(on) (2 T G.ml
[((((1)) (121)))2 112 1 (11) 3. (vz)
([8((11) So(22)ll
@ (uz)
1 (1 + (1 (1-ay) (1-aa) (1 + (1 - I 11 -22l2
1 ai I To,.1 11
12ll2.
(59)
Vai
22

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
where we apply the Lipschitz conditions of Pe+1, Vi++1, and sf in the assumptions and the Lipschitz condition of ut in Step 1. remains to show the second term in (57) is bounded and Lipschitz uniformly over m 2 0. Since Vi2+1 is L. 1,441 - Lipschitz, we deduce that
((m
T (v) -0-
S
11,11
1- 0tz
(v) - U
(v)
(
(y) 2
2
Vai
< L 1,++1
of
| (v) - I (1) 2
(60)
Br
where we apply the update rule (44) in the last equation. As
oi
11,41 S 1 - X, unrolling the recursion (60) leads to 3.
(e) 3 - - (1 - - 1 (v) gitt (v) /
< (1 - >m _ 1 1-a
|1/0 (1) (0)
Var
of 1
S (1 - -
1r p
11,
1 S(1 A1) mh P
L0,1+1-
(61)
Br
Moreover, for any y E Rdd, we have
(gim) ( Öy
(v) - 3. (w) || = (w)4, ( (v) - Winl (w)24 (w)2(" m0 (1))
wirl(() - W - (w) 2
+ Wint (w0l / (m)
(v) - -4," tm
11
(w)
(62)
The gradient expression (52) implies
| (v) 24, (m- 1) (v)
1/
1 1-ar
ai
Vus
(mm) (v) (v) Vul mmm (2)
of 0
(glm-")(s) (v) - -.: (2))((() (w)
(63)
2
Vat
Be
By
and (60)-(61) lead to the fact that
(w) - Winh 1) (w) |1. S
Well2]
1,1+1
1 1- at
(mm) (w) - ur (1)m))(())
(w)
5
E Wella
40,0+111 (1 - Am)" (64)
vot.
at
With (53), (54), (63) and (64), we can bound (62) with
i
( (v) "it""(())
18.00 E [1] Wellal
(1 (1 -o+) (1 (1-)))m
5
oy
Vaiae
of
0 (ōtm-")(y)
(w) - -0mm==1 (wm)
+ 3e 1,441
(i+11
(65)
Oy
Since of L 8, 1,041 S 1 - >s, unrolling the recursion (65) to have
a Dy
(Q.m) (v) Glm (w)) (w))
S
L
0 1 E Wleeelal
(1 (1+(1_0)(0))) mm1)=mm111--0 21)
a.)
vaioe
0 g:0 (v)
+ (1 - (1-),) (m))
By
8 E WllWellal
(1+ (1 a.) (m + 11(1 A,)"
(66)
Valge
where we utilize the fact that
|
a 5:0
(v)
1/2
14 (w))
<
L L0,41 EE (1WWlll)
(1+ (1 (1+(1-ax) -a=)
2
Va,o,
23

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence Hence, for any 1/1 and 1/2 in R (53)-(56), (61) and (66) imply that
Wit 2+1 (yn y (m T
@ (wn) - Winl (12) yu (m)) (uz) |12 9 (v1)
(Wirr (W (v2))
T
( (w2) - -9:41 (122)
Winl (w) )21 (m)) (v1)
04t (01)
|
Win (12)24 mm) (w2) |/2 sup "ER"
8 (c (v) -4 (w))
+
l1 321l2
Dy
1
at
Li (m)
+ (1 (1 (1++1-as)) (1- 2 1lW
40,41 (1 - A0) 1191 - -22112
Voi
+ (1 an.) 0.0. BB [[[lWalll] 11,41
+
p 10,c41(m + 1)(1 - "))"lly
1,) Ill1 -22ll2
0400
att Li + (m + 2) (1 + (1 - as )8)
2 E 11
LP, 4...t148.4114 - Aa) 11 - 92l12.
(67)
/
With there estimates in hand, are ready to prove the Lipschitz and gradient Lipschitz conditions of v.(m). Note that
(1 1-0x mz)) (m/) (v)
of wilm) (m)) (v).
(Vul
V (v)
3e 1+1 (w)
Vai
Consequently, for any E R., we deduce from (53), (54), (58) and (61) that
S (.(w) 3. (v)
||
+
Wiil
(1)
| (v) -
Be
1
(1 + (1 - az) 16.1) 00.411
of I.P
11,1411
(1+ (1 (1 mm 00..1
Vat
Be
Vas
1
L8.+1
(1 (1
(1+ (1 (1+ (1- m
Von
Finally, for any 81 and y2 in Rd it follows from (59) and (67) that
Vv, (m) (v)) vv, (m) (12)
2
- ((lr)) 1 (.))
(wn) - S 3l (w2)
+ of | (11)) T (Gim) (v1) -1 (v1))
6,
(Wim) (12)) ( (wz) T
(iml)(a)) (112) - -9/-1-11 (12)) 2
1-at Li. (.
-1 (1 + (1
+
|191 122l22
Vox:
0 1
1 a1 Li + (m+ + 2) (1 (1-11)))83
2 E [[]Well2]
P 18 - A1)ah 11 222
Br
1,841
Var
(1 (1- (1-o)) oz) (1 + (1 a1)
1 at -.i. V
10.411
Vai
+
(
(Tai
a4 L + (m
(m+ +2 (1+ ++((1 (1- av)Lo.))
2 E [[Wa
)
8 )dd.t1(1 (1-20 2)) 1 11 82ll2.
aroi
which finishes the proof.
Our next result characterizes the error of the control sequence obtained from the one-step update rule (44).
24

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
Lemma C.2 (One-step error analysis controls). Assume the same assumptions as in Lemma C.I. Moreover, suppose there is constant E+41 such that for all y that COPrii-VViLl (v) (w) S Ex41-
Let us (m)) and v/(n) be defined as (44) and (45), respectively. Then holds that for all E Rdd that
(w) (2) - u7 (000\l1 - (1- + X Eit1)) Var(1 - Ax)
(1- al)
Proof. Define the Bellman optimality operator T time as
varo? E
( (y (1 (1-aaa)u)) + +o.
(T) = sp (1)
VVi+i (Var
Moreover, define the approximate Bellman operator Te as
(Tr))() (1) (1 33
Vagoz E
Vizti
(1 (1- ac)u(y)) +
The update rule (44) implies ut
(m)+1) T (m)) Together with the optimality condition on Tuu we obtain that
(m+1) 1) (y) u7 uu(v)]
= (Feulm)) (1)) - (Tr
- (Fry (m)) (v) - -TTal (w) 2 + +T((Ta (()) (TTu))
(68)
Note that T, is contraction operator. Indeed, for any u, Rdd - Rdd we have
varo?
1
(y + (1-11a (1-a))(vv)) + o.W.)
(T) - (Tro) S -1- (1-o+)3
VVž+i
Vove
1
- -V
(y + (1 ar)x(()) + w.
Vo
2
of $ L [ [lu(y) v(y)ll2
B.
$ (1 Xt) [u(()) v(y)l12
(69)
where we apply the Lipschitz condition of VV and the choice of B.. Consequently, the second term in (68) is bounded with
(Teuh ( (v) (Te | (1 - (1> 1u (m)(())
(v) un (v) (10)11
(70)
Next, define
g (v) = E VV ( (y (1 (1-)aam) ac)u((m)))))
+ o...i)
and recall that
9L (v) = E ViCH1
1
( (B
(1 ( ar)ut" (a))a (v)) ar))((")()) oo..))
Vac
For the first term in (68), note that
Varo?
((Faa(m)) (b)) (Teu (
(m))) (v) (m) (v)
(1
or )
25

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
Valo? E (1 -+)
- vvci.) (rmm)(/)))) S valoz Ex+1,
(71)
(1 02
where s/ ~ ~N (vac y + (1 - (n))),(()))) of. la?
With (70) and (71), we can further bound (68) with
Varo?
[ue (Mim+()) (w) un (w) -w)
Bett + (1 - 200 - ui -
< (1 03,
(72)
Unrolling the recursion (72), we have
ui (y) - uz (v))
51 - 511-A1)"
| [v(0(v))
|u
ui ui(w)//
+ +1i² varo
(1 Ox Ett1
I ((1 - X₂) Lbiiii + A7 Fet1)
varo
(1 -
Var(1 de))
1 (11 Xo) (o + + Besl)
where we apply the condition that 1 o2 D. 1,++1 2 d This finishes the proof.
The last lemma in this section is on the error analysis of VV Lemma C.3 (One-step error analysis on the gradient of value function), Assume the same assumptions as in Lemma C.2. Then it holds that for all y E Rd
( (v) - vVr S C1,EE+1 + (2,(11-Ar
(73)
where
C₁.=== 1 1+ (1-11a0)) Loe and Cz = (1 (1 - o.)
10,0+1 + (1 + (1 av)) T.
LV L0
devar
Vot.
Van
Proof. Recall the definition that G(rl (v) = E (V++1 (1 (aa
y + (1 ar) + o..: It follows from (43) and
(57) that
VV/ ( (V) = 1
(Ia + (1 -1-00 (1, Vo. 175p0000 (2)) -
Vat
az)) vai
(Vu! ( (y) vop"r(()) (O.Ti( (v) "(
(74)
Observe that for any m 2 0, it holds that
| (w) Get
(w)) 5 (0) - (w)) + -| (v) ( Get (w)l
(75)
where we recall (lml( (v) E. (V Vas (u + (1 (1- (w)) of (75), we notice the following inequality holds:
+ w.) (w))
For the first term on the right hand side
| (v) s (w)
s E (1 (efx)>(///)) (/)1/]] s E
(76)
where y' N/ (N(aa, (y (1 (m) oolla ) For the second term in (75), the Lipschitz continuity of VV++ implies
| (v) - -c+1 < LL -
az
ut
(m)
(3)
u, (v)
(77)
26

---

Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
Consequently, Eq. (75) becomes
o (v) G S Ec+1 + 11,
1- at
[ui (v) - mz 2u00
(78)
var
With (53), (54), (61) and (78), we further bound (74) with
(v) si (BerI + L1,411
1 at
fm))
- ui( m((w))b))) + C2 d1)
Voi
where the coefficients are given by
1 (1+ and C22
1+ (1 - -ar 18 100.11.
Vai
Finally, we apply Lemma C.2 to obtain
(or1+ o.
v (v) vv
< C₁, E + 11,41 B.
(1 "", + X5 Bett)) + C2, 21)
1 (Eet1 + (1=) (1 (1- 1o) L0,44 + No =(t1)))) C2,(1 - 2e
B+t1 C22(11 - -
where
C2,= =, X, and C22 C22, + G
Therefore, we finish the proof.
Now we are ready to prove Theorem 3.1 with the results in Lemmas C.I, C.2 and C.3.
Proof of Theorem 3.1. It is straightforward to check that LY, 2 LYi for all S T. Thus, the choice of 3 guarantees that Theorem 2.8 holds. To prove the theorem, we begin with the error vv) (mm)) (v) - vV) Let < T be
fixed. To apply Lemma C.3, we choose V. (mm) for all t' > t. Unrolling recursion (73) and use the fact that ((( (y) r(y) (y), we obtain that
T < II
T-1 Σ k=t
(m)) - -V, (v)
Ci.
-v7)))
+
[ Es
C2 - m
k=t
C2 - m
E=t
Next, we apply Lemma C.2 to have
uu (y)
(me)) ui (1(00/1 I (11 X1) 10 X; V Me+ (v) (VZ)())1l>)
var(1 - 21)
(1 L.)
T-1 Σ
vai(l - Ar)
< (1- 21 + x
II C2e/
C2 - (2
a1)
kmt41
Ammt+1
which finishes the proof.
27


Reviews:
[{"review_id": "qm6CfvWtqZ", "reviewer": "Reviewer_Yp8J", "summary": "The authors propose a discrete-time stochastic optimal control framework with linear dynamics and KL regularization to model the problem of fine-tuning of diffusion models. They analyze well-posedness and regularity of the control formulation, propose a novel algorithmic scheme based on policy iteration, and they analyze it showing linear convergence. Ultimately they introduce a parametric extension and analyze the convergence of an associated policy gradient method.\n\n## update after rebuttal\nThe rebuttal and further discussion with the authors tackled properly some of my concerns and therefore I have increased the score from 2 to 3. Nonetheless, I still believe that some of the weaknesses mentioned in my review and further comments do still hold. In particular, I am not convinced by the authors justification for the current experimental evaluation proposed within the rebuttal. Although it certainly makes the work more complete than before, it arguably does not present a clear comparison with common methods in this space and might lead to a distorted/limited viewpoint/evaluation of the presented methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed feedback. We are glad to receive your positive feedback on our theoretical contributions and that you found our placement of the work clear and honest. Below please find our response to your questions.\n\n```\nIt is not clear to me if the policy iteration algorithm presented is supposed to work only on discrete-spaces, or, more precisely, how the control $u$ is represented within that algorithm. What settings does this algorithm actually cover? I am asking this since Sec. 4 seems to extend the control representation to be parametrized, but in doing so it seems to propose another algorithm. Did I misinterpret something?\n```\nWhile we consider a **discrete-time** model, the PI-FT algorithm is formulated on continuous state and action spaces, with both state and control taking values in $\\mathbb{R}^d$. This setting captures the essence of fine-tuning and indeed disparts us from the usual RL literature. Specifically, we focus on Markovian control policy $u_t: \\mathbb{R}^d \\to \\mathbb{R}^d$, and our algorithm iteratively updates $u_t(y)$ for each $y \\in \\mathbb{R}^d$. Section 4 introduced a linear parametrization of the control policy to enable an efficient implementation in practice. This parameterized setting should be viewed as a special version of the PI-FT algorithm, rather than a separate algorithm. Moreover, we believe the theoretical results regarding regularity and convergence established in Section 3 apply to this parameterized setting and a roadmap has been provided in Section 4. Please let us know if this explanation clarifies your concerns.\n\n```\nSec. 4 seems to propose an algorithm for continuous spaces where the control is represented via a linear parametrization. While a linear approximation would make sense to represent the forward process drift, this seems an excessive approximation for the reverse process control, which should have an order of complexity comparable to a score network. It seems to me that this object would be non-linear even in the easiest examples possible. Hence the question: are there some relevant examples (even simple) where this assumption holds?\n```\nWe refer to the parameterization in Section 4 as linear in the sense that the control policy is a linear combination, through parameter $K_t$, of (nonlinear) basis functions. Hence the resulting control can be possibly nonlinear in the state variable $y$. This is a flexible, powerful, and tractable parameterization often used in control and RL, where depending on the basis, it can have high expressivity. In particular, our parameterization includes a wide class of score approximators such as random features, kernel methods and even overparameterized neural networks under the NTK regime. Hence, despite being linear in parameter $K_t$, the policy class can be sufficiently rich to approximate highly nonlinear score functions. Moreover, our experimental results (which we discuss below) suggest our principled framework remains valid for generic neural networks. We will clarify these points in our revision.\n\n```\nAt times it is very unclear whether the algorithms are introduced within this work (e.g., the policy gradient one), or if this work aims to only perform theoretical analysis on them. Please clarify.\n```\nApologies for the confusion and thank you for raising this point. To clarify, both the policy iteration method (PI-FT) and its policy gradient variant are new and proposed in this work.  Existing methods such as DPOK and DDPO rely on generic RL developments such as PPO or REINFORCE and fail to fully leverage the fine-tuning structure. On the contrary, the specific setting considered, with linear dynamics and entropy-regularization, is tailored towards the development of efficient fine-tuning diffusion models. For this reason, the PI-FT algorithm and its parametric extension directly computes the policy gradient of a KL-regularized control objective. This principled design leads to a more efficient implementation in practice compared to prior works; see our experiments below. We will clarify all these points in our final revision. In particular, we will specifically highlight that both algorithms are proposed in this work and are our contribution. \n\n```\nNot presenting experimental evaluations renders impossible to evaluate the practical relevance of proposed algorithmic ideas.\n```\nWe included a new set of thorough experiments in this rebuttal. Please see link https://anonymous.4open.science/r/ICML-rebuttal-741C/figures.pdf for the results and see response to Reviewer 1H7D for detailed explanations on the experiments.\n\n```\nLine 434, 'replies' should be 'relies', I guess.\nLine 397, 'the' is repeated twice.\n```\nThank you for pointing out these typos. We will fix them in the revised manuscript.\n\n\nFinally, we would like to thank you for your constructive feedback. We hope our response answers your questions. If so, we would greatly appreciate it if you could reflect it in your evaluation of our paper.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "The paper proposes arguably novel Policy Iteration and Policy Gradient schemes for fine-tuning of diffusion models but does not perform any experiment (e.g., comparing with existing methods). This is not inherently problematic, as theoretically relevant papers without any experiment are fine, but as later pointed out, I have some doubts regarding the relevance of the presented theoretical results. And not presenting experimental validations of the presented algorithm fundamentally renders impossible to evaluate and/or them besides the theoretical results.", "theoretical_claims": "Only limited parts of Lemmas B.2 and B.3.", "experimental_designs_or_analyses": "There are no experiments.", "relation_to_broader_scientific_literature": "- To the best of my knowledge most works within this literature (of diffusion models fine-tuning) are fundamentally experimental. Overall I believe the authors of this work reported a very clear and honest representation of the current landscape of existing works. I find particularly interesting the focus on discrete-time and linear dynamics, which seems well motivated and renders it possible to connect the problem to common formulations in control theory. Having said so, I have doubts regarding the relevance of the theoretical results derived expressed within the weaknesses section.", "essential_references_not_discussed": "The work seems to properly mention all relevant references.", "other_strengths_and_weaknesses": "STRENGTHS:\n1. The paper focuses on a very timely and relevant problem.\n2. The authors analyze well-posedness and regularity properties of a practically relevant control problem, which I regard as a positive contribution.\n3. The authors present an arguably new algorithms based on policy iteration and policy gradient for fine-tuning of diffusion models and an extension to parametric controls.\n4. Formal derivations showing linear convergence guarantees of such algorithms are presented. \n\nWEAKENESSES:\nUnfortunately, I found the paper quite unclear in being able to distinguish between the policy iteration and policy gradient schemes presented, the limitations of each, and whether some limitations are purely of the theoretical analysis, or are actual algorithmic limitations. I list several concerns and doubts in the following.\n1. It is not clear to me if the policy iteration algorithm presented is supposed to work only on discrete-spaces, or, more precisely, how the control $u$ is represented within that algorithm. What settings does this algorithm actually cover? I am asking this since Sec. 4 seems to extend the control representation to be parametrized, but in doing so it seems to propose another algorithm. Did I misinterpret something? \n2. Sec. 4 seems to propose an algorithm for continuous spaces where the control is represented via a linear parametrization. While a linear approximation would make sense to represent the forward process drift, this seems an excessive approximation for the reverse process control, which should have an order of complexity comparable to a score network. It seems to me that this object would be non-linear even in the easiest examples possible. Hence the question: are there some relevant examples (even simple) where this assumption holds?\n3. At times it is very unclear whether the algorithms are introduced within this work (e.g., the policy gradient one), or if this work aims to only perform theoretical analysis on them. Please clarify.\n\n\nNot presenting experimental evaluations renders impossible to evaluate the practical relevance of proposed algorithmic ideas. Hence I have to evaluate only the theoretical contributions.", "other_comments_or_suggestions": "1. Line 434, 'replies' should be 'relies', I guess.\n2. Line 397, 'the' is repeated twice.", "questions_for_authors": "Please answer clearly regarding all 3 points mentioned within the weaknesses section. I am very open to changing my vote/decision based on these answers.", "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": "iOnLjx8xrA", "reviewer": "Reviewer_7pRd", "summary": "This paper proposes a stochastic control framework for fine-tuning diffusion models. The key contribution is establishing theoretical properties such as well-posedness, regularity, and linear convergence of a proposed policy iteration algorithm.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback. We are delighted that you found our theoretical results rigorous and our presentation clear. Below is our point-to-point response to your comments:\n```\nProper evaluation is missing: Algorithm 1 is defined and justified clearly but is purely theoretical; however, the practicality of the method is questionable…\n```\n```\nA significant weakness is the absence of any empirical validation…\n```\n```\nAs mentioned, Algorithm 1 appears overly idealized, and its practical utility in realistic scenarios remains highly questionable. In practical high-dimensional environment, taking expectations to obtain accurate value function can have difficulty…\n```\nWe included a new set of thorough experiments in this rebuttal. Please see https://anonymous.4open.science/r/ICML-rebuttal-741C/figures.pdf for the results and see response to Reviewer 1H7D for detailed explanations on the experiments.\n```\nI find following the flow of this work pretty easy and smooth. One suggestion would be: all theoretical results before Thm 2.8 are standard in the literature, that is to say, all contents up until Page 5. Results such as Lemma 2.5 and Eq. 8 (many remarks are also well-known) have long been discovered, see [1], [2]. These standard results take up too much space in the main text. I would suggest improving the presentation, for example, simply referring to these works instead of repeating all the details in the main text. Otherwise, it is also possible to only present the most important parts in the main text but defer burdensome calculation details and remarks to the appendix in the next version.\n```\nThank you for your constructive feedback. To the best of our knowledge, we are the first to use KL divergence over transition dynamics (on path space) to control the deviation of the fine-tuned model from the pre-trained model, whereas formulations in the literature consider KL between the terminal state distributions. For this reason, the results before Theorem 2.8 are novel and valuable to readers although not technically significant. Therefore, we decide to follow your suggestion and defer calculation details and remarks to the appendix in the revised manuscript.\n```\n…since this work is purely theoretical and does not have any kind of simulation studies. It is advised to spend more efforts in demonstrating the technical hardness of Thm 2.8 and Thm 3.1. It is confusing to me what the major difficulty is in presenting these results. It feels like Thm 2.8 is a direct calculation based on the Lipschitzness assumptions, which might not present enough technical novelty. I might be wrong, but can the authors comment on this?\n```\nWe appreciate the opportunity to clarify the technical novelty of Theorem 2.8 and Theorem 3.1. \n\nRegarding Theorem 2.8, the core challenge stems from the fact that Eq. (17) is an implicit equation, in which both sides involve the optimal control $u_t^*$. This prevents a direct application of Lipschitz assumptions on model parameters. In contrast,  we prove the Lipschitz property of $u_t^*$ through the detailed and nontrivial derivations in Line 726 - 736, where the choice of $\\beta_t$ ensures the invertibility of the coefficient. Additionally, the differentiability of $u_t^*$ relies on the Gaussian smoothing effect; see Eq. (30) – (32), which is novel analysis to the RL literature. The Lipschitz condition of $\\nabla u_t^*$ is especially non-trivial as it involves $\\nabla^2 V_{t+1}^*$, which may not be Lipschitz. We address this using the integration by parts formula in Eq. (34), which is a key technical step. We will discuss these in the manuscript.\n\nIn terms of Theorem 3.1, the convergence analysis of the PI-FT algorithm relies on **maintaining the regularity** of the value function throughout training, which is highly nontrivial. For example, prior work such as [3] proves convergence under the assumption that the value function remains regular during training (which is non-tractable); see Assumption 2 in [3]. In contrast, our result establishes the desirable regularity through a connection between $V_t^{(m)}$ and the optimal value function $V_t^*$. \n\nThanks to the reviewer’s question, we will revise the manuscript to better emphasize the technical challenges and contributions underlying these results.\n\n\\\nThank you for your time and for your valuable feedback. Your suggestions definitely helped us improve the quality of this manuscript and highlight our contributions. We hope our response addresses your questions and clarifies the challenges in our work. We also hope that the effort we invested in conducting thorough numerical experiments demonstrates the promising practical performance of our proposed algorithm. If so, we would be grateful if you could reflect it in your evaluation of our paper.\n\n[3] Zhou, Mo, and Jianfeng Lu. \"A policy gradient framework for stochastic optimal control problems with global convergence guarantee.\" arXiv preprint arXiv:2302.05816 (2023).", "claims_and_evidence": "The paper makes strong theoretical claims about the optimality, regularity, and linear convergence of their proposed algorithm. The theoretical analysis is mathematically rigorous with assumptions clearly presented.", "methods_and_evaluation_criteria": "Proper evaluation is missing: Algorithm 1 is defined and justified clearly but is purely theoretical; however, the practicality of the method is questionable. The algorithm’s iterative nature and associated computational overhead might limit its applicability in realistic high-dimensional problems, but no practical evaluation or even basic numerical simulations have been conducted.", "theoretical_claims": "I have roughly examined the theoretical analysis, particularly the convergence properties presented in Theorem 3.1 and Theorem 2.8. The proofs appear correct. (On the other hand, all results before Thm 2.8 are standard and well-known in the literature. I have concerns of ssuch presentation, see below)", "experimental_designs_or_analyses": "A significant weakness is the absence of any empirical validation. The paper does not include simulations or practical experiments demonstrating the behavior of Algorithm 1 or verifying the theoretical claims. In diffusion model fine-tuning, it is a common practice to test it on text-to-image models such as StableDiffusion v1.5. Pipelines for such implementations are also accessible and easy to work with. Can the authors provide clarifications why it is not conducted?", "relation_to_broader_scientific_literature": "The paper accurately positions itself within existing literature concerning fine-tuning diffusion models using reinforcement learning and stochastic control", "essential_references_not_discussed": "Reference discussion is proper and abundant to me.", "other_strengths_and_weaknesses": "As mentioned, Algorithm 1 appears overly idealized, and its practical utility in realistic scenarios remains highly questionable. In practical high-dimensional environment, taking expectations to obtain accurate value function can have difficulty. But the idealized algorithm overlooks such hardness.", "other_comments_or_suggestions": "I find following the flow of this work pretty easy and smooth. One suggestion would be: all theoretical results before Thm 2.8 are standard  in the literature, that is to say, all contents up until Page 5. Results such as Lemma 2.5 and Eq. 8 (many remarks are also well-known) have long been discovered, see [1], [2]. These standard results take up too much space in the main text. I would suggest improving the presentation, for example, simply referring to these works instead of repeating all the details in the main text. Otherwise, it is also possible to only present the most important parts in the main text but defer burdensome calculation details and remarks to the appendix in the next version.\n\nMore importantly, since this work is purely theoretical and does not have any kind of simulation studies. It is advised to spend more efforts in demonstrating the technical hardness of Thm 2.8 and Thm 3.1. It is confusing to me what the major difficulty is in presenting these results. It feels like Thm 2.8 is a direct calculation based on the Lipschitzness assumptions, which might not present enough technical novelty. I might be wrong, but can the authors comment on this?\n\n\n[1] https://arxiv.org/pdf/2402.15194\n[2] https://arxiv.org/abs/2403.06279", "questions_for_authors": "See above", "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": {"ethics_expertise_needed": ["Legal Compliance (e.g., GDPR, copyright, terms of use)"]}}, {"review_id": "4pewMwBchM", "reviewer": "Reviewer_1H7D", "summary": "This paper proposes a discrete-time stochastic control framework with linear dynamics and KL regularization for fine-tuning diffusion models. It establishes well-posedness, proves the regularity of the optimal value function, and develops a policy iteration algorithm (PI-FT) with guaranteed regularity and stability. The framework extends to parametric settings for efficient high-dimensional implementation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. We are glad that you found our theoretical results strong and rigorous. Below please find our response to your questions.\n```\nI have to say that the theory in this paper is very solid. However, the lack of empirical results is also the biggest drawback of this paper.\n```\nThank you for your suggestion. To demonstrate the practical efficiency and effectiveness of the proposed PI-FT method, we conduct thorough numerical experiments, which are included in https://anonymous.4open.science/r/ICML-rebuttal-741C/figures.pdf. \n\n**Experiment Set-up.** We fine-tune Stable Diffusion v1.5 for text-to-image generation using LoRA and ImageReward. Our prompts—“A green-colored rabbit,” “A cat and a dog,” “Four wolves in the park,” and “A dog on the moon”—assess color, composition, counting, and location. During fine-tuning, we generate 10 trajectories (50 transitions each) to calculate the gradient with 1K gradient steps (AdamW, LR = 3e-4), and KL regularization ($\\beta= 0.01$).\n\n**Evaluation.** We first compare ImageReward scores of images generated by the pre-trained model, DPOK and PI-FT. To ensure a fair evaluation, we perform 10 gradient steps with LR = 1e-5 per sampling step in DPOK. Each gradient step is computed using 50 randomly sampled transitions from the replay buffer. Consequently, 1K sampling steps and 10K gradient steps in DPOK result in a computational cost similar to that of PI-FT. As shown in Figure 1 in the link, PI-FT consistently **outperforms** both the pre-trained model and DPOK across all four prompts. Figure 2 further shows that PI-FT more accurately generates the correct number of wolves, and correctly places the dog on the moon. Additionally, PI-FT avoids mistakes such as depicting a rabbit with a green background as appeared in the pre-trained model and DPOK. The texture of the generated animals is also more natural-looking under PI-FT compared to other baselines.\n\n**Effect of KL regularization.** KL regularization is known to improve fine-tuning. We analyze its effect in the PI-FT method using the prompt “Four wolves in the park,” varying $\\beta \\in$ {0.01, 0.1, 1.0}. Figure 3 shows that the gradient norm decreases to 0 in all three settings, illustrating the convergence of our algorithm. In Figure 4, we see that the ImageReward score increases and eventually stabilizes when $\\beta$ is small. In contrast, the score exhibits limited improvement with large $\\beta$. This observation is consistent with Figure 5, where the KL divergence remains large throughout training for $\\beta = 0.01$ and stays significantly at a smaller level for $\\beta \\in$ {0.1, 1.0}. Figure 6 illustrates that smaller $\\beta$ produces images with nearly four wolves, while larger $\\beta$ yields fewer. These results underscore the importance of the KL coefficient in our framework.\n```\n… Could authors show more details about the comparison and similarity of their methods and your method.\n```\nIn the literature, DPOK and DDPO use generic RL algorithms such as PPO or REINFORCE to optimize the policy/score. Despite showing empirical promise on certain tasks, existing methods **overlook** the underlying structure of diffusion models, leaving significant room for improvement in efficiency and lacking any theoretical justification. In contrast, by fully utilizing the problem structure, we formulate fine-tuning as a KL-regularized stochastic control problem (with linear dynamics), and propose the PI-FT algorithm. This structure allows us to derive regularity and convergence guarantees, closing the gap left by prior works. We will revise lines 90-96 to clarify these points.\n```\n... show the convergence and the regularity of the optimal value function & the optimal control policy of your approach based on above mentioned empirical settings.\n```\nIn the newly added experiments, we do observe convergence of both the control policy and the value function during training, supporting our theoretical claims. \n```\n… I am curious could the conditions in above lemmas and theorems be realized in empirical settings…\n```\nEmpirically, we observe numerical results consistent with our theoretical claims. In Figure 3 (in the link), we observe that the algorithm converges **linearly** and the gradient norm of the U-Net stabilizes after approximately 200 sampling steps. Additionally, we compare PI-FT and DPOK under the same computational budgets; as shown in Figure 1, PI-FT achieves consistently higher ImageReward scores across all four prompts. This performance suggests that our approach converges more efficiently in practice. Although our theory assumes Lipschitz conditions and a specific range of the regularization coefficient, our experiments demonstrate that convergence remains robust even for small $\\beta = 0.01$ and relatively large learning rates. \n\nThank you for your detailed feedback. We hope our response addresses your concerns. If so, we would appreciate it if you could reflect it in your evaluation of our paper.", "claims_and_evidence": "The claims in this manuscript are supported by clear and convincing evidence, such as\n1. The illustration of the connection between DDPM and stochastic control\n2. The regularity and well-posedness of the optimal value function and the optimal control policy\n3. The framework/algorithm for fine-tuning diffusion models (PI-FT) and the analysis for its convergence\n4. linear parameterization to enable practical and efficient update for above algorithm", "methods_and_evaluation_criteria": "Both the regularity & well-posedness and the convergence analysis give the proposed method a strong theoretical foundation. However none of the above theoretical derivation results and convergence of the algorithm have been realized on a particular dataset and a specific task. I think adding some experimental results will make this approach more convincing.", "theoretical_claims": "This paper provides strong theoretical foundations for fine-tuning diffusion models, with rigorous derivations supporting its key claims. It establishes a clear connection between Denoising Diffusion Probabilistic Models (DDPM) and stochastic control, ensuring a principled approach to fine-tuning. The authors prove the well-posedness and regularity of the optimal value function and control policy, guaranteeing stability. They develop a policy iteration algorithm (PI-FT) and provide a detailed convergence analysis. Additionally, the framework extends to a parametric setting with linear parameterization, enabling efficient and practical updates for high-dimensional applications.", "experimental_designs_or_analyses": "Unfortunately, the method proposed in this paper has not yet been demonstrated on a specific dataset or task. Incorporating experimental results would further validate the approach and enhance its credibility.", "relation_to_broader_scientific_literature": "Recent empirical results provide compelling evidence to support proposed framework, which implements reinforcement learning (RL) or control-based fine-tuning algorithms closely aligned with their approach. These empirical studies have demonstrated the practical efficacy of KL-regularized control formulations in fine-tuning diffusion models.", "essential_references_not_discussed": "The paper provides a comprehensive discussion of related works. All essential references needed to understand its key contributions are adequately cited and analyzed.", "other_strengths_and_weaknesses": "I have to say that the theory in this paper is very solid. However, the lack of empirical results is also the biggest drawback of this paper.", "other_comments_or_suggestions": "Including experimental results would strengthen the paper by demonstrating the practical effectiveness of the proposed approach. Empirical validation would also make the theoretical contributions more compelling and applicable to real-world scenarios.", "questions_for_authors": "1. Line90-Line96 in introduction section, authors mention that recent empirical results provide compelling evidence to support their framework. Could authors show more details about the comparison and similarity of their methods and your method.\n2. Actually, authors can show the convergence and the regularity of the optimal value function & the optimal control policy of your approach based on above mentioned empirical settings.\n3. The parameterization trick in proposed algorithm a linear convergence rate. The linear convergence rate means convergence faster and obtaining better performance with same time/resources? I am curious could the conditions in above lemmas and theorems be realized in empirical settings. External experiment results to prove this point is better.", "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:

Meta-review: All reviewers appreciate the theoretical contribution of this paper. I also think in that way as well. However, I suggest citing related works in the proper place. For example, equations [7] and [8] in diffusion models were very explicitly derived in the following with a similar title paper:

Uehara, M., Zhao, Y., Black, K., Hajiramezanali, E., Scalia, G., Diamant, N. L., ... & Levine, S. (2024). Fine-tuning of continuous-time diffusion models as entropy-regularized control. arXiv preprint arXiv:2402.15194.

or Domingo-Enrich et al. (2024).
Decision: Accept (poster)