Value Matching: Scalable and Gradient-Free Reward-Guided Flow Adaptation
Authors: Cristian Perez Jensen
Luca Schaufelberger
Riccardo De Santi
Kjell Jorner
Andreas Krause
Venue: ICLR 2026 Poster (2026)
Keywords: diffusion models
flow models
black-box reward optimization
molecular design
image generation
stochastic optimal control
reinforcement learning
Primary area: generative models

Abstract: Adapting large-scale flow and diffusion models to downstream tasks through reward optimization is essential for their adoption in real-world applications, including scientific discovery and image generation. While recent fine-tuning methods based on reinforcement learning and stochastic optimal control achieve compelling performance, they face severe scalability challenges due to high memory demands that scale with model complexity. In contrast, methods that disentangle reward adaptation from base model complexity, such as Classifier Guidance (CG), offer flexible control over computational resource requirements. However, CG suffers from limited reward expressivity and a train-test distribution mismatch due to its offline nature. To overcome the limitations of fine-tuning methods and CG, we propose Value Matching (VM), an online algorithm for learning the value function within an optimal control setting. VM provides tunable memory and compute demands through flexible value network complexity, supports optimization of non-differentiable rewards, and operates on-policy, which enables going beyond the data distribution to discover high-reward regions. Experimentally, we evaluate VM across image generation and molecular design tasks. We demonstrate improved stability and sample efficiency over CG and achieve comparable performance to fine-tuning approaches while requiring less than 5% of their memory usage.

Published as a conference paper at ICLR 2026
VALUE MATCHING: SCALABLE AND GRADIENT-FREE REWARD-GUIDED FLOW ADAPTATION
Cristian Perez Jensen
Luca Schaufelbergerf1.a
Riccardo De Santi
Kjell Jorner
Andreas Krause!
'ETH Zürich
ETT AI Center
2NCCR Catalysis
ABSTRACT
Adapting large-scale flow and diffusion models to downstream tasks through reward optimization is essential for their adoption in real-world applications, including scientific discovery and image generation. While recent fine-tuning methods based on reinforcement learning and stochastic optimal control achieve compelling performance, they face severe scalability challenges due to high memory demands that scale with model complexity. contrast, methods that disentangle reward adaptation from base model complexity, such as Classifier Guidance (CG), offer flexible control over computational resource requirements. However, CG suffers from limited reward expressivity and train-test distribution mismatch due to its offline nature. To overcome the limitations of fine-tuning methods and CG, we propose Value Matching (VM), an online algorithm for learning the value function within an optimal control setting. VM provides tunable memory and compute demands through flexible value network complexity, supports optimization of non- differentiable rewards, and operates on-policy, which enables going beyond the data distribution discover high-reward regions. Experimentally, across image generation and molecular design tasks. We demonstrate improved stability and sample efficiency over and achieve comparable performance fine-tuning approaches while requiring less than 5% of their memory usage.
evaluate VM
INTRODUCTION
Large-scale generative foundation models have re- cently made remarkable progress. Among them, flow matching (Lipman et al., 2023; Albergo & Vanden- -ijnden,,, 2023; Liu et al., 2023a) and diffu- sion models (Ho et al., 2020; Sohl-Dickstein et al., 2015; Song et al., 2020) stand out for their ability to generate high-fidelity samples across wide range of domains, including images (Ho et al., 2020), chemistry (Hoogeboom et al., 2022), biology (Corso et al., 2023), and robotics (Chi et al., 2023). In many applications (e.g.. controllable image editing and drug discovery (Olivecrona et al., 2017)), it is essential to adapt these pre- pre-trained models to downstream rewards. However, existing reinforcement learning (RL) (Zhao et al., 2025; Black et al., 2023; Fan et al., 2023; Hu et al., 2025) and stochastic optimal control (SOC) (Domingo-Enrice Enrich et al., 2025; Uehara et al., 2024; Tang, 2024; Domingo-Enrian et al., 2024) approaches update model weights via full-model backpropagation, making them increasingly memory-intensin as model sizes scale to billions of parameters (Figure 1).
Table 1: Algorithm capabilities.
VM
CG
AM
CT-PPO
Memory-efficient Black-box rewards
x
x
X
Online data
Despite their promising performance (e.g.. Domingo-Enrich et al., 2025; Zhao et al., 2025), RL and SOC-based fine-tuning methods remain fundamentally limited by memory requirements that scale with model size. Moreover, many statc-of-the -fthe-of-theart approaches require differentiable rewards (e.g., Domingo-Entn -nnichh et al., 2025), restricting applicability to black-box settings where only function evaluations are available. This is particularly limiting in applications like drug discovery, where objectives often rely on costly simulators (Bannwarth al., 2019; Sholl & Steckel, 2009; Forli et al., 2016) or experimental measurements (Hughes et al., 201 1). This raises key question:
Coorespondence to Cristian Perez Jensen crist jensinnneeeeeeeeeeeee com. I Equal contribution.

---

