Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control

## Abstract

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

## 1 Introduction

Figure: Base model (Flow Matching) w/ Guidance
Refer to caption: /html/2409.08861/assets/figs/samples/cfg_prompt_49_image_1.jpg

Flow Matching (Lipman et al., ; Albergo and Vanden-Eijnden, ; Liu et al., ) and denoising diffusion
(Song and Ermon, ; Ho et al., ; Song et al., ; Kingma et al., )
models are being used for many generative modeling applications, including text-to-image (Rombach et al., ; Esser et al., ), text-to-video (Singer et al., ), and text-to-audio (Le et al., ; Vyas et al., ).
In most cases, the base generative model does not achieve the desired sample quality.
To improve the generated samples, it is common to resort to techniques such as classifier-free guidance (Ho and Salimans, ; Zheng et al., ) to get better text-to-sample alignment, or to fine-tune using human preference reward models to improve sample quality and realism (Wallace et al., ; Clark et al., ).

In the adjacent field of large language models, the behavior of the model is aligned to human preferences through fine-tuning with reinforcement learning from human feedback (RLHF).
Either explicitly or implicitly, RLHF methods (Ziegler et al., ; Stiennon et al., ; Ouyang et al., ; Bai et al., ) assume a reward model $r(x)$ that captures human preferences, with the goal of modifying the base generative model such that it generates the following *tilted distribution*:

$$ $\textstyle p^{*}(x)\propto p^{\mathrm{base}}(x)\exp(r(x)),$ (1) $$

where $p_{\mathrm{base}}$ is the base generative model’s sample distribution.

