Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control
Authors: Carles Domingo-Enrich
Michal Drozdzal
Brian Karrer
Ricky T. Q. Chen
Venue: ICLR 2025 Spotlight (2025)
Keywords: Reward fine-tuning
stochastic optimal control
flow matching
diffusion models
RLHF
adjoint method
Primary area: generative models
TL;DR: We introduce a reward fine-tuning framework for diffusion and flow matching models, based on stochastic optimal control (SOC), and Adjoint Matching, a new SOC algorithm.

Abstract: Dynamical generative models that produce samples through an iterative process, such as Flow Matching and denoising diffusion models, have seen widespread use, but there have not been many theoretically-sound methods for improving these models with reward fine-tuning. In this work, we cast reward fine-tuning as stochastic optimal control (SOC). Critically, we prove that a very specific *memoryless* noise schedule must be enforced during fine-tuning, in order to account for the dependency between the noise variable and the generated samples. We also propose a new algorithm named *Adjoint Matching* which outperforms existing SOC algorithms, by casting SOC problems as a regression problem. We find that our approach significantly improves over existing methods for reward fine-tuning, achieving better consistency, realism, and generalization to unseen human preference reward models, while retaining sample diversity.

Published as a conference paper at ICLR 2025
ADJOINT MATCHING: FINE-TUNING FLOW AND DIFFUSION GENERATIVE MODELS WITH MEMORYLESS STOCHASTIC OPTIMAL CONTROL
Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer & Ricky T. Q. Chen
Meta AI
cd2 cc27540nyu.eem edd /tt.r//aaaalppriaaaeeee.e...aoao
ABSTRACT
Dynamical generative models that produce samples through an iterative process, such Flow Matching and denoising diffusion models, have seen widespread use, but there have not been many theoretically-sound methods for improving these models with reward fine-tuning. In this work, we cast reward fine-tuning as stochastic optimal control (SOC). Critically, we prove that very specific mem- oryless noise schedule must be enforced during fine-tuning, in order account for the dependency between the noise variable and the generated samples. We also propose a new algorithm named Adjoint Matching which outperforms exist- ing SOC algorithms, by casting SOC problems as a regression problem. We find that our approach significantly improves over existing methods for reward fine- tuning, achieving better consistency, realism, and generalization to unseen human preference reward models, while retaining sample diversity.
Basc model (Flow Matching) w/ Guidance
Adjoint Matching (Ours)
Figure I: We introduce Adjoint Matching, a theoretically-driven yet simple algorithm for reward fine-tuning that works for large family of dynamical generative models, including for the first time, Flow Matching models. Text prompts: "Beautiful colorful sunset midst of building in Bangkok Thailand". "Beautiful grandma and granddaughter are mixing salad and smiling while cooking in kitchen", "The beautiful young woman in sunglasses is standing at the background of field and hill. She is smiling and looking over shoulder". "Chess, intellectual games, figure horse, chess board".

---

Published as a conference paper at ICLR 2025
INTRODUCTION
Flow Matching (Lipman et al.
2023) and denois- 2023 Albergo Vanden-Eijnden, 2023 Liu al
ing diffusion (Song & Ermon 2019] Ho et (2200 et al. 2021b; Kingma et al. 2021) models are being used for many generative modeling applications, including text-to-image (Rombach al.) Esser al. (024). text-to-video (Singer et al., (002)), and text-to-audio (Le et 20.4 Vyas eet al. 2023). In most cases, the base generative model does not achieve the desired sample qual- ity. To improve the generated samples, it is common to resort techniques such
2022
(lassffier-free guidance (Ho & Salimans) 2022; Zheng et al.) 2023) to get better text-to- soxt-te- aligmment, or to fine-tune using human preference reward models to improve sample quality and realism (Wallace
et al. 2023a Clark et al., 2024).
In the adjacent field of large language models, the behavior of the model is aligned to human prefer- ences through fine-tuning with reinforcement learning from human feedback (RLHF). Either explic- itly or implicitly, RLHF methods (Ziegler et al., 2020 [Siennon et al.) 2020 Ouyang et al.) 2022]
Bai et al. 2022) assume reward model r(z) that captures human preferences, with the goal of modifying the base generative model such that it generates the following tilted distribution:
p" (x) xc phase (x) exp(r(z)),
(1)
where Pbase is the base generative model's sample distribution.
Inspired by this, fine-tuning methods have been developed to improve denoising diffusion models
based human preference data; either using reward-based approach (Fan & Lee. 2023 Black
et al. 2024; Fan et al.) 2023 Xu et al. preference optimization (Wallace al. 2023a However, unlike the fine-tuning methods designed
2023 Clark et al. 2024 Uchara et al., 2024a aa b), or direct
for large language models, most of the existing methods to large degree ignore passe and focus solely on the reward model. Reward models can range from standard evaluation metrics such as ClipScore (Hessel et al.) [2021; Kirstain et al. 2023) to specialized models that have been trained on human preferences (Schuhmann & Beaumont 2022; Xu et al.) 2023 Wu et al. 2023c). As these are parameterized by neural networks, they fall pray to adversarial examples which lead to the generation of undesirable artifacts (Goodfellow et al. 2014 Mordvintsev et al. 2013). This has led some works to consider adding regularization during tine-tuning (Fan et al. [024]] Uehara al., 2024b) to incentivize staying close to the base model distribution; however, there does not yet exist simple, generic approach which actually provably generates from the tilted distribution (T).
The main contributions of our paper are as follows:
(i) We present a stochastic optimal control (SOC) formulation for reward fine-tuning of dynamical generative models. Importantly, we prove that the naïve approach considered by prior works lead to a value function bias problem that biases the fine-tuned model away from the tilted distribution (T]. This problem has also been observed by [eeaara al.] (2024b) but they propose a more complicated solution which involves training separate generative model.
(ii) Instead, we propose a very simple solution: the memoryless noise schedule. This is a unique noise schedule that completely removes the dependency between noise variables and the gen- erated samples, resulting in provable convergence to the tilted distribution. This allows us to fine-tune dynamical generative models in full generality, including being the first to fine-tune noiseless Flow Matching models.
(iii) We also propose new method for solving SOC problems, called Adjoint Matching, which com- bines the scalability of gradient-basee based methods and the simplicity of a least- squares regression objective. This can be applied to general SOC problems, beyond reward fine-tuning.
(iv) We perform extensive comparisons to baseline approaches, and analyze them from multiple perspectives such realism, consistency, and diversity. We find that proposed method provides generalization to unseen human preference reward models, better text-to-sampii con- sistency, and retains good diversity.
2 PRELIMINARIES ON DYNAMICAL GENERATIVE MODELS
We are interested in fine-tuning base generative models (X)) where samples are generated through the simulation of stochastic process. That is, these models transform noise variables into
2

---

Published as a conference paper at ICLR 2025
sample through iterative process. In particular, we discuss the specific constructions and sampling
processes of Flow Matching (Lipman et al. 2023 Liu et al. 2023 Liu 2022 Albergo Vanden- Eijnden 2023) and Denoising Diffusion Models (Ho et al. 2020; Song et al. 2021bba) The goal of this section is to provide background information on these methods.
Given random variables from an initial distribution Xo tributed according to some data distribution, we define the reference flow X =Xx)xe(oo1) where
= -0 1), and X1 which are dis-
X, = 3,X0 + X1,
where (at)rgfo.1). (Be)re(o.1] are functions such that ao B1 and = B. = 1. Diffusion models and Flow Matching construct generative Markov processes X, with initial distribution Xo ~ (o,, I) that result in flows X (Xe)xe(o.1) with the same time marginals as the reference flow X, i.e. the random variables X₁ and X, have identical distribution for all times E [0, 1j. This implies X₁ has the same distribution as the data distribution, so simulating the Markov process from random noise Xg is way generate artificial sample
Flow Matching. We focus on Flow Matching here, and defer the overview of denoising diffusion models (DDIM: [Song al. (2021a), DDPM: [H et al. (2020) App. In its simplest form, the generative Markov process of a Flow Matching model is an ordinary differential equation (ODE) of the form:
dX; u(X,, dt,
Xo ~ 10,
(3)
where ()), is a parametric velocity that is optimized match the derivative of the reference flow, i.e., f) arg min, (x,, aे Xell (see e.g. Lipman et al. (2023) for details on pre-training Flow Matching models). can then be proven that the solutron of the generative process
(3) has the same time marginals as the reference flow (Lipman et al. 2023 Liu 2022; Albergo &
Vanden-Eijnden 2023), and a commonly used choice 0g and p. consider family of stochastic differential equations (SDEs) with an arbitrary state- sute-independent diffusion coefficien
t. One can also
o())2 ax.- = (o(xe.t+ (. (x
a: x.)x) dt + (t)) dB,,
~ I), (4)
where (B:)220 is Brownian motion. The generative processes in and have the same time marginals. This can be seen by writing down the Fokker-Planck equations for (3) and (4), and observing that they are the same up a cancellation of terms (Maoutsa et al. 2020). The diffusion coefficient ol) in (4) is compensated by the second term in the drift.
Flow Matching in terms of the score function. continuous-time DDIM generative processes as:
We can unify both the Flow Matching and
dx,- = b(Xx, () dl + (1)
Xo ~ ~N(0 I),
(5)
where b(x, t) = k& + (t)) + sx, t),
Kt a. 1): = B B, 3))
(6)
where (ar, 3r) are coefficients of the reference flow (2). and s(x, is the score function cunction-defincl as the gradient of the log density of the random variable X₁. See App. C.4] and App. C.5 for the
derivation of (5) for DDIM and Flow Matching. In Subsec. [3.3] we rely on this characterization
to derive our fine-tuning procedure. This expression has been written before for DDIM, e.g. Bartosh
et al. 2024a b).
FINE-TUNING AS "MEMORYLESS" STOCHASTIC OPTIMAL CONTROL
We now discuss the crux of the problem: how to produce a fine-tuned generative model that pro- duces samples X, which follow the tilted distribution involving a reward model (I). An obvious direction is construct fine-tuming objective involving both the base generative model and the
In our derivations, we assume the base model has been trained perfectly during the pre-training phase. We use the common short-hand "over-dot" notation denote the time derivative, i.e., d; d 4z..
3

---

