Value Gradient Guidance for Flow Matching Alignment
Authors: Zhen Liu
Tim Z. Xiao
Carles Domingo-Enrich
Weiyang Liu
Dinghuai Zhang
Venue: NeurIPS 2025 poster ()
Keywords: Flow Matching
Alignment
Reward Finetuning
Reinforcement Learning
Optimal Control
Primary area: deep_learning
TL;DR: We propose a value gradient matching formulation for reward finetuning/alignment for flow matching models with the theory of optimal control, and empirically verify our method on the popular text-to-image flow matching model StableDiffusion3

Abstract: While methods exist for aligning flow matching models––a popular and effective class of generative models––with human preferences, existing approaches fail to achieve both adaptation efficiency and probabilistically sound prior preservation. In this work, we leverage the theory of optimal control and propose VGG-Flow, a gradient matching–based method for finetuning pretrained flow matching models. The key idea in this algorithm is that the optimal difference between the finetuned velocity field and the pretrained one should be matched with the gradient field of a value function. This method not only incorporates first-order information from the reward model but also benefits from heuristic initialization of the value function to enable fast adaptation. Empirically, we show on a popular text-to-image flow matching model, Stable Diffusion 3, that our method can finetune flow matching models under limited computational budgets while achieving effective and prior-preserving alignment.