Inspired by this,
fine-tuning methods have been developed to improve denoising diffusion models based on human preference data; either using a reward-based approach (Fan and Lee, ; Black et al., ; Fan et al., ; Xu et al., ; Clark et al., ; Uehara et al., , ),
or direct preference optimization (Wallace et al., ).
However, unlike the fine-tuning methods designed for large language models,
most of the existing methods to a large degree ignore $p^{\text{base}}$ and focus solely on the reward model.
Reward models can range from standard evaluation metrics such as ClipScore (Hessel et al., ; Kirstain et al., ) to specialized models that have been trained on human preferences (Schuhmann and Beaumont, ; Xu et al., ; Wu et al., ). As these are parameterized by neural networks, they fall pray to adversarial examples which lead to the generation of undesirable artifacts (Goodfellow et al., ; Mordvintsev et al., ). This has led some works to consider adding regularization during fine-tuning (Fan et al., ; Uehara et al., ) to incentivize staying close to the base model distribution; however, there does not yet exist a *simple* approach which actually provably generates from the tilted distribution ([1](#S1.E1)).

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 ([1](#S1.E1)).
This problem has also been observed by Uehara et al. () but they propose a more complicated solution which involves training a separate generative model for the optimal noise distribution.
- (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 generated 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 a new method for solving SOC problems, called *Adjoint Matching*, which combines the scalability of gradient-based methods and the simplicity of a least-squares regression objective. This is orthogonal to the reward fine-tuning application and can be applied to general SOC problems.
- (iv)
We perform extensive comparisons to baseline approaches, and analyze them from multiple perspectives such as realism, consistency, and diversity. We find that our proposed method provides generalization to unseen human preference reward models, better text-to-sample consistency, and retains good diversity.

In the following, sections are broken down as follows: [Section 2](#S2) summarizes the algorithms used for sampling from pre-trained Flow Matching and diffusion models, while [Section 3](#S3) provides a common notation that we will use throughout. [Sections 4](#S4) and [5](#S5) form the core of our contributions. [Section 4](#S4) details the value function bias problem and our proposed solution via the memoryless noise schedule. [Section 5](#S5) details the new Adjoint Matching algorithm for solving SOC problems.

## 2 Preliminaries on dynamical generative models

We are interested in fine-tuning base generative models $p^{\mathrm{base}}(X_{1})$ where samples are generated through the simulation of a stochastic process. That is, these models transform noise variables into a sample through an iterative process. In particular, we discuss the specific constructions and sampling processes of Flow Matching (Lipman et al., ; Liu et al., ; Liu, ; Albergo and Vanden-Eijnden, ) and Denoising Diffusion Models (Ho et al., ; Song et al., , ).
The goal of this section is to provide background information on these methods, which we will later unify into a single consistent notation in [Section 3](#S3).

Given random variables from an initial distribution $\bar{X}_{0}\sim p_{0}=\mathcal{N}(0,I)$, and $\bar{X}_{1}$ which are distributed according to some data distribution, we define the reference flow $\bm{\bar{X}}=(\bar{X}_{t})_{t\in[0,1]}$ where

$$ $\textstyle\bar{X}_{t}=\beta_{t}\bar{X}_{0}+\alpha_{t}\bar{X}_{1},$ (2) $$

where $(\alpha_{t})_{t\in[0,1]},(\beta_{t})_{t\in[0,1]}$ are functions such that $\alpha_{0}=\beta_{1}=0$ and $\alpha_{1}=\beta_{0}=1$.
Diffusion models and Flow Matching construct generative Markov processes $X_{t}$ with initial distribution $X_{0}\sim\mathcal{N}(0,I)$ that result in flows $\bm{X}=(X_{t})_{t\in[0,1]}$ with the same time marginals as the reference flow $\bm{\bar{X}}$, i.e., the random variables $X_{t}$ and $\bar{X}_{t}$ have identical distribution for all times $t\in[0,1]$.
This implies $X_{1}$ has the same distribution as the data distribution, so simulating the Markov process from random noise $X_{0}$ is a way to generate artificial samples(^1^11In our derivations, we will simply assume the base model has been trained perfectly during the pre-training phase.).

### 2.1 Flow Matching

In its simplest form, the generative Markov process of a Flow Matching model is an ordinary differential equation (ODE) of the form:

$$ $\textstyle\mathrm{d}X_{t}=v(X_{t},t)\,\mathrm{d}t,\qquad X_{0}\sim\mathcal{N}(0,I).$ (3) $$

where $v(X_{t},t)$ is a parametric velocity that is optimized to match the derivative of the reference flow, i.e.,
$v(X_{t},t)=\operatorname*{argmin}_{\hat{v}}\mathbb{E}\big{\|}\hat{v}(\bar{X}_{t},t)-\frac{\mathrm{d}}{\mathrm{d}t}\bar{X}_{t}\big{\|}^{2}$ (see e.g. Lipman et al. () for details on pre-training Flow Matching models).
It can then be proven that the solution of the generative process ([3](#S2.E3)) has the same time marginals as the reference flow (Lipman et al., ; Liu, ; Albergo and Vanden-Eijnden, ), and a commonly used choice is $\alpha_{t}=t$ and $\beta_{t}=1-t$.
One can also consider a family of stochastic differential equations (SDEs) with an arbitrary state-independent diffusion coefficient(^2^22We use the common short-hand “over-dot” notation to denote the time derivative, i.e., $\dot{x}_{t}=\frac{\mathrm{d}}{\mathrm{d}t}x_{t}$.):

$$ $\textstyle\mathrm{d}X_{t}=\left(v(X_{t},t)+\frac{\sigma(t)^{2}}{2\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}\left(v(X_{t},t)-\frac{\dot{\alpha}_{t}}{\alpha_{t}}X_{t}\right)\right)\,\mathrm{d}t+\sigma(t)\,\mathrm{d}B_{t},\qquad X_{0}\sim\mathcal{N}(0,I),$ (4) $$

where $(B_{t})_{t\geq 0}$ is a Brownian motion.
The generative processes in ([3](#S2.E3)) and ([4](#S2.E4)) have the same time marginals. This can be seen by writing down the Fokker-Planck equations for ([3](#S2.E3)) and ([4](#S2.E4)), and observing that they are the same up to a cancellation of terms (Maoutsa et al., ).
The diffusion coefficient $\sigma(t)$ in ([4](#S2.E4)) is compensated by the second term in the drift which scales proportionally as $\sigma(t)^{2}$.

### 2.2 Denoising Diffusion Models

We next discuss diffusion models, in particular the sampling scheme proposed by Denoising Diffusion Implicit Model (DDIM; Song et al. ()) which we will later relate to Denoising Diffusion Probabilistic Models (DDPM; Ho et al. ()) as a particular case of the former.
For sampling from a diffusion model, the DDIM update rule(^3^33We slightly depart from the notation in Song et al. () by flipping the direction of time and using $\bar{\alpha}_{k}$ which corresponds to the $\alpha_{k}$ in Song et al. () while it corresponds to the $\bar{\alpha}_{k}$ in Ho et al. ().) (Song et al. (), Eq. 12), typically stated in discrete time with $k\in\{0,\dots,K\}$, is:

$$ $\textstyle X_{k+1}=\sqrt{\bar{\alpha}_{k+1}}\big{(}\frac{X_{k}-\sqrt{1-\bar{\alpha}_{k}}\epsilon(X_{k},k)}{\sqrt{\bar{\alpha}_{k}}}\big{)}+\sqrt{1-\bar{\alpha}_{k+1}-\sigma_{k}^{2}}\epsilon(X_{k},k)+\sigma_{k}\varepsilon_{k},\qquad\varepsilon_{k}\sim\mathcal{N}(0,I),\ X_{0}\sim\mathcal{N}(0,I),$ (5) $$

where $\bar{\alpha}_{k}$ is an increasing sequence such that $\bar{\alpha}_{0}=0$, $\bar{\alpha}_{K}=1$, and the sequence $\sigma_{k}$ is arbitrary. That is, one samples an initial Gaussian random variable $x_{0}$, and applies the stochastic update ([5](#S2.E5)) iteratively $K$ times in order to obtain an artificial sample $X_{K}$. Updates can be interpreted as progressively denoising the iterate: $x_{0}$ is completely noisy and $x_{K}$ is fully denoised. The noise predictor model $\epsilon(x_{k},k)$ is trained to predict the noise of $x_{k}$ (see e.g. Ho et al. () for details on pre-training denoising diffusion models).

## 3 Flow Matching and diffusion models from a common perspective

We formulate Flow Matching and diffusion models in a unified framework, which we will later use throughout the paper.
Firstly, to simplify notation, we will be using continuous-time formulations. This will also directly enable fine-tuning methods inspired by the continuous-time paradigm, which we find tends to perform better than discrete-time counterparts in our empirical validations. Secondly, by consolidating notation, we will be able to discuss fine-tuning of dynamical generative models that follow the same time marginals as the reference flow ([2](#S2.E2)), pre-trained with either the Denoising Diffusion or Flow Matching framework, in full generality.

To convert DDIM to a continuous-time stochastic process, we can show that the DDIM update rule ([5](#S2.E5)), up to a first-order approximation, is equivalent to the Euler-Maruyama discretization of the following SDE:

$$ $\textstyle\mathrm{d}X_{t}$ $\textstyle=\big{(}\frac{\dot{\bar{\alpha}}_{t}}{2\bar{\alpha}_{t}}X_{t}-\big{(}\frac{\dot{\bar{\alpha}}_{t}}{2\bar{\alpha}_{t}}+\frac{\sigma(t)^{2}}{2}\big{)}\frac{\epsilon^{\text{base}}(X_{t},t)}{\sqrt{1-\bar{\alpha}_{t}}}\big{)}\mathrm{d}t+\sigma(t)\mathrm{d}B_{t},\qquad X_{0}\sim\mathcal{N}(0,I).$ (6) $$

See [Section 10.1](#S10.SS1) for the full derivation.
To go from ([5](#S2.E5)) to ([6](#S3.E6)), we assumed a uniform discretization of time, i.e. $t=\tfrac{k}{K}$.
This results in identifying the discrete-time process $(X_{k})_{k\in\{0,\dots,K\}}$ with a continuous-time process $(X_{t})_{t\in[0,1]}$, where $\bar{\alpha}_{k}:=\bar{\alpha}_{t}$, $\sigma_{k}:=\frac{1}{\sqrt{K}}\sigma(t)$, and $\epsilon(X_{k},k)$ with $\epsilon^{\text{base}}(X_{k},t)$.
In relation to the reference flow ([2](#S2.E2)), the generative process in ([6](#S3.E6)) has the same time marginals when $\alpha_{t}=\sqrt{\bar{\alpha}_{t}}$ and $\beta_{t}=\sqrt{1-\bar{\alpha}_{t}}$ (Ho et al., ).

Furthermore, when viewed up to first order approximations, the DDPM sampling scheme (Ho et al. (); Algorithm 2) can be seen as special instance of the DDIM sampling scheme when $\sigma(t)=\sqrt{\nicefrac{{\dot{\bar{\alpha}}_{t}}}{{\bar{\alpha}_{t}}}}$. This results in the following generative process:

$$ $\textstyle\mathrm{d}X_{t}$ $\textstyle=\big{(}\frac{\dot{\bar{\alpha}}_{t}}{2\bar{\alpha}_{t}}X_{t}-\frac{\dot{\bar{\alpha}}_{t}}{\bar{\alpha}_{t}}\frac{\epsilon^{\text{base}}(X_{t},t)}{\sqrt{1-\bar{\alpha}_{t}}}\big{)}\mathrm{d}t+\sqrt{\frac{\dot{\bar{\alpha}}_{t}}{\bar{\alpha}_{t}}}\mathrm{d}B_{t},\qquad X_{0}\sim\mathcal{N}(0,I),$ (7) $$

We can further consolidate notation by converting all quantities to the score function $\mathfrak{s}(x,t)$—defined as the gradient of the log density of the random variable $X_{t}$—which is possible when $X_{0}$ is Normal-distributed and under the affine reference flow ([2](#S2.E2)). In particular,
the velocity $v^{\text{base}}$ from Flow Matching can be expressed in terms of the score function:

$$ $\textstyle v^{\text{base}}(x,t)=\frac{\dot{\alpha}_{t}}{\alpha_{t}}x+\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})\mathfrak{s}(x,t).$ (8) $$

And the noise predictor $\epsilon^{\text{base}}$ also admits an expression in terms of the score function (see [Section 10.3](#S10.SS3)):

$$ $\textstyle\epsilon^{\text{base}}(x,t)=-\frac{\mathfrak{s}(x,t)}{\sqrt{1-\bar{\alpha}_{t}}}.$ (9) $$

Plugging these two equations into ([4](#S2.E4)) and ([6](#S3.E6)), respectively, and rewriting them in terms of only the $\alpha_{t}$ and $\beta_{t}$ in ([2](#S2.E2)), we can unify both the Flow Matching and continuous-time DDIM generative processes as:

$$ $\textstyle\mathrm{d}X_{t}$ $\textstyle=b(X_{t},t)\,\mathrm{d}t+\sigma(t)\,\mathrm{d}B_{t},\qquad X_{0}\sim\mathcal{N}(0,I),$ (10) $\textstyle\text{where}\ b(x,t)$ $\textstyle=\kappa_{t}x+\big{(}\frac{\sigma(t)^{2}}{2}+\eta_{t}\big{)}\mathfrak{s}(x,t),\quad\kappa_{t}=\frac{\dot{\alpha}_{t}}{\alpha_{t}},\quad\eta_{t}=\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})$ (11) $$

where $(\alpha_{t},\beta_{t})$ are coefficients of the reference flow ([2](#S2.E2)). We have hence expressed the generative process of a base model, whether it is a Flow Matching or a diffusion model, as an SDE of the form ([10](#S3.E10))-([11](#S3.E11)), unified by the choice of reference flow.

## 4 Fine-tuning as “memoryless” stochastic optimal control

We now discuss the crux of the problem: how to produce a fine-tuned generative model that produces samples $X_{1}$ which follow the tilted distribution involving a reward model ([1](#S1.E1)).
An obvious direction is to construct a *fine-tuning objective* involving both the base generative model and the reward model, where the optimal solution results in a fine-tuned generative model for the tilted distribution.
However, as we will explain, this turns out to be non-trivial, because a naïve formulation will introduce bias into the solution.

In [Section 4.1](#S4.SS1), we discuss the problem formulation of stochastic optimal control, a general framework for optimizing SDEs, and its relation to the maximum entropy reinforcement learning framework commonly used for RLHF fine-tuning.
Next, in [Section 4.2](#S4.SS2), we discuss the *initial value function bias* problem which plagues existing approaches and so far has seen no simple solution.
Finally, in [Section 4.3](#S4.SS3), we propose a novel simple solution that circumvents the bias problem, by enforcing a 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 ([1](#S1.E1)) without any statistical bias.

### 4.1 Preliminaries on the stochastic optimal control problem formulation

Stochastic optimal control (SOC; Bellman (); Fleming and Rishel (); Sethi ()) considers general optimization problems over stochastic differential equations, but we only need to consider a common instantiation, the control-affine problem formulation:

$\textstyle\min\limits_{u\in\mathcal{U}}\mathbb{E}\big{[}\int_{0}^{1}\big{(}\frac{1}{2}\|u(X^{u}_{t},t)\|^{2}+f(X^{u}_{t},t)\big{)}\,\mathrm{d}t+g(X^{u}_{1})\big{]},$
(12)
$\textstyle\begin{split}\textstyle\text{s.t.}\leavevmode\nobreak\ \mathrm{d}X^{u}_{t}=\left(b(X^{u}_{t},t)+\sigma(t)u(X^{u}_{t},t)\right)\,\mathrm{d}t+\sigma(t)\mathrm{d}B_{t},\qquad X^{u}_{0}\sim p_{0}\end{split}$
(13)

where in ([13](#S4.E13)), $X_{t}^{u}\in\mathbb{R}^{d}$ is the state of the stochastic process, $u:\mathbb{R}^{d}\times[0,1]\to\mathbb{R}^{d}$ is commonly referred to as the control vector field, $b:\mathbb{R}^{d}\times[0,1]\to\mathbb{R}^{d}$ is a base drift, and $\sigma:[0,1]\to\mathbb{R}^{d\times d}$ is the diffusion coefficient. These jointly define the *controlled process* $\bm{X}^{u}\sim p^{u}$ that we are interested in optimizing; often both $b$ and $\sigma$ are fixed and we only optimize over the control $u$.

As part of the objective functional ([12](#S4.E12)), we have an affine control cost $\frac{1}{2}\|u(X^{u}_{t},t)\|^{2}$, a running state cost $f:\mathbb{R}^{d}\times[0,1]\to\mathbb{R}$ and a terminal state cost $g:\mathbb{R}^{d}\to\mathbb{R}$.

The stochastic optimal control (SOC) objective ([12](#S4.E12)) can be decomposed recursively from the final time value. It is common to define the *cost functional* which is the expected future cost starting from state $x$ at time $t$:

$$ $\textstyle J(u;x,t):=\mathbb{E}_{\bm{X}\sim p^{u}}\left[\int_{t}^{1}\left(\frac{1}{2}\|u(X_{s},s)\|^{2}+f(X_{s},s)\right)\,\mathrm{d}s+g(X_{1})\;\big{|}\;X_{t}=x\right].$ (14) $$

From here, the *value function* is the optimal value of the cost functional(^4^44Note that there is a slight difference in terminology between SOC and reinforcement learning, where our cost functional is referred to as the state value function and our value function is the optimal state value function in RL.)
:

$$ $\textstyle V(x,t):=\min_{u\in\mathcal{U}}J(u;x,t)=J(u^{*};x,t),$ (15) $$

where $u^{*}$ is the *optimal control*, i.e., minimizer of ([12](#S4.E12)). Furthermore, a classical result is that the value function can be expressed in terms of the *uncontrolled* base process $p^{\text{base}}$ (Kappen (), see Domingo-Enrich et al. , Eq. 8, App. B for a self-contained proof):

$$ $\textstyle V(x,t)=-\log\mathbb{E}_{\bm{X}\sim p^{\text{base}}}\left[\exp(-\int_{t}^{1}f(X_{s},s)\mathrm{d}s-g(X_{1}))\;\big{|}\;X_{t}=x\right].$ (16) $$

A useful expression for the optimal control (which we will make use of in deriving the Adjoint Matching objective in [Section 5](#S5)) is that it is related to the gradient of the value function:

$$ $\textstyle u^{*}(x,t)=-\sigma(t)^{\top}\nabla_{x}V(x,t)=-\sigma(t)^{\top}\nabla_{x}J(u^{*},x,t).$ (17) $$

##### Relation to MaxEnt RL.

Stochastic optimal control with the control-affine formulation ([12](#S4.E12)) is the continuous-time equivalence of maximum entropy reinforcement learning (MaxEnt RL; Todorov (); Ziebart et al. ()) with a KL regularization instead of only an entropy regularization. In particular, by the Girsanov theorem ([Theorem 2](#Thmtheorem2)),
the affine control cost is equivalent to a Kullback–Leibler (KL) divergence between the base process $p^{\text{base}}$, when $u=0$, and the controlled process $p^{u}$, when conditioned on the same initial state $X_{0}$ (see [Section 11.4](#S11.SS4)):

$$ $\textstyle D_{\text{KL}}\left(p^{u}(\bm{X}|X_{0})\;\middle\|\;p^{base}(\bm{X}|X_{0})\right)=\mathbb{E}_{\bm{X}^{u}\sim p^{u}}\left[\int_{0}^{1}\frac{1}{2}\|u(X^{u}_{t},t)\|^{2}\mathrm{d}t\right],$ (18) $$

resulting in the KL-regularized RL interpretation of ([12](#S4.E12)):

$$ $\textstyle\max\limits_{u\in\mathcal{U}}\;\mathbb{E}_{X_{0}\sim p_{0}}\left[\mathbb{E}_{\bm{X}\sim p^{u}(\cdot|X_{0})}\big{[}\int_{0}^{1}-f(X_{t}^{u},t)\mathrm{d}t-g(X_{1}^{u})\big{]}-D_{\text{KL}}(p^{u}(\bm{X}|X_{0})\;\|\;p^{base}(\bm{X}|X_{0}))\right],$ (19) $$

where the negative state costs correspond to intermediate and terminal rewards in the RL interpretation. The KL divergence incentivizes the optimal solution to stay close to the distribution of the base process.

### 4.2 The initial value function bias problem

We next discuss why naïvely adding a KL regularization does not lead to the tilted distribution ([1](#S1.E1)).
From ([19](#S4.E19)), we can also show that the optimal distribution conditioned on $X_{0}$ is(^5^55Note ([20](#S4.E20)) is informal because densities over continuous-time processes are ill-defined; the formal statement is $\frac{\mathrm{d}\mathbb{P}^{*}}{\mathrm{d}\mathbb{P}^{\mathrm{base}}}(\bm{X}|X_{0})=\exp(-\int_{0}^{1}f(X_{t},t)\,\mathrm{d}t-g(X_{1}))$, where $\frac{\mathrm{d}\mathbb{P}^{*}}{\mathrm{d}\mathbb{P}^{\mathrm{base}}}$ denotes the Radon-Nikodym derivative. We treat this formally in the proofs.)

$$ $\textstyle p^{*}(\bm{X}|X_{0})\propto p^{\mathrm{base}}(\bm{X}|X_{0})\exp\big{(}-\int_{0}^{1}f(X_{t},t)\,\mathrm{d}t-g(X_{1})\big{)}.$ (20) $$

This is analogous to the exponentiated reward distribution in MaxEnt RL (Rawlik et al., ), but since we generalize the entropy regularization to a KL regularization, $p^{\text{base}}$ acts as a prior distribution.

Importantly, in order to relate this to the tilted distribution ([1](#S1.E1)) that we want to achieve for fine-tuning, we need to marginalize all time values and check the distribution of $p^{*}(X_{1})$. In order to do this, first notice that the normalization constant of ([20](#S4.E20)) is exactly the value function at $t=0$:

$$ $\textstyle\mathbb{E}_{\bm{X}\sim p^{\text{base}}(\bm{X}|X_{0})}\left[\exp\big{(}-\int_{0}^{1}f(X_{t},t)\,\mathrm{d}t-g(X_{1})\big{)}\right]=\exp\left(-V(X_{0},0)\right),$ (21) $$

where the equality is due to ([16](#S4.E16)). Therefore, we see that this normalization constant depends on $X_{0}$. Dividing ([20](#S4.E20)) by ([21](#S4.E21)) and multiplying by $p_{0}(X_{0})$, we obtain the normalized distribution over the full path $\bm{X}$,

$$ $\textstyle p^{*}(\bm{X})=p^{\mathrm{base}}(\bm{X})\exp\big{(}-\int_{0}^{1}f(X_{t},t)\,\mathrm{d}t-g(X_{1})+V(X_{0},0)\big{)}.$ (22) $$

Setting $f=0$ and $g=-r$, we arrive at an expression for the optimal distribution

$$ $\textstyle p^{*}(X_{0},X_{1})\propto p^{\mathrm{base}}(X_{0},X_{1})\exp\big{(}r(X_{1})+V(X_{0},0)\big{)}.$ (23) $$

This unfortunately does not lead to the tilted distribution ([1](#S1.E1)) because we have a bias in the optimal distribution that is due to the value function of the initial distribution $V(X_{0},0)$. That is to say, naïvely adding a KL regularization ([18](#S4.E18)) to the fine-tuning objective in the sense of ([19](#S4.E19)) leads to a biased distribution ([22](#S4.E22)) after fine-tuning and is not equivalent to the tilted distribution ([1](#S1.E1)). For instance, when the sampling procedure is noiseless, i.e., $\sigma(t)=0$, fine-tuning naïvely will not have any effect because $X_{0}$ completely determines $X_{1}$.

This is unlike the situation for large language models (Ouyang et al., ; Rafailov et al., ), where there is no dynamical process that samples $X_{1}$ iteratively and hence no dependence on the initial noise variable $X_{0}$.
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 ([19](#S4.E19)) has been explored in prior works (Fan et al., ), but its behavior was not well-understood and they did not relate the fine-tuned model to the tilted distribution ([1](#S1.E1)). Another direction that has been proposed is to learn the initial distribution $p_{0}$ to cancel out the bias (Uehara et al., ; Tang, ) but this simply shifts the work into tilting the initial distribution and requires an auxiliary model for parameterizing the optimal initial distribution.
In contrast, we show in the next section that it is possible to remove the value function bias by simply choosing a very particular noise schedule during the fine-tuning procedure.

### 4.3 The memoryless noise schedule for fine-tuning dynamical generative models

**Table 1: Diffusion coefficient $\sigma(t)$ and the factors $\kappa_{t}$, $\eta_{t}$ for the Flow Matching, Memoryless Flow Matching, DDIM, and DDPM generative processes. When the diffusion coefficient is $\sigma(t)=\sqrt{2\eta_{t}}$, the generative process is memoryless, i.e., samples $X_{1}$ will be independent of the initial noise $X_{0}$.**
|  | $\kappa_{t}$ | $\eta_{t}$ | Diffusion coefficient $\sigma(t)$ | Memoryless $X_{t}$ |
| --- | --- | --- | --- | --- |
| Flow Matching ([3](#S2.E3)) | $\frac{\dot{\alpha}_{t}}{\alpha_{t}}$ | $\beta_{t}\big{(}\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t}\big{)}$ | General (commonly 0 0) | No |
| Memoryless Flow Matching ([4](#S2.E4)) | $\frac{\dot{\alpha}_{t}}{\alpha_{t}}$ | $\beta_{t}\big{(}\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t}\big{)}$ | $\sqrt{2\eta_{t}}$ | Yes |
| DDIM ([6](#S3.E6)) | $\frac{\dot{\bar{\alpha}}_{t}}{2\bar{\alpha}_{t}}$ | $\frac{\dot{\bar{\alpha}}_{t}}{2\bar{\alpha}_{t}}$ | General (commonly 0 0) | No |
| DDPM ([7](#S3.E7)) | $\frac{\dot{\bar{\alpha}}_{t}}{2\bar{\alpha}_{t}}$ | $\frac{\dot{\bar{\alpha}}_{t}}{2\bar{\alpha}_{t}}$ | $\sqrt{2\eta_{t}}$ | Yes |

In this section, we propose a very simple method of turning ([23](#S4.E23)) into the tilted distribution ([1](#S1.E1)) through the use of a particular *memoryless* noise schedule.
Throughout, we provide an intuitive explanation of why this noise schedule is sufficient for fine-tuning while discussing the full theoretical result where we show that the memoryless noise schedule is actually not only sufficient but also necessary.

Intuitively, the main reason we cannot arrive at the tilted distribution from ([23](#S4.E23)) is due to the $p^{\text{base}}(X_{0},X_{1})$ distribution not factoring into $X_{0}$ and $X_{1}$. Hence, we define a memoryless generative process as follows:

###### Definition 1 (Memoryless generative process) .

A generative process of the form ([10](#S3.E10))-([11](#S3.E11)) is memoryless if $X_{0}$ and $X_{1}$ are independent, i.e., $p^{\text{base}}(X_{0},X_{1})=p^{\text{base}}(X_{0})p^{\text{base}}(X_{1})$.

When the base generative process is memoryless, this implies:

$$ $\textstyle p^{*}(X_{1})=\int p^{\text{base}}(X_{0})p^{\text{base}}(X_{1})\exp(r(X_{1})+V(X_{0},0))\mathrm{d}X_{0}\propto p^{\text{base}}(X_{1})\exp(r(X_{1})).$ (24) $$

That is, solving the SOC problem ([12](#S4.E12))-([13](#S4.E13)) with a memoryless base model will result in a fine-tuned model that generates samples $p^{*}(X_{1})$ according to the tilted distribution ([1](#S1.E1)).
This memoryless property is not satisfied generally by the family of generative processes captured by ([12](#S4.E12))-([13](#S4.E13)).
For instance, the Flow Matching and DDIM generative processes with zero diffusion coefficient (i.e., $\sigma(t)=0$) are definitely not memoryless due to $X_{0}$ and $X_{1}$ being theoretically invertible.
Below, we provide the sufficient and neccessary condition for the noise schedule in order to have a memoryless generative process.

###### Proposition 1 (Memoryless noise schedules) .

Within the family of generative processes ([10](#S3.E10))-([11](#S3.E11)), a generative process is memoryless if and only if the noise schedule is chosen as:

$$ $\textstyle\sigma(t)^{2}=2\eta_{t}+\chi(t),\text{ where }\chi:[0,1]\to\mathbb{R}\text{ is s.t. }\forall t\in(0,1],\quad\lim_{t^{\prime}\to 0^{+}}\alpha_{t^{\prime}}\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\chi(s)}{2\beta_{s}^{2}}\,\mathrm{d}s\big{)}=0.$ (25) $$

In particular, we refer to $\sigma(t)=\sqrt{2\eta_{t}}$ as the memoryless noise schedule.

Due to the endpoint constraints of $(\alpha_{t},\beta_{t})$ for the reference flow ([2](#S2.E2)), the memoryless noise schedule $\sigma(t)$ is infinite at $t=0$ and approaches zero at $t=1$. This provides a way for the generative process to mix when close to noise $X_{0}$ while stay steadying when close to the sample $X_{1}$. Hence, the sample will have no information about $X_{0}$ due to the enormous amount of mixing with a large diffusion coefficient.
Furthermore, while we have intuitively justified the memoryless noise schedule through its independence property, our theoretical result is actually even stronger: all generative models of the form ([10](#S3.E10))-([11](#S3.E11)) must be fine-tuned using the memoryless noise schedule.
We formalize this in the following theorem, which we prove in [Section 12.2](#S12.SS2):

###### Theorem 1 (Fine-tuning recipe for general noise schedule sampling) .

Within the family of generative processes ([10](#S3.E10))-([11](#S3.E11)), in order to allow the use of arbitrary noise schedules and still generate samples according to the tilted distribution ([1](#S1.E1)), the fine-tuning problem ([12](#S4.E12))-([13](#S4.E13)) with $f=0$ and $g=-r$ must be done with the memoryless noise schedule $\sigma(t)=\sqrt{2\eta_{t}}$.

[Theorem 1](#Thmtheorem1) states that we need to use the memoryless noise schedule for fine-tuning with the SOC objective—or equivalently, the KL regularized reward objective ([19](#S4.E19)).
This is the only noise schedule that retains the relationship between the velocity and score function, allowing the conversion to arbitrary noise schedules (e.g., $\sigma(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-time limit of the DDPM generative process ([7](#S3.E7)). However, the DDPM sampler (Ho et al., ) is not commonly used while the DDIM sampler (Song et al., ) and Flow Matching models typically generate samples using $\sigma(t)=0$, so an explicit conversion to the memoryless noise schedule is necessary for fine-tuning.
To the best of our knowledge, we are not aware of any existing works that have proposed a time-varying diffusion coefficient with theoretical guarantees.
[Table 1](#S4.T1) summarizes the memoryless schedule for diffusion and Flow Matching models, which we refer to as Memoryless Flow Matching. In [Figure 2](#S4.F2), we visualize fine-tuning a 1D model, where we see that constant $\sigma(t)$ leads to biased distributions whereas the memoryless noise schedule perfectly converges to the tilted distribution ([1](#S1.E1)).

Figure: (a) Pretrained FM $v^{\mathrm{base}}$
Refer to caption: /html/2409.08861/assets/figs/finetune1d/legend.png

For convenience, we plug the memoryless noise schedule into the controlled process for fine-tuning ([13](#S4.E13)), and express them in terms of each respective framework.
Let $\epsilon^{\mathrm{base}}$, $v^{\mathrm{base}}$ denote the pre-trained vector fields and $\epsilon^{\mathrm{finetune}}$, $v^{\mathrm{finetune}}$ the fine-tuned vector fields. Then we have the following expressions for the full drift $b(x,t)+\sigma(t)u(x,t)$ and control $u(x,t)$ when $\sigma(t)=\sqrt{2\eta_{t}}$:

DDIM / DDPM:

$\textstyle b(x,t)+\sigma(t)u(x,t)=\frac{\dot{\bar{\alpha}}_{t}}{2\bar{\alpha}_{t}}x-\frac{\dot{\bar{\alpha}}_{t}}{\bar{\alpha}_{t}}\frac{\epsilon^{\mathrm{finetune}}(x,t)}{\sqrt{1-\bar{\alpha}_{t}}},\quad\quad u(x,t)=-\sqrt{\frac{\dot{\bar{\alpha}}_{t}}{\bar{\alpha}_{t}(1-\bar{\alpha}_{t})}}(\epsilon^{\mathrm{finetune}}(x,t)-\epsilon^{\mathrm{base}}(x,t)).$
(26)

Memoryless Flow Matching:

$\textstyle b(x,t)+\sigma(t)u(x,t)=2v^{\mathrm{finetune}}(x,t)-\frac{\dot{\alpha}_{t}}{\alpha_{t}}x,\quad\quad u(x,t)=\sqrt{\frac{2}{\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}}(v^{\mathrm{finetune}}(x,t)-v^{\mathrm{base}}(x,t)).$
(27)

Thus, to solve the SOC problem ([12](#S4.E12))-([13](#S4.E13)) in practice, we parameterize the control $u$ in terms of $\epsilon^{\mathrm{finetune}}$ or $v^{\mathrm{finetune}}$ and optimize these vector fields instead. After plugging in ([26](#S4.E26))-([27](#S4.E27)), the SOC problem ([12](#S4.E12))-([13](#S4.E13)) can then be solved using any SOC algorithm in order to perform fine-tuning, and we proposed an especially effective algorithm next in [Section 5](#S5). After fine-tuning, $\epsilon^{\mathrm{finetune}}$ and $v^{\mathrm{finetune}}$ can simply be plugged back into their respective generative processes ([3](#S2.E3))-([7](#S3.E7)) to sample from the tilted distribution ([1](#S1.E1)) using any choice of diffusion coefficient.

###### Definition 1 (Memoryless generative process) .

###### Proposition 1 (Memoryless noise schedules) .

###### Theorem 1 (Fine-tuning recipe for general noise schedule sampling) .

## 5 Adjoint Matching for control-affine stochastic optimal control

We discuss existing methods and also propose a new method for optimizing control-affine SOC problems. The new Adjoint Matching method is a combination of the time-tested continuous adjoint method (Pontryagin, ) with recent developments on constructing least-squares objectives for solving SOC problems (Domingo-Enrich et al., ). In this section, we briefly discuss preliminaries on existing methods, their pros and cons, then detail the Adjoint Matching algorithm and its surprising connections to the prior methods. For numerical optimization, we now assume that the control $u$ is a parametric model with parameters $\theta$.

### 5.1 Existing methods for stochastic optimal control

#### 5.1.1 The adjoint method

The most basic method of optimizing the simulation of an SDE is to directly differentiate through the simulation using gradients from the SOC objective function. The adjoint method simply uses the objective:

$$ $\textstyle\mathcal{L}(u;\bm{X}):=\int_{0}^{1}\big{(}\frac{1}{2}\|u(X_{t},t)\|^{2}\!+\!f(X_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X_{1}),\qquad\bm{X}\sim p^{u}.$ (28) $$

This is a stochastic estimate of the control objective in ([12](#S4.E12)), and the goal is to take compute the gradient of $\mathcal{L}(u;\bm{X})$ with respect to the parameters $\theta$ of the control $u$.
Due to the continuous-time nature of SDEs, there are two main approaches to implementing this numerically. Firstly, the *Discrete Adjoint* method uses a “discretize-then-differentiate” 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 and Kappen (); Gómez et al. (); Hartmann and Schütte (); Kappen et al. (); Rawlik et al. (); Haber and Ruthotto ()).
This approach, however, uses an 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., ) to 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_{t}$, expressed as an adjoint ODE, and then applies a numerical method to simulate this gradient itself, hence it is referred to as a “differentiate-then-discretize” approach (Pontryagin, ; Chen et al., ; Li et al., ). We first define the *adjoint state* as:

$$ $\textstyle\begin{split}\textstyle&\textstyle a(t;\bm{X},u):=\nabla_{X_{t}}\big{(}\int_{t}^{1}\big{(}\frac{1}{2}\|u(X_{t^{\prime}},t^{\prime})\|^{2}\!+\!f(X_{t^{\prime}},t^{\prime})\big{)}\,\mathrm{d}t^{\prime}\!+\!g(X_{1})\big{)},\\ \textstyle&\textstyle\text{where }\bm{X}\text{ solves }\mathrm{d}X_{t}=\left(b(X_{t},t)+\sigma(t)u(X_{t},t)\right)\,\mathrm{d}t+\sigma(t)\mathrm{d}B_{t}.\end{split}$ (29) $$

This implies that $\mathbb{E}_{\bm{X}\sim p^{u}}\left[a(t;\bm{X},u)\;|\;X_{t}=x\right]=\nabla_{x}J(u;x,t)$,
where $J$ denotes the cost functional defined in ([14](#S4.E14)).
It can then be shown that this adjoint state satisfies
(^6^66Note we use the convention that a Jacobian matrix $J=\nabla_{x}v(x)$ is defined as $J_{ij}=\frac{\partial v_{i}(x)}{\partial x_{j}}$.):

$$ $\textstyle\begin{split}\textstyle\frac{\mathrm{d}}{\mathrm{d}t}a(t;\bm{X},u)&\textstyle=-\left[a(t;\bm{X},u)^{\mkern-1.5mu\mathsf{T}}{}\left(\nabla_{X_{t}}(b(X_{t},t)+\sigma(t)u(X_{t},t))\right)+\nabla_{X_{t}}\left(f(X_{t},t)+\frac{1}{2}\|u(X_{t},t)\|^{2}\right)\right],\end{split}$ (30) $\textstyle a(1;\bm{X},u)$ $\textstyle=\nabla g(X_{1}).$ (31) $$

The adjoint state is solved backwards in time, starting from the terminal condition ([31](#S5.E31)). Compututation of ([30](#S5.E30)) can be done with a vector-Jacobian product which can be efficiently done on automatic differentiation software (Paszke et al., ). Once the adjoint state has been solved for $t\in[0,1]$, then the gradient of $\mathcal{L}(u;\bm{X})$ with respect to the parameters $\theta$ can be obtained by integrating over the entire time interval:

$$ $\textstyle\frac{\mathrm{d}\mathcal{L}}{\mathrm{d}\theta}=\frac{1}{2}\int_{0}^{1}\frac{\partial}{\partial\theta}\lVert u(X_{t},t)\rVert^{2}\mathrm{d}t+\int_{0}^{1}\frac{\partial u(X_{t},t)}{\partial\theta}^{\mkern-1.5mu\mathsf{T}}{}\sigma(t)^{\mkern-1.5mu\mathsf{T}}{}a(t;\bm{X},u)\mathrm{d}t,$ (32) $$

where the first term is the partial derivative of $\mathcal{L}$ w.r.t. $\theta$ and the second term is the partial derivative through the sample trajectory $\bm{X}$. See [Proposition 6](#Thmproposition6) in [Section 13.1](#S13.SS1) for a statement and proof of this result.
The discrete and continuous adjoint methods converge to the same gradient as 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., , ; Li et al., ). As the adjoint methods are essentially gradient-based optimization algorithms applied on a highly non-convex problem, many have also reported they can be unstable empirically (Mohamed et al., ; Suh et al., ; Domingo-Enrich et al., ).

#### 5.1.2 Importance-weighted matching objectives for regressing onto the optimal control

An alternative is to consider regressing onto the optimal control $u^{*}$, which is the approach of the cross-entropy method (Rubinstein and Kroese, ; Zhang et al., ) and stochastic optimal control matching (SOCM; Domingo-Enrich et al. ()). These methods make use of path integral theory (Kappen, ) to express the optimal control through importance sampling, resulting in an *importance-weighted* least-squares objective function

$$ $\textstyle\mathcal{L}_{\text{SOCM}}(u;\bm{X}):=\int_{0}^{1}\lVert u(X_{t},t)-\hat{u}^{*}(X_{t},t)\rVert^{2}\mathrm{d}t\times\omega(u,\bm{X}),\qquad\bm{X}\sim p^{u},$ (33) $$

where $\omega$ is an importance weighting that approximates sampling from the optimal distribution $p^{*}$, and $\hat{u}^{*}$ is a stochastic estimator of the optimal control relying on having sampled from the optimal process. We defer to Domingo-Enrich et al. () for the exact details. The functional landscape of this objective is convex, which is argued to help yield stable training. However, the need for importance sampling renders this impractical for high dimensional applications: the variance of the importance weighting $\omega$ grows exponentially with dimension of the stochastic process, leading to catastrophic failure. This unfortunately means that such importance-weighted matching objectives are impractical for fine-tuning dynamical generative models; however, a least-squares objective is greatly coveted as it can lead to stable training and simple interpretations.

### 5.2 Adjoint Matching

We make two important observations which lead to our proposed method: (i) it is possible to construct a matching objective without any importance weighting, and (ii) there are unnecessary terms in the adjoint differential equation ([30](#S5.E30)) that can lead to higher variance at convergence.

Firstly, we notice that we can simply match the gradient of the cost functional under the current control. That is, while SOCM carefully constructs an importance-weighted estimator of the optimal control $u^{*}=-\sigma(t)^{\mkern-1.5mu\mathsf{T}}{}\nabla J(u^{*};x,t)$ ([17](#S4.E17)), we claim that we can actually just regress onto the target vector field $-\sigma(t)^{\mkern-1.5mu\mathsf{T}}{}\nabla J(u;x,t)$ where $u$ is the current control, and furthermore, this results in a gradient equal in expectation to the continuous adjoint method.
We formalize this in the following proposition, proven in [Section 13.2](#S13.SS2):

###### Proposition 2 .

Let us define, for now, the basic Adjoint Matching objective as:

$$ $\textstyle\begin{split}\textstyle\mathcal{L}_{\mathrm{Basic-Adj-Match}}(u;\bm{X})&\textstyle:=\frac{1}{2}\int_{0}^{1}\big{\|}u(X^{\bar{u}}_{t},t)+\sigma(t)^{\mkern-1.5mu\mathsf{T}}{}a(t;\bm{X}^{\bar{u}},\bar{u})\big{\|}^{2}\,\mathrm{d}t,\qquad\bm{X}\sim p^{\bar{u}},\quad\bar{u}=\texttt{stopgrad}(u),\end{split}$ (34) $$

where $\bar{u}=\texttt{stopgrad}(u)$ means that the gradients of $\bar{u}$ with respect to the parameters $\theta$ of the control $u$ are artificially set to zero. The gradient of $\mathcal{L}_{\mathrm{Basic-Adj-Match}}(u;\bm{X})$ with respect to $\theta$ is equal to the gradient $\frac{\mathrm{d}\mathcal{L}}{\mathrm{d}\theta}$ in equation ([32](#S5.E32)).
Importantly, the only critical point of $\mathbb{E}\left[\mathcal{L}_{\mathrm{Basic-Adj-Match}}\right]$ is the optimal control $u^{*}$.

Critical points of $\mathcal{L}$ are controls $u$ such that $\frac{\delta}{\delta u}\mathcal{L}(u)=0$, where $\frac{\delta}{\delta u}\mathcal{L}$ denotes the first variation of the functional $\mathcal{L}$.
In other words, [Proposition 2](#Thmproposition2) 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 is that it is a *consistency loss*. The Adjoint Matching objective is based off of the observation that the optimal control $u^{*}(x,t)$ is the unique fixed-point of the relation $u(x,t)=-\sigma(t)^{\mkern-1.5mu\mathsf{T}}{}\nabla_{x}J(u;x,t)$ (see [Lemma 6](#Thmlemma6) in [Section 13.2](#S13.SS2))
and so we are directly optimizing for a control that fits this relation, while using the adjoint state as a stochastic estimator of $\nabla_{x}J(u;x,t)$ (LABEL:eq:adjoint_state_defn).

The basic Adjoint Matching objective in [Proposition 2](#Thmproposition2) does not yet yield a novel algorithm for stochastic optimal control, because it produces the same gradient as the continuous adjoint method. This can be seen by taking the gradient w.r.t. $\theta$ after expanding the square in ([34](#S5.E34)) and removing terms that do not depend on $\theta$ to arrive exactly at the continuous adjoint method ([32](#S5.E32)).
However, it provides the means of deriving a simpler leaner objective function.

###### Proposition 2 .

##### 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

$$ $\textstyle u^{*}(x,t)=\mathbb{E}_{\bm{X}\sim p^{*}}\left[-\sigma(t)^{\mkern-1.5mu\mathsf{T}}{}a(t;\bm{X},u^{*})|X_{t}=x\right].$ (35) $$

Multiplying both sides by the Jacobian $\nabla_{x}u^{*}(x,t)$ and re-arranging, we get the relation

$$ $\textstyle\mathbb{E}_{\bm{X}\sim p^{*}}\left[u^{*}(x,t)^{\mkern-1.5mu\mathsf{T}}{}\nabla_{x}u^{*}(x,t)+a(t;\bm{X},u^{*})^{\mkern-1.5mu\mathsf{T}}{}\sigma(t)\nabla_{x}u^{*}(x,t)\;|\;X_{t}=x\right]=0.$ (36) $$

Notice that the terms inside the expectation in ([36](#S5.E36)) show up as part of the adjoint differential equation ([30](#S5.E30)), which we have now shown to have expectation zero at the optimal solution. Furthermore, the variance of the terms inside ([36](#S5.E36)) is non-zero even at the optimal solution, so the basic Adjoint Matching ([34](#S5.E34)) and hence *the continuous adjoint method will also have non-vanishing gradients even when $u=u^{*}$*.

Therefore, we motivate the definition of a *lean adjoint state* $\tilde{a}$ with the terms in ([36](#S5.E36)) removed. Plugging this lean adjoint back into the least-squares objective, we obtain our final proposed Adjoint Matching objective:

$\textstyle\mathcal{L}_{\mathrm{Adj-Match}}(u;\bm{X}):=\frac{1}{2}\int_{0}^{1}\big{\|}$
$\textstyle u(X_{t},t)+\sigma(t)^{\mkern-1.5mu\mathsf{T}}{}\tilde{a}(t;\bm{X})\big{\|}^{2}\,\mathrm{d}t,\qquad\bm{X}\sim p^{\bar{u}},\quad\bar{u}=\texttt{stopgrad}(u),$
(37)
$\textstyle\text{where }\quad\frac{\mathrm{d}}{\mathrm{d}t}\tilde{a}(t;\bm{X})$
$\textstyle=-(\tilde{a}(t;\bm{X})^{\top}\nabla_{x}b(X_{t},t)+\nabla_{x}f(X_{t},t)),$
(38)
$\textstyle\tilde{a}(1;\bm{X})$
$\textstyle=\nabla_{x}g(X_{1}).$
(39)

Equations ([38](#S5.E38))-([39](#S5.E39)) define the *lean adjoint state*, and ([37](#S5.E37)) is the complete Adjoint Matching objective.
The unique critical point of $\mathbb{E}[\mathcal{L}_{\mathrm{Adj-Match}}]$ is the optimal control, which we prove relying on [Proposition 2](#Thmproposition2) and equation ([36](#S5.E36)) (see [Proposition 7](#Thmproposition7) in [Section 13.3](#S13.SS3)).

Compared to the importance sampling methods ([Section 5.1.2](#S5.SS1.SSS2)), Adjoint Matching is a simple least-squares regression objective and has no importance weighting. This allows it to avoid the pitfalls of high variance importance weights and makes it as scalable as the adjoint methods while retaining the interpretation of matching a target vector field.

Compared to the adjoint method ([Section 5.1.1](#S5.SS1.SSS1)), 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., (LABEL:eq:adjoint_state_defn) is not true, except at the optimum when $u=u^{*}$.
Even at the optimal solution, since Adjoint Matching removes terms that have expectation zero, it can potentially exhibit better convergence and lower variance than the continuous adjoint method.
Additionally, computation of the lean adjoint state ([38](#S5.E38)) also exhibits a smaller computational cost due to the removal of the extra terms (no longer need the Jacobian of the control $\nabla_{x}u$).
We provide a rigorous derivation of Adjoint Matching and the above claims in [Section 13.3](#S13.SS3).

Adjoint Matching can be applied to reward fine-tuning of dynamical generative models through the memoryless SOC formulation discussed in [Section 4](#S4). We provide pseudo-code for this in [Algorithm 1](#alg1) for Flow Matching models and in [Algorithm 2](#alg2) in [Section 13.4](#S13.SS4) for denoising diffusion models.

Figure: Algorithm 1 Adjoint Matching for fine-tuning Flow Matching models

## 6 Related work

##### Fine-tuning from human feedback.

There are two main overarching approaches to RLHF: the reward-based approach (Ziegler et al., ; Stiennon et al., ; Ouyang et al., ; Bai et al., ) and direct preference optimization (DPO; Rafailov et al. ()).
The reward-based approach (Ziegler et al., ; Stiennon et al., ; Ouyang et al., ; Bai et al., ) consists in learning 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 a filtered dataset, and does not work explicitly with a reward model.
Furthermore, for flow and diffusion models specifically, it is possible to differentiate the reward function, so there is a larger emphasis on reward-based approaches.

##### Fine-tuning for diffusion models.

Among existing reward-based diffusion fine-tuning methods, Fan and Lee () interpret the denoising process as a multi-step decision-making task and use policy gradient algorithms to fine-tune diffusion samplers.
Black et al. () makes use of proximal policy gradients for fine-tuning but this does not make use of the differentiability of the reward model.
Fan et al. () also consider KL-regularized rewards ([19](#S4.E19)) but do not make the critical connection to the tilted distribution ([1](#S1.E1)) that we flesh out in [Section 4.2](#S4.SS2).
The fine-tuning algorithms of Xu et al. (); Clark et al. () directly take gradients of the reward model and use heuristics to try to stay close to the original base generative model, but their behavior is not well understood and unrelated to the tilted distribution: Xu et al. () takes gradients of the reward applied on the denoised sample at different points in time, and Clark et al. () backpropagates the reward function through all or part of the diffusion trajectory.
Finally, Uehara et al. () also fine-tune diffusion models with the goal of sampling from the tilted distribution ([1](#S1.E1)), but their approach is much more involved than ours as it requires learning a value function, and solving two stochastic optimal control problems.

##### 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 a guidance term to the velocity (Chung et al., ; Song et al., ; Pokle et al., ); 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, ; Wallace et al., ; Ben-Hamu et al., ); 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., ; Rout et al., ); these approaches aim to solve a separate control problem for each generated sample, instead of performing amortization (Amos et al., ) 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., ), SDEs (Li et al., ), and jump processes (Chen et al., ), enabled through the adjoint method.
They can be used to train arbitrary generative processes, but for simplified constructions these have fallen in favor to simulation-free matching objectives such as denoising score matching (Vincent, ) and Flow Matching (Lipman et al., ).
The optimal control formalism also has significance in sampling from un-normalized distributions (Zhang and Chen, ; Berner et al., ).
The inclusion of a state cost has been used to solve transport problems where intermediate path distributions are of importance (Liu et al., ; Pooladian et al., ).
These collective advances naturally lead to the consideration of the optimal control formalism for reward fine-tuning.

## 7 Experiments

**Table 2: Evaluation metrics of different fine-tuning methods for text-to-image generation. The second and third columns show the noise schedules $\sigma(t)$ used for fine-tuning and for sampling: $\sigma(t)=\sqrt{2\eta_{t}}$ corresponds to Memoryless Flow Matching, and $\sigma(t)=0$ to the Flow Matching ODE ([3](#S2.E3)). We report standard errors estimated over 3 runs of the fine-tuning algorithm on random sets of 40000 training prompts, each evaluated over a random set of 1000 test prompts.**
|  | Fine-tuning | Fine-tuning | Sampling | ClipScore$\,\uparrow$ | PickScore$\,\uparrow$ | HPS v2$\,\uparrow$ | DreamSim |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  | Method | $\sigma(t)$ | $\sigma(t)$ | Diversity$\,\uparrow$ |  |  |  |
|  | None | N/A | $\sqrt{2\eta_{t}}$ | 24.15$\pm$0.26 | 17.25$\pm$0.06 | 16.19$\pm$0.17 | 53.60$\pm$1.37 |
|  | (Base model) | 0 | 28.32$\pm$0.22 | 18.15$\pm$0.07 | 17.89$\pm$0.16 | 56.53$\pm$1.52 |  |
| Baselines | DRaFT-1 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 30.18$\pm$0.24 | 19.38$\pm$0.08 | 24.61$\pm$0.17 | 25.54$\pm$0.99 |
| 0 | 0 | 30.95$\pm$0.28 | 19.37$\pm$0.06 | 24.37$\pm$0.17 | 27.39$\pm$1.14 |  |  |
| DRaFT-40 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 26.94$\pm$0.28 | 18.34$\pm$0.19 | 19.98$\pm$1.02 | 41.98$\pm$2.14 |  |
| 0 | 0 | 30.07$\pm$0.39 | 19.45$\pm$0.08 | 24.06$\pm$0.24 | 36.53$\pm$1.69 |  |  |
| DPO | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 24.11$\pm$0.22 | 17.24$\pm$0.06 | 16.15$\pm$0.14 | 53.27$\pm$1.36 |  |
| 0 | 0 | 27.77$\pm$0.18 | 17.92$\pm$0.07 | 17.30$\pm$0.20 | 54.11$\pm$1.50 |  |  |
| ReFL | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 28.59$\pm$0.31 | 18.68$\pm$0.10 | 22.24$\pm$0.46 | 32.71$\pm$2.76 |  |
| 0 | 0 | 30.06$\pm$0.63 | 19.07$\pm$0.21 | 23.06$\pm$0.41 | 32.69$\pm$1.28 |  |  |
| Memoryless SOC | Cont. Adjoint | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 26.99$\pm$0.43 | 18.33$\pm$0.16 | 20.83$\pm$0.63 | 46.59$\pm$1.40 |
| $\lambda=12500$ | 0 | 29.49$\pm$0.32 | 18.98$\pm$0.16 | 21.34$\pm$0.53 | 48.41$\pm$1.44 |  |  |
| Disc. Adjoint | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 28.04$\pm$0.57 | 18.44$\pm$0.21 | 20.04$\pm$0.39 | 54.90$\pm$2.03 |  |
| $\lambda=12500$ | 0 | 29.28$\pm$0.17 | 18.82$\pm$0.14 | 19.73$\pm$0.17 | 53.36$\pm$2.48 |  |  |
| Adj.-Matching | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 30.36$\pm$0.22 | 19.29$\pm$0.08 | 24.12$\pm$0.17 | 40.89$\pm$1.50 |  |
| $\lambda=1000$ | 0 | 31.41$\pm$0.22 | 19.57$\pm$0.09 | 23.29$\pm$0.18 | 43.10$\pm$1.76 |  |  |
| Adj.-Matching | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 30.59$\pm$0.40 | 19.49$\pm$0.10 | 24.85$\pm$0.23 | 37.07$\pm$1.47 |  |
| $\lambda=2500$ | 0 | 31.64$\pm$0.21 | 19.71$\pm$0.09 | 24.12$\pm$0.27 | 39.88$\pm$1.59 |  |  |
| Adj.-Matching | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 30.62$\pm$0.30 | 19.50$\pm$0.09 | 24.95$\pm$0.28 | 34.50$\pm$1.33 |  |
| $\lambda=12500$ | 0 | 31.65$\pm$0.19 | 19.76$\pm$0.08 | 24.49$\pm$0.27 | 37.24$\pm$1.57 |  |  |

Figure: Adjoint Matching (Ours)
Refer to caption: /html/2409.08861/assets/figs/samples/adjmat1000_prompt_32_image_0.jpg

Figure: Text prompt: “Man sitting on sofa at home in front of fireplace and using laptop computer, rear view”
Refer to caption: /html/2409.08861/assets/figs/samples/adjmat0.0_prompt_37_image_3.jpg

Figure: Figure 5: Tradeoffs between different aspects of generative models: text-to-image consistency (ClipScore), sample diversity for each prompt (DreamSim Diversity), and generalization to unseen human preferences (HPS v2). Different points are obtained from varying values of $\lambda$ for Adjoint Matching and varying number of fine-tuning iterations for the DRaFT-1 baseline. Overall, we find our proposed method Adjoint Matching has the best Pareto fronts.
Refer to caption: /html/2409.08861/assets/x1.png

We experimentally validate our proposed method on reward fine-tuning a Flow Matching base model (Lipman et al., ). In particular, we use the usual setup of pre-training an autoencoder for 512$\times$512 resolution images, then training a text-conditional Flow Matching model on the latent variables with a U-net architecture (Long et al., ), similar to the setup in Rombach et al. (). We pre-trained our base model using a dataset of licensed text and image pairs. Then for fine-tuning, we consider the reward function:

$$ $\textstyle r(x):=\lambda\times\texttt{RewardModel}(x)$ (43) $$

corresponding to a scaled version of the reward model, which we take to be ImageReward (Xu et al., ). Different values of $\lambda$ provide different tradeoffs between the KL regularization and the reward model ([19](#S4.E19)).

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., ). For consistency, we make use of the standard ClipScore (Hessel et al., ) and PickScore (Kirstain et al., ); for generalization to unseen human preferences, we use the HPSv2 model (Wu et al., ); and for diversity, we compute averages of pairwise distances of the DreamSim features (Fu et al., ). More details are provided in [Section 15.4](#S15.SS4).

As our baselines, we consider the DPO (Wallace et al., ), ReFL (Xu et al., ), and DRaFT-K algorithms (Clark et al., ). DPO does not use gradients from the reward function, while ReFL and DRaFT make use of heuristic gradient stopping approaches to stay close to the base generative model.
Out of these baseline methods, we find that DRaFT-1 performs the best, so we perform additional ablation experiments comparing to this method.
Within the same SOC formulation as our method, we also consider the discrete and continuous adjoint methods. We provide full experimental details in [Section 15](#S15); an important implementation detail is that we slightly offset $\sigma(t)$ in order to avoid division by zero.

##### Evaluation results.

In [Table 2](#S7.T2) we report the evaluation metrics for the baselines as well as our proposed Adjoint Matching approach. We compare each method at roughly the same wall clock time (see the times and number of iterations in [Table 8](#S9.T8)).
We find that across all metrics, 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 in poorer consistency and human preference metrics.

##### Ablation: base model vs. reward tradeoff.

We note that the scaling in front of the reward model $\lambda$ determines how strongly the we should prefer the reward model over the base model. As such, we see a natural tradeoff curve: higher $\lambda$ 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 $\lambda$. 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 $\lambda$ in [Figure 3](#S7.F3), while for the DRaFT-1 baseline we show the effect of varying the number of fine-tuning iterations.

##### Ablation: classifier-free guidance.

We note that it is possible to apply classifier-free guidance (CFG; Ho and Salimans (); Zheng et al. ()) after fine-tuning. We use the formula $(1+w)v(x,t|y)-wv(x,t)$, where $w$ is the guidance weight, $v(x,t|y)$ is a 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 [Figure 5](#S7.F5) we show the evaluation metrics with classifier-free guidance applied. Comparing three different guidance weight values, we see a higher weight does improve text-to-image consistency, and to some extent, human preference, but this comes at the cost of being worse in terms of diversity. We show qualitative differences in [Figure 4](#S7.F4).

## 8 Conclusion

We investigate the problem of fine-tuning dynamical generative models such as Flow Matching and propose the use of a stochastic optimal control (SOC) formulation with a 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 optimal control problems, by casting SOC as a 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-based methods.
In summary, we are the first to provide a theoretically-driven algorithm for fine-tuning Flow Matching models, and we find that our approach significantly outperforms baseline methods across multiple axes of evaluation—text-to-image consistency, generalization to unseen human preference, and sample diversity—on large-scale text-to-image generation.

## 9 Additional Figures & Tables

**Table 3: Metrics for various fine-tuning methods for text-to-image generation. The second and third columns show the noise schedules $\sigma(t)$ used for fine-tuning and for inference: $\sigma(t)=\sqrt{2\eta_{t}}$ corresponds to Memoryless Flow Matching, and $\sigma(t)=0$ to the Flow Matching ODE ([3](#S2.E3)). Confidence intervals show standard errors of estimates; computed over 3 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-tuning | Fine-tuning | Sampling | ImageReward$\,\uparrow$ | ClipScore | PickScore | Total time (s) / |
| --- | --- | --- | --- | --- | --- | --- |
| loss | $\sigma(t)$ | $\sigma(t)$ | diversity$\,\uparrow$ | diversity$\,\uparrow$ | # iterations |  |
| None | N/A | $\sqrt{2\eta_{t}}$ | $-$1.384$\pm$0.040 | 28.07$\pm$1.40 | 1.63$\pm$0.08 | N/A |
| ($\mathrm{CFG}=1.0$) | 0 | $-$0.920$\pm$0.042 | 30.29$\pm$1.53 | 1.82$\pm$0.09 |  |  |
| DRaFT-1 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.357$\pm$0.039 | 16.86$\pm$0.98 | 1.21$\pm$0.07 | 140k$\pm$5.9k |
| 0 | 0 | 1.251$\pm$0.040 | 16.76$\pm$1.06 | 1.27$\pm$0.07 | / 4000 |  |
| DRaFT-40 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | $-$0.560$\pm$0.138 | 24.07$\pm$1.37 | 1.64$\pm$0.12 | 148k$\pm$4.2k |
| 0 | 0 | 0.424$\pm$0.042 | 20.99$\pm$1.54 | 1.67$\pm$0.08 | / 1500 |  |
| DPO | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | $-$1.386$\pm$0.033 | 27.80$\pm$1.40 | 1.62$\pm$0.08 | 118k$\pm$0.6k |
| 0 | 0 | $-$0.957$\pm$0.040 | 29.81$\pm$1.43 | 1.68$\pm$0.10 | / 1000 |  |
| ReFL | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.687$\pm$0.085 | 19.49$\pm$1.76 | 1.22$\pm$0.08 | 173k$\pm$10.9k |
| 0 | 0 | 0.709$\pm$0.080 | 18.39$\pm$1.11 | 1.31$\pm$0.10 | / 6000 |  |
| Cont. Adjoint | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | $-$0.448$\pm$0.135 | 26.97$\pm$1.37 | 1.82$\pm$0.09 | 153k$\pm$0.9k |
| $\lambda=12500$ | 0 | $-$0.249$\pm$0.116 | 26.25$\pm$1.30 | 1.90$\pm$0.10 | / 750 |  |
| Disc. Adjoint | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | $-$0.557$\pm$0.113 | 30.40$\pm$2.39 | 1.91$\pm$0.09 | 152k$\pm$1.5k |
| $\lambda=12500$ | 0 | $-$0.552$\pm$0.041 | 28.37$\pm$2.26 | 1.97$\pm$0.09 | / 1000 |  |
| Adj.-Matching | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.550$\pm$0.043 | 23.00$\pm$1.27 | 1.65$\pm$0.08 |  |
| $\lambda=1000$ | 0 | 0.454$\pm$0.055 | 22.76$\pm$1.40 | 1.73$\pm$0.09 |  |  |
| Adj.-Matching | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.755$\pm$0.040 | 21.33$\pm$1.71 | 1.55$\pm$0.08 | 156k$\pm$1.9k |
| $\lambda=2500$ | 0 | 0.671$\pm$0.047 | 21.42$\pm$1.54 | 1.64$\pm$0.08 | / 1000 |  |
| Adj.-Matching | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.882$\pm$0.058 | 20.49$\pm$1.48 | 1.50$\pm$0.09 |  |
| $\lambda=12500$ | 0 | 0.778$\pm$0.050 | 20.34$\pm$1.49 | 1.57$\pm$0.09 |  |  |

**Table 4: Evaluation metrics when using classifier-free guidance (CFG; Ho and Salimans ()).**
| $w$ | Fine-tuning | $\#$iter. | Fine-tun. | Sampl. | ImageReward$\,\uparrow$ | ClipScore$\,\uparrow$ | PickScore$\,\uparrow$ | HPS v2$\,\uparrow$ | DreamSim |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| loss | / $\lambda$ | $\sigma(t)$ | $\sigma(t)$ | diversity$\,\uparrow$ |  |  |  |  |  |
| 0.0 | None | N/A | N/A | $\sqrt{2\eta_{t}}$ | $-$1.384$\pm$0.040 | 24.15$\pm$0.26 | 17.25$\pm$0.06 | 16.19$\pm$0.17 | 53.60$\pm$1.37 |
| 0 | $-$0.920$\pm$0.042 | 28.32$\pm$0.22 | 18.15$\pm$0.07 | 17.89$\pm$0.16 | 56.53$\pm$1.52 |  |  |  |  |
| 0.0 | DRaFT-1 | 1000 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.913$\pm$0.068 | 29.80$\pm$0.22 | 19.16$\pm$0.06 | 23.63$\pm$0.16 | 35.21$\pm$1.93 |
| 0 | 0 | 0.626$\pm$0.195 | 30.48$\pm$0.32 | 18.91$\pm$0.34 | 21.92$\pm$1.63 | 38.52$\pm$2.01 |  |  |  |
| 2000 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.204$\pm$0.046 | 29.90$\pm$0.43 | 19.29$\pm$0.12 | 24.40$\pm$0.27 | 28.51$\pm$1.68 |  |  |
| 0 | 0 | 1.052$\pm$0.088 | 30.65$\pm$0.24 | 19.27$\pm$0.11 | 23.81$\pm$0.44 | 32.11$\pm$2.37 |  |  |  |
| 3000 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.307$\pm$0.041 | 29.96$\pm$0.22 | 19.31$\pm$0.06 | 24.42$\pm$0.13 | 26.57$\pm$1.32 |  |  |
| 0 | 0 | 1.173$\pm$0.058 | 30.86$\pm$0.25 | 19.37$\pm$0.06 | 24.17$\pm$0.23 | 29.69$\pm$1.30 |  |  |  |
| 4000 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.357$\pm$0.039 | 30.18$\pm$0.24 | 19.38$\pm$0.08 | 24.61$\pm$0.17 | 25.54$\pm$0.99 |  |  |
|  |  | 0 | 0 | 1.251$\pm$0.040 | 30.95$\pm$0.28 | 19.37$\pm$0.06 | 24.37$\pm$0.17 | 27.39$\pm$1.14 |  |
| 0.0 | Adj.-Match. | 1000 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.550$\pm$0.043 | 30.36$\pm$0.22 | 19.29$\pm$0.08 | 24.12$\pm$0.17 | 40.89$\pm$1.50 |
| 0 | 0 | 0.454$\pm$0.055 | 31.41$\pm$0.22 | 19.57$\pm$0.09 | 23.29$\pm$0.18 | 43.10$\pm$1.76 |  |  |  |
| 2500 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.755$\pm$0.040 | 30.59$\pm$0.40 | 19.49$\pm$0.10 | 24.85$\pm$0.23 | 37.07$\pm$1.47 |  |  |
| 0 | 0 | 0.671$\pm$0.047 | 31.64$\pm$0.21 | 19.71$\pm$0.09 | 24.12$\pm$0.27 | 39.88$\pm$1.59 |  |  |  |
| 12500 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.882$\pm$0.058 | 30.62$\pm$0.30 | 19.50$\pm$0.09 | 24.95$\pm$0.28 | 34.50$\pm$1.33 |  |  |
| 0 | 0 | 0.778$\pm$0.050 | 31.65$\pm$0.19 | 19.76$\pm$0.08 | 24.49$\pm$0.27 | 37.24$\pm$1.57 |  |  |  |
| 1.0 | None | N/A | N/A | $\sqrt{2\eta_{t}}$ | $-$0.269$\pm$0.050 | 30.41$\pm$0.22 | 18.74$\pm$0.07 | 20.47$\pm$0.18 | 43.82$\pm$1.24 |
| 0 | $-$0.123$\pm$0.041 | 31.83$\pm$0.17 | 19.28$\pm$0.07 | 20.95$\pm$0.16 | 42.59$\pm$1.23 |  |  |  |  |
| 1.0 | DRaFT-1 | 1000 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.123$\pm$0.051 | 32.06$\pm$0.19 | 19.69$\pm$0.06 | 24.56$\pm$0.17 | 28.25$\pm$1.55 |
| 0 | 0 | 0.856$\pm$0.167 | 32.32$\pm$0.25 | 19.38$\pm$0.34 | 22.88$\pm$1.54 | 29.98$\pm$1.86 |  |  |  |
| 2000 | 0 | 0 | 1.177$\pm$0.053 | 32.36$\pm$0.18 | 19.67$\pm$0.08 | 24.48$\pm$0.28 | 25.09$\pm$1.82 |  |  |
| 3000 | 0 | 0 | 1.255$\pm$0.038 | 32.36$\pm$0.19 | 19.70$\pm$0.06 | 24.64$\pm$0.17 | 23.24$\pm$1.19 |  |  |
| 4000 | 0 | 0 | 1.296$\pm$0.033 | 32.30$\pm$0.19 | 19.68$\pm$0.06 | 24.71$\pm$0.14 | 21.54$\pm$0.96 |  |  |
| 1.0 | Adj.-Match. | 1000 | 0 | 0 | 0.782$\pm$0.044 | 33.05$\pm$0.22 | 20.20$\pm$0.09 | 24.81$\pm$0.18 | 32.67$\pm$1.26 |
| 2500 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.027$\pm$0.038 | 32.85$\pm$0.21 | 20.08$\pm$0.08 | 25.88$\pm$0.20 | 29.83$\pm$1.00 |  |  |
| 0 | 0 | 0.910$\pm$0.040 | 33.20$\pm$0.17 | 20.29$\pm$0.09 | 25.39$\pm$0.24 | 30.34$\pm$1.51 |  |  |  |
| 12500 | 0 | 0 | 0.985$\pm$0.041 | 33.10$\pm$0.18 | 20.28$\pm$0.08 | 25.61$\pm$0.27 | 28.86$\pm$1.37 |  |  |
| 4.0 | None | N/A | N/A | $\sqrt{2\eta_{t}}$ | 0.277$\pm$0.043 | 32.68$\pm$0.18 | 19.50$\pm$0.07 | 22.29$\pm$0.16 | 35.12$\pm$0.92 |
| 0 | 0.209$\pm$0.046 | 32.83$\pm$0.17 | 19.79$\pm$0.07 | 22.30$\pm$0.17 | 32.05$\pm$1.05 |  |  |  |  |
| 4.0 | DRaFT-1 | 1000 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.062$\pm$0.045 | 32.29$\pm$0.16 | 19.48$\pm$0.06 | 23.67$\pm$0.13 | 25.03$\pm$1.32 |
| 0 | 0 | 0.604$\pm$0.395 | 31.80$\pm$0.86 | 19.09$\pm$0.53 | 21.69$\pm$2.10 | 25.92$\pm$2.57 |  |  |  |
| 2000 | 0 | 0 | 1.112$\pm$0.046 | 32.29$\pm$0.20 | 19.34$\pm$0.11 | 23.31$\pm$0.22 | 21.02$\pm$1.67 |  |  |
| 3000 | 0 | 0 | 1.151$\pm$0.036 | 32.31$\pm$0.21 | 19.36$\pm$0.06 | 23.29$\pm$0.14 | 19.53$\pm$1.24 |  |  |
| 4000 | 0 | 0 | 1.172$\pm$0.040 | 32.20$\pm$0.22 | 19.30$\pm$0.07 | 23.20$\pm$0.15 | 18.45$\pm$1.06 |  |  |
| 4.0 | Adj.-Match. | 1000 | 0 | 0 | 0.852$\pm$0.046 | 33.50$\pm$0.22 | 20.31$\pm$0.08 | 24.97$\pm$0.19 | 25.83$\pm$0.82 |
| 2500 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.052$\pm$0.039 | 33.51$\pm$0.19 | 20.15$\pm$0.07 | 25.56$\pm$0.18 | 26.21$\pm$0.73 |  |  |
| 0 | 0 | 0.942$\pm$0.042 | 33.61$\pm$0.19 | 20.35$\pm$0.08 | 25.34$\pm$0.21 | 24.30$\pm$0.86 |  |  |  |
| 12500 | 0 | 0 | 1.007$\pm$0.052 | 33.48$\pm$0.20 | 20.29$\pm$0.08 | 25.50$\pm$0.29 | 23.48$\pm$0.81 |  |  |

**Table 5: Metrics for alternative optimization hyperparameters (learning rate and Adam $\beta_{1}$).**
| LR / | Fine-tuning | Fine-tun. | Generat. | ImageReward$\,\uparrow$ | ClipScore$\,\uparrow$ | PickScore$\,\uparrow$ | HPS v2$\,\uparrow$ | DreamSim |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Adam $\beta_{1}$ | loss | $\sigma(t)$ | $\sigma(t)$ | diversity$\,\uparrow$ |  |  |  |  |
| $3\text{\times}{10}^{-5}$ | DRaFT-1 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.467$\pm$0.029 | 30.28$\pm$0.56 | 19.37$\pm$0.09 | 24.70$\pm$0.15 | 21.20$\pm$0.93 |
| / $0.97$ | Adj.-Match. | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.130$\pm$0.034 | 31.01$\pm$0.27 | 19.60$\pm$0.08 | 25.01$\pm$0.25 | 26.73$\pm$0.88 |
|  | $\lambda=12500$ |  |  |  |  |  |  |  |
| $2\text{\times}{10}^{-5}$ | Disc. Adj. | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | $-$1.186$\pm$0.553 | 21.95$\pm$4.29 | 16.94$\pm$0.95 | 12.34$\pm$4.40 | 28.33$\pm$10.26 |
| / $0.95$ | $\lambda=12500$ | 0 | 0 | $-$0.961$\pm$0.653 | 24.07$\pm$4.71 | 17.86$\pm$1.17 | 15.93$\pm$5.80 | 33.62$\pm$7.80 |

**Table 6: Comparison with an alternative fine-tuning noise schedule $\sigma(t)=1$. We see that the initial value function bias ([Section 4.2](#S4.SS2)) results in the model not having a high reward function (ImageReward is the reward function used for fine-tuning). Its performance on other metrics are also lower than when fine-tuning with the memoryless noise schedule, except for diversity.**
| Fine-tuning | Fine-tuning | Generative | ImageReward$\,\uparrow$ | ClipScore$\,\uparrow$ | PickScore$\,\uparrow$ | HPS v2$\,\uparrow$ | DreamSim |
| --- | --- | --- | --- | --- | --- | --- | --- |
| loss | $\sigma(t)$ | $\sigma(t)$ | diversity$\,\uparrow$ |  |  |  |  |
| Adj.-Matching | 1 | 1 | 0.009$\pm$0.077 | 29.18$\pm$0.51 | 18.66$\pm$0.09 | 20.75$\pm$0.32 | 41.33$\pm$1.24 |
| $\lambda=12500$ | 0 | 0.454$\pm$0.055 | 31.41$\pm$0.22 | 19.57$\pm$0.09 | 23.29$\pm$0.18 | 43.10$\pm$1.76 |  |
| Adj.-Matching | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.882$\pm$0.058 | 30.62$\pm$0.30 | 19.50$\pm$0.09 | 24.95$\pm$0.28 | 34.50$\pm$1.33 |
| $\lambda=12500$ | 0 | 0.778$\pm$0.050 | 31.65$\pm$0.19 | 19.76$\pm$0.08 | 24.49$\pm$0.27 | 37.24$\pm$1.57 |  |

**Table 7: 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 40 steps.**
| $\#$sampl. | Fine-tuning | Fine-tun. | Sampl. | ImageReward$\,\uparrow$ | ClipScore$\,\uparrow$ | PickScore$\,\uparrow$ | HPS v2$\,\uparrow$ | DreamSim |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| timesteps | loss | $\sigma(t)$ | $\sigma(t)$ | diversity$\,\uparrow$ |  |  |  |  |
| $10$ | None (Base) | N/A | $\sqrt{2\eta_{t}}$ | $-$2.279$\pm$0.001 | 13.99$\pm$0.12 | 14.98$\pm$0.05 | 7.37$\pm$0.10 | 5.07$\pm$0.13 |
| 0 | $-$1.386$\pm$0.040 | 26.26$\pm$0.24 | 17.64$\pm$0.07 | 14.92$\pm$0.17 | 51.26$\pm$1.38 |  |  |  |
| DRaFT-1 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.033$\pm$0.051 | 25.98$\pm$0.25 | 18.28$\pm$0.07 | 22.08$\pm$0.18 | 14.47$\pm$0.67 |  |
| 0 | 0 | 1.236$\pm$0.038 | 31.54$\pm$0.27 | 19.53$\pm$0.07 | 24.47$\pm$0.19 | 24.78$\pm$0.88 |  |  |
| Adj.-Match. | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | $-$2.104$\pm$0.074 | 17.12$\pm$0.56 | 15.76$\pm$0.20 | 11.48$\pm$1.03 | 9.88$\pm$0.81 |  |
| $\lambda=12500$ | 0 | 0.607$\pm$0.055 | 31.36$\pm$0.20 | 19.56$\pm$0.08 | 23.23$\pm$0.28 | 33.75$\pm$1.48 |  |  |
| $20$ | None (Base) | N/A | $\sqrt{2\eta_{t}}$ | $-$2.275$\pm$0.002 | 14.58$\pm$0.13 | 15.07$\pm$0.05 | 7.47$\pm$0.10 | 11.27$\pm$0.33 |
| 0 | $-$1.017$\pm$0.055 | 27.92$\pm$0.19 | 18.01$\pm$0.07 | 17.17$\pm$0.15 | 54.69$\pm$1.45 |  |  |  |
| DRaFT-1 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.301$\pm$0.039 | 27.09$\pm$0.24 | 18.93$\pm$0.07 | 23.78$\pm$0.20 | 21.05$\pm$1.12 |  |
| 0 | 0 | 1.255$\pm$0.038 | 31.14$\pm$0.25 | 19.43$\pm$0.06 | 24.52$\pm$0.16 | 26.15$\pm$1.11 |  |  |
| Adj.-Match. | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | $-$0.032$\pm$0.072 | 25.07$\pm$0.27 | 18.01$\pm$0.07 | 20.75$\pm$0.23 | 29.06$\pm$2.34 |  |
| $\lambda=12500$ | 0 | 0.768$\pm$0.048 | 31.70$\pm$0.17 | 19.73$\pm$0.08 | 24.30$\pm$0.26 | 35.90$\pm$1.52 |  |  |
| $40$ | None (Base) | N/A | $\sqrt{2\eta_{t}}$ | $-$1.384$\pm$0.040 | 24.15$\pm$0.26 | 17.25$\pm$0.06 | 16.19$\pm$0.17 | 53.60$\pm$1.37 |
| 0 | $-$0.920$\pm$0.042 | 28.32$\pm$0.22 | 18.15$\pm$0.07 | 17.89$\pm$0.16 | 56.53$\pm$1.52 |  |  |  |
| DRaFT-1 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.357$\pm$0.039 | 30.18$\pm$0.24 | 19.38$\pm$0.08 | 24.61$\pm$0.17 | 25.54$\pm$0.99 |  |
| 0 | 0 | 1.251$\pm$0.040 | 30.95$\pm$0.28 | 19.37$\pm$0.06 | 24.37$\pm$0.17 | 27.39$\pm$1.14 |  |  |
| Adj.-Match. | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.882$\pm$0.058 | 30.62$\pm$0.30 | 19.50$\pm$0.09 | 24.95$\pm$0.28 | 34.50$\pm$1.33 |  |
| $\lambda=12500$ | 0 | 0.778$\pm$0.050 | 31.65$\pm$0.19 | 19.76$\pm$0.08 | 24.49$\pm$0.27 | 37.24$\pm$1.57 |  |  |
| $100$ | None (Base) | N/A | $\sqrt{2\eta_{t}}$ | $-$0.881$\pm$0.041 | 27.83$\pm$0.19 | 18.10$\pm$0.07 | 18.43$\pm$0.17 | 57.21$\pm$1.50 |
| 0 | $-$0.881$\pm$0.036 | 28.65$\pm$0.18 | 18.22$\pm$0.06 | 18.20$\pm$0.17 | 57.73$\pm$1.68 |  |  |  |
| DRaFT-1 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.343$\pm$0.040 | 30.64$\pm$0.20 | 19.38$\pm$0.08 | 24.37$\pm$0.17 | 25.51$\pm$1.10 |  |
| 0 | 0 | 1.239$\pm$0.037 | 30.74$\pm$0.28 | 19.33$\pm$0.06 | 24.24$\pm$0.17 | 28.70$\pm$1.11 |  |  |
| Adj.-Match. | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.892$\pm$0.044 | 31.23$\pm$0.23 | 19.65$\pm$0.08 | 24.92$\pm$0.23 | 35.13$\pm$1.40 |  |
| $\lambda=12500$ | 0 | 0.779$\pm$0.048 | 31.64$\pm$0.17 | 19.76$\pm$0.08 | 24.57$\pm$0.25 | 38.26$\pm$1.65 |  |  |
| $200$ | None (Base) | N/A | $\sqrt{2\eta_{t}}$ | $-$0.848$\pm$0.048 | 28.37$\pm$0.21 | 18.27$\pm$0.08 | 18.56$\pm$0.19 | 58.00$\pm$1.58 |
| 0 | $-$0.871$\pm$0.036 | 28.50$\pm$0.18 | 18.23$\pm$0.06 | 18.25$\pm$0.14 | 57.84$\pm$1.60 |  |  |  |
| DRaFT-1 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 1.331$\pm$0.044 | 30.69$\pm$0.23 | 19.36$\pm$0.07 | 24.21$\pm$0.17 | 26.41$\pm$1.18 |  |
| 0 | 0 | 1.222$\pm$0.042 | 30.77$\pm$0.27 | 19.32$\pm$0.06 | 24.18$\pm$0.16 | 29.09$\pm$1.07 |  |  |
| Adj.-Match. | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.869$\pm$0.062 | 31.33$\pm$0.21 | 19.68$\pm$0.09 | 24.81$\pm$0.30 | 35.90$\pm$1.55 |  |
| $\lambda=12500$ | 0 | 0.766$\pm$0.050 | 31.61$\pm$0.16 | 19.75$\pm$0.08 | 24.52$\pm$0.24 | 38.60$\pm$1.38 |  |  |

**Table 8: Additional metrics for various fine-tuning methods for text-to-image generation, which complement the ones in [Table 2](#S7.T2) (both tables correspond to the same runs). The second and third columns show the noise schedules $\sigma(t)$ used for fine-tuning and for inference: $\sigma(t)=\sqrt{2\eta_{t}}$ corresponds to Memoryless Flow Matching, and $\sigma(t)=0$ to the Flow Matching ODE ([3](#S2.E3)).**
| Fine-tun. | Fine-tun. | Generat. | ImageReward$\,\uparrow$ | ClipScore$\,\uparrow$ | PickScore$\,\uparrow$ | HPS v2$\,\uparrow$ | DreamSim | Runtime/ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| loss | $\sigma(t)$ | $\sigma(t)$ | diversity$\,\uparrow$ | $\#$iter. |  |  |  |  |
| ReFL | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.459$\pm$0.096 | 28.46$\pm$0.25 | 18.77$\pm$0.09 | 22.54$\pm$0.17 | 37.51$\pm$3.50 | 43k$\pm$2.7k |
| 0 | 0 | 0.330$\pm$0.114 | 29.63$\pm$0.61 | 19.08$\pm$0.18 | 22.46$\pm$0.77 | 39.51$\pm$1.30 | / 1500 |  |
| DRaFT-1 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.913$\pm$0.068 | 29.80$\pm$0.22 | 19.16$\pm$0.06 | 23.63$\pm$0.16 | 35.21$\pm$1.93 | 35k$\pm$1.5k |
| 0 | 0 | 0.626$\pm$0.195 | 30.48$\pm$0.32 | 18.91$\pm$0.34 | 21.92$\pm$1.63 | 38.52$\pm$2.01 | / 1000 |  |
| Draft-40 | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | $-$1.427$\pm$0.267 | 23.39$\pm$1.72 | 17.24$\pm$0.45 | 15.72$\pm$1.80 | 41.98$\pm$2.14 | 49k$\pm$1.4k |
| 0 | 0 | $-$0.097$\pm$0.052 | 29.12$\pm$0.41 | 18.97$\pm$0.14 | 21.93$\pm$0.20 | 46.35$\pm$1.34 | / 500 |  |
| Adj.-Match. | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.107$\pm$0.046 | 29.37$\pm$0.25 | 19.05$\pm$0.07 | 22.79$\pm$0.20 | 46.38$\pm$1.36 |  |
| $\lambda=1000$ | 0 | 0.051$\pm$0.044 | 30.58$\pm$0.17 | 19.31$\pm$0.07 | 21.93$\pm$0.23 | 48.12$\pm$1.56 |  |  |
| Adj.-Match. | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.199$\pm$0.068 | 29.27$\pm$0.21 | 19.07$\pm$0.10 | 22.98$\pm$0.30 | 45.03$\pm$1.61 | 39k$\pm$0.5k |
| $\lambda=2500$ | 0 | 0.106$\pm$0.067 | 30.43$\pm$0.24 | 19.32$\pm$0.11 | 22.16$\pm$0.33 | 47.61$\pm$1.49 | / 250 |  |
| Adj.-Match. | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | 0.299$\pm$0.095 | 29.61$\pm$0.37 | 19.26$\pm$0.14 | 23.67$\pm$0.27 | 43.36$\pm$1.93 |  |
| $\lambda=12500$ | 0 | 0.224$\pm$0.051 | 30.70$\pm$0.23 | 19.52$\pm$0.11 | 22.93$\pm$0.21 | 44.62$\pm$1.79 |  |  |
| Cont. Adj. | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | $-$0.910$\pm$0.116 | 26.29$\pm$0.44 | 18.06$\pm$0.16 | 18.86$\pm$0.88 | 51.60$\pm$1.97 | 51k$\pm$0.3k |
| $\lambda=12500$ | 0 | $-$0.681$\pm$0.051 | 28.50$\pm$0.19 | 18.69$\pm$0.11 | 19.90$\pm$0.50 | 50.87$\pm$1.52 | / 250 |  |
| Disc. Adj. | $\sqrt{2\eta_{t}}$ | $\sqrt{2\eta_{t}}$ | $-$0.978$\pm$0.123 | 26.68$\pm$0.76 | 18.51$\pm$0.11 | 18.53$\pm$0.28 | 55.95$\pm$1.70 | 38k$\pm$0.4k |
| $\lambda=12500$ | 0 | $-$0.791$\pm$0.065 | 28.66$\pm$0.33 | 18.51$\pm$0.11 | 18.53$\pm$0.28 | 54.78$\pm$2.00 | / 250 |  |

Figure: Text prompt: “Man sitting on sofa at home in front of fireplace and using laptop computer, rear view”
Refer to caption: /html/2409.08861/assets/figs/samples/cfg0.0_prompt_37_image_3.jpg

Figure: Base Flow Matching model
Refer to caption: /html/2409.08861/assets/figs/samples_appendix_3/base_cfg_2_ode_prompt_6_image_1.jpg

Figure: Base Flow Matching model
Refer to caption: /html/2409.08861/assets/figs/samples_appendix_4/base_cfg_2_ode_prompt_12_image_0.jpg

Figure: Figure 9: Generated samples with CFG=0.0 and $\sigma(t)=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 bubbles. 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”.
Refer to caption: /html/2409.08861/assets/figs/samples_appendix_5/none_ode/prompt_606.jpg

Figure: Figure 10: Generated samples with CFG=0.0 and $\sigma(t)=\sqrt{2\eta_{t}}$ across seven selected prompts. Each row corresponds to a different finetuning algorithm. The prompts are the same as in [Figure 9](#S9.F9).
Refer to caption: /html/2409.08861/assets/figs/samples_appendix_5/none_sde/prompt_606.jpg

## 10 Results on DDIM and Flow Matching

### 10.1 The continuous-time limit of DDIM

The DDIM inference update (Song et al., , Eq. 12) is

$$ $\textstyle x_{k+1}=\sqrt{\bar{\alpha}_{k+1}}\big{(}\frac{x_{k}-\sqrt{1-\bar{\alpha}_{k}}\epsilon(x_{k},k)}{\sqrt{\bar{\alpha}_{k}}}\big{)}+\sqrt{1-\bar{\alpha}_{k+1}-\sigma_{k}^{2}}\epsilon(x_{k},k)+\sigma_{k}\epsilon_{k},\qquad x_{K}\sim N(0,I).$ (44) $$

If we let
$\Delta\bar{\alpha}_{k}=\bar{\alpha}_{k+1}-\bar{\alpha}_{k}$,
we have that

$$ $\textstyle\sqrt{\frac{\bar{\alpha}_{k+1}}{\bar{\alpha}_{k}}}=\sqrt{\frac{\bar{\alpha}_{k}+\bar{\alpha}_{k+1}-\bar{\alpha}_{k}}{\bar{\alpha}_{k}}}=\sqrt{1+\frac{\bar{\alpha}_{k+1}-\bar{\alpha}_{k}}{\bar{\alpha}_{k}}}=\sqrt{1+\frac{\Delta\bar{\alpha}_{k}}{\bar{\alpha}_{k}}}\approx 1+\frac{\Delta\bar{\alpha}_{k}}{2\bar{\alpha}_{k}},$ (45) $$

where we used the first-order Taylor approximation of $\sqrt{1+x}$. And

$$ $\textstyle\begin{split}\textstyle&\textstyle-\sqrt{\frac{\bar{\alpha}_{k+1}}{\bar{\alpha}_{k}}(1-\bar{\alpha}_{k})}+\sqrt{1-\bar{\alpha}_{k+1}-\sigma_{k}^{2}}=-\sqrt{\big{(}1+\frac{\Delta\bar{\alpha}_{k}}{\bar{\alpha}_{k}}\big{)}(1-\bar{\alpha}_{k})}+\sqrt{1-\bar{\alpha}_{k+1}-\sigma_{k}^{2}}\\ \textstyle&\textstyle=-\sqrt{1+\frac{\Delta\bar{\alpha}_{k}}{\bar{\alpha}_{k}}-\bar{\alpha}_{k}-\Delta\bar{\alpha}_{k}}+\sqrt{1-\bar{\alpha}_{k+1}-\sigma_{k}^{2}}=-\sqrt{1-\bar{\alpha}_{k+1}+\frac{\Delta\bar{\alpha}_{k}}{\bar{\alpha}_{k}}}+\sqrt{1-\bar{\alpha}_{k+1}-\sigma_{k}^{2}}\\ \textstyle&\textstyle=\sqrt{1-\bar{\alpha}_{k+1}}\big{(}-\sqrt{1+\frac{\Delta\bar{\alpha}_{k}}{\bar{\alpha}_{k}(1-\bar{\alpha}_{k+1})}}+\sqrt{1-\frac{\sigma_{k}^{2}}{1-\bar{\alpha}_{k+1}}}\big{)}\approx\sqrt{1-\bar{\alpha}_{k+1}}\big{(}-\big{(}1+\frac{\Delta\bar{\alpha}_{k}}{2\bar{\alpha}_{k}(1-\bar{\alpha}_{k+1})}\big{)}+1-\frac{\sigma_{k}^{2}}{2(1-\bar{\alpha}_{k+1})}\big{)}\\ \textstyle&\textstyle=-\big{(}\frac{\Delta\bar{\alpha}_{k}}{2\bar{\alpha}_{k}}+\frac{\sigma_{k}^{2}}{2}\big{)}\frac{1}{\sqrt{1-\bar{\alpha}_{k+1}}},\end{split}$ (46) $$

where we used the same first-order Taylor approximation. Thus, up to first-order approximations, ([44](#S10.E44)) is equivalent to

$$ $\textstyle x_{k-1}=\big{(}1+\frac{\Delta\bar{\alpha}_{k}}{2\bar{\alpha}_{k}}\big{)}x_{k}-\big{(}\frac{\Delta\bar{\alpha}_{k}}{2\bar{\alpha}_{k}}+\frac{\sigma_{k}^{2}}{2}\big{)}\frac{\epsilon(x_{k},k)}{\sqrt{1-\bar{\alpha}_{k+1}}}+\sigma_{k}\epsilon_{k},\qquad x_{K}\sim N(0,I).$ (47) $$

If we modify our notation slightly, we can rewrite this as

$$ $\textstyle X_{(k+1)h}=\big{(}1-\frac{h\dot{\bar{\alpha}}_{kh}}{2\bar{\alpha}_{kh}}\big{)}X_{kh}+\big{(}\frac{h\dot{\bar{\alpha}}_{kh}}{2\bar{\alpha}_{kh}}-\frac{h\sigma(kh)^{2}}{2}\big{)}\frac{\epsilon(X_{kh},kh)}{\sqrt{1-\bar{\alpha}_{kh}}}+\sqrt{h}\sigma(kh)\epsilon_{k},\qquad X_{0}\sim N(0,I).$ (48) $$

To go from ([47](#S10.E47)) to ([48](#S10.E48)), we introduced a continuous time variable and a stepsize $h=1/K$, and we regard the increment $h\bar{\alpha}_{k}$ as approximately equal to $h$ times the derivative of $\bar{\alpha}$. We also
identified $\sigma_{k}$ with $\sqrt{h}\sigma(kh)$, where $\sigma(kh)$ plays the role of a diffusion coefficient. Note that equation ([48](#S10.E48)) can be reverse-engineered as the Euler-Maruyama discretization of the SDE

$$ $\textstyle\mathrm{d}X_{t}=\big{(}-\frac{\dot{\bar{\alpha}}_{t}}{2\bar{\alpha}_{t}}+\big{(}\frac{\dot{\bar{\alpha}}_{t}}{2\bar{\alpha}_{t}}-\frac{\sigma(t)^{2}}{2}\big{)}\frac{\epsilon(X_{t},t)}{\sqrt{1-\bar{\alpha}_{t}}}\big{)}\mathrm{d}t+\sigma(t)\mathrm{d}B_{t},\qquad X_{0}\sim N(0,I).$ (49) $$

### 10.2 Forward and backward stochastic differential equations

Let ${(\kappa_{t})}_{t\in[0,1]}$ and ${(\eta_{t})}_{t\in[0,1]}$ such that

$$ $\textstyle\forall t\in[0,1],\quad\eta_{t}\geq 0,\qquad\qquad\int_{0}^{1}\kappa_{1-s}\,\mathrm{d}s=+\infty,\qquad\qquad 2\int_{0}^{1}\eta_{1-t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}=1.$ (50) $$

As shown in [Table 1](#S4.T1), DDIM corresponds to $\kappa_{t}=\frac{\dot{\alpha}_{t}}{2\alpha_{t}}$, $\eta_{t}=\frac{\dot{\alpha}_{t}}{2\alpha_{t}}$, and Flow Matching corresponds to $\kappa_{t}=\frac{\dot{\alpha}_{t}}{\alpha_{t}}$, $\eta_{t}=\beta_{t}\big{(}\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t}\big{)}$.

###### Lemma 1 (DDIM and Flow Matching fulfill the conditions ( 50 )) .

The choices of $(\kappa_{t})_{t\in[0,1]}$ and $(\eta_{t})_{t\in[0,1]}$ for DDIM and Flow Matching fulfill the conditions ([50](#S10.E50)). For DDIM, we have that

$$ $\textstyle\begin{split}\textstyle\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s=-\frac{1}{2}\log\alpha_{1-t}&\textstyle\implies\int_{0}^{1}\kappa_{1-s}\,\mathrm{d}s=+\infty,\\ \textstyle 2\int_{0}^{t}\eta_{t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}=1-\alpha_{1-t}&\textstyle\implies 2\int_{0}^{1}\eta_{t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}=1.\end{split}$ (51) $$

For Flow Matching,

$$ $\textstyle\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s=-\log\alpha_{1-t}$ $\textstyle\implies\int_{0}^{1}\kappa_{1-s}\,\mathrm{d}s=+\infty,$ (52) $\textstyle 2\int_{0}^{t}\eta_{t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}=\beta^{2}_{1-t}$ $\textstyle\implies 2\int_{0}^{1}\eta_{t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}=1.$ (53) $$

###### Lemma 1 (DDIM and Flow Matching fulfill the conditions ( 50 )) .

##### Forward and backward SDEs

Consider the forward and backward SDEs

$$ $\textstyle\mathrm{d}\vec{X}_{t}$ $\textstyle=-\kappa_{1-t}\vec{X}_{t}\,\mathrm{d}t+\sqrt{2\eta_{1-t}}\,\mathrm{d}B_{t},\qquad\vec{X}_{0}\sim p_{\mathrm{data}},$ (54) $\textstyle\mathrm{d}X_{t}$ $\textstyle=\big{(}\kappa_{t}X_{t}+2\eta_{t}\mathfrak{s}(X_{t},t)\big{)}\mathrm{d}t+\sqrt{2\eta_{t}}\,\mathrm{d}B_{t},\qquad X_{0}\sim N(0,I),$ (55) $$

where we let $\vec{p}_{t}$ be the density of $\vec{X}_{t}$, and we define the score function as $\mathfrak{s}(x,t):=\nabla\log\vec{p}_{1-t}(x)$. Similarly, we let $p_{t}$ be the density of $X_{t}$. $\vec{p}_{t}$ and $p_{t}$ solve the Fokker-Planck equations:

$$ $\textstyle\partial_{t}\vec{p}_{t}$ $\textstyle=\nabla\cdot\big{(}\kappa_{1-t}x\vec{p}_{t}\big{)}+\eta_{1-t}\Delta\vec{p}_{t},\qquad\vec{p}_{0}=p_{\mathrm{data}},$ (56) $\textstyle\partial_{t}p_{t}$ $\textstyle=\nabla\cdot\big{(}\big{(}-\kappa_{t}x-2\eta_{t}\nabla\log\vec{p}_{1-t}(X_{t})\big{)}p_{t}\big{)}+\eta_{t}\Delta p_{t},\qquad p_{0}=N(0,I).$ (57) $$

###### Lemma 2 (Solution of the forward SDE) .

Let $(\kappa_{t})_{t\geq 0}$, $(\eta_{t})_{t\geq 0}$ with $\eta_{t}\geq 0$, and $(\xi_{t})_{t\geq 0}$ be arbitrary. The solution $\vec{X}_{t}$ of the SDE

$$ $\textstyle\mathrm{d}\vec{X}_{t}$ $\textstyle=\big{(}-\kappa_{1-t}\vec{X}_{t}+\xi_{t}\big{)}\,\mathrm{d}t+\sqrt{2\eta_{1-t}}\,\mathrm{d}B_{t},\qquad\vec{X}_{0}\sim p_{\mathrm{data}}$ (58) $$

is

$$ $\textstyle\vec{X}_{t}=\vec{X}_{0}\exp\big{(}-\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}+\int_{0}^{t}\exp\big{(}-\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\xi_{1-t^{\prime}}\,\mathrm{d}t^{\prime}+\int_{0}^{t}\sqrt{2\eta_{1-t^{\prime}}}\exp\big{(}-\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}B_{t^{\prime}},$ (59) $$

which has the same distribution as the random variable

$$ $\textstyle\begin{split}\textstyle\hat{X}_{t}&\textstyle=\vec{X}_{0}\exp\big{(}-\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}+\int_{0}^{t}\exp\big{(}-\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\xi_{1-t^{\prime}}\,\mathrm{d}t^{\prime}+\sqrt{2\int_{0}^{t}\eta_{1-t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}}\epsilon,\\ \textstyle\epsilon&\textstyle\sim N(0,I).\end{split}$ (60) $$

Applying [Lemma 2](#Thmlemma2) with $\xi_{t}\equiv 0$, we obtain that $\vec{p}_{1}$ is also the distribution of

$$ $\textstyle\hat{X}_{1}=\vec{X}_{0}\exp\big{(}-\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}+\sqrt{2\int_{0}^{t}\eta_{1-t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}}\epsilon=\epsilon,$ (61) $$

where $\epsilon\sim N(0,I)$. The third equality in ([61](#S10.E61)) holds by ([50](#S10.E50)). Hence we obtain that $\vec{p}_{1}=N(0,I)$. Note also that

$$ $\textstyle\partial_{t}\vec{p}_{1-t}$ $\textstyle=-\nabla\cdot\big{(}\kappa_{t}x\vec{p}_{1-t}\big{)}-\eta_{t}\Delta\vec{p}_{1-t}=-\nabla\cdot\big{(}\big{(}-\kappa_{t}x-2\eta_{t}\nabla\log\vec{p}_{1-t}(x)\big{)}\vec{p}_{1-t}\big{)}+\eta_{t}\Delta\vec{p}_{1-t}$ (62) $$

Thus, $\vec{p}_{1-t}$ is a solution of the backward Fokker-Planck equation ([57](#S10.E57)), which proves the following:

###### Proposition 3 (Equality of marginal distributions) .

For any time $t\in[0,1]$, the densities of the solutions $\vec{X}_{t}$, $X_{t}$ of the forward and backward SDEs are equal up to a time flip: $p_{t}=\vec{p}_{1-t}$.

###### Lemma 2 (Solution of the forward SDE) .

###### Proposition 3 (Equality of marginal distributions) .

##### Forward and backward SDEs with arbitrary noise schedule

Next, we look at the following pair of forward-backward SDEs:

$$ $\textstyle\mathrm{d}\vec{X}_{t}$ $\textstyle=\big{(}-\kappa_{1-t}\vec{X}_{t}+\big{(}\frac{\sigma(1-t)^{2}}{2}-\eta_{1-t}\big{)}\mathfrak{s}(\vec{X}_{t},1-t)\big{)}\,\mathrm{d}t+\sigma(1-t)\,\mathrm{d}B_{t},\qquad\vec{X}_{0}\sim p_{\mathrm{data}},$ (63) $\textstyle\mathrm{d}X_{t}$ $\textstyle=\big{(}\kappa_{t}X_{t}+\big{(}\frac{\sigma(t)^{2}}{2}+\eta_{t}\big{)}\mathfrak{s}(X_{t},t)\big{)}\mathrm{d}t+\sigma(t)\,\mathrm{d}B_{t},\qquad X_{0}\sim N(0,I),$ (64) $$

Here, the score function $\mathfrak{s}$ is the same vector field as in ([64](#S10.E64)).
Remark that equations ([54](#S10.E54))-([55](#S10.E55)) are a particular case of ([63](#S10.E63))-([64](#S10.E64)) for which $\sigma(t)=\sqrt{2\eta_{t}}$. The Fokker-Planck equations for ([63](#S10.E63))-([64](#S10.E64)) are:

$$ $\textstyle\partial_{t}\vec{p}_{t}$ $\textstyle=\nabla\cdot\big{(}\big{(}\kappa_{1-t}x+\big{(}-\frac{\sigma(1-t)^{2}}{2}+\eta_{1-t}\big{)}\mathfrak{s}(X_{t},t)\big{)}\vec{p}_{t}\big{)}+\eta_{1-t}\Delta\vec{p}_{t},\qquad\vec{p}_{0}=p_{\mathrm{data}},$ (65) $\textstyle\partial_{t}p_{t}$ $\textstyle=\nabla\cdot\big{(}\big{(}-\kappa_{t}x-\big{(}\frac{\sigma(t)^{2}}{2}+\eta_{t}\big{)}\mathfrak{s}(X_{t},t)\big{)}p_{t}\big{)}+\frac{\sigma(t)^{2}}{2}\Delta p_{t},\qquad p_{0}=N(0,I).$ (66) $$

It is straight-forward to see that for any $\sigma$, the solutions $\vec{p}_{t}$ and $p_{t}$ of ([65](#S10.E65))-([66](#S10.E66)) are also solutions of ([56](#S10.E56))-([57](#S10.E57)). Hence, the marginals $\vec{X}_{t}$ and $X_{t}$ are equally distributed for all noise schedules $\sigma$, and they are equal to each other up to a time flip.

##### Equality of distributions over trajectories

The result in [Proposition 3](#Thmproposition3) can be made even stronger:

###### Proposition 4 (Equality of distributions over trajectories) .

Let $\vec{\bm{X}}$, $\bm{X}$ be the solutions of the SDEs ([63](#S10.E63))-([64](#S10.E64)) with arbitrary noise schedule. For any sequence of times $(t_{i})_{0\leq i\leq I}$, the joint distribution of $(\vec{X}_{t_{i}})_{0\leq i\leq I}$ is equal to the joint distribution of $(X_{1-t_{i}})_{0\leq i\leq I}$, or equivalently, that the probability measures $\vec{\mathbb{P}}$, $\mathbb{P}$ of the forward and backward processes $\vec{\bm{X}}$, $\bm{X}$ are equal, up to a flip in the time direction.

This result states that sampling trajectories from the backward process is equivalent to sampling them from the forward process and then flipping their order.

###### Proposition 4 (Equality of distributions over trajectories) .

#### 10.2.1 Proof of Lemma 1

As shown in [Table 1](#S4.T1), DDIM corresponds to $\kappa_{t}=\frac{\dot{\alpha}_{t}}{2\alpha_{t}}$, $\eta_{t}=\frac{\dot{\alpha}_{t}}{2\alpha_{t}}$. Thus, $\eta_{t}\geq 0$ because $\alpha_{t}$ is increasing, and

$$ $\textstyle\begin{split}\textstyle&\textstyle\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s=\int_{0}^{t}\frac{\dot{\alpha}_{1-s}}{2\alpha_{1-s}}\,\mathrm{d}s=-\frac{1}{2}\int_{0}^{t}\partial_{s}\log\alpha_{1-s}\,\mathrm{d}s=-\frac{1}{2}(\log\alpha_{1-t}-\log\alpha_{1})=-\frac{1}{2}\log\alpha_{1-t},\\ \textstyle&\textstyle\implies\int_{0}^{1}\kappa_{1-s}\,\mathrm{d}s=-\frac{1}{2}\log\alpha_{0}=+\infty\end{split}$ (67) $\textstyle\begin{split}\textstyle&\textstyle 2\int_{0}^{t}\eta_{t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}=\int_{0}^{t}\frac{\dot{\alpha}_{1-t^{\prime}}}{\alpha_{1-t^{\prime}}}\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\dot{\alpha}_{1-s}}{\alpha_{1-s}}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}\\ \textstyle&\textstyle=\int_{0}^{t}\frac{\dot{\alpha}_{1-t^{\prime}}}{\alpha_{1-t^{\prime}}}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\,\mathrm{d}t^{\prime}=\alpha_{1-t}\int_{0}^{t}\partial_{t^{\prime}}\big{(}\frac{1}{\alpha_{1-t^{\prime}}}\big{)}\,\mathrm{d}t^{\prime}=\alpha_{1-t}\big{(}\frac{1}{\alpha_{1-t}}-\frac{1}{\alpha_{1}}\big{)}=1-\alpha_{1-t},\\ \textstyle&\textstyle\implies 2\int_{0}^{1}\eta_{t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}=1-\alpha_{0}=1.\end{split}$ (68) $$

where we used that $\alpha_{1}=1$ and $\alpha_{0}=0$. And Flow Matching corresponds to $\kappa_{t}=\frac{\dot{\alpha}_{t}}{\alpha_{t}}$, $\eta_{t}=\beta_{t}\big{(}\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t}\big{)}$. We have that $\eta_{t}\geq 0$ because $\alpha_{t}$ is increasing and $\beta_{t}$ is decreasing, and

$$ $\textstyle\begin{split}\textstyle&\textstyle\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s=\int_{0}^{t}\frac{\dot{\alpha}_{1-s}}{\alpha_{1-s}}\,\mathrm{d}s=-\int_{0}^{t}\partial_{s}\log\alpha_{1-s}\,\mathrm{d}s=-(\log\alpha_{1-t}-\log\alpha_{1})=-\log\alpha_{1-t},\\ \textstyle&\textstyle\implies\int_{0}^{1}\kappa_{1-s}\,\mathrm{d}s=-\log\alpha_{0}=+\infty,\end{split}$ (69) $$

and

$$ $\textstyle\begin{split}\textstyle&\textstyle 2\int_{0}^{t}\eta_{1-t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}=2\int_{0}^{t}\beta_{1-t^{\prime}}\big{(}\frac{\dot{\alpha}_{1-t^{\prime}}}{\alpha_{1-t^{\prime}}}\beta_{1-t^{\prime}}-\dot{\beta}_{1-t^{\prime}}\big{)}\exp\big{(}-2\int_{t^{\prime}}^{t}\frac{\dot{\alpha}_{1-s}}{\alpha_{1-s}}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}\\ \textstyle&\textstyle=2\int_{0}^{t}\beta_{1-t^{\prime}}\big{(}\frac{\dot{\alpha}_{1-t^{\prime}}}{\alpha_{1-t^{\prime}}}\beta_{1-t^{\prime}}-\dot{\beta}_{1-t^{\prime}}\big{)}\big{(}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\big{)}^{2}\,\mathrm{d}t^{\prime},\end{split}$ (70) $$

To develop the right-hand side, note that by integration by parts,

$$ $\textstyle\begin{split}\textstyle&\textstyle\int_{0}^{t}\dot{\beta}_{1-t^{\prime}}\beta_{1-t^{\prime}}\big{(}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\big{)}^{2}\,\mathrm{d}t^{\prime}=-\int_{0}^{t}\partial_{t^{\prime}}\big{(}\frac{\beta_{1-t^{\prime}}^{2}}{2}\big{)}\big{(}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\big{)}^{2}\,\mathrm{d}t^{\prime}\\ \textstyle&\textstyle=-\big{[}\frac{\beta_{1-t^{\prime}}^{2}}{2}\big{(}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\big{)}^{2}\big{]}_{0}^{1}+\int_{0}^{t}\frac{\beta_{1-t^{\prime}}^{2}}{2}\partial_{t^{\prime}}\big{(}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\big{)}^{2}\,\mathrm{d}t^{\prime}=-\big{[}\frac{\beta_{1-t^{\prime}}^{2}}{2}\big{(}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\big{)}^{2}\big{]}_{0}^{t}+\int_{0}^{t}\beta_{1-t^{\prime}}^{2}\frac{\alpha_{1-t}^{2}\dot{\alpha}_{1-t^{\prime}}}{\alpha_{1-t^{\prime}}^{3}}\,\mathrm{d}t^{\prime}.\end{split}$ (71) $$

And if we plug this into the right-hand side of (LABEL:eq:second_term_FM), we obtain

$$ $\textstyle 2\int_{0}^{t}\eta_{1-t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}=\big{[}\beta_{1-t^{\prime}}^{2}\big{(}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\big{)}^{2}\big{]}_{0}^{t}=\beta_{1-t}^{2}-\beta_{1}^{2}\big{(}\frac{\alpha_{1-t}}{\alpha_{1}}\big{)}^{2}=\beta_{1-t}^{2},$ (72) $\textstyle\implies 2\int_{0}^{1}\eta_{1-t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}=\beta_{1}^{2}=1.$ (73) $$

where we used that $\beta_{1}=0$, $\alpha_{1}=1$.

#### 10.2.2 Proof of Lemma 2

We can solve this equation by variation of parameters. To simplify the notation, we replace $\kappa_{1-s}$, $\eta_{1-s}$ and $\xi_{1-s}$ by $\kappa_{s}$, $\eta_{s}$ and $\xi_{s}$. Defining $f(\vec{X}_{t},t)=\vec{X}_{t}\exp\big{(}\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}$, we get that

$$ $\textstyle\begin{split}\textstyle df(\vec{X}_{t},t)&\textstyle=\kappa_{1-t}\vec{X}_{t}\exp\big{(}\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}t+\exp\big{(}\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\mathrm{d}\vec{X}_{t}\\ \textstyle&\textstyle=\kappa_{1-t}\vec{X}_{t}\exp\big{(}\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}t+\exp\big{(}\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\big{(}(-\kappa_{1-t}\vec{X}_{t}+\xi_{1-t})\,\mathrm{d}t+\sqrt{2\eta_{1-t}}\,\mathrm{d}B_{t}\big{)}\\ \textstyle&\textstyle=\exp\big{(}\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\xi_{1-t}\,\mathrm{d}t+\sqrt{2\eta_{t}}\exp\big{(}\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}B_{t}.\end{split}$ (74) $$

Integrating from 0 to $t$, we get that

$$ $\textstyle\vec{X}_{t}\exp\big{(}\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}=\vec{X}_{0}+\int_{0}^{t}\exp\big{(}\int_{0}^{t^{\prime}}\kappa_{1-s}\,\mathrm{d}s\big{)}\xi_{1-t^{\prime}}\,\mathrm{d}t^{\prime}+\int_{0}^{t}\sqrt{2\eta_{1-t^{\prime}}}\exp\big{(}\int_{0}^{t^{\prime}}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}B_{t^{\prime}},$ (75) $\textstyle\iff\vec{X}_{t}=\vec{X}_{0}\exp\big{(}-\int_{0}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}+\int_{0}^{t}\exp\big{(}-\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\xi_{1-t^{\prime}}\,\mathrm{d}t^{\prime}+\int_{0}^{t}\sqrt{2\eta_{1-t^{\prime}}}\exp\big{(}-\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}B_{t^{\prime}}.$ (76) $$

Since

$$ $\textstyle\mathbb{E}\big{[}\big{(}\int_{0}^{t}\sqrt{2\eta_{1-t^{\prime}}}\exp\big{(}-\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}B_{t^{\prime}}\big{)}^{2}\big{]}=2\int_{0}^{t}\eta_{1-t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime},$ (77) $$

we obtain that $\int_{0}^{t}\sqrt{2\eta_{1-t^{\prime}}}\exp\big{(}-\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}B_{t^{\prime}}$ has the same distribution as $\sqrt{2\int_{0}^{t}\eta_{1-t^{\prime}}\exp\big{(}-2\int_{t^{\prime}}^{t}\kappa_{1-s}\,\mathrm{d}s\big{)}\,\mathrm{d}t^{\prime}}\epsilon$, where $\epsilon\sim N(0,1)$.

#### 10.2.3 Proof of Proposition 4

This is a result that has been used by previous works, e.g. (De Bortoli et al., , 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 a derivation which is still in discrete time, and hence not completely formal, but that corrects the gaps in the proof of De Bortoli et al. ().

We introduce the short-hand

$$ $\textstyle\vec{b}(x,t)$ $\textstyle=-\kappa_{1-t}x+\big{(}\frac{\sigma(1-t)^{2}}{2}-\eta_{1-t}\big{)}\mathfrak{s}(x,1-t),$ (78) $\textstyle b(x,t)$ $\textstyle=\kappa_{t}X_{t}+\big{(}\frac{\sigma(t)^{2}}{2}+\eta_{t}\big{)}\mathfrak{s}(X_{t},t),$ (79) $\textstyle\vec{\sigma}(t)$ $\textstyle=\sigma(1-t).$ (80) $$

Remark that $b(x,t)=-\vec{b}(x,1-t)+\sigma(t)^{2}\mathfrak{s}(X_{t},t)$.

Suppose that we discretize the forward process $\vec{X}$ using $K+1$ equispaced timesteps:

$$ $\textstyle x_{k+1}=x_{k}+h\vec{b}(x_{k},kh)+\sqrt{h}\vec{\sigma}(kh)\epsilon_{k},\qquad\text{with}\ \epsilon_{k}\sim N(0,1).$ (81) $$

It is important to remark that $x_{k+1}-x_{k}=O(h^{1/2})$. Throughout the proof we will keep track of all terms up to linear order in $h$, while neglecting terms of order $O(h^{3/2})$ and higher.
The distribution of the discretized forward process is:

$$ $\textstyle\vec{p}(x_{0:K})=\vec{p}_{0}(x_{0})\prod_{k=0}^{K-1}\vec{p}_{k+1|k}(x_{k+1}|x_{k}),\qquad\text{where}\qquad\vec{p}_{k+1|k}(x_{k+1}|x_{k})=\frac{\exp\big{(}-\frac{\|x_{k+1}-x_{k}-h\vec{b}(x_{k},kh)\|^{2}}{2h\vec{\sigma}(kh)^{2}}\big{)}}{(2\pi h\vec{\sigma}(kh)^{2})^{d/2}}$ (82) $$

Using telescoping products, we have that

$$ $\textstyle\begin{split}\textstyle\vec{p}(x_{0:K})&\textstyle=\vec{p}_{K}(x_{K})\prod_{k=0}^{K-1}\vec{p}_{k+1|k}(x_{k+1}|x_{k})\frac{\vec{p}_{k}(x_{k})}{\vec{p}_{k+1}(x_{k+1})}\\ \textstyle&\textstyle=\vec{p}_{K}(x_{K})\prod_{k=0}^{K-1}\vec{p}_{k+1|k}(x_{k+1}|x_{k})\exp\big{(}\log(\vec{p}_{k}(x_{k}))-\log(\vec{p}_{k+1}(x_{k+1}))\big{)}\end{split}$ (83) $$

We can use a discrete time version of Ito’s lemma:

$$ $\textstyle\log\vec{p}(x_{k+1},(k+1)h)$ $\textstyle\approx\log\vec{p}(x_{k},kh)+h\big{(}\partial_{t}\log\vec{p}(x_{k},kh)+\frac{\vec{\sigma}(kh)^{2}}{2}\Delta\log\vec{p}(x_{k},kh)\big{)}$ (84) $\textstyle\qquad+\langle\nabla\log\vec{p}(x_{k},kh),x_{k+1}-x_{k}\rangle+O(h^{3/2}).$ (85) $$

Using equation ([81](#S10.E81)) and a Taylor approximation, observe that

$$ $\textstyle\begin{split}\textstyle&\textstyle\langle\nabla\log p(x_{k},kh),x_{k+1}-x_{k}\rangle\\ \textstyle&\textstyle=\langle\nabla\log p(x_{k+1},(k+1)h)-\nabla^{2}\log p(x_{k+1},(k+1)h)(x_{k+1}-x_{k}),x_{k+1}-x_{k}\rangle+O(h^{3/2})\\ \textstyle&\textstyle=\langle\nabla\log p(x_{k+1},(k+1)h),x_{k+1}-x_{k}\rangle\\ \textstyle&\textstyle\qquad-\langle h\vec{b}(x_{k},kh)+\sqrt{h}\vec{\sigma}(kh)\epsilon_{k},\nabla^{2}\log p(x_{k+1},(k+1)h)\big{(}h\vec{b}(x_{k},kh)+\sqrt{h}\vec{\sigma}(kh)\epsilon_{k}\big{)}\rangle+O(h^{3/2})\\ \textstyle&\textstyle=\langle\nabla\log p(x_{k+1},(k+1)h),x_{k+1}-x_{k}\rangle-h\vec{\sigma}(kh)^{2}\Delta\log p(x_{k+1},(k+1)h)+O(h^{3/2}).\end{split}$ (86) $$

And since $\vec{p}$ satisfies the Fokker-Planck equation

$$ $\textstyle\partial_{t}\vec{p}_{t}=\nabla\cdot\big{(}(-\vec{b}(x,t)+\frac{\vec{\sigma}(t)^{2}}{2}\nabla\log\vec{p}_{t}(x))\vec{p}_{t}\big{)},$ (87) $$

we have that

$$ $\textstyle\begin{split}\textstyle\partial_{t}\log\vec{p}_{t}&\textstyle=\frac{\partial_{t}\vec{p}_{t}}{\vec{p}_{t}}=\frac{\nabla\cdot\big{(}(-\vec{b}(x,t)+\frac{\vec{\sigma}(t)^{2}}{2}\nabla\log\vec{p}_{t}(x))\vec{p}_{t}\big{)}}{\vec{p}_{t}}\\ \textstyle&\textstyle=-\nabla\cdot\vec{b}(x,t)+\frac{\vec{\sigma}(t)^{2}}{2}\Delta\log\vec{p}_{t}(x)+\langle-\vec{b}(x,t)+\frac{\vec{\sigma}(t)^{2}}{2}\nabla\log\vec{p}_{t}(x),\nabla\log\vec{p}_{t}(x)\rangle.\end{split}$ (88) $$

Hence,

$$ $\textstyle\begin{split}\textstyle&\textstyle\partial_{t}\log p(x_{k},kh)=\partial_{t}\log p(x_{k+1},(k+1)h)+O(h^{1/2})\\ \textstyle&\textstyle=-\nabla\cdot\vec{b}(x_{k+1},(k+1)h)+\frac{\vec{\sigma}((k+1)h)^{2}}{2}\Delta\log\vec{p}(x_{k+1},(k+1)h)\\ \textstyle&\textstyle\qquad+\langle-\vec{b}(x_{k+1},(k+1)h)+\frac{\vec{\sigma}((k+1)h)^{2}}{2}\nabla\log\vec{p}(x_{k+1},(k+1)h),\nabla\log\vec{p}(x_{k+1},(k+1)h)\rangle+O(h^{1/2}).\end{split}$ (89) $$

If we plug (LABEL:eq:nabla_log_diff_x) and (LABEL:eq:partial_t) into ([84](#S10.E84)), we obtain

$$ $\textstyle\begin{split}\textstyle&\textstyle\log p(x_{k+1},(k+1)h)-\log p(x_{k},kh)\\ \textstyle&\textstyle=h\big{(}\!-\!\nabla\cdot\vec{b}(x_{k+1},(k+1)h)\!+\!\langle-\vec{b}(x_{k+1},(k+1)h)\!+\!\frac{\vec{\sigma}((k+1)h)^{2}}{2}\nabla\log\vec{p}(x_{k+1},(k+1)h),\nabla\log\vec{p}(x_{k+1},(k+1)h)\rangle\big{)}\\ \textstyle&\textstyle\qquad+\langle\nabla\log p(x_{k+1},(k+1)h),x_{k+1}-x_{k}\rangle+O(h^{3/2})\\ \textstyle&\textstyle=\frac{\langle 2h\vec{\sigma}(kh)^{2}\nabla\log p(x_{k+1},(k+1)h),x_{k+1}-x_{k}-h\vec{b}(x_{k+1},(k+1)h)\rangle}{2h\vec{\sigma}(kh)^{2}}\\ \textstyle&\textstyle\qquad+h\big{(}-\nabla\cdot\vec{b}(x_{k+1},(k+1)h)+\frac{\vec{\sigma}((k+1)h)^{2}}{2}\|\nabla\log\vec{p}(x_{k+1},(k+1)h)\|^{2}\big{)}+O(h^{3/2}).\end{split}$ (90) $$

Applying a discrete time version of Ito’s lemma again, we have that

$$ $\textstyle\begin{split}\textstyle\vec{b}(x_{k},kh)&\textstyle=\vec{b}(x_{k+1},(k+1)h)-h\big{(}\partial_{t}\vec{b}(x_{k+1},(k+1)h)+\frac{\vec{\sigma}((k+1)h)^{2}}{2}\Delta\vec{b}(x_{k+1},(k+1)h)\big{)}\\ \textstyle&\textstyle\qquad+\nabla\vec{b}(x_{k+1},(k+1)h)^{\top}(x_{k}-x_{k+1})+O(h^{3/2})\\ \textstyle&\textstyle=\vec{b}(x_{k+1},(k+1)h)+\nabla\vec{b}(x_{k+1},(k+1)h)^{\top}(x_{k}-x_{k+1})+O(h).\end{split}$ (91) $$

where $\Delta\vec{b}$ denotes the component-wise Laplacian of $\vec{b}$. Thus,

$$ $\textstyle\begin{split}\textstyle&\textstyle\log\vec{p}_{k+1|k}(x_{k+1}|x_{k})\\ \textstyle&\textstyle=-\frac{d}{2}\log\big{(}2\pi h\vec{\sigma}(kh)^{2}\big{)}-\frac{\|x_{k+1}-x_{k}-h\vec{b}(x_{k},kh)\|^{2}}{2h\vec{\sigma}(kh)^{2}}\\ \textstyle&\textstyle=-\frac{d}{2}\log\big{(}2\pi h\vec{\sigma}(kh)^{2}\big{)}-\frac{\|x_{k+1}-x_{k}-h(\vec{b}(x_{k+1},(k+1)h)+\nabla\vec{b}(x_{k+1},(k+1)h)^{\top}(x_{k}-x_{k+1}))\|^{2}}{2h\vec{\sigma}(kh)^{2}}+O(h^{3/2})\\ \textstyle&\textstyle=-\frac{d}{2}\log\big{(}2\pi h\vec{\sigma}(kh)^{2}\big{)}-\frac{\|x_{k+1}-x_{k}-h\vec{b}(x_{k+1},(k+1)h)\|^{2}}{2h\vec{\sigma}(kh)^{2}}+\frac{\langle x_{k+1}-x_{k},\nabla\vec{b}(x_{k+1},(k+1)h)^{\top}(x_{k}-x_{k+1})\rangle}{\vec{\sigma}(kh)^{2}}+O(h^{3/2})\\ \textstyle&\textstyle=-\frac{d}{2}\log\big{(}2\pi h\vec{\sigma}(kh)^{2}\big{)}-\frac{\|x_{k+1}-x_{k}-h\vec{b}(x_{k+1},(k+1)h)\|^{2}}{h\vec{\sigma}(kh)^{2}}-\frac{h\vec{\sigma}(kh)^{2}\langle\epsilon_{k},\nabla\vec{b}(x_{k+1},(k+1)h)^{\top}\epsilon_{k}\rangle}{\vec{\sigma}(kh)^{2}}+O(h^{3/2})\\ \textstyle&\textstyle=-\frac{d}{2}\log\big{(}2\pi h\vec{\sigma}(kh)^{2}\big{)}-\frac{\|x_{k+1}-x_{k}-h\vec{b}(x_{k+1},(k+1)h)\|^{2}}{h\vec{\sigma}(kh)^{2}}-h\Delta\vec{b}(x_{k+1},(k+1)h)+O(h^{3/2})\end{split}$ (92) $$

Combining (LABEL:eq:log_difference) and (LABEL:eq:log_transition), we obtain that

$$ $\textstyle\begin{split}\textstyle&\textstyle\log\vec{p}_{k+1|k}(x_{k+1}|x_{k})-\big{(}\log p(x_{k+1},(k+1)h)-\log p(x_{k},kh)\big{)}\\ \textstyle&\textstyle=-\frac{d}{2}\log\big{(}2\pi h\vec{\sigma}(kh)^{2}\big{)}-\frac{\|x_{k+1}-x_{k}-h\vec{b}(x_{k+1},(k+1)h)+h\vec{\sigma}(kh)^{2}\nabla\log p(x_{k+1},(k+1)h)\|^{2}}{h\vec{\sigma}(kh)^{2}}+O(h^{3/2})\\ \textstyle&\textstyle=-\frac{d}{2}\log\big{(}2\pi h\vec{\sigma}((k+1)h)^{2}\big{)}-\frac{\|x_{k+1}-x_{k}-h\vec{b}(x_{k+1},(k+1)h)+h\vec{\sigma}((k+1)h)^{2}\nabla\log p(x_{k+1},(k+1)h)\|^{2}}{h\vec{\sigma}((k+1)h)^{2}}+O(h^{3/2}).\end{split}$ (93) $$

By Bayes rule, and taking the exponential of this equation, we obtain

$$ $\textstyle\begin{split}\textstyle\vec{p}_{k+1|k}(x_{k+1}|x_{k})&\textstyle:=\vec{p}_{k+1|k}(x_{k+1}|x_{k})\frac{\vec{p}_{k}(x_{k})}{\vec{p}_{k+1}(x_{k+1})}\\ \textstyle&\textstyle=\frac{\exp\big{(}-\frac{\|x_{k}-x_{k+1}+h\vec{b}(x_{k+1},(k+1)h)-h\vec{\sigma}((k+1)h)^{2}\nabla\log p(x_{k+1},(k+1)h)\|^{2}}{2h\vec{\sigma}((k+1)h)^{2}}\big{)}}{(2\pi h\vec{\sigma}((k+1)h)^{2})^{d/2}}+O(h^{3/2}).\end{split}$ (94) $$

Up to the $O(h^{3/2})$ term, the right-hand side is the conditional Gaussian corresponding to the update

$$ $\textstyle x_{k}=x_{k+1}+h\big{(}-\vec{b}(x_{k+1},(k+1)h)+\vec{\sigma}((k+1)h)^{2}\nabla\log p(x_{k+1},(k+1)h)\big{)}+\sqrt{h}\vec{\sigma}((k+1)h)\epsilon_{k+1},\ \epsilon_{k+1}\sim N(0,I).$ (95) $$

If we define $y_{k}=x_{K-k}$, and we use that $b(x,t)=-\vec{b}(x,1-t)+\vec{\sigma}(t)^{2}\nabla\log p(x,1-t)$, we can rewrite ([95](#S10.E95)) as

$$ $\textstyle y_{K-k}$ $\textstyle=y_{K-k-1}+h\big{(}-\vec{b}(y_{K-k-1},(K-k-1)h)+\vec{\sigma}((K-k-1)h)^{2}\nabla\log p(y_{K-k-1},(K-k-1)h)\big{)}$ (96) $\textstyle\qquad\qquad\quad+\sqrt{h}\vec{\sigma}((K-k-1)h)\epsilon_{k}=y_{K-k-1}+hb(y_{K-k-1},kh)+\sqrt{h}\sigma(kh)\epsilon_{K-k-1},$ (97) $\textstyle\implies y_{k+1}=y_{k}+hb(y_{k},kh)+\sqrt{h}\sigma(kh)\epsilon_{k}.$ (98) $$

And this is the Euler-Maruyama discretization of the backward process
$\textstyle\vec{}\mkern 4.0mu$
$\textstyle X$
.
If we plug ([94](#S10.E94)) into ([83](#S10.E83)), we obtain that

$$ $\textstyle\vec{p}(x_{0:K})$ $\textstyle\approx\vec{p}_{K}(x_{K})\prod_{k=0}^{K-1}\vec{p}_{k+1|k}(x_{k+1}|x_{k}).$ (99) $$

which concludes the proof, as $\vec{p}_{K}(x_{K})$ is the initial distribution of the backward process, and $\vec{p}_{k+1|k}(x_{k+1}|x_{k})$ are its transition kernels.

### 10.3 The relationship between the noise predictor ϵ italic-ϵ \epsilon and the score function

Applying [Lemma 2](#Thmlemma2) with the choices of $(\kappa_{t})_{t\geq 0}$ and $(\eta_{t})_{t\geq 0}$ for DDIM, we obtain that $\vec{X}_{t}$ has the same distribution as

$$ $\textstyle\hat{X}_{t}=\sqrt{\alpha_{1-t}}\vec{X}_{0}+\sqrt{1-\alpha_{1-t}}\epsilon,\qquad\epsilon\sim N(0,1).$ (100) $$

Since $\vec{X}_{t}$ and $\hat{X}_{t}$ have the same distribution, predicting the noise of $\vec{X}_{t}$ is equivalent to predicting the noise of $\hat{X}_{t}$. The noise predictor $\epsilon$ can be written as:

$$ $\textstyle\epsilon(x,t):=\mathbb{E}[\epsilon|\hat{X}_{1-t}=x]=\mathbb{E}\big{[}\epsilon|\sqrt{\alpha_{t}}\vec{X}_{0}+\sqrt{1-\alpha_{t}}\epsilon=x\big{]}=\mathbb{E}\big{[}\frac{x-\sqrt{\alpha_{t}}\vec{X}_{0}}{\sqrt{1-\alpha_{t}}}|\sqrt{\alpha_{t}}\vec{X}_{0}+\sqrt{1-\alpha_{t}}\epsilon=x\big{]}$ (101) $$

And the score function $\mathfrak{s}(x,t):=\nabla\log\vec{p}_{1-t}(x)$ admits the expression

$$ $\textstyle\mathfrak{s}(x,t)$ $\textstyle:=\nabla\log\vec{p}_{1-t}(x)=\frac{\nabla\vec{p}_{1-t}(x)}{\vec{p}_{1-t}(x)}=\frac{\nabla\mathbb{E}[\vec{p}_{1-t|0}(x|\vec{X}_{0})]}{\vec{p}_{1-t}(x)}=\frac{\mathbb{E}[\nabla\log\vec{p}_{1-t|0}(x|\vec{X}_{0})\vec{p}_{1-t|0}(x|\vec{X}_{0})]}{\vec{p}_{1-t}(x)},$ (102) $$

where

$$ $\textstyle\vec{p}_{1-t|0}(x|\vec{X}_{0})=\frac{\exp(-\|x-\sqrt{\alpha_{t}}Y_{1}\|^{2}/(2(1-\alpha_{t})))}{(2\pi(1-\alpha_{t}))^{d/2}}\implies\nabla\log\vec{p}_{t|1}(x|Y_{1})=-\frac{x-\sqrt{\alpha_{t}}Y_{1}}{1-\alpha_{t}}.$ (103) $$

Plugging this into the right-hand side of ([102](#S10.E102)) and using Bayes’ rule, we get

$$ $\textstyle\mathfrak{s}(x,t)=\mathbb{E}\big{[}-\frac{x-\sqrt{\alpha_{t}}\vec{X}_{0}}{1-\alpha_{t}}|\sqrt{\alpha_{t}}\vec{X}_{0}+\sqrt{1-\alpha_{t}}\epsilon=x\big{]}.$ (104) $$

Comparing the right-hand sides of ([101](#S10.E101)) and ([104](#S10.E104)), we obtain that $\epsilon(x,t)=-\frac{\mathfrak{s}(x,t)}{\sqrt{1-\alpha_{t}}}$.

### 10.4 The relationship between the vector field v 𝑣 v and the score function

By construction (Lipman et al., ; Albergo and Vanden-Eijnden, ; Albergo et al., ), we have that

$$ $\textstyle\begin{split}\textstyle v(x,t)&\textstyle=\mathbb{E}[\dot{\alpha}_{t}Y_{1}+\dot{\beta}_{t}Y_{0}|x=\alpha_{t}Y_{1}+\beta_{t}Y_{0}]\\ \textstyle&\textstyle=\mathbb{E}[\frac{\dot{\alpha}_{t}(x-\beta_{t}Y_{0})}{\alpha_{t}}+\dot{\beta}_{t}Y_{0}|x=\alpha_{t}Y_{1}+\beta_{t}Y_{0}]\\ \textstyle&\textstyle=\frac{\dot{\alpha}_{t}}{\alpha_{t}}x+(\dot{\beta}_{t}-\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t})\mathbb{E}[Y_{0}|x=\alpha_{t}Y_{1}+\beta_{t}Y_{0}],\end{split}$ (105) $$

where we used that
$Y_{1}=(x-\beta_{t}Y_{0})/\alpha_{t}$.
Also, we can write the score as follows

$$ $\textstyle\mathfrak{s}(x,t)\!:=\!\nabla\log p_{t}(x)\!=\!\frac{\nabla p_{t}(x)}{p_{t}(x)}\!=\!\frac{\nabla\mathbb{E}[p_{t|1}(x|Y_{1})]}{p_{t}(x)}\!=\!\frac{\mathbb{E}[\nabla p_{t|1}(x|Y_{1})]}{p_{t}(x)}\!=\!\frac{\mathbb{E}[p_{t|1}(x|Y_{1})\nabla\log p_{t|1}(x|Y_{1})]}{p_{t}(x)},$ (106) $$

where

$$ $\textstyle p_{t|1}(x|Y_{1})=\frac{\exp(-\|x-\alpha_{t}Y_{1}\|^{2}/(2\beta^{2}_{t}))}{(2\pi\beta^{2}_{t})^{d/2}}\implies\nabla\log\vec{p}_{t|1}(x|Y_{1})=-\frac{x-\alpha_{t}Y_{1}}{\beta^{2}_{t}}$ (107) $$

Plugging this back into the right-hand side of ([106](#S10.E106)), we obtain

$$ $\textstyle\begin{split}\textstyle\mathfrak{s}(x,t)&\textstyle=-\frac{\mathbb{E}[p_{t|1}(x|Y_{1})\frac{x-\alpha_{t}Y_{1}}{\beta^{2}_{t}}]}{p_{t}(x)}=-\frac{\int\vec{p}_{t|1}(x|Y_{1})p_{1}(Y_{1})\frac{x-\alpha_{t}Y_{1}}{\beta^{2}_{t}}\,dY_{1}}{\vec{p}_{t}(x)}\\ \textstyle&\textstyle=-\int p_{1|t}(Y_{1}|x)\frac{x-\alpha_{t}Y_{1}}{\beta^{2}_{t}}\,dY_{1}=-\mathbb{E}[\frac{x-\alpha_{t}Y_{1}}{\beta^{2}_{t}}|x=\alpha_{t}Y_{1}+\beta_{t}Y_{0}]=-\frac{\mathbb{E}[Y_{0}|x=\alpha_{t}Y_{1}+\beta_{t}Y_{0}]}{\beta_{t}}\end{split}$ (108) $$

The last equality holds because $(x-\alpha_{t}Y_{1})/\beta_{t}=Y_{0}$. Putting together ([105](#S10.E105)) and ([108](#S10.E108)), we obtain that

$$ $\textstyle v(x,t)=\frac{\dot{\alpha}_{t}}{\alpha_{t}}x+\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})\mathfrak{s}(x,t)\iff\mathfrak{s}(x,t)=\frac{1}{\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}\big{(}v(x,t)-\frac{\dot{\alpha}_{t}}{\alpha_{t}}x\big{)}$ (109) $$

Thus, the ODE ([3](#S2.E3))
can be rewritten like this:

$$ $\textstyle\frac{\mathrm{d}X_{t}}{\mathrm{d}t}$ $\textstyle=\frac{\dot{\alpha}_{t}}{\alpha_{t}}X_{t}+\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})\mathfrak{s}(X_{t},t),\qquad X_{0}\sim p_{0}.$ (110) $$

To allow for an arbitrary diffusion coefficient, we need to add a correction term to the drift:

$$ $\textstyle\mathrm{d}X_{t}$ $\textstyle=\big{(}\frac{\dot{\alpha}_{t}}{\alpha_{t}}X_{t}+\big{(}\frac{\sigma(t)^{2}}{2}+\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})\big{)}\mathfrak{s}(X_{t},t)\big{)}\mathrm{d}t+\sigma(t)\mathrm{d}B_{t},\qquad X_{0}\sim p_{0}.$ (111) $$

This can be easily shown by writing down the Fokker-Planck equations for ([110](#S10.E110)) and ([111](#S10.E111)), and observing that they are the same up to a cancellation of terms.
Finally, if we plug the right-hand side of ([109](#S10.E109)) into ([111](#S10.E111)), we obtain the SDE for Flow Matching with arbitrary noise schedule (equation ([4](#S2.E4))).

## 11 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 optimal control. We show that when the action space is Euclidean and the transition probabilities are conditional Gaussians, taking the limit in which the stepsize goes to zero on the KL-regularized RL problem gives rise to 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.

[Section 11.4](#S11.SS4) is providing a more direct, rigorous, continuous-time connection between SOC and MaxEnt RL, as it shows that the expected control cost is equal to the KL divergence between the distributions over trajectories, conditioned on the starting points (see equation ([18](#S4.E18))).

### 11.1 Maximum entropy RL

Several diffusion fine-tuning methods (Black et al., ; Uehara et al., ) are based on KL-regularized RL, also known as maximum entropy RL, which we review in the following.
In the classical reinforcement learning (RL) setting, we have an agent that, starting from state $s_{0}\sim p_{0}$, iteratively observes a state $s_{k}$, takes an action $a_{k}$ according to a policy $\pi(a_{k};s_{k},k)$ which leads to a new state $s_{k+1}$ according to a fixed transition probability $p(s_{k+1}|a_{k},s_{k})$, and obtains rewards $r_{k}(s_{k},a_{k})$.
This can be summarized into a trajectory $\tau=((s_{k},a_{k}))_{k=0}^{K}$.
The goal is to optimize the policy $\pi$ in order to maximize the expected total reward, i.e. $\max_{\pi}\mathbb{E}_{\tau\sim\pi,p}[\sum_{k=0}^{K}r_{k}(s_{k},a_{k})]$.

Maximum entropy RL (MaxEnt RL; Ziebart et al. ()) amounts to adding the entropy $H(\pi)$ of the policy $\pi(\cdot;s_{k},k)$ to the reward for each step $k$, in order to encourage exploration and improve robustness to changes in the environment: $\max_{\pi}\mathbb{E}_{\tau\sim\pi,p}[\sum_{k=0}^{K}r_{k}(s_{k},a_{k})+\sum_{k=0}^{K-1}H(\pi(\cdot;s_{k},k))]$ (^7^77The 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.). As a generalization, one can regularize using the negative KL divergence between $\pi(\cdot;s_{k},k)$ and a base policy $\pi_{\mathrm{base}}(\cdot;s_{k},k)$:

$$ $\textstyle\max_{\pi}\mathbb{E}_{\tau\sim\pi,p}[\sum_{k=0}^{K}r_{k}(s_{k},a_{k})-\sum_{k=0}^{K-1}\mathrm{KL}(\pi(\cdot;s_{k},k)||\pi_{\mathrm{base}}(\cdot;s_{k},k))],$ (112) $$

which prevents the learned policy to deviate too much from the base policy. Each policy $\pi$ induces a distribution $q(\tau)$ over trajectories $\tau$, and the MaxEnt RL problem ([112](#S11.E112)) can be expressed solely in terms of such distributions ([Lemma 3](#Thmlemma3) in [Section 11.3](#S11.SS3)):

$$ $\textstyle\max_{q}\mathbb{E}_{\tau\sim q}[\sum_{k=0}^{K}r_{k}(s_{k},a_{k})]-\mathrm{KL}(q||q^{\mathrm{base}}),$ (113) $$

where $q^{\mathrm{base}}$ is the distribution induced by the base policy $\pi_{\mathrm{base}}$, and the maximization is over all distributions $q$ such that their marginal for $s_{0}$ is $p_{0}$. We can further recast this problem as ([Lemma 4](#Thmlemma4) in [Section 11.3](#S11.SS3)):

$$ $\textstyle\min_{q}\mathrm{KL}(q||q^{*}),\qquad\text{where}\ q^{*}(\tau):=q^{\mathrm{base}}(\tau)\exp\big{(}\sum_{k=0}^{K}r_{k}(s_{k},a_{k})-\mathcal{V}(s_{0},0)\big{)},$ (114) $$

where

$$ $\textstyle\begin{split}\textstyle\mathcal{V}(s_{k},k)&\textstyle:=\log\big{(}\mathbb{E}_{\tau\sim\pi_{\mathrm{base}},p}[\exp\big{(}\sum_{k^{\prime}=k}^{K}r_{k^{\prime}}(s_{k^{\prime}},a_{k^{\prime}})\big{)}|s_{k}]\big{)}\\ \textstyle&\textstyle=\max_{\pi}\mathbb{E}_{\tau\sim\pi,p}\big{[}\sum_{k^{\prime}=k}^{K}r_{k^{\prime}}(s_{k^{\prime}},a_{k^{\prime}})-\sum_{k^{\prime}=k}^{K-1}\mathrm{KL}(\pi(\cdot;s_{k^{\prime}},k^{\prime})||\pi_{\mathrm{base}}(\cdot;s_{k^{\prime}},k^{\prime}))|s_{k}\big{]}\end{split}$ (115) $$

is the value function.
Problem ([114](#S11.E114)) directly implies that the distribution induced by the optimal policy $\pi^{*}$ is the tilted distribution $q^{*}$ (which has initial marginal $p_{0}$).

### 11.2 From maximum entropy RL to stochastic optimal control

The following well-known result, which we prove in [Section 11.3](#S11.SS3), shows that in a natural sense, the continuous-time continuous-space version of MaxEnt RL is the SOC
framework introduced in [Section 4.1](#S4.SS1). In particular, when states and actions are vectors in $\mathbb{R}^{d}$, policies are specified by a 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
$\mathbb{R}^{d}$,
- (ii)
Policies $\pi$ are specified as $\pi(a_{k};s_{k},k)=\delta(a_{k}-u(s_{k},kh))$, where $u:\mathbb{R}^{d}\times[0,T]\to\mathbb{R}^{d}$ is a vector field, and $\delta$ denotes the Dirac delta,
- (iii)
Transition probabilities are conditional Gaussian densities: $p(s_{k+1}|a_{k},s_{k})=N(s_{k}+h(b(s_{k},kh)+\sigma(kh)a_{k}),h\sigma(kh)\sigma(kh)^{\top})$, where $h=T/K$ is the stepsize, and $b$ and $\sigma$ are defined as in [Section 4.1](#S4.SS1).

Then, in the limit in which the number of steps $K$ grows to infinity, the problem ([112](#S11.E112)) is equivalent to the SOC
problem ([12](#S4.E12))-([13](#S4.E13)), identifying

- •
the sequence of states $(s_{k})_{k=0}^{k}$ with the trajectory $\bm{X}^{u}=(X^{u}_{t})_{t\in[0,1]}$,
- •
the running reward $\sum_{k=0}^{K-1}r_{k}(s_{k},a_{k})$ with the negative running cost $-\int_{0}^{T}f(X^{u}_{t},t)\,\mathrm{d}t$,
- •
the terminal reward $r_{K}(s_{K},a_{K})$ with the negative terminal cost $-g(X^{u}_{T})$,
- •
the KL regularization $\mathbb{E}_{\tau\sim\pi,p}[\sum_{k=0}^{K-1}\mathrm{KL}(\pi(\cdot;s_{k},k)||\pi_{\mathrm{base}}(\cdot;s_{k},k))]$ with $\frac{1}{2}$ times the expected $L^{2}$ norm of the control $\frac{1}{2}\mathbb{E}\big{[}\int_{0}^{T}\|u(X^{u}_{t},t)\|^{2}\,\mathrm{d}t\big{]}$,
- •
and the value function $\mathcal{V}(s_{k},k)$ defined in ([115](#S11.E115)) with the negative value function $-V(x,t)$ defined in [Section 4.1](#S4.SS1).

A first consequence of this result is that every loss function designed for generic MaxEnt RL problems 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 particular, we can differentiate the state and terminal costs.

A second consequence of [Proposition 5](#Thmproposition5) is that the characterization ([114](#S11.E114)) can be translated to the SOC setting. The analogs of the distributions $q^{*}$, $q^{\mathrm{base}}$ induced by the optimal policy $\pi^{*}$ and the base policy $\pi^{\mathrm{base}}$ are the distributions
$p^{*},p^{\mathrm{base}}$ induced by the optimal control $u^{*}$ and the null control.
For an arbitrary trajectory $\bm{X}=(X_{t})_{t\in[0,T]}$, the
relation between $\mathbb{P}^{*}$ and $\mathbb{P}^{\mathrm{base}}$ is given by

$$ $\textstyle\frac{\mathrm{d}\mathbb{P}^{*}}{\mathrm{d}\mathbb{P}^{\mathrm{base}}}(\bm{X})=\exp(-\int_{0}^{T}f(X_{t},t)\,\mathrm{d}t-g(X_{T})+V(X_{0},0))$ (116) $$

where $V$ is the value function as defined in [Section 4.1](#S4.SS1). Note that this matches the statement in ([22](#S4.E22)).

###### Proposition 5 .

### 11.3 Proof of Proposition 5 : from MaxEnt RL to SOC

Since the transition $p(s_{k+1}|a_{k},s_{k})$ is fixed, for each $\pi$ we can define

$$ $\textstyle\tilde{\pi}(a_{k},s_{k+1};s_{k},k)=\pi(a_{k};s_{k},k)p(s_{k+1}|a_{k},s_{k})\ \text{and}\ \tilde{\pi}_{\mathrm{base}}(a_{k},s_{k+1};s_{k},k)=\pi_{\mathrm{base}}(a_{k};s_{k},k)p(s_{k+1}|a_{k},s_{k}),$ (117) $$

and reexpress ([112](#S11.E112)) as (see [Lemma 3](#Thmlemma3))

$$ $\textstyle\min_{\tilde{\pi}}\mathbb{E}_{\tau\sim\tilde{\pi}}[\sum_{k=0}^{K}r_{k}(s_{k},a_{k})-\sum_{k=0}^{K-1}\mathrm{KL}(\tilde{\pi}(\cdot,\cdot;s_{k},k)||\tilde{\pi}_{\mathrm{base}}(\cdot,\cdot;s_{k},k))].$ (118) $$

Using the hypothesis of the proposition, we can write

$$ $\textstyle\begin{split}\textstyle\tilde{\pi}(a_{k},s_{k+1};s_{k},k)&\textstyle=\delta(a_{k}-u(s_{k},k\eta))N(s_{k}+\eta(b(s_{k},k\eta)+\sigma(k\eta)a_{k}),\eta\sigma(k\eta)\sigma(k\eta)^{\top})\\ \textstyle&\textstyle=\delta(a_{k}-u(s_{k},k\eta))\tilde{\pi}(s_{k+1};s_{k},k),\end{split}$ (119) $$

where $\tilde{\pi}(s_{k+1};s_{k},k)=N(s_{k}+\eta(b(s_{k},k\eta)+\sigma(k\eta)u(s_{k},k\eta)),\eta\sigma(k\eta)\sigma(k\eta)^{\top})$ is the state transition kernel. We set the base policy as $\pi_{\mathrm{base}}(a_{k};s_{k},k)=\delta(a_{k})$, and we obtain analogously that $\tilde{\pi}(a_{k},s_{k+1};s_{k},k)=\delta(a_{k})\tilde{\pi}_{\mathrm{base}}(s_{k+1};s_{k},k)$ with $\tilde{\pi}_{\mathrm{base}}(s_{k+1};s_{k},k)=N(s_{k}+\eta b(s_{k},k\eta),\eta\sigma(k\eta)\sigma(k\eta)^{\top})$. Now, if we take $K$ large, the trajectory $(s_{k})_{k=0}^{K}$ generated by $\tilde{\pi}$ can be regarded as the Euler-Maruyama discretization of a solution $X^{u}$ of the controlled SDE ([13](#S4.E13)), while the trajectory generated by $\tilde{\pi}_{\mathrm{base}}$ is the discretization of the uncontrolled process $X^{0}$ obtained by setting $u=0$. As a consequence

$$ $\textstyle\begin{split}\textstyle&\textstyle\lim_{K\to\infty}\mathbb{E}_{\tau\sim\tilde{\pi}}[\sum_{k=0}^{K-1}\mathrm{KL}(\tilde{\pi}(\cdot,\cdot;s_{k},k)||\tilde{\pi}_{\mathrm{base}}(\cdot,\cdot;s_{k},k))]\\ \textstyle&\textstyle=\lim_{K\to\infty}\mathbb{E}_{\tau\sim\tilde{\pi}}[\sum_{k=0}^{K-1}\mathrm{KL}(\tilde{\pi}(\cdot;s_{k},k)||\tilde{\pi}_{\mathrm{base}}(\cdot;s_{k},k))]=\mathbb{E}_{X^{u}\sim\mathbb{P}^{u}}[\log\frac{\mathrm{d}\mathbb{P}^{u}}{\mathrm{d}\mathbb{P}^{0}}(X^{u})],\end{split}$ (120) $$

where $\mathbb{P}^{u}$ and $\mathbb{P}^{0}$ are the measures of the processes $X^{u}$ and $X^{0}$, respectively. The Girsanov theorem ([Theorem 2](#Thmtheorem2)) implies that $\log\frac{\mathrm{d}\mathbb{P}^{u}}{\mathrm{d}\mathbb{P}^{0}}(X^{u})=-\int_{0}^{T}\langle u(X^{u}_{t},t),\mathrm{d}B_{t}\rangle-\frac{1}{2}\int_{0}^{T}\|u(X^{u}_{t},t)\|^{2}\,\mathrm{d}t$, which implies that $\mathbb{E}_{X^{u}\sim\mathbb{P}^{u}}[\log\frac{\mathrm{d}\mathbb{P}^{u}}{\mathrm{d}\mathbb{P}^{0}}(X^{u})]=-\frac{1}{2}\mathbb{E}_{X^{u}\sim\mathbb{P}^{u}}[\int_{0}^{T}\|u(X^{u}_{t},t)\|^{2}\,\mathrm{d}t]$. Setting the rewards $r_{k}(a_{k},s_{k})=\eta f(s_{k},k\eta)$ for $k\in\{0,\dots,K-1\}$ and $r_{K}(a_{K},s_{K})=\eta g(s_{k})$, where $f$ and $g$ are as in [Section 4.1](#S4.SS1), yields the following limiting object:

$$ $\textstyle\lim_{K\to\infty}\mathbb{E}_{\tau\sim\tilde{\pi}}[\sum_{k=0}^{K}r_{k}(s_{k},a_{k})]=\mathbb{E}_{X^{u}\sim\mathbb{P}^{u}}[\int_{0}^{T}f(X^{u}_{t},t)\,dt+g(X^{u}_{T})].$ (121) $$

Hence, the limit of the MaxEnt RL loss ([118](#S11.E118)) is the SOC loss ([12](#S4.E12)).

###### Lemma 3 .

Let $\tilde{\pi}(a_{k},s_{k+1};s_{k},k)$ and $\tilde{\pi}_{\mathrm{base}}(a_{k},s_{k+1};s_{k},k)$ be as defined in ([117](#S11.E117)). $\mathrm{KL}(\tilde{\pi}(\cdot,\cdot;s_{k},k)||\tilde{\pi}_{\mathrm{base}}(\cdot,\cdot;s_{k},k))]$ and $\mathrm{KL}(\pi(\cdot;s_{k},k)||\pi_{\mathrm{base}}(\cdot;s_{k},k))]$ are equal. Moreover, if $q$, $q^{\mathrm{base}}$ denote the distributions over trajectories induced by $\pi$, $\pi_{\mathrm{base}}$, we have that

$$ $\textstyle\mathrm{KL}(q||q^{\mathrm{base}})=\mathbb{E}[\sum_{k=0}^{K-1}\mathrm{KL}(\pi(\cdot;s_{k},k)||\pi_{\mathrm{base}}(\cdot;s_{k},k))].$ (122) $$

###### Proof.

We have that

$$ $\textstyle\begin{split}\textstyle&\textstyle\mathrm{KL}(\tilde{\pi}(\cdot,\cdot;s_{k},k)||\tilde{\pi}_{\mathrm{base}}(\cdot,\cdot;s_{k},k))]=\sum_{a_{k},s_{k+1}}\tilde{\pi}(a_{k},s_{k+1};s_{k},k)\log\frac{\tilde{\pi}(a_{k},s_{k+1};s_{k},k)}{\tilde{\pi}_{\mathrm{base}}(a_{k},s_{k+1};s_{k},k)}\\ \textstyle&\textstyle=\sum_{a_{k},s_{k+1}}\pi(a_{k};s_{k},k)p(s_{k+1}|a_{k},s_{k})\log\frac{\pi(a_{k};s_{k},k)p(s_{k+1}|a_{k},s_{k})}{\pi_{\mathrm{base}}(a_{k};s_{k},k)p(s_{k+1}|a_{k},s_{k})}\\ \textstyle&\textstyle=\sum_{a_{k},s_{k+1}}\pi(a_{k};s_{k},k)p(s_{k+1}|a_{k},s_{k})\log\frac{\pi(a_{k};s_{k},k)}{\pi_{\mathrm{base}}(a_{k};s_{k},k)}\\ \textstyle&\textstyle=\sum_{a_{k}}\pi(a_{k};s_{k},k)\big{(}\sum_{s_{k+1}}p(s_{k+1}|a_{k},s_{k})\big{)}\log\frac{\pi(a_{k};s_{k},k)}{\pi_{\mathrm{base}}(a_{k};s_{k},k)}\\ \textstyle&\textstyle=\sum_{a_{k}}\pi(a_{k};s_{k},k)\log\frac{\pi(a_{k};s_{k},k)}{\pi_{\mathrm{base}}(a_{k};s_{k},k)}=\mathrm{KL}(\pi(\cdot;s_{k},k)||\pi_{\mathrm{base}}(\cdot;s_{k},k))].\end{split}$ (123) $$

To prove ([122](#S11.E122)), by construction we can write

$$ $\textstyle q(\tau)=p_{0}(s_{0})\prod_{k=0}^{K-1}\tilde{\pi}(a_{k},s_{k+1};s_{k},k),\qquad\qquad q^{\mathrm{base}}(\tau)=p_{0}(s_{0})\prod_{k=0}^{K-1}\tilde{\pi}_{\mathrm{base}}(a_{k},s_{k+1};s_{k},k),$ (124) $$

which means that

$$ $\textstyle\begin{split}\textstyle\mathrm{KL}(q||q^{\mathrm{base}})&\textstyle=\mathbb{E}_{\tau\sim q}[\log\frac{q(\tau)}{q^{\mathrm{base}}(\tau)}]=\mathbb{E}_{\tau\sim q}[\sum_{k=0}^{K-1}\log\frac{\tilde{\pi}(a_{k},s_{k+1};s_{k},k)}{\tilde{\pi}_{\mathrm{base}}(a_{k},s_{k+1};s_{k},k)}]\\ \textstyle&\textstyle=\sum_{k=0}^{K-1}\mathbb{E}_{\tau\sim q^{0:(k+1)}}[\log\frac{\tilde{\pi}(a_{k},s_{k+1};s_{k},k)}{\tilde{\pi}_{\mathrm{base}}(a_{k},s_{k+1};s_{k},k)}]\\ \textstyle&\textstyle=\sum_{k=0}^{K-1}\mathbb{E}_{\tau\sim q^{0:k}}[\sum_{a_{k},s_{k+1}}\tilde{\pi}(a_{k},s_{k+1};s_{k},k)\log\frac{\tilde{\pi}(a_{k},s_{k+1};s_{k},k)}{\tilde{\pi}_{\mathrm{base}}(a_{k},s_{k+1};s_{k},k)}]\\ \textstyle&\textstyle=\sum_{k=0}^{K-1}\mathbb{E}_{\tau\sim q^{0:k}}[\mathrm{KL}(\tilde{\pi}(\cdot,\cdot;s_{k},k)||\tilde{\pi}_{\mathrm{base}}(\cdot,\cdot;s_{k},k))]\\ \textstyle&\textstyle=\sum_{k=0}^{K-1}\mathbb{E}_{\tau\sim q^{0:k}}[\mathrm{KL}(\pi(\cdot;s_{k},k)||\pi_{\mathrm{base}}(\cdot;s_{k},k))]\\ \textstyle&\textstyle=\mathbb{E}_{\tau\sim q^{0:k}}[\sum_{k=0}^{K-1}\mathrm{KL}(\pi(\cdot;s_{k},k)||\pi_{\mathrm{base}}(\cdot;s_{k},k))]\end{split}$ (125) $$

Here, the notation $q^{0:k}$ denotes the trajectory $q$ up to the state $s_{k}$.
∎

###### Lemma 4 .

The distribution-based MaxEnt RL formulation in ([113](#S11.E113)) is equivalent to the the following problem:

$$ $\textstyle\min_{q}\mathrm{KL}(q||q^{*}),\qquad\text{where}\ q^{*}(\tau):=\frac{q^{\mathrm{base}}(\tau)\exp\big{(}\sum_{k=0}^{K}r_{k}(s_{k},a_{k})\big{)}}{\frac{1}{p_{0}(s_{0})}\sum_{\{\tau^{\prime}|s^{\prime}_{0}=s_{0}\}}q^{\mathrm{base}}(\tau^{\prime})\exp\big{(}\sum_{k=0}^{K}r_{k}(s^{\prime}_{k},a^{\prime}_{k})\big{)}},$ (126) $$

where the minimization is over $q$ with marginal $p_{0}$ at step zero. The optimum of the problem is $q^{*}$, which satisfies the marginal constraint. The following alternative characterization of $q^{*}$ holds:

$$ $\textstyle q^{*}(\tau)$ $\textstyle=q^{\mathrm{base}}(\tau)\exp\big{(}\sum_{k=0}^{K}r_{k}(s_{k},a_{k})-\mathcal{V}(s_{0},0)\big{)},$ (127) $\textstyle\text{where}\ \mathcal{V}(x,k)$ $\textstyle=\max_{\pi}\mathbb{E}_{\tau\sim\pi,p}\big{[}\sum_{k^{\prime}=k}^{K}r_{k^{\prime}}(s_{k^{\prime}},a_{k^{\prime}})-\sum_{k^{\prime}=k}^{K-1}\mathrm{KL}(\pi(\cdot;s_{k^{\prime}},k^{\prime})||\pi_{\mathrm{base}}(\cdot;s_{k^{\prime}},k^{\prime}))|s_{k}=x\big{]}.$ (128) $$

###### Proof.

Let us expand $\mathrm{KL}(q||q^{*})$:

$$ $\textstyle\begin{split}\textstyle\mathrm{KL}(q||q^{*})&\textstyle=\mathbb{E}_{\tau\sim q}\big{[}\log\frac{q(\tau)}{q^{*}(\tau)}\big{]}\\ \textstyle&\textstyle=\mathbb{E}_{\tau\sim q}\big{[}\log q(\tau)-\log q^{\mathrm{base}}(\tau)-\sum_{k=0}^{K}r_{k}(s_{k},a_{k})\\ \textstyle&\textstyle\qquad\qquad+\log\big{(}\frac{1}{p_{0}(s_{0})}\sum_{\{\tau^{\prime}|s^{\prime}_{0}=s_{0}\}}q^{\mathrm{base}}(\tau^{\prime})\exp\big{(}\sum_{k=0}^{K}r_{k}(s^{\prime}_{k},a^{\prime}_{k})\big{)}\big{)}\big{]}\\ \textstyle&\textstyle=\mathrm{KL}(q||q^{\mathrm{base}})-\mathbb{E}_{\tau\sim q}\big{[}\sum_{k=0}^{K}r_{k}(s_{k},a_{k})\big{]}\\ \textstyle&\textstyle\qquad+\mathbb{E}_{s_{0}\sim p_{0}}\big{[}\log\big{(}\frac{1}{p_{0}(s_{0})}\sum_{\{\tau^{\prime}|s^{\prime}_{0}=s_{0}\}}q^{\mathrm{base}}(\tau^{\prime})\exp\big{(}\sum_{k=0}^{K}r_{k}(s^{\prime}_{k},a^{\prime}_{k})\big{)}\big{)}\big{]},\end{split}$ (129) $$

where the third equality holds because the marginal of $q$ at step zero is $p_{0}$ by hypothesis. Since the third term in the right-hand side is independent of $q$, this proves the equivalence between ([113](#S11.E113)) and ([126](#S11.E126)).

Next, we prove that the marginal of $q^{*}$ at step zero is $p_{0}$:

$$ $\textstyle\sum_{\{\tau|s_{0}=x\}}q^{*}(\tau):=\sum_{\{\tau|s_{0}=x\}}\frac{q^{\mathrm{base}}(\tau)\exp\big{(}\sum_{k=0}^{K}r_{k}(s_{k},a_{k})\big{)}}{\frac{1}{p_{0}(x)}\sum_{\{\tau^{\prime}|s^{\prime}_{0}=x\}}q^{\mathrm{base}}(\tau^{\prime})\exp\big{(}\sum_{k=0}^{K}r_{k}(s^{\prime}_{k},a^{\prime}_{k})\big{)}}=p_{0}(x).$ (130) $$

Now, for an arbitrary $s_{0}$, let $q_{s_{0}}$, $q^{*}_{s_{0}}$ be the distributions $q$, $q^{*}$ conditioned on the initial state being $s_{0}$. We can write an analog to equation ([129](#S11.E129)) for $q_{s_{0}}$, $q^{*}_{s_{0}}$:

$$ $\textstyle\begin{split}\textstyle\mathrm{KL}(q_{s_{0}}||q^{*}_{s_{0}})&\textstyle=\mathbb{E}_{\tau\sim q_{s_{0}}}\big{[}\log\frac{q_{s_{0}}(\tau)}{q^{*}_{s_{0}}(\tau)}\big{]}\\ \textstyle&\textstyle=\mathbb{E}_{\tau\sim q_{s_{0}}}\big{[}\log q_{s_{0}}(\tau)-\log q^{\mathrm{base}}_{s_{0}}(\tau)-\sum_{k=0}^{K}r_{k}(s_{k},a_{k})\\ \textstyle&\textstyle\qquad\qquad+\log\big{(}\frac{1}{p_{0}(s_{0})}\sum_{\{\tau^{\prime}|s^{\prime}_{0}=s_{0}\}}q^{\mathrm{base}}_{s_{0}}(\tau^{\prime})\exp\big{(}\sum_{k=0}^{K}r_{k}(s^{\prime}_{k},a^{\prime}_{k})\big{)}\big{)}\big{]}\\ \textstyle&\textstyle=\mathrm{KL}(q_{s_{0}}||q^{\mathrm{base}}_{s_{0}})-\mathbb{E}_{\tau\sim q_{s_{0}}}\big{[}\sum_{k=0}^{K}r_{k}(s_{k},a_{k})\big{]}\\ \textstyle&\textstyle\qquad+\log\big{(}\frac{1}{p_{0}(s_{0})}\sum_{\{\tau^{\prime}|s^{\prime}_{0}=s_{0}\}}q^{\mathrm{base}}(\tau^{\prime})\exp\big{(}\sum_{k=0}^{K}r_{k}(s^{\prime}_{k},a^{\prime}_{k})\big{)}\big{)},\end{split}$ (131) $$

Hence,

$$ $\textstyle\begin{split}\textstyle 0=\min_{q_{s_{0}}}\mathrm{KL}(q_{s_{0}}||q^{*}_{s_{0}})&\textstyle=-\max_{q_{s_{0}}}\{\mathbb{E}_{\tau\sim q_{s_{0}}}\big{[}\sum_{k=0}^{K}r_{k}(s_{k},a_{k})\big{]}-\mathrm{KL}(q_{s_{0}}||q^{\mathrm{base}}_{s_{0}})\}\\ \textstyle&\textstyle+\log\big{(}\frac{1}{p_{0}(s_{0})}\sum_{\{\tau^{\prime}|s^{\prime}_{0}=s_{0}\}}q^{\mathrm{base}}(\tau^{\prime})\exp\big{(}\sum_{k=0}^{K}r_{k}(s^{\prime}_{k},a^{\prime}_{k})\big{)}\big{)}.\end{split}$ (132) $$

And applying ([122](#S11.E122)) from ([122](#S11.E122)), we obtain that

$$ $\textstyle\begin{split}\textstyle&\textstyle\log\big{(}\frac{1}{p_{0}(s_{0})}\sum_{\{\tau^{\prime}|s^{\prime}_{0}=s_{0}\}}q^{\mathrm{base}}(\tau^{\prime})\exp\big{(}\sum_{k=0}^{K}r_{k}(s^{\prime}_{k},a^{\prime}_{k})\big{)}\big{)}\\ \textstyle&\textstyle=\max_{\pi}\mathbb{E}_{\tau\sim\pi,p}\big{[}\sum_{k=0}^{K}r_{k}(s_{k},a_{k})-\sum_{k=0}^{K-1}\mathrm{KL}(\pi(\cdot;s_{k},k)||\pi_{\mathrm{base}}(\cdot;s_{k},k))|s_{0}\big{]}=\mathcal{V}(s_{0},0),\end{split}$ (133) $$

which concludes the proof.
∎

###### Lemma 3 .

###### Proof.

###### Lemma 4 .

###### Proof.

### 11.4 Proof of equation ( 18 ): the control cost is a KL regularizer

###### Theorem 2 (Girsanov theorem for SDEs) .

If the two SDEs

$$ $\textstyle\mathrm{d}X_{t}$ $\textstyle=b_{1}(X_{t},t)\,\mathrm{d}t+\sigma(X_{t},t)\,\mathrm{d}B_{t},\qquad X_{0}=x_{\mathrm{init}}$ (134) $\textstyle dY_{t}$ $\textstyle=(b_{1}(Y_{t},t)+b_{2}(Y_{t},t))\,\mathrm{d}t+\sigma(Y_{t},t)\,\mathrm{d}B_{t},\qquad Y_{0}=x_{\mathrm{init}}$ (135) $$

admit unique strong solutions on $[0,T]$, then for any bounded continuous functional $\Phi$ on $C([0,T])$, we have that

$$ $\textstyle\begin{split}\textstyle\mathbb{E}[\Phi(\bm{X})]&\textstyle=\mathbb{E}\big{[}\Phi(\bm{Y})\exp\big{(}-\int_{0}^{T}\sigma(Y_{t},t)^{-1}b_{2}(Y_{t},t)\,\mathrm{d}B_{t}-\frac{1}{2}\int_{0}^{T}\|\sigma(Y_{t},t)^{-1}b_{2}(Y_{t},t)\|^{2}\,\mathrm{d}t\big{)}\big{]}\\ \textstyle&\textstyle=\mathbb{E}\big{[}\Phi(\bm{Y})\exp\big{(}-\int_{0}^{T}\sigma(Y_{t},t)^{-1}b_{2}(Y_{t},t)\,d\tilde{B}_{t}+\frac{1}{2}\int_{0}^{T}\|\sigma(Y_{t},t)^{-1}b_{2}(Y_{t},t)\|^{2}\,\mathrm{d}t\big{)}\big{]},\end{split}$ (136) $$

where $\tilde{B}_{t}=B_{t}+\int_{0}^{t}\sigma(Y_{s},s)^{-1}b_{2}(Y_{s},s)\,\mathrm{d}s$. More generally, $b_{1}$ and $b_{2}$ can be random processes that are adapted to filtration of $\bm{B}$.

Consider the SDEs

$$ $\textstyle\mathrm{d}X_{t}$ $\textstyle=b(X_{t},t)\,\mathrm{d}t+\sigma(t)\mathrm{d}B_{t},\qquad$ $\textstyle X_{0}=x_{0},$ (137) $\textstyle\mathrm{d}X^{u}_{t}$ $\textstyle=\left(b(X^{u}_{t},t)+\sigma(t)u(X^{u}_{t},t)\right)\,\mathrm{d}t+\sigma(t)\mathrm{d}B_{t},\qquad$ $\textstyle X^{u}_{0}=x_{0}.$ (138) $$

If we let $\mathbb{P}\rvert_{x_{0}}$, $\mathbb{P}^{u}\rvert_{x_{0}}$ be the probability measures of the solutions of ([137](#S11.E137)) and ([138](#S11.E138)), [Theorem 2](#Thmtheorem2) implies that

$$ $\textstyle\log\frac{\mathrm{d}\mathbb{P}\rvert_{x_{0}}}{\mathrm{d}\mathbb{P}^{u}\rvert_{x_{0}}}(\bm{X}^{u})=-\int_{0}^{1}u(X^{u}_{t},t)\,\mathrm{d}B_{t}-\frac{1}{2}\int_{0}^{1}\|u(X^{u}_{t},t)\|^{2}\,\mathrm{d}t.$ (139) $$

Hence,

$$ $\textstyle\begin{split}\textstyle D_{\text{KL}}\left(\mathbb{P}^{u}\rvert_{x_{0}}\;\middle\|\;\mathbb{P}\rvert_{x_{0}}\right)&\textstyle=\mathbb{E}\big{[}\log\frac{\mathrm{d}\mathbb{P}^{u}\rvert_{x_{0}}}{\mathrm{d}\mathbb{P}\rvert_{x_{0}}}(\bm{X}^{u})|X^{u}_{0}=x_{0}\big{]}=-\mathbb{E}\big{[}\log\frac{\mathrm{d}\mathbb{P}\rvert_{x_{0}}}{\mathrm{d}\mathbb{P}^{u}\rvert_{x_{0}}}(\bm{X}^{u})|X^{u}_{0}=x_{0}\big{]}\\ \textstyle&\textstyle=\mathbb{E}\big{[}\int_{0}^{1}u(X^{u}_{t},t)\,\mathrm{d}B_{t}+\frac{1}{2}\int_{0}^{1}\|u(X^{u}_{t},t)\|^{2}\,\mathrm{d}t|X^{u}_{0}=x_{0}\big{]}=\mathbb{E}\big{[}\frac{1}{2}\int_{0}^{1}\|u(X^{u}_{t},t)\|^{2}\,\mathrm{d}t|X^{u}_{0}=x_{0}\big{]},\end{split}$ (140) $$

where we used that stochastic integrals are martingales.

###### Theorem 2 (Girsanov theorem for SDEs) .

## 12 Proofs of Section 4.3 : memoryless noise schedule and fine-tuning recipe

### 12.1 Proof of Proposition 1 :
the memoryless noise schedule

We consider the forward-backward SDEs ([63](#S10.E63))-([64](#S10.E64)) with arbitrary noise schedule. By [Proposition 4](#Thmproposition4), the trajectories $\vec{\bm{X}}$, $\bm{X}$ of these two processes are equally distributed up to a time flip, which also means that their marginals satisfy $\vec{p}_{t}=p_{1-t}$, for all $t\in[0,1]$.
First, we develop an explicit expression for the score function $s(x,t)=\nabla\log p_{t}(x)$. By the properties of flow matching, we know that $p_{t}$ is the distribution of the interpolation variable $\bar{X}_{t}=\beta_{t}\bar{X}_{0}+\alpha_{t}\bar{X}_{1}$, where $\bar{X}_{0}\sim N(0,I),\bar{X}_{1}\sim p^{\mathrm{data}}$ are independent. Thus, $\frac{\bar{X}_{t}-\alpha_{t}\bar{X}_{1}}{\beta_{t}}\sim N(0,\mathrm{I})$, which means that we can express the density $p_{t}$ as

$$ $\textstyle p_{t}(x)=\int_{\mathbb{R}^{d}}\frac{\exp\big{(}-\frac{\|x-\alpha_{t}y\|^{2}}{2\beta_{t}^{2}}\big{)}}{(2\pi\beta_{t}^{2})^{d/2}}p^{\mathrm{data}}(y)\,\mathrm{d}y.$ (141) $$

Thus,

$$ $\textstyle s(x,t)=\nabla\log p_{t}(x)=-\frac{x}{\beta_{t}^{2}}+\frac{\alpha_{t}}{\beta_{t}^{2}}\frac{\int_{\mathbb{R}^{d}}y\exp\big{(}-\frac{\|x-\alpha_{t}y\|^{2}}{2\beta_{t}^{2}}\big{)}p^{\mathrm{data}}(y)\,\mathrm{d}y}{\int_{\mathbb{R}^{d}}\exp\big{(}-\frac{\|x-\alpha_{t}y\|^{2}}{2\beta_{t}^{2}}\big{)}p^{\mathrm{data}}(y)\,\mathrm{d}y}:=-\frac{x-\alpha_{t}\xi_{t}(x)}{\beta_{t}^{2}},$ (142) $$

where we defined

$$ $\textstyle\xi_{t}(x)=\frac{\int_{\mathbb{R}^{d}}y\exp\big{(}-\frac{\|x-\alpha_{t}y\|^{2}}{2\beta_{t}^{2}}\big{)}p^{\mathrm{data}}(y)\,\mathrm{d}y}{\int_{\mathbb{R}^{d}}\exp\big{(}-\frac{\|x-\alpha_{t}y\|^{2}}{2\beta_{t}^{2}}\big{)}p^{\mathrm{data}}(y)\,\mathrm{d}y}.$ (143) $$

Hence, we can rewrite the forward SDE ([63](#S10.E63)) as

$$ $\textstyle\mathrm{d}\vec{X}_{t}$ $\textstyle=\big{(}-\kappa_{1-t}\vec{X}_{t}-\big{(}\frac{\sigma(1-t)^{2}}{2}-\eta_{1-t}\big{)}\frac{\vec{X}_{t}-\alpha_{1-t}\xi_{1-t}(\vec{X}_{t})}{\beta_{1-t}^{2}}\big{)}\,\mathrm{d}t+\sigma(1-t)\,\mathrm{d}B_{t},\qquad\vec{X}_{0}\sim p_{\mathrm{data}}$ (144) $$

Hence, if we substitute $\kappa_{1-t}\leftarrow\kappa_{1-t}+\frac{\sigma(1-t)^{2}-2\eta_{1-t}}{2\beta_{1-t}^{2}}$, $\xi_{1-t}\leftarrow\frac{\alpha_{1-t}(\sigma(1-t)^{2}-2\eta_{1-t})}{2\beta_{1-t}^{2}}\xi_{1-t}(\vec{X}_{t})$ (where we ignore the dependency on $\vec{X}_{t}$), $\sqrt{2\eta_{1-t}}\leftarrow\sigma(1-t)$, we can apply [Lemma 2](#Thmlemma2), which yields

$$ $\textstyle\begin{split}\textstyle\vec{X}_{t}&\textstyle=\vec{X}_{0}\exp\big{(}-\int_{0}^{t}\big{(}\kappa_{1-s}+\frac{\sigma(1-s)^{2}-2\eta_{1-s}}{2\beta_{1-s}^{2}}\big{)}\,\mathrm{d}s\big{)}\\ \textstyle&\textstyle\qquad+\int_{0}^{t}\exp\big{(}-\int_{t^{\prime}}^{t}\big{(}\kappa_{1-s}+\frac{\sigma(1-s)^{2}-2\eta_{1-s}}{2\beta_{1-s}^{2}}\big{)}\,\mathrm{d}s\big{)}\frac{\alpha_{1-t^{\prime}}(\sigma(1-t^{\prime})^{2}-2\eta_{1-t^{\prime}})}{2\beta_{1-t^{\prime}}^{2}}\xi_{1-t^{\prime}}(\vec{X}_{t^{\prime}})\,\mathrm{d}t^{\prime}\\ \textstyle&\textstyle\qquad+\int_{0}^{t}\sigma(1-t^{\prime})\exp\big{(}-\int_{t^{\prime}}^{t}\big{(}\kappa_{1-s}+\frac{\sigma(1-s)^{2}-2\eta_{1-s}}{2\beta_{1-s}^{2}}\big{)}\,\mathrm{d}s\big{)}\,\mathrm{d}B_{t^{\prime}}.\end{split}$ (145) $$

We simplify the recurring expression:

$$ $\textstyle\kappa_{1-s}+\frac{\sigma(1-s)^{2}-2\eta_{1-s}}{2\beta_{1-s}^{2}}=\frac{\dot{\alpha}_{1-s}}{\alpha_{1-s}}+\frac{\sigma(1-s)^{2}-2\beta_{1-s}\big{(}\frac{\dot{\alpha}_{1-s}}{\alpha_{1-s}}\beta_{1-s}-\dot{\beta}_{1-s}\big{)}}{2\beta_{1-s}^{2}}=\frac{\sigma(1-s)^{2}}{2\beta_{1-s}^{2}}+\frac{\dot{\beta}_{1-s}}{\beta_{1-s}}$ (146) $$

Thus,

$$ $\textstyle\int_{t^{\prime}}^{t}\big{(}\kappa_{1-s}+\frac{\sigma(1-s)^{2}-2\eta_{1-s}}{2\beta_{1-s}^{2}}\big{)}\,\mathrm{d}s=\int_{t^{\prime}}^{t}\big{(}\frac{\sigma(1-s)^{2}}{2\beta_{1-s}^{2}}-\partial_{s}\log\beta_{1-s}\big{)}\,\mathrm{d}s=\int_{t^{\prime}}^{t}\frac{\sigma(1-s)^{2}}{2\beta_{1-s}^{2}}\,\mathrm{d}s-\big{(}\log\beta_{1-t}-\log\beta_{1-t^{\prime}}\big{)},$ (147) $$

which means that

$$ $\textstyle\exp\big{(}-\int_{t^{\prime}}^{t}\big{(}\kappa_{1-s}+\frac{\sigma(1-s)^{2}-2\eta_{1-s}}{2\beta_{1-s}^{2}}\big{)}\,\mathrm{d}s\big{)}=\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\sigma(1-s)^{2}}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\frac{\beta_{1-t}}{\beta_{1-t^{\prime}}},$ (148) $\textstyle\frac{\alpha_{1-t^{\prime}}(\sigma(1-t^{\prime})^{2}-2\eta_{1-t^{\prime}})}{2\beta_{1-t^{\prime}}^{2}}\xi_{1-t^{\prime}}(\vec{X}_{t^{\prime}})=\big{(}\frac{\sigma(1-t^{\prime})^{2}}{2\beta_{1-t^{\prime}}^{2}}+\frac{\dot{\beta}_{1-t^{\prime}}}{\beta_{1-t^{\prime}}}-\frac{\dot{\alpha}_{1-t^{\prime}}}{\alpha_{1-t^{\prime}}}\big{)}\xi_{1-t^{\prime}}(\vec{X}_{t^{\prime}}).$ (149) $$

If we define $\bar{\sigma}^{2}(1-s)$ such that $\sigma^{2}(1-s)=2\beta_{1-s}\big{(}\frac{\dot{\alpha}_{1-s}}{\alpha_{1-s}}\beta_{1-s}-\dot{\beta}_{1-s}\big{)}+\chi(1-s)$, we obtain that

$$ $\textstyle\begin{split}\textstyle&\textstyle\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\sigma(1-s)^{2}}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\frac{\beta_{1-t}}{\beta_{1-t^{\prime}}}=\exp\big{(}-\int_{t^{\prime}}^{t}\big{(}\frac{\dot{\alpha}_{1-s}}{\alpha_{1-s}}-\frac{\dot{\beta}_{1-s}}{\beta_{1-s}}+\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\big{)}\,\mathrm{d}s\big{)}\frac{\beta_{1-t}}{\beta_{1-t^{\prime}}}\\ \textstyle&\textstyle=\exp\big{(}\int_{t^{\prime}}^{t}\big{(}\partial_{s}\log\alpha_{1-s}-\partial_{s}\log\beta_{1-s}-\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\big{)}\,\mathrm{d}s\big{)}\frac{\beta_{1-t}}{\beta_{1-t^{\prime}}}=\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}},\end{split}$ (150) $\textstyle\big{(}\frac{\sigma(1-t^{\prime})^{2}}{2\beta_{1-t^{\prime}}^{2}}+\frac{\dot{\beta}_{1-t^{\prime}}}{\beta_{1-t^{\prime}}}-\frac{\dot{\alpha}_{1-t^{\prime}}}{\alpha_{1-t^{\prime}}}\big{)}\xi_{1-t^{\prime}}(\vec{X}_{t^{\prime}})=\frac{\chi(1-t^{\prime})}{2\beta_{1-t^{\prime}}^{2}}\xi_{1-t^{\prime}}(\vec{X}_{t^{\prime}})$ (151) $$

If we plug equations ([150](#S12.E150))-([151](#S12.E151)) into ([148](#S12.E148))-([149](#S12.E149)), and then those into ([145](#S12.E145)), we obtain that

$$ $\textstyle\begin{split}\textstyle\vec{X}_{t}&\textstyle=\vec{X}_{0}\exp\big{(}-\int_{0}^{t}\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\frac{\alpha_{1-t}}{\alpha_{1}}+\int_{0}^{t}\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\frac{\chi(1-t^{\prime})}{2\beta_{1-t^{\prime}}^{2}}\xi_{1-t^{\prime}}(\vec{X}_{t^{\prime}})\,\mathrm{d}t^{\prime}\\ \textstyle&\textstyle\qquad+\int_{0}^{t}\big{(}2\beta_{1-t^{\prime}}\big{(}\frac{\dot{\alpha}_{1-t^{\prime}}}{\alpha_{1-t^{\prime}}}\beta_{1-t^{\prime}}-\dot{\beta}_{1-t^{\prime}}\big{)}+\chi(1-t^{\prime})\big{)}\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\,\mathrm{d}B_{t^{\prime}}.\end{split}$ (152) $$

and if we take the limit $t\to 1^{-}$ and use that $\alpha_{1}=1$,

$$ $\textstyle\begin{split}\textstyle\vec{X}_{1}&\textstyle=\vec{X}_{0}\big{(}\lim_{t\to 1^{-}}\exp\big{(}-\int_{0}^{t}\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\alpha_{1-t}\big{)}+\lim_{t\to 1^{-}}\int_{0}^{t}\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\frac{\chi(1-t^{\prime})}{2\beta_{1-t^{\prime}}^{2}}\xi_{1-t^{\prime}}(\vec{X}_{t^{\prime}})\,\mathrm{d}t^{\prime}\\ \textstyle&\textstyle\qquad+\lim_{t\to 1^{-}}\int_{0}^{t}\big{(}2\beta_{1-t^{\prime}}\big{(}\frac{\dot{\alpha}_{1-t^{\prime}}}{\alpha_{1-t^{\prime}}}\beta_{1-t^{\prime}}-\dot{\beta}_{1-t^{\prime}}\big{)}+\chi(1-t^{\prime})\big{)}\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\,\mathrm{d}B_{t^{\prime}}.\end{split}$ (153) $$

The assumption on $\chi$ in ([25](#S4.E25)) is equivalent, up to a rearrangement of the notation and a flip in the time variable, to the statement that for all $t^{\prime}\in[0,1)$,

$$ $\textstyle\lim_{t\to 1^{-}}\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\alpha_{1-t}=0.$ (154) $$

Hence, under assumption ([25](#S4.E25)), the factor accompanying $\vec{X}_{0}$ in equation ([153](#S12.E153)) is zero. Moreover, this assumption also implies that

$$ $\textstyle\begin{split}\textstyle&\textstyle\lim_{t\to 1^{-}}\int_{0}^{t}\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\frac{\chi(1-t^{\prime})}{2\beta_{1-t^{\prime}}^{2}}\xi_{1-t^{\prime}}(\vec{X}_{t^{\prime}})\,\mathrm{d}t^{\prime}\\ \textstyle&\textstyle=\int_{0}^{1}\big{(}\lim_{t\to 1^{-}}\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\alpha_{1-t}\big{)}\frac{1}{\alpha_{1-t^{\prime}}}\frac{\chi(1-t^{\prime})}{2\beta_{1-t^{\prime}}^{2}}\xi_{1-t^{\prime}}(\vec{X}_{t^{\prime}})\,\mathrm{d}t^{\prime}=0.\end{split}$ (155) $$

If we plug ([154](#S12.E154)) and (LABEL:eq:coefficient_2) into ([153](#S12.E153)), we obtain that

$$ $\textstyle\vec{X}_{1}=\lim_{t\to 1^{-}}\int_{0}^{t}\big{(}2\beta_{1-t^{\prime}}\big{(}\frac{\dot{\alpha}_{1-t^{\prime}}}{\alpha_{1-t^{\prime}}}\beta_{1-t^{\prime}}-\dot{\beta}_{1-t^{\prime}}\big{)}+\chi(1-t^{\prime})\big{)}\exp\big{(}-\int_{t^{\prime}}^{t}\frac{\chi(1-s)}{2\beta_{1-s}^{2}}\,\mathrm{d}s\big{)}\frac{\alpha_{1-t}}{\alpha_{1-t^{\prime}}}\,\mathrm{d}B_{t^{\prime}},$ (156) $$

which shows that $\vec{X}_{1}$ is independent of $\vec{X}_{0}$. Next, we leverage that $\vec{\bm{X}}$ and $\bm{X}$ have equal distributions over trajectories ([Proposition 4](#Thmproposition4)). In particular, the joint distribution of $(\vec{X}_{0},\vec{X}_{1})$ is equal to the joint distribution of $(X_{1},X_{0})$. We conclude that $X_{1}$ and $X_{0}$ are independent, which is the definition of the memorylessness property.
Hence, the assumption ([25](#S4.E25)) is sufficient for memorylessness to hold.

It remains to prove that the assumption ([25](#S4.E25)) is necessary. Looking at equation ([152](#S12.E152)) we deduce that generally, for any $t\in[0,1)$, $\vec{X}_{0}$ and $\vec{X}_{t}$ are not independent, because the first two terms in ([152](#S12.E152)) are different from zero. Thus, if there existed a $t^{\prime}\in[0,1)$ such that the limit ([154](#S12.E154)) is different from zero, then $\vec{X}_{1}$ would not be independent from $\vec{X}_{t^{\prime}}$, which means that in general it would not be independent of $\vec{X}_{0}$ either.

### 12.2 Proof of Theorem 1 : fine-tuning recipe for general noise schedules

The proof of this result relies heavily on the properties of the Hamilton-Jacobi-Bellman equation:

###### Theorem 3 (Hamilton-Jacobi-Bellman equation) .

If we define the infinitesimal generator

$$ $\textstyle\mathcal{L}:=\frac{1}{2}\sum_{i,j=1}^{d}(\sigma\sigma^{\top})_{ij}(t)\partial_{x_{i}}\partial_{x_{j}}+\sum_{i=1}^{d}b_{i}(x,t)\partial_{x_{i}},$ (157) $$

the value function $V$ for the SOC
problem ([12](#S4.E12))-([13](#S4.E13)) solves the following Hamilton-Jacobi-Bellman (HJB) partial differential equation:

$$ $\textstyle\begin{split}\textstyle&\textstyle\partial_{t}V(x,t)=-\mathcal{L}V(x,t)+\frac{1}{2}\|(\sigma^{\top}\nabla V)(x,t)\|^{2}-f(x,t),\\ \textstyle&\textstyle V(x,T)=g(x).\end{split}$ (158) $$

Consider forward SDEs like ([63](#S10.E63)), starting from the distributions $p^{\mathrm{base}}$ and $p^{*}$, where $p^{*}(x)\propto p^{\mathrm{base}}(x)\exp(r(x))$.

$$ $\textstyle\mathrm{d}\vec{X}_{t}=\vec{b}(\vec{X}_{t},t)\,\mathrm{d}t+\sigma(t)\,\mathrm{d}B_{t},\qquad\vec{X}_{0}\sim p^{\mathrm{base}},$ (159) $\textstyle\mathrm{d}\vec{X}^{*}_{t}=\vec{b}^{*}(\vec{X}^{*}_{t},t)\,\mathrm{d}t+\sigma(t)\,\mathrm{d}B_{t},\qquad\vec{X}_{0}\sim p^{*}.$ (160) $$

where the drifts are defined as

$$ $\textstyle\begin{split}\textstyle\vec{b}(x,t)&\textstyle=-\kappa_{1-t}x+\big{(}\frac{\sigma(1-t)^{2}}{2}-\eta_{1-t}\big{)}\mathfrak{s}(x,1-t)=-\kappa_{1-t}x+\big{(}\frac{\sigma(1-t)^{2}}{2}-\eta_{1-t}\big{)}\nabla\log\vec{p}_{t}(x),\\ \textstyle\vec{b}^{*}(x,t)&\textstyle=-\kappa_{1-t}x+\big{(}\frac{\sigma(1-t)^{2}}{2}-\eta_{1-t}\big{)}\mathfrak{s}^{*}(x,1-t)=-\kappa_{1-t}x+\big{(}\frac{\sigma(1-t)^{2}}{2}-\eta_{1-t}\big{)}\nabla\log\vec{p}_{t}^{*}(x),\end{split}$ (161) $$

and $\vec{p}_{t}$, $\vec{p}_{t}^{*}$ are the densities of $X_{t}$, $\vec{X}_{t}$, respectively. $\vec{p}_{t}$, $\vec{p}_{t}^{*}$ satisfy Fokker-Planck equations:

$$ $\textstyle\begin{split}\textstyle\partial_{t}\vec{p}_{t}=\nabla\cdot(\vec{b}(x,t)\vec{p}_{t})+\nabla\cdot(\frac{\sigma(1-t)^{2}}{2}\nabla\vec{p}_{t}),\qquad\vec{p}_{0}=p^{\mathrm{base}},\\ \textstyle\partial_{t}\vec{p}^{*}_{t}=\nabla\cdot(\vec{b}^{*}(x,t)\vec{p}^{*}_{t})+\nabla\cdot(\frac{\sigma(1-t)^{2}}{2}\nabla\vec{p}^{*}_{t}),\qquad\vec{p}_{0}=p^{*}.\end{split}$ (162) $$

Plugging ([161](#S12.E161)) into ([162](#S12.E162)), we obtain

$$ $\textstyle\begin{split}\textstyle\partial_{t}\vec{p}_{t}=\nabla\cdot(\kappa_{1-t}x\vec{p}_{t})+\nabla\cdot\big{(}\eta_{1-t}\nabla\vec{p}_{t}\big{)},\qquad\vec{p}_{0}=p^{\mathrm{base}},\\ \textstyle\partial_{t}\vec{p}^{*}_{t}=\nabla\cdot(\kappa_{1-t}x\vec{p}^{*}_{t})+\nabla\cdot\big{(}\eta_{1-t}\nabla\vec{p}^{*}_{t}\big{)},\qquad\vec{p}_{0}=p^{*}.\end{split}$ (163) $$

We apply the Hopf-Cole transformation to obtain PDEs for $-\log\vec{p}_{t}$ (and $-\log\vec{p}^{*}_{t}$ analogously):

$$ $\textstyle\begin{split}\textstyle-\partial_{t}(-\log\vec{p}_{t})&\textstyle=\frac{\partial_{t}p_{t}}{p_{t}}=\frac{\nabla\cdot(\kappa_{1-t}x\vec{p}_{t})+\nabla\cdot\big{(}\eta_{1-t}\nabla\vec{p}_{t}\big{)}}{p_{t}}\\ \textstyle&\textstyle=\kappa_{1-t}\nabla\cdot x+\kappa_{1-t}\langle x,\nabla\log\vec{p}_{t}\rangle+\eta_{1-t}\frac{\nabla\cdot(\nabla\log\vec{p}_{t}\exp(\log p_{t}))}{p_{t}}\\ \textstyle&\textstyle=\kappa_{1-t}d+\kappa_{1-t}\langle x,\nabla\log\vec{p}_{t}\rangle+\eta_{1-t}\big{(}\Delta\log\vec{p}_{t}+\|\nabla\log\vec{p}_{t}\|^{2}\big{)}.\end{split}$ (164) $$

Hence, if we define $\mathscr{V}(x,t)=-\log\vec{p}_{t}(x)$, $\mathscr{V}^{*}(x,t)=-\log\vec{p}^{*}_{t}(x)$, then $\mathscr{V}$ and $\mathscr{V}^{*}$ satisfy the following Hamilton-Jacobi-Bellman equations:

$$ $\textstyle-\partial_{t}\mathscr{V}$ $\textstyle=\kappa_{1-t}d-\kappa_{1-t}\langle x,\nabla\mathscr{V}\rangle+\eta_{1-t}\big{(}-\Delta\mathscr{V}+\|\nabla\mathscr{V}\|^{2}\big{)},\qquad\mathscr{V}(x,0)=-\log p^{\mathrm{base}}(x),$ (165) $\textstyle-\partial_{t}\mathscr{V}^{*}$ $\textstyle=\kappa_{1-t}d-\kappa_{1-t}\langle x,\nabla\mathscr{V}^{*}\rangle+\eta_{1-t}\big{(}-\Delta\mathscr{V}^{*}+\|\nabla\mathscr{V}^{*}\|^{2}\big{)},\qquad\mathscr{V}^{*}(x,0)=-\log p^{*}(x).$ (166) $$

Now, define $\hat{\mathscr{V}}(x,t)=\mathscr{V}^{*}(x,t)-\mathscr{V}(x,t)$. Subtracting ([166](#S12.E166)) from ([165](#S12.E165)), we obtain

$$ $\textstyle\begin{split}\textstyle-\partial_{t}\hat{\mathscr{V}}&\textstyle=-\kappa_{1-t}\langle x,\nabla\hat{\mathscr{V}}\rangle+\eta_{1-t}\big{(}-\Delta\hat{\mathscr{V}}+\|\nabla\mathscr{V}^{*}\|^{2}-\|\nabla\mathscr{V}\|^{2}\big{)}\\ \textstyle&\textstyle=-\kappa_{1-t}\langle x,\nabla\hat{\mathscr{V}}\rangle+\eta_{1-t}\big{(}-\Delta\hat{\mathscr{V}}+\|\nabla(\hat{\mathscr{V}}+\mathscr{V})\|^{2}-\|\nabla\mathscr{V}\|^{2}\big{)}\\ \textstyle&\textstyle=-\kappa_{1-t}\langle x,\nabla\hat{\mathscr{V}}\rangle+\eta_{1-t}\big{(}-\Delta\hat{\mathscr{V}}+\|\nabla\hat{\mathscr{V}}\|^{2}+2\langle\nabla\mathscr{V},\nabla\hat{\mathscr{V}}\rangle\big{)}\\ \textstyle&\textstyle=\langle-\kappa_{1-t}x+2\eta_{1-t}\nabla\mathscr{V},\nabla\hat{\mathscr{V}}\rangle+\eta_{1-t}\big{(}-\Delta\hat{\mathscr{V}}+\|\nabla\hat{\mathscr{V}}\|^{2}\big{)}\\ \textstyle&\textstyle=\langle-\kappa_{1-t}x-2\eta_{1-t}\mathfrak{s}(x,1-t),\nabla\hat{\mathscr{V}}\rangle+\eta_{1-t}\big{(}-\Delta\hat{\mathscr{V}}+\|\nabla\hat{\mathscr{V}}\|^{2}\big{)},\\ \textstyle\hat{\mathscr{V}}(x,0)&\textstyle=-\log p^{*}(x)+\log p^{\mathrm{base}}(x)=-r(x)+\log\big{(}\int p^{\mathrm{base}}(y)\exp(r(y))\,\mathrm{d}y\big{)}.\end{split}$ (167) $$

Hence, $\hat{\mathscr{V}}$ also satisfies a Hamilton-Jacobi-Bellman equation. If we define $V$ such that $\hat{\mathscr{V}}(x,t)=V(x,1-t)$, we have that

$$ $\textstyle\partial_{t}V=\langle-\kappa_{t}x-2\eta_{t}\mathfrak{s}(x,t),\nabla V\rangle+\eta_{t}\big{(}-\Delta V+\|\nabla V\|^{2}\big{)},\qquad V(x,1)=r(x)-\log\big{(}\int p^{\mathrm{base}}(y)\exp(r(y))\,\mathrm{d}y\big{)}.$ (168) $$

Using [Theorem 3](#Thmtheorem3), we can reverse-engineer
$V$ as the value function of the following SOC
problem:

$$ $\textstyle\min\limits_{u\in\mathcal{U}}\mathbb{E}\big{[}\frac{1}{2}\int_{0}^{1}\|u(X^{u}_{t},t)\|^{2}\,\mathrm{d}t\!-\!r(x)\!+\!\log\big{(}\int p^{\mathrm{base}}(y)\exp(r(y))\,\mathrm{d}y\big{)}\big{]},$ (169) $\textstyle\begin{split}\textstyle\text{s.t.}\leavevmode\nobreak\ \mathrm{d}X^{u}_{t}\!=\!\big{(}\kappa_{t}x+2\eta_{t}\mathfrak{s}(x,t)\!+\!\sqrt{2\eta_{t}}u(X^{u}_{t},t)\big{)}\,\mathrm{d}t\!+\!\sqrt{2\eta_{t}}\mathrm{d}B_{t},\qquad X^{u}_{0}\sim p_{0}.\end{split}$ (170) $$

Note that this SOC problem is equal to the problem ([12](#S4.E12))-([13](#S4.E13)) with the choices $f=0$, $g=-r$, and $\sigma(t)=\sqrt{2\eta_{t}}$.
By equation ([17](#S4.E17)), the optimal control of the problem ([169](#S12.E169))-([170](#S12.E170)) is of the form:

$$ $\textstyle\begin{split}\textstyle u^{*}(x,t)&\textstyle=-\sqrt{2\eta_{t}}\nabla V(x,t)=-\sqrt{2\eta_{t}}\nabla\hat{\mathscr{V}}(x,1-t)=-\sqrt{2\eta_{t}}\big{(}\nabla\mathscr{V}^{*}(x,1-t)-\nabla\mathscr{V}(x,1-t)\big{)}\\ \textstyle&\textstyle=-\sqrt{2\eta_{t}}\big{(}-\nabla\log\vec{p}^{*}_{1-t}(x)+\nabla\log\vec{p}_{1-t}(x)\big{)}=\sqrt{2\eta_{t}}\big{(}\mathfrak{s}^{*}(x,t)-\mathfrak{s}(x,t)\big{)},\end{split}$ (171) $\textstyle\iff\mathfrak{s}^{*}(x,t)=\mathfrak{s}(x,t)+u^{*}(x,t)/\sqrt{2\eta_{t}}.$ (172) $$

As in ([64](#S10.E64)), the backward SDEs corresponding to the forward SDEs ([160](#S12.E160)) take the following form:

$$ $\textstyle\mathrm{d}X^{*}_{t}$ $\textstyle=\big{(}\kappa_{t}X^{*}_{t}+\big{(}\frac{\sigma(t)^{2}}{2}+\eta_{t}\big{)}\mathfrak{s}^{*}(X^{*}_{t},t)\big{)}\mathrm{d}t+\sigma(t)\,\mathrm{d}B_{t},\qquad X^{*}_{0}\sim N(0,I).$ (173) $$

If we plug ([172](#S12.E172)) into this equation, we obtain

$$ $\textstyle\mathrm{d}X^{*}_{t}$ $\textstyle=\big{(}\kappa_{t}X^{*}_{t}+\big{(}\frac{\sigma(t)^{2}}{2}+\eta_{t}\big{)}\big{(}\mathfrak{s}(X^{*}_{t},t)+\frac{u^{*}(X^{*}_{t},t)}{\sqrt{2\eta_{t}}}\big{)}\big{)}\mathrm{d}t+\sigma(t)\,\mathrm{d}B_{t},\qquad X^{*}_{0}\sim N(0,I),$ (174) $\textstyle\iff\mathrm{d}X^{*}_{t}$ $\textstyle=\big{(}b(X^{*}_{t},t)+\frac{\frac{\sigma(t)^{2}}{2}+\eta_{t}}{\sqrt{2\eta_{t}}}u^{*}(X^{*}_{t},t)\big{)}\mathrm{d}t+\sigma(t)\,\mathrm{d}B_{t},\qquad X^{*}_{0}\sim N(0,I).$ (175) $$

where we used that $b(x,t)=\kappa_{t}x+\big{(}\frac{\sigma(t)^{2}}{2}+\eta_{t}\big{)}\mathfrak{s}(x,t)$ by definition in equation ([11](#S3.E11)).

###### Theorem 3 (Hamilton-Jacobi-Bellman equation) .

##### The fine-tuned inference SDE for DDIM

Now, for DDIM, we have that $u^{*}(x,t)=-\sqrt{\frac{\dot{\alpha}_{t}}{\alpha_{t}(1-\alpha_{t})}}(\epsilon^{*}(x,t)-\epsilon^{\mathrm{base}}(x,t))$ by ([26](#S4.E26)). Hence,

$$ $\textstyle\frac{\frac{\sigma(t)^{2}}{2}+\eta_{t}}{\sqrt{2\eta_{t}}}u^{*}(x,t)=-\frac{\frac{\sigma(t)^{2}}{2}+\frac{\dot{\alpha}_{t}}{2\alpha_{t}}}{\sqrt{\frac{\dot{\alpha}_{t}}{\alpha_{t}}}}\sqrt{\frac{\dot{\alpha}_{t}}{\alpha_{t}(1-\alpha_{t})}}(\epsilon^{*}(x,t)-\epsilon^{\mathrm{base}}(x,t))=-\frac{\frac{\sigma(t)^{2}}{2}+\frac{\dot{\alpha}_{t}}{2\alpha_{t}}}{\sqrt{1-\alpha_{t}}}(\epsilon^{*}(x,t)-\epsilon^{\mathrm{base}}(x,t)),$ (176) $\textstyle\begin{split}\textstyle&\textstyle\implies b(x,t)+\frac{\frac{\sigma(t)^{2}}{2}+\eta_{t}}{\sqrt{2\eta_{t}}}u^{*}(x,t)=\frac{\dot{\alpha}_{t}}{2\alpha_{t}}X_{t}-\big{(}\frac{\dot{\alpha}_{t}}{2\alpha_{t}}+\frac{\sigma(t)^{2}}{2}\big{)}\frac{\epsilon^{\mathrm{base}}(X_{t},t)}{\sqrt{1-\alpha_{t}}}-\frac{\frac{\sigma(t)^{2}}{2}+\frac{\dot{\alpha}_{t}}{2\alpha_{t}}}{\sqrt{1-\alpha_{t}}}(\epsilon^{*}(x,t)-\epsilon^{\mathrm{base}}(x,t))\\ \textstyle&\textstyle\qquad\qquad\qquad\qquad\qquad=\frac{\dot{\alpha}_{t}}{2\alpha_{t}}X_{t}-\big{(}\frac{\dot{\alpha}_{t}}{2\alpha_{t}}+\frac{\sigma(t)^{2}}{2}\big{)}\frac{\epsilon^{*}(X_{t},t)}{\sqrt{1-\alpha_{t}}}.\end{split}$ (177) $$

We obtain that the fine-tuned inference SDE for DDIM is

$$ $\textstyle\mathrm{d}X^{*}_{t}$ $\textstyle=\big{(}\frac{\dot{\alpha}_{t}}{2\alpha_{t}}X^{*}_{t}-\big{(}\frac{\dot{\alpha}_{t}}{2\alpha_{t}}+\frac{\sigma(t)^{2}}{2}\big{)}\frac{\epsilon^{*}(X^{*}_{t},t)}{\sqrt{1-\alpha_{t}}}\big{)}\mathrm{d}t+\sigma(t)\,\mathrm{d}B_{t},\qquad X^{*}_{0}\sim N(0,I),$ (178) $$

which is matches the SDE ([6](#S3.E6)) with the choice $\epsilon=\epsilon^{*}$.

##### The fine-tuned inference SDE for Flow Matching

For Flow Matching, we have that $u^{*}(x,t)=\sqrt{\frac{2}{\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}}(v^{*}(x,t)-v^{\mathrm{base}}(x,t))$ by ([27](#S4.E27)). Hence,

$$ $\textstyle\begin{split}\textstyle&\textstyle\frac{\frac{\sigma(t)^{2}}{2}+\eta_{t}}{\sqrt{2\eta_{t}}}u^{*}(x,t)=\frac{\frac{\sigma(t)^{2}}{2}+\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}{\sqrt{2\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}}\sqrt{\frac{2}{\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}}(v^{*}(x,t)-v^{\mathrm{base}}(x,t))\\ \textstyle&\textstyle\qquad\qquad\qquad\ =\big{(}1+\frac{\sigma(t)^{2}}{2\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}\big{)}(v^{*}(x,t)-v^{\mathrm{base}}(x,t)).\end{split}$ (179) $\textstyle\begin{split}\textstyle&\textstyle\implies b(x,t)+\frac{\frac{\sigma(t)^{2}}{2}+\eta_{t}}{\sqrt{2\eta_{t}}}u^{*}(x,t)=v^{\mathrm{base}}(x,t)+\frac{\sigma(t)^{2}}{2\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}\big{(}v^{\mathrm{base}}(x,t)-\frac{\dot{\alpha}_{t}}{\alpha_{t}}x\big{)}\\ \textstyle&\textstyle\qquad\qquad\qquad\qquad\qquad\qquad\quad+\big{(}1+\frac{\sigma(t)^{2}}{2\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}\big{)}(v^{*}(x,t)-v^{\mathrm{base}}(x,t))\\ \textstyle&\textstyle\qquad\qquad\qquad\qquad\qquad\qquad\quad=v^{*}(x,t)+\frac{\sigma(t)^{2}}{2\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}\big{(}v^{*}(x,t)-\frac{\dot{\alpha}_{t}}{\alpha_{t}}x\big{)}.\end{split}$ (180) $$

We obtain that the fine-tuned inference SDE for Flow Matching is

$$ $\textstyle\mathrm{d}X^{*}_{t}=\big{(}v(X^{*}_{t},t)+\frac{\sigma(t)^{2}}{2\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}\big{(}v^{*}(X^{*}_{t},t)-\frac{\dot{\alpha}_{t}}{\alpha_{t}}X^{*}_{t}\big{)}\big{)}\,\mathrm{d}t+\sigma(t)\,\mathrm{d}B_{t},\qquad X^{*}_{0}\sim N(0,I),$ (181) $$

which matches equation ([4](#S2.E4)) with the choice $v=v^{*}$.

## 13 Loss function derivations

### 13.1 Derivation of the Continuous Adjoint method

###### Proposition 6 .

The gradient $\frac{\mathrm{d}\mathcal{L}}{\mathrm{d}\theta}$ of the adjoint loss $\mathcal{L}(u;\bm{X})$ defined in ([28](#S5.E28)) with respect to the parameters $\theta$ of the control can be expressed as in ([32](#S5.E32)).

###### Proof.

First, note that we can write

$$ $\textstyle\begin{split}\textstyle&\textstyle\nabla_{\theta}\mathbb{E}\big{[}\int_{0}^{T}\big{(}\frac{1}{2}\|u_{\theta}(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\big{]}\\ \textstyle&\textstyle=\mathbb{E}\big{[}\int_{0}^{T}\nabla_{\theta}u_{\theta}(X^{u_{\theta}}_{t},t)u_{\theta}(X^{u_{\theta}}_{t},t)\,\mathrm{d}t\big{]}+\nabla_{\theta}\mathbb{E}\big{[}\int_{0}^{T}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\big{]}\rvert_{v=\mathrm{stopgrad}(u_{\theta})}.\end{split}$ (182) $$

To develop the second term, we apply [Lemma 5](#Thmlemma5). Namely, by the Leibniz rule and equation ([187](#S13.E187)), we have that

$$ $\textstyle\begin{split}\textstyle&\textstyle\nabla_{\theta}\mathbb{E}\big{[}\int_{0}^{T}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\big{]}\rvert_{v=\mathrm{stopgrad}(u_{\theta})}\\ \textstyle&\textstyle=\mathbb{E}\big{[}\nabla_{\theta}\big{(}\int_{0}^{T}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\big{)}\rvert_{v=\mathrm{stopgrad}(u_{\theta})}\big{]}\\ \textstyle&\textstyle=\mathbb{E}\big{[}\int_{0}^{T}(\nabla_{\theta}u_{\theta})(X^{u_{\theta}}_{t}(\omega),t)^{\top}\sigma(t)^{\top}a_{t}(\omega)\,\mathrm{d}t\big{]}.\end{split}$ (183) $$

Plugging the right-hand side of this equation into (LABEL:eq:first_eq_cont_adjoint) concludes the proof.
∎

###### Lemma 5 .

Let $v$ be an arbitrary fixed vector field.
The unique solution of the ODE

$$ $\textstyle\begin{split}\textstyle\frac{\mathrm{d}}{\mathrm{d}t}a(t;\bm{X}^{u},u)&\textstyle=-\left[\left(\nabla_{X^{u}_{t}}(b(X^{u}_{t},t)+\sigma(t)u(X^{u}_{t},t))\right)^{\mkern-1.5mu\mathsf{T}}{}a(t;\bm{X}^{u},u)+\nabla_{X^{u}_{t}}\left(f(X^{u}_{t},t)+\frac{1}{2}\|v(X^{u}_{t},t)\|^{2}\right)\right],\end{split}$ (184) $\textstyle a(1;\bm{X}^{u},u)$ $\textstyle=\nabla g(X^{u}_{1}),$ (185) $$

satisfies:

$$ $\textstyle\begin{split}\textstyle&\textstyle a(t;\bm{X}^{u},u):=\nabla_{X^{u}_{t}}\big{(}\int_{t}^{1}\big{(}\frac{1}{2}\|u(X_{t^{\prime}}^{u},t^{\prime})\|^{2}\!+\!f(X_{t^{\prime}}^{u},t^{\prime})\big{)}\,\mathrm{d}t^{\prime}\!+\!g(X^{u}_{1})\big{)},\\ \textstyle&\textstyle\text{where }\bm{X}^{u}\text{ solves }\mathrm{d}X^{u}_{t}=\left(b(X^{u}_{t},t)+\sigma(t)u(X^{u}_{t},t)\right)\,\mathrm{d}t+\sigma(t)\mathrm{d}B_{t}.\end{split}$ (186) $$

Moreover, when $u=u_{\theta}$ is parameterized by $\theta$ we have that

$$ $\textstyle\begin{split}\textstyle\nabla_{\theta}\big{(}\int_{0}^{T}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\big{)}=\int_{0}^{T}(\nabla_{\theta}u_{\theta})(X^{u_{\theta}}_{t}(\omega),t)\sigma(t)^{\top}a_{t}(\omega)\,\mathrm{d}t.\end{split}$ (187) $$

###### Proof.

We use an approach based on Lagrange multipliers which mirrors and extends the derivation of the adjoint ODE (Domingo-Enrich et al., , Lemma 8). For shortness, we use the notation $\tilde{b}_{\theta}(x,t):=b(x,t)+\sigma(t)u_{\theta}(x,t)$. Define a process $a:\Omega\times[0,T]\to\mathbb{R}^{d}$ such that for any $\omega\in\Omega$, $a(\omega,\cdot)$ is differentiable. For a given $\omega\in\Omega$, we can write

$$ $\textstyle\begin{split}\textstyle&\textstyle\int_{0}^{T}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\\ \textstyle&\textstyle=\int_{0}^{T}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\\ \textstyle&\textstyle\qquad-\int_{0}^{T}\langle a_{t}(\omega),(dX^{u_{\theta}}_{t}(\omega)-\tilde{b}_{\theta}(X^{u_{\theta}}_{t}(\omega),t)\,\mathrm{d}t-\sigma(t)\,\mathrm{d}B_{t})\rangle.\end{split}$ (188) $$

By stochastic integration by parts (Domingo-Enrich et al., , Lemma 9), we have that

$$ $\textstyle\begin{split}\textstyle\int_{0}^{T}\langle a_{t}(\omega),dX^{u_{\theta}}_{t}(\omega)\rangle=\langle a_{T}(\omega),X^{u_{\theta}}_{T}(\omega)\rangle-\langle a_{0}(\omega),X^{u_{\theta}}_{0}(\omega)\rangle-\int_{0}^{T}\langle X^{u_{\theta}}_{t}(\omega),\frac{da_{t}}{dt}(\omega)\rangle\,\mathrm{d}t.\end{split}$ (189) $$

Hence, if $X^{u_{\theta}}_{0}=x_{0}$ is the initial condition, we have that(^8^88Unlike (Domingo-Enrich et al., , Lemma 8), we use the convention that a Jacobian matrix $J=\nabla_{x}v(x)$ is defined as $J_{ij}=\frac{\partial v_{i}(x)}{\partial x_{j}}$. Their definition of $\nabla_{x}v$ is the transpose of ours.)

$$ $\textstyle\begin{split}\textstyle&\textstyle\nabla_{x_{0}}\big{(}\int_{0}^{T}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\big{)}\\ \textstyle&\textstyle=\nabla_{x_{0}}\big{(}\int_{0}^{T}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\\ \textstyle&\textstyle\qquad\quad-\langle a_{T}(\omega),X^{u_{\theta}}_{T}(\omega)\rangle+\langle a_{0}(\omega),X^{u_{\theta}}_{0}(\omega)\rangle+\int_{0}^{T}\big{(}\langle a_{t}(\omega),\tilde{b}_{\theta}(X^{u_{\theta}}_{t}(\omega),t)\rangle+\langle\frac{da_{t}}{dt}(\omega),X^{u_{\theta}}_{t}(\omega)\rangle\big{)}\,\mathrm{d}t\\ \textstyle&\textstyle\qquad\quad+\int_{0}^{T}\langle a_{t}(\omega),\sigma(t)\,\mathrm{d}B_{t}\rangle\big{)}\\ \textstyle&\textstyle=\int_{0}^{T}\nabla_{x_{0}}X^{u_{\theta}}_{t}(\omega)^{\top}\nabla_{x}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t}(\omega),t)\big{)}\,\mathrm{d}t+\nabla_{x_{0}}X^{u_{\theta}}_{T}(\omega)^{\top}\nabla_{x}g(X^{u_{\theta}}_{T}(\omega))\\ \textstyle&\textstyle\qquad\quad-\nabla_{x_{0}}X^{u_{\theta}}_{T}(\omega)^{\top}a_{T}(\omega)+\nabla_{x_{0}}X^{u_{\theta}}_{0}(\omega)^{\top}a_{0}(\omega)\\ \textstyle&\textstyle\qquad\quad+\int_{0}^{T}\big{(}\nabla_{x_{0}}X^{u_{\theta}}_{t}(\omega)^{\top}\nabla_{x}\tilde{b}_{\theta}(X^{u_{\theta}}_{t}(\omega),t)^{\top}a_{t}(\omega)+\nabla_{x_{0}}X^{u_{\theta}}_{t}(\omega)^{\top}\frac{da_{t}}{dt}(\omega)\big{)}\,\mathrm{d}t\\ \textstyle&\textstyle=\int_{0}^{T}\nabla_{x_{0}}X^{u_{\theta}}_{t}(\omega)^{\top}\big{(}\nabla_{x}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t}(\omega),t)\big{)}+\nabla_{x}\tilde{b}_{\theta}(X^{u_{\theta}}_{t}(\omega),t)^{\top}a_{t}(\omega)+\frac{da_{t}}{dt}(\omega)\big{)}\,\mathrm{d}t\\ \textstyle&\textstyle\qquad\quad+\nabla_{x_{0}}X^{u_{\theta}}_{T}(\omega)^{\top}\big{(}\nabla_{x}g(X^{u_{\theta}}_{T}(\omega))-a_{T}(\omega)\big{)}+a_{0}(\omega).\end{split}$ (190) $$

In the last line we used that $\nabla_{x_{0}}X^{u_{\theta}}_{0}(\omega)=\nabla_{x_{0}}x_{0}=\mathrm{I}$.
If choose $a$ such that

$$ $\textstyle\begin{split}\textstyle da_{t}(\omega)&\textstyle=\big{(}-\nabla_{x}\tilde{b}_{\theta}(X^{u_{\theta}}_{t}(\omega),t)^{\top}a_{t}(\omega)-\nabla_{x}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t}(\omega),t)\big{)}\big{)}\,\mathrm{d}t,\\ \textstyle a_{T}(\omega)&\textstyle=\nabla_{x}g(X^{u_{\theta}}_{T}(\omega)),\end{split}$ (191) $$

which is the ODE ([184](#S13.E184))-([185](#S13.E185)), then we obtain that

$$ $\textstyle\nabla_{x_{0}}\big{(}\int_{0}^{T}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\big{)}=a_{0}(\omega)$ (192) $$

Without loss of generality, this argument can be extended from $t=0$ to an arbitrary $t\in[0,1]$, which proves the first statement of the lemma.

To prove ([187](#S13.E187)), we similarly write

$$ $\textstyle\begin{split}\textstyle&\textstyle\nabla_{\theta}\big{(}\int_{0}^{T}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\big{)}\\ \textstyle&\textstyle=\nabla_{\theta}\big{(}\int_{0}^{T}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\\ \textstyle&\textstyle\qquad\quad-\langle a_{T}(\omega),X^{u_{\theta}}_{T}(\omega)\rangle+\langle a_{0}(\omega),X^{u_{\theta}}_{0}(\omega)\rangle+\int_{0}^{T}\big{(}\langle a_{t}(\omega),\tilde{b}_{\theta}(X^{u_{\theta}}_{t}(\omega),t)\rangle+\langle\frac{da_{t}}{dt}(\omega),X^{u_{\theta}}_{t}(\omega)\rangle\big{)}\,\mathrm{d}t\\ \textstyle&\textstyle\qquad\quad+\int_{0}^{T}\langle a_{t}(\omega),\sigma(t)\,\mathrm{d}B_{t}\rangle\big{)}\\ \textstyle&\textstyle=\int_{0}^{T}\nabla_{\theta}X^{u_{\theta}}_{t}(\omega)^{\top}\nabla_{x}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t}(\omega),t)\big{)}\,\mathrm{d}t+\nabla_{\theta}X^{u_{\theta}}_{T}(\omega)^{\top}\nabla_{x}g(X^{u_{\theta}}_{T}(\omega))\\ \textstyle&\textstyle\qquad\quad-\nabla_{\theta}X^{u_{\theta}}_{T}(\omega)^{\top}a_{T}(\omega)+\nabla_{\theta}X^{u_{\theta}}_{0}(\omega)^{\top}a_{0}(\omega)\\ \textstyle&\textstyle\qquad\quad+\int_{0}^{T}\big{(}\nabla_{\theta}X^{u_{\theta}}_{t}(\omega)^{\top}\nabla_{x}\tilde{b}_{\theta}(X^{u_{\theta}}_{t}(\omega),t)^{\top}a_{t}(\omega)+\nabla_{\theta}\tilde{b}_{\theta}(X^{u_{\theta}}_{t}(\omega),t)^{\top}a_{t}(\omega)+\nabla_{\theta}X^{u_{\theta}}_{t}(\omega)^{\top}\frac{da_{t}}{dt}(\omega)\big{)}\,\mathrm{d}t\\ \textstyle&\textstyle=\int_{0}^{T}\nabla_{\theta}X^{u_{\theta}}_{t}(\omega)^{\top}\big{(}\nabla_{x}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t}(\omega),t)\big{)}+\nabla_{x}\tilde{b}_{\theta}(X^{u_{\theta}}_{t}(\omega),t)^{\top}a_{t}(\omega)+\frac{da_{t}}{dt}(\omega)\big{)}\,\mathrm{d}t\\ \textstyle&\textstyle\qquad\quad+\nabla_{\theta}X^{u_{\theta}}_{T}(\omega)^{\top}\big{(}\nabla_{x}g(X^{u_{\theta}}_{T}(\omega))-a_{T}(\omega)\big{)}+\int_{0}^{T}(\nabla_{\theta}\tilde{b}_{\theta})(X^{u_{\theta}}_{t}(\omega),t)^{\top}a_{t}(\omega)\,\mathrm{d}t.\end{split}$ (193) $$

In the last line we used that $\nabla_{\theta}X^{u_{\theta}}_{0}(\omega)=\nabla_{\theta}x=0$.
When $a$ satisfies ([191](#S13.E191)), we obtain that

$$ $\textstyle\begin{split}\textstyle&\textstyle\nabla_{\theta}\big{(}\int_{0}^{T}\big{(}\frac{1}{2}\|v(X^{u_{\theta}}_{t},t)\|^{2}\!+\!f(X^{u_{\theta}}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{u_{\theta}}_{T})\big{)}\\ \textstyle&\textstyle=\int_{0}^{T}(\nabla_{\theta}\tilde{b}_{\theta})(X^{u_{\theta}}_{t}(\omega),t)a_{t}(\omega)\,\mathrm{d}t=\int_{0}^{T}(\nabla_{\theta}u_{\theta})(X^{u_{\theta}}_{t}(\omega),t)^{\top}\sigma(t)^{\top}a_{t}(\omega)\,\mathrm{d}t.\end{split}$ (194) $$

The last equality holds because $\tilde{b}_{\theta}(x,t):=b(x,t)+\sigma(t)u_{\theta}(x,t)$.
∎

###### Proposition 6 .

###### Proof.

###### Lemma 5 .

###### Proof.

### 13.2 Proof of Proposition 2 : Theoretical guarantees of the basic Adjoint Matching loss

Let $\bar{u}=\texttt{stopgrad}(u_{\theta})$. We can rewrite equation ([32](#S5.E32)) as:

$$ $\textstyle\nabla_{\theta}\mathcal{L}(u_{\theta};\bm{X}^{\bar{u}})$ $\textstyle=\frac{1}{2}\int_{0}^{1}\nabla_{\theta}\lVert u_{\theta}(X^{\bar{u}}_{t},t)\rVert^{2}\mathrm{d}t+\int_{0}^{1}\nabla_{\theta}u(X^{\bar{u}}_{t},t)^{\mkern-1.5mu\mathsf{T}}{}\sigma(t)^{\mkern-1.5mu\mathsf{T}}{}a(t;\bm{X}^{\bar{u}},\bar{u})\mathrm{d}t$ (195) $\textstyle=\frac{1}{2}\int_{0}^{1}\nabla_{\theta}\lVert u_{\theta}(X^{\bar{u}}_{t},t)+\sigma(t)^{\mkern-1.5mu\mathsf{T}}{}a(t;\bm{X}^{\bar{u}},\bar{u})\rVert^{2}\mathrm{d}t=\nabla_{\theta}\mathcal{L}_{\mathrm{Basic-Adj-Match}}(u_{\theta};\bm{X}^{\bar{u}})$ (196) $$

This proves the first statement of the proposition. To prove that the only critical point of the expected basic Adjoint Matching loss is the optimal control, we first compute the first variation of $\mathbb{E}[\mathcal{L}_{\mathrm{Basic-Adj-Match}}]$. Letting $v:\mathbb{R}^{d}\times[0,T]\to\mathbb{R}^{d}$ be arbitrary, we have that

$$ $\textstyle\begin{split}\textstyle&\textstyle\frac{\mathrm{d}}{\mathrm{d}\epsilon}\mathbb{E}[\mathcal{L}_{\mathrm{Basic-Adj-Match}}(u+\epsilon v;\bm{X}^{\bar{u}})]=\frac{\mathrm{d}}{\mathrm{d}\epsilon}\mathbb{E}\big{[}\frac{1}{2}\int_{0}^{T}\|(u+\epsilon v)(X^{\bar{u}}_{t},t)+\sigma(t)^{\top}a(t,X^{\bar{u}},\bar{u})\|^{2}\,\mathrm{d}t\big{]}\\ \textstyle&\textstyle=\mathbb{E}\big{[}\int_{0}^{T}\langle v(X^{\bar{u}}_{t},t),u(X^{\bar{u}}_{t},t)+\sigma(t)^{\top}a(t,X^{\bar{u}},\bar{u})\rangle\,\mathrm{d}t\big{]}\\ \textstyle&\textstyle=\mathbb{E}\big{[}\int_{0}^{T}\langle v(X^{\bar{u}}_{t},t),u(X^{\bar{u}}_{t},t)+\sigma(t)^{\top}\mathbb{E}\big{[}a(t,X^{\bar{u}},\bar{u})|X^{\bar{u}}_{t}\big{]}\rangle\,\mathrm{d}t\big{]}\\ \textstyle&\textstyle\implies\frac{\delta}{\delta u}\mathbb{E}[\mathcal{L}_{\mathrm{Basic-Adj-Match}}(u)(x,t)=u(x,t)+\mathbb{E}\big{[}a(t,X^{\bar{u}},\bar{u})|X^{\bar{u}}_{t}=x\big{]}\end{split}$ (197) $$

Hence, critical points satisfy that

$$ $\textstyle\begin{split}\textstyle u(x,t)&\textstyle=-\sigma(t)^{\top}\mathbb{E}[a(t,X^{u},u)|X^{u}_{t}=x]=-\sigma(t)^{\top}\mathbb{E}\big{[}\nabla_{X^{v}_{t}}\int_{t}^{T}\big{(}\frac{1}{2}\|v(X^{v}_{t},t)\|^{2}\!+\!f(X^{v}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{v}_{T})|X^{v}_{0}=x\big{]}\\ \textstyle&\textstyle=-\sigma(t)^{\top}\nabla_{x}\mathbb{E}\big{[}\int_{t}^{T}\big{(}\frac{1}{2}\|v(X^{v}_{t},t)\|^{2}\!+\!f(X^{v}_{t},t)\big{)}\,\mathrm{d}t\!+\!g(X^{v}_{T})|X^{v}_{0}=x\big{]}=-\sigma(t)^{\top}\nabla J(u;x,t),\end{split}$ (198) $$

In this equation, the second equality holds by equation (LABEL:eq:adjoint_grad_x) from [Lemma 5](#Thmlemma5), and the third equality holds by the Leibniz rule.

[Lemma 6](#Thmlemma6) shows that any control $u$ that satisfies ([198](#S13.E198)) is equal to the optimal control, which concludes the proof.

###### Lemma 6 .

Suppose that for any $x\in\mathbb{R}^{d}$, $t\in[0,T]$, $u(x,t)=-\sigma(t)^{\top}\nabla_{x}J(u;x,t)$. Then, $J(u;\cdot,\cdot)$ satisfies the Hamilton-Jacobi-Bellman equation (LABEL:eq:HJB_setup). By the uniqueness of the solution to the HJB equation, we have that $J(u;x,t)=V(x,t)$ for any $x\in\mathbb{R}^{d}$, $t\in[0,T]$. Hence, $u(x,t)=-\sigma(t)^{\top}\nabla_{x}V(x,t)$ is the optimal control.

###### Proof.

Since $J(u;x,t)=\mathbb{E}\big{[}\int_{t}^{T}\big{(}\frac{1}{2}\|u(X^{u}_{t},t)\|^{2}+f(X^{u}_{t},t)\big{)}\,ds+g(X^{u}_{T})|X^{u}_{t}=x\big{]}$, we have that

$$ $\textstyle J(u;x,t)=\mathbb{E}\big{[}J(u;X^{u}_{t+\Delta t},t+\Delta t)|X_{t}=x\big{]}+\mathbb{E}\big{[}\int_{t}^{t+\Delta t}\big{(}\frac{1}{2}\|u(X^{u}_{t},t)\|^{2}+f(X^{u}_{t},t)\big{)}\,ds|X_{t}=x\big{]},$ (199) $$

which means that

$$ $\displaystyle 0=\frac{\mathbb{E}[J(u;X^{u}_{t+\Delta t},t+\Delta t)|X_{t}=x]-J(u;x,t)}{\Delta t}+\frac{\mathbb{E}\big{[}\int_{t}^{t+\Delta t}\big{(}\frac{1}{2}\|u(X^{u}_{t},t)\|^{2}+f(X^{u}_{t},t)\big{)}\,ds|X_{t}=x\big{]}}{\Delta t}$ (200) $$

Recall that the generator $\mathcal{T}^{u}$ of the controlled SDE ([13](#S4.E13)) takes the form:

$$ $\textstyle\begin{split}\textstyle\mathcal{T}^{u}f(x,t)&\textstyle:=\lim_{\Delta t\to 0}\frac{\mathbb{E}\big{[}f(X^{u}_{t+\Delta t},t)|X_{t}=x\big{]}-f(x,t)}{\Delta t}\\ \textstyle&\textstyle=\partial_{t}f(x,t)+\langle\nabla f(x,t),b(x,t)+\sigma(t)u(x,t)\rangle+\mathrm{Tr}\big{(}\frac{\sigma(t)\sigma(t)^{\top}}{2}\nabla^{2}f(x,t)\big{)}\end{split}$ (201) $$

Hence, if we take the limit $\Delta t\to 0$ on equation ([200](#S13.E200)), we obtain that:

$$ $\textstyle\begin{split}\textstyle 0&\textstyle=\mathcal{T}^{u}J(u;x,t)+\frac{1}{2}\|u(x,t)\|^{2}+f(x,t)\\ \textstyle&\textstyle=\partial_{t}J(u;x,t)+\langle\nabla J(u;x,t),b(x,t)+\sigma(t)u(x,t)\rangle+\mathrm{Tr}\big{(}\frac{\sigma(t)\sigma(t)^{\top}}{2}\nabla^{2}J(u;x,t)\big{)}+\frac{1}{2}\|u(x,t)\|^{2}+f(x,t).\end{split}$ (202) $$

Now using that $u(x,t)=-\sigma(t)^{\top}\nabla_{x}J(u;x,t)$, we have that

$$ $\textstyle\begin{split}\textstyle\langle\nabla J(u;x,t),\sigma(t)u(x,t)\rangle+\frac{1}{2}\|u(x,t)\|^{2}&\textstyle=-\|\sigma(t)^{\top}\nabla_{x}J(u;x,t)\|^{2}+\frac{1}{2}\|\sigma(t)^{\top}\nabla_{x}J(u;x,t)\|^{2}\\ \textstyle&\textstyle=-\frac{1}{2}\|\sigma(t)^{\top}\nabla_{x}J(u;x,t)\|^{2}.\end{split}$ (203) $$

Plugging this back into ([202](#S13.E202)), we obtain that

$$ $\textstyle\begin{split}\textstyle 0&\textstyle=\partial_{t}J(u;x,t)+\langle\nabla J(u;x,t),b(x,t)\rangle+\mathrm{Tr}\big{(}\frac{\sigma(t)\sigma(t)^{\top}}{2}\nabla^{2}J(u;x,t)\big{)}-\frac{1}{2}\|\sigma(t)^{\top}\nabla_{x}J(u;x,t)\|^{2}+f(x,t).\end{split}$ (204) $$

And since $J(u;x,T)=g(x)$ by construction, we conclude that $J(u;x,t)$ satisfies the HJB equation (LABEL:eq:HJB_setup).
∎

###### Lemma 6 .

###### Proof.

### 13.3 Theoretical guarantees of the Adjoint Matching loss

###### Proposition 7 (Theoretical guarantee of the Adjoint Matching loss) .

The only critical point of the loss $\mathbb{E}[\mathcal{L}_{\mathrm{Adj-Match}}]$ is the optimal control $u^{*}$.

###### Proof.

If $\tilde{a}(t,X^{v}):=\tilde{a}(\omega,t)$ is the solution of the Lean Adjoint ODE ([38](#S5.E38))-([39](#S5.E39)),
we obtain that $\mathbb{E}[\tilde{a}(t,X^{v})|X^{v}_{t}]$ is a solution of

$$ $\textstyle\frac{d\mathbb{E}[\tilde{a}(t,X^{v})|X^{v}_{t}]}{dt}$ $\textstyle=-\nabla_{x}b(X^{v}_{t},t)\mathbb{E}[\tilde{a}(t,X^{v})|X^{v}_{t}]-\nabla_{x}f(X^{v}_{t},t),$ (205) $\textstyle\mathbb{E}[\tilde{a}(T,X^{v})|X^{v}_{T}]$ $\textstyle=\nabla g(X^{v}_{T}),$ (206) $$

Let us rewrite $\mathbb{E}[\mathcal{L}_{\mathrm{Adj-Match}}]$ as follows:

$$ $\textstyle\begin{split}\textstyle\mathbb{E}[\mathcal{L}_{\mathrm{Adj-Match}}(u)]&\textstyle:=\mathbb{E}\big{[}\int_{0}^{T}\big{\|}u(X^{v}_{t},t)+\sigma(t)^{\top}\mathbb{E}\big{[}\tilde{a}(t,X^{v})|X^{v}_{t}\big{]}\big{\|}^{2}\,\mathrm{d}t\big{]}\rvert_{v=\mathrm{stopgrad}(u)}\\ \textstyle&\textstyle\qquad+\mathbb{E}\big{[}\int_{0}^{T}\big{\|}\sigma(t)^{\top}\big{(}\mathbb{E}\big{[}\tilde{a}(t,X^{v})|X^{v}_{t}\big{]}-\tilde{a}(t,X^{v})\big{)}\big{\|}^{2}\,\mathrm{d}t\big{]}\rvert_{v=\mathrm{stopgrad}(u)},\end{split}$ (207) $$

Now, suppose that $\hat{u}$ is a critical point of $\mathbb{E}[\mathcal{L}_{\mathrm{Adj-Match}}]$. By definition, this implies that the first variation of $\mathbb{E}[\mathcal{L}_{\mathrm{Adj-Match}}]$ is zero. Using ([207](#S13.E207)), we can write this as follows:

$$ $\textstyle 0=\frac{\delta}{\delta u}\mathbb{E}[\mathcal{L}_{\mathrm{Adj-Match}}(\hat{u})](x)=2\big{(}\hat{u}(x,t)+\sigma(t)^{\top}\mathbb{E}[\tilde{a}(t,X^{\hat{u}})|X^{\hat{u}}_{t}=x]\big{)},$ (208) $\textstyle\implies\hat{u}(x,t)=-\sigma(t)^{\top}\mathbb{E}[\tilde{a}(t,X^{\hat{u}})|X^{\hat{u}}_{t}=x].$ (209) $$

Hence, we have

$$ $\textstyle\begin{split}\textstyle-\nabla_{x}\hat{u}(X^{\hat{u}}_{t},t)\sigma(t)^{\top}\mathbb{E}[\tilde{a}(t,X^{\hat{u}})|X^{\hat{u}}_{t}]-\nabla_{x}{\hat{u}}(X^{\hat{u}}_{t},t){\hat{u}}(X^{\hat{u}}_{t},t)=0.\end{split}$ (210) $$

Adding this to the right-hand side of ([205](#S13.E205)), we obtain that $\mathbb{E}[\tilde{a}(t,X^{\hat{u}})|X^{\hat{u}}_{t}]$ also solves the ODE

$$ $\textstyle\begin{split}\textstyle\frac{d\mathbb{E}[\tilde{a}(t,X^{\hat{u}})|X^{\hat{u}}_{t}]}{dt}&\textstyle=-\nabla_{x}(b(X^{\hat{u}}_{t},t)+\sigma(t){\hat{u}}(X^{\hat{u}}_{t},t))\mathbb{E}[\tilde{a}(t,X^{\hat{u}})|X^{\hat{u}}_{t}]\\ \textstyle&\textstyle\quad-\nabla_{x}(f(X^{\hat{u}}_{t},t)+\frac{1}{2}\|{\hat{u}}(X^{\hat{u}}_{t},t)\|^{2}),\end{split}$ (211) $\textstyle\mathbb{E}[\tilde{a}(T,X^{\hat{u}})|X^{\hat{u}}_{T}]$ $\textstyle=\nabla g(X^{\hat{u}}_{T}),$ (212) $$

If $a(t,X^{\hat{u}}):=a(\omega,t)$ is the solution of the Adjoint ODE ([30](#S5.E30))-([31](#S5.E31)), we obtain that $\mathbb{E}[a(t,X^{\hat{u}})|X^{\hat{u}}_{t}]$ is a solution of

$$ $\textstyle\begin{split}\textstyle\frac{d\mathbb{E}[a(t,X^{\hat{u}})|X^{\hat{u}}_{t}]}{dt}&\textstyle=-\nabla_{x}(b(X^{\hat{u}}_{t},t)+\sigma(t){\hat{u}}(X^{\hat{u}}_{t},t))\mathbb{E}[a(t,X^{\hat{u}})|X^{\hat{u}}_{t}]\\ \textstyle&\textstyle\quad-\nabla_{x}(f(X^{\hat{u}}_{t},t)+\frac{1}{2}\|{\hat{u}}(X^{\hat{u}}_{t},t)\|^{2}),\end{split}$ (213) $\textstyle\mathbb{E}[a(T,X^{\hat{u}})|X^{\hat{u}}_{T}]$ $\textstyle=\nabla g(X^{\hat{u}}_{T}),$ (214) $$

Remark that ([213](#S13.E213))-([214](#S13.E214)) is the same ODE as ([211](#S13.E211))-([212](#S13.E212)). By uniqueness of ODE solutions, we obtain that $\mathbb{E}[\tilde{a}(t,X^{\hat{u}})|X^{\hat{u}}_{t}]=\mathbb{E}[a(t,X^{\hat{u}})|X^{\hat{u}}_{t}]$ for all $t\in[0,T]$ when ${\hat{u}}$ is a critical point of $\mathbb{E}[\mathcal{L}_{\mathrm{Adj-Match}}]$.
Since we can reexpress the basic Adjoint Matching loss as

$$ $\textstyle\begin{split}\textstyle\mathbb{E}[\mathcal{L}_{\mathrm{Basic-Adj-Match}}(u)]&\textstyle:=\mathbb{E}\big{[}\int_{0}^{T}\big{\|}u(X^{v}_{t},t)+\sigma(t)^{\top}\mathbb{E}\big{[}a(t,X^{v})|X^{v}_{t}\big{]}\big{\|}^{2}\,\mathrm{d}t\big{]}\rvert_{v=\mathrm{stopgrad}(u)}\\ \textstyle&\textstyle\qquad+\mathbb{E}\big{[}\int_{0}^{T}\big{\|}\sigma(t)^{\top}\big{(}\mathbb{E}\big{[}a(t,X^{v})|X^{v}_{t}\big{]}-a(t,X^{v})\big{)}\big{\|}^{2}\,\mathrm{d}t\big{]}\rvert_{v=\mathrm{stopgrad}(u)},\end{split}$ (215) $$

we obtain that when ${\hat{u}}$ is a critical point of $\mathcal{L}_{\mathrm{Adj-Match}}$,

$$ $\textstyle\frac{\mathrm{d}}{\mathrm{d}u}\mathbb{E}[\mathcal{L}_{\mathrm{Basic-Adj-Match}}(\hat{u})](x)$ $\textstyle=2\big{(}\hat{u}(x,t)+\sigma(t)^{\top}\mathbb{E}[a(t,X^{\hat{u}})|X^{\hat{u}}_{t}=x]\big{)}$ (216) $\textstyle=2\big{(}\hat{u}(x,t)+\sigma(t)^{\top}\mathbb{E}[\tilde{a}(t,X^{\hat{u}})|X^{\hat{u}}_{t}=x]\big{)}=0.$ (217) $$

Thus, we deduce that the critical points of $\mathbb{E}[\mathcal{L}_{\mathrm{Adj-Match}}]$ are critical points of $\mathbb{E}[\mathcal{L}_{\mathrm{Basic-Adj-Match}}]$. By [Proposition 2](#Thmproposition2), $\mathbb{E}[\mathcal{L}_{\mathrm{Basic-Adj-Match}}]$ has a single critical point, which is the optimal control $u^{*}$, which concludes the proof of the statement for $\mathbb{E}[\mathcal{L}_{\mathrm{Adj-Match}}]$.
∎

###### Proposition 7 (Theoretical guarantee of the Adjoint Matching loss) .

###### Proof.

### 13.4 Pseudo-code of Adjoint Matching for DDIM fine-tuning

Figure: Algorithm 2 Adjoint Matching for fine-tuning DDIM

Note that for each pair of equations ([218](#S13.E218))-([219](#S13.E219)), ([220](#S13.E220))-([221](#S13.E221)),
([222](#S13.E222))-([223](#S13.E223)), the first equation corresponds to the updates in the DDPM paper, while the second equation is an Euler-Maruyama / Euler discretization of the continuous-time object.
To check that both discretizations are equal up to first order, remark that

$$ $\textstyle\sqrt{\frac{\bar{\alpha}_{k+1}}{\bar{\alpha}_{k}}}=\sqrt{1+\frac{\bar{\alpha}_{k+1}-\bar{\alpha}_{k}}{\bar{\alpha}_{k}}}\approx 1+\frac{\bar{\alpha}_{k+1}-\bar{\alpha}_{k}}{2\bar{\alpha}_{k}}+O((\bar{\alpha}_{k+1}-\bar{\alpha}_{k})^{2}).$ (224) $$

## 14 Adapting diffusion fine-tuning baselines to flow matching

### 14.1 Adapting ReFL (Xu et al., 2023 ) to flow matching

Reward Feedback Learning (ReFL) is a diffusion fine-tuning algorithm introduced by Xu et al. () which tries to increase the reward on denoised samples. Namely, if $\bm{X}=(X_{t})_{t\in[0,1]}$ is the solution of the DDPM SDE ([7](#S3.E7)), we can denoise $X_{t}$ as

$$ $\textstyle\hat{X}_{1}(X_{t})=\frac{X_{t}-\sqrt{1-\bar{\alpha}_{t}}\epsilon(X_{t},t)}{\sqrt{\bar{\alpha}_{t}}}.$ (225) $$

This equation follows from the stochastic interpolant equation ([2](#S2.E2)) if we replace $\bar{X}_{0}$ with the noise predictor $\epsilon(X_{t},t)$. And then, the ReFL optimization update is based on the gradient:

$$ $\textstyle\nabla_{\theta}r(\hat{X}_{1}(X_{t}))=\nabla_{\theta}r\big{(}\frac{X_{t}-\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta}(X_{t},t)}{\sqrt{\bar{\alpha}_{t}}}\big{)},$ (226) $$

where the trajectories have been detached.

To adapt ReFL to Flow Matching, we need to express the denoiser map in terms of the vector field $v$. We have that

$$ $\textstyle\begin{split}\textstyle v(x,t)&\textstyle=\mathbb{E}\big{[}\dot{\beta}_{t}\bar{X}_{0}+\dot{\alpha}_{t}\bar{X}_{1}\big{|}\beta_{t}\bar{X}_{0}+\alpha_{t}\bar{X}_{1}=x\big{]}=\mathbb{E}\big{[}\frac{\dot{\beta}_{t}}{\beta_{t}}\big{(}\beta_{t}\bar{X}_{0}+\alpha_{t}\bar{X}_{1}\big{)}+\big{(}\dot{\alpha}_{t}-\frac{\dot{\beta}_{t}}{\beta_{t}}\alpha_{t}\big{)}\bar{X}_{1}\big{|}\beta_{t}\bar{X}_{0}+\alpha_{t}\bar{X}_{1}=x\big{]}\\ \textstyle&\textstyle=\frac{\dot{\beta}_{t}}{\beta_{t}}x+\big{(}\dot{\alpha}_{t}-\frac{\dot{\beta}_{t}}{\beta_{t}}\alpha_{t}\big{)}\hat{X}_{1}(x,t).\end{split}$ (227) $$

where we defined the denoiser map $\hat{X}_{1}(x,t):=\mathbb{E}\big{[}\bar{X}_{1}|\beta_{t}\bar{X}_{0}+\alpha_{t}\bar{X}_{1}=x\big{]}$.
Hence,

$$ $\textstyle\hat{X}_{1}(x,t)=\frac{v(x,t)-\frac{\dot{\beta}_{t}}{\beta_{t}}x}{\dot{\alpha}_{t}-\frac{\dot{\beta}_{t}}{\beta_{t}}\alpha_{t}}.$ (228) $$

### 14.2 Adapting Diffusion-DPO (Wallace et al., 2023a ) to flow matching

The Diffusion-DPO loss assumes access to ranked pairs of generated samples $x_{1}^{w}\succ x_{1}^{l}$, where $x^{w}$ and $x^{l}$ are the winning and losing samples. For DDPM, the loss implemented in practice reads (Wallace et al., , Eq. 46):

$$ $\textstyle\begin{split}\textstyle L_{\mathrm{DPO}}(\theta)&\textstyle=-\mathbb{E}_{(x_{1}^{w},x_{1}^{l})\sim\mathcal{D},k\sim U[0,K],x_{kh}^{w}\sim q(x_{kh}^{w}|x_{1}^{w}),x_{t}^{l}\sim q(x_{kh}^{l}|x_{1}^{l})}\big{[}\\ \textstyle&\textstyle\qquad\log S\big{(}-\frac{\tilde{\beta}}{2}\big{(}\|\varepsilon^{w}-\epsilon_{\theta}(x_{kh}^{w},kh)\|^{2}-\|\varepsilon^{w}-\epsilon_{\mathrm{ref}}(x_{kh}^{w},kh)\|^{2}\\ \textstyle&\textstyle\qquad\qquad\qquad\quad-\big{(}\|\varepsilon^{l}-\epsilon_{\theta}(x_{kh}^{l},kh)\|^{2}-\|\varepsilon^{l}-\epsilon_{\mathrm{ref}}(x_{kh}^{l},kh)\|^{2}\big{)}\big{)}\big{)}\big{]},\end{split}$ (229) $$

where $S(x)=\frac{1}{1+e^{-x}}$ denotes the sigmoid function, and $q(x_{kh}^{*}|x_{1}^{*})$ is the conditional distribution of the forward process, i.e. $x_{kh}^{*}$ is sampled as $x_{kh}^{*}=\sqrt{\gamma_{kh}}x_{1}^{*}+\sqrt{1-\gamma_{kh}}\epsilon$, $\epsilon\sim N(0,I)$. Following the derivation of the Diffusion-DPO loss in (Wallace et al., , Sec. S4), we observe that the term $-\frac{\tilde{\beta}}{2}\|\varepsilon^{w}-\epsilon_{\theta}(x_{kh}^{w},kh)\|^{2}$ arises from

$$ $\textstyle-\frac{\tilde{\beta}}{2\frac{1-\gamma_{kh}}{\gamma_{kh}}}\|\hat{x}_{1}(x_{kh}^{w})-x_{1}^{w}\|^{2},$ (230) $$

up to a constant term in $\theta$. If we switch to the more general flow matching scheme, the analog of this term is

$$ $\textstyle-\frac{\tilde{\beta}}{2\frac{\beta^{2}_{kh}}{\alpha^{2}_{kh}}}\|\hat{x}_{1}(x_{kh}^{w})-x_{1}^{w}\|^{2}.$ (231) $$

Using the expression of the denoiser map in terms of the vector field $v$ in equation ([228](#S14.E228)), we can rewrite ([231](#S14.E231)) as:

$$ $\textstyle-\frac{\tilde{\beta}}{2\frac{\beta^{2}_{kh}}{\alpha^{2}_{kh}}}\big{\|}\frac{v(x^{w}_{kh},kh)-\frac{\dot{\beta}_{kh}}{\beta_{kh}}x^{w}_{kh}}{\dot{\alpha}_{kh}-\frac{\dot{\beta}_{kh}}{\beta_{kh}}\alpha_{kh}}-x_{1}^{w}\big{\|}^{2}=-\frac{\tilde{\beta}}{2}\big{\|}\frac{v(x^{w}_{kh},kh)-\frac{\dot{\beta}_{kh}}{\beta_{kh}}x^{w}_{kh}}{\frac{\dot{\alpha}_{kh}}{\alpha_{kh}}\beta_{kh}-\dot{\beta}_{kh}}-\frac{\alpha_{kh}}{\beta_{kh}}x_{1}^{w}\big{\|}^{2}.$ (232) $$

Thus, the Diffusion-DPO loss for Flow Matching reads

$$ $\textstyle\begin{split}\textstyle L_{\mathrm{DPO}}(\theta)&\textstyle=-\mathbb{E}_{(x_{1}^{w},x_{1}^{l})\sim\mathcal{D},k\sim U[0,K],x_{kh}^{w}\sim q(x_{kh}^{w}|x_{1}^{w}),x_{t}^{l}\sim q(x_{kh}^{l}|x_{1}^{l})}\big{[}\\ \textstyle&\textstyle\qquad\log S\big{(}-\frac{\tilde{\beta}}{2}\big{(}\big{\|}\frac{v_{\theta}(x^{w}_{kh},kh)-\frac{\dot{\beta}_{kh}}{\beta_{kh}}x^{w}_{kh}}{\frac{\dot{\alpha}_{kh}}{\alpha_{kh}}\beta_{kh}-\dot{\beta}_{kh}}-\frac{\alpha_{kh}}{\beta_{kh}}x_{1}^{w}\big{\|}^{2}-\big{\|}\frac{v_{\mathrm{ref}}(x^{w}_{kh},kh)-\frac{\dot{\beta}_{kh}}{\beta_{kh}}x^{w}_{kh}}{\frac{\dot{\alpha}_{kh}}{\alpha_{kh}}\beta_{kh}-\dot{\beta}_{kh}}-\frac{\alpha_{kh}}{\beta_{kh}}x_{1}^{w}\big{\|}^{2}\\ \textstyle&\textstyle\qquad\qquad\qquad\quad-\big{(}\big{\|}\frac{v_{\theta}(x^{l}_{kh},kh)-\frac{\dot{\beta}_{kh}}{\beta_{kh}}x^{l}_{kh}}{\frac{\dot{\alpha}_{kh}}{\alpha_{kh}}\beta_{kh}-\dot{\beta}_{kh}}-\frac{\alpha_{kh}}{\beta_{kh}}x_{1}^{l}\big{\|}^{2}-\big{\|}\frac{v_{\mathrm{ref}}(x^{l}_{kh},kh)-\frac{\dot{\beta}_{kh}}{\beta_{kh}}x^{l}_{kh}}{\frac{\dot{\alpha}_{kh}}{\alpha_{kh}}\beta_{kh}-\dot{\beta}_{kh}}-\frac{\alpha_{kh}}{\beta_{kh}}x_{1}^{l}\big{\|}^{2}\big{)}\big{)}\big{)}\big{]},\end{split}$ (233) $$

(Wallace et al., , Sec. 5.1) claim that $\beta\in[2000,5000]$ yields good performance on Stable Diffusion 1.5 and Stable Diffusion XL-1.0, which if we translate to our notation corresponds to $\tilde{\beta}\in[4000,10000]$.

When we have access to the reward function $r$, instead of a winning sample $x^{w}_{1}$ and a losing sample $x^{l}_{1}$, we have a pair of samples $(x^{a}_{1},x^{b}_{1})$ with winning weights $S(r(x^{a}_{1})-r(x^{b}_{1}))=\frac{1}{1+\exp\big{(}r(x^{b}_{1})-r(x^{a}_{1})\big{)}}$, $S(-(r(x^{a}_{1})-r(x^{b}_{1})))=\frac{1}{1+\exp\big{(}-(r(x^{b}_{1})-r(x^{a}_{1}))\big{)}}$. Hence, the loss ([233](#S14.E233)) becomes:

$$ $\textstyle\begin{split}\textstyle L_{\mathrm{DPO}}(\theta)&\textstyle=-\mathbb{E}_{(x_{1}^{a},x_{1}^{b})\sim\mathcal{D},k\sim U[0,K],x_{kh}^{a}\sim q(x_{kh}^{a}|x_{1}^{a}),x_{t}^{b}\sim q(x_{kh}^{b}|x_{1}^{b})}\bigg{[}\sum_{s\in\{\pm 1\}}S\big{(}s(r(x^{a}_{1})-r(x^{b}_{1}))\big{)}\times\\ \textstyle&\textstyle\qquad\log S\big{(}-\frac{s\tilde{\beta}}{2}\big{(}\big{\|}\frac{v_{\theta}(x^{a}_{kh},kh)-\frac{\dot{\beta}_{kh}}{\beta_{kh}}x^{a}_{kh}}{\frac{\dot{\alpha}_{kh}}{\alpha_{kh}}\beta_{kh}-\dot{\beta}_{kh}}-\frac{\alpha_{kh}}{\beta_{kh}}x_{1}^{a}\big{\|}^{2}-\big{\|}\frac{v_{\mathrm{ref}}(x^{a}_{kh},kh)-\frac{\dot{\beta}_{kh}}{\beta_{kh}}x^{a}_{kh}}{\frac{\dot{\alpha}_{kh}}{\alpha_{kh}}\beta_{kh}-\dot{\beta}_{kh}}-\frac{\alpha_{kh}}{\beta_{kh}}x_{1}^{a}\big{\|}^{2}\\ \textstyle&\textstyle\qquad\qquad\qquad\quad-\big{(}\big{\|}\frac{v_{\theta}(x^{b}_{kh},kh)-\frac{\dot{\beta}_{kh}}{\beta_{kh}}x^{b}_{kh}}{\frac{\dot{\alpha}_{kh}}{\alpha_{kh}}\beta_{kh}-\dot{\beta}_{kh}}-\frac{\alpha_{kh}}{\beta_{kh}}x_{1}^{b}\big{\|}^{2}-\big{\|}\frac{v_{\mathrm{ref}}(x^{b}_{kh},kh)-\frac{\dot{\beta}_{kh}}{\beta_{kh}}x^{b}_{kh}}{\frac{\dot{\alpha}_{kh}}{\alpha_{kh}}\beta_{kh}-\dot{\beta}_{kh}}-\frac{\alpha_{kh}}{\beta_{kh}}x_{1}^{b}\big{\|}^{2}\big{)}\big{)}\big{)}\bigg{]},\end{split}$ (234) $$

## 15 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\text{\times}{10}^{-5}$ and parameters $\beta_{1}=0.95$, $\beta_{2}=0.999$, $\epsilon=$1\text{\times}{10}^{-8}$$, weight decay $1\text{\times}{10}^{-2}$, gradient norm clipping value $1$. For Discrete Adjoint, these hyperparameters resulted in fine-tuning instability (see [Table 5](#S9.T5)); the results that we report in all other tables for Discrete Adjoint were obtained with learning rate $1\text{\times}{10}^{-5}$.
- •
Bfloat16 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 [Table 3](#S9.T3). 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 a total set of 100k prompts.

### 15.1 Noise schedule details

Since we use $K=40$ discretization steps, the timesteps are $t\in\{0,0.025,0.05,0.075,0.1,\dots,0.95,0.975\}$. To sample $X_{t+h}$ from $X_{t}$ we use equation ([40](#S5.E40)). We use the choices $\alpha_{t}=t$, $\beta_{t}=1-t$, which means that $\sigma(t)=\sqrt{2\beta_{t}(\frac{\dot{\alpha}_{t}}{\alpha_{t}}\beta_{t}-\dot{\beta}_{t})}=\sqrt{2(1-t)(\frac{1-t}{t}+1)}=\sqrt{\frac{2(1-t)}{t}}$.

Note that if we plug $t=0$ into this expression, we obtain infinity, and if we plug $t\lessapprox 1$, we obtain $\sigma(t)\approx 0$. For obvious reasons, the former issue requires a fix: we simply add a small offset to the denominator of $\sigma(t)$, replacing $\sqrt{1/t}$ by $\sqrt{1/(t+h)}$ (note that $h:=1/K=0.025$). But the latter issue is also not completely satisfactory from a practical standpoint, because looking at the adjoint matching loss ([37](#S5.E37)), we observe that $u(X^{\bar{u}}_{t},t)$ is trained to approximate the conditional expectation of $\sigma(t)^{\mkern-1.5mu\mathsf{T}}{}\tilde{a}(t;\bm{X}^{\bar{u}})$. Thus, if we set $\sigma(t)$ very close to zero for $t\lessapprox 1$, we are forcing the control $u$ to be close to zero as well, or equivalently preventing $v^{\mathrm{finetune}}$ from deviating from $v^{\mathrm{base}}$. While this is the right thing to do from a theoretical perspective, we concluded experimentally that setting $\sigma(t)$ just slightly larger results in substantially faster fine-tuning, thanks to the additional leeway provided to $v^{\mathrm{finetune}}$ to deviate from $v^{\mathrm{base}}$. In particular, we added a small offset to the factor $1-t$ in the numerator $1-t$ of $\sigma(t)$: we replaced $1-t$ by $1-t+h$. Thus, the expression that we used to compute the diffusion coefficient in our experiments is

$$ $\textstyle\sigma(t)=\sqrt{\frac{2(1-t+h)}{t+h}}.$ (235) $$

When solving the lean adjoint ODE ([38](#S5.E38))-([39](#S5.E39)) backwards in time via the Euler scheme ([41](#S5.E41)), the timesteps we use are $t\in\{1,0.975,0.95,0.925,0.9,\dots,0.05,0.025\}$. We do not actually initialize the adjoint state as $\nabla_{x}g(X_{1})$, but rather as $\nabla_{x}g(\hat{X}_{1})$, where $\hat{X}_{1}:=X_{1-h}+hv^{\mathrm{base}}(X_{1-h},1-h)$. That is, $\hat{X}_{1}$ is obtained by performing a final noiseless update, instead of using noise $\sigma(1-h)=\sqrt{4h}$ given by equation ([235](#S15.E235)). The reason for this is that the regular final iterate $X_{1}$ contains some noise that was added in the final step, and that can distort the gradient $\nabla_{x}g(X_{1})$. By setting $\tilde{a}(1;\bm{X})=\nabla_{x}g(X_{1})$, we get rid of this bias. Note that in the continuous time limit $h\to 0$, $\hat{X}_{1}=X_{1}$, which means that this small trick is consistent.

### 15.2 Selection of gradient evaluation timesteps

In [Algorithm 1](#alg1), equation ([42](#S5.E42)), we state that the term $\big{\|}\frac{2}{\sigma(t)}\big{(}v^{\mathrm{finetune}}_{\theta}(X_{t},t)-v^{\mathrm{base}}(X_{t},t)\big{)}+\sigma(t)\tilde{a}_{t}\big{\|}^{2}$ must be computed for all $K$ steps in $\{0,\dots,1-h\}$. However, the gradient signal provided by backpropagating through this expression for consecutive times $t$ and $t+h$ is quite similar. In the interest of computational efficiency, we sample a subset $\mathcal{K}$ of timesteps, and we only compute and backpropagate the terms $\big{\|}\frac{2}{\sigma(t)}\big{(}v^{\mathrm{finetune}}_{\theta}(X_{t},t)-v^{\mathrm{base}}(X_{t},t)\big{)}+\sigma(t)\tilde{a}_{t}\big{\|}^{2}$ for those timesteps. We construct $\mathcal{K}$ by sampling ten timesteps uniformly without repetition among $\{0,\dots,0.725\}$, and always sampling the last ten timesteps $\{0.75,\dots,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 as important.

### 15.3 Loss function clipping: the LCT LCT \mathrm{LCT} hyperparameter

Note that the magnitude of $\sigma(t)^{\mkern-1.5mu\mathsf{T}}{}a(t;\bm{X}^{\bar{u}},\bar{u})$ is much larger for times $t\gtrapprox 0$ than for times $t\lessapprox 1$. The reason is two-fold:

- •
As discussed in [Section 15.1](#S15.SS1), $\sigma(t)$ is much larger for $t\gtrapprox 0$ than for $t\lessapprox 1$.
- •
The magnitude of the lean adjoint state $\tilde{a}$ grows roughly exponentially as $t$ goes backward in time. In fact, if we assumed that $\nabla_{x}b(X_{t},t)$ is constant in time, this statement would be exact.

Observe that when $\sigma(t)^{\mkern-1.5mu\mathsf{T}}{}a(t;\bm{X}^{\bar{u}},\bar{u})$ is large, the gradient $\nabla_{\theta}\big{\|}\frac{2}{\sigma(t)}\big{(}v^{\mathrm{finetune}}_{\theta}(X_{t},t)-v^{\mathrm{base}}(X_{t},t)\big{)}+\sigma(t)\tilde{a}_{t}\big{\|}^{2}$ also has a high magnitude. Including such terms in our gradient computation decreases the signal to noise ratio of the gradient. Even more so, as discussed in [Section 15.2](#S15.SS2) 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\lessapprox 0$ in our gradients can muffle these other important, low-magnitude terms.

To fix this issue, we clip the terms such that $\big{\|}\frac{2}{\sigma(t)}\big{(}v^{\mathrm{finetune}}_{\theta}(X_{t},t)-v^{\mathrm{base}}(X_{t},t)\big{)}+\sigma(t)\tilde{a}_{t}\big{\|}^{2}>\mathrm{LCT}$, where $\mathrm{LCT}$ stands for the loss clipping threshold. That is, the adjoint matching loss that we use in our experiments is of the form:

$$ $\textstyle\hat{\mathcal{L}}_{\mathrm{Adj-Match}}(\theta)=\sum_{t\in\mathcal{K}}\min\big{\{}\mathrm{LCT},\big{\|}\frac{2}{\sigma(t)}\big{(}v^{\mathrm{finetune}}_{\theta}(X_{t},t)-v^{\mathrm{base}}(X_{t},t)\big{)}+\sigma(t)\tilde{a}_{t}\big{\|}^{2}\big{\}},$ (236) $$

where $\mathcal{K}$ is the random timestep subset described in [Section 15.2](#S15.SS2).

For adjoint matching, we set $\mathrm{LCT}=1.6\times\lambda^{2}$. Remark that $\mathrm{LCT}$ needs to grow quadratically with $\lambda$, because the magnitude of the lean adjoint $\tilde{a}$ grows quadratically with $\lambda$. We set the constant 1.6 through experimentation; all or almost all of the terms for the last ten timesteps fall below $\mathrm{LCT}$, but only a fraction of the terms ($\approx 25\%$) for the first ten steps fall below $\mathrm{LCT}$. The constant for $\mathrm{LCT}$ is a relevant hyperparameter that needs to be tuned to obtain a similar behavior.

We also used loss function clipping on the continuous adjoint loss. For that loss we set $\mathrm{LCT}=1600\times\lambda^{2}$. 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).

### 15.4 Computation of evaluation metrics

We used the [open_clip](open_clip) library (Ilharco et al., ) to compute ClipScores. We computed ClipScore diversity as the variance of Clip embeddings of 40 generations for a given prompt, averaged across 25 prompts. Namely,

$$ $\textstyle\mathrm{ClipScore_Diversity}=\frac{1}{40}\sum_{k=1}^{40}\frac{2}{25\cdot 24}\sum_{1\leq i<j\leq 25}\|\mathrm{Clip}(g^{k}_{i})-\mathrm{Clip}(g^{k}_{j})\|^{2},$ (237) $$

where $g^{k}_{i}$ denotes the $i$-th generation for the $k$-th prompt.

We used the [transformers](transformers) library to compute the PickScore processor and model (Kirstain et al., ). PickScore diversity is computed in analogy with ClipScore diversity.

We used the [hps](hps) library to compute values of Human Preference Score v2 (Wu et al., ).

To compute Dreamsim diversity we use the [dreamsim](dreamsim) library (Fu et al., ). Dreamsim diversity is computed in analogy with ClipScore diversity.

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: What is the complexity/computation cost of computing/sampling adjoints in Equation (26) and (27)? 

In 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?
soundness: 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.
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.
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. 

**Minors**:

1）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.

2）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.

3）In proposition 2, it is better to first define earlier of the adjoint matching objective instead of combining the definition with the proposition.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: see above.
soundness: 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.

The proposed Adjoint Matching algorithm for SOC, casting it as a least-squares regression problem, is novel and effective.

This paper is well-written, clearly structured and easy to follow.
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.
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.

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

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - Did the authors compare the proposed approach to the work by Uehara et. al. [1]?
- Did the authors test how adjoint matching works with only a few diffusion steps?
- How did the authors select the time points to discretize? Uniform in the interval [0,1]?
- 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?
- Are the authors planning to make the code public?

---

- [1] Uehara et al. (2024). Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control. *arXiv preprint arXiv:2402.15194*
- [2] Vargas, F., Grathwohl, W., & Doucet, A. (2023). Denoising diffusion samplers.  *ICLR 2023*. 2024.
- [3] Zhang, Q., & Chen, Y. (2021). Path integral sampler: a stochastic control approach for sampling. *arXiv preprint arXiv:2111.15141*
- [4] Nusken, Nikolas, et al. "Transport meets variational inference: Controlled Monte Carlo diffusions." *The Twelfth International Conference on Learning Representations: ICLR 2024*. 2024.
- [5] Richter, L., & Berner, J. (2023). Improved sampling via learned diffusions. *ICLR 2024*. 2024.
soundness: 4
strengths: - The paper is well-written and structured.
- Strengths of memory-less noise schedule:
    - 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].
    - The proposed approach is compatible with both DDPM and Flow/Bridge-Matching models.
- Strengths of (lean) adjoint matching:
    - Simple regression-based objective with circumvents memory problems associated with the discrete adjoint method
    - Simulating the adjoint ODE does not require control evaluations, making it more scalable than the continuous adjoint method.
    - Compatible with general SOC problems.
- Numerical evaluations show that the proposed approach outperforms competing methods on a variety of evaluation criteria.
summary: This paper aims to fine-tune a pre-trained diffusion model. The contribution of this paper is two-fold:

First, 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.

Second, 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.
weaknesses: - Weaknesses of the memory-less noise schedule:
    - 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
    
    $$
    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)) 
    $$
    
    which es ensures that 
    
    $$
    p_{\text{target}}(X_1) = \exp(r(X_1)).
    $$
    
    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].
    
- Weaknesses of (lean) adjoint matching:
    - It is not clear (at least to me) which objective function lean adjoint matching is optimizing. 
    - Lean/Continous Adjoint-matching requires simulating another differential equation which may result in computational overhead.
- Code is not public

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See weaknesses.
soundness: 4
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:

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


```
[1] Denker et al., Efficient Finetuning of Conditional Diffusion Models by Learning the Generalised h-transform.
```
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.
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.


```
[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.
[3] Domingo-Enrich et al., Stochastic Optimal Control Matching.
```

Meta-review: This paper addresses the challenge of reward fine-tuning in dynamical generative models, such as Flow Matching and denoising diffusion models, by framing it as a stochastic optimal control (SOC) problem. It introduces a theoretically grounded approach requiring a specific memoryless noise schedule during fine-tuning to account for noise-sample dependencies. The proposed algorithm, Adjoint Matching, reformulates SOC as a regression problem, outperforming existing methods by improving consistency, realism, generalization to unseen human preferences, and sample diversity. Following the author response and author-reviewer discussions, this paper has received unanimous support from the reviewers. Therefore, I recommend acceptance. In the final version, the authors are encouraged to include a discussion on SPIN-Diffusion, a closely related work focused on fine-tuning diffusion models, which has demonstrated greater effectiveness compared to Diffusion-DPO.

Decision: Accept (Spotlight)