Published as a conference paper at ICLR 2025
reward model, where the optimal solution results in fine-tuned generative model for the tilted dis- tribution. However, as we will explain, this turns out to be non-trivial, because a naive formulation will introduce bias into the solution.
In Subsec. 3.1 we discuss the problem formulation of stochastic optimal control, a general frame- work for optimizing SDEs, and its relation to the maximum entropy reinforcement learning frame- work commonly used for RLHF fine-tuning. Next, in Subsec. 3.2 we discuss the initial value function bias problem which plagues existing approaches and so far has seen no simple solution. Finally, in Subsec. [.3] we propose a novel simple solution that circumvents the bias problem, by enforcing particular diffusion coefficient, the memoryless noise schedule, to be used during fine- tuning. This results in an extremely simple fine-tuning objective that provably converges to a model which generates the tilted distribution without any statistical bias.
3.1 PRELIMINARIES ON THE STOCHASTIC OPTIMAL CONTROL PROBLEM FORMULATION
Stochastic optimal control Bellmal 1957); Fleming & Rishel 2012 ; Sethi 2018 considers general optimization problems over stochastic differential equations, but we only need to consider a common instantiation, the control-affine problem formulation:
min E So (2 +) + f(X,,,)) t)) dt + +(x))
(7)
wEM
s.t. dxpp" = (b(x,, t) + (t) t)) dt + ot))dB,
X ~ po
(8)
where in X Rd is the state of the stochastic process, u R2 x [0, + Rd commonly referred to as the control vector field, : Rd x [0, 1] Rd is base drift, and [o, + Rdxd is the diffusion coefficient. These jointly define the controlled process X" ~ p" that we are interested in optimizing: often both b and a are fixed and we only optimize over the control u.
we have an affine control cost 1 [XX",XX, t)I?, ; a running state cost f: ddd x [0, 1] + IR and a terminal state cost g Rdd + R.
As part of the objective functional
The stochastic optimal control (SOC) objective 7 can be decomposed recursively from the final time value. It is common to define the cost functionall which is the expected future cost starting from state & at time t:
J(u x. := Ex-pp Exp
[ (2 (lllu(X, 8)1 + f(X,, 8)) ds + g(Xi) X,
(9)
From here, the value functiin is defined as the optimal value of the cost functional V(x. t) = minnet Ju, JI"" æz t), where u² is the optimal control, i.e., minimizer of Further- more, a classical result is that the value function can be expressed in terms of the uncontrolled base process y hase (Kappen (2005). see Domingo-Enn EnrrchhErcchhhh al. 2023 Eq. 8, App. B for self-contained proof):
V(r, 1) = log E Ex [exp(- S(X,, s)ds g(Xi)) X, =
(10)
A useful expression for the optimal control (which we will make use of in deriving the Adjoint Matching objective in Sec. 4 is that it is related to the gradient of the value function:
ut (x, () (() V VV, = (() V.J( .r,).. ().
(11)
Relation to MaxEnt RL. Stochastic optimal control with the control-affine formulation is the continuous-time equivalence of maximum entropy reinforcement learning (MaxEnt RL; Todorov 2006);; Ziebart et al.) (2088) with KL regularization instead of only an entropy regularization. In particular, by the Girsanov theorem (Thm. 2). the affine control cost equivalent to a Kull- back-Leibler (KL) divergence between the base process pase when process y", when conditioned on the same initial state Xo (see App. D.4p:
O. and the controlled
DEL (x"(p(XX
|| buse (X ( Ex"xp=
[ 1 [XX, aat
(12)
NNot that there is slight difference in terminology between SOC and reinforcement leamiig. where our cost functional is referred to the state value function and our value function is the optimal state value function in RI.
4

---

Published as a conference paper at ICLR 2025
resulting in the KL-regularized RL interpretation of 7
max E Xxoppo [Expr -x(;,, t)dt ()dt-g(X()] DRL (X I (xxeexXXXx) (x(x)))
uell
(13)
where the negative state costs correspond to intermediate and terminal rewards in the RL interpreta- tion. The KL divergence forces the optimal process to stay close to the base process.
3.2
THE INITIAL VALUE FUNCTION BIAS PROBLEM
We next discuss why naïvely adding a KL regularization does not lead to the tilted distribution From 13 we can show that the optimal distribution conditioned on Xu is
p'(X Xy) base X Xx) exp ( (- So [(X,,, di ((x)))
(14)
This is analogous to the exponentiated reward distribution in MaxEnt RL (Rawlik et al. 2013). but the entropy regularizer generalized to KL. regularization with respect to a prior distribution base In order to relate this to the tilted distribution ① that we want to achieve for fine-tuning, first notice that the normalization constant of the right-hame side (RHS) of (I4) is exactly the value function at t == 0:
Ex (X exp (- dt
exp (-v 0))
(15)
where the equality is due to (T)). Dividing the RHS of (14 by 15 and multiplying by po (Xo). obtain the normalized distribution over the full path X,
p" (X) pX) ( (- f(X,, t) dt 9(X1) + O)).
(16)
Setting f = and -r, we arrive at an expression for the optimal distribution
pX X1) p hase (Xo. X1) exp (r(X1) + V(Xo, 0)).
(17)
This unfortunately does not lead to the tilted distribution 1 because we have bias in the optimal distribution that due to the value function of the initial distribution V(Xg. 0). That to say, naively adding KL regularization (12) the fine-tuning objective in the sense of (13) leads to a biased distribution (16) after fine-tuning and is not equivalent to the tilted distribution ( For instance, when the sampling procedure noiseless, i.e., o()) = 0, fine-tuning naïvely will not have any effect because X, completely determines X1.
This is unlike the situation for large language models Ouyang et al. 2022 Rafailov et al. 202.3 where there is no dynamical process that samples X2 iteratively and hence no dependence on the initial noise variable Xo. Although this KL regularization is a common objective for RLHF of large language models, it has seen seldom use in fine-tuning diffusion models, likely due to this issue of the initial value function bias. In the context of diffusion models, KL regularization (13) has been explored in prior works (Fan et al. 2024). but its behavior was not well-understood and they did not relate the fine-tuned model to the tilted distribution (.. Another direction that has been proposed is to learn the initial distribution Pp to cancel out the bias (Uehara et al., 2024b) Tang. 2024) but this simply shifts the work into tilting the initial distribution and requires parameterizing the optimal initial distribution. In contrast, we show in the next section that it is possible remove the value function bias by simply choosing very particular noise schedule during the fine-tuning procedure.
auxiliary model for
3.3 THE MEMORYL.ESS NOISE SCHEDULE TO FINE-TUNE DYNAMICAL GENERATIVE MODELS
In this subsection, we propose a very simple method of turning (17) into the tilted distribution (T) through the of memoryless noise schedules. We provide an intuitive explanation of why such noise schedules are sufficient for fine-tuning, and show that if we want to sample the fine-tuned model with an arbitrary noise schedule, we must use particular memoryless noise schedule.
Intuitively, the main reason we cannot arrive at the tilted distribution from (7) is due to the base basee X1) distribution not factoring into Xo and X1. Hence, we define memoryless gener- ative process as follows:
NNot 14 is informal because densities over continuous-time processes are ill-defined; the formal statement exp(- f f(X,,t) t) dt s(Xi)), where dD" Tibuc denotes the Radon-Nikodym derivative.
dF
[X
We this formally in the proofs.
5

---

Published as a conference paper at ICLR 2025
Kig
Diffusion coefficient o(t)
Memoryless X
Flow Matching 3
du
9e St B.)
General (commonly 0)
No
91
Memoryless Flow Matching
dx
Br (1 Bo
B.)
22
Yes
DDIM 29
&1.
åt.
General (commonly 0)
No
28,
DDPM 30
de.
dx.
V2TH
Yes
Table I: Diffusion coefficient (t) and the factors Ki, T)t for the Flow Matching, Memoryless Flow Matching, DDIM, and DDPM generative processes. When the diffusion coefficient is o(t) v2n the generative process is memoryless, i.e., samples X1 will be independent of the initial noise Xo-
Definition (Memoryless generative process). A generative process of the form 55 66 is memory- less if Xo and X1 are independent, i.e., ypase (Xo, X1) = place (Xo)p (X1).
When the base generative process is memoryless, this implies:
p" (X)) = plsc (No) (x1 exp(r(Xi) + V(Xo, 0))XX, x p exp(r(X1)). (18)
That is, solving the SOC problem (7)
8 with a memoryless base model will result in fine-tuned model that generates samples p" ("x1) according to the tilted distribution (T). This memoryless property is not satisfied generally by the family of generative processes captured by 7 For instance, the Flow Matching and DDIM generative processes with zero diffusion coefficient (i.e. (1) 0) are definitely not memoryless due to Xo and X being theoretically invertible. Below, we provide the sufficient and necessary condition for the noise schedule in order to have memoryless generative process.
Proposition 1 (Memoryless noise schedules). Within the family of generative processes generative process is memoryless if and only if the noise schedule is chosen as:
5
60.
(t) = 2n, + x(t), where X : [0.1] + Risst Vt E (0, lime -0+ de cxp (- J
x(s) 23
ds)
= 0,
(19)
where ti defined in (6. In particular, we refer to (t) 22n as the memoryless noise schedule.
Due to the endpoint constraints of (a,, 3,) for the reference flow (2), the memoryless noise schedule (t) is infinite at f = 0 and approaches zero at f = 1. This provides a way for the generative process to mix when close to noise X, while stay steadying when close to the sample X1. Hence, the sample will have no information about Xo due to the enormous amount of mixing with large diffusion coefficient. Furthermore, while we have intuitively justified the memoryless noise schedule through its independence property, theoretical result is actually even stronger: all generative models of the form (5)-(() must be fine-tuned using the memoryless noise schedule. We formalize this in the following theorem, which prove in App. F.2
Theorem (Fine-tuning recipe for general noise schedule sampling). Within the family of gener- ative processes ()) in order allow the of arbitrary noise schedules and still generate samples according to the tilted distribution ( the fine-tuning problem
with f = and
g must be done with the memoryless noise schedule oI1) 22n
Thm. states that we need to use the memoryless noise schedule for fine-tuning with the SOC objective-or equivalently, the KL regularized reward objective (13). This is the only noise sched- ule that retains the relationship between the velocity and score function, allowing the conversion to arbitrary noise schedules (e.g., (t) = 0) after fine-tuning. It is worth noting that when using the memoryless noise schedule for DDIM, this recovers what we derived as the continuous- -ime limit of the DDPM generative process (30). However, the DDPM sampler (Ho et al. 2020) is not com- monly used as the DDIM sampler (Sog et al. 2021a) and Flow Matching models typically generate samples using a (t) = 0, so explicit conversion to the memoryless noise schedule is necessary for fine-tuning. Tab. | summarizes the memoryless schedule for diffusion and Flow Matching models, which we refer to as Memoryless Flow Matching, In Fig. we visualize fine-tuning a ID model, where we see that constant or (f) leads to biased distributions whereas the memoryless noise schedule perfectly converges to the tilted distribution (.. In App. E. express the base drift b and the control u in terms of the base and fine-tuned Flow Matching vector fields hase and finetune and do the same for DDIM.

---

Published as a conference paper at ICLR 2025
pbase(x)
exp(r(x))
(x) ()see(x)xxp((x)
(a) Pretrained FM 1 base
(b) Fine-tuned FM (luetonn(c)) Fine-tuned FM (lnetunc Fine-tuned FM 1 Anetune with (t) = 0.2
with (t) = 1.0
with memoryless ot)) /2m
Figure 2: Visualization of Thm. ssowing that fine-tuning must be done with the memoryless noise schedule to ensure convergence to the tilted distribution (I). (a) Shows the base Flow Matching model. (b, c) Fine-tuning using a constant (t) leads biased distributions. (d) Fine-tuning using the memoryless noise schedule leads to the correct tilted distribution. Note that sample generation use any noise schedule after fine-tuning, including (1) O.
ADIOINT MATCHING FOR STOCHASTIC OPTIMAL CONTROI.
Although several deep learning methods have been proposed to solve SOC problems using deep learning (Domingo-Enrich et al. 2023 Nüsken & Richter 202T), the preferred approach the ad- joint method, which performs gradient- based optimization on the control objective (7) with respect a parameterized control function. There are two approaches which yield the same gradient in the small step size limit (see App. F.1.I for more detail): the discrete adjoint method, where the control objective is discretized and an automatic differentiation engine used to backpropagate through it (Han & E, 2016), and the continuous adjoint method, where the continuous-time structure exploited by solving the adjoint ODE backwards in time:
da (t: X, u) ar
a(t: X, U)T (V
(f(X,, t) + fllu(X,, t)]
(20)
with initial condition (11 X, u) (X)) The gradient of the continuous adjoint loss is given by
dc d0 = 1/2 So 06 |u 1 + fo Bu(X,.)) 09
att a(t: X, u)dt,
(21)
The following proposition introduces and studies loss whose gradient is also 21 Proposition 2. Let us define, for now, the basic Adjoint Matching objective as:
L Adj-Mazch (u: X): [ (o x]l di, X pd, á stopgrad(u),
(22)
where u = stopgrad(() means that the gradients of u with respect to the parameters @ of the con- trol u are artificially set to zero. The gradient of Lrasie- Adj- Match (u; X) with respect to equal to the gradient dC in equation 21.. Importantly, the only critical point of [Caasic-Ade Match) is do the optimal control u*.
Critical points of C. are controls u such that 5 L()) O, where Su L denotes the first variation of the functional L. In other words, Prop. Prop.[]statee states that the only control that satisfies the first-order optimality condition for the basic Adjoint Matching objective is the optimal control, which provides theoretical grounding for gradient-based optimization algorithms. An intuitive way to understand the basic Adjoint Matching objective that it is a consistency loss. The Adjoint Matching objective is based off of the observation that the optimal control u= [x. is the unique fixed- point of the relation u(x t) = - -o()) Ju; x, t) (see Lemma 6 in App. F33 and so we are directly optimizing for a control that fits this relation, while using the adjoint state stochastic estimator of V. t).
7

---

Published as a conference paper at ICLR 2025
We can that the basic Adjoint Matching objective produces the same gradient w.r.t. 8 as the continuous adjoint method (2T) by expanding the square in (22) and removing terms that do not depend on f. And while the basic Adjoint Matching not entirely novel, it provides the means of deriving a simpler leaner objective function.
The "Lean" Adjoint. The minimizer of a least-squares objective is the conditional expectation of the regression target, so for the Adjoint Matching objective, at the optimum we have that
u' (z, t) = Ex~p" [-o(t)] a(t; X, xx) = z
(23)
Multiplying both sides by the Jacobian V (zz t) and re-arranging, get the relation
Ex-p" [u (z, un (I, at: X, u Vt) (x, X; = x|
(24)
Notice that the terms inside the expectation in 24 show up as part of the adjoint ODE 20 There-
fore, we motivate the definition of lean adjoint state à with the terms in (24) removed. Plugging this lean adjoint back into the squares objective, we obtain our final proposed Adjoint Matching objective:
Cadj (u; X) 1/' [acllllx.. (((,.tt) a(t) X)
1 dt. X ~ p", ū stopgrad(u))
(25)
where
a X) (ā(t; Veb(X,,)) VV, Vo
(26)
á(1: X Vxg(X)).
(27)
define the lean adjoint state, and 25 is the complete Adjoint Matching objec- tive. The unique critical point of E. Lddj Match is the optimal control, which prove relying on
Equations 26
(
Prop. 2 and equation (24) (see Prop. 7 in App. R44..
Compared to the adjoint method (App. F.I.I. Adjoint Matching produces a different gradient in expectation than the continuous adjoint. This is because the lean adjoint state is not related to the gradient of the cost functional anymore, i.e., (171) is not true, except at the optimum when u = u". Even at the optimal solution, since Adjoint Matching removes terms that have expectation zero, can potentially exhibit better convergence and lower variance than the continuous adjoint method. Additionally, computation of the lean adjoint state (26) also exhibits a smaller computational cost due to the removal of the extra terms (no longer need the Jacobian of the control Vz ()). We provide a rigorous derivation of Adjoint Matching and the above claims in App. F44
Adjoint Matching can be applied to reward fine-tuning of dynamical generative models through the memoryless SOC formulation discussed in Sec. 3 In App. F:5 we provide pseudo-code for Flow Matching models (Alg. T and for denoising diffusion models (Alg.
5 EXPERIMENTS
We experimentally validate our proposed method on reward fine-tuning a Flow Matching base model (Lipman et al. 2023). In particular, use the usual setup of pre-training autoencoder for 512x512 resolution images, then training a text-conditional Flow Matching model on the latent variables with U-net architecture (Long et al., 2015), similar to the setup in Romobach et al. (202)).
We pre-trained our base model using a dataset of licensed text and image pairs. Then for fine-tuning, consider the reward function: r(x) = À Rewa rdMode 1 (2)) corresponding to a scaled version of the reward model, which we take to be ImageReward (Xu et al. 2023). Different values of provide different tradeoffs between the KL regularization and the reward model (13). For evaluation and benchmarking purposes. we report metrics that separately quantify text-to-image consistency, human preference, and sample diversity, capturing the tradeoff between each aspect of generative models (Astolfi et al.] 2024). For consistency, we make use of the standard ClipScore (Hessel et al.) 2021) and PickScore (Kirstain et al. 2023); for generalization to unseen human preferences, we use the HPSv2 model (Wu et al. (223b); and for diversity, we compute averages of pairwise distances of the DreamSim features (Fu et al.] (023). More details are provided in App. H.4 As our baselines, consider the DPO (Wallace et al. 2023a), ReFL (Xu et al., 2023). and DRaFT-K algorithms (Clark et al. 2024). DPO does not use gradients from the reward function, while ReFL and DRaFT make of heuristic gradient stopping approaches to stay close to the base generative model. Out
8

---

Published as a conference paper at ICLR 2025
Fine-tuning
Fine-tuning (t))
Sampling a(t)
DreamSim
ClipScore
PickScore t
HPP v21
Method
Diversity
None (Base model)
24.15
17.25
16. 19.17
13.60
N/A
28.32-40.2
18.15
17.894a16
56.53
22
/21jt
30.18=0.24
19.38
24.61+0017
25.54+0099
DRaFT-1
30.95
19.37
24.37
27.39
Basclines
22
/2
26.94=2022
18.34
19.98
41. +1.98
DRaFT-40
30.07=0.3
19.45
24.06+00.4
36.53
27p
2
24.1 =0.22
17.24
16.15
53.27
DPO
10.11
17.92:1007
17.30
54.1
V2nc
/2m
28.59
18.68
22.24
ReFL.
30.06.00.3
19.07
23.06:1041
1..69
Cont. Adjoint 12500
22
26.99
18.3322066
20.83
46.59=1.An
22
29.4920.32
18.98:n
21.34-10.55
48.41
Mcmoryless SOC
Adjoint
/21jt
28.04
18.44-40.1
20.044:0.9
54.90++2.00
V2nc
12500
29.28
18.82
19.73
53.36--2.44
Adj-Matching 1000
/21
30.36
19.29
24.12+0.17
40.89
22m
31,41 40.22
19.57
23.29440.8
43.10
Adj.MMatching
221
30.59.0.40
19.49
24.85
37.07
/2
2500
31.6420.21
19.71 :0.09
24. 22:1..22
39.88:11.99
Adj.MMatc
30.62.+0.00
19.50
24.95
34.50+1155
/2
12500
31.65-=0.19
19.76
24.49+027
37.24
Evaluation metrics of different fine-tuning methods for text-to-image generation. The second and third columns show the noise schedules a()) used for fine-tuning and for sampling: (t) 22n corresponds to Memoryless Flow Matching, and (t) 0 the Flow Matching ODE
Table 2:
3). We report standard errors estimated over runs of the fine-tuning algorithm on random sets of 40000 training prompts, each evaluated over a random set of 1000 test prompts.
1000
1000 itrs.
2500
2000 itr..
15500
4000 itss
Adjoint Matching (Ours)
DRaFT-1
Figure Our proposed Adjoint Matching using the memoryless SOC formulation introduces much more principled way of trading off how close to stay to the base model while optimizing the reward model. In contrast, baseline methods such as DRaFT-1 only optimize the reward model and must rely early stopping to perform this trade off, resulting in much more sensitive hyperparam- eter. Samples are produced using (f) 0 with the same noise sample. Text prompts: "Handsome Smiling man in blue jacket portrait" and "Quinoa and Feta Stuffed Baby Bell Peppers".
of these baseline methods, we find that DRaFT-1 performs the best, so we perform additional ab- lation experiments comparing to this method. Within the same SOC formulation, also consider the discrete and continuous adjoint methods. We provide full experimental details in App. H an important implementation detail is that we slightly offset (t) in order to avoid division by zero.

---

Published as a conference paper at ICLR 2025
Adj. Match. w-0 DRaFT-1 w-0
Adj. Match. w-1 DRaFT-1 w-1
Adj. Match. w-4 DRaFT-1 wm4
215
225
100
510
245
ClipScore
12.
HPS V2
24.0
ClipScore
$7.
22.0
233
32.0
11.5-
23.
31.5
225
135
25
215
30 2
225 22.0 225
22.0 215 26. ads 250 ais
DreamSim Diversity
DreamSim Diversity
HPS
V2
Figure 4: Tradeoffs between different aspects of generative models: text-to-image consistency (Clip- Score), sample diversity for each prompt (DreamSim Diversity), and generalization to unseen human preferences (HPS v2). Different points are obtained from varying values of X for Adjoint Matching and varying number of fine-tuning iterations for the DRaFT-1 baseline. Overall, we find our pro- posed method Adjoint Matching has the best Pareto fronts.
Evaluation results. In Tab. 2 we report the evaluation metrics for the baselines as well as our proposed Adjoint Matching approach. We compare each method roughly the same wall clock time (see the times and number of iterations in Tab. B and comments in App. H.5). D.. We find that across all metries, our proposed memoryless SOC formulation outperforms existing baseline methods. The choice of SOC algorithms also obviously favors Adjoint Matching over continuous and discrete adjoint methods, which result poorer consistency and human preference metries.
A Ablation: base model vs. reward tradeoff. We note that the scaling in front of the reward model determines how strongly the we should prefer the reward model over the base model. As such, we see a natural tradeoff curve: higher À results in better consistency and human preference, but lower diversity in the generated samples. Overall, we find that Adjoint Matching performs stably across all values of A. Our method of regularizing the fine-tuning procedure through memoryless SOC works much better than baseline methods which often must employ early stopping. We show the qualitative effect of varying À in Fig. 3 while for the DRaFT-1 baseline we show the effect of varying the number of fine-tune tuning iterations.
Ablation: classifier-free guidance. We note that it is possible to apply classifier-free guidance (CFG; Ho & Salimans (2022); /Zheng et al. (2023)) after fine-tuning. We use the formula (1 + w)v(2, (ly) - wu(xr, f). where is the guidance weight, v1x, fly) is fine-tuned text-to-image model while v(x t) is an unconditional image model. This is not principled as only the conditional model is fine-tuned, but generally it is unclear what distribution guided models sample from anyhow. In Fig. 4 we show the evaluation metrics with classifier- guidance applied. Comparing three different guidance weight values, and to some extent, human preference, but this comes at the cost of being worse terms of diversity. We show qualitative differences in Fig. (App. A).
higher weight does improve text-to-image consistency,
CONCLUSION
We investigate the problem of fine-tuning dynamical generative models such as Flow Matching and propose the use of stochastic optimal control (SOC) formulation with memoryless noise schedule. This ensures we converge to the same tilted distribution that the large language modeling literature uses for learning from human feedback. In particular, the memoryless noise schedule corresponds to DDPM sampling for diffusion models and a new Memoryless Flow Matching generative process for flow models. In conjunction, we propose a novel training algorithm for solving stochastic opti- mal control problems, by casting SOC as regression problem, which we call the Adjoint Matching objective. Empirically, we find that our memoryless SOC formulation works better than multiple existing works on fine-tuning diffusion models, and our Adjoint Matching algorithm outperforms related gradient- methods. In summary, we are the first to provide a theoretically-dri rithm for fine-tuning Flow Matching models, and we find that our approach significantly outperforms baseline methods across multiple axes of evaluation tettuttton-eetttto teet-tte-iiaaeeeeeeeeet tomi consistency, generalization to unseen human preference, and sample diversity large-scale text-to-image generation.
algo-
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijndent Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.0877 2023.
Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023.
Brandon Amos al. Tutorial on amortized optimization. Foundations and Trends@ in Machine Learning, 16(5):592-732, 2023.
Pietro Astolfi, Marlene Careil, Melissa Hall, Oscar Mañas, Matthew Muckley, Jakob Verbeek, Adri- ana Romero Soriano, and Michal Drozdzal. consistency-ddiersitt-eellll pareto fronts of condi- tional image generative models. arXiv preprint arXiv:2406. 10429, 2024.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, Nicholas Joseph, Saurav Kadavath, Jackson Kernion, Tom Conerly, Sheer El-Showk, Nelson Elhage. Zac Hatfield-Dodds, Danny Hernan- dez, Tristan Hume, Scott Johnston, Shauna Kravec, Liane Lovitt, Neel Nanda, Catherine Olsson, Dario Amodei, Tom Brown, Jack Clark, Sam McCandlish, Chris Olah, Ben Mann, and Jared Ka- plan. Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint orXiv::204. 0058622,0....... 2022.
Grigory Bartosh, Dmitry Vetrov, and Christian A. Naesseth. Neural diffusion models. arXiv preprint 10X.0:2310.08337, 20244.
Grigory Bartosh, Dmitry Vetrov, and Christian A. Naesseth. Neural flow diffusion models: Learn- able forward process for improved diffusion modelling, arXiv preprint mxXv:2204.112400 2024b.
Richard Bellman. Dynamic programming. Princeton Landmarks in Mathematics. Princeton Univer- sity Press, Princeton, NJ, 2010., 1957.
Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Kamer,, Uriel Singer, and Yaron Lipman. D-flow: Differentiating through flows for controlled generation. arXiv preprint arXiv:2402.14017, 2024.
Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion-based generative modeling. arXiv preprint arXiv:2211.0133 2023.
Joris Bierkens and Hilbert J Kappen. Explicit solution of relative entropy weighted control. Systems & Control Letters, 72:36-43, 43,66443222 2014.
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine Training diffusion models with reinforcement learning. In The Twelfth International Conference on Learning Rep- resentations, 2024.
Joan Bruna and Jiequn Han. Posterior sampling with denoising oracles via tilted transport. arXiv preprint 2xX1v:2407.00745, 2024.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Ricky T. Q. Chen, Brandon Amos, and Maximilian Nickel. Learning neural event functions for ordinary differential equations. In International Conference on Learning Representations, 2021.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear
memory cost. arXiv preprint arXiv:1604.06177 066644 2016.
Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Mare Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209. 14687, 2022.
Kevin Clark, Paul Vicol, Kevin Swersky, and David J. Fleet. Directly fine-tuning diffusion models on differentiable rewards. In The Twelfth International Conference on Learning Representations, 2024.
11

---

Published as a conference paper at ICLR 2025
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. In Advances in Neural Information Processing Systems, volume 34, pp. 17695-17709. Curran Associates, Inc., 2021.
Alexander Denker, Francisco Vargas, Shreyas Padhy, Kieran Didi, Simon Mathis, Vincent Dutordoir, Riccardo Barbano, Emile Mathieu, Urszula Julia Komorowska, and Pietro Lio. Deft: Efficient finetuning of conditional diffusion models by learning the generalised h-transform. arXiv preprint 7,111 2024.
Carles Somingo-EEnichh Jiequn Han, Brandon Amos, Joan Bruna, and Ricky T. Q. Chen. Stochastic optimal control matching. arXiv preprint 22X00 2023.
Yuanqi Du, Michael Plainer, Rob Brekelmans, Chenru Duan, Frank Noć, Carla P. Gomes, Alan Apsuru-Guzik, and Kirill Neklyudov. Doob's lagrangian: A sample-efficient variational approach to transition path sampling. arXiv preprint arXiv:2410.07974, 2024.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024.
Ying Fan and Kangwook Lee. Optimizing ddpm sampling with shortcut fine-tuning, In International Conference on Machine Learning, 2023.
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Picter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lec. Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models. arXiv preprint arXiv:2305. /6381, 2023.
Ying Fan, Olivia Watkins, Yuqing Du Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Reinforcement learning for fine- tuning text-to-image diffusion models, Advances in Neural Information Processing Systems, 36, 2024.
W.H. Fleming and R.W. Rishel. Deterministic and Stochastic Optimal Control. Stochastic Mod- elling and Applied Probability. Springer New York, 2012.
Stephanie Fu, Netanel Tamir, Shobhita Sundaram, Lucy Chai, Richard Zhang, Tali Dekel, and Phillip Isola. Dreamsim: Learning dimensions of human visual similarity using synthetic data. arXiv preprint arXiv:2306.09344. 2023.
Vicenç Gómez, Hilbert J Kappen, Jan Peters, and Gerhard Neumann. Policy search for path inte- gral control. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 482-497. Springer, 2014.
Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv:1412.6572, 2014.
Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse problems, 34(1):014004, 2017.
Jiequn Han and Weinan E. Deep learning approximation for stochastic control problems. arXiv preprint arXiv: 1.04422, 2016.
Carsten Hartmann and Christof Schütte. Efficient rare event simulation by optimal nonequilibrium forcing. Journal of Statistical Mechanics: Theory and Experiment, 2012(11):P1 1004, 2012.
Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Le Bras, and Yejin Choi. Clipscore: A reference- free evaluation metric for image captioning. arXiv preprint arXiv:2104.087/6 2021.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv:2207.12558. 2022.
arXiv preprint
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33. Curran Associates, Inc., 2020.
12

---

Published as a conference paper at ICLR 2025
Yujia Huang, Adishree Ghatare, Yuanzhe Liu, Ziniu Hu, Qinsheng Zhang, Chandramouli S Sas- try, Siddharth Gururami, Sageev Oore, and Yisong Yue. Symbolic music generation with non- differentiable rule guided diffusion. arXiv preprint arXiv:2402.14285, 2024.
Gabriel Ilharco, Mitchell Wortsman, Ross Wightman, Cade Gordon, Nicholas Carlini, Rohan Taori, Achal Dave, Vaishaal Shankar, Hongseok Namkoong, John Miller, Hannaneh Hajishirzi, Ali Farhadi, and Ludwig Schmidt. Openelip, July 2021.
HJ Kappen. Path integrals and symmetry breaking for optimal control theory, Journal of Statistical Mechanics: Theory and Experiment, 2005(11), 1). nov 2005.
Hilbert J Kappen, Vicenç Gómez, and Manfred Opper. Optimal control as a graphical model infer- ence problem. Machine learning, 87(2):159-182 2012.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. On density estimation with diffu- sion models. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Pick- a-pic: An open dataset of user preferences for ext-to-image generation. In Thirty-seventh Con- ference Neural Information Processing Systems, 2023.
Matthew Le, Apoorv Vyas, Bowen Shi, Brian Karrer, Leda Sari, Rashel Moritz, Mary Williamson, Vimal Manohar, Yossi Adi, Jay Mahadeokar, et al. Voicebox: Text-guided multilingual universal speech generation at scale. Advances neural information processing systems, 36, 2024.
Dongzhuo Li. Differentiable gaussianization layers for inverse problems regularized by deep gener- ative models. arXiv preprint arXiv:21 2021.
Xuechen Li, Ting- Kam Leonard Wong, Ricky Q. Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pp. 3870-3882, PMLR, 2020.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos Theodorou, and Ricky T. Q Chen. Generalized schrödinger bridge matching. In The Twelfth International Con- ference on Learning Representations, 2024.
Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint 14571777709 2022.
Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Jonathan Long, Evan Shelhamer, and Trevor Darrell. Fully convolutional networks for semantic segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3431-3440, 2015.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment gflownets. arXiv preprint arXiv:2201.132555 2023.
Dimitra Maoutsa, Sebastian Reich, and Manfred Opper. Interacting particle solutions of fokker- planck equations through pradient-logddesiie estimation. Entropy, 22(8):802, 2020.
Shakir Mohamed, Mihaela Rosca. Michael Figurnov, and Andriy Mnih. Monte carlo gradient esti- mation machine learning. Journal of Machine Learning Research, 21(132):1-62, 2020.
Alexander Mordvintsev, Christopher Olah, and Mike Tyka. Inceptionism: Going deeper into neural
networks. Google research blog, 20(1)):5 2015.
13

---

Published as a conference paper at ICLR 2025
Nikolas Nüsken and Lorenz Richter. Solving -igh-iimensional Hamilton-Jacon Deecobi-Bellmann pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial differential equations and applications, 2:1-48, 2021.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kel- ton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In Advances in Neural Information Processing Systems, volume 35, pp- 27730-27744 Curran As- sociates, Inc., 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Ashwini Pokle, Matthew Muckley, Ricky T. Q. Chen, and Brian Karrer. Training-free linear image inversion via flows. arXiv preprini arXiv:2310.0443 2023.
L.S. Pontryagin. The Mathematical Theory of Optimal Processes. Interscience Publishers, 1962.
Aram-Alexandre Pooladian, Carles Domingo-Fnien Enrich, Ricky T. Q. Chen, and Brandon Amos. Neural optimal transport with lagrangian costs. arXiv preprint arXiv:2406.00 00288, 2024.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. In Thirty- seventh Conference on Neural Information Processing Systems, 2023.
Konrad Rawlik, Marc Toussaint, and Sethu Vijayakumar. On stochastic optimal control and rein- forcement learning by approximate inference. In Twenty-Third International Joint Conference on Artificial Intelligence, 2013.
Lorenz Richter and Julius Berner. Improved sampling via learned diffusions. In The Twelfth Inter- national Conference on Learning Representations, 2024.
Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco Ruiz, and Omer Deniz Akyildiz, Var- Grad: A low-variance gradient estimator for variational inference. Advances in Neural Informa- tion Processing Systems, 33, 2020.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Litu Rout, Yujia Chen, Nataniel Ruiz, Abhishek Kumar, Constantine Caramanis, Sanjay Shakkottai, and Wen-Sheng Chu. Rb-modulation: Training-fiee free personalization of diffusion models using stochastic optimal control. arXiv preprint
2024.
Reuven Y Rubinstein and Dirk P Kroese. The cross-entropy method: unified approach to combina- torial optimization, Monte-Carlo simulation and machine learning. Springer Science & Business Media, 2013.
Christoph Schuhmann and Romain Beaumont. Laion-acsthetics, 2022.
S.P. Sethi. Optimal Control Theory: Applications to Management Science and Economics. Springer International Publishing, 2018.
Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin, Jie An, Songyang Zhang, Qiyuan Hu, Harry Yang, Oron Ashual Oran Gafni, et al. Mike-a-video: Text-to-vided generation without text-video 2022.
data. arXiv preprint arXiv:2209.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- tional Conference on Learning Representations, 2021a. la
Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Excudoinverse-guidee diffusion models for inverse problems. In International Conference Learning Representations, 2023.
14

---

Published as a conference paper at ICLR 2025
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. arXiv preprint arXiv:1907.05660 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations (ICLR 2021), 2021b.
Nisan Stienmon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning summarize with human feedback. In Advances in Neural Information Processing Systems, volume 33, pp. 3008-3021, Curran Associates, Inc., 2020.
Hyung Ju Suh, Max Simchowitz, Kaiqing Zhang, and Russ Tedrake. Do differentiable simulators give better policy gradients? In International Conference on Machine Learning, pp. 20668- 20696. PMLR, 2022.
Wenpin Tang. Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond. arXiv preprint arXiv:2403.06679, 2024.
Emanuel Todorov. Linearly-solvable markov decision problems. Advances neural information processing systems, 19, 2006.
Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. arXiv:1903.01608 2019.
Masatoshi Uehara, Yulai Zhao, Tommaso Biancalani, and Sergey Levine. Understanding rein- forcement learning-based fine-tuning of diffusion models: A tutorial and review. arXiv preprint arXiv:2407.13734, 2024a.
Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine-tuning of continuous- time diffusion models as rntropy-regullrizze control. arXiv preprint 2rXxx:2402.15194 2024b.
Francisco Vargas, Andrius Ovsianas, David Lopes Fernandes, Mark Girolami. Neil D Lawrence, and Nikolas Nüsken, Bayesian learning via neural Schrödinger-föllme flows. In Fourth Symposium Advances in Approximate Bayesian Inference, 2022.
Francisco Vargas, Will Sussman Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. In The Eleventh International Conference on Learning Representations, 2023.
Pascal Vincent. lation, connection between score matching and denoising autoencoders. Neural compu-
2011.
Apoorv Vyas, Bowen Shi, Matthew Le, Andros Tjandra, Yi-Chiao Wu, Baishan Guo, Jiemin Zhang, Xinyue Zhang, Robert Adkins, William Ngan, et al. Audiobox: Unified audio generation with natural language prompts. arXiv preprint 23. 2023.
Bram Wallace, Meihua Dang, Rafael Rafailov. Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusson model alignment using direct preference optimization. arXiv preprint arXiv:231 1.222/1..12000 /2908, 2023a.
Bram Wallace, Akash Gokul, Stefano Ermon, and Nikhil Naik. End-to-end diffusion latent opti- mization improves classifier guidance. arXiv preprint arXiv:2303. 2023b.
Luhuan Wu, Brian Trippe, Christian Naesseth, David Blei, and John P Cunningham. Practical and asymptotically exact conditional sampling in diffusion models. In Advances in Neural Informa- tion Processing Systems, volume 36. pp. 31372-31444 Curran Associates, Inc., 2023a.
Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text-to- image synthesis. arXiv preprint arXiv:2306.093.. 2023b.
Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text-to- image synthesis. arXiv preprint arXiv:2306.098: 093411 2023c.
15

---

Published as a conference paper at ICLR 2025
Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li. Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. In Thirty-seventh Conference Neural Information Processing Systems, 2023.
Dinghuai Zhang, Yizhe Zhang, Jiatao Gu, Ruixiang Zhang, Josh Susskind, Navdeep Jaitl,, and Shuangfei Zhai. Improving gflownets for texi-to-image diffusion alignment. arXiv preprint arXiv:2406.0066 00633, 2024.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: stochastic control approach for sam- pling. In International Conference Learning Representations, 2022.
Wei Zhang, Han Wang, Carsten Hartmann, Marcus Weber, and Christof Schütte. Applications of the cross-entropy method to importance sampling and optimal control diffusions. SIAM Journal Scientific Computing, 3666) 2014.
Qinqing Zheng, Matt Le, Neta Shaul, Yaron Lipman, Aditya Grover, and Ricky T. Q. Chen. Guided flows for generative modeling and decision making. arXiv preprint arXiv:2311.13444 2023.
Brian D Ziebart, Andrew I. Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. Aaai, volume 8, pp. 1433-1438. Chicago, II., USA, 2008
Daniel M. Ziegler, Nisan Stienmon, Jeffrey Wu, Tom B. Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv:1909.08540 2020.
16

---

Contents
Additional Figures & Tables
18
Related work
27
Results on DDIM and Flow Matching
28
C.1
Denoising Diftusion Models
28
The continuous time limit DDIM
28
C.3
Forward and backward stochastic differential equations
29
C.3.J
Proof of Lemmall
C,3,2
Proof of Lemmal
31
C.3.3
Proof of Prop.
32
C.4
The relationship between the noise predictor and the function
34
C.5
The relationship between the vector field & and the score function
34
Stochastic optimal control as maximum entropy RL in continuous space and time
35
D.I
Maximum entropy RL
35
D.2
From maximum entropy RI. to stochastic optimal control
36
D.3 Proof of Prop. from MaxEnt RL to SOC.
37
D.4
Proof of equation
[12]: the control cost is KL regularizer
39
Proofs of 'Subsec. 3.3 memoryless noise schedule and fine-tuning recipe
39
E.1
Proof of Prop. the memoryless noise schedule
39
E.2
Fine-tuning recipe for general noise schedules)
E.2.1
Expressing b, u in terms of
41
F.2.2
Proof Thm.
42
F
Methods to solve SOC problems
44
F.1
Existing methods
44
F.1.1
The adjoint method
44
F.1.2
Importance-weighter matching objectives for regressing onto the optimal control
45
F2
Derivation of the Continuous Adjoint method
45
F3
Proof of Prop. 2
Theoretical guarantees of the basic Adjoint Matching loss
47
F.4
Theoretical guarantees of the Adjoint Matching loss
48
R.5
Pseudo-code of Adjoint Matching for Flow Matching and DDIM fine- tuning
50
Adapting diffusion fine-tuning baselines flow matching
50
G.I
Adapting ReFL (Xu et al. (203) to flow matching
50
G.2
Adapting Diffusion-DPO (Wallace al. (023a) to flow matching
H
Experimental details
53
H.1
Noise schedule details
53
H.2
Selection of gradient evaluation timesteps
54
H.3
Loss function clipping: the LOT hyperparameter
H.4
Computation of evaluation metrics
55
H.5
Remarks on computational costs
55
H.6
Remarks on number of sampling timesteps
55
17

---

Published as a conference paper at ICLR 2025
A
ADDITIONAL FIGURES & TABLES
15
107
65
0.5
DROPT-L DRaFT-12 DRO
Control
t03
DRaPT- -RDDRT--0
PPPL DPO
-15
Adj Match A- 2530
-Ad Maton A 2500
Corr 22522500
Cort
17990
25000
70000
75O00
10000C
125020
3500GD
:7560C
23000
30000
T5000
300000
123000
150000
17500C
Wall block time bd)
Wall clock time (a)
30-
Oppcoo
28
-0hAFT-40
Mates 1=5000
Metcon
Metoss 1-000
Curi =11250
125000
173000
Wall-dock
Figure 5: Average values of ImageReward (reward function), control cost (f ([ t)ll dt), and ClipScore vs. wall-clock time for Adjoint Matching and baselines. Lines show averages three fine-tuning runs, evaluating on separate test datasets size 200. Confidence intervals show standard errors of estimates.
O'
PEEDZO
Text prompt: "Man sitting on sofa home front of fireplace and using laptop computer, rear view"
Text prompt: "3D World Food Day Morocco"
Figure Generated samples from varying classifier-free guidance weight w. from Adjoint Matching fine-tuned model. Higher guidance increases text-to-image consistency but loses diversity and has use cases for generating highly structured images such 3D renderings. Corresponding samples from the base model can be found in Fig. 7
18

---

Published as a conference paper at ICLR 2025
Fine-tuning
Fine-tuning (t)
Sampling x(t)
ClipScore diversity 1
PickScore diversity 1
Total time # iterations
ImageReward f
loss
None
/21H
-0.384
28.07
1.33
N/A
N/A
(CFG = 1.0)
0.920
30.29
1.82+009
v 2mz 0
V 22
1.357 +0009
16.86=0098
1.21 +000
440k
DRaFT-1
1.251 1:0000
16.76.1.06
1.27
4000
/2n
22m
0.560
1..07:1.37
1.64:10.12
148k
DRaFT-40
0
0.42440042
20.9
1.67
/ 1500
2mj
27p
1.386=0.035
27.80=1.44
1.62
118k
DPO
0
0..57
29.81 1.443
1.68
1000
/2n=
/2nc
0.687 10085
1974
1.22.10.8
173k
ReFL
0.709 100000
18.39
1.31
6000
Cont. Adjoint
V 227z
-0.48
26.97
1.82+am
153k
2m
= 12500
-.2299
26.25
1.90
750
Disc. Adjoint
/2n
-0.557
30,40++2.39
1.91
152k
22m
À = 12500
-0.552+0.01
28.37++2.26
1.97
/ 1000
Adj.- Mj.-MMatchnng
v /27
0.55010043
13.00
1.65
22n
X = 1000
0.454
22.76
1.73
AdjtMMatching
/
0.755 +0000
21.33++1.71
1.55 +0.0s
156k
22n
À = 2500
0
0.67 1 +0047
21.42
1.64+
/ 1000
Adi--Matching
/21lc
0.882
20.49
1.50+000
22m
À = 12500
0.77810050
20.34:1.49
1.57
Table 3: Metrics for various fine-tuning methods for text-to-image generation. The second and third columns show the noise schedules (t) used for fine-tuning and for inference: (t)) /21z, corre- sponds to Memoryless Flow Matching, and (t) = O to the Flow Matching ODE 33. Confidence intervals show standard errors of estimates; computed over runs of the fine-tuning algorithm on
separate fine-tuning prompt datasets of size 40000 each. Test prompt sets are of size 1000, and also different for each run.
Fine tun. loss
Fine tun. (tt)
Generat. ((t)
DreamSim diversity f
Runtime/ #iter.
ImageReward f
ClipScore
PickScore
HPS v2
22n
2m
0.4592000%
28.46 28.4620025
18.774009
22.54::0.17
37.51
43k-t:277
ReFL
0
0
0.330
29.63
19.08
22.46:20.77
1.35
15500
/2n:
/2m
0.91.20000
20.20
19.16:2006
23.633:00 16
35.21
35k 1k
DRaFT-1
0
0
0.626220 195
00.48-0.12
18.91:1034
21.921
38.5222.01
1000
/2
mm
1.427:10.267
23.39
17.24:145
15.72:t 8D
21.98=2.144
49k- 1kk
Draft-40
0
0.097:10003
29. 12:441
18.97.0144
21.93:1:000
1..35.1.344
500
Adj-Matthh
22p
0.107
29.37
19.05
22.79:+0.00
46.38
v 2mp
À 1000
0
0.051
30.58-00117
19.311007
21.93:10233
1.55
Adj-MMachh.
2mp
0.199:00008
29.27=0.21
19.07-t0.10
22.98::0.00
15.03=1161
v 2m
2500
0.106:1:006
30.43::0.24
19.32-10.1
22.16::0.33
47.6
/50
Adj- Match.
2m
0.299400005
20.61
19.26:20.14
23.67:20227
13.36
2m
0.22440051
19.5240.11
22.93-:0:1
44.62-1.79
Cont. Adj.
/2m
0.91024.116
26.29
18.06:2016
18.86:20.88
51.6021.97
/2n
09.5
19.90:10.50
50.87
/50
Disc. Adj.
/2m
0.978
26. 26.6820.76
18.511011
18.53:20
15.95
38k:1004
21J
12500
10006
18.514011
18.53:10:8
24.78
/50
Table 4: Additional metrics for various fine-tuning methods for text-to-image generation, which complement the ones in Tab. 2 (both tables correspond to the same runs). The second and third columns show the noise schedules (t)) used for fine-tuning and for inference: o(t) 22m, corre- sponds to Memoryless Flow Matching, and a (t) = to the Flow Matching ODE B
19

---

Published as a conference paper at ICLR 2025
Fine- tuming
Witer.
Fine-tum tun. (tt)
Sampl.
DreamSim diversity f
ImageReward
ClipScore 1
PickScore T
HPS v2
loss
/
a(t)
221
- -,084
24.15:40226
17:25-40.06
16.19-40.17
53,60++1.31
0.0
None
N/A
N/A
0
0.920=00004
28.32:1003
18.15:0..07
17.8920.16
1.553:1
V2
/2
0.913=t0.068
29.80++0.22
19.16+
23.63
35.21+
1000
0.626=t0.195
30.48+0322
18.91
21.924:.66
38.524:201
0.204
09
19.29
24.40=0027
18.51
0.0
DRaFT-1
2000
L052440008
30.65
19.27400.1
23.81
32.11
2m
1.37==00001
29.96
19.31
24.42
2.357
3000
1.173-000088
30.8640.2
19.37=00.0
24.17
29.69:1.30
/2
1.357
30.182024
19.38:20008
24.61 10111
25.54
4000
2.000
24.37
27.39
/2n
229
0.550
30.36++0.22
19.29
24.12:40.17
1..89.+-.50
1000
0.454
31.41:023
19.57:00.0
23.29:00188
13.10:11.76
/ /2
< 22
0.755++0.040
30.59++0.00
19.49+-0.10
24.85+60.25
37.07++1447
0.0
Adj-MMrchh.
2500
0.004
31.641021
19.71
24.12:20.27
39.88:11.39
22p
V 21
0.882:00.088
30.62.40.0
19.50++0.09
24.95
34.50:1:33
12500
0.778:t0.150
31.65-+0.19
19.76::0.08
24.49:20.77
37.24
0.269=0.050
18.74+0007
20.47:0.18
43.8241.24
1.0
None
N/A
N/A
0
0.123
31.83
19.28
20.95:20.16
42.59:11.2
/ /2m
/200c
1.13320051
32.0640
19.6920.06
24.5620.17
28.25
1000
0
0.856 32.324025
19.38:2.334
22.88
29.98:11.86
1.0
DRaFT-1
2000
0
0.007 32.36:10.18
19.67
24.48:20.28
25.09:11.82
3000
0
1.2555+0.038 32.36-+0.19
19.70+t0.05
24.64-t0.17
23:24-+-119
4000
0
1.296220.033 32.30.20
19.68:2000
24.71
21.54410:66
1000
0
0.782:20.044
33.05:1022
20.20:10009
2:011
32.67:11.26
2/7
/2
1.027=00.088
32.85++021
20.08+t0.08
25.88==0.20
29.83 29.83+t:000
1.0
AdathMtthh.
2500
0
0.910 33.20:1017
20.29:10.09
25.39:2024
1..34:1..51
12500
0
0.985=t0.041
33.10+0.18
20,28:+0.08
25.61
28.86:
/2mc 0
0.277:10.043
32.68:10.18
19,50:10.07
22.29:10116
35. 12:1:29
4.0
None
N/A
N/A
0.209:
32.83+0.117
19.79
22.30
32.05
/ /2
22m
0.062
32.29:10.16
19.48:10.05
23.67:t:0.13
25.00 :::003::1.22
1000
D
0.604
31 31.802386
19.09
21.69
25.92::2:57
4.0
DRaFT-1
2000
0
1.112+0.066
32.29-40200
19.34
23.3144-022
21.02++1667
3000
0
32.31
19.36+-0.05
23.29:40.14
19.53:41.24
4000
1.172:20.040
32.20:1022
19.3010007
23.20:20155
18.45:1:006
1000
0
0.852+40.046
33.50-1-0.12
20.31
24.97-40.19
25.8340032
/ /2
1.052#0.099
33.51::0.19
20.15::0.07
25.56::0.18
26.21
4.0
Adj -arch.
2500
0.94240.042
33.61
20.35-.0.0
25.34-..211
24.30::066
12500
0
1.007
33.48-1-.200
20.29::0.08
25.50=00.99
23.48
Table 5: Evaluation metrics when using classifier-free guidance (CFG; Ho & Salimans 2022
LR
Fine-tuning
Fine-tun. ((t)
Generat. (tt)
DreamSim
ImageReward
ClipScore
PickScore
HPS
Adam 81
loss
diversity
3 x 10
DRaFT-1
22m
V2m
1.67-40029
30.28++0.56
19.37=0a9
24.70=0.15
21.20++0s2
0.997
Adj-MMatch.
v 22m
/ 22
1.13010034
31,01:00.7
19.60=0.05
25.01
26:73:1088
X = 12500
2 x 10
Disc. Adj.
/2m
/2m
1.18620.53
21.95
16.94=0.95
12.34=4.40
10
/0.95
X = 12500
0
0.96140053
24.07+44.1
17.86-1.17
15.93=5.80
33.62447.0
Table 6: Metrics for alternative optimization hyperparameters (learning rate and Adam B1).
20

---

Published as a conference paper at ICLR 2025
Fine-tuning loss
Fine-tunie tine-tuning (tt)
Generative a(t)
DreamSim diversity 7
ImageReward
ClipScore f
PickScore
HPS v2
MajcMing
0.009::0077
29.18:051
18.6622..9
20.75=0.32
41.33::124
12500
0
0.454
31.41 20.22
19.57
23.29:20.18
13.10::1.76
Matching
22
0.882:00088
30.6240.0
19.50
24.95=0.28
12500
O
0.778:20:000
31.6520.19
19.7622000
24.49-0.27
37.24
Table 7: Comparison with alternative fine-tuning noise schedule (t) = 1. We see that the initial value function bias (Subsec. 3.2 results the model not having high reward function (ImageReward is the reward function used for fine-tuning). Its performance other metrics are also lower than when iin-tuniig with the memoryless noise schedule, except for diversity.
#sampl. timesteps
Fine-tuning loss
Fine-tun.
Sampl. a(t)
DreamSim
ImageReward f
ClipScore f
PickScore ?
HPS v21
ott)
diversity
22m
1.000
13.99:10.12
14.98
7.37:10.10
5.07
None (Base)
N/A
1 1.386220000
26.26::0:24
17.64=0.07
14.92:20.17
51,26++
v 22e
22
1.005
25.98
18.28:10007
22.08:10.18
14.47:10.67
DRaFT-1
0
1 1.236:20.08
31.54:0027
19.53-0.07
24.47:20119
24.78:0888
Adj--Match.
/2
2..04
17.12:1:0.5
15.76
11.48:
9.88:1081
27
À 12500
40.055
31.36:1-0200
19.56
23.23-20028
33.75-+1.488
21
4.000
14.58 1-013
15.07
7.47:400110
11,27:4033
None (Base)
N/A
-.017++00055
27.92+-0.19
18.01 0-007
17.17=40.15
14.69+
v 2rjs
1.301::0.083
27.09+
18.93
23.78:+0.00
21.05
DRaFT-1
0
1.255
31.14++0.25
19.43=0006
24.52:40.16
26.15+
Adj.-Mutch.
2m
0.032:10:022
25.07:1:027
18.01 0.007
20.75:10.23
29.06:12.4
A 12500
v 27j
0.768
31.70-t0.17
19.73=0.08
24.30=0026
35.90+
22
1.384.100000
24.15:1:0:2
0.005
16.19:10117
53.60:11.7
None (Base)
N/A
0.920:10.042
28.32:1:022
18.15
17.89.10.16
56,53:11.52
40
V2ne
/2
0.003
30.18:1::24
1:0.8
24.61
25.54:1:0.99
DRaFT-1
0
1..01
30.95
1903
24.37:1011
1..19
Adj.-Match.
0.8824400588
30.62++0.30
19.50 0.009
2.295
34.50++123
À 12500
27
0.778
31.65++-0.19
1..76
24.49-40.27
37.24++1.57
0.881++0.011
27.83++-0199
18. 000
18.43-40.17
$7.21-+1.50
(Base)
N/A
0.881+00006
28.65 28.65-40.118
0..2
18.20=0.17
57.73++168
100
22ps
22)
1.343:100000
30.641:0200
19,38
24.37
25.511+:1100
DRaFT-1
0
1. 1.239:20.057
30.74+:0.4
0.03
24.24=00117
28.70:
Adj.-Mathh.
21)
0 0.892:10.044
31.23:1:223
19.65
24.92:10.23
35.11:1:1.40
2rps
12500
0.779.10.048
31.64:10.17
19.76..0088
24.57:10.25
38.26:1:.65
22
0.848
28.37
18.27:00.00
18.56:10.19
1.550
None (Base)
N/A
0.801
28.50
1..0
18.25-10.14
1.684
200
v 27yz
1.331 0.004
30.69
19.36
24.21
26.41-4-1.18
DRaFT-1
2.224400042
30.77++0.7
19.32
24.18-+0.16
29.09++:.07
Adj.Matthh.
0. 0.869
31.33440:21
19.68 6668-800000
24.81
35.90++1.55
/2m
12500
0. 0.766-40.000
31.61-10.16
19.75
24.52+40.4
38.60++128
Table 8: Performance metrics for different number of sampling steps. Only the number of sampling steps is ablated; the fine-tuned models used in all cases are the ones fine-tuned using steps.
21

---

Published as a conference paper at ICLR 2025
10
vbf
ee
Text prompt: "Man sitting sofa home in front of fireplace and using laptop computer, view"
Text prompt: "3D World Food Day Morocco"
Figure 7:
Generated samples from varying classifier-free guidance weights, from the pre-trained Flow Matching model. Corresponding samples from the fine-tuned model be found in Fig. 6
22

---

Published as a conference paper at ICLR 2025
Base Flow Matching model
Adjoint Matching (Ours)
DRaFT-1
Figure 8: Generated samples with classifier-free guidance (w = 1) and (tt) = 0 across ten selected prompts. Each row corresponds to different prompt and each image corresponds to different random seed consistent across models.
23

---

Published as a conference paper at ICLR 2025
Base Flow Matching model
Adjoint Matching (Ours)
DRaFT-1
Figure 9: Generated samples with classifier-free free guidance (w 1) and (t) 0 across ten selected prompts with people. Each row corresponds to different prompt and each image corresponds to different random seed consistent across models.
24

---

Published as a conference paper at ICLR 2025
None (Base)
DRaFT-1
DRaFT-40
ReFL
Cont. Adj.
A 12500
Disc. Adj.
A = 12500
Adj. match.
À = 1000
Adj. match.
A 2500
Adj. match.
À = 12500
Figure 10: Generated samples without guidance (w 0) and (tt) 0 across seven selected prompts. Each row corresponds to a different finetuning algorithm. Prompts: "Seaside view poster with palm trees vector image". "Cayucos Beach Inn". "Happy Summer Life- Aloha Flowers and Melon Pattern Metal Print", "Castle Square, Warsaw Old Town". "Funny girl blowing soap bub- bles. High quality photo", "Colombian man with sweatshirt over yellow wall listening to something by putting hand on the ear", "man in the hood black mask masquerade".
25

---

Published as a conference paper at ICLR 2025
None (Base)
DRaFT-I
DRaFT-40
ReFL
Cont. Adj.
À = 12500
Disc. Adj.
À = 12500
Adj. match.
À = 1000
Adj. match.
À = 2500
Adj. match.
À 12500
Figure 11: Generated samples without guidance (w = 0) and (t) 22 across seven selected prompts. Each row corresponds different finetuning algorithm. The prompts are the same in
Fig. 10
26

---

Published as a conference paper at ICLR 2025
B RELATED WORK
Fine-tuning from human feedback. There are two main overarching approaches to RLHF: the
reward-based approach (Ziegler et al. 2020 Stiennon et al. 2020; Ouyang et al. 2022 Bai et al.) 2022] and direct preference optimization (DPO; Rafailov et al. (2023 D. The reward based approach (Ziegler et al., [020, [tiennon et al. 2020, Ouyang et al. 2022 | Bai et al. 2022) consists in learn- ing the reward model r(x) from human preference data, and then solving a maximum entropy RL problem with rewards produced by r(x).. DPO merges the two previous steps into one: there is no need to learn r(x) as human preference data is directly used to fine-tune the model. However, DPO is typically only applied with filtered dataset, and does not work explicitly with reward model. Furthermore, for flow and diffusion models specifically, it is possible to differentiate the reward function, so there is larger emphasis reward-based approaches.
Fine-tuning for diffusion models. Among existing reward-based diffusion fine-tuning methods, Fan & (2023) interpret the denoising process as multi-step decisionnmakinng task and use policy gradient algorithms to fine-tune diffusion samplers. Black et al. (2024) makes use of proximal policy gradients for fine-tuning but this does not make use of the differentiability of the reward model. Fan
jet al. (2023) also consider KL-regularized rewards (33) but do not make the critical connection to the tilted distribution (I) that we flesh out in Subsec. 3.2 The fine-tuning algorithms of Xu
let al.] (2023); Clark et al. (2024) directly take gradients of the reward model and use heuristics to try stay close to the original base generative model, but their behavior is not well understood and unrelated to the tilted distribution: Xu et al.] (2023) takes gradients of the reward applied on the denoised sample at different points in time, and Clark et al. (2024) backpropagates the reward function through all or part of the diffusion trajectory. Finally, Uehara et al. (2024b) also fine-tune diffusion models with the goal of sampling from the tilted distribution (I). but their approach is much more involved than ours as it requires learning value function, and solving two stochastic
optimal control problems. Additional reward fine-tuning works include Bruna & Han (2024). that provide theoretical guarantees to sample from the tilted distribution when the reward is a quadratic
function, and Zhang et al. (2024), that propose a reward fine-tuning algorithm for the GFlowNet architecture.
Inference-time optimization methods. Some have proposed methods that do not update the base model but instead modify the generation process directly. One approach is to add guidance term to the velocity (Chung et [al]|2222) (ong et al.j [023) Pokle al.) (023); however, this is a heuristic and it is not well-understood what particular distribution is being generated. Another approach is to directly optimize the initial noise distribution (Li) 2021: Wallace al., 2023b (Bnn-Hauu et al., 2024]; this is taking an opposite approach to the inital value bias problem than us by moving all of the work into optimizing the initial distribution. A more computationally intensive approach is to perform online estimation of the optimal control, for the purpose of heuristically solving an optimal control problem within the sampling process (Huang et al., 2024, Rout et al. 2024); these approaches aim to solve separate control problem for each generated sample, instead of performing amortization (Amos et al. 2023) to learn a fine-tuned generative model.
Optimal control in generative modeling. Methods from optimal control have been used to train dynamical generative models parameterized by ODEs (Chen et al.) 2018), SDEs (Li et al.) 2020). and jump processes (?), enabled through the adjoint method. They can be used to train arbitrary generative processes, but for simplified constructions these have fallen in favor of simulation free matching objectives such as denoising score matching (Vincent 2011) and Flow Matching (Lipman) jet al. 2023). The optimal control formalism also has significance in sampling from un- normmllzdl
distributions (Zhang& & Chen 2022 Berner al., [023]] Vargas et al. 2023 2022 Richter & Berner
2024 Tzen & Raginsky] (201)). The inclusion of a state cost has been used to solve transport
problems where intermediate path distributions are of importance (Liu et al. 2024 Pooladian al.)
2024]. These collective ad vlvnnces naturally lead to the consideration of the optimal control formalism
For reward fine-tuning.
Conditional sampling in inverse problems. Denker et al. (2024) and Wu et al. (202aa) indepen- dently consider pre-trained diffusion model p(z), and observation y the generated sample x, as well as the analytic likelihood p(y(z)) Their aim is to sample from the posterior p(z) (z), and their applications include inpainting, class-conditional generation, super-resolution, phase retrieval,
27

---

Published as a conference paper at ICLR 2025
non-linear deblurring, computed tomography, and protein design. Their setting reduces partic- ular case of our reward finettnnig framework by setting r(æ) = log ply|z). Denker et al. (2024) formulate an SOC problem, and they solve it via the log oog-varince loss (Richter et al. (2020); Küsken
& Richter (2 and the moment loss (Nüsken & Richter 2021). which they refer to as the tra- jectory bal loss (Malkin et al., 2023). Wu al. (2023a) propose Twisted Diffusion Sampler, an
algorithm based on Sequential Monte Carlo that uses increased inference- time compute to reduce
bias. A third work that also tackles the conditional sampling problem Du al. (2024), which use a Lagrangian formulation that they solve approximately using Gaussian paths.
C RESULTS ON DDIM AND FL.w MATCHING
C.I DENOISING DIFFUSION MODELS
We next discuss diffusion models, in particular the sampling scheme proposed by Denoising Diffu- sion Implicit Model (DDIM: (Song et al. (2021a)) which will later relate to Denoising Diffusion Probabilistic Models (DDPM; Ho et al. (2020)) as particular case of the former. For sampling from a diffusion model, the DDIM update rule (Song al. (2021a), Eq. 12), typically stated in discrete time with & E {O.....K) is:
Xk-1 v XE-VI-MAee (X V 0K-1 (x + OKEK
(28)
Vog
where EK ~ I). Xo ~ I), (āx) is an increasing sequence such that ag = 0, ā = 1, and the sequence og is arbitrary. That is, one samples an initial Gaussian random variable Tp. and applies the stochastic update (28) iteratively K times in order to obtain an artificial sample Xx- Updates can be interpreted as progressively denoising the iterate: Fn is completely noisy and æ is fully denoised. The noise predictor model e( k) is trained to predict the noise of I& (see e.g. Ho al. (2020) for details on pre-training denoising diffusion models).
To convert DDIM continuous stochastic process, we can show that the DDIM update rule (28). up to first-order approximation, is equivalent to the Euler-Maruyama discretization of the tollowing SDE:
dx, ( è X, 2
(
à. 2
+ a(t)
(o di + (I)dB,, No~ No0, I). VI-ac
(29)
See App. C for the full derivation. To go from 28 to (2). we assumed uniform discretiza- tion of time, iee. = - This results in identifying the discrete-time process (X..) with a continuous-time process (X)))Efo.1), where ā₃ := â. 0x = 1 (t), and e ((K,, k) with base (X,, (X,,). In relation to the reference flow (2) the generative process in 29 has the same time marginals when a, = vā, and Br vI ai (Ho et al. 2220))
Furthermore, when viewed up to first order approximations, the DDPM sampling scheme (Ho et al. 2020 Algorithm 2) can be seen as special instance of the DDIM sampling scheme when (t)
a /a.a This results in the following generative process:
dX2= =
dr 20
X1 å Joee (X..f) dt + V L dB
Xo ~ ~(0,1,,
(30)
VI-ai
C.2 THE CONTINUOUS TIME LIMIT OF DDIM
The DDIM inference update (Song et al. 2021a Eq. 12) is
Fk+1 = Vok++1 Zh. T-
v āktt (rs, k) + okER
IK ~ N(00 T).
(31)
If we let Adg Oik++1 aū, we have that
dig 00++041110
x 1 1+ 4%
Ade 2
(32)
We slightly depart from the notation in Song et al. 2021aa) by flipping the direction of time and using ag which corresponds to the ag Song et al. 202 while it corresponds the ā in (o et al. 2020).
28

---

Published as a conference paper at ICLR 2025
where we used the first-order Taylor approximation of /1 + And
Ga+1 a6 - ax) + V1 -k++1 0% ---v-V
(1
Ane )(1 (ik) + v 0k+1 - 0
+ āi
Adu
- Oik - Dikk + vI -x++1
Kk =
ai+t + Ane
aa
+ v - ck+1
0
V1 āxti
Ddu /1 11+ ae(1-āz+)
+
of
1-ūe+1
ax+1
20
Ang
ax+1l
+1
2 (11-6+1)
Ade 2
VT-
(33)
where we used the same first-order Taylor approximation. Thus, up to first-order approximations, is equivalent to
Tk- (1
Ade 200
Ade 2
+
c(mg..) V1-āu-
UK ~ (),
(34)
we modify our notation slightly, we can rewrite this as X(k ((k=11A = =(1
hāgp Zöte
Xah
+
(or()))
c(Xen.kh)
+ (ho(kh)
Xo ~ NN(O I).
2
(35)
To go from 34 to 33 we introduced a continuous time variable and a step size h = 1/K, and we regard the increment hās approximately equal h times the derivative of ā. We also identified a with (hh() where a (kh) plays the role of diffusion coefficient. Note that equation (35) can be reverse-engineered as the Fuler-Maruyama discretization of the SDE
dX, = (-
å 20z
c 2
e(X...) dt + at)dBB
Xo ~ 1).
(36)
2
C.3 FORWARD AND BACKWARD STOCHASTIC DIFFERENTIAL EQUATIONS
Let (KO) and (nt)eeja,1) such that
vt [0, 1], ne
So Ki-s ds t00,
2 2 m-- exp (-2 R1- ds) dt' 1.
(37)
As shown in Tab. DDIM corresponds to d. 20,,"
à and Flow Matching corresponds to 20
Nr 21 1/ = 3g Lemma 1 (DDIM and Flow Matching fulfill the conditions
cs mr Br - -c).
The choices of (Kx (ee)) and For DDIM, we have that
(ndkejo.1) for DDIM and Flow Matching fulfill the conditions
37
Jo Ki-s ds - log ā1-t = So Ki-s ds +00,
(38)
2 St Tu exp ( (-2 fú N₂ ds) dt' = a1-t = So no exp - 2 Ni ds) dt' 1.
For Flow Matching,
S0 f1-= ds = -lg loga1-t a1 = So fi-- ds = t00,
(39)
2 f ne exp (-2 f Kg ds) dt' 87_. = 2 So ny exp 2 2e N₂ ds) dt' = 1.
(40)
Forward and backward SDEs Consider the forward and backward SDEs
dX, -1 dt + /Bn dx, (,, + 27 X,, 1)) + 2mh dB,
No Pdata:
(41)
Xo N(O, 7),
(42)
where we let p. be the density of X,, and define the score function as s(x, 1) V log ( Similarly, we let be the density of X₂. ph and pa solve the Fokker-Planck equations:
0,7² V (kI-1177) + ? Ap..
pn
Pdata:
(43)
0.p: V ((-K,z) K,, 27 V log (x))) m Ape,
Po 1).
(44)
29

---

Published as a conference paper at ICLR 2025
Lemma 2 (Solution of the forward SDE). Let (Kz)120. (ndeoo with n 2 0, and (E))22 be arbi-
mrrry. The solution X, of the SDE
xx + Ex) dt /2m Xo Pdata
(45)
X; = X,=Xo exp (- k1-s d5) +o exp (-Se ki-s ds) dt' + 22 exp (- (-Se hi-8 ds) dBr,
(46)
which has the same distribution as the random variable
X-= No exp(pp -fo ki-s ds) So" exp ( (-fe K1-s EI de' + 22 m- exp (-2 [e ki ds) dl'e,
<~ N(..
(47)
Applying Lemma 2 with & 0, we obtain that şí is also the distribution of
X1 Xo exp (- in - ds)
/ 2 'n - exp (- J K1-s ds) dl'e =;
(48)
where E ~ (N, 1). The third equality in 48 holds by (37). Hence we obtain that pi = ((,
Note also that
0.p1- V (HATTI-1) -- - - ((-Ker KyT 2n, V Voggi + TJ -pi
(49)
Thus, -1 is solution of the backward Fokker-Plancee equation 44). which proves the following:
Proposition 3 (Equality of marginal distributions). For any time E [0, 1}, the densities of the solutions X X₂ of the forward and backward SDEs are equal up to time flip: pt 71-t-
Forward and backward SDEs with arbitrary noise schedule Next, look at the following pair of borkwrd-backward SDEs:
dX, (- (-K1-XX, + (0(1-)2 m)))s(X,, +)) dt + +o
Xo Pdatas (50)
dX, (R,Xe + (000 + m)),(X, t) o(t) dB,,
Xo N),,
(51)
Here, the score function s is the same vector field as in (5I). Remark that equations particular case of (50)-(51 (51) for which (1) = 22n. The Fokker-Pland equations for 50 5
are
are:
0.p = V ((K1-10 +
a(1-t)) (-
7 t)) + 171 Ap
P0
Pdata
(52)
0.p. = V K1 (((t) + s(,, 1)) + a(t)² Apt.
Po = N(O, I).
(53)
It is straight-forward to see that for any the solutions 55 and of 52
53 are also solutions of 45 (44). Hence, the marginals X₁ and X; are equally distributed for all noise schedules a, and they are equal to each other up to a time flip.
Equality distributions over trajectories
The result Prop. 3 can made even stronger:
Proposition 4 (Equality of distributions over trajectories). Let X, X be the solutions of the SDES with arbitrary noise schedule. For any sequence of times (4))osig,, joint distribution of is equal to the joint distribution of (X1-t))osis) measures P, of the forward and backward processes X. X are equal, up to a flip the time direction.
551-51
equivalently, that the probability
This result states that sampling trajectories from the backward process is equivalent to sampling them from the forward process and then flipping their order.
30

---

Published as a conference paper at ICLR 2025
C.3.1 PROOF OF LEMMA
As shown in Tab. 1 DDIM corresponds to increasing, and
de. 20,
7/
â. 26.
Thus, 2 because dz is
S0 1i ds = -To à 201-6 ds = = 0, log 11 ds = 21log 01-t - log (1) = 2 log Ū1-t1
So ki- ds - llog āo - +00
(54)
2 260 1/ exp ( 2 Ks ds) dl' = =Jo àب
âل_ exp (- -& âi- ds) dt'
ā1_.
= =Jo an-e! Pi-d'
ài- ai- dt' = a1 So du (2,-1) dt' = 0n- )=)==
=1 01-tt
=1 111-01-
(55)
2 20 N1 exp (-2 fe Kis ds) dt' =1 1 Onp 1.
where we used that ã1 = and āy = O. And Flow Matching corresponds to (e 3r 3t). We have that vi 2 0 because a₁ is increasing and Br is decreasing, and
ai"
de
178
So 1K-o ds = So س1
01-a ds fo 0g log 01-a ds - -llog 01 log 01) = log 011to (56)
So ks ds - log +00,
and
2 2.00 n-- exp ( (-2 fe -1-s ds) di' 2 2o 31-1 du
aa Br- 3 exp ( (-2 ài- ds) dil
dil
a._.
=2 31-1
(1 ع 021
an BI- 31-v) ( o-
() ) dt',
(57)
To develop the right-hand side, note that by integration by parts,
fo 31 1 - ai-e 2 dt' = -o Oy
al-e )² dt'
aa-ee 211
2 "a..""
= c ( 1o + fo Or (
2 dt' =
10 39 of مس dt'.
o
(58)
And if we plug this into the right-hand side of (57). we obtain
fo m- exp - 2 f k1-s ds) dt' = [ (0i-e) 911 = 89_ ()(=)) =B? =P_
(59)
2 m-r exp (-2 fe k1-s ds) dt' 67 =1
(60)
where we used that 0, 01 =
C.3.2 PROOF OF LEMMA
We can solve this equation by variation of parameters. To simplify the notation, we replace K1-s> Th and E1 by Ks, 7ss and &.. Defining f(X,, t) = Xx exp (So ki-s ds). we get that
df, t) -K1 exp So Ki-s ds) dt + exp (Jo ki ds)
= exp (So -i ds) dt +ex (.o Ki-s ds) ((-x1-eXe dt v
= =exp (Jo k1- EI dt + 2m exp (No ki-, ds) dB,.
(61)
Integrating from 0 to f, we get that
X, exp (Jo k1-s ds) = No lo exp (.o² ds) dt' +1o /2m-u exp (fo ki-s ds) d
(62)
- X, = Xxexpl ( - 1/e Ki dss + fo exp ( - fe 1i (l) di'
(63)
+ So 22 exp (- f Ki ds) dBr.
(64)
31

---

Published as a conference paper at ICLR 2025
Since
B[CJo v 22m - exp S ki-s do) dBB.) 2o nn m- exp (-2 So
kn - ds) de'
(65)
we
obtain that Sdd -2 exp ( fC Ki-s ds) dBy has
same distribution
as
/2 2 S m exp (-2 f N1-a ds) dt'c, where e N(0, 1).
C.3.3 PROOF OF PROP.
This is a result that has been used by previous works, e.g. (De Bortoli et al. 2021, Sec. 2.1), but their derivation lacks rigor as it uses some unexplained approximations. While natural, the result is not common knowledge in the area. We provide derivation which is still in discrete time, and hence not completely formal, but that corrects the gaps in the proof of Bortoli et al. (202)).
We introduce the short-hand
Mar 1) - (allz -m-)a((s,)) 1),
(66)
b(x,t) t) H (=if)² (x,,
(67)
Ø(t) o(1 t).
(68)
Remark that b(r, -1==-ip Mx, 1 1) + (()))((X,,)))
Suppose that discretize the forward process X using K + equispaced timesteps:
Øh+1 = 2b + hb(zk, kh) + (hō(kh) with c [(, 1).
(69)
is important to remark that 1 Iy = 12). Throughout the proof we will keep track of all terms up to linear order in h, while neglecting terms order (113/2) and higher. The distribution of the discretized forward process is:
p(ao.k) = (o())) IT=o0 PK+ 1j (#8+1)œ4), where Pk+ (2 [xk) = exp (- ling41 Abl AM112
2ha
(70)
Using telescoping products, we have that
(( (xx) IIK Pk+1ja (14-110
(xx) Peti (mx+1)
(wx)
(71)
= FK (x< K-1 Fk+1jK (#1+11 exp (1og(P((*x)) log(pk+ (I+++1))
We can use a discrete time version of Ito's lemma:
log p( (k+ 1)h) ss log pl kh) (( log kh) a(RA) A log kh))
(72)
(V log kh), #z+1 - xn) + (113/2).
(73)
Using equation 669 and a Taylor approximation, observe that
(V log p( #x41 - m
(V logp(1<+1, (k: (k+ 1)) v log (k: + 1)4)(xa+1 - xR), x4+) x④) + 0(1312
(Vlogp(re>i 1)h), 24+1 -xx)
(hō(,z kh) + (ho((hh(eeeee.. 2 pl (k 1)4) (hō(aa, kh) + (ho((h))e))) + 0(13/2)
(VV [Vlogp( (k+ + rx+1 zx) log (k+ 1)h) + 0(h32))
(74)
And since p satisfies the Bokker-Plande equation
Oep. = V ((-b(x, t) + 8(02
v log (x)))..
(75)
2
we have that
O log Pa
8.8 Ft.
v-(1-E(z.0)+ ووي
log (w))
(76)
= =-V =-V b(x, + 0(1)2 A (z) + (-b(xzt) + 3002 V (og V (z))
32

---

Published as a conference paper at ICLR 2025
Hence,
O logp((r kh) O. logp(a (k+ 1)h) 0(14/2)
=-V b(zg444 (ks 1)h) + n()+ A log (k 1)h)
+ (-b(#x41, (k 1)h) (((++ logp(ra+1, (k 1)h), V Vlogp(ra+1, (k 1)h)) + (h
(77)
If plug 74) and 77 into 72
we obtain
log (k + 1)h) - log kh)
- (1 b( + (-))) (A: 1)h)
{( V log M (k: ))) V log Mx (k + 1)h)))
(78)
+ V log p 1. (k: + 1)h), Fx+1 mn) + (h
(khd log pize+ (k+1)h) -g- hb(2 (6+ 1)42
2hd (RAj"
11( MaTei, (k + 1)4) + I((K+1)))' V log M(x+1, (I: + 1)4)1?) + 0(19/2).
2
Applying a discrete time version of Ito's lemma again, we have that
b(I,, kh) = b( (k + 1)h) - h (0, (k 1)h) + #(c++1)Aj Ab(xn- (k: 1)h)
+ vb(0x+1, (k + 1)4) [ (zx - IK++1) + 0(h3/2)
= b( (k + 1)h) + Vb(xz+1, (k I) (xx IR+1) + O(h).
(79)
where Ab denotes the component-wise Laplacian of Thus,
log Pk+ (1 (xx)
_ log (kh)))
|| Ab)0
Y/ log (kh)))
| mm Aeh(aa + DAAH 1 (oy ( 2111)) O(13/2)
2ha(kh)²
d log (2thō(kh))
| 11l1-a0 Ti.. hh (k+ 11M0I
(ay 1- 1I (A (36- + (
2 d log (2thh(kk))
2ha(kh)²
d(ch)²
||
hō (K+1)A011
RO(RA) tek, (1tta) +0(13/2)
2
hdd (Kh))
d(kh)²
d log (2tha((h))
||
h8(pp (K+N)
(k + +)+))))+0(187)
hd (Kh))
(80)
Combining 78 and (8))
we obtain that
log 11 (# (zw) log p(opt1 1, (k 1)h) log, kh))
log (2hhō(k))??)
| =xt hoizz (1. (KL)) po 0(h3/2)
ha(kb)²
/ log (2thā + 1)h)3) | hō(zz (K+ 1 + O(h/22).
ho((k+1)))=
(81)
By Bayes rule, and taking the exponential of this equation, we obtain
(1 (xx) == Fh+ (1 (20) (mm)
F (8441)
1 Aac (++ - (1 logrirs 11 (1
O(h3/2).
(K+
(82)
Up to the 3(13/2) term, the right-hand side is the conditional Gaussian corresponding to the update
2k = Fe+1 + h( - 6( (I + 1)h) + a(k + 1)h)² log (k + 1)1)
(83)
+ Vhō((k: + 1)h)cot1,
Ck+1 T).
33

---

Published as a conference paper at ICLR 2025
If we define Mk zK kss and use that b(x, -b(x; - o t), we rewrite (83) as
9K
UK
- - htgp::hh(00--------0------------ Vhō((( - k 1)b)eE -k-k-1 hb (hh
= Mh+1 = yk + hb(yk. kh) + v (kh)eK
(84)
And this is the Buler-Maruyama discretization of the backward process X. If we plug (22 into
we obtain that
p( (K( I - (1
(85)
which concludes the proof, as PK (zx) is the initial distribution of the backward process, and Pk++ (#0441 (24) are its transition kernels.
C.4 THE RELATIONSHIP BETWEEN THE NOISE PREDICTOR e AND THE SCORE FUNCTION
Applying Lemma with the choices of (kz)120 and (n:)20 for DDIM, we obtain that X, has the same distribution as
X, Vai a E e~ 10,
(86)
Since X, and X, have the same distribution, predicting the noise X, is equivalent to predicting the noise of &, The noise predictor can be written as:
x] <(x, 1)= =|=======)=
+ VI age = =x]=
Z-. /o, Xo
Va, + vT - age = x
(87)
And the score function s(x,t) V log ((x) admits the expression
(2)-((2) V3Di
B loglo (x
s(x,t) =
(88)
pl (1))
F1)(2))
Fi Fi))(a))
where
(=_ expl-1xz-ve Va, Yi 112(1-6)))
V log ( (i) = Vi
(89)
(2m (1-a)))
Plugging this into the right-hand side of (88) and using Bayes' rule, we get
sxr t) =E
Var X
VarKo vT- ane wr
(90)
Comparing the right-hand sides of (87 and (90). we obtain that s(x,
c(x..). Vi-a;
C.5 THE RELATIONSHIP BETWEEN THE VECTOR FIEL AND THE SCORE FUNCTION
By construction (Limman et al.
2023 Albergo & Vanden-Fijnden 2023
Albergo al.
2023)
have that
v(x, () ElarY, + 3,Yo a.Y, 3.
B - + 3,Yo 1 = a,Yi + 3,Yo
(91)
ài x + (3 a. E a.Yi + 8.Y).
ai
where we used that Yi (x B.Yo)/a4. Also, we can write the score as follows
s(x, t)= V logp:(x)=
Vp.(a) (r(m)
((pan) log (
(92)
Pc(2)
Pr(
Pe(w)
where
expl oz Vi /(2871)
Pr (x
log Pep (x
(93)
(2m 14y
34

---

Published as a conference paper at ICLR 2025
Plugging this back into the fight-hand side of (92, obtain
s(x, t) = (x|x))
"1
Fel (a Vi)pi " dYi
of
Pf
Pe(a)
(x)
= Pijo((l) -
g-a Vi dY E
a, + B,Yo]
R am oy V1 + Yol
8
(94)
The last equality holds because (x ai\BBBBB Yo- Putting together (91 and 94
we obtain that
u [z, t)
G. x + (a Be (,)s(1, t) = s(x, = Be( (e- (o(x, () at x))
(95)
Thus, the ODE 3 can be rewritten like this:
dx: Xx 3e (:t Br (x),(X.,(); Xo po.
(96)
To allow for an arbitrary diffusion coefficient, we need to add a correction term to the drift:
dx, = (1) X₂ + [ ( B (x), d) +()dB,
Xu ~ Pp-
(77)
This can be easily shown by writing down the Fokker-Planck equations for (96 and (97), and ob- serving that they are the same up to a cancellation of terms. Finally, if we plug the right-hand side of (95) into (97). obtain the SDE for Flow Matching with arbitrary noise schedule (equation (4).
D STOCHASTIC OPTIMAL CONTROL AS MAXIMUM ENTROPY RL IN CONTINUOUS SPACE AND TIME
In this section, we bridge KL-regularized (or MaxEnt) reinforcement learning and stochastic opti- mal control. We show that when the action space is Euclidean and the transition probabilities are conditional Gaussians, taking the limit in which the step size goes to zero on the KL-regularized RL problem gives rise the SOC problem. A consequence of this connection is that all algorithms for KL-regularized RL admit an analog for diffusion fine-tuning. This is not novel, but it may be useful for researchers that are familiar with RL fine-tuning formulations.
App. D.4 is providing a more direct, rigorous, continuous- time connection between SOC and Max- Ent RL, as it shows that the expected control cost is equal to the KL divergence between the distri- butions over trajectories, conditioned on the starting points (see equation 12)
D.1 MAXIMUM ENTROPY RL
Several diffusion fine-tuning methods (Black et al., 2024; Uehara et al. 2024b) are based on KL- regularized RL, also known maximum entropy RL, which we review in the following. In the classical reinforcement learning (RL) setting. we have an agent that, starting from state so ~ Po. iteratively observes state 8k, takes action ag according to policy л va 8k; 8k;k) which leads to a new state 8k+1 according to a fixed transition probability p(sk441 [ak, 8k), and obtains rewards (sk,,, ag). This can be summarized into a trajectory policy π in order to maximize the expected total reward. i.c. max, EE, E, P T a+)]]
((Sk, ag) k=o- The goal is optimize the
Maximum entropy RL (MaxEnt RL; Ziebart et al. 20088) amounts to adding the entropy H ( of the policy <- Sk. k) the reward for each step in order to encourage exploration and improve robustness to changes in the environment: max, Eyyx,pl .lEk Tk(S a&) + -K=1 HTT::s sy, k)) o
As a generalization, one can regularize using the negative KL divergence between (T|: Sky k) and base policy Tbase (:; 8k, k): EK
max E k (5k, ax) - -K- KL(<(: sk. k)]l Thase (:: Sk. k))],
(98)
which prevents the learned policy to deviate too much from the base policy. Each policy induces a distribution g(T) over trajectories т, and the MaxEnt RL problem (98) can expressed solely in terms of such distributions (Lemma App. D.3
max 16, 1e Σ k= (K( ax)] - (Llglga
(99)
Thee entropy terms are usually multiplied by a factor to tune their magnitude, but one can equivalently rescale the rewards, which is why we do not add any factor.
35

---

Published as a conference paper at ICLR 2025
where 9 ghase is the distribution induced by the base policy Tbase and the maximization is over all distributions q such that their marginal for su is po- We can further recast this problem as (Lemma in App. D.3
min, KL(allo"),
where (r) == ( exp (E (sk, ax) - V(so: 0)),
(100)
where
ks), ;= log (By p eep ( E The (s (x)))a))
= max, E Ex The (sw,, ox) )- SK-1 KI(T(;; Taabe (1 8ks, [s)
(101)
is the value function. Problem (100) directly implies that the distribution induced by the optimal policy x² the tilted distribution q (which has initial marginal po).
D.2 FROM MAXIMUM ENTROPY RL TO STOCHASTIC OPTIMAL CONTROL
The following well-known result, which we prove in App. D.3 shows that in a natural sense, the continuous-timg comtinuous-space version of MaxEnt RL is the SOC framework introduced in Sub- sec. 3.1 In particular, when states and actions are vectors in R policies are specified by vector field u (the control), and transition probabilities are conditional Gaussians, the MaxEnt RL problem becomes an SOC problem when the number of timesteps grows to infinity.
Proposition 5. Suppose that
(i) The state space and the action space are Rdd,
(ii) Policies a are specified as (ox; 8k.
8k.k) k) = ō(as w(kk kk),,kh)), where u Rd [0,7] - Rdd is vector field, and 5 denotes the Dirac delta,
(ii) Transition probabilities are conditional Gaussian densities: p(ss+1 aa, sa) (sk + ((b(sk. kh) o (kk)) ax). ho ha (kh) )), where T/K K is the step size, and b and a are defined in Subsec. 3.1
Then, in the limit in which the number of steps K grows to infinity, the problem 98 is equivalent to the SOC problem 18. identifying
the sequence of states (s with the trajectory X" (X(")
the running reward Ek= K-1 Tk(Sk, ug) with the negative running cost -o s(XF, 1) dl,
the terminal reward F (SKK with the negative terminal cost -9(X)).
the KL regularization E EEK K- KL(T(; sk, k) lTbsse ( with } times the K=o pected I?? norm of the control [
[|u(XV, f)]l at).
and the value function (88k, k) defined in 101 with the negative value function -V(r,, t) defined in Subsec. 3.1
A first consequence of this result is that every loss function designed for generic MaxEnt RL prob- lems has a corresponding loss function for SOC problems. The geometric structure of the latter allows for additional losses that do not have an analog in the classical MaxEnt RL setting; in partic- ular, we can differentiate the state and terminal costs.
A second consequence of Prop. 5 is that the characterization 10O) be translated to the SOC setting. The analogs of the distributions 9", base induced by the optimal policy л and the base policy base the distributions p" p", base induced by the optimal control u" and the null control. For an arbitrary trajectory X (Xt)1E(o,T]. the relation between p and pbase is given by
dF"
X
exp(- So- f(X,, g(xr)) + V))
(102)
where V is the value function defined in Subsec. 3.1 Note that this matches the statement in 16
36

---

Published as a conference paper at ICLR 2025
D.3 PROOF OF PROP. 5 FROM MAXENT RL TO SOC
Since the transition pl 8 ags Sk) is fixed, for each x we can define
Tak, 81+ Sle, k) π (axi Shs k)p((k laa s4) and base (ag, S4415 8k, k) Thase (0g5 8p, k)p((ss (xk 8k);
(103)
and reexpress (98) (see Lemma
min E E - Nk (sk,
EkE KL(#(;; 13 kk, R)ll 8k: k:)]]
(104)
Using the hypothesis of the proposition, we can write
it (akk Sk+1i8k: k) = ō(az u(sk. kT)) + n(b(sk kn) + (o(n))))))))))))))) (105)
5(ax u(sk. km) (*R+1i 8k, k),
where 7( k) (ss + (b((s, kn) + o (kn))((S,, kn))), (kn))) is the state transition kernel. We set the base policy as Taase(ass 8k, k) 6(az), and we obtain anal- ogously that a(ag, Sk+1i k) 8(ax base (Sk+1; Sk, 5k, k) with Abase (5k+1: Sk,
N(s + 1b($2, no (knn)o) T). Now, if we take K large, the trajectory (#x) generated by ft can be regarded as the Euler-Maruyama discretization of a solution X of the controlled SDE (8), while the trajectory generated by Thasn is the discretization of the uncontrolled process X obtained by setting 0. As consequence
8k, limk E E E K--I- KT(((,, (8 8ks basse ( k)]]
(106)
= =limx E Ek- KL(#(:; 8k: k)]l base (:; 8k, k))] Ex- [log dF" deo (X")),
where Ipu and IPo are the measures of the processes X" and Xº, respectively. The Girsanov theorem
(Thm. 2p implies that log dP" dD=
So (u(XF, t), dB,) - Jo [|u(XX,, t)l2 dt, which implies _ Ex=~p" 1fO [|([X] ,)) f)l dt). Setting the rewards (ak, sk) (sx: kn) for & E {O..... K - 1) and "x(a,, 8k) = (g(s)). where f and g are as in Subsec. 3.1
(X")
deT (X")]
that Ex==r= [log dPo
yields the following limiting object:
limg +o0 E, E,wa E ax)) = Ex==p= So f(X,, t) dt g(X+)..
(107)
Hence, the limit of the MaxEnt RL loss (T04) is the SOC loss (7).
Lemma 3.
a(ag, 8k+1; Sk. k) and Thasa (ak, 8k ; 8k; k) be as defined (I 88, "hnell (: 888 k))l and KI(T(:; 8k. k)]ldrasel ( 8h. k))] are equal.
103 Moreover,
if g, q hase denote the distributions over trajectories induced by T,, Taase, we have that
KL(glla bane) EEEK- (11 Sk: k)] Tbase (:: k))]
(108)
Proof. We have that
KL(#(;,;; Sk, 1)]l busee k))] Ear a(ak, Sk+1; 8k, k) log R(os 1804.884
Whosen 0L 9641188
= -Eaisess n Taas 8k, k))((s lax, 8k) log kklusss.k)p(sss (as.ss)
= -aaioo maks 8k, ()) lav; 8k) log =(a4:on.k)
(109)
Thann [k
-Ea, Tlani 8k, k) Exs p(52+1 ( sk)) log r(a
Thnee (axise
-aa T(a;: Sk. k)) log x(aw:se.k)
KI(T(:: k:) Tosse 8k, k))].
To prove (08). by construction we can write
q(T) = (s))
I मे (ap, Sk41 Sk, k).
base base = -Po(so) [ Abase (ake SK+13 Sk. k).
(110)
37

---

Published as a conference paper at ICLR 2025
which means that
(L
= BBrgglog flog (bw)(r) E, Brag Ex=o log #(og:9611198) Frase (ax (8x.51+2)se..k)
((r)
=ΣK- E (a [log &ban (ax .k 1;0x.k))
(1+t))
ΣK-1 Broggtt Eec.oo # #(ax, 8b-11 8k, k) log Rbown #(01,0241100.
(111)
-E E Brig [KL :; Sk. k) Tbaee (,, 8k, kk,
EK-1 E Braggrr [ (T: , R:] busec (:; 8k; k))]
Brngutt k (-( $2, k) Tase Sky k:)]
Here, the notation y 90th denotes the trajectory y up the state 8k-
Lemma 4, The distribution-bau based MaxEnt RL. formuumuttionn in (99 is equivalent to the the following problem:
bame (1) exp EE (z(aE.
ming (L(g)la"), where (T) ==
(112)
Po(Po] (1 q (r") exp (Ek ( r(((),a)))))
where the minimization is over q with marginal po at step zero. The optimum of the problem q", which satisfies the marginal constraint. The following alternative characterization of q" holds:
(T(T) - (T) exp E Fk(Sk, ax) - V(50, 0));
(113)
where (xz max, Erwx, Σ = Th (sw, ER' K-1
K
KL(T(:; sk. IT (1; sk' k))) 8y x]
(114)
Proof. Let us expand (g) y²"):
KL(glla")= Er-ug log
a(r) (F)
Brgg log ((T) log base (() Ta
log po(co] Σ (T11666 q (T') exp
ab))]]
(115)
=KL(gllabawn)
E Ennng [EK TK(sk, ax)]
+ Esoopo
[log po(eo) Σ 1 ( exp (Eko r of))],
where the third equality holds because the marginal of q at step zero is po by hypothesis. Since the third term in the right-hand side is independent of this proves the equivalence between (9) and 112p.
Next, we prove that the marginal of q² at step zero po:
(ra exp EK ra(s).0)))
((rloo=)) (T) = E(r Pola) E(r))) g exp E ra (s(..()))
= Po(x))
(116)
Now, for an arbitrary let qsu Tig
be the distributions g, q" conditioned on the initial state being 80. We can write an analog to equation
(15) for qmo In0
KL KL(gaoll4tp)
= Eroogo [log
Lan (r)
4i
Eraggo
[log Joo (T) log base (T)
Ek rkk (sk, at)
+ -log (ootoo E E(T(lab (77116==)) 980 (as)((r) exp (x
(117)
KL I 1ybo base)
Eragzo EE T'k (k( ( ax)]
log (potao) E (7)166=00) (ase exp o
rk(ske d())),
Hence,
0 = milgoo KL(gxoll = maXq.p EEroag Ex (Sk, ax)l - KL( [1gbsse]
(118)
+ llog (po(ao) Σ E(r(166=00) q base (ase(7") exp T ak))).
Th(sk,
38

---

Published as a conference paper at ICLR 2025
And applying (088) from 108 obtain that
log ( po())) Σ ( 9 hase (7) exp (Ex=o ( d(b))
= , B,Nx,p Ex=o (8k: ax) - ER= - KI(T(:; 8k. R)]l (;:8: 8: [o) = V(80. 0),
(119)
which concludes the proof.
D.4 PROOF OF EQUATION THE CONTROL COST KI. REGULARIZER
Theorem (Girsanov theorem for SDEs). If the two SDEs
dX; = b, (X2, t) dt + () dBr, Xo = Tinit
(120)
dY, (Y, 1) + b2 t)) dt oY,, t) dB,, admit unique strong solutions [0. T) for any bounded continuous functional & on C(O0, T)), we have that
= Finit
(121)
[ = E 4 Y exp ( 10- -Y.,,)))16 62,, t) dB; 5 (X(,,)))
E ( exp ( -o- oYY..t) b2, by t) dB. + 21.0 to).)))) by t)ll dt)
(122)
where B, = B (Y,, -b 8) ds. More generally, by and by can be random processes that are adapted to filtration of B.
Consider the SDEs
dX, = b(X,, (I)dB.,
Xo = Jo
(123)
dxT (b(X), + ol))(Xp) () de + o(I)B
X IQ.
(124)
If we let IP mo P" be the probability measures of the solutions of [23 and 124 Thm.
implies
that
dF (X") = - "XXx(t()))) 1 Su || 2) dr. a. Iap
log
(125)
Hence,
DKL (IP"
Izo
IP|...) = E log dPIzo
dP" lan (X (I = z0] = - log -F
(X |X = mo]
1l (x
E So u(X) t) dB, + 1/2 So' ||x(X;, t)| dt| zo]
(126)
E [ So [|u(X," 1 xt = _o].
where we used that stochastic integrals are martingales.
E PROOFS OF SUBSEC. 3.3; MEMORYLESS NOISE SCHEDULE AND FINE-TUNING RECIPE
E. PROOF OF PROP. 1 THE MEMORYLESS NOISE SCHEDULE
51 with arbitrary noise schedule. By Prop. the trajectories X, X of these two processes are equally distributed up to time flip, which also means that their marginals satisfy p. P1-1, for all t E [0, 1]. First, we develop an explicit expression for the score function (ax, t) = V log (x). By the properties of flow matching, we know that p. the distribution of the interpolation variable X, = 3,X, + 04 X, where X, ~ N(0, 7), X, ~ p are independent. Thus, 3,
We consider the forward-backward backward SDEs 50
X N(O, I), which means that we can express the density p. as
exp 1
(x(x)= Suu
287
data (y) dy.
(127)
(2m
39

---

Published as a conference paper at ICLR 2025
Thus,
Sad yexp
(- 1i=-au11 12
p ddnnn (2) dy
s(x,t) = log
2 5f
2,1
5 (.
(128)
Sad exp
11
o 2
(ythan(()) dy
8f
20/
where we defined
Ja y eexp
Is-owy 201
2
10
data (v) dy
(xr)
(129)
Sau exp
[ 207
ull"
pdistin (2) dy
Hence, we can rewrite the forward SDE (50p as
dX, = (- (of1-0² 1/1 m/-1) X -
di + o(1 - -l)dB,
Xo ~ Pdata
(130)
Hence, if we substitute Ki-t Ki-t + p(1 ((1-1)8-m11- f1-t & (0== -mm
(X,)
23
28
(where we ignore the dependency on X.). 2111 o(1 - 1), we can apply Lemma yields
2 which
X, = exexp (- S0 (mI-s + a(1 ((1-s))-211 27110 ds)
2/
So exp ( S (Ni-o + o(1-))2-21) ds) 01 (01--(--)))-- (Xx) dt' (131)
209
209
+ +I0 al - 11)) exp ( (- Se (ki-s o(1-9)² o(1-9)²-211- - ) dss Br.
23
We simplify the recurring expression:
di-a ((1-s)--28) ( (0--1 B -B
K1-s+ + 011-8) x)² -2
a(l=s)²
A-
+
(132)
237-.
28
28
Bi-
Thus,
Se (Ki-a + 0(1-9)2-20 2m )) ds J/ 237
(0(1-0)²
O. log B1- ds
o(1-s)² ds (log 3 log 31-v), 239
23
(133)
which means that
exp ( - f (kI-0 + o(1-a)² 2211 0(1-0)2-211
ds) = exp ( 1
a(l-a]² ds) B
(134)
203
20f
a1-r(a(1--))) -211--1----))) 2th 201-")
E1 (Xe) a1-1' (0(1-112
+ Bn-
B aq (1-)) (135)
257_1"
23 23
If we define x(1 - 5) such that 0² - 8) = 2,31 -8 (a) B1_ - /1-5) + x(1 - $), we obtain that
exp (-Je aلا-اا² -
ds) 8i-L cxp (- -/' aiir
Bi- Bi-e
+ x(1-s) ds)
$1-1
2 Bi-e!
2p
Fi-t'
= exp (Jr (0, log a - O, logog Bi_-
x(1=s]
ds)
B1-t exp (- 1 28
x(1=s)
ds
23
B,-e!
(136)
0111' (
(+ 28 3i-e!
ex(1 17)
(
(137)
2
If we plug equations
(36)
137
into
134
[155] and then those into [31 we obtain that
X, Xxexp( ( ( (o x(1-s) 28,_, an
ds)
ai-t 01 So exp ( - -p x(1-5) ds) x(1
S r (Xx) di
28 28
So (2 di-t'' B1- //rr + x(1 - t1)) cxp (- S 203 a. (138)
x(1-)
ds)
dBv.
40

---

Published as a conference paper at ICLR 2025
and if we take the limit 1 and use that a1 1,
X1 Xo [iim--1 exp (- So 29
x(I=o) ds) 01-1) lim a1 exp ( Se x(1=0) ds) x(1
23 20 E1 v (Xx)
dt
+ lim So (231- ليي B 31- + x(1 - 1)) exp ( (- Si ((اا ds) بيي dB _8 aj-ز (139)
The assumption on in 19 is equivalent, up to a rearrangement of the notation and a flip in the time variable, to the statement that for all 1' [0, 1).
lim exp
- l x(1=) ds) -1 = =0 0.
(140)
2pr_.
Hence, under assumption 19 the factor accompanying Xo in equation 133 is zero. Moreover, this assumption also implies that
lim 01 Sü exp (- Jer x(1=s) ds) x(1-12
f1- ( (Xx) de
20 23
(141)
(lim,-1- exp ( J x(1-s) ds) dss)n-1)
x(1 a(1+t),
Xx dt' = 0. E1
-
28
If we plug (140) and [41 into (139). we obtain that
X1 - Jo (281 ( B - Bi-e) + + xx(1-t))
t')) exp (- -t x(1-s) ds) tnt dB (142)
201 arーd
which shows that X, is independent of Xg. Next, we leverage that X and X have equal distributions over trajectories (Prop. 4 In particular, the joint distribution of (Xo, X1) ) is equal to the joint distribution of (X1, Xo). We conclude that Xi and Xo are independent, which is the definition of the memorylessness property. Hence, the assumption (19) is sufficient for memorylessness to hold.
It remains to prove that the assumption (19) is necessary. Looking at equation (138) we deduce that generally, for any t E [0, 1). Xo and X, are not independent, because the first two terms in (138) are different from zero. Thus, if there existed t' E [0, 1) such that the limit (140) is different from zero, then X1 would not be independent from Xy,, which means that in general it would not be independent of Xo either.
E.2 FINE-TUNING RECIPE FOR GENERAL NOISE SCHEDULES
E.2. EXPRESSING b, u IN TERMS OF o OR
Before proving the result we clarify the relation between the DDIM and Flow Matching vector fields u. and the base drift b and the control Let base base denote the pre-trained vector fields and ciintuuee ffnnesuue the fine-tuned vector fields. Then we have the following expressions for the full drift b(x, t) + (1)u(x, t) and control u(x, t) when (1)) 22nc:
DDIM DDPM:
b(z,t) ((t)))))()
t)=
ân x &e Jiuctune 200
=(=,t)=-- V 8e (1-a)))) (e Ainetune (x., t2,t)= f) base (z,, t)).
&e vI-
(143)
Memoryless Flow Matching:
2 bxx (oe,t)+o(0 (())))) fnetume (x, 1)- der. u(x,t)= Aal
Snetunee (xr, b (3r, ())-
ar
%.
(144)
Thus, to solve the SOC problem 7 in practice, we parameterize the control u in terms of finetuuee gftuctune and optimize these vector fields instead. After plugging in (143)- (144), the
SOC problem -7-8) { can then be solved using any SOC algorithm in order to perform fine-tuning,
and we propose especially effective algorithm in Sec. 4 Adjoint Matching. After fine-tuning, cfinctuee and fineetuee can simply be plugged back into their respective generative processes (30)
to sample from the tilted distribution (I) using any choice of diffusion coefficient.
41

---

Published as a conference paper at ICLR 2025
E.2.2 PROOF OF THM.
The proof of Thm. relies heavily on the properties of the Hamilton-Jacoii equation:
Theorem (Hamilton-Jacon) Jacob--ellla equation). we
we define the infinitesimal generator
L: 2= 2 (a0))))))))))) (t)(t),,, Ei
t)0z.: t,
(145)
the value function V for the SOC problem 7-8 solves the following Hamilton-Jaccob-Belllmn
(HJB) partial differential equation:
0,V 1) = LVV (x, ()
211(a VVV) D112 - (xx,t),
(146)
V(z,T)=g((a g(
Consider forward SDEs like 50p, starting from the distributions phase and pt, where p¹ (x) x pbase (x) exp(r(x)).
dx, 6(X,, t) dt (t) dB,
No ~ pbase, Xo Xa~pi pp.
(147)
ax; 5* (X; 1) dt + (t)
(148)
where the drifts are defined as
bizz - + (1(1)) 1/1 1/-1 5(zz 1 - t) --1-1- + (o(1-1)
7/1 V log (xx,
8' (z, t) = -- + (1-1)² 7=-))5 (2, - t) = - (1-1) 1/2 V log P/ (z).
(149)
and pe. P?, are the densities of X,, X,, respectively. P., Pt satisfy fokker-Planeh equations:
d.p: = (b(a, t) + V (لاام) Vp.):
Po = p base
2
(150)
0.P, V (v, 1P +
Vp
Po = p"
2
Plugging
149) into 150 we obtain
0.p. V (kI-)
Vp.),
Po p base
(151)
0.p, (k11)1p)) V ()-+
Pu p"
We apply the Hopf-Cole transformation to obtain PDEs for log pt (and log p. analogously):
log = d.pt (-( (n-
= =1 V x + (0, V log + 1/1- log exp[log pc))
(152)
=k + (I, V log ph)) 171 1 log + |/ log Pe 13)
Hence, if we define v(x, ()= = -og (xx 1) = - -lgg% (x), then & and y satisfy the following
equations:
-. = -I- K1 -1- VX) + + AY + |3)))))
V(x, 0) - -log base (z)
(153)
-v,y --- - 11-
V + 77 AY + ||
(' 0) - log p" (1:).
(154)
Now, define (x,t) = 1 (x, t) - (xx t). Subtracting 154) from (53).
we obtain
-0,8 -KI --
1
-- VX) + -- ||
= -x +
(155)
+ <(
= (, -
f(, 0) llogp' (x) + logg (x) -x) + log
p ba exp(r())) dy).
42

---

Published as a conference paper at ICLR 2025
Hence, & also satisfies Betillton--acbi--Belllll equation. If we define V such that (x,
V(x, 1- t), we have that
8,V = (-KE - 21,, (), VV) + TM AV + [[VV]?) V(x, 1) = r(z) - loo (J h (y) exp(r(y)) dy).
(156)
Using Thm.
we can reverse-engineer V the value function of the following SOC problem:
min B So [|u((xr, dl log (ʃ p (v) exp(r()) well
(157)
s.t. dx,"" = (N², + (z, t) (X t)) dt /2
XØ ~ Po-
(158)
Note that this SOC problem is equal to the problem 7 8 with the choices 0. g = -, and
(t) 2nh- By equation (11. the optimal control of the problem (157) (558) is of the form:
(x, = - 2mm (V, - 2m (z, 1 - t) 2 VV (z, - (V
- 21) ( V log (x) + V log logpi-((2)) = - (zz)))
(159)
s* (z, t) = s(r, t) + u² (z, t)/ 2
(160)
As in 5T. the backward SDEs corresponding to the forward SDEs (148) take the following form:
dx; = (mi:; + (o +
o(t)
X N(O.)).
(161)
If we plug (160) into this equation, we obtain
dx;- = (M.X; + (0)2 2 + m) (s(Xi, "x(x+,) dd (()) 12m
XØ ~ N(O,
(162)
a()) + ur (X;, t) dt + (t) dB,,
= dx; (b(x;,t) t)
X ~ N(O,
(163)
where we used that b(xz, t) = Ke0 + (=ih)² + Te s(x, t) by definition in equation 6
The fine-tuned inference SDE for DDIM Now, for DDIM, we have that u² (I,
(E, base (x, t)) by 143). Hence,
2(12 t,
[())
u (I,
(e" (e"(x, base base (x, 1))
(E (x, () - e base base 1)),
22m
(164)
= b(x,t) +
2
trtt
u² (Ir, t)
de
X;
du
o(1)
base (Xi.))
2011
(E" (x, t) base (I, t))
V2n
20r
200
v1-0r
20re
X,
20r
+ 2
21)2
c" (X..t)
VT-ac
(165)
We obtain that the fine-tuned inference SDE for DDIM is
dx
Xi
de.
(X7 di + (1)
X5 ~ ~N(O,,))
(166)
Zore
vI-ai
which is matches the SDE
29 with the choice 6"
43

---

Published as a conference paper at ICLR 2025
The fine-tuned inference SDE for Flow Matching For Flow Matching, have that u (x, t)
(v, 1) basse (x,()) by (14)). Hence,
B. Bc)
2
0603 +B( 8: Be)
2 + u (2)
Del ue 8. B.) (2, t) - 2 buse (., t))
2
2t
V /2 2e B1 (
(167)
(1
o11)²
(v" (2r,1) 1) v base (5,
28. ce 8 3-3,)
=02 + u (x, 1) = y base (z. 1) + 20r( mq -e- (1 (1,, 1) ar.r)
= bxx +
r(t)²
22m
+ (1 + 23x
16'(1,,4) - b.see (ba ()) 21 9, 9.)
a(t)²
(168)
"d (z, t)
o1)²
(w' (x,,
âi 17).
20. ( 3.)
on
We obtain that the fine-tuned inference SDE for Flow Matching is
dx;- = (v(Xi,))+
(() 29.( B BJ
(0 (Xi, ae xi)) dt + o(t) dB,
Xō ~ 1),
(169)
which matches equation 4 with the choice v².
F METHODS TO SOLVE SOC PROBLEMS
F.1 EXISTING METHODS
F.1.1 THE ADJOINT METHOD
The most basic method of optimizing the simulation of an SDE is to directly differentiate through simulation using gradients from the SOC objective function. The adjoint method simply uses the objective:
C(u; = So ()) dt g( X ~ p"
(170)
This is stochastic estimate of the control objective in 77. and the goal is to take compute the gradient of L(u: X) with respect to the parameters 9 of the control u. Due to the continuous- time nature SDEs. there are two main approaches implementing this numerically. Firstly, the Discrete Adjoint method uses diiieretize-thenddiifeeeettt
approach, where the numerical solver for simulating the SDE is simply stored in memory then differentiated through, and it has been studied extensively (e.g., Bierkens & Kappen (2014); Gomez et al.] (2014); Harrmann & Schütte 2012); Kappen et al. (2012); Rawlik et al. (2013); Haber & Ruthotto (017)). This approach, however, uses extremely large amount of memory as the full computational graph of the numerical solver must be stored in memory and implementations often must rely on gradient checkpointing
(Chen et al. 2016) reduce memory usage.
Secondly, the Continuous Adjoint method exploits the continuous-time nature of SDEs and uses an analytical expression for the gradient of the control objective with respect to the intermediate states X,, expressed as an adjoint ODE, and then applies a numerical method simulate this gradient itself, hence it is referred to as
approach (Pontryagin 1962. Chen
al. 2018 Li et al. 2020). We first define the adjoint state as:
att: X) = Vx, (S. (3 (MXe, + (Xe, 11) di' +g((1)),, where X solves dXx; (b(Xe, ((t)a(Xe, 1)) dt + ot)dB.
(171)
44

---

Published as a conference paper at ICLR 2025
This implies that Ex-p" la(t; X, |X, x| (u; t), where J denotes the cost func-
tional defined in 2 It can then be shown that this adjoint state satisfies
att; X, u) = at(: x, u)
+ VVx, (f(X,,t) 2 +) 1)
de
(172)
a(1 X, x
(173)
The adjoint state is solved backwards in time, starting from the terminal condition 173 Com- pututation of (122) be done with vector-Jacobian product which can be efficiently done
automatic diffcrentiation software (Paszke et al., 2019). Once the adjoint state has been solved
for & E [0, 1J. then the gradient of Z(u; X) with respect to the parameters 0 be obtained by
integrating over the entire time interval:
de = 2 Jo a (X,, dr + fo 80
au(X..)) at)) X, u)dt, 09
(174)
where the first term is the partial derivative of L w.r.t. 0 and the second term is the partial derivative through the sample trajectory X. See Prop. c App. F.2 for statement and proof of this result. The discrete and continuous adjoint methods converge to the same gradient the step size of the numerical solvers go to zero. Both are scalable to high dimensions and have seen their fair share of usage in optimizing neural ODE/SDEs (Chen et al. [018 2021 et 2020). As the adjoint meth- ods are essentially gradient- optimization algorithms applied on a highly con-convex problem,
many have also reported they be unstable empirically (Mohamed et al. 2020 Suh et al. 2022 Domingo-Enrich et al. 2023
F.1.2 IMPORTANCE OPTIMAL CONTROL
-EERRTAANE-WEIHHEED WEIGHTED MATCHING OBJECTIVES FOR REGRESSING ONTO THE
An alternative is consider regressing onto the optimal control "² which is the approach of the cross-entropy method (Rubinstein & Kroese| 2013 Zhang et al. 2014) and stochastic optimal con- trol matching (SOCM: Domingo-Enin EnrrccEEicchhhi et al. (2023)). These methods make use of path integral theory (Kappen] 2005) to express the optimal control through importance sampling, resulting in an inportance-weighted least-squares objective function
(u; X) = I|u(X.,t) a" (X,, 1) w(u, X),
X~p",
(175)
where w is an importance weighting that approximates sampling from the optimal distribution p" and û" is a stochastic estimator the optimal control relying on having sampled from the optimal process. We defer to Comingo-Ennich et al. (2023) for the exact details. The functional landscape of this objective is convex, which argued help yield stable training. However, the need for importance sampling renders this impractical for high dimensional applications: the variance of the importance weighting w grows exponentially with dimension of the stochastic process, leading to catastrophic failure. This unfortunately means that such importance-weighte matching objectives are impractical for fine-tuning dynamical generative models; however, a least-squares objective greatly coveted as it can lead to stable training and simple interpretations.
F.2 DERIVATION OF THE CONTINUOUS ADJOINT METHOD
Proposition 6. The gradient de dø of the adjoint loss L(u; X) defined in 170 with respect to the
parameters of the control can be expressed as in (74
Proof. First, note that we can write
VO S (Xll ( +SX?" 6) d(+g(Xx"))
E [ (Xuo t)) .t) + VoE[Jh (XI'(XX",)))P+++xXXX +(X)++(XX'))) dt dt+g(XX))) |==stopgrad(v))
(176)
Note use the convention that Jacobian matrix J V.v(() is defined as Juz (x,
45

---

Published as a conference paper at ICLR 2025
To develop the second term, we apply Lemma 5 Namely, by the Leibniz rule and equation we have that
181
VOE [L (2ll +(X;" 6)) dt dg (==stopgrad(ua) V (H) 1 1)) dl+g( (x7)) (t=soograd( = (Vou)))) t) (t)) aa (w) dt)
(177)
Plugging the right-hand side of this equation into (76) concludes the proof.
Lemma 5, Let v be an arbitrary fixed vector field. The unique solution of the ODE
dt alll; X u) = [ (xxxX,)(XX,)) +(0(())0))))))))) a(l; x", + Vxr (S(X), ()
(178)
(11: u) Vg((f),
(179)
satisfies:
att: X := Vxr
dt drr+
(180)
where X solves dX, = (b(X),)) + (o())((X) 0)) dt + o())d
Moreover, when u == ug is parameterized by 0 we have that
volSo 100XX ()) dt + = 5 (Vous) (w), ()((t) ae(w) dt.
(181)
Proof. We use an approach based on Lagrange multipliers which mirrors and extends the derivation
of the adjoint ODE Domingo-Enn Ennich et al. 2023 Lemma 8). For shortness, we use the notation
bz( t) == b(x, t) + (t) t). Define a process a 12 x [0, TI - Rd such that for any w E ? .)w, is differentiable. For a given w 92 we can write
So" 1-' (6 + ()) (x+)
= = T 2 ( + t)) dt+
(182)
S (a,(w,, (dx;" (w) by (x, (w), dt o(t))) d())
By stochastic integration by parts (Domingo-Farichee et al. 2023 Lemma 9), we have that
So (eelw),X,"" (w)) (aT(w), X (w)) (an(w), X (w))
-o' (X! (w), das
di (w)) dt. (183)
Hence, if X Tu the initial condition, we have that
Va (5o (H)lo(X7)) -00=) dx+x)())))
V.. ( ( dt+
- (ar(w), X (w)) + (aolw), X," (w)) + So bol(w)),b((x"" (x)")))))) + dan (w))) dt
+ + (t) dB;))
T V N0 X," (w))
|o (X." +[ (w) dt + Vao (w)) V= (w))
V X (w) [T V X."h (w) ao(w)
S (V
bo ( (w), 1) (ww + V.o (w) des (w:) dt
di (w)) dt 5 V Va, X (w) V. lo ( 1) 20 (w), 1)) + V. bx (w),t) ar(w)) + day
+ V X (w) (V.g(X (w)) ar(w)) + an(w)
(184)
Unliie Domingo-Enmi Enrich al.
2023 Lemma 8), we use the convention that a Jacobian matrix Their definition of V is the transpose of ours.
V Vx0(2) is defined as Jaj Bay
46

---

Published as a conference paper at ICLR 2025
In the last line used that Vao X00 (w) V z0 = I. If choose such that
dae (w) = ( V. (x² (w),t) (a(Ww=-(w---))))))))(()))))))[ - V ( (x .00.",1))+++0 ar(w) = (XT (w)),
dt,
(185)
which is the ODE (178)- then we obtain that
xo (.6" ())l(((2,,))) f(X?" 1)) dt dt++g(X²)) = ao(w)
(186)
Without loss of generality, this argument can be extended from ( = 0 to an arbitrary 1 E [0, 11. which proves the first statement of the lemma.
To prove 181. we similarly write
Va So ,0I b)
dx+g
Vo ,) t)) gx+)
(w), ( dt - (ar(w), (w)) + (ao(w), _ (w)) + 1o¹ ((a(( + (w),
+ (ai(w) (B))
(o) (X (w), dt Vo. (w) ( V=g(<>> (w))
- VeX*
+ + (w) [x) (4),
(x(ww) VoXt (w) de (w) di
dos
fo VoX," (w) (V ,00, (w) + Vabo (w), (w(w) + du (w:) dt + VoXt (Vz( (X7 (w)) (T((()) + S (Vabs) (w) dt.
(187)
In the last line we used that VoXo (w) Vox = 0. When satisfies 185). we obtain that
Vo (Hllv(x," .0IXX""))))+ 111 ((X ()) dl d(+g(X7)))
(188)
= (Vo(ooo)) (Ww(()) dt = fo (X (()) aT
The last equality holds because by (Ir 1) = b(x, 1) + o1))ua(1
F.3 PROOF OF PROP. 2 THEORETICAL GUARANTEES OF THE BASIC ADJOINT MATCHING LOSS
Let stopgrad(uo). We can rewrite equation 174 as:
VaL(uo; x²) _/2 fo V Juo X2 dl + So' Vau 1)T (o) a(t; X a)jd = 1
(189)
2 fo V (X, + +tt a(t: X () = VoLBasic- Adj---d------- Match (uo: X²)
(190)
This proves the first statement of the proposition. To prove that the only critical point of the ex- pected basic Adjoint Matching loss is the optimal control, we first compute the first variation of ElLnasic- Adj Matchh) Letting v Rdd x [0, Tj
+ Rd be arbitrary, we have that
E[L Adj- Match (u + ev; X } = 1 lo I1(u + (x)) + +tt) X", a] de
d
= [[[ (u(XF, t), uxx o att, _ x)) dt]
_E[So" (o(XF, f), u(XP, ( T X x) XF] di)
5 Au
E(Casicc Match (u) u(z, t) E a(t, X x) =
(191)
Hence, critical points satisfy that
u(x, t) =-a(t) TE (aaat, X", x) = = 1." (4I ( (
= =-o(1) (XTx(X))))++0000 (x) 1, 1), (192)
47

---

Published as a conference paper at ICLR 2025
In this equation, the second equality holds by equation (880) from Lemma and the third equality holds by the Leibniz rule.
Lemma 6 shows that any control u that satisfies 192 is equal to the optimal control, which con- cludes the proof.
Lemma 6. Suppose that for any E R L [0, u(a, () -o V [u; l). Then, [u; ,.) satisfies the Hamilton-Jacobi Belblan equation [146.. By the uniqueness of the solution to the
HJB equation, we have that [u, (0, 1) V(x, 1) for any x Rd, [O, Hence, u(x, ()
-o(t) (a,t) the optimal control.
Proof. Since J Ju; =[ (allu(Xy, +)12 + f(X) 1) ds + 9(x;) = æt], we have that Ju; E [J +t]|x, 2] [. 1 (bllu(X) s]]| f(X"'' s)) ds X, =
=
(193)
which means that
[[J(u; X Art + At) 2] 2u;
E efte (2 s)1? - f(X", 8)) ds X, x)
At
A
(194)
Recall that the generator T of the controlled SDE 8 takes the form:
(T,(x,)) lim AC-0 -|x, (xx,))
Ar
(195)
-,,((x, 1) + (V (x, (), b(x, + (1)u(x, 1)) + Tr( al (x)(x)))
2
Hence, if we take the limit At - 0 on equation 1944 we obtain that:
T", xr, + (2, 12 tx)
J,;(u; x, () [V](u; 1), (x,, 1) + (1) () + Tr( oitaai" 42 (2J(uxx + =1lu(x 1) + f(x, 1).
2
(196)
Now using that u(x, () (1) T V V. (u; t). have that ()) + (1) 1)12 (1) (o(())TV.c V.,(,, 1) V 0)
(l) V. [u,
(197)
Plugging this back into 196.. obtain that
0 = (u; + (VJ(u; (u; x, t), b(r, t)) + Tr( (t))(()) 2
V Ju,(x, f)) - [llo(t)T Ju x, t)l + f(x, t).
(198)
since Ju;,,, 7) ( by construction, we conclude that Ju,,r, satisfies the HJB equation
146).
F.4 THEORETICAL GUARANTEES OF THE ADJOINT MATCHING LOSS
Proposition 7 (Theoretical guarantee of the Adjoint Matching loss). The only critical point of the loss Mathh) is the optimal control w*.
Proof. Let e be an arbitrary control. If ātt; is the solution of the Lean Adjoint ODE 26) 27 it satisfies the integral equation
a(t;X*) STT (VaM(X; Vo,( VV
(199)
Hence,
B[ā(t: X [ =E[,," (V_b(X) +x V (x
=E[J," x) X - Vx](X,, do [x({))
(200)
48

---

Published as a conference paper at ICLR 2025
where we used the tower property of conditional expectation in the second equality.
Similarly, if all; X", v) is the solution of the Adjoint ODE 172 173). it satisfies the integral equation
all; x = -ee =
+ (x(((, V {ll s)] ds Vg((;),
(201)
and its expected value satisfies
E|a(t; X xy]
=E[
als; + +) (Xxx)))))))) ds [
=E[J (Vx(, ",dd +) 9))] ds+ V (xx))) X
(002)
Let us rewrite Matchl as follows:
B[Cadj Match (u)) = [[Lo || t) + (t) E[ X() Xx]]] dt| Immstopgrad(()
+ +E (1(()) (E(a(t, x x.] a(t, xx)) 1 at) (==stogrrad(s))
(203)
Now, suppose that a critical point of LaLi Maccll By definition, this implies that the first variation of [[[Aadj-Matt zero. Using (203). we can write this as follows:
0 = Su
R Match (Q) = 2 ()) + (1) E[a(t, X = -
(204)
ū(æ, t) -o()) Bla(t, x |X = xr.
(205)
Hence, we have
V (x ( xx),, u 1) (xx;t)) t) 0, [ x X + [|(((X)))))1) ds|x?]
(206)
E[J," (V=()(s) (BBns))))BBBBA
(207)
If we set = á in equation (00). and add (207) to its right-hand side, we obtain that E(ā(t, X x also solves the integral equation
a(t x | =E
B[a(s:
8) + s)((3) ds Vg(X+)) (208)
Note that this integral equation is the same one as equation 202) when we in the Prop. g states that the solution of the integral equation is unique, which means that E [a([:; X 4) X for all E [0, T).
latter.
E [[a(t;
X
X
Since we can reexpress the basic Adjoint Matching loss as
(u)] = E [ |u(XF, to(t) T x') I| dt] l==stopgrad(()) E [[J (Io(t) (B[a(t u)|x]] a(t: X",
x", v)) I| dt) stopgrad(w)>
(209)
obtain that when à is critical point of TE[CAdj-Matech
du E Li
Adj Adj- Match (+)) 2( t) (t) E[(a(t x x) = =x])
(210)
2(((x, - a [E x ] = x]) = 0,
where the second equality holds because E [[t x x E[a(t; X a) X and the third equality holds by equation (205). Thus, deduce that the critical points of A/LCAd--Matd are critical points of E[L Adj -attcihhh By Prop. B[C Adj Match] has a single critical point, which is the optimal control u², which concludes the proof of the statement for [f Adthdj-Mathh) n
49

---

Published as a conference paper at ICLR 2025
Proposition 8. Let u be an arbitrary control. Consider the integral equation:
Y, == (V,("(X,, ())))))(()) "x) V
( s)])) ds+ Vg(;;) x
(211)
where & E [0, This equation has a unique solution, i.e. ifY?, Y are two solutions then Y = Yg.
Proof. Let Y!, Y² be two solutions of the integral equation. We have that
Y - Y?? = =E[ (OX, - -Y3) ( s)) X;].
(212)
Thus,
|X] - 1
S [Bl Ef VV 8))d. s
[|ds]
[;]
B [L, I || -Y?1 V.b(X, 9) ds [;] = -," B[IlY - -Y?1 VV 8))
x ds
$ S (E 12 [ s) (xxx))))
(213)
And this implies that
(o. (Bil ,,11'1x;.1)) 1/2
$ J (B ill! - x331((1x;1] 1/2 ([X .) 1)|]p> 1s2ds
(214)
$ J su ( - r2112 1/2 supre(o..) (E V b Xx; 8) || x 1/2 ds.
[B ,,"1"1,;))) "/
Applying Grönwall's inequality on the function f(t) s
we obtain that supree[o.] (E[IIy? - x111) 1/2 = o for all E [0, which means
that Y! Y,? almost surely. And since |A/ Y?1 s f (Blll - 1211)x;))
(B V,bXX,, x|xxx) ds O, obtain that Y Y
F.5 PSEUDO-CODE OF ADJOINT MATCHING FOR FLOW MATCHING AND DDIM FINE-TUNING
Note that for each pair of equations 218 219 (20) 221). 222 223). the first equation corre-
sponds the updates in the DDPM paper, while the second equation is Euler-Maruyama / Euler discretization of the continuous-time object. To check that both discretizations are equal up to first order, remark that
1
ay 1-0 2O 1 &%
4% 1- + O((ā++1 2)
(224)
C
20
G ADAPTING DIFFUSION FINE-TUNING BASELINES TO FLOW MATCHING
G.1 ADAPTING REFL ET AL. 2023) TO FLOW MATCHING
Reward Feedback Learning (ReFL) is a diffusion fine-tuning algorithm introduced by Xu et al. (2023) which tries to increase the reward on denoised samples. Namely, X =(X))c=jo.) is the solution of the DDPM SDE (30), we denoise X,
xx,),) X X
(225)
This equation follows from the stochastic interpolant equation (2 if we replace Xo with the noise predictor e(),,t). And then, the ReFL optimization update is based the gradient:
(x)()((X,)) = Ver( X--vI-areaXX..
(226)
vāc
where the trajectories have been detached.
50

---

Published as a conference paper at ICLR 2025
Algorithm 1 Adjoint Matching for fine-tuning Flow Matching models
Input:
Pre-trained FM velocity field base step size h, number of fine-tuning iterations N. Initialize fine-tuned vector fields: y flnetuuee = = base with parameters 0. for n (0, .O..... N 1} do
Sample m trajectories X = (X..))( with memoryless noise schedule o()) = 2 (at 3r - Br). e.g.
X14h Xx + (20g Aincuuee (X,, (X,,,
de X1) (io(e))e E1 N(0,1)
No ~o0,
as
(215)
For each trajectory, solve the lean udjoint ODE 26
backwards time from to egg:
ū ã. + hā? V Note that X, and ā. should computed without gradients, i.e., X stopgrad(X1), ā. = stopgrad(ā.).
(2 base (X,, () X.
ã1 vxir(Xi).
(216)
For each trajectory, compute the Adjoint Matching objective
Cadj-Match (0) = Eie (o..........AA h)
[(8) (ng Suetune (X,, t) e base (X,, f)) + att)a_11
(217)
Compute the gradient V. (()) and update using favorite gradient descent algorithm. end
Output: Fine-tuned vector field u Anetune
To adapt ReFL to Flow Matching, we need express the denoiser map in terms of the vector field e. We have that
u(xx t) = Eli,No + ai X1 3: + xxX1= x
= E (8:X0 x,) + (de - -a) 3: + ax = z]
(227)
- + (ar a.) 1 (x.
where we defined the denoiser map 81 (z, E [x BeXy axx1 = x) Hence,
Xi (x, ((x,t)- de
(228)
Lea
G.2 ADAPTING DIFFUSION-DPO (WALLACE ET AL.. 2023A TO FLOW MATCHING
The Diffusion- DPO loss assumes access to ranked pairs of generated samples æ x/. where x" and x' are the winning and losing samples. For DDPM, the loss implemented in practice reads Wallace et al. 2023a Eq. 46):
(fro(
-E
log 5( _ (1s5 - 6 kh)ll - 11e0 - ever (2 kh)ll
111l"-eso - -(1ss (o co kb)ll2 -1es Cref (x kh) [?)))]
(229)
where S(x) T denotes the sigmoid function, and g( (xt) is the conditional distribution of the forward process, i.c. zinn is sampled z 2 + TAkE. (0, I). Following the derivation of the Diffusion-DPO loss in (Wallace et al. 2023a Sec. S4), we observe that the term 3
2 |eee cx(x%) k2) arises from
6 (mt)) 2
(230)
7uh
up to a constant term in 0. If we switch to the more general flow matching scheme, the analog of this term is
Li (x
xx 11.
(231)
51

---

Published as a conference paper at ICLR 2025
Algorithm 2 Adjoint Matching for fine-tuning DDIM
Input:
Pre-trained denoiser E base number of fine-tuning iterations N.
Initialize fine-tuned denoiser: Anetume E for n (0,... :0,,...4) .N- do Sample trajectories X (X.))E(0.....)) according DDPM, e-g.:
base with parameters 0.
X441 = V iitt (Xe 1- --/// Euetune k), + V/
1-
(1
)ek
(218)
VI-8a
1-62
nx+1
X4+1 Xk +
Rk+1-a4
25
Xx
aLVI-ag
6x41 Bg finetune (Xu, k) V
Ak
Ek!
(219)
where En (0, 1). Xu ~ .V(0, 1).
For each trajectory, solve the lean adjoint ODE 26 backwards in time from k: = K to e.g.:
.x+t + a Vx.
ax+1
(Xx
1-ap /oe// base (Xx; Xx)
äx = Vxxx((X)
v1-og
(220)
ux+1 + alew Vx.
1x+11-8x 201
Xx
0xe1 -gg hase (Xe, āk -xxx(xx)
SLVImax"
(221)
Note that Xx and should be computed without gradients, ie.,, Xx stopgrad(Xe).. stopgrad(ás).
For each trajectory, compute the Adjoint Matching objective
25
Ladi- (0) = (O.......KKI)
V
1
Bontume (Xn,t
base (Xn, k))
ux+1)
0k+0 1-
(1
Txx
(222)
= or LAdj-Match = (o.... K-1)
G611-B4 xx(1-az))
Annsune (X k) hase (Xe,
6611 1.111-40 Bgg 0s &
(223)
Compute the gradient Vof and update using favorite gradient descent algorithm.
end
Output: Fine-tuned vector field y Anetune
Using the expression of the denoiser map in terms of the vector field in equation (228) we can rewrite (231) as:
B
Ma .ch)-
Pan Ven
2
Pad oW TEA akh Pan - Bun ashh
Ork 2 Bahh
2
(232)
Fet Feh
Thus, the Diffusion- DPO loss for Flow Matching reads
(D) = -E (ay
L (o Ea 1x )ar we (uf)
log -
(1)
no mo(o kh)
Buh
Mees Meee (1r kch) Noh) Heh
Bun 1
Dach x1 Bka
aihh Ben Ben
Ben
dihh 3ah - Bah
muh co(xt,, kh) Soa
GRR Ban
11
Beed Rred .bh)
PEa whe
Pen
auhh 2
[13)
dett Ban
Geh BBA Ban
Dach
(233)
Wallace et al. 2023al Sec. 5.1) claim that 8 E |2000, 5000] yields good performance on Stable Diffusion 1.5 and Stable Diffusion XL-1.0, which if we translate our notation corresponds to 3 [4000, 10000].
When we have access to the reward function r, instead a winning sample x;. and a los- ing sample 21. we have pair of samples (1) with winning weights S(r(2f) r(x))) Hence, the loss 233) be-
1texp (+=)))-(())) SS (r-(r((f)- - r(x)))) 1texp (((()))-r((p)))
52

---

Published as a conference paper at ICLR 2025
comes:
[0) -efr..t))DDD (Dit))D,,.. (0 k (of (of) let (wg)
E 1911?
log Su-
a3
uo ( )
Pkb
Vrat (aa
Beh
Bun
zan
Aeke Ben Bun
Beh
Made Bk.. Ban
Den
(1)
my ( .hh)
Ath ah
Par
Akh Moor Mour (1 .kh) - Pen she
aun a
(11)))) (234)
-
Lidk Bah
Bach
Den
Akh PRA Bun
Bac
Bigh
We want emphasize that despite the similarities, even though the loss LDPO that we use (equation (344) is very similar to the one implemented by Walllace et al.] (2023a), the preference data pairs that we use are very different from theirs. We sample the preference data from the current model, which results in imperfect samples, while they consider off-policy high-quality, curated preference samples. The reason for this discrepancy that the starting point of our work is reward model, not set of preference data, and we only benchmark against approaches that leverage reward models for an apples-to-apples to-apples comparison. Our experimental results on DPO (Tab. 2 Fig. 5 Tab, 3 show that the resulting model performs like the base model, bit worse according some metrics. Hence, we conclude that DPO not a competitive alternative for on- on-policy fine-tune when the base model not already good.
H EXPERIMENTAL DETAILS
Unless otherwise specified, we used the same hyperparameters across all fine-tuning methods. Namely, we used:
K 40 timesteps.
Adam optimizer with learning rate 2 x 10-5 and parameters 31 = 0.95, B2 0.999, 1 x 10-5, weight decay 10-2, gradient norm clipping value 1. For Discrete Adjoint, these hyperparameters resulted in fine-tuning instability (see Tab. 6 the results that we report in all other tables for Discrete Adjoint were obtained with learning rate
1 x 10-5
1 precision.
Effective batch size 40; for each run we used two 80GB A100 GPUs with batch size 20 each.
A set of 40k fine-tuning prompts taken from a licensed dataset consisting of text and image pairs (note that we disregarded the images). Thus, each epoch lasts 1000 iterations; see the total amount of fine-tuning iterations for each algorithm in Tab. For each of the three runs that we perform for each data point that we report, the set of 40k prompts is sampled independently among total set of 100k prompts.
H.1 NOISE SCHEDULE DETAILS
Since we use K = 40 discretization steps, the timesteps are t {0, 0.025, 0.05, 0.075, 0.1, 0.95, 0.975). To sample Xeth from X, we use equation (215)
We the choices a4 8c 1 which means that (()) 2/ dh B, (.)
V 2(1 ((=
2(1-()
Note that if we plug ( = 0 into this expression, obtain infinity, and if we plug L E 1, we obtain (t) s O. For obvious reasons, the former issue requires a fix: we simply add small offset to the denominator of (t), replacing V1/t by V1/(t= + h) (note that == 1/K = 0.025). But the latter issue is also not completely satisfactory from a practical standpoint, because looking the adjoint matching loss 25), we observe that u(X;, t) is trained to approximate the conditional expectation of (t) att X). Thus, if we set (t) very close to zero for S 1, we are forcing the control to be close to zero as well, equivalently preventing finetuee from deviating from base While
53

---

Published as a conference paper at ICLR 2025
this is the right thing to do from theoretical perspective, concluded experimentally that setting (t) just slightly larger results in substantially faster fine-tuning thanks to the additional leeway provided to y finetune to deviate from base In particular, we added small offset to the factor t in the numerator - of (1): we replaced & by 1 - + h. Thus, the expression that we used to compute the diffusion coefficient in experiments is
20-++4) 1+h
(235)
When solving the lean adjoint ODE 22 27 backwards in time via the Euler scheme (216). the timesteps use are E {1, 0.975, 0.95, 0.925, 0.9, ................................ 0.05, 0.025]. We do not actually initialize the adjoint state as (.))( but rather as (XI( where Xi = XI-A + hy (X1-A, h). That is, X₁ obtained by performing final noiseless update, instead of using noise (1-h) = / given by equation (235). The reason for this is that the regular final iterate X1 contains some noise that was added in the final step, and that can distort the gradient (X)). By setting ā(1 X) = V Xx) ), we get rid of this bias. Note that in the continuous time limit h - 0, X1 = X1, which means that this small trick is consistent.
H.2 SELECTION OF GRADIENT EVALUATION TIMESTEPS
In Alg. equation 217 we state that the term || 2 (0g fnnetune (Xeet) base (baso(XX,,)) (X,, + @t)
must be computed for all K steps in {O....,I ., - h). However, the gradient signal provided by backpropagating through this expression for consecutive times t and + h is quite similar. In the interest of computational efficiency, we sample a subset IC of timesteps, and we only compute and backpropagate the terms || ( (of Jufinetuee (X,, t) u t)) + (t)),l for those timesteps. We construct K by sampling ten timesteps uniformly without repetition among , 0.725). and always sampling the last ten timesteps {0.75, ..... 0.975). This is because fine-tuning the last ten steps (25% of the total) well is critical for good empirical performance, while the initial steps are not important.
H.3
LOSS FUNCTION CLIPPING: THE LCT HYPERPARAMETES
Note that the magnitude of att) a(t; X ū) is much larger for times t 2 0 than for times E 1. The reason two-fold:
As discussed in App. H.I (t) is much larger for t 2 0 than for t $ 1.
The magnitude of the lean adjoint state à grows roughly exponentially as 1 goes backward in time. In fact, if we assumed that Vab(X,, t) is constant time, this statement would be exact.
Observe that when a(t) X uu) is large, the gradient Voll 2 ((3) (ug
(ffincuue (X,, t)
buse (X,,t)) + (f)āe]]² also has a high magnitude. Including such terms in our gradient com- putation decreases the signal to noise ratio of the gradient. Even more so, as discussed in App. H.2 for good practical performance it is critical to get a good gradient signal from the last 25% steps. Hence, including the high-magnitude terms for t S 0 in our gradients muffle these other impor- tant, low-magnitude terms.
To fix this issue, we clip the terms such that || o (ufinetume (X,,t - u base b t)) + (t)ā,|l > LCT, where LCT stands for the loss clipping threshold. That is, the adjoint matching loss that we use in our experiments of the form:
(basse(e,,)) Lxdj- (0) = Eiek min {LCT, l| ( (uguetume (Xe,)) - basse (X., 0) + (1)ae_1) (236)
where IC is the random timestep subset described in App.
For adjoint matching, we set LCT = 1.6 x X Remark that LCT needs to grow quadratically with X, because the magnitude of the lean adjoint à grows quadratically with X. We set the constant 1.6 through experimentation; all or almost all of the terms for the last ten timesteps fall below LCT, but only a fraction of the terms (*s 25%) for the first ten steps fall below LCT. The constant for LCT is relevant hyperparameter that needs to be tuned to obtain similar behavior.
54

---

Published as a conference paper at ICLR 2025
We also used loss function clipping on the continuous adjoint loss. For that loss set LCT 1600 x X². The reason is that the magnitude of the regular adjoint states is significantly larger than the magnitude of the lean adjoint states (which is a big reason why adjoint matching outperforms the continuous adjoint).
H.4 COMPUTATION OF EVALUATION METRICS
We used the open_ clip library Mlaarco et al. 2021 to compute ClipScores. We computed Clip- Score diversity as the variance of Clip embeddings of 40 generations for a given prompt, averaged across 25 prompts. Namely,
2 ClipScore Diversity = 40 k=1 25-21
1 E
Eisioses [CCip(([)) Clip(s))1?
(237)
where of denotes the i-th generation for the k-th prompt.
We used rans library to compute the PickScore processor and model Kirstain et al.,
2023). PickScore diversity is computed in analogy with ClipScore diversity.
We used the hps library to compute values of Human Preference Score v2 (Wu et al. 2023b
To compute Dreamsim diversity we use the theddremmin library (Fu et al. 2023
Dreamsim diversity
computed in analogy with ClipScore diversity.
H.S REMARKS ON COMPUTATIONAL COSTS
Observe from the figures reported in Tab, that the per iteration wall-clock time of Adjoint Matching (156 seconds) is very similar that of the Discrete Adjoint loss (152 seconds). We report hardware and hyperparameter details at the beginning of App. H The reason for which both algorithms take a similar time is that they perform a similar amount of forward and backward passes on the flow matching model and the reward model. Namely, for each sample in the batch, both algorithms perform forward passes on the flow model to obtain the trajectories. In order to compute the gradient of the training loss, the Discrete Adjoint loss does K additional forward passes to evaluate the base flow model, one forward and backward pass on the reward model, and K backward passes avoid memory overflow.
on the current flow model, which typically use gradient checkpointing In the case of Adjoint Matching, solving the lean adjoint ODE requires one forward and backward pass on the reward model, and K backward passes on the base flow model. Finally, computing the gradient of the loss takes K/2 additional backward passes if we evaluate at only half of the timesteps as we do, although this computation is much quicker because it can be fully parallelized.
Meanwhile, computing the gradient the Continuous Adjoint loss takes 204 seconds. With respect to Adjoint Matching, Continuous Adjoint performs additional backward passes to compute the gra- dients V x. (X,, t) when solving the adjoint ODE. Finally, we observe that models that directly fine-tune the reward are quicker, but that comes with its own set of issues that discuss throughout the paper.
H.6 REMARKS ON NUMBER OF SAMPLING TIMESTEPS
In our experiments and all baselines, used 40 timesteps in the fine-tuning procedure (h 1/40 in Alg. T. The experiments reported in all tables and figures except for Tab. g were performed at 40 inference timesteps. In Tab. g (App. A), show experimental results TO, 20, 40, 100, and 200 inference timesteps, for the base model and the models fine-tuned with adjoint matching and DRaFT-1. We make the following observations about the results:
The metrics for Adjoint Matching at 100 and 200 timesteps are statistically equal to the ones for 40 timesteps, with slight increases Dreamsim diversity. This suggests that fine- tuning at large numbers of timesteps is a good idea if we want to perform inference at a large number of timesteps, as otherwise the capabilities of the model are limited by the number of fine-tuning timesteps instead of the inference compute. Also, at 100 and 200 timesteps the difference in performance of Adjoint Matching relative to DRaFT-1 increases.
The metries for Adjoint Matching at 10 and 20 timesteps are worse than at 40 timesteps, especially for 10. The difference in performance between Adjoint Matching and DRaFT-1
55

---

Published as a conference paper at ICLR 2025
vanishes at 10 timesteps for all metrics except for diversity, for which Adjoint Matching is still clearly better.
56


Reviews:
[{"review_id": "2CC4rdfNQY", "reviewer": "Reviewer_4cB8", "summary": "This paper provides theorectical insights on why optimizing a KL-regularized reward objective (which is popular and dominant in RLHF for LLM) could lead to a bias in the optimal solution for diffusion models, and how to address this issue by a proper choice of noise schedule. The paper also provides solutions on how to solve the stochastic optimal control problem by adjoint methods, and provide a more efficient alternative than the classical methods and prove its equivalence. Empirical examples are further provided to show the algorithm effectiveness.", "questions": "What is the complexity/computation cost of computing/sampling adjoints in Equation (26) and (27)? \n\nIn addition, for Diffusion-DPO, how the preference pairs are sampled for further fine-tuning? Can authors explain why Diffusion-DPO tuned models yield a decrease in performance?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "This paper is theorectically well written, with detailed introduction to the literature, explanation of motivations, several interesting theorems and propositions, and provide theorectically-driven approaches for diffusion models fine-tuning. Diffusion Models fine-tuning or RLHF for diffusion models are an important direction which already contributes to improving the performance to SOTA diffusion models. This paper indeed provides a novel model-based SOC method for diffusion models alignment, which is theorectically sound and also yields good performance for tuning Flow Matching based models.", "weaknesses": "The main weakness is that the paper can be benefited from more comparisons with baseline methods empirically, specifically, there lacks a baseline in experiments in directly optimizing objective (17) (though theorectically there is value bias as shown in the paper), using stochastic control methods like adjoint matching proposed in this paper. There is comparison for this in the synthetic examples in Figure 2, but more practical downstream tasks evaluations are needed to show that the noise schedule proposed in this paper can indeed yield better performance. \n\n**Minors**:\n\n1）On line 229-230, the expression \"Dividing (14) by (15)\" is odd as (14) is not an equality, \"Plug in normalization constant (15) to (14)\" might be better.\n\n2）The adjoint method in Section 4 needs more introduction or discussion, or an earlier pointer to the part in appendix, including more clarification on what is adjoint on line 350-352, what is this loss on Equation (21); The reference in Appendix should be referred earlier instead of being deferred to until line 422.\n\n3）In proposition 2, it is better to first define earlier of the adjoint matching objective instead of combining the definition with the proposition.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DVsq14q33C", "reviewer": "Reviewer_7e67", "summary": "This paper studies reward-based fine-tuning for diffusion models. The authors frame the reward fine-tuning problem as stochastic optimal control, and point out an “initial value function bias” problem that exists in previous RLHF fine-tuning approaches. The authors propose using a memoryless noise schedule for fine-tuning in order to turn the learned distribution into the desired reward-tilted distribution without bias. Furthermore, a novel algorithm named adjoint matching is proposed to solve the stochastic optimal control problem. Experimental results show that fine-tuning a flow matching base model with adjoint matching outperforms baselines such as DRaFT, ReFL, and DPO.", "questions": "see above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper provides a theoretically sound framework for the reward fine-tuning problem, viewing it as a stochastic optimal control problem. The observation of the value function bias problem in previous approaches and proposal of using “memoryless” noise schedules are based on this view.\n\nThe proposed Adjoint Matching algorithm for SOC, casting it as a least-squares regression problem, is novel and effective.\n\nThis paper is well-written, clearly structured and easy to follow.", "weaknesses": "The main paper presented experimental results on fine-tuning a Flow Matching model, and provided pseudo-code for fine-tuning denoising diffusion models; it would be more convincing if results on denoising diffusion models are provided.\n\nThe experiments with classifier-free guidance do not seem comprehensive. It would be better if there are similar quantitative comparisons with other baselines than selected DRaFT-1.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LaAeVUumzM", "reviewer": "Reviewer_MA7f", "summary": "This paper aims to fine-tune a pre-trained diffusion model. The contribution of this paper is two-fold:\n\nFirst, the paper explains the problem of overoptimization when fine-tuning diffusion models based on a learned reward function. To avoid the initial value function bias problem, the paper proposes using a memoryless noise schedule, which is defined as a noise schedule that ensures that the joint distribution over the initial and final states is independent.\n\nSecond, the paper introduces (lean) adjoint matching for stochastic optimal control methods which resolves memory constraints of the discrete adjoint method and makes the continuous adjoint method more scalable by making the simulation of the adjoint ODE cheaper to compute.", "questions": "- Did the authors compare the proposed approach to the work by Uehara et. al. [1]?\n- Did the authors test how adjoint matching works with only a few diffusion steps?\n- How did the authors select the time points to discretize? Uniform in the interval [0,1]?\n- Do the authors have an intuition as to why the lean adjoint matching objective outperforms the discrete/continuous adjoint matching objective? I read the explanation that the continuous adjoint needs clipping and is unstable, but are there other reasons?\n- Are the authors planning to make the code public?\n\n---\n\n- [1] Uehara et al. (2024). Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control. *arXiv preprint arXiv:2402.15194*\n- [2] Vargas, F., Grathwohl, W., & Doucet, A. (2023). Denoising diffusion samplers.  *ICLR 2023*. 2024.\n- [3] Zhang, Q., & Chen, Y. (2021). Path integral sampler: a stochastic control approach for sampling. *arXiv preprint arXiv:2111.15141*\n- [4] Nusken, Nikolas, et al. \"Transport meets variational inference: Controlled Monte Carlo diffusions.\" *The Twelfth International Conference on Learning Representations: ICLR 2024*. 2024.\n- [5] Richter, L., & Berner, J. (2023). Improved sampling via learned diffusions. *ICLR 2024*. 2024.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and structured.\n- Strengths of memory-less noise schedule:\n    - The proposed approach is a less complex and provides an arguably more elegant solution to the initial value function bias compared to the work by Uehara et. al. [1].\n    - The proposed approach is compatible with both DDPM and Flow/Bridge-Matching models.\n- Strengths of (lean) adjoint matching:\n    - Simple regression-based objective with circumvents memory problems associated with the discrete adjoint method\n    - Simulating the adjoint ODE does not require control evaluations, making it more scalable than the continuous adjoint method.\n    - Compatible with general SOC problems.\n- Numerical evaluations show that the proposed approach outperforms competing methods on a variety of evaluation criteria.", "weaknesses": "- Weaknesses of the memory-less noise schedule:\n    - The statement “however, there does not yet exist a simple approach which actually provably generates from the tilted distribution” is, to the best of my knowledge not true: Take for example the approach in [2, 3] and use the base measure as reference process. Then, we have the terminal cost\n    \n    $$\n    g(X_1) =p^{\\text{base}}(X_1)/p_{\\text{target}}(X_1) = p^{\\text{base}}(X_1)/(p^{\\text{base}}(X_1)\\exp(r(X_1)) = 1/\\exp(r(X_1)) \n    $$\n    \n    which es ensures that \n    \n    $$\n    p_{\\text{target}}(X_1) = \\exp(r(X_1)).\n    $$\n    \n    The ‘memoryless’ property thus eludes to the fact, that we need a time-reversal of the base process. A more general discussion compared to [2, 3] can be found in [4, 5].\n    \n- Weaknesses of (lean) adjoint matching:\n    - It is not clear (at least to me) which objective function lean adjoint matching is optimizing. \n    - Lean/Continous Adjoint-matching requires simulating another differential equation which may result in computational overhead.\n- Code is not public", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wbZKbzWUlm", "reviewer": "Reviewer_pgS7", "summary": "The authors propose a SOC formulation for fine-tuning flow-based generative models, demonstrating that naive approaches can introduce inherent bias. They introduce a memoryless noise schedule to ensure convergence to the tilted distribution and present the Adjoint Matching objective as a scalable training objective for SOC problems. Their comparisons show improved generalization, consistency, and diversity in the fine-tuned generative models.", "questions": "See weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The concept of using stochastic optimal control (SOC) to fine-tune diffusion-based generative models is not new, as [1] formulates fine-tuning for diffusion models (specifically forward-reverse models) through Doob's h-transform, which involves optimally controlled diffusion processes. The novelty lies in:\n\n* Reformulating fine-tuning of flow-matching based generative models as an SOC problem, the authors introduce a suitable cost functional for the control policy parameterized by a training neural network. Additionally, they incorporate a memoryless noise schedule to ensure factorizability, thereby eliminating inherent bias.\n* Proposing an adjoint-matching objective to solve the above SOC problem. It is well known that solving SOC problems is computationally challenging due to the need for continuous gradient graph caching. Extending the adjoint methods commonly used in dynamical learning as a dynamic solver for SOC objectives is, in my view, a significant contribution.\n\n\n```\n[1] Denker et al., Efficient Finetuning of Conditional Diffusion Models by Learning the Generalised h-transform.\n```", "weaknesses": "* While I may have missed it in the appendix, it appears that the experiments in this paper primarily address the quality of fine-tuning, without providing quantitative results on aspects of the proposed adjoint objective, such as convergence plots or memory usage statistics comparing the adjoint matching loss to traditional SOC objectives, as shown in studies like [2, 3]. From a theoretical standpoint, although the author suggests that the approach could be sufficiently applied, it would strengthen the claim to include experiments validating the numerical effectiveness of the new SOC objective.\n\n\n```\n[2] Nüsken and Ritcher, Solving high-dimensional Hamilton-Jacobi-Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space.\n[3] Domingo-Enrich et al., Stochastic Optimal Control Matching.\n```", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)