Score as Action: Fine Tuning Diffusion Generative Models by Continuous-time Reinforcement Learning
Authors: Hanyang Zhao
Haoxian Chen
Ji Zhang
David Yao
Wenpin Tang
Venue: ICML 2025 poster (2025)
Keywords: Continuous-time Reinforcement Learning
Diffusion Models Fine-tuning
Reinforcement Learning from Human Feedback
Primary area: reinforcement_learning
TL;DR: A continuous-time RL method for diffusion models fine-tuning, being more stable, efficient and robust to time discretizations

Abstract: Reinforcement learning from human feedback (RLHF), which aligns a diffusion model with input prompt, has become a crucial step in building reliable generative AI models. Most works in this area uses a *discrete-time* formulation, which is prone to induced errors, and often not applicable to models with higher-order/black-box solvers. 
The objective of this study is to develop a disciplined approach to fine-tuning diffusion models using *continuous-time* RL, formulated as a stochastic control problem with a reward function that aligns the end result (terminal state) with input prompt. The key idea is to treat score matching as controls or actions, and thereby connecting to policy optimization and regularization in continuous-time RL. To carry out this idea, we lay out a new policy optimization framework for continuous-time RL, and illustrate its potential in enhancing the value networks design space via leveraging the structural property of diffusion models. We validate the advantages of our method by experiments in downstream tasks of fine-tuning large-scale Text2Image models, Stable Diffusion v1.5.