Published as a conference paper at ICLR 2026
How can leverage SOC for RL) machinery flow and diffusion models to black-box rewards with efficient resource usage and competitive performance?
Limitations of fine-tuning methods suggest reconsidering approaches like Classifier Guid- ance (CG) (Dhariwal & Nichol, 2021) that offer computational advantages by avoiding updates to base model parameters and enables adaptation with black-box rewards. However, CG an offline algorithm that trains only on samples from the pre-trained policy, leading to limited exploration beyond the training distribution. We address this shortcoming by introducing Value Matching (VM), an online algorithm that learns the value function, enabling discovery of high-reward regions beyond the data distribution. Further, VM keeps the base generative model fixed, resulting in memory-efficient flow adaptation, and supports black-box reward functions. Lastly, show that VM achieves comparable performance with fine-tuning methods at a significantly lower cost.
SDXL (2024)
SD3-M (2024)
VRAM (GB)
750
SSD-IB (2024)
500
SD2.1 (2022)
250
DIT (2023)
SDI.5 (2022)
0
0
2
Parameters (B)
Figure I: The recent trend toward more complex flow and diffusion generative models leads prohibitively high fine-tuning memory (VRAM) requirements.
Our contributions:
A control-theoretic viewpoint shedding light on how the offline nature of classifier guidance hinders the discovery of high-reward samples beyond the data distribution (Section 4).
online method for reward adaptation of flow models through value function learning, highlighting advantages compared to AM and CT-PPO (Section 5).
Value Matching (VM), a
An experimental evaluation across image and molecular generation tasks using non-differentiable rewards, showing that VM (i) achieves superior stability and sample efficiency compared to classifier guidance, (ii) reduces resource requirements by 95% relative to fine-tuning methods while main- taining comparable performance, and (iii) enables effective adaptation of molecular and large-scale eex-tt--mmae models through value networks significantly smaller than the base model (Section 6).
2 BACKGROUND AND NOTATION
Flow models. In this work, we consider the problem of adapting flow (Lipman et al., 2023; Albergo & Vanden-Fijnden 2023; Liu al., 2023a) and diffusion models (Ho al., 2020; Sick-Dickstein et al., 2015; Song et al., 2020), which have emerged as leading approaches for generative modeling across various domains, including images (Rombach et al., 2022; Peebles & Xie, 2023; Podell et al., 2024; Gupta et al., 2024; Esser et al., 2024), text (Gat et al., 2024), and molecular design (Hoogeboom al., 2022; Dunn & Koes, 2024a;b). Typically, flow models are sampled through an ordinary differential equation (ODE) (Lipman et al., 2023), however, they can also be sampled via stochastic differential equation (SDE) with equal time marginals (Maoutsa et al., 2020). We sample by simulating SDE:
(1) where b: Rd J0, 1] Rd is the drift, [0, 1] + Radd is the diffusion, and B, is Brownian tion. For unbiased adaptation, we use memoryless parameterization (Domingo-Eme) Enrich et al., 2025):
dx, = b(x,, t) dt + (t) dB,, Xo Po-
ài x + logp:.)) o² 23r ( Br -Be) -
6(x, t) A
aa
ar
with (a,, B:) defined the flow schedule (Lipman et al., 2023) and (e)) being time derivatives.
Stochastic optimal control. For our approach, we frame KL-regularized adaptation as an SOC problem (Bellman & Dreyfus, 2015; Fleming & Rishel, 2012), general framework that deals with optimization over stochastic processes. Specifically, we restrict ourselves to quadratic cost control-affine Bolza problem (Bolza, 1904) finite time horizon [0, 1] with dynamics:
dx; (b(x,, t) (t))u(xx t)) dt + (tt)
Xo Pu-
(3)
Further, the cost functional J is defined as the total cost of a control u starting from a point (x, t). composed of quadratic running cost 3 [u(x,, 2) and arbitrary terminal cost g : Rd - R. The
2

---

Published as a conference paper at ICLR 2026
control problem is to find the control that minimizes the cost functional J at every point (x, t):
u" E arg min J (u; t) Eyn [ [. |lu(xe, s)l? ds g(x1) x =
(4)
wEll
From here, the value function is defined as the optimal value the cost functional:
V(x,)) 1) - inf J(u; x, 1) = Ju""; xx
(5)
well
Further, V can be defined through the base distribution DDomingo-Enrich al., 2024, Appendix B): V(x, 1) = - log p (exp(-g((x) X = x).
(6)
PROBLEM SETTING
We consider the generative optimization problem (Li et al., 2024b; De Santi et al., 2025a:b) of
adapting pre-trained flow or diffusion model p to maximize reward function R - R
in expectation, weighted by À E Rgo, and remain close to pere in terms of Kullback-Leiber (KL) divergence. Formally, we optimize over policies ı with induced last-timestep marginal pf:
arg max E [Ar(x1)] - (pL p! s.t. dx, b" (x,, t) dt o(t dB,.
(7)
The optimal solution x* of Problem (7) induces the tilted distribution pi (x) x pp (x) exp(Ar(x)) if Enrich et al., 2025). In flow models, this objective is equivalent to the following quadratic cost control-affine control problem (Tang, 2024):
a follows a memoryless noise schedule (Domingo-Fie Eomiingg--nrcch
2 J: [|u(x,, s)1 ds >r(x1)) x =
arg min w: [0011--R
E
(8)
s.t. dx: = (xp_ () + o (1)((x),)))))))))))(())))) ()) dl + o(1) dB,-
Existing fine-tuning methods (e.g., Black et al., 2023; Fan al., 2023; Nomingo-Enrich et al., 2025; Zhao et al., 2025) encounter scalability challenges because they require updating the base model parameters. To perform these updates, the gradients of the loss function must be backpropagated through the entire model. This process is highly memory-intensivee as calculating the gradients requires storing all intermediate activations of the model. As result, the memory footprint scales with model size, posing major bottleneck as these models grow to billions of parameters. Importantly, we focus on the black-box setting, where the reward is accessible only through function evaluations. This generalization makes the problem considerably more challenging than the setting of previous works that leverage gradient information (e.g., Xu et al., 2023; Clark al., 2024; Domingo-Entn Enrich et al., 2025; Wang et al., 2025). Next, we introduce approach that addresses these issues.
DISENTANGLING OPTIMIZATION THROUGH VALUE FUNCTION LEARNING
In this work, we advocate for learning the value function V of Equation (5) to solve Problem (7), which offers compelling advantages over current fine-tuning methods: (1) support of non-dificrentiable rewards (ef. Xu et al., 2023; Clark et al., 2024; Domingo-Enen Ennichh et al., 2025) and (2) controllable resource usage. Once learned, the value function allows us to find the optimal control u² through Pontryagin's minimum principle (Pontryagin. 1962), where the first-order optimality condition gives:
u' u'(x,t) = -oT V(, t).
(9)
A key insight for Problem (7) is that the value function retains differentiability even when the reward function is not, ensuring well-defined optimal control. Intuitively, this occurs because stochastic noise acts smoothing kernel; the value function V time averages over all noise realizations from x, to X1, effectively regularizing reward discontinuities (see Figure 3). We formalize this as follows. Propesition 1. Under the memoryless noise schedule and assuming that is bounded and measurable, the value function V is differentiable in x at < 1.
Proof outline. We write Vxx,t) = log (xx t) where ((xx t) = Epare [expp(xxxx)x x; = x]. Then apply the chain rule and show that vi > 0 and that vi is differentiable (see Appendix A.3).
This proposition enables the optimization of non-differentiabl reward functions commonly used in practical applications. It proves particularly valuable in settings where reward func- tions are treated as black boxes, providing only function evaluations without access gradi- ents or structural information. For instance, in molecular generation, reward functions often rely on external simulators (Bannwarth et al., 2019; Sholl & Steckel, 2009; Forli et al., 2016)
3

---

Published as a conference paper at ICLR 2026
High reward
(a) Base model pr
(b) Reward r.
(c) Optimal pi.
(d) CG training data. (c) VM training data.
Figure 2: Training distribution evolution for VM and CG a 2D environment, showing that VM
aligns its training distribution with the optimal distribution throughout training, whereas CG does not. (2d, 2e) Points with lower opacity represent earlier training stages.
or experimental measurements (Hughes et al., 2011) that only return scalar values. In such cases, value function learning makes reward adaptation possible, as opposed to methods that rely reward gradients (e.g., Xu et al., 2023; Clark et al., 2024; Domingo-Enrie -nnrihh et al., 2025).
Beyond handling non-differentiable rewards, learning the value function disentangles reward adaptation from base model op- timization. Consequently, the dominating computational cost shifts from base model training to base model inference and value function learning. This leads to controllable resource usage by allowing for the flexible choice of a value network architecture. As we show in Section 6, the value function can be made significantly smaller than the base model, resulting in a substantial reduction in cost, while achieving comparable performance to fine-tuning methods. Because of this significant cost reduction, the value function learning approach is highly practical for the adaptation of lcage-scale models where fine-tuning would be prohibitively expensive.
0.66 0.74 t 0.81
0.88 = 0.94 t = 1.00
Figure 3: Evolution of value func- tion for discontinuous reward.
Having established the advantages of value function learning, we next recall how classifier guidance can be viewed as an offline algorithm to learn the value function V.
Classifier guidance as offline value function learning. As previously established by Pandey al. (2025), Classifier guidance (CG) (Dhariwal & Nichol, 2021) admits an interpretation as value function learning. The approach involves training a classifier Pv]+(y | x) over noisy samples at timesteps /. E [o. 1] and using its log-gradient V log Pr]( x) to guide generation. This procedure can be understood leveraging Equation (6) solve Problem (7) with the reward function r(x) = log Py (y x) and value function (xx = - log PY x) (see Appendix A.5 for the derivation), which in turn motivates a generalized loss function for arbitrary rewards:
1 2
f/
Lea ((o,1)
(xxp(-Vo(x,,11 (xx(())xx))
(10)
dx; = jrre (xt, 1) di (1) dB:,
X ~ Po-
Throughout the remainder of this work, we use CG to refer to this loss formulation. Next, we discuss issues with this approach for generative optimization (Li et al., 2024b; Santi et al., 2025a;b).
Limitations of classifier guidance for generative optimization. The CG formulation in Equa- tion (10) represents an offline value function learning approach because it trains on samples from the fixed, pre-trained distribution peee rather than the current policy distribution p`" This offline nature inherent to the loss formulation itself: the expectation in Equation (6) is computed over the pre-trained distribution, necessitating that the value function is trained on samples from this fixed source.
However, for generative optimization tasks such as reward maximization, we wish to sample high reward designs beyond regions of high data availability (De Santi et al., 2025a;b; 2026; Klarner et al., 2024). This objective reveals fundamental limitation of the offline approach: distribution mismatch arises between the fixed training distribution pt and the target distribution P that the learned policy encounters during inference (see Figure 2). As the policy shifts probability mass toward higher-reward regions, the training samples from pt become increasingly less informative. reducing sample efficiency and potentially limiting the method's ability to discover optimal solutions.
4

---

Published as a conference paper at ICLR 2026
Algorithm 1 Value Matching algorithm.
Require: Pre-trained base model with sampling SDE dx; = b(x,, dt (t) dBt, Reward function
r Rd -> IR, Untrained value function approximator V08 Rd x [0, 1] + R, Number of iterations
N E N, Trajectories per iteration E N, Timestep-dependent weighting [0, 1] I2
1: for N iterations do
Sample trajectories under the current policy:
dx (b( o² (t) dt + (tt) dB,, x,,0 pu: i [m].
3: Estimate the cost functional for each timestep each trajectory with 0 stopgrad
Fr 1a=22 2/2 [. (lxs))VV((x s)ll ds - Ar(x,,1,, tE 10, 1], i [m].
2
4: Compute the loss function:
( 1
/ w(t) (Vo(x.,) 2 dt.
2m
5: Make optimization step with (L))
6: end for
Additionally, the exponential terms in Equation (10) create numerical stability issues during training, due to overflows when λr(x) > 90 under 32-bit floating-point precision. This constrains the method's expressivity small reward scalings A, show in Section 6. To address these limitations, leverage the control-theoretic theoretic viewpoint to next introduce an online value function learning approach that aligns the training and inference distributions while maintaining numerical stability.
5 VALUE MATCHING: SCALABLE AND GRADIENT-FREE DEWARD-GUIDER ADAPTATION
We introduce Value Matching (VM), an online method for learning the value function that overcomes a fundamental limitation of CG by training on trajectories from the current policy. As detailed in Algorithm 1, VM leverages Equations (5) and (9) to estimate the value function. This is achieved by it- eratively regressing the approximator V onto the cost functional J, computed using the current policy,
A 1 { w(t) (Vo(x,, - -)
(0;
2 /
Vx Voi X 1))
1
(11)
(1;)) t)
|lu(x 3) ds -x((x)), = stopgrad(
2 :
dx; = (bpre (x,, t) (t t)) dt + (t) dB,, X ~ Po
The VM algorithm follows a simple iterative procedure. Each iteration begins by sampling trajectories using the control policy u(x, t) -o V t), derived from the current value function approx- imator Vg. These trajectories serve a dual purpose in the training loop. First, they are used to compute single-sampi Monte Carlo estimates J₁ of the cost functional J, which serves as the regression target:
1 S. o (3s)) s)l Ar(x1)- 2 [.
j.
(12)
Here, 0 = stopgrad(O) prevents gradients from flowing through the target, ensuring that Je is treated as a fixed target value during backpropagation. Second, the states x, along these same trajectories provide the input data for the value function. The network parameters 0 are updated by regressing the model's prediction Vo onto the target j, using an loss:
{ w((t) [Vo t) J.12 dt.
(( 2
(13)
5

---

Published as a conference paper at ICLR 2026
To stabilize training under the memoryless schedule where (tt) - as t 0, incorporate a weighting w : [0, 1] + Rşo. Empirically, we find the following scheme effective (see Figure 9):
w(t)
(14)
x2 1 (3())
This scheme normalizes rewards by the scaling factor À and Nown-weight, timesteps with high future variance. For models employing multiple schedulers, weights are averaged across schedulers.
Finally, the weights A are updated using a gradient descent step with V VL()). For numerical imple- mentation, the underlying SDE is simulated with the Buler-Maruyama method (Kloeden & Platen, 1992) using I' steps, and the integrals are approximated as Riemann sums (Anton, 1999) over the same discretization. The procedure to efficiently compute Equation (12) is outlined in Appendix B.3, In the following result, show that this approach converges to the true value function in expectation. Proposition 2.
The value function V is the unique critical point of B[[vv].
Proof outline. We first establish that V is stationary point by computing the functional derivative of B[CvM]. Uniqueness then follows as standard result in SOC (see Appendix A.4).
This proposition provides theoretical justification for using gradient-basee methods to optimize Vo. ensuring that VM converges to the correct value function under appropriate conditions. To build intuition for VM, we next show how it relates to current state-of-the-art fine-tuning algorithms.
Value Matching the gradient-free analogue of Adjoint Matching. Conceptually, Adjoint Matching (AM) (Domingo-Enrich et al., 2025) can be understood as learning the value function gradient Vx V by iteratively matching it single-sample Monte Carlo estimates of Vx.. Our method represents a zeroth-order analogue, where is learned by regressing onto estimates of J, and the gradient is obtained via backpropagation. Thus, AM and VM are procedurally very similar.
Zhao et (2025) introduced a Continuous- Time Proximal Policy Optimization (CT-PPO) algorithm that learns the optimal control by iteratively alternating between training value function and using it to optimize an actor network, starting from the pre-trained model (sce Algorithm 2). We argue that by setting the actor to gPre (x, V Vo. f), the actor optimization step becomes redundant and the VM algorithm emerges. This substantially simplifies the algorithm and eliminates the need for fine-tuning the base model. Moreover, VM requires fewer hyperparameters to achieve optimal performance (see Appendix E.
Value Matching simplifies continuous-time PPO.
6 RESULTS
We now evaluate Value Matching (VM), aiming showcase four primary insights: (i) we verify that VM recovers the tilted distribution in an illustrative environment; (ii) we demonstrate the scalability of VM to dise-ssimnii image and molecular domains; (iii) in the image design task, we show VM is more sample efficient and expressive than CG; and (iv) across image and molecular design tasks, VM reduce resource requirements by over 95% compared to fine-tuning methods, while achieving superior performance.
VM recovers the tilted distribution with non-differentiable re- wards. To confirm that VM optimizes the intended objective, we test it in simple one-dimensional setting, where the base distri- bution is a Gaussian mixture and the reward is binary and non- differentiable. We show that VM successfully converges the tilted distribution, the optimal solution (Figure 4). Next, we consider chal- lenging high-dimensional image and molecular generation settings.
p
p
pre
Pi
VM effectively adapts large-scale image generation models. To demonstrate the general efficacy of VM on more challenging settings, we apply it to the Diffusion Transformer (DiT) (Peebles & Xie, 2023) trained on the 256x256 ImageNet dataset (Deng et al., 2009), and the text-to-image model Stable Diffusion (SD2) (Rombach et al., 2022). For training prompts, randomly selected 40k captions
Figure 4: ID toy experiment.

---

Published as a conference paper at ICLR 2026
6.50
4.86
5.12
13.87
10.26
13.44
(a) SD2, compression.
(b) SD2, incompression.
(c) FlowMol, dipole moment.
Figure 5: Samples with same random seed. Top: Base model. Bottom: VM model.
4
(k || pn)
Reward f
3.5
0.3
VM
CG
0.2
3
Diverged
2.5
0.1
Base
0
50
100
200
400
Reward Scale (A)
Training Steps
Figure 6: Comparison between VM and CG in high- and low-dimensional settings. Left: On the CIFAR image model with aesthetic reward, CG is unstable under moderate reward scaling. Right: simple 2D environment, VM converges significantly faster than CG to the tilted distribution in terms of KL.
from the LVIS dataset (Schuhmann & Bevan, 2023). The reward functions compression and incom- pression, which correspond to minimizing and maximizing the bits per pixel (BPP) of the sample's JPEG-compressed version, quality level 85. These rewards non-differentiable and therefore cannot be directly optimized by gradient-based methods, such as Adjoint Matching. By learning exploit JPEG's frequency-bast based method, we find that VM generates less detailed, low-frequency images under the compression reward and high-frequency Moiré patterns under the incompression reward (Appendix D). Next, investigate the performance of VM on molecular generation.
VM can effectively adapt molecular generation models. In molecular de- sign, we evaluate VM on the continuous FlowMol model (Dunn Koes, 2024b), pre-trained on the GEOM-Drugs dataset (Axclrod Nomeez-Bombarell,, 2022).
Table 2: Performance VM on FlowMol GEOM-Drugs with QED reward. Metries computed over 10k samples.
Atom stable
Stable (%) f
Valid (%) ?
QED (0-1)
Method (A)
(%) f
None (N/A)
98.0
49.5
48.3
0.42
We consider two (non-differentiable)) re-
ward functions: the dipole moment com-
VM (250)
98.4
56.7
50.0
0.49
puted with GFN2-xTB (Bannwarth et al.,
VM (500)
99.0
67.6
51.7
0.49
2019) and QED computed with rdkit (Landrum et al., 2016). In both cases, molecular geometries are first relaxed using GFN-FF (Spicher Grimme, 2020) before evaluating the reward. For the dipole moment reward, fragmented molecules receive score of zero, and for the QED reward, invalid molecules are assigned a score of zero. Because atom types. edges, and formal charges are discrete and geometry relaxation introduces non-smooth transformations, both reward functions are inherently non-differentiable. Across 10K samples, VM increases the average dipole moment to 7.5 Debye from the base model's 6.4 Debye, while simultaneously reducing the fragmentation rate from 31% to 28%. We find that optimizing the dipole moment reward increases the frequency of heteroatoms and halogens, for example yielding a five-fold rise highly electronegative fluorine (Figure 23). Further, under the QED reward, VM simultaneously improves molecular stability, validity, and QED (Table 2). Taken together, these results demonstrate that VM can successfully optimize target molecular properties without resorting reward hacking. Next, compare VM against classifier guidance.
7

---

Published as a conference paper at ICLR 2026
Base
DDPO
CT-PPO
DPOK
SVDD
VM
Compression
Incompression
Aesthetic
-2
9
3.5
Reward
8
7
2.5 6
Base FID 4
300
300
200
200
200
100
100
100
Diversity 1
30
30
30
20
20
20
10
10
10
0
50
100
50
100
0
50
100
Reward Scale (A)
Figure 7: Comparison of VM with fine-tuning and inference-time schemes across various reward functions and scalings with CIFAR base model. Fine-tuning methods suffer from mode collapse, whereas VM does Mean reward. Middle: FID to base model. Bottom: Vendi diversity with CLIP embeddings.
shows more predictable and stable behavior, outperforming SVDD. Top:
VM is more sample efficient and expressive than classifier guidance. In this comparison, we demonstrate two key advantages of VM over higher reward expressivity and greater sample efficiency. To evaluate the first, we use a 32x32 CIFAR-10 base model and the LAION aesthetics reward (Schuhmann, 2022). This experiment reveals that CG becomes unstable at moderate reward scales (Figure 6, left), significant practical limitation given that meaningful optimization often requires higher values. In contrast, VM maintains stable training and leads to much higher rewards. We then evaluate sample efficiency in one-dimensional environment by tracking the KL divergence to the optimal distribution during training. The results show VM consistently converging to superior optima (Figure 6, right), indicating more effective of training samples. The combination of enhanced stability and improved sample efficiency makes VM more robust and practical alternative CG for reward adaptation tasks, especially in resource-constrained settings. Next, compare VM against fine-tuning methods.
VM demonstrates superior stability compared to fine-tuning methods. To assess how VM compares with fine-tuning approaches, benchmark it against fine-tuning approaches CT-PPO (Zhao et al., 2025), DDPO (Black et al., 2023), and DPOK (Fan et al., 2023) on the CIFAR base model across three reward functions: compression, incompression, and aesthetics (Schuhmann, 2022), which commonly used in prior work. Our goal is to characterize the behavior of all methods under varying reward scales X. As shown Figure 7, VM exhibits stronger robustness and practical reliability than competing approaches. On the compression task, all fine-tuning approaches suffer from mode collapse, whereas remains stable. Across the incompression and aesthetics tasks, the diversity metric indicates that both DDPO and DPOK also collapse to single mode as À increases, and they consistently underperform VM terms of reward. On the other hand, CT-PPO achieves competitive performance, but only after extensive hyperparameter grid search (Appendix E) which is not necessary for VM because it has only a single hyperparameter. Moreover, VM offers more predictable and controllable behavior than CT-PPO under reward scaling. Next, consider the computational cost of VM.
8

---

Published as a conference paper at ICLR 2026
Mcmory (GB)
200
GPU Time (h)
800
150
600
VM, CG C DDDDO
100
400
50
200
0
CIFAR DiT
SD2 FlowMol
CIFAR DiT
SD2 FlowMol
Figure 8: Resource requirements for adaptation methods across base models (16-bit, A100 GPUs).
Left: Memory requirements in GB. Right: Training wall-clock time in hours.
Table Scaling ablation on the CIFAR base model using the Aesthetic score as the reward function. The results are estimates over 3 runs with 1000 samples each, and the standard error is reported. Depth corresponds to the number of blocks per stage, and width the base channel multiplier.
Params (M)
Memory (GB)
Config
Depth
Width
Reward f
Base FID +
Diversity 1
None
N/A
N/A
N/A
N/A
2.31 i: 0.01
17.23 1 0.11
30.15 1 0.10
A
32
0.5
3.2
3.77 i 0.07
179.69
14.43
27.44 1: 2.72
B
1
64
1.8
3.7
4.00 :: 0.09
208.32 :: 4.36
21.81
+ 0.84
C
2
64
3.8
0.009
154.72
28.55 3.40
D
2
128
15.1
5.7
0.11
197.60
21.35
22.85 3.86
E
3
128
23.1
6.6
2.87 = 0.46
88.69 = 52.80
25.61
3.43
F
3
256
92.3
11.2
3,26
0.200
94.29 21.91
32.74 = 0.59
VM is remarkably resource-efficient. As shown in Figure 8, VM demonstrates exceptional com- putational efficiency in comparison to fine-tuning methods. requires less than 12 GB memory across all evaluated models, whereas fine-tuning methods demand up 250 GB for fine-tuning SD2, a reduction of over 95%. The time requirements show similar advantages: VM completes training in under 35 hours for all models, while fine-tuning methods would require up to 800 GPU-hours for SD2. This efficiency gap widens with model scale: while the resource cost for fine-tuning methods grows substantially from CIFAR to SD2, VM maintains a consistently low overhead. These results establish that VM can be orders of magnitude more efficient than fine-tuning alternatives. Next, we consider varying value network sizes.
Small value networks are sufficient. To understand how the complexity the value network influences performance, we conduct scaling ablation on the CIFAR base model using the LAION aesthetic reward (Schuhmann, 2022) with 100. evaluate six configurations (A-F), varying both the depth and width of the CNN architecture, spanning models from 0.5M to 92M parameters. We find that smaller networks are sufficient to achieve strong performance in this setting (Table 3). Increasing model size does not consistently improve reward in these experiments, smaller architectures in the rest this paper. Next, we compare VM against inference-time schemes.
focused the
Value function guidance is significantly putationally efficient than inference- -imee schemes. Unlike inference-time methods such SVDD (Li et al., 2024a), VM adds only minimal overhead during sam- pling, with runtime OT G)), where B is the base model cost and G the cost of computing V Vo t). This yields only modest 1-30% increase in sampling time (Table 4). In contrast, SVDD evaluates the base model and reward for incurring ( + R?)) complexity and more than 40x slowdown for compression and 600x for aesthetics. Moreover, VM achieves higher aesthetic rewards while preserving diversity (Figure 7). Thus, value function guidance offers strong controllability with minimal additional inference cost, making VM far more practical than inference-time approaches.
Table 4: Inference-time overhead of VM, measured as the time to sample a batch of 128 on RTX 4090.
Model
Base time (s)
VM time (s)
CIFAR DiT
3.36
4.00
89.14 0.05
90.37 0.008
SD2 FlowMol
122.02 0.02
127.44 4.002
20 candidates at each step,
12.77 0.14
15.83 1 0.13

---

Published as a conference paper at ICLR 2026
7 RELATED WORK
Reward-guided flow fine-tuning for generative optimization. Recent work has explored fine-tuning flow and diffusion models for objectives beyond likelihood estimation, with leading approaches formulating the problem through RL and SOC frameworks. In this view, the generation process is a sequential decision problem where a policy is learned to steer the model toward desirable outcomes. An early approach, DDPO (Black et al., 2023) applies a policy gradient method to directly optimize for arbitrary rewards but often suffers from "reward collapse". where overfits to a few high-reward samples the cost of diversity. To counter this, DPOK (Fan et al., 2023) incorporated KI. regularization to preserve diversity, though the KI. term was approximated by an upper bound. key insight was that the KL divergence can be computed with quadratic running cost, enabling a control-theoretic interpretation. Leveraging this insight, SOCM (Domingo-Enrich et al., 2024) casts the control problem as an smpporaanee-eigghee regression task. Further advancing this line, Adjoint Matching (Domingo-Enrich et al., 2025) resolved a critical value function bias in earlier methods, enabling provably unbiased reward adaptation. In an effort to address limitations of discretization, Zhao et al. (2025) introduced a continuous-time RL framework. Complementary to these methods, LaSRO (Jia et al., 2025) improves reward adaptation in few-step diffusion models by learning a differentiable latent-space surrogate reward. Recent work also renders it possible maximize rewards while preserving information from prre more generally than KL, as well as enabling risk-averse and risk-sensitive reward optimization (De Santi et al., 2025a; Wang al., 2026). Our work advances this research line by introducing an algorithm that preserves the online nature of control-theoree schemes, while lowering the memory requirements significantly, thereby easing its practical adoption.
Classifier(ffree) guidance. A widely used alternative to fine-tuning is to steer the generative process at inference time. Classifier guidance (Dhariwal & Nichol, 2021) leverages the gradients of a sep- arately trained classifier push the sampling trajectory toward samples that exhibit desired attributes. To eliminate the need for an external model, classifier-free guidance (Ho & Salimans, 2022) modifies the training of the generative model itself to learn both a conditional and an unconditional distribution. At inference, the model is guided by amplifying the difference between the two, effectively steering generations toward the desired condition. These guidance mechanisms are foundational in diffusion model research, and improving upon them has become an active field of study (e.g.. Karras et al., 2024; Sadat al., 2024; 2025; Rajabi al., 2025; Perez Jensen & Sadat, 2025). In this work, we establish connection by showing that VM can be viewed as an online generalization of classifier guidance.
Inference-time schemes beyond guidance. Another family of methods performs reward optimiza- tion at inference-ting time through local, step-wise decisions. Many of these approaches can be understood as approximating an optimal denoising process by leveraging the pre- pre-trained model as look-ahead function to predict future rewards (Uehara et al., 2025). For instance, at each denoising step, methods like SVDD (Li et al., 2024a) and SCG (Huang et al., 2024) evaluate multiple candidate states and select the next state based on these predictions, employing strategics such as resampling or greedy selection. In contrast to such local methods, FlowGrad (Liu et al., 2023b), D-Flow (Ben-Hamm et al., 2024), and OC-Flow (Wang et al., 2025) adopt a global perspective, optimizing the entire trajectory by leveraging reward gradients. The downside of these methods is substantially increased wall-clock time for generation. VM does not suffer from this by amortizing the optimization cost during training.
8 CONCLUSION AND OUTLOOK
We introduce Value Matching (VM), a scalable and efficient online method for adapting pre-trained flow models to arbitrary reward functions. Drawing from fundamental insights in optimal control theory, VM learns the value function, which yields several key advantages. First, the resource requirements are controllable by flexible choice of the value network architecture. Second, by learning the value function, VM naturally handles non-differentiable rewards, a crucial capability for black-box optimization.
Our experiments on image and molecular generation tasks demonstrate these benefits empirically, VM reduces memory and compute requirements by up to 95% compared to fine-tuning methods while achieving comparable performance. Furthermore, VM shows higher reward expressivity and proves more sample efficient than classifier guidance. By providing theoretically grounded and practical framework for seward-guidd adaptation, VM opens up promising opportunities for future research, such applying VM more complex, real-world problems.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We provide comprehensive details to ensure the reproducibility of our work. For all algorithms intro- duced, pseudocode is included, and benchmarks are performed against existing, publicly documented methods. In Appendix B, we detail the value network architectures used in this work. Further, in Appendix B.3, show how to compute the integrals practice. Moreover, in Appendix C, give descriptions of the evaluation metrics, including how they are computed. Lastly, in Appendix E. we report the hyperparameters used for the CT-PPO experiments. For other baselines (DDPO, DPOK, and SVDD), we use the hyperparameters specified in their respective paper.
ACKNOWLEDGMENTS
This publication was made possible by the ETH AI Center doctoral fellowship to Riccardo De Santi. The project has received funding from the Swiss National Science Foundation under NCCR Catalysis (grant number 180544 and 225147) and NCCR Automation (grant agreement 5INF40 180545). This work was supported by ETH Zurich Research Grant.
REFERENCES
Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023.
Howard Anton. Calculus: A new horizon. New York: Wiley, 1999.
Simon Axelrod and Rafael Gomez-Bombaree Geom, energy- annegtt-d molecular conformations for property prediction and molecular generation. Scientific Data, 9(1):185, 2022.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06455) 2016.
Christoph Bannwarth, Sebastian Ehlert, and Stefan Grimme. Gfn2-xtb an accurate and broadly parametrized self-consistent tight-binding quantum chemical method with multipole electrostatics and density-dependen dispersion contributions. Journal of chemical theory and computation, 15 (3):1652- 1671, 2019.
Richard E Bellman and Stuart E Dreyfus. Applied dynamic programming. Princeton university press, 2015.
Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman. D-flow: Differentiating through flows for controlled generation. In Forty-first International Conference Machine Learning, 2024.
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305. 13301, 2023.
Oskar Bolza. Lectures the Calculus of Variations, volume 14. University of Chicago Press, 1904.
Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal of Robotics Research, pp. 02783649241273668. 2023.
Kevin Clark, Paul Vicol, Kevin Swersky, and David J. Fleet. Directly fine-tuning diffusion models on differentiable rewards. In The Twelfth International Conference on Learning Representations, 2024.
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi S. Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. In The Eleventh International Conference Learning Representations, 2023.
Riccardo De Santi, Marin Vlastelica, Ya-Ping Hsieh, Zebang Shen, Niao He, and Andreas Krause. Flow density control: Generative optimization beyond entropy-regularized fine-tuning. In The Exploration in Al Today Workshop at ICMI. 2025, 2025a.

---

Published as a conference paper at ICLR 2026
Riccardo De Santi, Marin Vlastelica, Ya-Ping Hsieh, Zebang Shen, Niao He, and Andreas Krause. Provable maximum entropy manifold exploration via diffusion models. In Forty-second Interna- tional Conference on Machine Learning, 2025b.
Riccardo De Santi, Kimon Protopapas, Ya-Ping Hsieh, and Andreas Krause. Verifier-constraineed flow expansion for discovery beyond the data, 2026. URL httpps///aa.x..ro///.as.oe. 15984.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 JEEE conference on computer vision and pattern recognition, pp. 248-255. Ieee, 2009.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-888 10944 2021.
Carles Domingo- Ennggh,,,,ch,,,,, Jiequn Han, Brandon Amos, Joan Bruna, and Ricky TQ Chen. Stochastic optimal control matching. Advances in Neural Information Processing Systems, 37:1 12459-111111 112004, 2024.
Carles Domingo-Enrich 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 The Thirteenth International Conference on Learning Representations, 2025.
Ian Dunn and David Ryan Koes. Exploring discrete flow matching for 3d novo molecule generation.
ArXiv, pp. arXiv-2411, 2024a.
Ian Dunn and David Ryan Koes. Mixed continuous and categorical flow matching for 3d de novo molecule generation. ArXiv, pp. arXiv-2404, 2024b.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference on machine learning, 2024.
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Reinforcement learning for fine-tuning text-to-image diffusion models. In Thirty-seventh Conference Neural Information Processing Systems (NeurIPS) 2023. Neural Information Processing Systems Foundation, 2023.
Wendell H Fleming and Raymond W Rishel. Deterministic and stochastic optimal control, volume 1. Springer Science & Business Media, 2012.
Wendell H Fleming and H Mete Soner. Controlled Markov processes and viscosity solutions. Springer, 2006.
Stefano Forli, Ruth Huey, Michael E Pique, Michel F Sanner, David S Goodsell, and Arthur J Olson. Computational protein-ligand docking and virtual drug screening with the autodock suite. Nature protocols, 1155:9955-919, 2016.
Avner Friedman. Stochastic differential equations and applications. In Stochastic differential equations, pp. 75-148. Springer, 1975.
Dan Friedman and Adji Bousso Dieng. The vendi score: A diversity evaluation metric for machine learning. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. Advances in Neural Information Processing Systems, 37: 133345-133388 2024.
Yatharth Gupta, Vishnu V Jaddipal, Harish Prabhala, Sayak Paul, and Patrick Von Platen. Progressive knowledge distillation of stable diffusion xl using layer level loss. arXiv preprint 4rX2v72401.02677 2024.
Dan Hendrycks and Kevin Gimpel. Gaussian error arXiv: 16006084455 2016.
linear units (gelus).
arXiv preprint
12

---

Published as a conference paper at ICLR 2026
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12590 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 2020.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation 3d. In International conference on machine learning, pp. 8867-8887. PMLR, 2022.
Zijing Hu, Fengda Zhang, Long Chen, Kun Kuang, Jiahui Li, Kaifeng Gao, Jun Xiao, Xin Wang, and Wenwu Zhu. Towards better alignment: Training diffusion models with reinforcement learning against sparse rewards. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 23604-260232 23614,, 2025.
Yujia Huang, Adishree Ghatare, Yuanzhe Liu, Ziniu Hu, Qinsheng Zhang, Chandramouli Shama Sastry, Siddharth Gururani, Sageev Oore, and Yisong Yue. Symbolic music generation with non- differentiable rule guided diffusion. Forty-first International Conference on Machine Learning, 2024.
James P Hughes, Stephen Rees, S Barrett Kalindjian, and Karen L Philpott. Principles of early drug discovery. British journal of pharmacology, 162(6):1239-12))) 2011.
Zhiwei Jia, Yuesong Nan, Huixi Zhao, and Gengdai Liu. Reward fine- fine-tuning two-step diffusion models via learning differentiable latent-space surrogate reward. In Proceedings of the Computer Vision and Putlern Recognition Conference, pp. 12912-12922, 2025.
Mark Kac. On distributions of certain wiener functionals. Transactions of the American Mahematical Society, 65(1):1-13, 1949.
Tero Karras, Miika Aittala, Tuomas Kynkäänniemi, Jaakko Lehtinen, Timo Aila, and Samuli Laine. Guiding diffusion model with a bad version of itself. Advances in Neural Information Processing Systems, 37:52996-53021, 2024.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1eXiv::412.69990 2014.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations, 2017.
Leo Klarner, Tim GJ Rudner, Garrett M Morris, Charlotte M Deane, and Yee Whye Teh. Context-guided diffusion for out-of-distributio molecular and protein design. arXiv preprint arXiv:2407.11942, 2024.
Peter E Kloeden and Eckhard Platen. Numerical solution of stochastic differential equations. Springer Berlin, Heidelberg, 1992.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical report, 2009.
Landrum et Rdkit: Open-source cheminformatics. https::/rdkit..n 2016.
Xiner Li, Yulai Zhao, Chenyu Wang, Gabriele Scalia, Gokcen Eraslan, Surag Nair, Tommaso Biancalani, Shuiwang Ji, Aviv Regev, Sergey Levine, et al. Derivative-free guidance in continuous and discrete diffusion models with soft balued decoding. arXiv preprint 2488:2408.00222 2024a.
Zihao Li, Hui Yuan, Kaixuan Huang, Chengzhuo Ni, Yinyu Ye, Minshuo Chen, and Mengdi Wang. Diffusion model for data-driven black-box optimization. CoRR, 2024b.
13

---

Published as a conference paper at ICLR 2026
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen, David Lopez-Paz. Heli Ben-Hamu, and Itai Gat. Flow matching guide and code. arXiv preprint MrXiv:2412.06264. 2024.
Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023a.
Xingchao Liu, Lemeng Wu, Shujian Zhang, Chengyue Gong, Wei Ping, and Qiang Liu. Flowgrad: Controlling the output of generative odes with gradients. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24335-2222222 22344,, 2023b.
Dimitra Maoutsa, Sebastian Reich, and Manfred Opper. Interacting particle solutions of fokker- planck equations through gradient-log-denstyy estimation. Entropy, 22(8):802, 2020.
Mareus Olivecrona, Thomas Blaschke, Ola Engkvist, and Hongming Chen. Molecular de-novo design through deep reinforcement learning. Journal of cheminformatics, 9(1):48, 2017.
K. Pandey, F. Sofian, F. Draxler, T. Karaletsos, and S. Mandt. Variational control for guidance in diffusion models. International Conference on Machine Leurning (ICML), 2025.
Amey P Pasarkar and Adji Bousso Dieng. Cousins of the vendi score: A family similarity- bassed diversity metries for science and machine learning. In International Conference on Artificial Intelligence and Statistics, pp. 3808-3816 16. PMLR, 2024.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 4195-4205, 2023.
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with general conditioning layer. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018
Cristian Perez Jensen and Seyedmorteza Sadat. Efficient distillation of classifier-fre guidance using adapters. Transactions Machine Learning Research, 2025.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. SDXL: Improving latent diffusion models for high-resolution image synthesis. In The Twelfth International Conference on Learning Representations, 2024.
Lev Semenovich Pontryagin. Mathematical theory of optimal processes. Routledge, 1962.
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. International conference on machine learning, pp. 8748-876 PMLR, 2021.
Javad Rajabi, Soroush Mehraban, Seyedmorteza Sadat, and Babak Taati. Token perturbation guidance for diffusion models. arXiv preprint arXiv:2506. 10036, 2025.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Seyedmorteza Sadat, Otmar Hilliges, and Romann M Weber. Eliminating oversaturation and artifacts of high guidance scales in diffusion models. In The Thirteenth International Conference on Learning Representations, 2024.
Seyedmorteza Sadat, Manuel Kansy, Otmar Hilliges, and Romann M. Weber. No training, problem: Rethinking classifier free guidance for diffusion models. In The Thirteenth International Conference Learning Representations, 2025.
14

---

Published as a conference paper at ICLR 2026
Christoph Schuhmann. Laion aesthetics predictor vl, 2022. URL https://gittub..cco
LALOO--Al/asstheti-ppeeedd
Christoph
Schuhmann
and Peter
Bevan.
Gpt4vision
captions for
livis, 2023.
URL https: httpss//hhggingaac...c co http://fwa...ccctttttt/ts/sseeee dataatts/sssssninn 220k-GT44ission wwwpionnnpppiins--ooooooo LIVIS.
David S Sholl and Janice A Steckel. Density functional theory: practical introduction. John Wiley & Sons, 2009.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference machine learning, pp. 2256-2265. pmlr, 2015.
Yang Song, Jascha -ihh-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 2020.
Sebastian Spicher and Stefan Grimme. Robust atomistic modeling of materials, organometallic, and biochemical systems. Angewandle Chemie International Edition, 19(36):15665-1556)) 2020.
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Du- mitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1-9, 2015.
Wenpin Tang. Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond. arXiv preprint arXiv:2403.06273 2024.
Masatoshi Uchara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathanicl Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine-tuning of continuous-tima diffusion models as entropy-regularizzed control. arXiv preprint arXiv:2402.15190 2024.
Masatoshi Uehara, Yulai Zhao, Chenyu Wang, Xiner Li, Aviv Regev, Sergey Levine, and Tommaso Biancalani. Inference- time alignment in diffusion models with reward-guide generation: Tutorial and review. arXiv preprint arXiv:2501.096 09685, 2025.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017
Luran Wang, Chaoran Cheng, Yizhen Liao, Yanru Qu, and Ge Liu. Training free guided flow- dow-matching with optimal control. In The Thirteenth International Conference on Learning Representations, 2025.
Zifan Wang, Riccardo De Santi, Xiaoyu Mo, Michael M. Zavlanos, Andreas Krause, and Karl H. Johansson. Efficient tail-aware generative optimization via flow model fine-tuning, 2026. URL https: /ttss://aaxiiii org/ abs 2222 16796.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp. 3-19, 2018
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. Advances in Neural Information Processing Systems, 56035593-55 2023.
Richard Zhang. Making convolutional networks shift-invariant again. In International conference on machine learning, pp. 7324-7334, PMLR, 2019.
Hanyang Zhao, Haoxian Chen, Ji Zhang, David Yao, and Wenpin Tang. Score as action: Fine tuning diffusion generative models by continuous-time reinforcement learning. In Forty-second International Conference Machine Learning, 2025.
15

---

Published as a conference paper at ICLR 2026
APPENDICES
CONTENTS
A Proofs
17
A.1 Assumptions
17
A.2 Useful Lemmas
17
A.3 Proposition
18
A.4 Proposition 2
18
A.5 Derivation of Classifier Guidance Value Function
19
B
Experimental Details
20
B.1
Value Network Architectures
20
B.2
Efficiently Computing Rewards
20
B.3
Efficiently Computing the Cost Functional Estimate
20
B.4
Weighting Functions
21
Evaluation Metrics
22
C.1 Fréchet Inception Distance
22
C.2 Vendi Score
22
Samples and Training Curves
23
D.1 Diffusion Transformer
23
D.2
Stable Diffusion
................................
24
D.3 GEOM-Drugs
27
Baselines
29
E.1
Continuous-Timm Proximal Policy Optimization (CT-PPO)
29
Flow Models
31
F.1
Flow Matching
31
E2 Diffusson Models
32
F3 Diffusion Models an Instance of Flow Matching
33
16

---

Published as a conference paper at ICLR 2026
A PROOFS
A.1 ASSUMPTIONS
Assumption 1. a(t) a (t) is uniformly elliptic.
Assumption The base drift Rd [0. 1] R Lipschitz continuous in x and continuous in t.
Assumption 3. The norm of the base drift |]b]] bounded.
Assumption 4. The reward function r: I - R is bounded.
A.2 USEFUL LEMMAS
Lemma 1 (Application of the Feynman-Kan Kac formula (Kac, 1949)), Let V be the value function defined in Equation (5), then:
V(x,t) (xx,t)= = - log Eppre (exeeeex(((((x)))
x].
(15)
Lemma 2 (Friedman (1975); Chapter 6, Theorem 4.5). Under Assumptions 1 to 3, the transition density Ps el x) of the uncontrolled SDE satisfies the following upper bound on its norm for 0 ttsss
VxPsjely x)|| S -t) exp (-c lly -x1xx
(16)
s-t
constants and x,y RR.. Further, VxPolely x), V2Pajely x), and Oipale(y x) are uniformly continuous.
where C,C
Lemma 3 (Fleming & Soner (2006), Chapter 5, Theorem 9.1). Consider the following Hamilton- equation:
ax, 1) H(x, 1, VxW(x, V =0,
(17)
where in our case the Hamiltonian, H. is:
H(x,t, t, A. = - tr(a(t)A) - (b(x, t). + (t))1
(18)
Assume Assumptions to 4. Let W be a bounded viscosity subsolution and be a bounded viscosity supersolution. Then,
sup (x.1)ER*x(O.1)
(W () V(x, = sup xER
(19)
Lemma (Uniqueness). Under the assumptions of Lemma 3, the viscosity solution to Equation (17) unique.
Proof. We show this result by comparison principle. Assume that Equation (17) has two viscosity solutions Vi and 1/2 with terminal condition Vi (x, 1) = V2(x, 1) -xr(x). Since they are viscosity solutions, they are also viscosity sub- and supersolutions. By their terminal condition, we know that:
sup (W 1) Vx, 1)) sup -Xr(x) Ar(x) 0. XERd
(20)
XERd
We will first show that Vi 222 Apply Lemma 3 with W Vi and V V2, then have: ()) sup (Vi 1) - V2 1)) 0. xER
Vi t) 12(x,t) t) < sup(Vi
(21)
(x,t)
As such have V((x,,) t) < V2 t) for all (x,
Now we show that 1/2 v V1. Again, apply Lemma with W V/2 and = V1, then we have:
12 t) (i t) < sup
V1 (xx
sup (V2 1) V1 1)) 0. (22)
(x.t)
xeRd
Thus we also have V2(x, t) < Vi (x, t) for all (x, t).
In conclusion, we have Vi 1/2 0. meaning that they are equal.
Putting all together, we have that the following HJB equation has a unique solution:
0, O,W(x, () + ur(a(l)V3W(x, ()) + (b(x, 4), VxW) 1 ((x, 1) = -xr
[1 00
(23)
17

---

Published as a conference paper at ICLR 2026
A.3 PROPOSITION
Under Assumptions to 4, the value function V Rd x [0, T defined in Equation (5) is differentiable in x xfr
Proof. Let f < 1. Define (xxx t) A Epera (exx(()) X, = x|. Then from Lemma 1, we have V(x, t) log (x,t), Thus, it suffices show that (1) and (2) is differentiable x:
1. We assume that is bounded, exp(Ar(x)) O. Hence, (xx, > 0.
Writing d an integral we have:
(xx,t) = exp(Ar(y))P1)(( dy.
(24)
Using that r is bounded such that exp(Ar) M for some M and Lemma 2, we can show that the gradient norm of the integrand dominated by an integrable function:
V exp(Ar(Y))) x)] = (xp(Ar(>))lV>XPet)
x)]l
(25)
lly -x12
< MC - -) exp
(26)
1-1
where M, C, c, > 0 are constants. Thus, we can differentiate under the integral:
Vi((x,t= = f exp(Ar(y)) dy.
(27)
Further using Lemma the transition density is continuously differentiable in Thus, ab is differentiable.
This fails for t 1 since r might be non-differentiable and we have V(x, 1) -Ar(x). In conclusion, by the chain rule:
VV(x,t) = (x) t)
(28)
xxx,t) t)
Therefore, V is continuously differentiable in for t < 1.
A.4 PROPOSITION 2
The value function V is the unique critical point of [[[c))
Proof. Let W I2d x [0, 1] - IR be value function approximator and denote W = stopgrad (W) where the argument of stopgrad is treated as constant w.r.t. differentiation. In this proof, assume that any trajectory X| is sampled from the current policy without gradients w.r.t. weights:
dx; = (b(x,, t) (t) (t) t)) dt + ot) dB,.
(29)
Critical point. In order to find the critical points, we will derive the functional derivative of [[[vM]. Let C R8d [0, E be an arbitrary function, then:
d
[[[vM(W cC; <CCxx xoo.))
(30)
de
c=
(W -((oo)) xo., (
E w(()
-
(31)
de
[
2 P w(t) del
d (W - cc)((x,,))) - j(( v xo.).
EE
dt
(32)
==
//' C(x,,t) -(tt))w((): (W(x,,t) -(-- V(oo
E
(33)
18

---

Published as a conference paper at ICLR 2026
Using the tower property of expectation:
E % ((x,,t)-wlf
(o-o))) xxx] |x1) dt
(34)
So, the functional derivative is:
8
W [F(Cm(WW) O] =(l) (W(x, -B(( (oo.._) xx = x)) xx=x]
(35)
SW Thus, any critical point (a point where the functional derivative equals zero) must satisfy:
(x,, E J -ooooW" X =x
(36)
1 L. | VX"" s)1 ds - Xr(x1))
x; = x
(37)
2
By plugging Equation (9) into Equation (5), we know that the value function can be written as:
V(x, x(,)
(38)
= =J[-oVV;;x,)
(39)
E 2 -B12/
1
f.
|0 (sV ds Ar(x))
x, =x
(40)
Therefore V is a critical point of [F[
Uniqueness. As shown, critical point W must satisfy the fixed-point:
1 L. (xx,)))(( ds Ar(xi) xx =x=
W(x,4 () E
(41)
where the expectation is over trajectories from the controlled SDE: dx,- = (b(x,, s) a(s) VW(),, s)) ds + (s) dB,,
x; =
(42)
By the Feynman-Kan Kac formula, W satisfies the following PDE:
1 (r(a + [ - = 0. W(x, 1) -xr(x).
a.W + (b aWW, vW)
(43)
2
Noticing that (VW, wVW))w) = VVV the PDE simplifies to:
0.W + (b. VW) + 1 tr(a - 1o 0, W(x, 1) -Ar(x).
(44)
Using Lemma we know that this HJB equation has a unique solution. This concludes the proof of Proposition 2: V is the unique critical point of B[CvM].
A.5 DERIVATION OF CLASSIFIER GUIDANCE VALUE FUNCTION
In this setting, we have r(x) = log p (u) x) for some class label and = 1. From Lemma 1, we have:
v(x,t) - llgg Epaco (exp(Ar(x1)) x; x]
(45)
= - lloo Eparo (y) x1) X4 = x]
(46)
llo / ppf
ppre (x1 x)) (y) xx) dx
(47)
We have ylx; X1. so by the chain rule:
- =olog /
pre ( dxi
(48)
By marginalization:
log p x).
(49)
This concludes the derivation of the statement.
19

---

Published as a conference paper at ICLR 2026
B EXPERIMENTAL DETAILS
All VM experiments employ the Adam optimizer (Kingma Ba, 2014) with a learning rate of I 10 batch size of 128, and 100 SDE discretization steps. Unless otherwise specified, image experiments utilize a 2.M--arameter CNN and molecular experiments employ a --rammmeee graph neural network (GNN) to parameterize the value function approximator Vg. Additionally, to normalize the CG loss function (as we do for VM by adding the 1/x² term to wo()), we normalize Lcg by dividing it by exp(2A).
B. VALUE NETWORK ARCHITECTURES
Convolutional neural network. We employ standard CNN architecture consisting of input convolution, three downsampling stages, an adaptive average pool, and a final linear head. Timesteps embedded using sinusoidal timestep embedder (Vaswani al., 2017). Each downsampling stage comprises two layers with the following structure: convolution with a 3x3 kernel - group normalization (Wu & He, 2018) > FiLM (Perez al., 2018) to incorporate timestep information + sigmoid linear umit (Hendrycks & Gimpel, 2016) activation function. Finally, the input is added residually and the result is downsampled using blur pool (Zhang, 2019). The convolutional layers in the downsampling stages use base hidden dimensionality of 64, which doubles at each stage.
Graph neural network. The GNN architecture follows a similar design to the CNN architecture (excluding downsampling), where replace the input convolution with linear layer, convolutions with graph convolutions (Kipf& & Welling, 2017), and group normalization with layer normalization (Ba et al., 2016). We utilize all node information available from the FlowMol model: atom type, and formal charge. We also incorporate edge data by linearly transforming the edge features and adding the mean of all incoming edge features to the node features after the input linear layer. Each block uses hidden dimensionality of 256 across 6 stages.
position,
B.2 EFFICIENTLY COMPUTING REWARDS
To efficiently compute reward functions for latent diffusion models, we decode samples individually. This approach significantly reduces VRAM requirements, as decoded samples are typically very large. We find that this strategy does not result in substantially increased wall-clock time.
B.3 EFFICIENTLY COMPUTING THE COST FUNCTIONAL ESTIMATE
We discretize the time horizon into 7 evenly spaced points. On this discretization, we perform the Maruyama method for sampling. Thus, each step, the gradient V computed. Based on this gradient, we compute the running cost at every step:
1 1 V
L,
(50)
At the final time step, the reward R = r(xi) ) is received. The estimated cost functional J, is then com- puted by summing the running costs from time L onward and subtracting the scaled terminal reward:
Je T
Σ
L f/r XR.
(51)
In total, computing the cost functional estimate involves computing d-dimensional norms and adding 7' scalars using reverse cumulative
20

---

Published as a conference paper at ICLR 2026
B.4 WEIGHTING FUNCTIONS
Optimal Transport
DDPM
FlowMol
mltt)
0.5
0.5
0.5
0
0
0
0
0.2
0.4 0.6
0.8
1
0
0.2 0.4 0.6 0.8 1
0
0.2 0.4 0.6 0.8
Timestep (t)
Figure 9: Weighting function w(t) under various (ax, (echeels-s À 1, and the memoryless noise schedule.
Figure 9 displays the weighting function defined in Equation (14) for the schedules of models considered in this work. can seen, it deighte earlier timesteps, which intuitively have highest variance.
21

---

Published as a conference paper at ICLR 2026
C EVALUATION METRICS
Throughout this work, we employ three key metrics to assess performance: average reward, Fréchet inception distance (FID) (Heusel al., 2017), and Vendi diversity (Friedman Dieng, 2023; Pasarkar & Dieng, 2024). The average reward measures an algorithm's 's ability to exploit the reward function effectively, while FID relative to the base model captures the extent of deviation from the base model required to achieve this performance. Vendi diversity quantifies the variety within generated samples. Our objective to achieve high reward and diversity while maintaining low FID. However, an inherent trade-off exists between reward optimization and sample diversity. In this section, we detail the computation of these metries. For each metric, we assume access to dataset of n samples.
C.I FRÉCHET INCEPTION DISTANCE
FID (Heusel al., 2017) is computed by first embedding each data point through pre-trained Inception network (Szegedy et al., 2015) and extracting feature activations from the final layer. The Fréchet distance computes the means (M (42) and covariance matrices (E1, Σ2) of both datasets, then calculates:
dr dx X2)
|| 12ll + tr tr + E2 -
(52)
Typically, X₁ represents reference dataset and X₂ contains samples from the generative model. In this work, however, we set the reference dataset to samples from the base model and X₂ to samples from the neward-adapted model. This provides measure of how much the reward-adapp adaped version has deviated from the base model.
C.2 VENDI SCORE
The Vendi score (Friedman Dieng, 2023; Pasarkar & Dieng, 2024) diversity metric that requires only positive semi-definite similarity function k X' X - R with k(x,, x) 1 for all X. It computes pairwise similarities between all samples and organizes them into matrix K E g=** where kij = ((x;,xxx) Xy). The Vendi score is defined as the exponential of the entropy of the eigenvalues of K/n:
n
VSp((x),.... x,,,),,,,) a exp --E As log A
(53)
In this work, we employ the following similarity function:
k(x,y) clipy)))
(54)
where clip(.) represents a CLIP image encoder (Radford et al., 2021) that produces normalized
embeddings.
22

---

Published as a conference paper at ICLR 2026
D SAMPLES AND TRAINING CURVES
The plotted costs reflect the deviation of the fine-tuned model from the base model; they correspond the KL divergence between the base and controlled processes p" (Domingo-Enrich al., 2025):
[1 DKL (p" (x90.1) (x),))) = =Ep- /
|| t)l dt
(55)
D.1 DIFFUSION TRANSFORMER
5
(a) Base model (mean size: 1.89 bits/pixel).
(b) Compression reward (A 25; mean size: 0.49 bits/pixel).
Incompression reward 25; size: 3.31 bits/pixel).
Figure 10: Samples from Diffusion Transformer generated under the same random seed. Inference with CFG weight 2, whereas training done without CFG.
23

---

Published as a conference paper at ICLR 2026
0
Running Cost
30
Reward
20
10
-2
0
1,000
2,000
1,000
2,000
Training Steps
Figure 11: Training curves for VM on the DiT model with compression reward (A 25).
Running Cost
80
60
Reward
40
4
20
0
1,000
2,000
1,000
2,000
Training Steps
Figure 12: Training curves for VM on the DiT model with incompression reward (A = 25).
D.2 STABLE DIFFUSION
Figure 13: Prompt: A chocolate cake on a plate with decorative pattern, a fork beside it, giving off sense of indulgence or celebration. Reward: Compression (A 2500). CFG weight: 4.0.
24

---

Published as a conference paper at ICLR 2026
Figure 14; Prompt: Skiing scene with multiple individuals dressed in ski gear, engaging in skiing activities amidst snowy surroundings. suggesting a resort or slope ambiance. Reward: Compression (A 2500). CFG weight: 4.0.
Figure 15: Prompt: Dog seated on red lounge chair in a cozy, sophisticated room with a painting, various decorations, and multiple lampshades while wearing collar. Reward: Compression (A 2500). CFG weight: 4.0.
Figure 16: Prompt: surfer in a wet suit performs a carving turn by a pier, on a beach break with no other surfers or boats present. Reward: Incompression (A = 2500). CFG weight: 4.0.
25

---

Published as a conference paper at ICLR 2026
Figure 17: Prompt: An orange cat with blue hut featuring logo, resting on dark-colored
background. Reward: Incompression (A = 2500). CFG weight: 4.0.
Figure 18: Prompt: freight train with cargo containers passes through a railroad crossing. Reward: Incompression (A 2500). CFG weight: 4.0.
-0.5
Running Cost
2,000
Reward
1,000
-1.5
-2
200
400
600
800
200
400
600
800
Training Steps
Figure 19: Training curves for VM the SD2 model with compression reward (A 2500).
26

---

Published as a conference paper at ICLR 2026
Running Cost
1,500
Reward
1,000
500
0
200
400
600
800
200
400
600
800
Training Steps
Figure 20: Training curves for VM on the SD2 model with incompression reward (A 2500).
D.3 FEOM-DRUGE
Fragmented
6,50
5.53
4.22
5.24
Fragmented
4.86
4.50
5.05
Fragmented
5.14
9.26
4.89
8.63
(a) Base model (mean reward: 4.56).
Fragmented
13.87
6.12
Fragmented
4.60
Fragmented
10.26
4.55
5.35
6.41
13.44
3.21
4.55
16.01
(b) VM on dipole moment reward (A
100; mean reward: 6.31).
Figure 21: Samples from the continuous GEOM-Drugs FlowMol base and VM model under the same random seed.
Running Cost
400
M Mwwww//wwwwwaommo
Reward
200
4
0
0
1,000
2,000
1,000
2,000
Training Steps
Figure 22: Training curves for VM on the continuous GEOM-Drugs FlowMol model with dipole moment reward (A 100).
27

---

Published as a conference paper at ICLR 2026
Base Model
Value Matching
60
% Atoms
40
20
C O
S F BI
CNOS F B
Figure 23: Atom type frequency distributions for molecules generated by FlowMol base model and VM with dipole moment reward (A 100). Results shown for 10K generated samples.
28

---

Published as a conference paper at ICLR 2026
E BASELINES
E. CONTINUOUS- TIME PROXIMAL POLICY OPTIMIZATION (CT-PPO)
Algorithm 2 Continuous-Time PPO (CT-PPO) algorithm.
Require: Initial policy parameters 0 (pre-trained), Initial value function parameters Ø, Reward function r TRd + IR, Number iterations N E N. Training steps per iteration K E N. N, Batch size B N, Exploration level E Ryo, Scaling parameter E 3. Clipping parameter c E Rso-
Trajectories per iteration
1: for N iterations do 2: Fix current policy = 0.
Sample trajectories under the current policy:
à ( xx + (1)at) dt + (t)) a, =x((x,)).
dx;
04+
4: Compute returns:
Rs = r(x1) L: (8))sss((x) s spre s)1 2 ds.
5: Initialize dataset Dv {(t,,x,, R4) Re(o, with all trajectories.
6:
repeat K times Sample B DDv with batch-size B.
7:
8:
Compute loss:
((Φ)
Σ (r)E,
(Vo(x,,t) R2)²
9:
Make optimization step with (L()).
10:
end repeat
11:
Sample exploration noise e N(, I) independently for each timestep and trajectory.
12:
Compute psupdessamples and advantages:
a, a + 0Et
1 qi (Vo(xi + no (t)(,)), Vx((x,,))
17
13:
Initialize dataset D, = {1,,,,, x,, āc, gr )ejo.. with all trajectories.
14:
repeat K times
15:
Sample B C Dr, with batch-size B.
16:
Compute likelihood ratio:
( xt,
p
To(a N(x
To àg x),
17:
Compute loss:
((
Σ
min (in(plg.. clip(p,, c, + x)
B
18:
Make an optimization step with (C))).
19: end repeat
20: end for
We set K = [m/B]] such that each point is seen once. For the actor and critic, we use learning rates 3 1 and 1 1 respectively. For data collection, we standardize the number of trajectories and batch size using m = 512, B 128, and N 250. This configuration processes 128K trajectories during training, consistent with other methods in this work.
29

---

Published as a conference paper at ICLR 2026
Hyperparameter ablations. To ensure fair comparison, conduct comprehensive hyper- parameter optimization for CT-PPO through an extensive grid search over the clipping parame- ter 6, exploration level and scale 7). Specifically, conducted grid search on (, n) (0.05, 0.1, 0.2) x {0.01, 0.2, x {0.001, 0.005, 0.01). While this additional tuning effort could be considered part of CT-PPO's computational overhead, it ensures optimal performance for evaluation. In contrast, both VM and AM do not have any hyperparameter search cost.
0.05
= 0.1
0.2
8
0.01
6.5
7.5
7.6
6.4
7.6
7.7
6.5
7.6
7.8
7.5
Scale (n)
0.005
6.3
7.2
7.4
6.3
7.3
7.5
6.3
7.4
7.6
7
6.5
0.001
6.1
6.6
6.7
6.1
6.7
6.8
6.1
6.8
7
6
0.01
0.1
0.2
0.01
0.1
0.2
0.01
0.1
0.2
Exploration Level (a)
(a) Average reward.
0.05
0.1
0.2
0.01
28.4
116.5
118.3
26.6
126.3
133
30
128.2
132.3
Scale (m)
100
0.005
21.6
87.6
99.7
23.1
100.7
99.9
23.5
116.5
131.2
50
0.001
17.5
32.6
39.7
17.7
40.1
47.5
17.5
47.7
56.9
0.01
0.1
0.2
0.01
0.1
0.2
0.01
0.1
0.2
Exploration Level (a)
(b) FID to the base model.
é 0.05
e = 0.1
0.2
31
0.01
30
27.
27.5
29.5
27.2
27.3
30.4
27.7
27.8
Seale (n))
30
0.005
30.1
27.9
26.8
29.7
26.7
28.3
30
27.5
27.1
29
28
0.001
29.7
28.4
28.2
28.9
28.1
27.3
29.9
27.5
27.4
27
0.01
0.1
0.2
0.01
0.1
0.2
0.01
0.1
0.2
Exploration Level (ar
Vendi score (diversity).
Figure 24: Base model: CIFAR. Reward: Incompression (A 100).
30

---

Published as a conference paper at ICLR 2026
F FLOW MODELS
For completeness, this section we provide an overview of flow matching (Lipman et al., 2023; Albergo & Vanden-Eijnden, 2023; Liu et al., 2023a) and how diffusion models (Ho et al., 2020; Sohl-Diee Sohl-Dickstein et al., 2015; Song et al., 2020) can viewed as an instance of it. Further, we show how to sample flow matching models through an SDE with equivalent time marginals. Lastly, show how to sample flow matching or diffusion models from a common perspective through the score function.
F.1 FLOW MATCHING
Given a source distribution and target distribution g, the flow matching framework aims to solve the flow matching problem:
Find the velocity field v : I x [0,1] + Rd generating marginal distributions Po, where Po = p and P1 q.
The flow matching framework solves this problem by the following steps:
1. Identify known source distribution p and unknown target distribution g, of which have finite samples.
Define probability path Pt that interpolates po p and p1 = q.
3. Learn the velocity field by neural network
Sample the learned model by solving an ODE:
dx, ug (x,, 8) dt.
(56)
In general, we could use a coupled data distribution X X1 ~ Pa,1, however, we will only be considering the case where Xo ~ (O, L2). Further, q is unknown, but do assume that have access to a dataset of samples from this distribution. E.g., we might want to model a distribution of images and take the 32 x32 CIFAR-I0 dataset (Krizhevsky et al., 2009) samples from this distribution.
Next, need define a probability path (pe)rejo.1) that interpolates between Po p and pı q- This is done by a conditional strategy, which involves defining Ptj1. We can then construct the marginal probability path by:
(x = f Pqj[(x xi xx)g(x1) dxx)
(57)
We will consider an affine parameterization of the conditional probability path:
(xji( x1) = =N (L
(58)
where au Br J0, 1] + [0, 1] are smooth functions satisfying au B1 0, a1 3o 1, and
de >> 3, for (0, 1). (The dot-notation denotes the time-derivative.) We can sample from this distribution as follows:
X 04Xi + B,x0, Xo Po-
(59)
Commonly, the optimal transport schedule is used where a; and Br L.
Differentiating w.r.t. t gives the associated marginal velocity field:
u(x, Eldxx1 BexO xx x].
(60)
Thus, we can train using the flow matching loss:
(FM(O) A Er.x, [x), - v(x, {)]
(61)
((100 1]), ~ Pt-
However, this is (almost always) intractable, because we do not know the velocity field yet and cannot sample Ps. In order to alleviate these issues, conditioning on the target sample X1
drastically simplify the loss by
(0) A Ex X. [1leo (xe, v(xe, & xx)] ~ U((0, 1)), X1 ~ Pi.xt ~ Paul xi).
(62)
31

---

Published as a conference paper at ICLR 2026
Algorithm 3 Flow Matching training.
Require: Untrained velocity model ve, Schedule ar, 3o [0, 1]
[0, 1], Source distribution p.
Target distribution
1: while not converged do
2: Sample X ~ p. x1 ~ q, and / u(0,, 1])-
3: Compute X x + fix0 and dxpp + $ix0.
4: Compute the loss CCFM (x,(,,,) vell?.
5: Do an optimization step with VOLCFM.
6: end while
7: return Ug
where the conditional velocity is
v(x,t x1) = 6xx1 + Bix0;
x; = x
(63)
Br
(x o4,11)
(64)
132
0rge
X1 +
x.
(65)
Br
Remarkably, these two loss functions have the same gradient w.r.t. the parameters (Lipman et al., 2023):
VoLrm VoLCrM-
(66)
This justifies applying gradient- optimization methods on the conditional loss, which is tractable, because it will lead to the same parameter updates. See Algorithm 3 for the training algorithm. Refer to (Lipman et al., 2024) for an in-depth treatment of flow matching models.
Lastly, instead of sampling from deterministic ODE, we can also consider sampling from family of SDEs:
dx, ((x,,)) t) + o²(t) (u(x,, t) - -x,,,,, dt + (t) dB,
(67)
2m
where B, is a Brownian motion, o [0, 1] + Raxd an arbitrary state-independent diffusion coefficient, and
2 B, d B. -
a Ki
(68)
at
can be shown that the generative processes in Equation (56) and Equation (67) have equivalent time marginals (Maoutsa al., 2020). In the memoryless noise schedule, we have (() /2n.
(Domingo-Emnan Enrich---cccchhe et al., 2025).
F.2 DIFFUSION MODELS
Diffusion models take (slightly) different perspective than flow matching. They view sampling as the reversal of a data destruction (or noising) process. For this, we must first define the noising process:
X++1
Txx 1 -1--et MEt 766
G N0 In),
(69)
where Yh follows some schedule from 0 to T such that XT I.).. As such, starting from Xg = x ~ g, the data gets progressively more like Gaussian noise. Using Gaussian arithmetic, the above process can be computed in closed form:
X -1 I,)
(70)
where 7: = -I-=1] Yo. The denoising process from time to can be computed as follows:
1 7t E, Xx VT 7
+ aiz z NO,
(71)
Here the only unknown Et so we will train a network to approximate see Algorithm
'Generally, (Yt, Ne) are denoted by ā.). This notation is used here to avoid confusion with flow matching schedules.
32

---

Published as a conference paper at ICLR 2026
Algorithm 4 Diffusion model training.
Require: Untrained epsilon model co, Schedule (NO) Target distribution !
1: while not converged do
2: Sample Xg ~ q, E (N In) and t U(TT).
3: Compute X, v 1: + /I -E
4: Compute the loss LDM = (x,( Ell².
5: Do an optimization step with VoLDM-
6: end while
7: return cg
F.3 DIFFUSION MODELS AS AN INSTANCE OF FLOW MATCHING
We can sample diffusion model with the DDIM schedule through the following SDE (Domingo- Enrich et al., 2025);
7t
7t 25
(²(1))
c(x,, 1)
dx;
2
X
+ 2
-I-
dt + (tt dB,
(72)
In order to consolidate diffusion models and flow matching models into a common framework where Po (0, Ia), we will be working with the score function:
sxx, t) 2
Vx log (x
(73)
1
s sxx, f)
(v(xx, t) ki())
(74)
c(x, 1)
s(x, 6)
(75)
VI -16
We can now sample either a diffusion model or flow matching model by converting their parametriza- tion to the score function and sampling the following SDE:
(o2(1)
s(xx,1) di dB,,
dx,
KixXe
me
(76)
In the case of diffusion models, we have
0t v 2t
3,=
/
(77)
with associated time derivatives:
à
8,
(78)
2/77
2-1 - 5t
We will use the convention of flow matching models. Generally, in diffusion models, we have that time is discrete from 0-K and decreases when sampling. Thus, we have the following conversion between the two conventions:
7 [L((1 -]]
(79)
t. = K (7 K(1 - - - - 7K (1-t)))) t)l)
(80)
One can easily verify that this is equivalent to sampling from DDIM.
33


Reviews:
[{"review_id": "WwJw1N2NTx", "reviewer": "Reviewer_1EuJ", "summary": "The authors propose an online algorithm for learning the value function of the flow matching models by formulating the reward-guided generation as an optimal control problem. By drawing the analogy that classifier guidance is seen as offline value function learning, the method enables reward-guided generation without reward fine-tuning the base flow matching model. Specifically, it trains a separate value model as the value function. This admits flexible reward model design (in terms of both architectures and model sizes) and significantly reduces memory usage compared to directly fine-tuning the base model. The method achieves comparable performance with PPO.", "questions": "See weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "+ The connection between CG and VM is intuitive (VM viewed as an online generalization of CG)\n+ The results of VM achieving comparable performance with PPO on image and molecule generation tasks look promising.", "weaknesses": "+ One of the major benefits of using reward guidance over directly doing reward fine-tuning (if the reward is non-differentiable, one can use an approach similar to [1]) is the reduced memory footprint. However, one can adopt LoRA to effectively reduce the memory requirement of the latter. I suggest the author compare different LoRA setups to show the trade-off of memory usage vs. fine-tuned performance of the baseline method to better illustrate how the proposed method does in preservation performance while using much less memory.\n+ Since the value model is decoupled from the base model, it would be good to perform a bit of a scaling study on the scale of the value model.\n+ Optimization with more reward functions will make the results more convincing.\n+ It would be good to add [1] to related work as it also incorporates value function learning for reward fine-tuning flow-matching models.\n\n[1] Reward Fine-Tuning Two-Step Diffusion Models via Learning Differentiable Latent-Space Surrogate Reward, CVPR 2025", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gY1abvkSh1", "reviewer": "Reviewer_UG8D", "summary": "The paper proposes Value Matching (VM), a novel, scalable, and gradient-free algorithm for reward-guided adaptation of large-scale flow and diffusion models. VM uses less memory than finetuning methods and offers improved stability compared to classifier guidance.", "questions": "see above", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. By decoupling the value learning from the base model, VM requires a separate, small value network. This reduces memory usage by over 95% compared to fine-tuning.\n2. VM operates on-policy, which gives it enhanced reward expressivity and stability compared to CG.\n3. The presentation of this work is clear and easy to follow, with abundant theoretical justifications.", "weaknesses": "1. The main quantitative results in Figures 8 and 9 focus only on Reward, KL, and FID. Why not directly compare the generated samples' diversities?\n2. The paper contrasts VM's simplicity with the \"extensive hyperparameter search\" required by CT-PPO. However, this work does not include essential ablations studies on the architecture and size of the separate value network. It is important to know the sensitivity of VM's final performance to variations in the value network. Does VM also require a careful choice of a value network to support acceptable performance?\n3. A concern lies in the selection of inference-time techniques. While CG is a classic baseline, its performance and controllability are not comparable to other non-fine-tuning guidance methods. Why not compare with those in experiments?\n4. Figure 6 does not compare with any baselines. I think any training or non-training methods can perform well for the compressibility task.\n5. In Figure 9, why is VM only compared with CT-PPO? For example, for aesthetic scores, many direct propagation algorithms can easily achieve >7 aesthetic reward in 10 epochs, but they are not mentioned. The current results are also confusing. It seems unclear what this means \"VM demonstrates performance comparable to CT-PPO but with more predictable and stable behavior.\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "F4X7pyNRVc", "reviewer": "Reviewer_9Esi", "summary": "This paper proposes value matching (VM), an online approach that relies on a small learned value network for guiding the model towards higher rewards. The proposed method is derived from optimal control theory and effectively learns a vector field network and a value network that simultaneously approximate the running cost. On downstream generation tasks, the proposed approach achieved better performance across multiple objective rewards.", "questions": "Please refer to the list of weaknesses above. In addition:\n- VM is an *on-policy* algorithm (Algorithm 1), which means it discards past trajectories after each update. On-policy methods are generally known to be sample-inefficient. While the paper shows VM is more efficient than CG, how does its absolute sample efficiency compare to (hypothetical) offline value-based methods?", "limitations": "", "rating": 2, "confidence": 5, "soundness": 1, "presentation": 3, "contribution": 1, "strengths": "- The paper's proposed method is built upon adjoint matching (AM), which is intuitive and easy to follow given the AM framework.\n- The paper, built upon AM, also shares some desirable mathematical properties or guarantees.\n- The proposed method has better running time compared to gradient-based AM models, but the performance comparison remains unclear (see weaknesses).", "weaknesses": "- Although the theoretical part in the paper seems solid to me, the experimental evaluations have **very limited baselines and focus only on small-scale datasets or reward functions** to convincingly demonstrate the effectiveness of the proposed approach.\n  - **Regarding the number of baselines**. For each task, only the CFG and a single model were used as the baseline, even though many existing works were mentioned in the paper. For example, for image generation, there are gradient-based methods such as Adjoint Matching, SOC, FlowGrad [1], and DFlow [2], as well as other RL approaches such as ReFL [3] and [4] (gradient-free) and DRaFT [5] (gradient-based). For molecule generation, DFlow and OCFlow [6] achieved decent generations with almost perfect stability and validity. Current results do not fully demonstrate the superiority of the proposed approach over existing work.\n  - **Regarding the evaluation metric and reward function**. The reward function used in the experiments in this paper is either a toy example from existing work or a non-standard metric not established in previous work on the same task, further weakening the paper's claims, as most metrics are not necessarily comparable. \n    - For example, the standard task in previous flow matching guidance papers on image generation almost all focus on guiding text-to-image models like SD2/3 and use robust metrics like CLIP score, PickScore, or HPSv2 to prevent easily hacking the rewards (e.g., [6] demonstrated that the compression metric can be easily hacked to almost perfect but meaningless generations). For molecule generation, it also remains unclear why the authors did not follow the standard molecule generative modeling evaluation setup in [7] (uses CFG), DFlow, and OCFlow to evaluate the generation and compare the results with these existing works easily. \n    - For molecule generation, it is widely known that quantum chemical calculations are either computationally expensive (ab initio methods) or highly inaccurate (empirical or semiempirical methods). It remains unclear what class the calculation method used in the experiments belongs to, how accurate and generalizable it is (GEOM-Drug is known to have some unreasonable configurations), and what the computational time is, which is crucial for reproducibility. In addition, crucial molecular properties like stability and validity were never mentioned or compared in the paper. Therefore, I am highly skeptical about why the authors did not follow the standard and easy approach in existing work but opted for a seemingly far more complex setup. I would highly discourage such an approach when comparing with baselines for fairness.\n\nTo summarize, for a paper emphasizing the scalability of the proposed method, I believe the existing experiments are, in contrast, limited in scope and poorly credible in supporting its fundamental claims. \n\n- **The theoretical contributions in this paper are limited to me**. The core idea is almost identical to [VGG-Flow](https://openreview.net/forum?id=6MmOy2Ji8V). The scale of the experiments and the number of baselines in this paper fall significantly short of those in VGG-Flow, even if the latter is to be considered concurrent. In addition, there are existing works that have explored the role of the value function or its equivalent, the Q-function, for generative modeling, such as [8] and [9]. Despite different application domains, the underlying core ideas are pretty similar to me. Given that the theoretical results primarily come from the adjoint-matching paper, the theoretical contributions are limited.\n \n- The method's scalability hinges on the value network being \"significantly smaller than the base model.\" But what happens when the reward function is extremely complex? A small network may fail to accurately model the true value function, creating a new performance bottleneck. Is there a trade-off between VM's memory savings and its ability to represent a complex reward landscape? This weakness also echoes in the paper's limited, small-scale evaluation, as thoroughly mentioned in the first part.\n\n- The algorithm learns by regressing the value network's predictions $V_{\\theta}(x_t, t)$ onto a Monte Carlo estimate of the cost functional $\\hat{J}$ (Eq. 12). This target $\\hat{J}$ is based on a single sample trajectory and also depends on the current value function itself. This can be a very high-variance target, which is known to make value-based RL difficult to stabilize. The paper uses a weighting scheme, but the inherent stability of this learning process, especially for very long trajectories, is a potential concern.\n\n- To sample from VM, one must run both the large base model and the (smaller) value network at every step to compute the guiding gradient $\\nabla V_{\\theta}$. While this is still much faster than other gradient-based inference-time optimization schemes, it's not \"free\" and adds computational overhead compared to using a single fine-tuned model. Other gradient-free approaches, such as ReFT and [6], should be benchmarked to support a more credible claim. Additionally, this issue may be coupled with the expressive power of the value net for more complex rewards mentioned above, and it may not be easy to find a balance. \n\n[1] Liu, Xingchao, et al. \"Flowgrad: Controlling the output of generative odes with gradients.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.\n\n[2] Ben-Hamu, Heli, et al. \"D-flow: Differentiating through flows for controlled generation.\" arXiv preprint arXiv:2402.14017 (2024).\n\n[3] Luong, Trung Quoc, et al. \"Reft: Reasoning with reinforced fine-tuning.\" arXiv preprint arXiv:2401.08967 (2024).\n\n[4] Fan, Jiajun, et al. \"Online reward-weighted fine-tuning of flow matching with wasserstein regularization.\" The Thirteenth International Conference on Learning Representations. 2025.\n\n[5] Clark, Kevin, et al. \"Directly fine-tuning diffusion models on differentiable rewards.\" arXiv preprint arXiv:2309.17400 (2023).\n\n[6] Wang, Luran, et al. \"Training free guided flow matching with optimal control.\" arXiv preprint arXiv:2410.18070 (2024).\n\n[7] Hoogeboom, Emiel, et al. \"Equivariant diffusion for molecule generation in 3d.\" International conference on machine learning. PMLR, 2022.\n\n[8] Zhang, Shiyuan, Weitong Zhang, and Quanquan Gu. \"Energy-weighted flow matching for offline reinforcement learning.\" arXiv preprint arXiv:2503.04975 (2025).\n\n[9] Alles, Marvin, et al. \"FlowQ: Energy-Guided Flow Policies for Offline Reinforcement Learning.\" arXiv preprint arXiv:2505.14139 (2025).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Dh6QDlVMon", "reviewer": "Reviewer_U1zm", "summary": "This work proposes an online algorithm called Value Matching (VM) to learn a value function within an optimal control framework, addressing the scalability challenges faced when adapting large-scale flow and diffusion models to downstream tasks via reward optimization. VM enables flexible control over the complexity of the value network, offering adjustable memory and computational requirements. It supports optimization with non-differentiable rewards and operates in an on-policy manner, allowing it to go beyond the training data distribution and explore high-reward regions. Experiments on image generation and molecular design tasks demonstrate that VM achieves better stability and sample efficiency compared to Classifier Guidance (CG).", "questions": "1.Please clarify how your work differs from the following approaches, and discuss the performance gap:\nInference-Time Alignment Control for Diffusion Models with Reinforcement Learning Guidance\nEfficient Controllable Diffusion via Optimal Classifier Guidance\n2.Although your method reduces fine-tuning costs, can it outperform the following full fine-tuning approaches?\nLarge-scale Reinforcement Learning for Diffusion Models\nTraining Diffusion Models with Reinforcement Learning\nDPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1.It learns the scalar value function directly rather than estimating guidance gradients, leading to more stable training.\n2.It offers significant theoretical and practical value by enabling low-cost reinforcement learning (RL) fine-tuning.", "weaknesses": "1.The evaluation is insufficient, with comparisons limited to only CT-PPO.\n2.The core contribution is primarily encapsulated in Equation 11, but the paper devotes excessive space to background, making it hard to follow.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)