Value Gradient Guidance for Flow Matching Alignment
Zhen Liu
Tim Z. Xiao** Carles Domingo-Enrich"
Weiyang Liu
Dinghuai Zhang
"The Chinese University Hong Kong (Shenzhen) Microsoft Research
2 University of Tübingen
"The Chinese University of Hong Kong
Mila vgg-flow thub. io
Quebec Al Institute
Equal contribution *Corresponding author
L. [[u](xe, (san (i1 - r(x:) di
min Es, (o(x))) 2
21 (1 (Tte,
Euler Step
Control
r(21)
(V. vx(á.) (N(r.
Matching
Un t]
do (sta.0)
Number of Finetuning Steps
Figure 1: Left: Illustration of proposed VGG-Flow algorithm. The velocity field is trained to match the value gradient field obtained from the optimal control problem. The value gradient field is parametrized as the reward gradient field of one-step Euler prediction plus learnable residual field. Right: Evolution of samples (with fixed seeds and prompts) during the course of finetuning on the reward model of Aesthetic Score.
Abstract
While methods exist for aligning flow matching models a popular and effective class of generative models - with human preferences, existing approaches fail to achieve both adaptation efficiency and probabilisticaly sound prior preservation. In this work, we leverage the theory of optimal control and propose VGG-Flow, a gradient-matching- based method for finetuning pretrained flow matching models. The key idea behind this algorithm is that the optimal difference between the finetuned velocity field and the pretrained one should matched with the gradient field of value function. This method not only incorporates first-order information from the reward model but also benefits from heuristic initialization of the value function to enable fast adaptation. Empirically, image flow matching model, Stable Diffusion 3, that our method can finetune flow matching models under limited computational budgets while achieving effective and phirr-preservig alignment.
show popular text-to-
Introduction
Flow matching models [1, 36, 38] one of the most effective methods in modeling high-dimensionall real-world continuous distributions and widely used for the generation of images [17], videos [65], 3D objects [39, 40, 69, 78], etc. These models, compared to diffusion models that rely simulation with stochastic differential equations (SDEs), are trained to sample with deterministic ordinary differential equations (ODEs) of which sampling paths are often straighter and easier to model.
Similar to the motivations for performing alignment for diffusion models [6, 19], natural to finetune flow matching models with reward models so that the generated samples are more aligned
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

with human preferences. While existing methods have already achieved fast, effective, diversity- preserving and prior-preserving alignment for diffusion models through approaches, the ODE sampling paths of flow matching models pose challenges in applying these methods. The key challenge is that, with flow matching models, one typically has access to neither reference path (unless one has access to the large-scale pretraining dataset) nor the probability flow. Since it is non- con-trivial to obtain the probability flow and to incorporate the learned prior from base models for flow matching models, it is harder to align flow matching models in an efficient yet probabilistic way.
matching- based
To address this issue, we take inspiration from the theory of optimal control and consider relaxed objective: we optimize the target reward but with the accumulated cost-to-go defined as the lz distance between the velocity fields of the finetumed model and the base model. The optimal solution of this optimization program described by the Mamilton-Jacohi-Blllmee be shown in our formulation equivalent to two conditions: a gradient matching condition that the residual velocity field matches the gradient of the value function, and a value consistency condition that ensures correct estimation of value functions. In light of this result, we propose our finetuning method, dubbed VGG-Flow (short for Value Gradient Guidance for Flow Matching Alignment), that finetunes the flow matching model via "matching with value gradient guidance" thee the difference between the velocity fields of the finetuned model and the base model is expected to be the gradient of the value function- while the value function can be solved with a consistency equation. Such formulation allows us not only to directly propagate the reward gradient to the matching target through the value consistency equation in amortized and memory-efficiee initialization of the value gradient for fast convergence. We empirically show that VGG-Flov can effectively and robustly finetune large flow matching models like Stable Diffusion 3 [17] within limited computational resources.
(HJBB equation and can
way but also to use a heuristic
To summarize, our contributions are
With a relaxed objective, we leverage the HJB equation from optimal control theory to propose VGG-Flow, an efficient and effective alignment method for flow matching models that matches the residual velocity field with the guidance signal of value function gradient.
We propose to parametrize the value gradients with forward-looking technique, which eases the difficulty in learning accurate value gradients limited time and thus accelerates convergence.
We empirically demonstrate the effectiveness of VGG-Flow matching model, Stable Diffusion 3, and show that VGG-Flow achieves better reward convergence, sample diversity, and prior preservation compared other alignment baselines.
large-scale text-to- -mage flow
2 Related Work
General alignment strategies. Since large generative models are typically trained on uncurated massive datasets, their sample distributions are typically far from human preferences. A common approach to solve this problem is through reinforcement learning from human feedback (RLHF) [45], in which one first trains a reward model from human preference datasets and later finetunes a generative model with reinforcement learning methods such that it samples from this reward model. In the alignment of large language models, it is common to use simple policy gradient methods such as PPO [56] and GRPO [57]. While they are general enough to be applicable to continuous generative models, they can be less efficient because they do not leverage the differentiable nature of both reward models and generative models. Similar to traditional RL [73] methods, the framework of generative flow networks (GFlowNets) [4, 46, 74, 75, 76], which are highly correlated with soft RL methods, can be used to finetune diffusion models [41, 71, 72, 77]. Alternatively, one may simple rewerd- methods [18, 19, 34] for the same purpose. It is also under exploration to perform test-time scaling on diffusion models via methods like sequential Monte Carlo [27, 58], parallel tempering [28], and search [42] without any model finetuning.
Differentiable RLHF for continuous foundation generative models. Diffusson models and flow matching models, commonly used to build foundation models in the continuous domain, exhibit different properties due to their differentiable and sequential sampling process. For diffusion models, since each sampling step is stochastic (probably excluding the last step), one may finetune these models using stochastic optimal control [61] which typically requires extra steps for learning surrogate functions. A recent work inspired by the framework of generative flow networks [41] builds a gradient- informed finetuning strategy that efficiently aligns diffusion models with gradient-matchinn like
2

---

losses in a probabilistic way. For flow matching models, these approaches are not applicable because they require the transitions to be stochastic. One way that applies to both diffusion models and flow matching models is to treat the sampling process as a computational graph with which we directly optimize differentiable rewards [12, 70]. However, such a strategy by design fails to align with the target distribution but only aims to find some modes in the distribution. More principled approaches for aligning flow matching models include the recent method of Adjoint Matching [15], which by turning flow matching models into equivalent SDEs computes gradient matching target for the velocity field with an adjoint ODE, Solving such an ODE is however expensive, especially in cases of foundation models where accurate adjoint ODE solving requires smaller time steps in ODE solver.
Optimal control and machine learning. Optimal control (OC) is concerned with steering systems subject to random fluctuations so as to minimize given cost. OC methods, including the subset of stochastic optimal control (SOC), have been employed in a broad range of areas, including the simulation of rare events in molecular dynamics [24, 25, 30, 80], modeling in finance and economics [21, stochastic filtering and data assimilation [43, 54), tackling nonconvex optimization problems [10], management of power systems and energy markets [3, 50], robotic control [23, 60], analysis of
mean-field games [9], optimal transport theory [63, 64], the study of backward stochastic differential equations [8], and large-deviation principles [20]. Relevant and recent applications of SOC in machine learning include performing reward fine-tuning of diffusion and flow matching models [15, 41, 61, 77] and conditional sampling of diffusion processes [13, 48, 66). There is also growing literature on SOC methods for sampling from unnormalized densities, as an alternative to MCMC methods [2, 5, 7, 11, 26, 62, 79]. Additionally, there have been string of methodological works exploring deep learning loss functions for SOC [14, 16, 44].
3 Preliminaries
3.1 Flow Matching Models
Flow matching models are a class of generative models that are trained to generate samples sequen- tially following some reference paths. Specifically, one generates samples 21 with flow matching model by simulating a trajectory from an initial state Ip ~ N(00 7) with the dynamics è = vo t). The velocity field ve (zz t) is learned with the flow matching loss:
[)) = E, Unitfoo [| () - u(<x(<r))) |=))
where u(x1(211) is reference conditional velocity field. A popular choice is u(zz(=1)) (1 t) 1 adopted by a variant called rectified flows [38].
The probability flow p(r, 1) corresponding to the velocity field v(ar, 4) satisfies the so-called continuity equation:
a at p(.. 1) + V (p(z, t)o ()) = 0.
(2)
Since flow matching model is modeled as ordinary differential equation (ODE), one may use any ODE solver to generate samples, including the simplest Fuler sampler: T++At = d4 + Alv(sr, I) where A/. is the step size.
3.2 Optimal Control
In optimal control theory, aim find optimal control signal u" u"(x:, f) under known time- varying dynamics Æ = fxx u², t) with the initial state z(0)= = 20 such that a cost functional is minimized. The standard control objective is defined as
arg min J[], Ju] A [" L(x(t), u(t),t) dt + (x(T)),
(3)
a
where (--) is the terminal cost and L(x, u, t) is the running cost.
The Jamilton-Jacobi-BBe Bellmmn (HJB) equation is indeed the continuous counterpart of the Bellman equation the discrete domain, where the transitions are defined by a graph instead of dynamical system i = u(x, t). The solution of this optimal control problem satisfies the HJB equation:
@v (x, 1) min [L u, 1) + VV T)T [(x, u, ()))
(4)
Øt
3

---

Algorithm VGG-Flow algorithm
Require: Pretrained flow matching model Passe (), given reward function r(21), value gradient model 9o(2. 1) parameterized by Equation 16.
Ensure: Finetuned flow matching model (2r 1)
Initialize flow matching model Vyg f- Ubase.
while Stopping criterion not met do
Collect trajectories {x+} Je via solving the current neural ODE u = vi (I'z, t). Update value gradient model (0(x, t) with loss Lconsistency (Ø) + b'ouuuuyrry (Ø).
Update velocity field model by (z, t) with loss Lmatching (9).
end while
in which Vax, t) is the value function or the minimal cost-to-go from state x at time
v(a,t) min {T" J." L(x(s), ds +(x(T)) xx
(5)
Method
4.1 Gradient Matching for Aligning Flow Matching Models
Given reward function r(-), we want to train our generative model to achieve high reward scores for the generated samples and also preserve the prior distribution of the pretrained model. We can then define the following formulation for training flow matching model dr, va( () dt, / E [o, 1] that transforms standard Gaussian distribution po = = /) to a target distribution po-
We start with the following optimal control formulation for flow matching alignment problems
X
2 f. (10, t) dt- de-r(xx)) r(aa) u0 t) A Uhase (Ft, t) + C0 t), (6)
min Beo do (2)
where 0g = Up Ubase is the residual velocity field and is the reward multipliertemperature With such relationship, we interchangeably This program can interpreted as a control problem where we want to find a deterministic control parameterized by ū that minimizes the expected cost of the system, which is defined the sum of the terminal reward function r(x1) and the running cost (i.e., regularization term) Su' [] 4) dt. Remark 1 (Connection to Equation 3). In our reward funeting setup, the control from the general optimal control formulation in Section 3.2 is denoted as v and the terminal time T set to 1. Our dynamics here á f(z, v, A v(z, t), the running cost is defined in L(x, v,t) & ||o () Vhase (ar, 1) the terminal cost is (xx()) a -r(x1), T 1, and the value function f) & min, f(2 3/ [l:(,,,)) Uhase (I., }) ds r(x)) ) for dynamic starting with x,
vg and üg denote the parameterized flow matching model.
The corresponding HJB equation for the above objective is:
0,V t) + min (V t)
(Doase (x. t) + ग(z, t)) + || t)| =0,
2
With the first-order condition of the minimization program of e in the HJB equation, we obtain the following optimal control law:
(-, = -V
(Value Gradient Matching)
(8)
This optimal control law can be interpreted as a gradient matching criterion, where the residual velocity field (z t) should match value function gradient VV t) state at time t. If an oracle value function is provided, then alignment of the flow matching model can simply be achieved through "gradient matching" loss between the residual velocity field and the oracle value gradient.
4.2 Solving HJB Equation with Value Gradient Guidance
With the optimal control law (Equation 8), the HJB equation reduces to
(Value Consistency)
V(z,t)
VV t) VV(x, (2,
(9)
0t
2X

---

While we could in principle solve this equation by parametrizing Vix, t) with neural net, it is better that we directly parametrize (z, t) since is considerably more effective and robust, as shown diffusion model and energy-based model literature [55, 59]. With (z t) & VVa t), we may write the equivalent gradient-version HJB equation by taking gradients on both sides:
0r x V Dø - gib I" "basse (Ir, 1) - Chase (2r, [)]" 9
(10)
= -Voll (7 Phace (2, t) V Uhuse (z, t)]² 9
(11)
with the boundary condition go (z, 1) - -r((x) at terminal time.
With 5 = 1/x, we write the following set of losses to update value function gradient model g0 t):
Leonsistency (+)
0r + Vooll (Drace
- (Thace
-gge) + [ Vihase] 9b
(12)
Lboundaay ()) (0. (10 1) vr(x1))
(13)
In practice, this consistency loss based on Equation 9 can be efficiently implemented with finite difference methods and Jacobian-vector products in PyTorch.
Furthermore, with decently learned value gradient model that captures the optimal control, regress our residual velocity field to it to learn our flow matching model ug:
Lmatching (0) Ezo (too ( () + 8gs (1r, I)l
(14)
Notice that we only use this objective to update 0, not Φ. This makes the total training objective Crotal
Lunal (0, 6) Lmatching (0) + Counsistency (Ø) OoLonundayy (Ø),
(15)
where a is a coefficient to tune the importance of boundary condition loss in the training.
Efficient parametrization of value function gradients. Solving the consistency equation for the value function gradient model Io (2. t) can take a non-trivial amount of time. For flow matching models, especially variants like rectified flows, the value of ze can be well approximated by the reward of the single-Euler-step prediction è = 21 (Ie, t) å I2 + (1 t) stop-gee cooo---rrdientt(((ee f)), in which the stop gradient operation is inspired by DreamFusion [49] and helps improve results. Therefore, propose parametrize (x, t) with
g (2x,t) t) A -/; stop- ceop-grradeent (V,rr(d(a t))) + us (Te,
(16)
where is a positive weighting scalar and (o t) is a learnable error correction term which is supposed be close to zero when 1 +
Putting everything together. At each training step, our VGG-Flow algorithm simulates trajectories i (1 with ODE solver, and use the obtained trajectory data to update the value gradient model 9ø and velocity field model ug. We summarize the proposed method in Algorithm 1.
5
Experiments
5.1 Experiment Settings
Base model. Throughout the paper, we consider the popular open-sourced text-conditioned flow matching model Stable Diffusion 3 [17] and 20-step Fuler solver to sample trajectories.
Reward model. We consider three reward models learned from large-scale human preference datasets: Aesthetic Score [33], Human Preference Score (HPSv2) [67, 68], and PickScore [32].
Prompt dataset. For Aesthetic Score, we use set of simple animal prompts used in the original DDPO paper [6]; for HPSv2, consider photo+painting prompts from the human preference dataset (HPDv2) [67]; for PickScore, we use the prompt set in the Pic-a-Pick dataset [32].
Metrics. We follow previous works [15, 41] and compute the variance of latent features (both DreamSim features [22] and CLIP features [29, 53]) extracted from a batch of generated images (we batch of size 16) to measure sample diversity. To measure the degree of prior preservation,
5

---

Pretrained Roward 6.09
VGG-Flow Reward 7,75
ReFI. Reward 7.59
DRaFT Reward 7.46
AM Reward 7.23
VGG-Flow-PMP
Reward 7 7,86
Figure Comparison on samples generated by models finetuned different methods. models finetuned with a maximum of 400 update steps and for fair qualitative comparison we pick the model checkpoints that yield the best rewards without significant collapsing in image semantics (as ReFL and DRaFT are more prone to overfitting) For each of images produced by each method, display their average reward on the left.
Aesthetic
IIPSv2 Diversity DreamSim (f. 1)=
Method
Reward (1)
Diversity DreamSim
HD
Reward
HID
Reward
Diversity DreamSim
FID
10100-3
(t. 10
(t. 10
Rase (SD3)
5.99
23.12
212 2,80 22.22
558
21.81 27.81 589
ReFL
10.00 0.31 5.59 1.39 1338 191 388 0.1 14.08 0.55 1195 21 23.19 0.5 17.71 0.77 997 +
DRaFT
9.54
7.78
1518
3,76
15.05
1177
968
AM
6.87 0.19 22.34 + 2:39
465 93 3.59 0.03 14.11
0.26 1246 24 22.78 0.04 19.70 L.8 1033
VGO-How-PPMP
VGG-Flow
8.24 +: 0.00 22.12 +: 0.17
375 +
25 3.86 0.03
18.40
1.12 1161 19 23.21 008 20.93 0.98 1058 31
Table I: Comparison between the models finetuned with proposed method VGG-Flow and the baselines. All models are finetuned with 400 update steps. Since there are inherent trade-offs between reward and other metrics. values at final update step not fully capture the differences between methods. therefore refer readers to the Pareto front comparisons in Figs. 6, 8, and for more comprehensive evaluation.
compute the per-prompt FID score between image sets generated by the finetuned model and the base model and use the average per- per-prompt FID score as the prior preservation metric.
Baselines. We consider two types of basclines: the generic ones that rely on direct reward maximiza- tion with truncated computation graphs, including ReFL [70] and DRaFT [12], and the adjoint-based Adjoint Matching method [15]. Specifically, ReFL samples trajectory and take the truncated com- putational graph stop-grose (r > Fi+ with randomly sampled time step t. The model is finetuned to maximize the reward r(ē (zz. t)) of the single-step prediction è (2₁, t). Similarly, DRaFT truncate the full inference computational graph at some random time step K A/ and perform backpropagation on this length K graph with the differentiable reward signal r(x1). For optimal-control baselines, we consider adjoint matching [15], which finetunes flow matching model under stochastic settings (details in Appendix A.1). Additionally, consider variant

---

Pretrained
VGG-Flow
ReFL
DRaFT
AM
VGG-Flow-PMP
Prompt: Cute conceptual style.
Prompt: with bunny
Prompe: Footage of astronaut tropical beach.
Figure 3: Qualitative results on HPSv2.
Pretrained
VGG-Flow
ReFL
DRaFT
AM VGG-Flow-PME
Prompt: An antroooomephicc white rabbit, male wizard face, dressed in black and white, fine art, award-winning intricate, elegani, sharp focus, lighting. highly decailed, digital painting. concept by guweiz gu, masterplece, trending Sk.
Prompt: surfing steampunk giraffe a rainforest.
Prompt: Flying horses with wings, sunset at the Lower Galilee.
Figure 4: Qualitative results on PickScore.
of VGG-Flow which, derived with the Pontryagin's Maximum Principle [35, PMP], follow an sigjiintmatccinn-lieee like algorithm but with slightly different evolution equations (Appendix A.I).
Experiment settings and implementation details. We LoRA parametrization [31] on attention layers of the finetuned flow matching model with LoRA rank of 8. The value gradient network in VGG-Flow set with tiny weights in the final output layers. Since Stable Diffusion 3 is a latent flow matching model, reward for a sampled image 4i r(decode(11) where decode(-) is the VAE decoder of the Stable Diffusion 3 model. This decoder is always frozen and we only finetune the LoRA parameters. For all experiments, we use random seeds. For Aesthetic Score, HPSv2 and PickScore experiments, set the default inverse temperature terms 3 1/A 5e4, 3c7 and 5c5, respectively;
scaled-down version of the Stable Diffusion-vi.55 U-Net, initialized

---

VGG-Flow DRaFT
ReFL
VGG-
10
50
159
108
326
310
30
Figure 5: Convergence curves of different metrics for different methods throughout the finetuning process on Aesthetic Score. Finetuning with our proposed VGG-Flow converges faster than the non non-gradient- non-gradient-nformmed informed methods and with better diversity- and prior-preserving capability.
Bdc
Figure 6: Trade-offs between reward, diversity preservation and prior preservation for different reward finetuning methods on Aesthetic Score. Dots represent the evaluation results of models checkpoint saved after every 5 iterations of finetuning, where ones with greater reward, greater diversity scores and smaller FID scores considered better.
for all ablation studies with Aesthetic Score, we set 3 = 1e4. We set the boundary loss coefficient to 10000 for all experiments. We use an effective batch size of 32 for all methods, and only use on-policy samples without any replay buffer. For VGG-Flow, we sub-sample the collected trajectories by uniformly splitting each into bins and then taking transition out of each; also clip the computed reward gradients in Eqn. 16 at the 80th percentile of the gradient norms of the corresponding training batches. For ReFL, sample the truncation time step between 15 and 20. We follow prior work [51, 70] use ReLU(r(z)) as the reward model for both ReFL and DRaFT adjoint matching (AM), we use 4 GPUs for each and set
for stable training. For experiments the inverse temperature 33 to 5 x 10 3 x 105, x 10a for Aesthetic, HPS, PickScore, respectively, based on the same hyperparameter choosing protocol from [15]. All AM experiments use float32 computation and drop samples that result in too large gradient norms.
5.2 Results
General experiments. We show in Figure 2 the visualization of samples produced by both the base model and the models finetuned on Aesthetic Score. As ReFL and DRaFT finetuning can easily lead to reward hacking, we perform early stopping and pick the model checkpoints without major loss of image semantics and with the highest reward values possible. Compared to ReFL and DRaFT, out proposed VGG-Flow produces higher rewards with better preservation of semantic prior from the base Stable Diffusion model. Our method also works well on other reward models, including HPSv2 and PickScore, as shown in Figure 3 and 4. To further illustrate the advantage of our method and the tradeoffs between reward convergence and other metrics, we present quantitative results in Table 1, Figure 5 and Figure 6. Specifically, we observe that VGG-Flow achieves comparable speed with respect to direct reward maximization methods (ReFL and DRaFT) but better maintains sample diversity (measured by DreamSim and CLIP diversity score) and base model prior (measure by FID score). We observe that ReFL and DRaFT on Aesthetic Score easily achieves reward values close to 9, of which value typically indicates complete forgetting of base model prior [12]. Furthermore, the Pareto front figures of reward values, diversity scores and FID scores show that our VGG-Flor achieves better diversity/IID scores at the same level of reward values demonstrating that our VGG-Flow outperforms the baselines even if we perform early stopping.
Effect of reward temperature. We conduct an ablation study on Aesthetic Score with different (5000, 10000, 50000} and show in Figure 11 the effect of reward temperature. We observe that for all reward temperatures, the reward smoothly increases at speed proportional to B. The sample diversity and prior preservation capability are generally worse with greater 3 values. Greater 3 values also leads to worse trade-off on FID vs. reward but no significant difference in diversity vs. reward.
8

---

VGG-Flow
DRaFT
ReFL
VGG-Flow-
012
500
200
325
Figure 7: Convergence of different metrics different methods throughout the finetuning process on HPSv2.
US0
VCC
Figure 8: Trade-offs between metrics for different reward finetuning methods (experiments on HPSv2).
VGG Flow
DRaFT
RefL
AM
VGG-Flow PMP
C22
22.0
900
210
100
Ubder
100
700
Uppaecttpp
400
Up 95p
Figure 9: Convergence of different metrics for different methods throughout the finetuning process on PickScore.
Truel Pulcow
Tmeal Puldoue
Fiera Smant Puldiin
kil
AM
AM
Figure 10: Trade-off' between metries for different reward finetuning methods (experiments PickScore).
1/A-5e4
1/A=1c4
1/A=5o3
CLIF Peeeee
110
300
Sppaer Kay
Figure 11: Evolution of metrics for different reward temperature (experiments on Aesthetic Score). Higher temperature 3 leads to faster convergence but with less diversity and less prior preservation.
Effect of 7) schedule. We observe that by setting 7) t,, the convergence speed faster than our default choice of quadratic schedule п = t² (Fig. 13 and 14). Both schedule yields nearly identical trade-offs between metrics, which not only suggest that relative independence of the final performance of trained models on the choice of the parameterization of the learned value gradient model.
Effeet of transition subsampling rate. We also investigate if lower subsampling rate, with which the variance of estimated parameter gradients are lower, leads to better performance. Fig. 15 and 16, observe that there is significant difference between subsampling rates of 25% and 50%.

---

ED Seree Jonken
Figure 12: Trade-offss between metrics for different reward temperatures (experiments on Aesthetic Score).
6 Discussions
HJB PMP. Another way characterize the optimal control is through Pontryagin's Maximum Principle [35, PMPJ. With the control formulation in Section 3.2, we can define the Hamiltonian
H u, L a) _ L(x, u) + af u, () and the adjoint state (also known as co-state) a()) that satisfies
a(t) = - -V (xH(((t)), u(t) S.t. a(T) = (T))) 2 f(r,
(17)
Essentially, the PMP states that the optimal control u² satisfies (²) arg max, (x" (t), u. t, a (t)). where x" and a' are the solutions to equation 17 for the op- timal control u" With the cost functional and dynamics in our setting (Equation 3), we have = -2 Ubasell 1|2 + uTaa] and (1 t) + a(t) = 0. By comparing it with Equation 8, we have a(t) = (2),,) for any trajectory (0, from the dynamics ă u2z t). While mathematically equivalent, solving this adjoint equation not only requires the expensive (and often unaffordable) computation of V multiple times per trajectory but also is prone to accumulated errors in solving the adjoint equation. In contrast, our HJB-based method is more efficient and robust because it 1) solves for VV in an amortized approach with the howward-looking parametrization of VV and 2)
allows for efficient transition subsampling. See Section A. for more details.
Connection with adjoint matching. Adjoint matching [15] reaches matching objective similar to the one in the above PMP discussion. However, their framework is based on stochastic optimal control instead of deterministic optimal control. While the stochastic setting allows them to sample from a simple tilted distribution Piuse (z) c(z). their algorithm requires modifying the flow matching ODE into an SDE with equal marginals. Our proposed algorithm fine-tunes directly the ODE dynamics with deterministic control. Computationally, adjoint matching relies on solving the adjoint ODE, which requires taking one backward pass through the model for each time step. VGG-Flow relies on the value function gradient model of current step and is thus more computationally tractable.
Limitations. Since our method relies relaxed objective, the finetuned distribution approximates the ideal KL-regularized distribution well only in the case of a relatively small A. Implementation- wise, finite differences approximate the first-order gradients of the value gradient estimator and disable all second-order gradients during backpropagation, which inevitably leads to biases. Furthermore, method suffers from the same challenge of exploration-explootttion tradeoff as in common reinforcement learning settings. As we aim for fast convergence within limited computa- tional resources, our hyperparameter settings in theory more prone mode collapse. Furthermore, we do not explore better architecture designs, which is shown important for efficient and stable finetuning of foundation models [37, 52].
7 Conclusion
We propose VGG-Flow, an efficient and robust method for performing alignment of flow matching models with some reward model. By leveraging a relaxed objective and the HJB equation in optimal control theory, we derive gradient matching method that allows us to finetune flow matching models with probabílistic guarantees and memory-efficiee computation. We empirically demonstrate the effectiveness of our VGG-Flow on Stable Diffusion 3, popular large-scale text-conditioned flow matching model, with common image-input reward functions. As for broader impact, we point out that improving the alignment of flow matching models enhances their ability to reliably follow human instructions, contributing to the development of more trustworthy and controllable Al systems that can better serve societal needs in education, healthcare, and decision support.
10

---

References
[1] Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Fijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023. 1
[2] Michael S Albergo and Eric Vanden- preprint arXiv:2410.02711, 2024. 3
Vanden-Eijnden, Eijnden. Nets: A non-equilibrium transport sampler. arXiv
[3] Andrea Belloni, Luigi Piroddi, and Maria Prandini. A stochastic optimal control solution to the energy management of microgrid with storage and renewables. In American Control Conference (ACC), 2016. 3
[4] Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet foundations. Journal of Machine Learning Research, 2023. 2
[5] Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion- based generative modeling. Trunsactions Machine Learning Research, 2024. 3
[6] Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. In ICLR, 2024. 1,5
[7] Denis Blessing, Julius Berner, Lorenz Richter, and Gerhard Neumann. Underdamped diffusion bridges with applications to sampling. In ICLR, 2025. 3
[8] René Carmona. Lectures on BSDEs, stochastic control, and stochastic differential games with financial applications, volume 1. SIAM, 2016. 3
[9] René Carmona, François Delarue, et al. Probabilistic Theory of Mean Field Games with Applications 1-11. Springer, 2018.
[10] Pratik Chaudhari, Adam Oberman, Stanley Osher, Stefano Soatto, and Guillaume Carlier. Deep relaxation: partial differential equations for optimizing deep neural networks. Reseurch in the Mathematical Sciences, 2018.
[11] Junhua Chen, Lorenz Richter, Julius Berner, Denis Blessing, Gerhard Neumann, and Anima Anandkumar. Sequential controlled langevin diffusions. In ICLR, 2025. 3
[12] Kevin Clark, Paul Vicol, Kevin Swersky, and David J. Fleet. Directly fine-tuning diffusion models on differentiable rewards. In ICLR, 2024. 3, 6,
[13] Alexander Denker, Francisco Vargas, Shreyas Padhy, Kieran Didi, Simon Mathis, Vincent Dutordoir, Riccardo Barbano, Emile Mathieu, Urszula Julia Komorowska, and Pietro Lio. Deft: Efficient finetuning of conditional diffusion models by learning the generalised h-transform. arXiv preprint arXiv:2406.0178/ 2024. 3
[14] Carles Domingo-Enrich, A taxonomy of loss functions for stochastic optimal control. arXiv
preprint arXiv:24100.00345, 2024. 3
[15] Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, and Ricky T. Q. Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. In ICLR, 2025. 3, 5, 6, 8, 10, 17, 18, 19
[16] Carles Domingo-Enrich, Jiequn Han, Brandon Amos, Joan Bruna, and Ricky T. Q. Chen. Stochastic optimal control matching. In NeurIPS, 2024.
[17] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, Kyle Lacey, Alex Goodwin, Yannik Marek, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthesis. In ICML, 2024. 2,
[18] Jiajun Fan, Shuaike Shen, Chaoran Cheng, Yuxin Chen, Chumeng Liang, and Ge Liu. Online weigrt-weighted fine-tuning of flow matching with wasserstein regularization. In ICLR, 2025.
[19] Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models. In NeurIPS, 2023. 1,2
11

---

[20] Jin Feng and Thomas G Kurtz. Large deviations for stochastic processes. Number 131.
American Mathematical Soc., 2006. 3
[21] Wendell H Fleming and Jerome Stein. Stochastic optimal control, international finance and debt. Journal of Banking & Finance, 28(5).979-99) -969,999699 2004. 3
[22] Stephanie Fu, Netanel Y Tamir, Shobhita Sundaram, Lucy Chai, Richard Zhang, Tali Dekel, and Phillip Isola. Dreamsim: learning new dimensions of human visual similarity using synthetic data. In NewrIPS, 2023. 5
[23] Alex Gorodetsky, Sertac Karaman, and Youssef Marzouk. High-dimensional stochastic optimal control using continuous tensor decompositions. International Journal of Robotics Research, 2018. 3
[24] Carsten Hartmann, Ralf Banisch, Marco Sarich, Tomasz Badowski, and Christof Schütte. Characterization of rare events in molecular dynamics. Entropy, 16(1):350-377 2014.
[25] Carsten Hartmann and Christof Schütte. Efficient rare event simulation by optimal nonequi- librium forcing. Journal of Statistical Mechanics: Theory and Experiment, 2012(1 (012(11):P1I0))) 1004,
2012. 3
[26] Aaron Havens, Benjamin Kurt Miller, Bing Yan, Carles Commngg--Enriche Anuroop Sriram, Brandon Wood, Daniel Levine, Bin Hu, Brandon Amos, Brian Karrer, Xiang Fu, Guan-Horng Liu, and Ricky T. Q. Chen. Adjoint sampling: Highly scalable diffusion samplers via adjoint matching. arXiv preprint 1rXiv:2504.11713, 2025. 3
[27] Jiajun He, José Miguel Hernández-Lobate Yuanqi Du, and Francisco Vargas. Rne: plug-and- play diffusion inference-tine control and energy- training, 2025. 2
[28] Jiajun He, Paul Jeha, Peter Potaptchik, Leo Zhang, José Miguel Hernández-Lobato, Yuanqi Du, Saifuddin Syed, and Francisco Vargas. Crepe: Controlling diffusion with replica exchange, 2025. 2
[29] Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Le Bras, and Yejin Choi. Clipscore: A reference-free evaluation metric for image captioning. EMNLP, 2021. 5
[30] Lars Holdijk, Yuanqi Du, Ferry Hooft, Priyank Jaini, Bernd Ensing, and Max Welling. Stochastic optimal control for collective variable free sampling of molecular transition paths. In NeurIPS, 2023.
[31] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In ICLR, 2022.
[32] Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Pick-a-pic: An open dataset of user preferences for text-to-image generation. In NewrIPS, 2023. 5
[33] LAION. Laion aesthetic score predictor. httpa://laio..a//llgllliooooc
2024. Accessed: 2024-09-27.
[34] Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins, Yuqing Du, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, and Shixiang Shane Gu. Aligning text-to-image models using human feedback. arXiv preprint arXiv:2302. 12192, 2023. 2
[35] Daniel Liberzon. Calculus of variations and optimal control theory: a concise introduction.
Princeton university press, 2011. 7, 10
[36] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Lc. Flow matching for generative modeling. In ICLR, 2023.
[37] Weiyang Liu, Zeju Qiu, Yao Feng, Yuliang Xiu, Yuxuan Xue, Longhui Yu, Haiwen Feng, Zhen Liu, Juyeon Heo, Songyou Peng, Yandong Wen, Michael J. Black, Adrian Weller, and Bernhard Schölkopf. Parameter-efficient orthogonal finetuning via butterfly factorization. In ICLR, 2024. 10
12

---

[38] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In ICLR, 2023. 3.3
[39] Zhen Liu, Yao Feng, Michael J. Black, Derek Nowrouzzezarrii Liam Paull, and Weiyang Liu. Meshdiffusion: Score-based generative 3d mesh modeling. In ICLR, 2023.
[40] Zhen Liu, Yao Feng, Yuliang Xiu, Weiyang Liu, Liam Paull, Michael J. Black, and Bernhard Schölkopf. Ghost the shell: An expressive representation of general 3d shapes. In ICLR, 2024.
[41] Zhen Liu, Tim Z. Xiao, Weiyang Liu, Yoshua Bengio, and Dinghuai Zhang. Efficient diversity- preserving diffusion alignment via gradient-informen gflownets. In ICLR, 2025. 3,5
[42] Nanye Ma, Shangyuan Tong, Haolin Jia, Hexiang Hu, Yu-Chuan Su, Mingda Zhang, Xuan Yang, Yandong Li, Tommi Jaakkola, Xuhui Jia, and Saining Xie. Inference-time scaling for diffusion models beyond scaling denoising steps. arXiv preprint arXiv:2501.09732 2025. 2
[43] Sanjoy K Mitter. Filtering and stochastic control: A historical perspective. IEEE Control Systems Magazine, 16(3):67-76, 1996. 3
[44] Nikolas Nüsken and Lorenz Richter. Solving high-dimensional Hamilton-lacobi-elllaan pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Purtial differential equations and applications. 2:1-48, 2021. 3
[45] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. In NeurlPS, 2022.
[46] Ling Pan, Dinghuai Zhang, Moksh Jain, Longbo Huang, and Yoshua Bengio. Stochastic generative flow networks. In UAI, 2023. 2
[47] Huyên Pham. Continuous-time stochastic control and optimization with financial applications, volume 61. Springer Science & Business Media, 2009. 3
[48] Jakiw Pidstrigach, Elizabeth Baker, Carles Domingo-Enrich, George Deligiannidis, and Nikolas Nüsken. Conditioning diffusions using malliavin calculus. In ICML, 2025. 3
[49] Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall, Dreamfusion: Text-to-3d using
2d diffusion. In ICLR, 2023. 5
[50] Warren B. Powell and Stephan Meisel. Tutorial on stochastic optimization in energy- part i:
Modeling and policies. IEEE Transactions on Power Systems, 31(2):1559-11477 2016. 3
[51] Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, and Katerina Fragkiadaki. Aligning text-to- image diffusion models with reward backpropagation. arXiv preprint arXiv:2310.00533, 2023.
8
[52] Zeju Qiu, Weiyang Liu, Haiwen Feng, Yuxuan Xue, Yao Feng, Zhen Liu, Dan Zhang, Adrian Weller, and Bernhard Schölkopf. Controlling text-to-image diffusion by orthogonal finetuning. In NewrlPS, 2023. 10
[53] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In /CML, 2021.
[54] Sebastian Reich. Data assimilation: The Schrödinger perspective. Acta Numerica, 28:635-711,
2019.
[55] Tim Salimans and Jonathan Ho. Should EBMs model the energy or the score? In Energy Based Models Workshop ICLR 2021, 2021. 5
[56] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
13

---

[57] Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.033300 2024. 2
[58] Raghav Singhal, Zachary Horvitz, Ryan Teehan, Mengye Ren, Zhou Yu, Kathleen McKeown, and Rajesh Ranganath. general framework for inference-tinge scaling and steering of diffusion models. arXiv preprint 2xXiv:2501006843 2025. 2
[59] Yang Song and Diederik P Kingma. How to train your -assed models. arXiv preprint arXiv:2101.03288 2021.
[60] Evangelos Theodorou, Freek Stulp, Jonas Buchli, and Stefan Schaal. An iterative path integral stochastic optimal control approach for learning robotic tasks. IFAC Proceedings Volumes, 44(1)111594-1161, 2011. 18th IFAC World Congress.
[61] Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng. Tommaso Biancalani, and Sergey Levine. Fine- tuning of continuous-time diffusion models as entropy control. arXiv preprini arXiv:2402. 15194, 2024. 2,3
[62] Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. Transport meets variational inference: Controlled monte carlo diffusions. In ICLR, 2024. 3
[63] C. Villani. Topics in Optimal Transportation. Graduate studies in mathematics. American
Mathematical Society, 2003. 3
[64] C. Villani, Optimal Transport: Old and New. Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 2008. 3
[65] Team Wan, Ang Wang, Baole Ai, Bin Wen, Chaojie Mao, Chen-Waer Xie, Di Chen, Feiwu Yu, Haiming Zhao, Jianxiao Yang, et al. Wan: Open and advanced large-scale video generative models. arXiv preprint (ixXix:2033203/11 2025,
[66] Luhuan Wu, Brian Trippe, Christian Naesseth, David Blei, and John P Cunningham. Practical and asymptotically exact conditional sampling in diffusion models. In NeurIPS, 2024.
[67] Xiaoshi Wu, Yiming Hao. Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng
Li. Human preference score v2: A solid benchmark for evaluating human preferences of text-to-image synthesis. arXiv preprint arXiv:2306.099.. 2023. 5
[68] Xiaoshi Wu, Keqiang Sun, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score: Better aligning text-to-image models with human preference. In CVPR, 2023. 5
[69] Jianfeng Xiang, Zelong Lv, Sicheng Xu, Yu Deng, Ruicheng Wang, Bowen Zhang, Dong Chen, Xin Tong, and Jiaolong Yang. Structured 3d latents for scalable and versatile 3d generation. In CVPR, 2025.
[70] Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. In NeurlPS, 2024. 3, 6,
[71] Taeyoung Yun, Dinghuai Zhang, Jinkyoo Park, and Ling Pan. Learning to sample effective and diverse prompts for text-to- text-to-image generation. In CVPR, 2025. 2
[72] Dinghuai Zhang, Ricky TQ Chen, Cheng-Hao Liu, Aaron Courville, and Yoshua Bengio. Diffuusinn generative flow samplers: Improving learning signals through partial trajectory optimization. In ICLR, 2024. 2
[73] Dinghuai Zhang, Aaron Courville, Yoshua Bengio. Qinqing Zheng, Amy Zhang, and Ricky Q. Chen. Latent state marginalization as a low-cost approach for improving exploration. In ICLR, 2023. 2
[74] Dinghuai Zhang, Hanjun Dai, Nikolay Malkin, Aaron C Courville, Yoshua Bengio, and Ling Pan. Let the flows tell: Solving graph combinatorial problems with gflownets. In NeurIPS, 2023. 2
14

---

[75] Dinghuai Zhang, Jie Fu, Yoshua Bengio, and Aaron Courville. Unifying likelihood-fiii inference with black-box sequence design and beyond. In ICLR, 2022.
[76] Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. In ICML, 2022.
2
[77] Dinghuai Zhang. Yizhe Zhang. Jiatao Gu, Ruixiang Zhang, Joshua M. Susskind, Navdeep Jaitly, and Shuangfei Zhai. Improving GFlownets for text-to-image diffusion alignment. Transactions on Machine Learning Research, 2025. 2
[78] Longwen Zhang, Ziyu Wang, Qixuan Zhang, Qiwei Qiu, Anqi Pang, Haoran Jiang, Wei Yang, Lan Xu, and Jingyi Yu. Clay: controllable large-scale generative model for creating high-quality 3d assets. ACM Transactions on Graphics (TOG), 2024.
[79] Qinsheng Zhang and Yongxin Chen. Path integral sampler: stochastic control approach for sampling. ICLR, 2022.
[80] Wei Zhang, Han Wang, Carsten Hartmann, Marcus Weber, and Christof Schütte. Applications of the cross-entropy method to importance sampling and optimal control of diffusions. SIAM Journal on Scientific Computing, 3666) 2014.
15

---

Appendix
Table of Contents
A Theoretical Connections between Optimal Control Formulations
17
A.1
Deriving adjoint matching from Pontryagin's 's maximum principle
A.2 The adjoint method vs. adjoint matching for deterministic optimal control
B Bounds of Resulted Distributions
20
B.1
Bounding the Wasserstein-2 Distance
20
B.2 Bounding the KL Divergence..
Experiment Details
C.I Finite difference for value consistency
22
C.2 More implementation details
22
D Additional Figures
23
E Evolution of Generated Samples
24
F More Generated Samples
25
16

---

A Theoretical Connections between Optimal Control Formulations
A.1 Deriving adjoint matching from Pontryagin's maximum principle
For general dynamical system ă = f(x, u, t) and control problem min, So L(x(t), u(t), t) de
((x(T)), we can define the Hamiltonian (x, f, a. & L(x, u) + af t) and the adjoint state a(t) that satisfies a(t) H and a(T) = (T)) The Pontryagin's maximum principle (PMP) states that the optimal control u² satisfies u (t) arg max, 11 (xx (t), u t, a (t)).
(2 setup and the control formulation in Remark 1, we have For our flow matching model ii
f. 1110(12, t) Ubase (2e, t)ll r(a2) dt
min Es, = 8
(18)
t,, A |1 t) Ubuse (2., t)]1 + a vo (2)
(19)
X v (110(22, Vaaso t)]l + a(x, (octt))vo(aa vo(z,t)).
ä(z,t) t) - 2
(20)
with the terminal constraint a( 1) - vr(z1). In order to solve v* arg max, (x, u, a),
with such quadratic form we have
a
Ubase
X
0.
(21)
Therefore, we can have an intuitive algorithm following the common practice of solving PMP:
1. Solve forward ODE, i; ua( t) (), to obtain (xx))=o)
2. Solve backward ODE in Equation 20, à = - V, a(x1, 1) -r((x),, to obtain {ar (a))
3. Update velocity field vi by minimizing matching loss
(()) = [A t) + aill 2
(22)
This is already very similar the adjoint matching algorithm proposed in [15] (but for stochastic control settings), where the authors obtain their algorithm from a different derivation and slightly different assumptions.
Comparing this Equation 21 with Equation 8, we can easily get
axx, L) (v(x,), (x,
(23)
which, maybe not that surprisingly, indicates adjoint matching and VGG-Flow share the same vector matching objective (Equation 22 and Equation 8). They differ in the way of how to obtain the matching target (either a, or VV).
A.2
The adjoint method vs. adjoint matching for deterministic optimal control
Here we give a more general discussion between the adjoint method (for open-loop control) and adjoint matching (for closed-loop control) the deterministic control setting.
A.2.1 Open-loop vs closed-loop formulations of deterministic optimal control
Consider the deterministic control problem equation with a quadratic cost and affine control. That is, L(x, u. 1 <1ulll@(e [xx t), where Q(t) Rddd is positive-defiite matrix, and
[|ull2 = u T Qu, and the drift is fann(æ, u,t) b(x, t) + u. Hence, the control problem reads
("(lluco)Be + s(Xe.t) dt + 2(X)),
min Mlc A u: +R
(24)
S.L. X, = b(x, u(1),
Xo z0.
(25)
This is an open-loop control problem, because the control u(l)) does not depend explicitly on the state X₁. Note however that since both the starting point and the dynamics are deterministic, given the
17

---

function u j - R it is possible to determine the state X1, and hence u(t) can be defined to depend implicitly on X,. Alternatively, consider the control problem
J Juu] J." [[lm(X,,0)] 11200 + s(Xx,)) dt ((X)),
min (0.
(26)
X, = b(X,,t) t) t),, Xo 20-
(27)
In this case, is a function that depends explicitly on the state X which makes this closed-loop control problem.
In general, closed-loop control problems are more general than open-loop problems, but in our case both problems are actually equivalent because the deterministic dynamics and initial conditional make it possible to make u(t) depend on X, implicitly. In other words, for any open-loop control w: [0, 7] + Rd, we can define a closed-loop control à Rd x [o, + R4 by setting ū(x, () u(I). And for any closed-loop control
Rd x [0, T + R we can define an open-loop control u: [0, 7] + Rd by setting u(t) u(x,, t), where X (X,)=e(o.1) satisfies the ODE in equation 27.
Thus, finding the solution to the open-loop problem in equations 24-25 is equivalent to finding the solution the closed-loop problem in equations 26-27. As naturally gives rise to the adjoint matching loss for deterministic optimal control, while the latter yields the basic adjoint matching loss, which is simply a reformulation of the adjoint method.
next, the former formulation
A.2.2 Solving the closed-loop problem: the adjoint method for deterministic optimal control
Using argument similar to the used in [15, Prop. 2] for stochastic optimal control, we derive the continuous-time version of the basic adjoint matching loss for deterministic control: Proposition 2 (Basic adjoint matching for deterministic control). Consider the adjoint ODE
dt (t: X, u) (Vx, - (Vx,, ((;; X1; (1(X).))+ (015,,t)+11(((() (l,(X,,
(28)
(1: V
(29)
Suppose that the control IRd x [0,T] + Rd is parameterized by 8, and let Juu be the closed-loop control objective in equation 26. Then, the gradiens VoJ[u] is equal to the gradient of this loss:
Lasic - -atchaatch (1) L Q(0) ( dt,
(30)
X s.t. X, b(X,,t) + x({,,t), = s=sop-ggaddeen(())
where ū stop-gradien control u are artificially set to zero.
means that the gradients of with respect to the parameters 0 of the
Proof. The proof mirrors the proof of [15, Prop. 2]. If we define the adjoint state
alt, X, x) Vx. (Hlll () + ((xx,t) dt + (((x))
(31)
where X a solution of X, = b(Xe,t) u(X,,t))
(32)
we have that alt, X,,) u) satisfies the adjoint ODE in equations 28-29, In analogy with equation 32 of [15], we have that
d j [ d0
T 0 Il llu(Xe
0) dt + (- Ou((x.:))
00 alt, u) dt,
(33)
00
18

---

where and de have that
denote the total and partial derivatives with respect to 6. Completing the square, we
Be
0 00
tu(X,,t)
(ux,,,)))l +
00
1 0
Ou(Xe,t) Q()))2(()))//(((. X, (34)
2 00 ((()),,))0()))))()()))) (Xx,1)))
09
1 0
|u (u) + + ((t))
2 00
where Q(t)1/2 is defined as the matrix with the same cigenvectors as Q(t) and eigenvalues equal to the square root of the eigenvalues of (t).. and stopgrad(u). Notice adjoint a does depend on 0. Plugging equation 34 into equation 33, we can finally rewrite the gradient as the gradient of Basic--Ase Adj- Match-
A.2.3 Solving the open-loop problem: the adjoint matching loss for deterministic optimal control
Proposition 3 (Adjoint matching for deterministic control). Consider the lean adjoint ODE:
dt a(t: X) [V (Vx,,(X,, a(t: xxxx
d
(35)
(11:)) (X).
(36)
Suppose that the control u. [0, RRd is parameterized by and let Huu be the open-loop control objective in equation 26. Then, the gradient [u][u] is equal to the gradient of this loss:
Lndj- (u)
[u(t) +(())-1/2a(tt dt,
(37)
X S.t. X, = b(X,,t) t) + a(t), u = stop- groppggradeent(u)
where ū = stop-graadse gradient(u) means that the gradients of a with respect to the parameters of the control are artificially set zero.
Proof. The proof mirrors the proof of [15, Prop. 2), and the proof of our Prop. 2. If we define the adjoint state
+ f(X,,t) dt + ((x)))
x)=
Vx. =
(38)
where X is a solution of X, = b((,, u()).
(39)
we have that att, X) satisfies the adjoint ODE in equations 28-29. Note that unlike in Prop. 2, the control u({) does not depend the state X., which simplifies expressions substantially as V Vx, u(() = 0. In analogy with equation 32 of [15] and our equation 33, we have that
d Jum d0
("
a 2e (1(()) di S" @u(() 00
a(l, X)d1, dl,
(40)
and completing the square as in equation 34, we obtain that
8 Iu (t)l
Bu(t) 20
alt, X)
2 30
80 uu(() + ((t)>>/att X)ll Plugging this equality into equation 40 concludes the proof. 0
19

---

B Bounds of Resulted Distributions
B.1 Bounding the Wasserstein-2 Distance
We first analyze the relationship between our objective of VGG-Flow and the 2- 2-Wasserstein distance (W2). We show that our objective minimizes strong upper bound on W2(P1. 91).
Let po = qo be the initial distribution. Consider the two flows, coupled by their initial condition:
z (o t), Zo Po 11 = Ubase (Ve, t), yoo 2o
Pt 40 y1 q:
By definition, the squared W2 distance is the minimum expected squared distance over all possible couplings. Our choice of zu = yo is one such coupling, so it provides an upper bound:
W2(Pe, 9:)² = ||| yall²)
Proposition (W₂ Bound via Grönwall's Inequality). Assume the base vector field Ubase is L-- Lipschitz in x. Then the W2 distance is bounded by the L2 FM loss:
W2(p1. S C
Ep. ( f)
where C 222++ is constant.
Proof. Let A, = Jt. - and u(t) B|] We have u(0) = El]l20 yo 0. The time derivative is å = ug (o( Daase (ye, t). Let üg (z, 0e (2, t) Ubasc (,r, t).
A 22(,; Â.)
-, [ng (ze: t) Dbase (2 t)] + [Dbase (2e, Dbase (yt, t)]) = -((,, (1 t)) + 2(A Vome (2 t) Ubase (Bs, t))
We apply the Cauchy- Sauchh-Schwwrz inequality the first term and the L-Lipschitz condition the second:
_ 11..1 - 21l f) + (
Using Young's 's inequality (2ab < a² + 122) the first term gives:
| s (1| + (,,, ., 1) 2. 11 (2L + 1)11A,11 [1 t)l
Taking the expectation and letting ((t) = Ep. || (z [) we have the differential inequality:
ú(t) < (2L + 1)t(t) b(t)
By the integral form of Grönwall's inequality. which states that
with u(()= = 0, then u(t) s L call-o)((s)de
it) < au(() t) + b(t)
Applying this with (2L 1), the solution at 1 is:
u(1) < L (2L+1)0->))))d 2 L( b(8)ds
Since W2(pi (1).g))3 u(1), we arrive at the bound:
W2(p1. 24+1 [( Ep, (0l0000 t) dt
(41)
This result confirms that our objective a theoretically sound one for minimizing an upper bound on the W2 distance.
20

---

B.2
Bounding the KL Divergence
We now analyze the KL divergence. Unlike the W2 distance, the KL divergence is sensitive to changes in density, which are governed by the divergence of the vector field.
The marginal densities satisfy the continuity equations (0,1))
(s, t) -V (2, ()) t))
(eqe(r, - (ge (I., 1) (z. 1))
Proposition 5 (KL Divergence Identity for ODEs). Assume the vector fields Vg, bhase and densities Pt. 9t are sufficiently smooth and have sufficient decay at infinity such that all boundary terms from integration by parts vanish. Then, the exact identity for the final KL divergence is:
DKL (pi (oill - - Ey, [To(xr, 1)- log dt - -f0 Ep. V (0( d)
where ig 00
Ubuse-
Applying a bound to the first term (as we did for the W2 proof) gives the final inequality:
2 [ Ep, (00 {)l²) +
1
DKL (pillq1) -
C(p.g)
Ep. V (0 t)] dt (42)
Path-Dependeus Tomn
(A) Value Gradient Matching
(C) Divergence
where C(p, q) = 1 So px, VI log qe |=] dt is a functional that depends on both the target path and the learned path pt.
Remark 6 (Justification for the L2 Proxy Objective). Equation 42 shows that the KL divergence bounded by the L2 value gradient matching loss (Term A), a path-dependent term (Term B), and a divergenneedeeeeddett term (Term C). Term (B) depends
both the learned path ps and the target path qt- It can be bounded, for example, if the target score function has a uniform bound (i.e., || log qe (x:] M, for all x, L), which would imply E,, [l| log |/]] < M2. The primary challenge is that Term (A) and Term (C) are geometrically independent, and Term (C) computationally expensive to estimate. We therefore use the value gradient matching loss computationally efficient proxy objective. We empirically justify this choice, as our finetuned models produce high-quality samples. This success suggests that for our network architecture and problem setup, minimizing Term (A) is sufficient, and the "missing" divergence term (Term C) is implicitly regularized or remains small, likely due the implicit bias of the neural network.
21

---

C Experiment Details
C.1
Finite difference for value consistency
Our value consisteney loss requires the costly computation of second-order gradients during back- propagation. To save memory and time, we instead use finite differences to approximate the terms (with u
Ubase
g0(x) cu, g0(
Ot (o(,t)
(43)
(VV9a)" (Ubuse
1 9 (x))))
g0 + + 1)1, - Do gol - -eVlu(ar, (),,))
(44)
20
Pbeao( + e Wlggs 1)1,
Ubaso (Ie
V (
(Vodbeel 96) (x)..))
(45)
2e
where (.) stop-gre (roop-gradien((). The stop gradient operations nested function calls prevent
second-order gradients during backpropagation. Empirically, we find this approximation works well.
C.2 More implementation details
In our experiments, we choose 7/2 = f2 in Equation 16 if not otherwise specified. We use CFG scale of werG = 5.0 for all experiments, and the velocity fields of both the base and finetuned models are CFG-composited as v(x, f; c) (1 + WCFG) (xz,) t; c) WCrGu( t; Ø). We stop the gradients on v(xx, t; Ø) as we found this leads faster convergence. We use the best learning rates (in terms of fast stable reward convergence) for each method instead of fixed ones,
observe that methods like ReFL and DRaFT can be unstable for very large learning rates. Specifically, we use 5e 4 for VGG-Flow all reward models, 5e 5 for VGG-Flow- -MM HPSv2 and PickScore, and le 4 for all others. We use the standard Adam W optimizer with 31 = 0.9, 32 = 0.999 and weight decay 1e 2. We clip the norm of network update gradients to 1. We use bfloat 1 computation for the flow matching model but float32 for the reward model due to numerical precision issues.
22

---

D
Additional Figures
quadratic r schedule
linear ni schedule
244
CLP
M1
212
921
200
300
Updnns
Updane
Figure 13: Evolution of metries for different schedule (experiments on Aesthetic Score). The linear schedule of 1) leads faster convergence.
Figure 14: Trade-offs between metrics for different 1/ schedule (experiments Aesthetic Score).
Subsampling Rate-0.25
Subsampling Rate-0.5
242
340
202
0:1
Up Aep
Up Sep
1b600 Stp
Figure 15: Evolution of metrics for different transition subsampling rates (experiments on Aesthetic Score).
ETP Halletaa Seard Sechel
Figure 16: Trade-off' between metrics for different transition subsampling rates (experiments on Aesthetic Score).
23

---

E Evolution of Generated Samples
We show in Figure 17 that our method is more capable of preserving the prior from the base model during the finetuning process.
Pretrained
VGG-Flow
ReFL
DRaFT-LV
AM
VGG-Flow-PMP
Epoch 40
Epoch 150
Epoch 250
Figure 17: The degradation of image quality of baselines, compared to the evolution sequence of results produced by our method.
24

---

More Generated Samples
Pretrained
VGG-Flow
ReFL
DRaFT
AM
VGG-Flex.PMP
Figure 18: More qualitative results on Aesthetic Score.
25

---

Pretrained
VGG-Flow
ReFL.
DRaFT
AM
VGG-Flow-PMPP
Prompt: street with cars lined with poles and wires.
Prompt: man and dogs riding scooter.
Prompt: Two chill the bathtub, laying down.
Figure 19: More qualitative results on HPSv2.
Pretrained
VGG-Flow
ReFL.
DRaFT
AM
VGG-Flow-PMM
Prompt: black lab catching tennis ball.
Prompt: with blue blue hill.
Prompt: style minecratt.
Figure 20: More qualitative results PickScore.
26

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The claims are accurately reflected through corresponding sections.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We incorporate a limitation discussion in the discussion Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
27

---

Justification:
We provide rigorous proof to all theoretical results in Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide the algorithm pseudocode in Algorithm and experimental details in Section 5 which ensures reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
28

---

Answer: [Yes]
Justification:
We release the whole set of code that can reproduce our algorithm.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /Co///eodimmsseeeeeeee
/ /nips. wwtppclcccc/ccccccccccccccccclct /tteeguudddeeeeeeuumiiieeeee ssssssssiiciiio for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
We provide all experimental details in Section and Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Following common practice, provide standard deviation based repeated random seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
29

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Our compute resources are written in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http///eerrp...eee http://///rpppppppp/t/a///////// htttc///tthiiiiiieeieeeeieeaeeee ines?
Answer:
[Yes]
Justification:
The research conducted in the paper conform with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discussed the potential societal impact in the Section 7.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
30

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Not applicable for our case.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Every asset use is properly credited and used.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
31

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We will provide detailed documentation for the released code.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details of the datasettcooeemmodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve human subjects crowdsourcing. All experiments are conducted using publicly available datasets or synthetic data without interaction with human participants.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
32

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs are only used for writing polish and formatting purposes.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy (https: http:://eeuupppeeeeee cct.cc///////. httprripps.cc//oneeeenneeee///// for what should or should not be described.
33


Reviews:
[{"review_id": "55waLcMTzr", "reviewer": "Reviewer_FPr9", "summary": "This paper proposes VGG-Flow, a method for aligning flow matching models with reward models using a value gradient derived from optimal control theory. The approach estimates a value gradient via the Hamilton-Jacobi-Bellman equation and uses it to guide residual velocity updates. Experiments show improvements over existing baselines in reward alignment while preserving prior and diversity.", "questions": "Can the authors clarify how sensitive the method is to the choice of reward model? For instance, do different reward functions lead to different optimization dynamics or mode collapse rates?\n\n\nHow does the proposed method behave when the reward model is noisy or weakly aligned with human judgment? Could this be mitigated through smoothing or ensemble strategies?", "limitations": "The method assumes access to pretrained, differentiable, and trustworthy reward models, which may not always be available in other domains (e.g., text, video).\n\nAlthough the proposed approach is efficient, its performance on much larger models (e.g., >7B parameters) is not yet demonstrated.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- The paper proposes a principled formulation based on deterministic optimal control, which is novel in the context of aligning flow matching models.\n- The idea of parametrizing and matching value gradients offers a computationally efficient alternative to adjoint-based methods.\n- Empirical results on reward score, sample diversity, and prior preservation are solid and consistently outperform baselines.\n\nWeaknesses:\n- The method assumes access to a pretrained reward model, which may be hard to obtain or validate in real-world settings. The paper does not discuss the impact of reward model quality on the results.\n- The gradient-matching framework is conceptually similar to existing diffusion alignment methods; the novelty lies mainly in adapting it to flow matching models.\n- While the proposed method preserves prior better than other baselines, the trade-off between reward maximization and sample quality remains heuristic, and more theoretical or empirical analysis of this tension would strengthen the work.", "was_revised": true, "final_justification": "Thank you to the authors for the clear and detailed rebuttal. The explanations regarding the availability of reward models, the theoretical basis of the trade-off parameter, and the principled derivation from optimal control addressed my main concerns. These clarifications improved my understanding and confidence in the work. While some questions remain (e.g., robustness to weak reward models), I believe the core contribution is solid, and I have raised my score accordingly.", "rebuttal": "We appreciate the reviewer for their acknowledgement on our contribution and raising their concerns.\n\n> The method assumes access to a pretrained reward model, which may be hard to obtain or validate in real-world settings. The paper does not discuss the impact of reward model quality on the results.\n\n> The method assumes access to pretrained, differentiable, and trustworthy reward models, which may not always be available in other domains (e.g., text, video).\n\nWe argue that the pretrained models, even the differentiable ones, are common and relatively easy to obtain these days. Pretrained reward models (with differentiable neural nets) are prevalent in the field of foundation models, and many of them are open-sourced and available on HuggingFace. There are publicly available preference datasets (such as HPDv2) with more than 100k text-image pairs. People from the industry are working on larger scale datasets; for instance, the CVPR best paper “Rich human feedback for text-to-image generation” [1] builds a pixel-aligned dense reward model (i.e., they have a map of rewards, not just one single scalar) from a large-scale dataset. On training these reward models: people typically assume the Bradley–Terry preference model, which turns out to be a simple logistic regression problem: telling which one is preferred in a pair of data points with binary classification. Beyond these models, we may also adopt pretrained VLMs to extract differentiable likelihood values (e.g., [2]) as rewards.\n\nOur paper aims to build a generic RL finetuning method that works on any given differentiable reward model and therefore we believe a detailed discussion on reward models themselves are beyond the scope of this particular work.\n\n> The gradient-matching framework is conceptually similar to existing diffusion alignment methods; the novelty lies mainly in adapting it to flow matching models.\n\nThe resulting algorithm indeed seems similar to other existing works in high level, and one can surely design such heuristics and try them out. The key differentiating factors of our paper are that 1) our method is derived with a principled approach from an optimal control perspective where the gradient of value function is naturally the velocity, 2) we show what the correction term for the naive heuristics should be.\n\n> While the proposed method preserves prior better than other baselines, the trade-off between reward maximization and sample quality remains heuristic, and more theoretical or empirical analysis of this tension would strengthen the work.\n\nThank you for pointing it out. We model the finetuning objective in Equation 9 in line 134 in which the $\\lambda$ parameter controls the trade-off between reward following and prior preservation. We accidentally used $\\beta=1/\\lambda$ to indicate this trade-off in the ablation study (Figure 6).\n\nFrom a mathematical side, the $\\beta=1/\\lambda$ parameter on the reward model (from Eq. 9), to some extent, is similar to the classifier-free guidance scale (typically set to 7.5 or 5.0) which is something heuristic even at this point. Similar to prior works such as adjoint matching, when $\\beta$ is large, the method focuses more on reward maximization and vice versa. The trade-off is inevitably empirical. One may prefer better reward following instead of sample quality/prior. For instance, if one aims for generating sci-fi looking images, then imposing a very strong prior of natural images makes less sense. \n\n\n> Can the authors clarify how sensitive the method is to the choice of reward model? For instance, do different reward functions lead to different optimization dynamics or mode collapse rates?\n\nEmpirically we do not find the model very sensitive to reward models, though the margin compared to other baseline methods may vary. For instance, we observed that finetuning with Aesthetic Score is prone to overfitting and our method is more robust and more efficient; on HPSv2 and PickScore, the baseline methods do not suffer that much.\n\n> How does the proposed method behave when the reward model is noisy or weakly aligned with human judgment? Could this be mitigated through or ensemble strategies?\n\nA prior work [3] does apply reward smoothing (by averaging the rewards of sampled neighboring points) due to concerns about the non-smoothness of the reward function. \nIn our work we explored this trick in an early stage but did not find it useful in our setup. We have not explored ensemble strategies because we focus on the setting where a fixed reward function is given. That being said, DRaFT (one of the baselines, [4]) has an interesting experiment where they try to use a linear combination of different reward functions to achieve a mixed effect.\n\n ---\n[1] Rich Human Feedback for Text-to-Image Generation. Liang et al. CVPR 2024. https://arxiv.org/abs/2312.10240\n\n[2] Beyond Yes and No: Improving Zero-Shot LLM Rankers via Scoring Fine-Grained Relevance Labels. Zhuang et al. NAACL 2024. https://arxiv.org/abs/2310.14122\n\n[3] Efficient Diversity-Preserving Diffusion Alignment via Gradient-Informed GFlowNets. Liu et al. ICLR 2025. https://arxiv.org/abs/2412.07775\n\n[4] Directly Fine-Tuning Diffusion Models on Differentiable Rewards. Clark et al. ICLR 2024. https://arxiv.org/abs/2309.17400"}, {"review_id": "IMlh6aoUJY", "reviewer": "Reviewer_UNtT", "summary": "The authors consider a composite loss to finetune Flow Matching models following a reward function designed to better align generated images with human preferences.\nThis loss is a trade-off between a $\\ell_2$ term that keeps the fine-tuned model closed to the pretrained velocity field (the prior model)  and a term that maximizes in expectation the reward on the generated images.\nTo circumvent the need to generate images in the loss (which involves solving the ODE), the authors leverage optimal control theory: they instead use a surrogate network that models the gradient of the value function in the Hamilton-Jacobi-Bellman (HJB) equation.", "questions": "First, I think the paper deserves a better positioning in the literature:\n- One motivation stated by the paper is that there is a strong difference between diffusion and flow matching in order to leverage OC. The sentence \"the key challenge\" l25 should be clarified.\n- I think the introduction does not clearly present the goal of the paper. As I understand, it is to avoid methods both offline RL methods  (where one would need the pretrained training set) or online (where one would need to generate new samples on the fly).\n\n**On the preliminaries on FM models** There seems to be an error in $u(x_t | x_1)$. I guess it should be $\\frac{x_1 - x_t}{1-t}$. I also wonder why the authors specify a variant (rectified flows) so early while some more crucial details are missing for a clear introduction to FM: define how $x_t$ is built (linear interpolation $x_t = (1-t)x_0 + tx_1$ ?).\nI understand the need for brevity but I think this subsection could be handled a bit more carefully, especially given some confusion on $u(x_t |x_1)$ seem to arise again later in the paper (l204: the factors are now inverted).\n\n**On the preliminaries on Optimal Control**\n- As far as I understand, the use of optimal control theory in this ODE FM setting is the main highlight of the paper: then, this subsection is important and has to be more pedagogical.  Notably, when writing the continuity equation as a HJB equation: precisely define what is $V$ ($- \\log p_t$ ? ), the control signal, ... or provide a reference.\n- The sentence from l.121 to l.123 is  unclear: is it that the CE does not enable to access the $p_t$ given a pretrained velocity field or is it taht the maximum likelihood is not a good objective to train a velocity field  (due to the infinite number of paths, from which comes Flow Matching which explicitly regress against a given conditional velocity field).\n\n**4.1: The method**\nThere are still many points that remain unclear.\n- I think the setup could be more rigorously exposed (differentiable reward function, etc).\n- While the original loss looks very natural and looks quite common (similar to [20]), I have trouble understanding how the reparameterization through the HJB equation makes things computationally easier (if that is the goal). How can training a new model (which is also a Stable Diffusion UNET) can make things go quicker than generating samples (with OT/rectified FM models, this can be achieved with few Euler steps as noted by the authors themselves). I feel like that while being elegant, it introduces the need for many approximations afterwards, one of them being that at the end you still generate points (but you do it with one Euler step).\n\n**On the experiments**  In Table 1, it seems that concurrent models achieve lower FID which is desirable and yet are not put in bold characters.\n\n- While the visual results look quite convincing, there is always this trading off between the two terms in the loss (the one maximizing the reward and the one forcing to stay close to the base model): I would like to see of this parameter impacts the metrics, and same for concurrent methods that must have some kind of similar tradeoff.\n As it is in the paper, it is still unclear the impact of each parameter on the tradeoff (is it $\\beta$, $\\lambda$, $\\alpha$, $\\eta_t$: especially given by Figure 6 where there seems to be a discrepancy between the $\\beta$ in the caption and the $\\lambda$ in the legend ?)\n\n- The main missing point for me regarding the experiments is whether you achieve faster training time (you gain in terms of iterations but what about computation time ?).  Among the benchmarks models you pick, is there any methods where images are generated (since, as i understand you claim to circumvent this aspect).", "limitations": "Yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\nThe overall structure of the paper is quite clear and easy-to-follow. The writing style is good.\nThe fact that the paper takes its roots in the optimal control theory is quite elegant.\n\n**Weaknesses**\n\nThe point that puzzles me the most is the additional complexity the method seems to require: it involves an extra network jointly trained in the loss. This then imply many approximations and extra hyperparemeters to maintain the method tractable (e.g. parameterization of $g_\\phi$, stop gradient, finite order methods), while the initial objective doesn't seem at first overly complex. See questions below for more details.", "was_revised": true, "final_justification": "I raised my score to weak accept, as the authors have provided a strong rebuttal, with additionnal ablation studies and explanations on their method. I believe the interest of the paper lies in the new use of the HJB equation for FM alignment with a given reward function. I do not give a higher score because the practical implementation of the method still relies on several approximations.  Moreover, it does not eliminate the need for empirical tuning to balance reward maximization and sample quality.", "rebuttal": "We appreciate the reviewer for their time and efforts in reviewing our paper and their constructive questions.\n\n> ...many approximations and extra hyperparameters to maintain the method tractable (e.g. parameterization of g_phi, stop gradient, finite order methods), while the initial objective doesn't seem at first overly complex.\n\nThank you for raising this issue. We do agree that there are probably engineering details in implementing our method. In the meantime, we would like to highlight that the initial objective is not an easy one: tracking the exact value function (which indicates the probability flow) itself is hard. In discrete time, obtaining the value function is often done with Bellman equation / temporal difference (TD) methods and people have designed many tricks to stabilize RL training: target network (with gradients detached) with exponential moving average, trust region / KL clip (as in PPO), replay buffer, etc. The continuous optimal control shares the same nature as the discrete counterpart and therefore cannot be much simpler. In addition, the HJB equation involves derivatives in high dimensional spaces and therefore some approximation has to be done there (another potential approach is that we turn this problem back to the discrete-time RL problem).\n\n> One motivation stated by the paper is that there is a strong difference between diffusion and flow matching in order to leverage OC. The sentence \"the key challenge\" l25 should be clarified. \n\nWe apologize for this vague statement. We meant that in diffusion models, we have the (approximate) score function $s_\\theta(x_t, t) = \\nabla \\log p_t(x_t)$ and therefore performing distribution matching (with the reward function) is relatively easy. For flow matching models which is trained by matching the velocity $v_\\theta(x_t,t)$ with some predefined reference path $u(x_t; x_0)$, the \"probability\" of a state $(x_t, t)$ is not obvious for general flow matching models (for instance, models other than rectified flow models). It is even harder in the case that we are only given model weights but not how it is trained. This is why we turn to the optimal control formulation and solve for the value function (that indicates the probability flow) in an online manner.\n\nOn rectified flow models: for a pretrained rectified flow model, we know an easy way to compute probability flow $p_t(x_t)$; however, once this model is finetuned without a dataset, the model starts to deviate from being a rectified flow model and the probability flow in theory cannot be computed that way.\n\n> I think the introduction does not clearly present the goal of the paper. As I understand, it is to avoid methods both offline RL methods (where one would need the pretrained training set) or online (where one would need to generate new samples on the fly).\n\nSorry for the confusion. The purpose of the paper is to find an online and computationally-efficient RL method that finetunes any flow matching model with any differentiable reward model. Offline methods typically assume the presence of trajectory data for training, while here we only have a reward function without any dataset given.\n\n> On the preliminaries of FM method.\n\nWe apologize for the typo and will fix it in later versions.\n\n> On the preliminaries on Optimal Control. As far as I understand, the use of optimal control theory in this ODE FM setting is the main highlight of the paper: then, this subsection is important and has to be more pedagogical. Notably, when writing the continuity equation as a HJB equation: precisely define what is V(-log pt), the control signal, ... or provide a reference. \n\nWe will provide more pedagogical details about the foundation of optimal control in Sec 3.2 in the final version. $V$ is the value function of the control problem, which we define in Eq. 5. Currently, we refer to the Sec 4.1 of [1] paper for a self-contained preliminary of optimal control formulation.\n\n> The sentence from l.121 to l.123 is unclear: is it that the CE does not enable to access the pt given a pretrained velocity field or is it that the maximum likelihood is not a good objective to train a velocity field (due to the infinite number of paths, from which comes Flow Matching which explicitly regress against a given conditional velocity field).\n\nThank you for raising this point and sorry for the unclear explanation here. We wanted to mention that with only terminal rewards or distributions, it is typically hard to train a model and the L2 running cost (the soft constraint between the finetuned and pretrained models) can be seen as a component to stabilize training. This is similar to diffusion models and flow matching models where a noising process or a reference velocity field is used. We will polish the sentence and explain the intuition in a clearer way.\n\n> 4.1: There are still many points that remain unclear. I think the setup could be more rigorously exposed (differentiable reward function, etc).\n\nThanks for pointing this out and we will explicitly explain the setup in the last two paragraphs in the intro section and the beginning of the method section. \n\nOn our setting:\n\nGiven a differentiable reward function, we want to train our generative model to achieve high reward scores for the generated samples and also to preserve the prior distribution of the pretrained model. To be more specific, we formulate this to Eq. 9 via the optimal control framework and then focus on solving the control problem. We refer to our Sec 4.1 for a description of the setting, which we will make sure to refine in our final version.\n\n> ...how the reparameterization through the HJB equation makes things computationally easier (if that is the goal)..\n\nGood question! Here are our intuitions:\n\n- By using the HJB formulation, we know that the optimal velocity is simply the value gradient (up to a sign). The value function itself can be well approximated by single-Euler-step guess. The additional U-Net is used to correct the errors in this approximation so that the model, after long periods of training, is not too biased.\n- Optimizing the value function in an amortized way (with some approximations like finite differences) gives us benefits in computational efficiency: instead of solving for the adjoints, we may subsample transitions for training, especially in the case when the number of inference steps is large. This is similar to the motivation of using temporal difference (TD) methods in RL (for instance, soft actor-critic [1]) instead of tracking a very long trajectory.\n\n> On the experiments In Table 1, it seems that concurrent models achieve lower FID which is desirable and yet are not put in bold characters.\n\nSorry for this mistake. We indeed have revised our experiments by setting the subsampling rate in our method to 50% (previously 25%). Please fine the table below:\n\n**Aesthetic Score**\n\n| Method     | Reward | CLIPScore | DreamSim | FID |\n|------------|--------------|------------|--------------|------------|\n| Ours       | 8.02     | 10.92       | 20.95 | 341 |\n| ReFL       | 10.00     | 5.30       | 5.59  | 1338 |\n| DRaFT       | 9.54     | 7.47       | 7.77 | 1246 |\n| Adjoint Matching       | 6.87    | 15.82       | 22.34 | 465 |\n\n**HPSv2**\n\n| Method     | Reward | CLIPScore | DreamSim | FID |\n|------------|--------------|------------|--------------|------------|\n| Ours       | 3.86     | 15.39       | 18.40 | 1161 |\n| ReFL       |   3.87   | 12.13       | 14.08  | 1195 |\n| DRaFT       | 3.76     | 12.79       | 15.05 | 1177 |\n| Adjoint Matching       | 3.59    | 11.76       | 14.11 | 1247 |\n\n**PickScore**\n\n| Method     | Reward | CLIPScore | DreamSim | FID |\n|------------|--------------|------------|--------------|------------|\n| Ours       | 23.01     | 18.87       | 23.36 | 967 |\n| ReFL       |   23.19   | 15.06       | 17.71  | 997 |\n| DRaFT       | 23.00     | 15.49       | 19.03 | 968 |\n| Adjoint Matching       | 22.78    | 16.12       | 19.69 | 1033 |\n\n> While the visual results look quite convincing, there is always this trading off between the two terms in the loss (the one maximizing the reward and the one forcing to stay close to the base model): I would like to see of this parameter impacts the metrics, and same for concurrent methods that must have some kind of similar tradeoff.\n\nWe apologize that we accidentally use $\\beta = 1 / \\lambda$ to indicate the trade-off between reward maximization and prior preservation, for which the results are shown in Figure 6. For $\\eta_t$, please find the results on our new ablation study on Aesthetic Score below:\n\n| Method     | Reward | CLIPScore | DreamSim | FID |\n|------------|--------------|------------|--------------|------------|\n| $\\eta_t = (1-t)^2$   |  8.02     | 10.92       | 20.95 | 341 |\n| $\\eta_t = (1-t)$       | 8.13     | 10.06       | 17.75  | 484 |\n\n> ...whether you achieve faster training time (you gain in terms of iterations but what about computation time ?). Among the benchmarks models you pick, is there any methods where images are generated (since, as i understand you claim to circumvent this aspect).\n\nPlease find the runtime table below (running for 400 iterations; 2 GPUs, batch size 64 with gradient accumulation). We would like to note that Adjoint Matching is implemented with a replay buffer while ours is pure on-policy. We may adopt the same trick to accelerate training but we leave that to future work.\n\n| Method     | Time (hours) |\n|------------|--------------|\n| Ours (subsampling 50%)      | 34       | \n| Ours (subsampling 25%)       | 22       | \n| ReFL       |   10       | \n| DRaFT       | 10       | \n| Adjoint Matching          | 14    |\n\n---\n[1] Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control. Domingo-Enrich et al. ICLR 2025. \n\n[2] Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor. Haarnoja et al. ICML 2018."}, {"review_id": "iVzmq89ptn", "reviewer": "Reviewer_Voux", "summary": "This work proposes VGG-Flow, a method for efficiently aligning pretrained flow matching models with reward functions by leveraging optimal control theory. The authors formulate a relaxed control objective whose optimality conditions are given by the Hamilton–Jacobi–Bellman (HJB) equation, showing that the difference between the fine-tuned and base velocity fields should match the gradient of a value function. In practice, they parametrize this value-gradient via a forward-looking approximation, i.e., using the reward gradient of a one-step Euler prediction plus a small learnable residual, and jointly train the velocity field and value-gradient models over simulated trajectories.. Empirically, VGG-Flow is applied to Stable Diffusion 3 under various reward models (e.g., aesthetic score, HPSv2, PickScore) and demonstrates faster increase of rewards while better preserving sample diversity and base-model priors compared to baselines such as ReFL, DRaFT, and Guidance Matching.", "questions": "1. I am curious about the choice of the value function network, i.e., using a module initialized from SD v1.5 UNet. Is there any reason why the authors choose this architecture? \n2. I think the paper’s most relevant baseline should be adjoint matching, which uses memoryless SDE to derive the value term, and update the velocity. I saw the comparison in supplementary files, but the results are not in the main paper. Is there any reason why adjoint matching is not on the main paper?\n3. Furthermore, regarding the results in supplementary, adjoint matching shows better diversity and FID, while other methods such as ReFL or GM shows the highest reward. I can understand that ReFL, GM or DRaFT achieve the highest reward as they are prone to reward hacking. But, when compared to adjoint matching, the better diversity or FID is partly due to the SDE formulation. What do you think about it? If this is true, can the proposed method also be applied with SDE?\n4. If one considers multi-reward optimization, how would this be the case? Do we need to compute all the Jacobians for each reward? Also, if we are using VLMs for reward modeling, can the proposed method scale to such a case?", "limitations": "They have mentioned the limitations of the work, e.g., the limitations on using first-order gradient estimators, lack of exploration on architectural designs.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "### Strength\n The strength of this paper is its theoretically grounded and practical algorithm that introduces the HJB framework, where the method provides a clear probabilistic interpretation and memory-efficient gradient-matching objective. The forward-looking parametrization of the value gradient not only sidesteps costly adjoint ODE solvers but also accelerates adaptation, yielding substantial reward improvements within a tight computational budget. Moreover, the extensive ablations on reward temperature, boundary-loss weighting, and subsampling strategies show the method’s robustness across different reward signals, while quantitative results (e.g., Pareto fronts of reward vs. diversity and FID) convincingly illustrate that VGG-Flow achieves better trade-offs than direct reward-maximization and gradient-matching baselines .\n\n### Weakness\n Despite its elegance, the paper’s theoretical contributions lack rigorous convergence guarantees or bounds on the approximation error introduced by finite-difference estimation of the value gradient, thus relying on heuristic weighting and early-stop selection to avoid reward hacking. Furthermore, technically the method requires approximation of (jacobian of) value function and update the residual with learned value function, which also could induce bias throughout the learning. The empirical evaluation is narrowly focused on a single T2I model and human-preference rewards, while neither multi-objective scenarios (e.g., balancing aesthetics and alignment) nor perceptual user studies are explored, leaving open whether the numeric gains translate to genuinely better image quality.", "was_revised": true, "final_justification": "The authors' clarified most of my concerns, I raise my score to 5 accept.", "rebuttal": "We thank the reviewer for their acknowledgement on our contributions and their time and efforts invested in the reviewing process.\n\n> Despite its elegance, the paper’s theoretical contributions lack rigorous convergence guarantees or bounds on the approximation error introduced by finite-difference estimation of the value gradient, thus relying on heuristic weighting and early-stop selection to avoid reward hacking. \n> Furthermore, technically the method requires approximation of (jacobian of) value function and update the residual with learned value function, which also could induce bias throughout the learning. \n\nThank you for pointing out the approximation issue. Yes, this is indeed a coarse approximation, though our empirical results show that the finetuning process does not break due to this approximation. In the meantime, we would like to argue that if we aim for fast finetuning, then we have to accept some level of bias unless the cost of carefully finetuning the value function with small learning rates is affordable.\n\nOn theoretical guarantees: thank you for your suggestion. We believe that the convergence bound can be shown with classical results on HJB or some recent papers (for instance, [1] on approximation bounds for distributional HJB). Indeed, from the control perspective out setting is simply that we use a different set of costs and dynamics. We will include more analysis and details on this aspect in the next version of our paper.\n\n> The empirical evaluation is narrowly focused on a single T2I model and human-preference rewards, while neither multi-objective scenarios (e.g., balancing aesthetics and alignment) nor perceptual user studies are explored, leaving open whether the numeric gains translate to genuinely better image quality.\n\nWe acknowledge that our paper does not explore multi-reward scenarios nor present user studies. We made this decision because 1) our paper focuses on the finetuning method itself rather than the properties of specific reward functions and 2) we follow the same evaluation setting in many other related papers in this domain, such as the baselines (ReFL, DRaFT and Adjoint Matching).\n\n> I am curious about the choice of the value function network, i.e., using a module initialized from SD v1.5 UNet. Is there any reason why the authors choose this architecture?\n\nThe choice of this value network is due to the convenience in implementation: it takes few lines to copy the model. Other architectures work for our purpose as well.\n\n> Furthermore, regarding the results in supplementary, adjoint matching shows better diversity and FID, while other methods such as ReFL or GM shows the highest reward. I can understand that ReFL, GM or DRaFT achieve the highest reward as they are prone to reward hacking. But, when compared to adjoint matching, the better diversity or FID is partly due to the SDE formulation. What do you think about it? If this is true, can the proposed method also be applied with SDE?\n\nWe argue that this is not due to the SDE formulation. Our method directly applies to the ODE since we always assume $dx = v(x,t)dt$. Adjoint matching, in contrast, turns the ODE into an SDE with the same probability flow because it relies on the assumption of “memoryless” process.\n\nPartially to show why the SDE formulation is not a must, we also show the results of “adjoint matching method without SDE” (see Section A.1 in the supplementary, derived with PMP) in the table below. Even without the equivalent SDE, adjoint matching (or more precisely, optimal control with direct simulation of the adjoint equation) leads to some performance.\n\n(Aesthetic Score)\n\n| Method     | Reward | CLIPScore | DreamSim | FID |\n|------------|--------------|------------|--------------|------------|\n| Adjoint Matching       | 6.87    | 15.82       | 22.34 | 465 |\n| \"Deterministic Adjoint Matching\"       | 7.53    | 8.26       | 11.17 | 1171 |\n\n\n> If one considers multi-reward optimization, how would this be the case? Do we need to compute all the Jacobians for each reward? Also, if we are using VLMs for reward modeling, can the proposed method scale to such a case?\n\nIf we use a linear combination of multi-rewards (e.g., weighted sum), then we do need to compute all the Jacobians as the differentiation of sum equals to sum of differentiations. \n\nRegarding VLM models, it depends on how we utilize the VLM to get a score value. For example, one approach [2] is to ask the the model to give a binary answer about “whether the generated image is aligned with the given text prompt; only answer ‘yes’ or ‘no’”. We may extract the likelihood of the \"yes\" token. Still, it remains an issue to compute the reward gradient with respect to the input $x_t$. It can possibly be solved by stochastic approximations (e.g., with randomly sampled directions, similar to what Hutchinson's trick is used for estimating divergence).\n\n---\n\n[1] Tractable Representations for Convergent Approximation of Distributional HJB Equations. Alhosh et al. https://arxiv.org/abs/2503.05563\n\n[2] Beyond Yes and No: Improving Zero-Shot LLM Rankers via Scoring Fine-Grained Relevance Labels. Zhuang et al. NAACL 2024. https://arxiv.org/abs/2310.14122"}, {"review_id": "b9CYx4uCqw", "reviewer": "Reviewer_jDnT", "summary": "The paper proposed a new approach, VGG-Flow, for flow matching alignment. The proposed method was deduced from the theory of optimal control, and effectively learn a vector field network and a value network that approximates the gradient of the running cost simultaneously. On text-to-image generation tasks, the proposed approach exhibited better performance in terms of multiple objective rewards.", "questions": "Please refer to the weaknesses listed above. In addition:\n- For text-to-image alignment, it is unclear why the CLIP score is missing in the reward. In all previous work, like Adjoint Matching and OCFlow, the CLIP score was regarded as the primary metric due to its robustness to adversarial attacks.\n- The network $g$ looks exactly the same as the adjoint state in Adjoint Matching, or the costate in OCFlow. Can you further elaborate on their connections?", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strength\n- The paper is inspired by the theory of optimal control, which has a solid mathematical background.\n- The final algorithm has a clear correspondence to the running cost in the optimal control objective, which offers a fine-tuning method with gradient initialization from the reward function.\n\n## Weakness\n- ** Proposition 1, to the best of my knowledge, is incorrect for ODE-based flow matching**. The authors referred to an external source for this proposition. After checking the original paper, I only found the proof for the SDE setup (i.e., diffusion/score matching), whose bound contains the variance of the diffusion (noise) term. Therefore, the adaptation to ODE-based flow matching is non-trivial.\n- **The motivations behind the proposed approach seem to focus on the wrong aspect**. For example, the authors claimed that solving the optimal control objective is \"intractable\". This is clearly an overstatement, as existing approaches like SOC [1] and Adjoint Matching [2] for stochastic flow (SDE), and OCFlow [3] for deterministic flow (ODE) have all proposed practically feasible algorithms for alignment. Instead, the major difference of the proposed approach is that it does not require backpropagating the reward gradient through the ODE/SDE solver, as it learns a model to approximate it. Efficiency (and therefore the practical time benchmark) shall be one of the major contributions.\n- Although the authors have noted the close relation to the OC-based method in Section 6, they were never compared as baselines.\n- The proposed approach involved more time-consuming operations like JVP and additional training parameters (of the gradient value network). However, a benchmark on the running cost was not provided to give a direct comparison with the existing baselines.\n- The results in Table 1 were **unsatisfactory** in demonstrating the superior performance, with **clear errors in highlighting the best model**. For example, HPSv2 reward and FID metrics are erroneously highlighted. The PickScore FID is also incorrect.\n\n[1] Domingo i Enrich, Carles, et al. \"Stochastic optimal control matching.\" Advances in Neural Information Processing Systems 37 (2024): 112459-112504.\n\n[2] Domingo-Enrich, Carles, et al. \"Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control.\" arXiv preprint arXiv:2409.08861 (2024).\n\n[3] Wang, Luran, et al. \"Training Free Guided Flow Matching with Optimal Control.\" arXiv preprint arXiv:2410.18070 (2024).", "was_revised": true, "final_justification": "Please refer to my discussions with the authors.", "rebuttal": "We appreciate the reviewer for their time and efforts to review our paper and to raise constructive questions.\n\n> ** Proposition 1, to the best of my knowledge, is incorrect for ODE-based flow matching**. The authors referred to an external source for this proposition. After checking the original paper, I only found the proof for the SDE setup (i.e., diffusion/score matching), whose bound contains the variance of the diffusion (noise) term. Therefore, the adaptation to ODE-based flow matching is non-trivial.\n\nThank you for pointing this important issue out and we apologize for not carefully checking the cited paper. You are right, this claim is incorrect: without any additional assumptions, we can only claim that the Wasserstein-2 distance is bounded, not the KL divergence. While we lose track of the exact probabilistic formulation of the objective function, finetuning to minimize W2 distance is still a valid setting.\n\nIn practice, we may assume that the velocity field $v(x,t)$ is Lipchitz continuous for $t \\in [0, 1]$ or $t \\in [\\epsilon, 1 - \\epsilon]$ ($\\epsilon > 0$). With this additional mild assumption, we may show that the KL divergence in Proposition 1 is bounded.\n\nWe will revise our statements and elaborate the technical details in the next version of our draft.\n\n> The motivations behind the proposed approach seem to focus on the wrong aspect. For example, the authors claimed that solving the optimal control objective is \"intractable\". This is clearly an overstatement, as existing approaches like SOC [1] and Adjoint Matching [2] for stochastic flow (SDE), and OCFlow [3] for deterministic flow (ODE) have all proposed practically feasible algorithms for alignment. Instead, the major difference of the proposed approach is that it does not require backpropagating the reward gradient through the ODE/SDE solver, as it learns a model to approximate it. Efficiency (and therefore the practical time benchmark) shall be one of the major contributions.\n\nSorry to possible misunderstanding here: in our paper, we do not claim that OC objectives are intractable (indeed, we have presented empirical results on Adjoint Matching.), but instead argue that our amortized objective is computationally more tractable. We will revise our claims and clearly point out that the major motivation of using HJB instead of ODE/SDE solver is efficiency.\n\nNevertheless, we would like to point out that we observe that to use Adjoint Matching we need to carefully treat the outliers (for instance, throwing out samples when adjoints get too large) and empirically it does not work that well compared to our method.\n\nPlus, we appreciate the reviewer for mentioning OCFlow and we will cite and explain the connections to OCFlow in our revised draft.\n\n> Although the authors have noted the close relation to the OC-based method in Section 6, they were never compared as baselines.\n\n> The results in Table 1 were unsatisfactory in demonstrating the superior performance, with clear errors in highlighting the best model. For example, HPSv2 reward and FID metrics are erroneously highlighted. The PickScore FID is also incorrect.\n\n> For text-to-image alignment, it is unclear why the CLIP score is missing in the reward. In all previous work, like Adjoint Matching and OCFlow, the CLIP score was regarded as the primary metric due to its robustness to adversarial attacks.\n\nWe are sorry that we did not have the chance to present the empirical results on OC-based methods. Our original plan was to show Adjoint Matching (since it is the more relevant and recent one) but it was not open-sourced at that point plus the experiments in the AM paper are done in some different setting. \n\nWe apologize for the mistakes in the tables in the main text and thank you for suggesting CLIPScore. Here we present the most recent results after some refinement on the hyperparameters:\n\n**Aesthetic Score**\n\n| Method     | Reward | CLIPScore | DreamSim | FID |\n|------------|--------------|------------|--------------|------------|\n| Ours       | 8.02     | 10.92       | 20.95 | 341 |\n| ReFL       | 10.00     | 5.30       | 5.59  | 1338 |\n| DRaFT       | 9.54     | 7.47       | 7.77 | 1246 |\n| Adjoint Matching       | 6.87    | 15.82       | 22.34 | 465 |\n\n**HPSv2**\n\n| Method     | Reward | CLIPScore | DreamSim | FID |\n|------------|--------------|------------|--------------|------------|\n| Ours       | 3.86     | 15.39       | 18.40 | 1161 |\n| ReFL       |   3.87   | 12.13       | 14.08  | 1195 |\n| DRaFT       | 3.76     | 12.79       | 15.05 | 1177 |\n| Adjoint Matching       | 3.59    | 11.76       | 14.11 | 1247 |\n\n**PickScore**\n\n| Method     | Reward | CLIPScore | DreamSim | FID |\n|------------|--------------|------------|--------------|------------|\n| Ours       | 23.01     | 18.87       | 23.36 | 967 |\n| ReFL       |   23.19   | 15.06       | 17.71  | 997 |\n| DRaFT       | 23.00     | 15.49       | 19.03 | 968 |\n| Adjoint Matching       | 22.78    | 16.12       | 19.69 | 1033 |\n\n> The proposed approach involved more time-consuming operations like JVP and additional training parameters (of the gradient value network). However, a benchmark on the running cost was not provided to give a direct comparison with the existing baselines.\n\nThank you for raising this question. Please find the runtime table below (running for 400 iterations; 2 GPUs, batch size 64 with gradient accumulation). We would like to note that Adjoint Matching is implemented with a replay buffer while ours is pure on-policy. We may adopt the same trick to accelerate training but we leave that to future work.\n\nRuntime\n| Method     | Time (hours) |\n|------------|--------------|\n| Ours (subsampling 50%)      | 34       | \n| Ours (subsampling 25%)       | 22       | \n| ReFL       |   10       | \n| DRaFT       | 10       | \n| Adjoint Matching          | 14    | \n\n> The network g looks exactly the same as the adjoint state in Adjoint Matching, or the costate in OCFlow. Can you further elaborate on their connections?\n\nThey are indeed very closely connected. For smooth optimal control problems, the adjoint states are the same as $g = \\nabla V$ **on the optimal trajectory**. Here we would like to underscore that HJB is that it is a globally true condition: even without knowing the optimal trajectories it must hold everywhere. Therefore, one benefit of this HJB formulation is that it gives room for off-policy training."}]

Author rebuttal:

Meta-review: Reviewers agree that this is a significant step forward in the field. The rebuttal and post-rebuttal discussion clarified initial questions. Congratulations, this submission is accepted.
Decision: Accept (poster)