Score as Action:
Fine-Tuning Diffusion Generative Models
by Continuous-time Reinforcement Learning
Hanyang Zhao
Haoxian Chen
Ji Zhang?
David D. Yao
Wenpin Tang'
Abstract
(Xu et al., 2022). And, the emergence of duman-interactivv platforms like ChatGPT (Ouyang et al., 2022) and Stable Diffusion (Rombach et al., 2022) has further increased the demand for diffusion models to align with human preference or feedback.
Reinforcement learning from human feedback (RLHF), which aligos a diffusion model with in- put prompt, has become a crucial step in building reliable generative AI models. Most works in this area a discrete-time formulation, which is prone induced discretization errors, and often not applicable to models with higher-order/black- box solvers. The objective of this study is to de- velop a disciplined approach to fine-tune diffusion models using continuous-time RL, formulated as stochastic control problem with a reward func- tion that aligns the end result (terminal state) with input prompt. The key idea is to treat score match- ing controls or actions, and thereby making connections to policy optimization and regular- ization in continuous-ting RL. To carry out this idea, we lay out a new policy optimization frame- work for continuous-time RI., and illustrate its potential in enhancing the value networks design space via leveraging the structural property of diffusion models. We validate the advantages of our method by experiments in downstream tasks of fine-tuning large-scale Text2Image models of Stable Diffusion v1.5.
To meet such demands, (Hao al.. 2022) proposed a natu- ral way to fine-tune diffusion models using reinforcement learning (RL, (Sutton & Barto, 2018)). Indeed, RL has al- ready demonstrated empirical successes in enhancing the performance of LLM (large language models) using hu- man feedback (Christiano et al., 2017; Ouyang et al., 2022; Bubeck et al., 2023), and (Fan & Lee, 2023) is among the first to utilize RL-like methods to train diffusion models for better image synthesis. Moreover, (Lee et al., 2023; Fan et al., 2023; Black et al., 2023) have improved the text-to- image (T2I) diffusion model performance by incorporating reward models to align with human preference (e.g., CLIP (Radford et al., 2021), BLIP (Li et al., 2022), ImageReward (Xu et al., 2024)).
Notably, all studies referenced above that combine diffusion models with RL are formulated as discrete-ting sequential optimization problems, such as Markov decision processes (MDPs, (Puterman, 2014)), and solved by discrete-time RL algorithms such as REINFORCE (Sutton et al., 1999) or PPO (Schulman et al., 2017).
Yet, diffusion models are intrinsically continuous as they were originally created to model the evolution of ther- modynamics Scohl-Dickstein et al., 2015). Notably, the continuous-time formalism of diffusion models provides unified framework for various existing discrete- time al- gorithms as shown in (Song et al., 2021b): the denoising steps in DDPM (Ho et al., 2020) can be viewed as a discrete approximation of stochastic differential equation (SDE) and are implicitly score-based under a specific variance- preserving SDE (Song al., 2021b); and DDIM (Song et al., 2020), which underlies the success of Stable Diffu- sion (Rombach et al., 2022), can also be seen as numerical integrator of an ODE (ordinary differential equation) sam- pler (Salimans & Ho, 2022). Awareness of the continuous- time nature informs the design structure of the discrete-time SOTA large-scale T2I generative models (e.g.,(Dhariwall Nichol, 2021; Rombach et al., 2022; Esser et al., 2024)), and
1. Introduction
Diffusion models "Sohl-Dickstein et al., 2015), with the initial distribution
capacity to turn a into a desired target distribution through a well-designed denoising process (Ho et al., 2020; Song et al., 2020; 2021b), have recently found applications in diverse areas such as high-quality and creative image generation (Ramesh al., 2022; Shi et al., 2020; Saharia al., 2022; Rombach al., 2022), video synthesis (Ho et al., 2022), and drug design
'Deaarmment of IEOR, Columbia University, New York, USA l DDepartment of CS, Stony Brook University, New York, USA. Correspondence to: <hz2684, yao, wt2319@ 223319@cluumbia.cuu.
Proceedings of the 4 International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Score as Action: Fine-tuning Diffusion Models by Continuous- time Reinforcement Learning
enables simple controllable generations by classifier guid- ance to solve inverse problems (Song et al., 20)) It also motivates more efficient diffusion models with continuous- time samplers, including the ODE-governed probability (normalizing) flows ( (Papamakarios et al., 2021; Song et al., 2021b) and rectified flows (Liu et al., 2022; 2023) underpin- ning Stable Diffusion v3 (Esser et al., 2024). A discrete-time formulation of RL algorithms for fine-tuning diffusion mod- els, if/when directly applied to continuous-tima diffusion models via discretization, can nullify the models' contin- uous nature and fail to capture or utilize their structural properties.
property of diffusion models generate tractable optimiza- tion problems and to enhance the design space of value networks. This includes transforming the KL regularization to a tractable running reward over time, and a novel design of value networks that involves "sample prediction" (also known as z-prediction in diffusion literature) by sharing parameters with policy networks and fine-tuned diffusion models. Through experiments, we demonstrate the drastic improvements over naive value network designs.
Third, we provide new theory for RL in continuous-time and space, which leads to the first scalable policy optimiza- tion algorithm (for continuous-timg RL), along with esti- mates of policy gradients via generated samples. Compared with existing works in continuous-time PPO, we consider special case of state-independent diffusion coefficients cater for the diffusion models design, which yields sharper bounds and closed-form advantage-rate functions instead of estimations. (See Section 4.)
For fine-tuning diffusion models, discrete-time RL algo- rithms (such as DDPO) require a prior chosen time dis- cretization in sampling. We thus examine the robustness of fine-tuned model to the inference time discretization, and observe an "overfitting" phenomenon as illustrated in Figure 1. Specifically, improvements observed during infer- ence at alternative discretization timesteps (25 and 100) are significantly smaller than that of sampling timestep (50) in
1.1. Related Works
RL.
Papers that relate to our work are briefly reviewed below.
0.02
Continuous- time RL. (Wang et al., 2020) models the noise or randomness in the environment dynamics as following an SDE, and incorporates an antropy-based regularizer into the objective function to facilitate the peploration-eploiaatatie tradeoff. Follow-up works include designing model-free methods and algorithms under either finite horizon (Jia & Zhou, 2022a;b; 2023) or infinite horizon (Zhao et al., 2024b).
0.01
0.46
-0.01
-0.02
-0.03
-0.04
Step
Figure 1. Reward curve of model checkpoints sampling under dif- ferent discretization steps (25, 50, 100): After training Stable Diffussion v1.4 for fixed prompt with 60 training steps by DDPO (Black et al., 2023) with 50 discretization steps, the average re- ward of images generated by the checkpoints obtained (under 50 discretization steps) evaluated by ImageReward (Xu et al., 2024) increases by 0.046, while the average reward of images generated with 100 discretization steps only increases by less than 0.016.
RL for fine-tuning T2I diffusion models. DDPO (Black et al., 2023) and DPOK (Fan et al., 2023) both discrete the time steps and fine-tune large pretrained T2I diffusion mod- els through the reinforcement learning algorithms. More- over, (Ren et al., 2024) introduces DPPO, policy gradient- based RL framework for fine-tuning diffusion- based policies in continuous control and robotic tasks.
In addition, for high-order solvers (such as 2nd order Heun used in EDM (Karras et al., 2022)), discrete-time RL meth- ods will require solving high-dimension root-finding prob- lem for each inference step, which is inefficient in practice.
Other Preference Optimizations for diffusion models. (Wallace et al., 2024) proposes an adaptation of Direct Pref- erence Optimization (DPO) for aligning T2I diffusion mod- els like Stable Diffusion XI. to human preferences. (Yuan et al., 2024) proposes novel fine-tuning method for diffu- sion models that iteratively improves model performance through self-play, where a model competes with its previous versions to enhance human preference alignment and visual appeal. See Section 4.5 in (Winata et al., 2024) for a review.
Main contributions. To address the above issues, we de- velop a unified continuous-time RL framework to fine-tune score-based diffusion models.
Our first contribution a continuous-time RL framework for fine-tuning diffusion models by treating score functions actions. This framework naturally accommodates discrete- time diffusion models with any solver as well as continuous- time diffusion models, and overcomes the afore-mentioned limitations of discrete-tima time RL methods. (See Section 3.)
Stochastic Control. (Uchara et al., 2024), which also for- continuous-
mulated the diffusion models alignment time stochastic control problem with different parameter- ization of the control; (Tang, 2024) also provides a more rigorous review and discussion. (Domingo-Enrich et al., 2024) proposes to use adjoint methods instead to solve the
Second, we illustrate the promise of leveraging the structural
2

---

Score as Action: Fine-tuning Diffusion Models by Continuous- time Reinforcement Learning
similar control problem. In concurrent work ours, (Gao et al., 2024) uses q-learning (Jia & Zhou, 2023) for infer- ring the score of diffusion models (instead of fine tuning pretrained model). whose formulation relies on an earlier version (Zhao et al., 2024a) of this paper.
The value function associated with the feedback policy
(a( 1,x0 RR4 is
V(1,;)):= B [ r(s, (x X;
The rest of the paper is organized as follows. In Sec- tion 2, we review the preliminaries of continuous-ting RL and score-based diffusion models. Section presents our continuous-ting framework for fine fine-tuning diffusion mod- els using RLHF, with the theory and algorithm for policy optimization detailed in Section 4, and the effectiveness of the algorithm illustrated in Section 5. Concluding remarks and discussions are presented in Section 6.
F. F(s, (11,, + X =B [."
(4)
The performance metric is V* := SRx V(O, x; p))(d)) and V := max, V The task to construct a sequence of (feedback) policies Tg., k 2. recursively such that the performance metric is non-decreasing in k.
2. Preliminaries
q- Valuee Following the definition in (Jia & Zhou, 2023), given a policy т and (t,x,a) E (0) x R x A, we construct a "perturbed" policy, denoted by t:: It takes the action a E AA on [t, + At), and then follows т on [t + At, oo). Specifically, the corresponding state process X given X; = x, breaks into two pieces: on [t, + At), it is X following (1) with a (i.e., (t,, 2, a) 1); while on [L + Al, oo), it is X following (3) but with the initial time-state pair (t + A
2.1. Continuous-time RL
Diffusion Process. We consider the state space Rd, and denote by A the action space. Let x( (, be a feedback policy given E [0, and x E Rd. The state dynamics (X7, 0 S t < T) is governed by the following SDE:
dxx" = bt, X,,, a.) dt + (t)dB, XI p. (1)
The q-value measures the rate of the performance difference between the two policies when At - 0, and is shown in (Jia & Zhou, 2023) to take the following form:
X
where (B,, t 2 0) is d-dimensional Brownian motion; b R, x Rdd, x A + Rdd and σ E, + R, are given functions; the action a, follows the distribution π (- (t, by external randomization; and p is the initial distribution over the state space.
OV Öt (t,
at,, a;))
OV
8V Or² (t, (t.a;))) (5)
Performance Metric. Our goal to find the optimal feed- back policy nt that maximizes the expected reward over finite time horizon:
H 1, a dx
(t, (1,;;1)
y, A) = b(t, a) y fo²(t) Au + r(t, x, a) is the (generalized) Hamilton function in stochas- tic control theory (Yong & Zhou, 1999).
where H(tt a,
[." r (t, at dt + h(x7)
V = max E
X5 ~ p
(2) where r: R2+ x Rdd A- + R and h d R are the running and terminal rewards respectively. Given policy
2.2. Score-Based Diffusion Models
Forward and Backward SDE. We follow the presentation in (Tang & Zhao, 2024). Consider the following SDE that governs the dynamics of process (X,, 0 x t T) in Rd (Song et al., 2021b).
x()) let a(,, x(-)) Sa MI. x, (a)da. We consider the following equivalent representation of (I):
dx, b (t. 12, x.)) dt (t)) Xo ~~pp P.
(3) in the sense that there exists probability measure P that sup- ports a d-dimensional Brownian motion (B., 2 0), and for each 2 O, the distribution of X, under ip agrees with that of X, under P defined by (I). Note that the dynamics (3) does not require external randomization. Accordingly, set (t,, xr, x) := JA rtt, xz, a) (a)
dx, = f(1, x,)dt + g())d No ~ Pdmaa.)); (6)
where (B,, t 1 0) is a d-dimensional Brownian motion, S : R x Rd - Rd and و : R. - R, are two given functions (up to the designer to choose), and the initial state X, follows a distribution with density Paau(.), which is shaped by data yet unknown priori. Denote by pa(-) the probability density of Xt.
FFor our applications here we assume that the diffusion co- efficient att) only depends on time t. Note, however, that the general continuous-time RL theory also holds for time-, state- and action-dependent a (t,, x, a), see (Jia & Zhou, 2022a;b).
Run the SDE in (6) until a given time > 0, to obtain X7 ~ .T. Next, consider the "time reversal" of X,, denoted X700, such that the distribution of X agrees with
3

---

Score as Action: Fine-tuning Diffusson Models by Continuous- time Reinforcement Learning
that of - on [0, T). Then, (Xfe,, 0 S l < T) satisfies the following SDE under mild conditions on and g:
Scores as Actions. First, to broaden the application context of the diffusion model, we add parameter the score function, interpreted as "class" index or label (e.g. for imput prompts). Then, the backward SDE in (9) becomes:
dxp"" (T, + g²(T - - log Po dt + 9[T
dx;- = (- + + t) (T - t, Xi. c)) dt + gT -
where V Vlog Pt (3) is known as Stein's 's score functiion Below we will refer to the two SDE's in (6) and (7), respectively, as the forward and the backward SDE.
(10)
Next, comparing the continuous RL process in (3) and the in- ference process (10), we choose and a in the RL. dynamics in (3) as:
For sampling from the backward SDE, we replace pr() with some Puoise (.) as an approximation. The initialization Phoise (.) is commonly independent of pdam(.), which is the reason why diffusion models are known for generating data from "noise".
{ ((t) :- -T 6 0t) = x, + g t)a.
(11)
Score Matching. Since the score function v log Pr (x) in (7) is unknown, the idea is to learn the score Spn (t, x) R V Pxr),, which is often referred to as pretraining. It boils down solving the following denoising score match- ing (DSM) problem (Vincent, 2011) 2:
In the sequel, we will stick this definition of b and ø.
Define a specific feedback control, of 80,0 (T t, Xi c)2 and the backward SDE in (10) is expressed as:
dx; = (1,,, aq 9m dt (t)dB,
(12)
JOSM (0) = TE A(1 (1, xx)) - V log (og (8)
This way, the score function is replaced by the action (or control/policy), and finding the optimal score becomes a policy optimization problem in RL. Denote by pa (t, ,))
where 22 ~ (-( and X [0,7] + Rs0 is a chosen positive-valued weight function.
the probability density of X;- in (12).
Inference Process. Once the best approximation $ is obtained, we use to replace (x)) in (7). The cor- responding approximation to the reversed process XTe, de- .then noted as Xi_ then follows the SDE:
Exploratory SDEs. As will deal with the time-reversed process Xē exclusively from now on, the superscript "- will be dropped to lighten the notation. To enhance exploration, we will use a Gaussian control:
dx;- = (- - x, + g - 1)00,m (T - 1, )) de + - g
at ~ x,, =
c), Ex).
(13)
(9)
Specifically, the dependence on @ is through that of the mean function nº, while the covariance matrix E, only depends on time t, representing a chosen exploration level at t. For brevity, write X for the tüme-reversede process X;" driven is governed by by the policy T Then (X?, 0 S SLS ST) the SDE:
with Xi- (-).sse-). At time t = T, the distribution of Xx= is expected to be close to (-) The well-known DDPM (Ho et al., 2020) can be viewed as a discretized version of the SDE in (9). This has been established in (Song et al., 2021b; Salimans & Ho, 2022: Zhang Chen, 2022; Zhang al., 2022); also refer to further discussions in Appendix A. Throughout the rest of the paper, we will focus the continuous formalism (via SDE).
xxx; = [- + 22(TT t) (t, c)] dt x ~.
+
(14)
Denote by (t) the probability density of X--
3. Continuous-time RL for Diffusion Models Fine Tuning
Objective Function. The objective function of the RL problem consists of two parts. The first part is the terminal reward, i.c. a given reward model (RM) that a function of both XT and c. For instance, if the task is T2I generation, then RM(X,,e) represents how well the generated image XT aligns with the input prompt The second part is a penalty (i.c., regularization) term, which takes the form of the KL. divergence between pd (T, :cc and its pretrained counterpart. This is similar in spirit to previous works on fine-tuning diffusion models by discrete-time RI,, see e.g..
Here we formulate the task of fine-tuning diffusion models as continuous-time stochastic control problem. The high- level idea it to treat the score function approximation as a control process applied to the backward SDE.
'There are several existing score matching methods, among which the DSM is the most tractable one because (=(-)1o)) is ac- cessible for wide class of diffusion processes; in particular, it is conditionally Gaussian (6) is linear SDE (Song et al., 2021b).
4

---

Score as Action: Fine-tuning Diffusion Models by Continuous-ting Reinforcement Learning
(Ouyang et al., 2022; Fan al., 2023). As for exploration, note that it has been represented by the Gaussian noise in af refer to (13), and more on this below. So, here is the problem we want to solve:
diffusion models:
V' (1, (1 2:
Cakip() (1) RM (1,3, + Com(t) Fo(t, x,
(20)
max E [MM X - KI. (V' Clp (T,y,))
reward mean predictor
residual term corrector
(15)
where va denotes a family of functions parameterized by (0, Φ), and
where 3 > 0 is (given) penalty cost.
To connect the problem in (15) to the objective function of the RI. model in (2), we need the following explicit expression for the KL divergence term in (15).
Lo x, = 1 (o)sa(t,z,)) +)) a22
(21)
with or and or being noise schedules of diffusion models (see Appendix A.2 for details). When 0 = ppc. de predicts denoised sample given the current x: and the score estimate so (t, x, c), which is known as Tweedie formula. To treat the second term in (18), our intuition comes from that
Theorem 3.1. For any given e, the KL divergence between g and prre is:
KL
(L (.T.,c))lppr (c)
(T -t)
B
11 (t, (1,
1 2 (1, (t, c) dt.
0
2
RM (M(E((XT X))) =s ((RM(Xx)) |
(22)
(16)
if we are allowed to exchange the conditional expectation and the reward model score (though generally it's not true). Fo x, c) are effectively approximations the residual term, which can be seen composition of the possible reward error and the first term in (18).
Proof Sketch. The full proof is given in Appendix B.I.
As a remark, it is important to use the "reverse" KL diver- gence KT. (p (T, (c),, pp)lp"rr (T,,,))), because yields the expectation under the current policy T0 that can be estimated from sample trajectories. By Theorem 3.1, the objective function in (15) is equivalent to the following:
We refer these two parts to as reward mean predictor and residual corrector. There Cakip(t) and Cour (t) are differen- tiable functions such that Cakip (T) = and Cour = 0. so the boundary condition (19) is satisfied. Notably, similar parametrization trick has also been used train successful diffusion models such as EDM (Karras et al., 2022) and consistency models (Song al., 2023).
17 E
3 g (T +)
pr 1
r(1.22a0f)
E ERMMX c),
(17)
For learning the value function, we use trajectory-wise Monte Carlo estimation to update &i by minimizing the mean square error (MSE). In our experiments, we observe that choosing Cski (t) = cos( y t) and Coml (t) sin 2 t) yields the smallest loss (see Table 1). Also refer Section 5.2 for more architecture details.
M(Xg.c)
where we abbreviate (t, X². c) and yr XP. c) by o
and Me Pere respectively. Thus, maximizing the objective function in (15) aligns with the RL model formulated in (2). We can also define the corresponding value function as:
Original Onginal-Dnnoseed
2; = [
B g -t)ll14 - Opre 1|
Square Errur
Denoised- Origgnall Denoised-Denoleed Denoised+ with sin
+ +RM(X.,c) X =x m
(18)
5
4"
3
Value Network Design. We also adopt function approxi- mation to learn the value function (i.e., the critic). For the value function (t 2; associated with policy x², there is the boundary condition:
2
0
5
10
15
Number of itterations
(T) = E [RM(X).)) X² = x] = RM(x, c). (19)
Figure 2. Pretraining Value Function with Different Architecture.
To meet this condition, we propose the following parametrization that leverages the structural property of
Note that is possible to learn the value function by either solving the associated Hamilton-Jacoh Bellman equation,
5

---

Score as Action: Fine-tuning Diffusion Models by Continuous- time Reinforcement Learning
Architecture
Predictor
Corrector
Coul
MSE
Contimuous-time TRPO/PPO. We also derive the continuous-time finite horizon analogies of TRPO and PPO for the discrete RL in the finite horizon setting (Schulman et al., 2015; 2017), and in the continuous-ting infinite hori- zon setup (Zhao et al., 2024b). The Performance Difference Lemma (PDL) is as follows.
Baseline
RM(2, c)
F(a.c) c)
1-cos( 77
2.63
Org+Denoised
RM( c)
F(e9,c) c)
1-cos(
2.51
27
Denoiscd+Orig
RM [ia, c)
F(z, c)
1-cos(
0.67
Denoised-Denoised
RM c)
C(
1-cos(
pr
0.66
Denoised+Orig
RM (8)
c(
sin( ()
0.29
Lemma 4.2. We have that:
Table 1. Comparison of different architecture configurations: de is ahbreviated for 5o fo Ir c).
v - vo - S'" (t, xt, dz; ( ydt.
=B
(24)
or minimizing Bellman's 's error rate (which can be seen as the continuous-time analog of temporal difference). However, both approaches will yield supervised learning objective that contains the second-order derivative of the value func- tion, which is hard to optimize. We leave for future work the investigation of policy evaluation methods that are based on partial differential equations or temporal difference rates.
The proof is similar to Theorem 2 in (Zhao et al., 2024b). In the same essence of (Kakade & Langford, 2002; Schul- man et al., 2015; Zhao et al., 2024b), we define the local approximation function to v by
6 (af I, X 2 (aflt, x²)
(0(0)= v +E +
g X at 12:7) ddtt. (25)
Observe that
4. Continuous-time Policy Optimization
(i) ()) - vo, (ii V L (6) 10-0 Vave l==
To efficiently optimize the continuous-time RL problem raised above, we further develop the theory of policy opti- mization in continuous time and space for fine-tuning dif- fusion models. Different from the general formalism in the literature (Schulman et al., 2015; Zhao et al., 2024b), we focus on the case of (1) KL regularized rewards, and (2) state-independent diffusion coefficients in the continuous- time setup, which yield new results not only in the analysis but also in the resulting algorithms.
i.e. the local approximation function and the true perfor- mance objective share the same value and the same gradi- ent with respect to the policy parameters. Thus, the local approximation function can be regarded as the first order approximation to the performance metric.
Now we provide analysis on the gap - [0), which guarantees the policy improvement (similar to approaches for discounted/average reward MDP (Schulman et al., 2015; Zhang & Ross, 2021), and for continuous -tim RL in the infinite horizon (Zhao et al., 2024b)).
Policy Gradient. We first show that the continuous-time policy gradient can be directly computed without any prior discretization of the time variable.
Assumption 4.3. (Bounded Reward and q function): There exists M 0 such that for any c, and a, [RM(z, c) s M and (t(t, xx, 0; x x²) I M for any t,, n². 9.
Theorem 4.1. The gradient of an admissible policy n² parameterized by 0 takes the form:
VoVO E. 1 Volog R (a] (1,1,,xx))x ((t, Xi, al: T )dt
Theorem 4.4. Under Assumption 4.3, then for any policy 0 such that |ln( [In(_"(e)) (lbe)))))
(ar) ) is bounded, and KL(p ( S 1,
(23) there exists a constant C. > O such that:
where n², at and are as defined in (13) and (5).
1 10)S
Proof Sketch. The full proof is given in Appendix B.2.
/ (-1,,
E T
C
x()aat (1,t,
Note that the only terms in the q-value function that involve action are (the second order term is irrelevant to action a):
C
(26)
Proof Sketch. Different from the proof' in (Schulman et al., 2015; Zhao et al., 2024b), which bound the difference through PDL., our proof relies on the structural property of diffusion models by bounding:
0 (t, x) (t,
(T
Ox
In addition, the value function approximation can be com- puted by Monte Carlo or the martingale approach in Jia & Zhou (2022a), and then Oz can be evaluated by backward propagation. Since the reward can be con-differentii.llc and also for the sake of efficient computation, we can approxi- mate j" (t, x, a) x (V(t, x + ? - t)a) - V (tt.x)) /n. where 7) is a scaling parameter.
@V
v and (0) v
separately. The detailed proof is given in Appendix B.3.
By Theorem 4.4, we can apply the same technique as in PPO (Schulman et al., 2017) by clipping the ratio and replacing !
6

---

Score as Action: Fine-tuning Diffusion Models by Continuous-time Reinforcement Learning
with (which is equivalent to adapting baseline function). This yields the policy update rule as:
DxMI
CTRL Mean FID
clip (pi, (dide",clip ( min (Dise",
df") dl,
0+ max 8
20
(27)
where the advantage rate function and the likelihood ratio are defined by
12
1000
2000
3000
4000
5000
6000
Step
T It, p.
Bn (. -.- = q(t, X
7he (a (t X("))
Figure 3. Training curves of DxMI and continuous-tine RL.
The surrogate objective can then be optimized by stochastic gradient descent. The pseudo-code of our algorithm is listed in Appendix C.2.
5. Experiments
5.1. Enhancing Small-Steps Diffusion Models
Setup. We evaluate the ability of our proposed algorithm to train short-run diffusion models with significantly reduced generation steps T', while maintaining high sample quality. In the experiment, we take T 10. Our experiments are conducted on the CIFAR-10 10 (32x32) dataset (Krizhevsky et al., 2009). We fine-tune pretrained diffusion model back- bone using DDPM (Ho al., 2020). The primary evaluation metric the Fréchet Inception Distance (FID) (Heusel et al., 2017), which measures the quality of generated samples.
Figure 4. DxMI samples at the 6000-th step
Figure 5. Continuous time RL samples at the 6000-th step
Ladent Timestamp Prompt
To benchmark our method, we compare it against DxMI (Yoon et al., 2024), which formulates the diffusion model training as an inverse reinforcement learning (IRL) problem. DxMI jointly trains a diffusion model and an energy-based model (EBM), where the EBM estimates the log data den- sity and provides reward signal to guide the diffusion process. To ensure fair comparison, replace the policy improvement step in DxMI with our continuous-time RL counterpart, maintaining consistency while evaluating the effectiveness of our approach. We set the learning rate of the value network to x 10-5 and U-net to 3 x 10-7
roward mnaa predictor
Prempt
Latent
WE Dectalo
Xe(C...C)
X
VA-
CC(c)
Imsge DCx
Image
SJP
c2(e)
MPP Ourput
Ouput
residual conrector
F(z,
(tM
Result. Figure 3 shows our approach converges signifi- cantly faster than DxMI, and achieves consistently lower FID scores throughout training. The samples from the two fine-tuned models are shown in Figures 4 and 5. In compar- ison, the samples generated from the model fine-tuned by continuous-time RL have clearer contours, better aligned with world features, and exhibit superior aesthetic qual- ity.
Figure 6. We adopt the similar backbone of ImageReward for two parts in value network, both by adding an ML. layer over the BLIIP encoded latents.
v1.5 We adopt the pretrained ImageReward (Xu al., 2024) the reward signal during RL, it has been shown in previous studies to achieve better alignment with human preferences to other metrics such as aesthetic scores, CLIP and BLIP scores.
5.2. Fine- Line-Tuning Stable Diffusion
Setup. We also validate our proposed algorithm for fine- uning large-scale T2I diffusion models, Stable Diffusion
h http:://ttgginnicccccooooaaaaaat diffusion-vi-5

---

Score as Action: Fine-tuning Diffusion Models by Continuous-ting Reinforcement Learning
We train the value networks with full parameter tuning. while use LoRA (Hu et al., 2021) for tuning the U-nets of diffusion models. We adopt a learning rate of 10-7 for optimizing the value network, 3 x 10-6 for optimizing the U-net and B = 5 x 10-5 for regularization. We train the models on H200 GPUs with 128 effective batch sizes.
Value Network Architecture. Since we fix the reward model as ImageReward, we design the value network by using similar backbone to the ImageReward model, which is composed of BLIP and a MLP header (see Figure 6). To ensure the boundary condition, we fix the parameters (i.e., BLIP and MLP) in the left part (skyblue) of the value network and only tune 30% of the parameters of BLIP in the right part (green). The VAE Decoder on both parts is fixed for efficiency and stabilized training.
Figure 8. Model generations with prompt "A unicorn in clear- ing. It has a single shining horn. Volumetric light." a) Top: Base model Stable Diffusion v1.5; b) Mid: Continuous-time RI after 50 training steps: c) Bot: Continuous- RL after 100 training steps.
As remark, replacing with (f,x)) in the "residual cor- rector" leads to minimum gain, compared to the drastic improvement brought forth by using (t, z) as the input in the "reward mean predictor". See Figure 2 and Table for our ablation of network architecture and MSE statistics.
models fine-tuning methods like DRaFT (Clark et al., 2023) and AlignProp (Prabhudesai et al., 2023) Appendix C.I.
Policies trained by Continuous-timt RL are robust to time discretization. We find that the policies trained by continuous-ting time RL achieve coherent performance in terms of the reward mean evaluated by IeageReward. In Figure 7, three line plots that correspond to 25, 50, and 100 steps almost always overlap after 20 epochs of training.
1.90
1.85
1.80
Reward Mean
1.75
1.70
1.65
1.90
1.60
1.85
1.55
CTRL
Reward Mean
1.80
1.50
1.75
10
20
30
60
Number of epochs
1.70
50 steps
1.65
Figure 9. Performance of contimuous-time RL against discrete-time RL under the same 50 discretization timesteps.
1.60
25 steps
1.55
100 steps
1.50
Why continuous-time approaches show better performance? Here we provide a heuristic explanation. Discrete-time RL methods optimize the objective with a priori time- discretization, which induces an error such that the result- ing optimal policy can be significantly away from the true optimum in continuous time. Continuous-time RL meth- ods, the other hand, only require time-discretization in estimating the policy gradient. The error caused by this discretization - the gap between the resulting optimum and the true (continuous-ting objective) optimum - is bounded by a polynomial of the step size (in gradient estimation) under suitable regularity conditions.
10
20
30
40
50
Number of epochs
Figure 7. Performance continuous-time RL's checkpoints with respect discretization timesteps.
This showcases that our continuous-time RL trained policy is robust to time discretization, which is consistent with our theoretical analysis. Qualitative examples with the same prompt generated by the base model, checkpoints of 50 steps and 100 steps after continuous-tine time RL training can be found in Figure 8.
Continuous-time RL outperforms Discrete-time RL baseline methods in both efficiency and stability. We also compare the reward curves of discrete-time RL with our continuous-time RL algorithms. In Figure 9, the per- formance of the continuous-tin -iime RL is much more stable, and is more efficient in achieving a high average reward. We include the comparison with other baseline diffusion
6. Discussion and Conclusion
We have proposed in this study a continuous-tim reinforce- ment learning (RL) framework for fine-tuning diffusion models. Our work introduces novel policy optimization theory for RI. in continuous time and space, alongside a
8

---

Score as Action: Fine-tuning Diffusion Models by Continuous-ting Reinforcement Learning
scalable and effective RL algorithm that enhances the gen- eration quality of diffusion models, validated by our experiments.
fine-tuning diffusion models differentiable rewards. arXiv preprint arXiv:2309. /7400, 2023.
Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
In addition, our algorithm and network designs exhibit a striking versatility that allows us to incorporate and leverage some of the advantages of prior works diffusion mod- els design, so as to better exploit model structures and to improve value network architectures. view of this, we be- lieve the continuous-time RL, in providing cross-pollination between diffusion models and RLHF, presents highly promising direction for future research.
Aamingo-Enriche C., Drozdzal, M., Karrer, B., and Chen, T. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprint arXiv:2409.0886/ 2024.
Esser, P., Kulal, S., Blattmann, A., Entezari, R.. Müller, J.,
Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F,, et al. Scaling rectified flow transformers for high-resolution image synthesis. arXiv preprint arXiv:2403.03206 2024.
Impact Statement
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Fan, Y. and Lee, K. Optimizing ddpm sampling with short- cut fine-tuning. arXiv preprint arXiv:2301 2023.
Fan, Y., Watkins, O., Du, Y., Liu, H., Ryu, M., Boutilier, C., Abbeel, P., Ghavamzadeh, M., Lee, K., and Lee, K. Dpok: Reinforcement learning for fine- fine-tuning text-to- image dif- fusion models. arXiv preprint arXiv:2305. 16381, 2023.
Acknowledgement
Hanyang Zhao, Haoxian Chen and Wenpin Tang are sup- ported by NSF grant DMS-211377 133799 and DMS-2206038. Haoxian Chen is supported by the Amazon CAIT fellow- ship. Wenpin Tang receives support from the Columbia Innovation Hub grant, and the Tang Family Assistant Pro- fessorship. The works of Haoxian Chen, Hanyang Zhao, and David Yao are part of a Columbia-City collabo- rative project that is supported by InnoHK Initiative, The Government of the HKSAR and the AIFT Lab.
Gao, X., Zha, J., and Zhou, X. Y. Reward-directed score- based diffusion models via q-learning. arXiv preprint arXiv:2409.0443 2024.
Hao, Y., Chi, Z.,, Dong, L., and Wei, F. Optimizing prompts for text-to-image generation. arXiv preprint arXiv:2212.096111 2022.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
References
Black, K., Janner, M., Du. Y., Kostrikov, I., and Levine, S. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.13301 2023.
Ho, J., Jain, A., and Abbeel, Denoising diffusion proba- bilistic models. In Neurips, volume 33, pp. 6840-6851, 2020.
Bubeck, S., Chandrasekaran, V., Eldan, R., Gehrke, J., Horvitz, E., Kamar, E., Lee, P., Lee, Y. T., Li, Y., Lundberg, S., al. Sparks of artificial general intel- ligence: Early experiments with gpt-4. arXiv preprint 12222 2023.
Ho, J., Chan, W., Saharia, C., Whang, J., Gao, R., Gritsenko, A., Kingma, D. P., Poole, B., Norouzi, M., Fleet, D. J., et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02300, 2022.
Chen, S., Chewi, S., Li, J., Li, Y.. Salim, A., and Zhang, A. R. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. arXiv preprint arXiv:2209. /215, 2022.
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y.. Wang. S., Wang, L., and Chen, W. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.0966 2021.
Christiano, P. F., Leike, J., Brown, T., Martic, M., Legg, S., and Amodei, D. Deep reinforcement learning from human preferences. Advances in neural information pro- cessing systems, 30, 2017.
Jia, Y. and Zhou, X. Y. Policy evaluation and temporal- difference learning in continuous time and space: A mar- tingale approach. J. Mach. Learn. Res., 23(154):1-55, 2022a.
Clark, K., Vicol, P., Swersky, K., and Fleet, J. Directly

---

Score as Action: Fine-tuning Diffusion Models by time Reinforcement Learning
Jia, Y. and Zhou, X. Y. Policy gradient and actor-critic learn- ing in continuous time and space: Theory and algorithms. J. Mach. Learn. Res., 23(275):1-50 2022b.
Prabhudesai, M., Goyal, A., Pathak, D., and Fragkiadaki, K. Aligning text-to-image diffusion models with reward backpropagation. 2023.
Jia, Y. and Zhou, X. Y. q-learning in continuous time. J. Mach. Learn. Res., 24(161):1-61, 2023.
Puterman, M. L. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Kakade, S. M. and Langford, J. Approximately optimal approximate reinforcement learning. In Sammut, C. and Hoffmaann, A. G. (eds.), Machine Learning, Proceedings of the Nineteenth International Conference (ICML 2002), University of New South Wales, Sydney, Australia, July 8-12, 2002, pp. 267-274. Morgan Kaufmann, 2002.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763. PMLR, 2021.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion-based generative models. Advances in Neural Information Processing Systems, 35: 06565-26577, 2022.
Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.0612 06125, 1(2):3 2022.
Ren, A. Z.., Lidard, J., Ankile, L. L., Simeonov, A., Agrawal, P. Majumdar, A., Burchfiel, B., Dai, H., and Simchowitz, M. Diffusion policy policy optimization. arXiv preprint arXiv:2409.0055 00588, 2024.
Kingma, D., Salimans, T., Poole, B., and Ho, J. Varia- tional diffusion models. Advances in neural information processing systems,
2021.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEF/CVF con- ference computer vision and pattern recognition, pp. 10684-10695, 2022.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
Lee, K., Liu, H., Ryu, M., Watkins, O., Du, Y., Boutilier, C.,, Abheel, P., Ghavamzadeh, M., and Gu, S. Aligning text- to-image models using human feedback. arXiv preprint arXiv:2302.12192, 2023.
Saharia, C., Chan, W., Saxena, S., Li, L., Whang, J., Denton, E. L., Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., Salimans, T., et al. Photorealistic rext-to-image dif- fusion models with deep language understanding. Ad- vances in Neural Information Processing Systems, 35: 36479-36499 2022.
Li, J., Li, D., Xiong, C., and Hoi, S. Blip: Bootstrapping language- trggage--mage pre-training for unified languae understanding and generation. In International confer- ence on machine learning. pp. 12888-12906 PMLR, 2022.
Salimans, T. and Ho, J. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512, 2022.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint (iix:220900300) 2022.
Schulman, J., Levine, S., Abbeel. P., Jordan, M., and Moritz, P. Trust region policy optimization. In International conference machine learning, pp. 1889 1897. PMLR, 2015.
Liu, X., Zhang, X., Ma, J., Peng, J., et al. Instaflow: One step is enough for high-quality diffusion-basse text-to- image generation. In The Twelfth International Confer- ence on Learning Representations, 2023.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint 1rXiv:170..06347, 2017.
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C.,, Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35:27730-27744, 2022.
Shi, Z., Zhou, X., Qiu, X., and Zhu, X. Improving image captioning with better use of captions. arXiv preprint arXiv:2006. 11807, 2020.
Papamakarios, G., Nalisnick, E., Rezende, D. J., Mohamed, S., and Lakshminarayanan, B. Normalizing flows for probabilistic modeling and inference. The Journal of Machine Learning Research, 22(1)22617-2680, 2021.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International conference on machine learning, pp. 2256-2265. PMLR, 2015.
10

---

Score as Action: Fine-tuning Diffusion Models by Continuous-ting time Reinforcement Learning
Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
Xu, J., Liu, X., Wu, Y., Tong, Y., Li, Q., Ding, M., Tang, J., and Dong. Y. Imagereward: Learning and evaluating hu- man preferences for text-to-image generation. Advances in Neural Information Processing Systems, 36, 2024.
Song. Y. Shen, L., Xing, L., and Ermon, S. Solving inverse problems in medical imaging with score-based generative models. arXiv preprint 2Xiv:2111.08000 202
Xu, M., Yu, L.,, Song, Y., Shi, C., Ermon, S., and Tang, J. Geodiff: A geometric diffusion model for molecular con- formation generation. arXiv preprint 2022.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score- based generative modeling through stochastic differential equations. In ICLR, 2021b.
Yong, J. and Zhou, X. Y. Stochastic controls: Hamiltonian systems and HJB equations, volume 43. Springer Science & Business Media, 1999.
Song, Y., Dhariwal, P., Chen, M., and Sutskever, Consis- tency models. arXiv preprint arXiv:2303.01469 2023.
Yoon, S., Hwang, H., Kwon, D., Noh, Y.-K., and Park, F. C. Maximum entropy inverse reinforcement learning of difftusion models with energy-based models. arXiv preprint arXiv:2407.0064. 2024.
Sutton, R. S. and Barto, A. G. Reinforcement learning: An introduction. MIT press, 2018.
Sutton, R. S., McAllester, D., Singh, S., and Mansour, Y. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Yuan, H.,, Chen, Z., Ji, K., and Gu, Q. Self-play fine-tuning of diffusion models for eox-tto-mmaee generation. arXiv preprint arXiv:2402.10210, 2024.
Tang. W. Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond. arXiv:2403.06275 2024.
Zhang, Q. and Chen, Y. Fast sampling of diffusion
models with exponential integrator. arXiv preprint arXiv:2204. 23922224..39,,,,. 2022.
Tang, W. and Zhao, H. Score-based diffusion models via stochastic differential equations-a technical tutorial. arXiv preprint arXiv:2402.07487, 2024.
Zhang, Q., Tao, M., and Chen, Y. gddim: Generalized denoising diffusion implicit models. arXiv preprint arXhv:2206.05564, 2022.
Uehara, M., Zhao, Y., Black, K., Hajiramezanali, E., Scalia, G., Diamant, N. L., Tseng, A. M., Biancalani, T., and Levine, S. Fine-tuning of continuous-time diffusion models as arXiv:2402. 15194, 2024.
Zhang. Y. and Ross, K. W. On-policy deep reinforce- ment learning for the average-reward criterion. In In- ternational Conference on Machine Learning, pp. 12535- 12545. PMLR, 2021.
control. arXiv preprint
Zhao, H., Chen, H., Zhang, J., Yao, D. D., and Tang, W. Scores as actions: framework of fine-tuning diffusion models by continuous-time reinforcement learning, arXiv preprint 140Xv:::4090884000 2024a.
Vincent, A connection between score matching and de- noising autoencoders. Neural computation, 23(7):1661- 1674, 2011.
Zhao, H., Tang, W., and Yao, D. Policy optimization for continuous reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024b.
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 Computer Vision and Pattern Recognition, pp. 8228 8238,82288188 2024.
Wang, H., Zariphopoulou, T., and Zhou, X. Y. Reinforce- ment learning in continuous time and space: A stochastic control approach. J. Mach. Learn. Res., 21(198):1-34, 2020.
Winata, G. I., Zhao, H., Das, A., Tang, W., Yao, D. D., Zhang, S.-X., and Sahu, S. Preference tuning with human feedback on language, speech, and vision tasks: A survey. arXiv preprint arXiv:2409. /564, 2024.

---

Score as Action: Fine-tuning Diffusion Models by Continuous-timg Reinforcement Learning
A. Connection between discrete-time and continuous-time sampler
In this section, we summarize the discussion popular samplers like DDPM, DDIM, stochastic DDIM and their continuous- time limits being a Variance Preserving (VP) SDE.
A.1. DDPM sampler is the diseretization of VP-SDE
We review the forward and backward process in DDPM, and its connection to the VP SDE following the discussion in (Song al., 2021b; Tang & Zhao, 2024). DDPM considers a sequence of positive noise scales 0 < 31, ... .5N < For each training data point ze Pdata (z), discrete Markov chain {2o. 21. NNN is constructed such that:
F2 /1 Br B/22-11 i= 1. .N
(28)
where 2;-1 (0, T), thus p (xt 11-1) N - We can further think of x; as the 7th point of T
uniform discretization of time interval [0, 7] with discretization stepsize At N. i.e. Tidt Ij; and also z 2,. To
obtain the limit of the Markov chain when N + oo, we define a function 3 : [0, + Rt assuming that the limit exists:
((t) Arm 3./2r At with t/At. Then when is small, we get:
Je V1- 3()AAr+++ + VB(I)aiZ NE 1 3()), ((00)A+//(()A))) /3(1AAZ..
Further taking the limit At + 0, this leads to:
dx, [(t) /B(t)dB.. OstsT,
and we have:
f(t,z) = ((t)), g(t) = -3()))(())) B(t).
Through reparameterization, we have Pa, (a: zn) = N (a: Vā,zo. (1 - ā;)) 1). where ā; := IT j=1 (1 - -B3)- For the backward process, a variational Markov chain in the reverse direction parameterized with Pa (x:-1 r.))
N 1i-1i - (o+ + Biso (i, z1)) ...I). and trained with a re-weighted variant of the evidence lower bound (ELBO):
rr arg min Σ (1 1--6 --a)) Epan (n) Bpa (ā(m) [I
V log Pa, (5 =)][[]]
After getting the optimal model sg- (i, samples can be generated by starting from uy N(0, 7) and following the estimated reverse Markov chain as:
T
v1 B (x+ + B,s0- (1, + Bie, i= N
(29)
Similar discussion as for the forward process, the equation (29) can further be rewritten as:
1
(1-1)) V1 Biss (DAAI + ((iAt) sø- ((tt, + VBi.,.
(1 BIaIst) (IIS: + B(iAt) -g- (iAt, - V Bizs.
(30)
2
2 (1
2 BISIAI) + At Sg- (ist, + BBz
when Bist is small. This is indeed the time discretization of the backward SDE:
- G + 3T t) (T - t, x, dt + /B(t)dBB. ax; ax; =
(31)
= =--((T - -1, + (T - t) (T t, xT) dt + g(T -t)dB,
12

---

Score as Action: Fine-tuning Diffusion Models by Continuous-time Reinforcement Learning
A.2. DDIM sampler is the discretization of ODE
We review the backward process in DDIM, and its connection the probability flow ODE following the discussion in (Song et al., 2021b; Kingma al., 2021; Salimans & Ho, 2022; Zhang & Chen, 2022).
(i) DDIM update rule: The concrete updated rule in DDIM paper (same as in the implementation) adopted the following rule (with = 0 in Equation (12) of (Song et al., 2020)):
(6) vI Cg (Ii))
It-1
v
2t
V
(t) ã Eo (Iz)
(32)
Ott
"direction pointing to 2,"
"predicted zio"
To show the correspondence between DDIM parameters and continuous-time SDE parameters, we follow one derivation in (Salimans & Ho, 2022) by considering the "predicted ro": note that define the predicted zo paramcterization as:
ia (1,
VI 0eeg (x) or, .0
(6) (x)
I v vā (t,
Van
vI āt
above (32) can be rewritten as:
/1 - ōc-1
(x+ Vānto (t, (t) + vāa-i 20 (t,x)
(33)
V1 αt
Using parameterization 0e v1 - ar and 04 = -ot, we have for t. - = < t:
X, ot [X, G (t, X,)l + agão p (t, X,),
(34)
which is the same derived in (Kingma al., 2021; Salimans & Ho, 2022).
A.2.1. ODE EXPLANATION BY ANALYZING THE DERIVATIVE
We further assume a VP diffusion process with af = 1
- of = sigmoid (Ax) for A log [o]/o?] in which X is known
as the signal-to-noise ratio. Taking the derivative of (34) with respect to Xę, assuming again a variance preserving diffusion process, and using das
foros and da dx
a2 gives
dA
Xx.
d
[X, ar Xt)l das
to (t, Xt)
dig
d)g di
d)g
1 0
1
a
[x - anto (t, Xi)l + aa Xt).
2 ox
2
Evaluating this derivative at s = then gives
Xx
1
- 1 ax20 (t, +
+20ll+a00000 _0, (1,
dig
s=t
2 12 a2 (t,
2 ax (1 -a3) do (t. Xx)
(35)
Jario e (t, Xx) azXal.
Recall that the forward process in terms of SDE is defined as:
dX, = ft. dt) + g(t)dB,, E [0,
and (Song et al., 2021b) shows that backward of this diffusion process is an SDE, but shares the same marginal probability density of associated probability flow ODE (by taking T f):
1 dx, = [(, } 9 v (t))x.logp(, x3) dt, t te (T.0)
13

---

Score as Action: Fine-tuning Diffusion Models by Continuous-ting Reinforcement Learning
where in practice v (l) is approximated by learned denoising model using
Va a 1o = se(t, 0r00 (t, x) - x
(x)
(36)
of
with two chosen noise scheduling parameters and os, and corresponding drift term f(t, = dlog and diffusion term (1)
do?
2 ddllog 0i of.
de
Further assuming vP diffusion process with ot = 1 of = sigmoid (Ax) for X log [aî/o /o.) we get
dlog dt
d logog a
1 (1 - a/
d)
1 2 dd
2((,2)
dt 2
dd dt
2
dt
2 1
dt
1.
Similarly, we get
?
d log at
d
dd
di
di
di
(t)
2
o
of
(ot
-o,)
= 0
di
di
dA
di
di
dl
dl
di
Plugging these into the probability flow ODE then gives
dX,- =
S(1, X,) - 2 log ((, di
(37)
of [X, V logp(t, Xx)l dde.
Plugging in our function approximation from Equation (36) gives
1
Xx a,00 (t, (t, - X1 -xx) ) ddi
dX.-
2
of
1
larã (t, X1) + -1) X.] dd
(38)
2
1
larto (t,, X,) - az Xe] dd,-
2
Comparison this with Equation (35) now shows that DDIM follows the probability flow ODE up to first order, and can thus be considered as an integration rule for this ODE.
A.2.2. EXPONENTIAL INTEGRATOR EXPLANATION
In (Zhang& & Chen, 2022) that the integration role above is referred as "exponential integrator" of (38). We adopt two ways of derivations:
(a) Notice that, if we treat the ãe (t, Xt) as constant in (38) (or assume that it does not change w.r.p. along the ODE trajectory), we have:
ax, + 1 afxx so (t, Xt) Laddde 2
1
(39)
Both sides multiplied by 1/0, and integrate from l to yields:
X,
X. = so (t. Xi)
(cxp(>>)))
1 xo) (x) = Do (t,
a2
at
(40)
which is thus
X, = 0+ Xx + a4 0z 0+
to (1, Xx),
(41)
which is the same DDIM continuous-time intepretation as in (34).
(b) We also notice that can also simplify the whole proof by treating the scaled score (same as in (Zhang & Chen, 2022)):
a,30 (t,x))
o, V log
0))
(42)
14

---

Score as Action: Fine-tuning Diffusion Models by Continuous-time Reinforcement Learning
constant in (37) (or assume that it does change w.r.p. along the ODE trajectory). Notice that from backward ODE, we have:
1 o X V.V p(t pttt,t,g Xi) dx,
dde.
(43)
2
oi
Both sides multiplied by 1/ and integrate from / to yields:
X
Xx
[1 (t, Xi
(-
0z + Tt
(44)
at
0t.
ax
Ott
which is thus
Ge X1 +
0e
X,
io (t, Xx),
(45)
0+.
which is the same DDIM continuous-time intepretation as in (34).
As a summary, treating the denoised mean or the noise predictor as the constants will both recovery the rule of DDIM. Usually, for ODE flows, the denoised mean assumption naturally holds; however, why the scaled score leads to the same integration rule remains to be an interesting question, probably comes from the design property of DDIM, see e.g. discussions in (Karras al., 2022).
B. Theorem Proofs
B.1. Proof of Theorem 3.1
The main proof technique relies on Girsanov's Theorem, which is similar to the argument in (Chen et al., 2022). First, we recall a consequence of Girsanov's theorem that can be obtained by combining Pages 136-139, Theorem 5.22, and Theorem 4.13 of Le Gall (2016).
Theorem B.1. For t [0, 7, let L, = S by dB, where B is Q-Brownian motion. Assume that EQ Jo l|b.1l ds < co. Then, a Q-martingale in (Q). Moreover, if
1 (QE(C)== 1, where E(C); := exp (f.r b. dB.
2 l' |1b.11 ds ds)
(46)
then E(C) also a Q-martingale, and the process
t + By S by ds
(47)
a Brownian motion under P == E(C)QQ. the probability distribution with density E {C)r w.r. Q.
If the assumptions of Girsanov's 's theorem are satisfied (i.e., the condition (46)), we can apply Girsanov's theorem to Q the law of the following reverse process (we omit e for brevity),
dX, (- (T -t)0, - X() dt + - t)dBr, Xo p.o(.)
(48)
and
= -[ - ( -
[T x,, 80,mo (,
(49)
where /. E [0, T). This tells us that under P= (, there exists a Brownian motion (B:)eeto.r s.t.
dB, g(T t) [s X,) sepro (T t, Xe]] dt + d.3
(50)
Plugging (50) into (48) we have P-a.s.,
dx, = (-J +2) -t)so(T x,),)) g t)d3r, Xo pool.)
(51)
15

---

Score as Action: Fine-tuning Diffusion Models by Continuous-timg time Reinforcement Learning
In other words, under P. the distribution of X is the the distribution generated by current policy parameterized by Le., po(.) Pr Q. Therefore,
DEL (pollpopre) Ep In dPr Ep, E())
dQr
2 / ||b. 2
1
EErr
by dB.
Epr
by dB. +
2 ||b.11
1
(T 1) Er |0 (T t, Xx) 80gna (T (1 xx) di
2 Ju
Thus we can bound the discrepancy between distribution generated by the policy 0 and the pretrained parameters Ppe as
1 DKI, (poo,po S S /'' g² - t)
(52)
B.2. Proof Theorem 4.1
First we include the policy gradient formula theorem for finite horizon in continuous time from (Jia Zhou, 2022b):
Lemma B.2 (Theorem 5 of (Jia & Zhou, 2022b) when R = 0). Under some regularity conditions, given an admissible parameterized policy no, the policy gradient of the value function V (t, x; #0) admits the following representation:
8 X.²) (s, ;m v(1,;;9) =E² J: (( { log " (a
( 30
00
(53)
[R(s, a (s, XT= x0)) ds)} Xi = 20. (t,x) E [0, Rd
-
which we denote the regularized reward
rR(1, Xi" = - 50 (1, x.)11
First, by applying Itô's formula to Xt), we have:
0v
1
0v
aV X,)
X,)
(1)
(1, Xi) di + oz (1, X,)
(54)
0t
d
Further recall that:
0V g(t,r,a;))= Øt (t, + 76 t,z,a, dr (t,z;æ), Ox2 (t.x;)) -BV (t,
0V
(55)
this implies that (similar discussion also appeared in (Jia & Zhou, 2023))
q (t. X7 a7; TT) dt dJ (f, x); + r X
X7, (.) dt BJ (t, dt (...B...
(56)
Plug this equality back in (53) yiclds:
i 00 v(t,z;=') =BP {." (((--) ② log x (a,
|s, X3" ) (t, X7 :) ds Xx" (57)
00
Let f 0, B a and further taking expectation the initial distribution yields Theorem 4.1.
16

---

Score as Action: Fine-tuning Diffusion Models by Montinuous-time Reinforcement Learning
B.3. Proof of Theorem 4.4
We recall the definition of terms and rewrite the bound in Theorem 4.4 for convenience. It suffices prove that there exists constant 1j and C2 such that
V - | ScI E l
(e/(" KL( (-11, (1)
(58)
and
T ( (-1t, (
(0) - -- s C2.
KI(
(59)
To prove (58), we need to utilize the KL-regularized objective we study for this paper, as this might not hold under general cases. Explicitly writing the definition of 8 and V we have (excluding contents e for simplicity):
v V ((()M(X))) E(RM (x+)) +3 KL. (p' (T. )1p pre (T, KI. (T, (T, 11p (T,
reward difference
(ii KL difference
For bounding (i), we have:
ii - -((. S - I N KKL (T' (T, :))
where the last equality due to Pinsker's inequality. For bounding (ii), standard identity for Kullback-Leibllr divergences gives
(L 8 ( KL(P" (|p" -KL(p" 11 /LS( p((x)] In (mim
(((x)
) dx.
Hence, by the triangle inequality,
kL" ) KLL [| S KL(x"111 + 1 10-- -p") ln( fimm))
By the bounded log-ratio assumption,
In (
(x)
) - C2
so
1100 (pp-pp) -p²) Inf (rmn)
S C3 1|Ie 10
Finally, Pinsker's inequality implies 110² - p0112 S v 222 |1x²)
Combining these bounds yields
(L' pp (p | s KL(P" l1 + Ca 2KL pd).
Further given the assumption that pd close to pd such that KL(P 11 < 1, we have
1
BCCCSVVKA /KL(P )|
Further applying the same argument in Theorem 3.1 proves (58). To prove (59), it suffices to bound
- (aflt, X!) (afl,, x q(t, X ai )d
E
(60)
17

---

Score as Action: Fine-tuning Diffusion Models by Continuous- time Reinforcement Learning
By importance sampling, have
(af x2 T (a x() 9(l, X af =0)] Eøla(t, X2 ap;a af;
Eg
Hence the left side equals So Egla(- ]dk, Assume that |al S M, we get
B Ro [a] I M 1| 1 M (2 11
by Pinsker's inequality. Therefore.
Esla(-..) < Eola(--.)] + M V 22L((T ((0))
Integrate over from 0 to T, and the result follows that:
$
(L) v0 s Co IE I (1,, d()) s Or 1 % KL (1, (1)l (16,
(61)
where the last inequality follows from Jensen's 's ineqality and Cauchy-Schward inequality.
C. More Experimental Details
C.1. Comparison with other discrete-time baselines for diffusion models fine-tuning
We also include the comparison with non RL based discrete time diffusion models fine-tuning baselines, DRaFT and AlignProp. The implementation of DRaFT is modified from AlignProp given their similarity. As in Figure 10, CTRL method outperforms our discrete time baselines.
1.9
1.8
Reward Mean
1.7
1.6
CTRL
DDPO
1.5
DRaFT
AlignProp
1.4
10
20
30
40
50
60
Number of epochs
Figure 10. Performance of continuous-time RL against discrete-tine time baselines under the same 50 discretization timesteps.
C.2. Algorithm Pseudocode
We present the algorithm pseudocode as below.
18

---

Score as Action: Fine-tuning Diffusion Models by Continuous-time Reinforcement Learning
Algorithm CTRI. for Diffusion Models
Input: Initial policy parameters Oo, value function parameters do (after pretraining), clip parameter learning rates Old. ag, number of epochs K, number of trajectories N, a fixed exploration level a, number exploration actions M, scaling parameter 17.
for 0, 1, 2....... do
Collect N trajectories {(li, X dz P ) by running policy 8M
Compute returns RG = Ei=,, 1d for each trajectory Initialize value function dataset Dv = {(t,, x R.))
# Train value function
1.2..... K do for k
Sample batch By C Dv Update Φn+1 t- On a V 1 Byl BBvll
Me.,, (V, -R)2
end for
# Compute advantage rate function estimates
for each (t,, x.. .a..) in collected trajectories do
Sample M samples of random noise Ej ~ ~0,0 T). 8n
Compute M pseudo samples d(..
8n = dt, + oej.
Compute advantage C.... (V(t,, + n (T t)e,) - V(t.. /n end for
Initialize policy optimization dataset Dz = [tt,, X(a A:i {(rg)} # Update policy using PPO objective for = 1,2,..., K do
Sample batch B,, Dx
* <o (0,311)))))
Compute likelihood ratios pl.
for all (t, x(" aty (b)) E Br
Compute clipped objective:
I(0)- |+ (t, )...).. min(pl 8n clip(pf,j, 1- 1, 0g7
Update 0 < + as VoL
end for
0,- f- 0
end for
19


Reviews:
[{"review_id": "VSjJC5aN5G", "reviewer": "Reviewer_b3WA", "summary": "This paper proposes a novel continuous-time reinforcement learning (RL) framework for fine-tuning diffusion models, reframing the learned “score” function as the policy’s action. Unlike discrete-time RL methods for text-to-image (T2I) or other diffusion settings, the paper leverages the underlying continuous-time SDE nature of diffusion models to avoid artifacts from fixed time-discretization and to better exploit structural properties.", "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": "We sincerely thank you for providing thoughtful suggestions to improve our paper. Please find our responses to your questions below:\n\n ***Q: Comparisons to other important baselines, such as RAFT and DRaFT***\n\n***A***: We conduct additional experiments comparing our method, DDPO, DRaFT [1] and AlignProp [2]. Since DRaFT code is not released, we implemented it by revising the codebase of AlignProp. Please find the reward curves comparisons [in the link here](https://ibb.co/KxX6SxhF). Our experimental results show that DRaFT and AlignProp perform similarly, but they underperform our CTRL approach, or even DDPO when using ImageReward as both the reward signal and the evaluation metric. \n\nWe would like to comment that we agree these are important diffusion models fine-tuning baseline methods, but our paper primarily focuses on RL-based methods, so we didn't pursue these baseline in our current submitted version. Nevertheless, we agree that these comparisons could strengthen our results, and we will update the experiments of these non-RL baseline methods in the revised version. Due to the time constraints, we haven't tested the RAFT [3] based on ImageReward, but we will also add the corresponding results and comparisons in the later version.\n\n***References:***\n\n[1] DRaFT: Directly fine-tuning diffusion models on differentiable rewards\n\n[2] AlignProp: Aligning text-to-image diffusion models with reward backpropagation\n\n[3] RAFT:  Reward rAnked Fine-Tuning Algorithm", "claims_and_evidence": "Yes, the claims proposed are supported by experiments shown in Section 5.2.", "methods_and_evaluation_criteria": "Yes. The proposed methods make sense for the problem and the benchmark models are also reasonable (DxMI and SD v1.5).", "theoretical_claims": "I briefly checked the proofs and they look sound.", "experimental_designs_or_analyses": "The experimental design is generally strong:\n\n- With T=10 steps, they show that continuous-time RL plus carefully chosen sampling steps can drastically improve generation quality (lower FID) compared to a baseline IRL approach.\n\n- Stable Diffusion Fine-Tuning: They measure the ImageReward score across training and across different sampling step counts (25, 50, 100) to confirm time-discretization invariance and consistent improvement brought by the continuous-time RL.", "relation_to_broader_scientific_literature": "The paper belongs to topics related to RL for diffusion models. The paper references existing results on continuous-time RL. The authors refine these ideas for the specific scenario of KL-regularized reward and time-homogeneous diffusion coefficients.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n\n- By treating the score as an action in a stochastic control problem, the approach neatly sidesteps time-discretization artifacts.\n- The paper presents a clear policy-gradient theorem (Theorem 4.1) and a TRPO/PPO-inspired surrogate bound (Theorem 4.4) specifically for continuous-time, state-independent diffusion coefficients.\n\nWeaknesses:\n- The paper only compares its method with DDPO while missing comparisons to other important baselines, such as RAFT and DRaFT, which makes the results less convincing.\n- The paper uses only SDv1.5 as the base model; including more models for comparison would improve the evaluation.", "other_comments_or_suggestions": "N/A", "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": {}}, {"review_id": "AqdQl5AojV", "reviewer": "Reviewer_4zQ6", "summary": "The paper reformulates fine-tuning diffusion generative models as a continuous-time reinforcement learning problem by treating the score function as a control action in the backward SDE. This continuous-time framework enables direct computation of policy gradients and leads to novel continuous-time analogs of popular RL algorithms (like TRPO/PPO), which are shown to yield more robust and efficient fine-tuning. Additionally, the paper proposes a value network design that leverages the structural properties of diffusion models, and experiments on CIFAR-10 and Stable Diffusion v1.5 demonstrate improved convergence and generation quality over traditional discrete-time methods.\n\n## update after rebuttal\nI confirm my score. Authors addressed comments and added clarity and results to the original submission.", "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": "We sincerely thank you for providing thoughtful suggestions to improve our paper. Please find our responses to your questions below:\n\n ***Q1: add pseudocode, figure of method***\n\n***A1***: Please find the pseucode of our algorithm [here](https://ibb.co/KjhQ4s3m). We will update this part to the revised version.\n\n ***Q2: provide detaild description and more evidence of \"Continuous-time RL outperforms Discrete-time RL baseline methods in both efficiency and stability because CTRL only require time discretization for estimating the policy gradient.\"***\n\n***A2***: As illustrated in the earlier Section 4, we are esstentially motivated to approximate the policy gradient formula in Equation (23), since PPO can be understood as more efficient implementations of such policy gradient by introducing the surrogate loss that shares the same gradient and clipping mechanism. This is shown in the definition of loss objective in the [pseucode](https://ibb.co/KjhQ4s3m). \n\nAlso we showed in the experiments of fine-tuning stable diffusion v1.5 (the second environment) as our method yields higher rewards training on the same amount of samples compared to DDPO, and the variance of the reward curves also appear to be much smaller. That's why we claim \"efficiency and stability\", and we will add this heuristic explanation part in the revised version.\n\n***Q3: only compared with 2 models, and add more set of the number of discretization timesteps compared with DPPO.***\n\n***A3***: We conduct additional experiments comparing our method, DDPO, DRaFT [1] and AlignProp [2]. Since DRaFT code is not released, we implemented it by revising the codebase of AlignProp. Please find the reward curves comparisons [in the link here](https://ibb.co/KxX6SxhF). Our experimental results show that DRaFT and AlignProp perform similarly, but they both underperform our CTRL approach, or even DDPO when using ImageReward as both the reward signal and the evaluation metric. Due to time constraints, we haven't tested the more set of the number of discretization timesteps of our methods compared with DPPO since they would require more hyparameter tuning for a fair comparison, but we will also add the corresponding results and comparisons in the later version.\n\n***Q4: How long does it take of finetune SD v1.5 with CTRL?***\n\n***A4***: The current version of our implemented CTRL takes twice as much time as DDPO for our current implementation, since it requires the value network update and inference which is similar to PPO vs REINFORCE methods in LLMs.\n\n***References:***\n\n[1] DRaFT: Directly fine-tuning diffusion models on differentiable rewards\n\n[2] AlignProp: Aligning text-to-image diffusion models with reward backpropagation", "claims_and_evidence": "In Section 5.2, the paper mentioned that “Continuous-time RL outperforms Discrete-time RL baseline methods in both efficiency and stability because CTRL only require time discretization for estimating the policy gradient. ”, which should provide detaild description and more evidence.", "methods_and_evaluation_criteria": "The methods and evaluation criteria are well-suited to the problem. The continuous-time RL formulation aligns naturally with the intrinsic continuous nature of diffusion processes, and evaluating on benchmarks like CIFAR-10 and Stable Diffusion v1.5 is appropriate for demonstrating improvements in fine-tuning generative models", "theoretical_claims": "There are no apparent flaws in the theoretical arguments as presented.", "experimental_designs_or_analyses": "**Section 5.1**\n\nCompare DxMI with CTRL: DxMI is an inverse reinforcement learning based method, which is different from the direct use of reinforcement learning fine-tuning.", "relation_to_broader_scientific_literature": "The paper related to the field of alignment, diffusion model and reinforment learning.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "**Strengths:**\n\n- A novel continuous time RL finetuning framework\n\n**Weaknesses:**\n\n- Clarity of the paper: add pseudocode, figure of method.\n- Small experiments: only compared with 2 models, and add more set of the number of discretization timesteps compared with DPPO.", "other_comments_or_suggestions": "Add more experient to show CTRL is practical.", "questions_for_authors": "How long does it take of finetune SD v1.5 with CTRL?", "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": "j1n5ypbY2m", "reviewer": "Reviewer_2RBC", "summary": "The paper proposes a continuous-time RL approach for fine-tuning diffusion generative models. It develops a policy optimization framework tailored for continuous-time RL, and empirically validate on the Stable Diffusion v1.5 model.", "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": "We sincerely thank you for providing detailed and thoughtful feedbacks to improve our paper. Please find our responses to your questions below:\n\n ***Q1: Clarification on the motivation in Figure 1 and Continuous-RL framework***\n\n***A1***:  We are sorry if our wording has caused confusion, but we didn't claim on such of \"mode collapse\". What we claimed  or observed is: the discrete RL trained policy, as a diffusion model, \"overfits\" to using 50 time steps as inference steps, as illustrated in Figure 1. Diffusion models has a nice natural tradeoff of generation quality and speed, as more inference steps could typically lead to more satisfactory images while causing extra time for inference. The most advanced models typically still haves such tradeoffs. We are motivated to let the fine-tuned models not lose such property (which is a dimension of freedom for users of different purposes), as it's pitiful that the performance improvement brought by RL fine-tuning like DDPO when inferencing the fine-tuned models on other time discretization steps like 25 and 100 could be discounted compared to the improvement typically reported only on fixed time steps like 50.\n\nComparing our framework to [1] and [2], we agree that the continuous-time control formulation already exists and is similar to our learning objective, but 1) we adopt different parameterizations that focus on treating the score function as the first class citizen and as the primary action space, which motivates us to incorporate the structural property in value function design, 2) the methods in these two works are pure stochastic optimal control methods, which are fundamentally different from our RL-based approaches. The effectiveness and scalability of RL methods have been largely investigated in LLM fine tuning which makes them quite popular, and we believe that our work can help improve the study of RL-based fine-tuning methods for diffusion models by the continuous-time framework with better leverage of structural property and thus resulting scalable, more efficient and more stable algorithm.\n\n ***Q2: The implications and interpretations of Figure 2***\n\n***A2***: To summarize, Figure 2 and Table 1 mainly ablate on how to smartly incorporate the diffusion model policy network in the value function design space and partial effects of the link function. As we already theorectically explained our motivations after Equation (20), we conduct experiments in Figure 2 and Table 1 to empirically showcase how things differ when putting the denoised mean $\\hat{x}_{\\theta}(t,x)$ in the first reward mean predictor term or the second residual corrector term. Our experiments indeed showcase the promise of incorporating such structural property in our design space, as shown by the reduced MSE of green line compared to the orange line. Incoporating the denoised mean in the second part almost does not improve performance while introducing extra computational burden. The performance of the green line and purple line also showcases that the choice of link function also matters. We agree that there could be other link functions that might outperform sin and cos, however it's impossible for us to exhaust all. We would like to pursue possible theorectical guidance of chosing this as future work.\n\n***Q3: Comparisons against common RL-finetuning baselines such as DRaFT and AlignProp.***\n\n***A3***: We would like to remark that DRaFT and AlignProp are algorithms for fine-tuning diffusion models, but they are not RL-finetuning methods, which might partically deviate from our primary study on RL-based methods. Nevertheless, we agree that they are baselines for fine-tuning diffusion models and we conduct additional experiments comparing our method, DDPO, DRaFT and AlignProp. Since DRaFT code is not released, we implemented it by revising the codebase of AlignProp. Please find the reward curves comparisons [in the link here](https://ibb.co/KxX6SxhF). Our experimental results show that DRaFT and AlignProp perform similarly (the time costs are also similar), but they both underperform our CTRL approach, or even DDPO when using ImageReward as both the reward signal and the evaluation metric. This demonstrates the advantage of our proposed method.\n\n***Q4: If continuous RL can further boost a reward direct propagation-based method***\n\n***A4***: Thank for raising up a direction that could be insightful. Since DRaFT and AlignProp are not RL-based algorithms, we agree and believe that such methods could have continuous time extensions. This is by high level related to our paper of utilizing continuous-time RL over discrete RL, but we believe that the overall framework will be quite different and will not be RL-related anymore. We would add this point to possible direction which can further showcase the insights brought by our framework; we would like to pursue this as future work.", "claims_and_evidence": "1. Most importantly, it is still not clear to me why continuous RL is a preferred framework. Finetuning diffusion models in continuous SDEs has been widely studied in the literature. But regarding continuous RL, it is still not made entirely clear why it is needed.\nIn the Figure 1, it is claimed that \"discretizing into 50 steps has more overfitting than discretized with 25 or 100 steps\". This is confusing. If I understand correctly, overfitting refers to 'mode collapse'. However, such a claim is not evident. How many runs have you conducted to draw this plot? Is it just for one run? Besides, the reward differences are not significant. Finally, even if this plot is properly drawn, what is the goal of this plot anyway? As far as I am concerned, in practice no matter what continuous formulation is used in modeling, in the final inference stage you would need to choose a discretization number of steps to denoise samples. \nIn Figure 7, it seems the claim is model fine-tuned by continuous RL and does not exhibit performance difference when using discretization steps. But it is not entirely convincing to me as the experiments can have randomness. Besides, in practice if one discretization step works well it is already sufficient, what is the actual motivation to have several discretization steps with similar performances?\n\n2. The implications and interpretations of Figure 2 are confusing. The authors compare several value network architectures. Firstly the training of value networks with MC across time steps is standard in RL finetuning. The novelty seems to be showing the representative power of introducing sin functions. However I dont think this representativeness can be definitely generalized to all domains of diffusion models. For example if the goal is to train a value function for sequence diffusion model, there are some embedding forms that might outperform sin and cos. Such a randomness and specialty makes me doubt the meaning of showing figure 2 and table 1, which are not generalizable.", "methods_and_evaluation_criteria": "The methods and evaluation criteria generally make sense", "theoretical_claims": "Claims are mostly sound.", "experimental_designs_or_analyses": "A major concern is about the baseline selection. If the continuous-time RL method proposed in this work intends to demonstrate superiority over existing fine-tuning methods, only comparing against DDPO appears insufficiently justified, as DDPO is relatively slow and computational heavy (Besides, Figure 9 does not seem to indicate a significant difference).\nWhy not include comparisons against common RL-finetuning baselines such as DRaFT and AlignProp , which are already effective on simple \"discrete RL\"? Even if the original reward is nondifferentiable, the parametrized value functions would be differentiable so implementation would not be a problem. It is advised to test if continuous RL can further boost a reward direct propagation-based method.\n\nAdditionally, the authors completely omit comparisons to a line of inference-time techniques, i.e., Sequential Monte Carlo (SMC)-based methods [1-2]. These inference-time techniques are critical baselines since they do not even incur the computational overhead associated with fine-tuning and can still have good performances. At the cost of fine-tuning, it is necessary to show better performances over the finetuning-free ones. \n\nAt minimum, clarifying clear superiority over these other potential baselines would substantially strengthen the importance of using 'continuous RL' versus ordinary \"discrete RL\".\n\n[1] https://arxiv.org/pdf/2402.06320 [2] https://arxiv.org/pdf/2408.08252", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strength:\n\n1. I appreciate the theoretically sound framework of policy optimization in the continuous RL formulation presented in Section 4. The contents are self-contained and implicative for future potential works studying continuous RL settings.", "other_comments_or_suggestions": "1. Eq 10-12 builds that finding the optimal score becomes a policy optimization problem in RL. However, this is known in the literature; see examples [1-2]. References are lacking. Additionally, how are these results different from the known results? Similarly, Thm 3.1 is also a known result.  These take up much space in the main text.\nCan the authors comment on why they are novel and what the underlying difficulty of proving them is?\n\n2. It is recommended to enhance interpretations of results (make figures and tables clear in their limitations as much as possible to avoid over-statements that might not definitely be true in other scenarios), add essential experimental baselines (at least explain why some important ones are missing).\n\n[1] https://arxiv.org/pdf/2402.16359 [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": {}}, {"review_id": "gcKpluKOkt", "reviewer": "Reviewer_HDpz", "summary": "The paper proposes a continuous-time reinforcement learning framework to fine-tune diffusion models by treating score functions as actions in a stochastic control problem. The authors derive a policy gradient theorem for continuous-time RL and connects KL regularization to a tractable running reward. The authors also develop a policy optimization theory by adapting ideas from TRPO/PPO to the continuous setting.\n\n## update after rebuttal  \nThanks for the response. I will keep my score.", "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": "We sincerely thank you for providing thoughtful suggestions to improve our paper. Please find our responses to your questions below:\n\n ***Q1: intuitive explanations for the value function design in Equation (20)***\n\n***A1***: The design of our value function in Equation (20) is motivated to capture the structural property of score-based diffusion models (we explain in (a) below) and to meet the boundary condition of the value function (we explain in (b) below). \n\n- **(a)** Recall that the value function $V(t,x)$ is defined as the expected reward of final generated image (if regularization penalty parameter $\\beta$ is 0) conditional on the current time $r$ and state/latent $x$, so an natural and naive approximation of it can be the reward of the denoised image $\\hat{x}_\\theta(t,x)$ of the current state $x$, as the denoised image can be understood as \"a prediction of the image at the end of SDE trajectory\". This illustrate our first term reward mean predictor; of course such approximation to the value function is not accurate, so we use the second term with the same network architecture to further approximate the possible residual term which can improve the performance by reducing the mean square error. \n\n- **(b)** The value function has a boundary condition that at the final time $T$, $V(T,x)=R(x)$, so we choose two proper coefficient functions $c_{\\text{skip}}(t)$ and $c_{\\text{out}}(t)$ to satisfy such needs: we let $c_{\\text{skip}}(T)=1$ and $c_{\\text{out}}(T)=0$, thus the boundary condition is naturally satisfied. \n\nWe hope that this clarifies our design space and we are happy to take further questions.\n\n***Q2: How does the closed-form advantage-rate function to generate images with higher quality?***\n\n***A2***: We illustrate this mainly through the RL prospective as images with higher quality can be reflected as higher average reward for RL. As far as we know, the most relevant work about the conventional continuous-time PPO approaches is in [1], which estimates the advantage rate function through a model-free manner, but requires sampling many more examples to make the estimation accurate and algorithms stable. In our method, thanks to the closed-form advantage-rate function and further subtraction of second order term (which is hard to compute) as the baseline function (please find more detailed discussion about this in Line 309-322 after Theorem 1), we get a much simpler way to compute the advantage function which is also much more sample efficient, since we do not need to query the generative model for the next state to go. We believe the leverage of such structural property is one of the key reasons we get quite stable training reward curves as in illustrate in orange curve of Figure 9.\n\n[1]. Zhao et al. 2023. Policy optimization for continuous reinforcement learning. NeurIPS 2023", "claims_and_evidence": "The authors claim that the proposed continuous-time RL framework, specialized value function design and policy optimization algorithm can improve the quality of diffusion generative model’s output. Empirical results on both small-scale (CIFAR-10) and large-scale (Stable Diffusion v1.5) settings provide convincing evidence of the method’s effectiveness.", "methods_and_evaluation_criteria": "YES", "theoretical_claims": "The authors claim that their continuous-time policy optimization leads to a tighter bound and a closed-form advantage-rate function compared with conventional continuous-time PPO approaches. There is no obvious error in the proofs.", "experimental_designs_or_analyses": "The experiment design is mostly reasonable, but it would be better to elaborate more on the advantages of their policy optimization algorithm.", "relation_to_broader_scientific_literature": "Most of the prior works of leveraging RL to fine-tune diffusion generative models use discrete-time formulation. The paper’s key idea of connecting diffusion generative models to continuous-time reinforcement learning is novel.", "essential_references_not_discussed": "There is no essential reference not mentioned.", "other_strengths_and_weaknesses": "Strengths:\n1.\tThe idea of treating scoring function as action and aligning diffusion model fine-tuning with continuous-time RL is novel. \n2.\tThe theoretical derivations provide a robust foundation for the proposed method and help to connect with the continuous framework.\n\nWeaknesses:\n1.\tThere is limited discussion on the additional computational cost that might arise from adopting a continuous-time formulation compared to conventional discrete-time methods.", "other_comments_or_suggestions": "Nil", "questions_for_authors": "1.\tCould you please provide some intuitive explanations for the value function design in Equation (20)?\n2.\tThe authors claim that they provide a tighter bound and closed-form advantage-rate function compared with conventional continuous-time PPO approaches. How does this help to generate images with higher quality?", "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: **Summary:** This work investigates the use of continuous-time reinforcement learning (CTRL) for fine-tuning diffusion models. The authors present a rigorous theoretical framework that establishes connections between CTRL and diffusion fine-tuning, including several insightful lemmas characterizing the KL divergence between state distributions and its relation to the score function. Experimental results demonstrate that CTRL is both robust to discretization strategies and more sample-efficient than existing discrete-time baselines.

**Justification:** All reviewers agree that this work makes a meaningful contribution to the area of diffusion model fine-tuning. However, some reviewers noted confusion in terminology and figure captions. I recommend the authors clarify these points to improve the overall readability and accessibility of the paper.
Decision: Accept (poster)