WTF?! 基于Wasserstein倾斜流图的无模拟强化学习

arXiv cs.LG 论文

摘要

本文提出Wasserstein-Tilted Flow Maps(WTF),一种无需模拟的强化学习算法,用于微调基于流的生成模型以增强奖励对齐效果。该方法可在减少高达280倍计算量的情况下,实现比基线方法更高的奖励值。

arXiv:2609.27033v1 Announce Type: new Abstract: Reward fine-tuning aims to update a pre-trained flow-based generative model to improve the downstream reward of its generated samples. Existing methods typically formulate this problem as sampling from a reward-tilted distribution, the solution to a KL-regularized reward-maximization problem. Here, we introduce an optimal transport regularizer built directly from the pre-trained drift. Unlike KL reward tilting, the resulting objective transports individual samples toward higher reward rather than reweighting the base distribution. We show that the resulting problem is equivalent to a deterministic optimal control problem on the flow. Given a pre-trained flow map, this equivalence yields a simulation-free reinforcement learning algorithm for fine-tuning generative flows. We call the resulting framework Wasserstein-Tilted Flow Maps (WTF), the first end-to-end fine-tuning recipe native to flow maps. The output is a fine-tuned flow map that retains strong reward-aligned performance at few-step inference budgets without post-hoc distillation. Experiments on ImageNet-256 and text-to-image show that WTF achieves higher reward with comparable or higher diversity than baselines, while requiring up to $280\times$ less training compute. More broadly, we argue that accelerated samplers such as flow maps are essential infrastructure for efficient post-training, and that the dominant KL-regularized formulation is only one of many choices worth revisiting.
查看原文
查看缓存全文

缓存时间: 2026/09/24 09:35

# WTF?!Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps
Source: [https://arxiv.org/html/2609.27033](https://arxiv.org/html/2609.27033)
Jerry Y\. HuangJustin LinPartha KaushikSheel ShahKartik NairYee Whye TehNicholas M\. Boffi

###### Abstract

*Reward fine\-tuning*aims to update a pre\-trained flow\-based generative model to improve the downstream reward of its generated samples\. Existing methods typically formulate this problem as sampling from a reward\-tilted distribution, the solution to a KL\-regularized reward\-maximization problem\. Here, we introduce anoptimal transportregularizer built directly from the pre\-trained drift\. Unlike KL reward tilting, the resulting objective transports individual samples toward higher reward rather than reweighting the base distribution\. We show that the resulting problem is equivalent to adeterministic optimal control problemon the flow\. Given a pre\-trained flow map, this equivalence yields a simulation\-free reinforcement learning algorithm for fine\-tuning generative flows\. We call the resulting frameworkWasserstein\-Tilted Flow Maps \(WTF\), the first end\-to\-end fine\-tuning recipe native to flow maps\. The output is a fine\-tuned flow map that retains strong reward\-aligned performance at few\-step inference budgets without post\-hoc distillation\. Experiments on ImageNet\-256 and text\-to\-image show that WTF achieves higher reward with comparable or higher diversity than baselines, while requiring up to280×280\\timesless training compute\. More broadly, we argue that accelerated samplers such as flow maps are essential infrastructure for efficient post\-training, and that the dominant KL\-regularized formulation is only one of many choices worth revisiting\.

††∗Equal contribution\. Correspondence toabbas\.mammadov@stats\.ox\.ac\.ukandjerryhua@andrew\.cmu\.edu\.![Refer to caption](https://arxiv.org/html/2609.27033v1/figures/front_hero.png)Figure 1:Reward\-aligned generation at few steps\.Flow maps fine\-tuned with WTF for HPSv2\[[1](https://arxiv.org/html/2609.27033#bib.bib41)\], a learned human\-preference reward\. WTF at44steps \(text\-to\-image, left\) and11step \(ImageNet\-256, right\) matches the quality of its own5050\-step generation; the base model is shown at5050steps\. Prompts are given in[SectionI\.2](https://arxiv.org/html/2609.27033#A9.SS2)\.## Introduction

Flow\-based generative models are state\-of\-the\-art for high\-fidelity synthesis across continuous and discrete modalities including images\[[2](https://arxiv.org/html/2609.27033#bib.bib21),[3](https://arxiv.org/html/2609.27033#bib.bib27)\], video\[[4](https://arxiv.org/html/2609.27033#bib.bib22)\], protein structure\[[5](https://arxiv.org/html/2609.27033#bib.bib29)\], materials\[[6](https://arxiv.org/html/2609.27033#bib.bib30)\], and text\[[7](https://arxiv.org/html/2609.27033#bib.bib52),[8](https://arxiv.org/html/2609.27033#bib.bib50),[9](https://arxiv.org/html/2609.27033#bib.bib51)\]\. These models are pre\-trained on large corpora of data to enable generation from the distribution observed during training, but in practice, we rarely desire samples without further preference\. Applications in creative generation, language modeling, and scientific design generically require samples that are likely under the data distribution but which also attain a high score under areward functionrr, which quantifies, for example, alignment with human aesthetic preferences\[[10](https://arxiv.org/html/2609.27033#bib.bib42),[1](https://arxiv.org/html/2609.27033#bib.bib41)\], adherence to safety guidelines or cultural norms\[[11](https://arxiv.org/html/2609.27033#bib.bib47),[12](https://arxiv.org/html/2609.27033#bib.bib63)\], or suitability for downstream tasks such as binding affinity\[[13](https://arxiv.org/html/2609.27033#bib.bib28),[14](https://arxiv.org/html/2609.27033#bib.bib64)\]\. The goal of*reward fine\-tuning*111We write “fine\-tuning” for “reward fine\-tuning” throughout\.is to leverage an additional post\-training phase to update the model’s weights so that its samples improverrwhile remaining close to its base distribution\. The supervision comes from the reward, with no target samples provided, making this a reinforcement learning problem\. Most of these methods require expensive rollouts from the pre\-trained model to score terminal samples underrr, which makes rollout cost a major component of the post\-training budget\.

The prevailing theoretical formulation casts reward fine\-tuning as sampling from a*reward\-tilted*distributionρ~1​\(x\)∝eλ​r​\(x\)​ρ1​\(x\)\\tilde\{\\rho\}\_\{1\}\(x\)\\propto e^\{\\lambda\\,r\(x\)\}\\,\\rho\_\{1\}\(x\), whereρ1\\rho\_\{1\}is the model’s sampling distribution andλ\>0\\lambda\>0is the inverse temperature reward scale\. This reward tilt is the solution to a KL\-regularized reward\-maximization problem\. Qualitatively, this definition reweights samples fromρ1\\rho\_\{1\}according to the reward, making samples with higher reward more likely even if they were rare under the original flow\. Significant recent effort has gone into algorithms that approximate this distribution\[[15](https://arxiv.org/html/2609.27033#bib.bib38),[16](https://arxiv.org/html/2609.27033#bib.bib37),[17](https://arxiv.org/html/2609.27033#bib.bib44),[18](https://arxiv.org/html/2609.27033#bib.bib46),[19](https://arxiv.org/html/2609.27033#bib.bib39)\]through an equivalent diffusion process, rather than operating on the flow itself\. While seemingly awkward, this occurs because the corresponding pathwise KL regularizer can be computed efficiently for diffusions via Girsanov’s theorem\[[20](https://arxiv.org/html/2609.27033#bib.bib26)\]\. Existing methods therefore adopt an approach in which the flow is first converted into a diffusion, the diffusion is fine\-tuned, and the fine\-tuned diffusion is converted back into a flow\.

The fine\-tuning procedure is therefore chiefly diffusion\-based, while the field has primarily moved towards the use of deterministic flows\. Furthermore, increasing attention has been placed on the distillation of flows into few\-step*flow maps*, which amortize the inference process and produce a sample in as few as one network evaluation\[[21](https://arxiv.org/html/2609.27033#bib.bib9),[22](https://arxiv.org/html/2609.27033#bib.bib45),[23](https://arxiv.org/html/2609.27033#bib.bib5),[24](https://arxiv.org/html/2609.27033#bib.bib8),[25](https://arxiv.org/html/2609.27033#bib.bib3)\]\. Recently, Flow Map Reward Guidance \(FMRG\)\[[26](https://arxiv.org/html/2609.27033#bib.bib55)\]suggested that this deterministic structure can be exploited for efficient reward alignment without relying on reward tilting, achieving few\-step inference\-time alignment\. For reward fine\-tuning, however, existing procedures remain reliant on diffusion and cannot exploit these accelerated state transitions during training\. A flow\-centric fine\-tuning paradigm could leverage these distilled generators for dramatically accelerated rollouts, significantly reducing the central computational burden of post\-training\.

This mismatch between flow map deployment and existing fine\-tuning procedures motivates the central question of our work:

*Is there an efficient framework for direct end\-to\-end fine\-tuning of a deterministic flow map?*

Wasserstein tilt \(ours\)arg⁡maxν​\{λ​𝔼ν​\[r\]−𝒯b​\(ρ0,ν\)\}\\arg\\max\_\{\\nu\}\\,\\big\\\{\\,\\lambda\\,\\mathbb\{E\}\_\{\\nu\}\[r\]\-\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu\)\\,\\big\\\}xxlocal transport toward nearby high\-reward regionsReward tiltargmaxν\{λ𝔼ν\[r\]−KL\(ν∥ρ1\)\}∝eλ​rρ1\\arg\\max\_\{\\nu\}\\,\\big\\\{\\,\\lambda\\,\\mathbb\{E\}\_\{\\nu\}\[r\]\-\\mathrm\{KL\}\\left\(\\nu\\\|\\rho\_\{1\}\\right\)\\,\\big\\\}\\propto e^\{\\lambda r\}\\rho\_\{1\}xxreweighting within existing supportFigure 2:Overview\.\(Left\)Our Wasserstein\-tilted approach transports each base sample locally to improve its reward, with the displacement magnitude regularized by the transport cost𝒯b\\mathcal\{T\}\_\{b\}\.\(Right\)Standard reward tilting targetsρ~1∝eλ​r​ρ1\\tilde\{\\rho\}\_\{1\}\\propto e^\{\\lambda r\}\\rho\_\{1\}, which keeps mode positions fixed and only reweights mass within them\.To construct such a framework, we propose to measure deviation from the base distribution via an alternative optimal transport cost built directly from the pre\-trained drift\. Unlike KL reward tilting, the resulting objective transports individual samples toward higher reward rather than reweighting the base distribution\. This choice is natural for a deterministic flow because the cost measures how much the dynamics must be steered away from the pre\-trained drift\. We show that the resulting fine\-tuning problem admits an equivalent deterministic optimal control formulation\. Leveraging this perspective, we arrive at a new class of efficient algorithms in which the flow map emerges as a core component, enablingsimulation\-freebackpropagation through model rollouts by amortizing model inference\. We call the resulting framework Wasserstein\-tilted flow maps \(WTF\)\.

Ourmain contributionsare:

1. 1\.We introduce an optimal transport regularizer between the prior and the candidate terminal distribution that we build directly from the pre\-trained drift\. We show that its induced fine\-tuning objective decouples per particle, moving each sample toward higher reward in contrast to the*reweighting*performed by the reward tilt\.
2. 2\.We prove that fine\-tuning under our regularizer is equivalent to a deterministic optimal control problem with terminal reward, and identify it as the rescaled small\-noise limit of stochastic optimal control\-based algorithms for reward tilting\.
3. 3\.We devise the first end\-to\-end fine\-tuning algorithm for flow maps\. Our method exploits the flow map’s ability to advance the dynamics in a single network evaluation\. In particular, we introduce a Monte Carlo estimate of the deterministic value function, so that terminal samples and value estimates requireO⁡\(1\)O\(1\)flow map evaluations instead ofO⁡\(T\)O\(T\)network evaluations for aTT\-step rollout\. The Monte Carlo value estimator is unbiased and critic\-free, and the fine\-tuned model retains the flow map’s few\-step inference at deployment\.
4. 4\.We evaluate WTF on high\-resolution class\-conditional image synthesis on ImageNet\-256 and text\-to\-image modeling with the TiM\-T2I checkpoint\[[27](https://arxiv.org/html/2609.27033#bib.bib67)\]\. Across both benchmarks, WTF matches or surpasses baselines on reward\. At text\-to\-image scale, WTF reaches matched reward with up to280×280\\timesand47×47\\timesless compute than zeroth\-order and first\-order baselines, respectively\.

## Background

### Flow and flow map\-based generative models

In this work, we assume access to a pre\-trained flow\-based generative model specified by a velocity fieldb:\[0,1\]×ℝd→ℝdb:\[0,1\]\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}\. The resulting probability flow is given by

x˙t=bt​\(xt\),x0∼ρ0,\\dot\{x\}\_\{t\}=b\_\{t\}\(x\_\{t\}\),\\qquad x\_\{0\}\\sim\\rho\_\{0\},\(1\)which transports samples from a tractable prior distributionρ0=𝒩⁡\(0,I\)\\rho\_\{0\}=\\mathcal\{N\}\(0,I\)to the pre\-trained distributionρ1\\rho\_\{1\}\. The velocitybtb\_\{t\}is typically learned over a stochastic interpolant such asIt=\(1−t\)​x0\+t​x1I\_\{t\}=\(1\-t\)x\_\{0\}\+tx\_\{1\}withx0∼ρ0x\_\{0\}\\sim\\rho\_\{0\}andx1∼ρ1x\_\{1\}\\sim\\rho\_\{1\}by minimizing the flow matching objective\[[28](https://arxiv.org/html/2609.27033#bib.bib20),[29](https://arxiv.org/html/2609.27033#bib.bib19),[30](https://arxiv.org/html/2609.27033#bib.bib7),[31](https://arxiv.org/html/2609.27033#bib.bib6)\], but can also be obtained from the probability flow for a diffusion model\[[32](https://arxiv.org/html/2609.27033#bib.bib23),[33](https://arxiv.org/html/2609.27033#bib.bib14),[34](https://arxiv.org/html/2609.27033#bib.bib18)\]\. Samples from the pre\-trainedρ1\\rho\_\{1\}can be obtained by solving[Equation1](https://arxiv.org/html/2609.27033#S2.E1)numerically, which typically requires tens to hundreds of network evaluations and is computationally demanding\[[34](https://arxiv.org/html/2609.27033#bib.bib18)\]\.

To avoid the cost of this numerical integration, recent work has centered on learning theflow mapX:\[0,1\]2×ℝd→ℝdX:\[0,1\]^\{2\}\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}\[[21](https://arxiv.org/html/2609.27033#bib.bib9),[22](https://arxiv.org/html/2609.27033#bib.bib45),[23](https://arxiv.org/html/2609.27033#bib.bib5),[24](https://arxiv.org/html/2609.27033#bib.bib8),[25](https://arxiv.org/html/2609.27033#bib.bib3),[35](https://arxiv.org/html/2609.27033#bib.bib32)\], which is the solution operator of[Equation1](https://arxiv.org/html/2609.27033#S2.E1)\. By definition, the flow map satisfies thejump conditionalong trajectories for any\(s,t\)∈\[0,1\]2\(s,t\)\\in\[0,1\]^\{2\},

Xs,t​\(xs\)=xt\.X\_\{s,t\}\(x\_\{s\}\)=x\_\{t\}\.\(2\)GivenXs,tX\_\{s,t\}, we can generate a sample fromρ1\\rho\_\{1\}in a single evaluationx1=X0,1​\(x0\)x\_\{1\}=X\_\{0,1\}\(x\_\{0\}\)\. More generally, we can use the semigroup propertyXs,t=Xu,t∘Xs,uX\_\{s,t\}=X\_\{u,t\}\\circ X\_\{s,u\}to implement higher\-accuracy multi\-step sampling on an arbitrary grid\[[21](https://arxiv.org/html/2609.27033#bib.bib9),[22](https://arxiv.org/html/2609.27033#bib.bib45)\]\. Recent applications of this approach obtain samples in≤4\\leq 4steps that match the performance of many\-step flows\[[25](https://arxiv.org/html/2609.27033#bib.bib3),[36](https://arxiv.org/html/2609.27033#bib.bib10),[37](https://arxiv.org/html/2609.27033#bib.bib65)\]\.

A complete self\-contained review on both flows and flow maps is provided in[AppendixA](https://arxiv.org/html/2609.27033#A1)\.

### Reward fine\-tuning

Existing approaches for reward fine\-tuning and alignment of flow\- and diffusion\-based generative models take the reward\-tilted distribution with an inverse temperatureλ\>0\\lambda\>0

ρ~1​\(x\)=1Z​eλ​r​\(x\)​ρ1​\(x\),Z=∫eλ​r​\(x\)​ρ1​\(x\)​𝑑x,\\tilde\{\\rho\}\_\{1\}\(x\)=\\frac\{1\}\{Z\}e^\{\\lambda\\,r\(x\)\}\\,\\rho\_\{1\}\(x\),\\quad Z=\\int e^\{\\lambda r\(x\)\}\\rho\_\{1\}\(x\)dx,\(3\)as the fundamental object of interest, and design algorithms to approximately sample it\[[15](https://arxiv.org/html/2609.27033#bib.bib38),[16](https://arxiv.org/html/2609.27033#bib.bib37),[19](https://arxiv.org/html/2609.27033#bib.bib39),[17](https://arxiv.org/html/2609.27033#bib.bib44),[18](https://arxiv.org/html/2609.27033#bib.bib46),[38](https://arxiv.org/html/2609.27033#bib.bib31)\]\.

A classical result based on the Doobhh\-transform\[[39](https://arxiv.org/html/2609.27033#bib.bib33),[20](https://arxiv.org/html/2609.27033#bib.bib26)\]establishes that one can sample fromρ~1\\tilde\{\\rho\}\_\{1\}by estimating the gradient of a stochastic value functionU:\[0,1\]×ℝd→ℝU:\[0,1\]\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\},

Ut​\(x\)=log⁡𝔼⁡\[eλ​r​\(X1\)\|Xt=x\],U\_\{t\}\(x\)=\\log\\mathbb\{E\}\\\!\\left\[\\,e^\{\\lambda\\,r\(X\_\{1\}\)\}\\,\\big\|\\,X\_\{t\}=x\\,\\right\],\(4\)where the conditional expectation is taken over a*memoryless*stochastic process

dXt=bt\(Xt\)dt\+12σ2\(t\)∇logρt\(Xt\)dt\+σ\(t\)dWt,dX\_\{t\}=b\_\{t\}\(X\_\{t\}\)dt\+\\tfrac\{1\}\{2\}\\sigma^\{2\}\(t\)\\,\\nabla\\log\\rho\_\{t\}\(X\_\{t\}\)dt\+\\sigma\(t\)dW\_\{t\},\(5\)withσ⁡\(t\)=2​\(1−t\)/t\\sigma\(t\)=\\sqrt\{2\(1\-t\)/t\}constructed so thatX0X\_\{0\}andX1X\_\{1\}are independent\[[16](https://arxiv.org/html/2609.27033#bib.bib37)\]\. In[Equation5](https://arxiv.org/html/2609.27033#S2.E5),ρt=Law⁡\(xt\)\\rho\_\{t\}=\\mathrm\{Law\}\(x\_\{t\}\)is the density of the probability flow[Equation1](https://arxiv.org/html/2609.27033#S2.E1)\. Given this value function, the modified stochastic process

dX~t=bt\(X~t\)dt\+12σ2\(t\)∇logρt\(X~t\)dt\+σ2\(t\)∇Ut\(X~t\)dt\+σ\(t\)dWt,d\\tilde\{X\}\_\{t\}=b\_\{t\}\(\\tilde\{X\}\_\{t\}\)dt\+\\tfrac\{1\}\{2\}\\sigma^\{2\}\(t\)\\nabla\\log\\rho\_\{t\}\(\\tilde\{X\}\_\{t\}\)dt\+\\sigma^\{2\}\(t\)\\nabla U\_\{t\}\(\\tilde\{X\}\_\{t\}\)dt\+\\sigma\(t\)dW\_\{t\},\(6\)then satisfies thatX~1∼ρ~1\\tilde\{X\}\_\{1\}\\sim\\tilde\{\\rho\}\_\{1\}is a sample from the reward\-tilted measure[Equation3](https://arxiv.org/html/2609.27033#S2.E3)\. The fine\-tuned diffusion process[Equation6](https://arxiv.org/html/2609.27033#S2.E6)can then be converted back into an equivalent fine\-tuned flow,

x~˙t=bt\(x~t\)\+12σ2\(t\)∇Ut\(x~t\),x~0∼ρ0,\\dot\{\\tilde\{x\}\}\_\{t\}=b\_\{t\}\(\\tilde\{x\}\_\{t\}\)\+\\tfrac\{1\}\{2\}\\sigma^\{2\}\(t\)\\,\\nabla U\_\{t\}\(\\tilde\{x\}\_\{t\}\),\\\>\\\>\\tilde\{x\}\_\{0\}\\sim\\rho\_\{0\},\(7\)whose terminal lawLaw⁡\(x~1\)=ρ~1\\mathrm\{Law\}\(\\tilde\{x\}\_\{1\}\)=\\tilde\{\\rho\}\_\{1\}is also the reward\-tilted measure[Equation3](https://arxiv.org/html/2609.27033#S2.E3)\. We refer the reader to[AppendixB](https://arxiv.org/html/2609.27033#A2)for the full derivation of these processes and why they sample from the reward tilt\.

## Wasserstein\-tilted flow maps

While conceptually elegant, the above recipe is clunky in practice, as it revolves around fine\-tuning a diffusion as a surrogate for the flow\. Computationally, the value gradient∇xUt​\(x\)\\nabla\_\{x\}U\_\{t\}\(x\)depends on a terminal\-time conditional expectation, which requires expensive rollouts of the auxiliary process[Equation5](https://arxiv.org/html/2609.27033#S2.E5)all the way tot=1t=1\. These rollouts are slow because stable integration requires many small steps, and they can be high variance because the exponentialeλ​r​\(X1\)e^\{\\lambda r\(X\_\{1\}\)\}can be dominated by a small number of samples with high reward\. To address these pathologies, we develop a fine\-tuning procedure that operates directly on the flow and uses the flow map to amortize deterministic value gradient rollouts\.[Figure3](https://arxiv.org/html/2609.27033#S3.F3)compares the resulting value estimators\. Throughout this section, all stated results assume the standard regularity conditions onrr,bb, andρ0\\rho\_\{0\}collected in[SectionD\.1](https://arxiv.org/html/2609.27033#A4.SS1)\.

### An optimal transport regularizer

At a qualitative level, reward fine\-tuning aims to find a new terminal distributionρ~1\\tilde\{\\rho\}\_\{1\}that maximizes the expected reward while “staying close” to the pre\-trained flow\. This can be made precise as the regularized reward\-maximization problem

ρ~1=argmaxν∈𝒫⁡\(ℝd\)λ​𝔼x∼ν​\[r⁡\(x\)\]−ℛ⁡\(ν\),\\tilde\{\\rho\}\_\{1\}=\\argmax\_\{\\nu\\in\\mathcal\{P\}\(\\mathbb\{R\}^\{d\}\)\}\\ \\lambda\\,\\mathbb\{E\}\_\{x\\sim\\nu\}\[r\(x\)\]\-\\mathcal\{R\}\(\\nu\),\(8\)whereν\\nuranges over candidate terminal distributions onℝd\\mathbb\{R\}^\{d\}andℛ:𝒫⁡\(ℝd\)→ℝ≥0\\mathcal\{R\}:\\mathcal\{P\}\(\\mathbb\{R\}^\{d\}\)\\to\\mathbb\{R\}\_\{\\geq 0\}is a regularizer that penalizes the departure ofν\\nufrom the base\. Intuitively, the choice ofℛ\\mathcal\{R\}determines our definition of “close”\. Settingℛ\(ν\)=KL\(ν∥ρ1\)\\mathcal\{R\}\(\\nu\)=\\mathrm\{KL\}\\left\(\\nu\\\|\\rho\_\{1\}\\right\)recovers the reward\-tilted distribution[Equation3](https://arxiv.org/html/2609.27033#S2.E3)of[Section2\.2](https://arxiv.org/html/2609.27033#S2.SS2)\. Settingℛ\\mathcal\{R\}to the entropic Schrödinger\-bridge cost against the auxiliary process[Equation5](https://arxiv.org/html/2609.27033#S2.E5)recovers the approach of[Domingo\-Enrich et al\. \[16\]](https://arxiv.org/html/2609.27033#bib.bib37)\.[SectionD\.2](https://arxiv.org/html/2609.27033#A4.SS2)gives this connection in detail\. Existing algorithms for these choices are stochastic and rely on a diffusion process\.

We instead seek a regularizer defined directly by the deterministic flow\.

As the flow’s primitive operation is transport, we argue that it is natural to search forℛ\\mathcal\{R\}over optimal\-transport problems\. The general Kantorovich functional\[[40](https://arxiv.org/html/2609.27033#bib.bib16)\]with costc:ℝd×ℝd→ℝ≥0c:\\mathbb\{R\}^\{d\}\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}\_\{\\geq 0\},

𝒯c​\(μ,ν\)=infπ∈Π⁡\(μ,ν\)∫c⁡\(x,y\)​π​\(d​x​𝑑y\),\\mathcal\{T\}\_\{c\}\(\\mu,\\nu\)=\\inf\_\{\\pi\\in\\Pi\(\\mu,\\nu\)\}\\int c\(x,y\)\\,\\pi\(dx\\,dy\),\(9\)defines a family of candidate regularizersℛ⁡\(ν\)=𝒯c​\(μ,ν\)\\mathcal\{R\}\(\\nu\)=\\mathcal\{T\}\_\{c\}\(\\mu,\\nu\)parameterized by the costccand the source measureμ\\mu, whereΠ⁡\(μ,ν\)\\Pi\(\\mu,\\nu\)denotes the set of couplings with marginalsμ\\muandν\\nu\.

#### Defining the regularizer\.

The choicec⁡\(x,y\)=12​∥x−y∥2c\(x,y\)=\\tfrac\{1\}\{2\}\\lVert x\-y\\rVert^\{2\}defines the standardW2W\_\{2\}distance\[[40](https://arxiv.org/html/2609.27033#bib.bib16)\], but neitherℛ⁡\(ν\)=12​W22​\(ρ1,ν\)\\mathcal\{R\}\(\\nu\)=\\tfrac\{1\}\{2\}W\_\{2\}^\{2\}\(\\rho\_\{1\},\\nu\)norℛ⁡\(ν\)=12​W22​\(ρ0,ν\)\\mathcal\{R\}\(\\nu\)=\\tfrac\{1\}\{2\}W\_\{2\}^\{2\}\(\\rho\_\{0\},\\nu\)is suitable for our goals\. In the former case, the optimizer is a transport map betweenρ1\\rho\_\{1\}andν\\nuthat must be composed with the base model at inference, leaving the fine\-tuned model a two\-stage object\. In the latter, the cost knows nothing aboutbborρ1\\rho\_\{1\}\. We therefore propose a regularizerℛ⁡\(ν\)=𝒯cb​\(ρ0,ν\)\\mathcal\{R\}\(\\nu\)=\\mathcal\{T\}\_\{c\_\{b\}\}\(\\rho\_\{0\},\\nu\), wherecbc\_\{b\}measures theresidual controlrequired tosteerthe base flow from an initial pointx∈ℝdx\\in\\mathbb\{R\}^\{d\}to a target pointy∈ℝdy\\in\\mathbb\{R\}^\{d\}\. The sourceρ0\\rho\_\{0\}ensures that the optimizer is a flow from the same prior as the base model\.

For a candidate velocityvv, we define a Lagrangianℒb:\[0,1\]×ℝd×ℝd→ℝ≥0\\mathcal\{L\}\_\{b\}:\[0,1\]\\times\\mathbb\{R\}^\{d\}\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}\_\{\\geq 0\}to measure the instantaneous deviation ofvvfrom the base drift\. We then define our cost functioncb:ℝd×ℝd→ℝ≥0c\_\{b\}:\\mathbb\{R\}^\{d\}\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}\_\{\\geq 0\}as the integrated Lagrangian over pathsω:\[0,1\]→ℝd\\omega:\[0,1\]\\to\\mathbb\{R\}^\{d\}joiningxxtoyy,

ℒb​\(t,x,v\)=12​∥v−bt​\(x\)∥2,cb​\(x,y\)=infω0=x,ω1=y∫01ℒb​\(t,ωt,ω˙t\)​𝑑t,\\mathcal\{L\}\_\{b\}\(t,x,v\)=\\tfrac\{1\}\{2\}\\lVert v\-b\_\{t\}\(x\)\\rVert^\{2\},\\quad c\_\{b\}\(x,y\)=\\inf\_\{\\omega\_\{0\}=x,\\,\\omega\_\{1\}=y\}\\ \\int\_\{0\}^\{1\}\\mathcal\{L\}\_\{b\}\(t,\\omega\_\{t\},\\dot\{\\omega\}\_\{t\}\)\\,dt,\(10\)which measures the minimum residual energy required to steer a particle fromxxtoyywhile followingbtb\_\{t\}as closely as possible\. We then take our regularizer to be the corresponding Kantorovich problem,

R⁡\(ν\)=𝒯b​\(ρ0,ν\)=infπ∈Π⁡\(ρ0,ν\)∫cb​\(x,y\)​π​\(d​x​𝑑y\),R\(\\nu\)=\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu\)=\\inf\_\{\\pi\\in\\Pi\(\\rho\_\{0\},\\,\\nu\)\}\\int c\_\{b\}\(x,y\)\\,\\pi\(dx\\,dy\),\(11\)

which defines the minimum total residual energy needed to steer the base lawρ0\\rho\_\{0\}to a candidate terminal lawν\\nualong the dynamicsbtb\_\{t\}\. Whenb≡0b\\equiv 0,cb​\(x,y\)=12​∥x−y∥2c\_\{b\}\(x,y\)=\\tfrac\{1\}\{2\}\\lVert x\-y\\rVert^\{2\}and𝒯b​\(ρ0,ν\)=12​W22​\(ρ0,ν\)\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu\)=\\tfrac\{1\}\{2\}W\_\{2\}^\{2\}\(\\rho\_\{0\},\\nu\), so that𝒯b\\mathcal\{T\}\_\{b\}generalizes the Benamou–Brenier formulation ofW22W\_\{2\}^\{2\}\[[41](https://arxiv.org/html/2609.27033#bib.bib17)\]to a non\-trivial reference\. This construction has been studied under the name*optimal transport with prior*by[Chen et al\. \[42\]](https://arxiv.org/html/2609.27033#bib.bib49),[Chen et al\. \[43\]](https://arxiv.org/html/2609.27033#bib.bib48),[Chen et al\. \[44\]](https://arxiv.org/html/2609.27033#bib.bib25); we adapt their framework here to fine\-tuning of deterministic generative flows\.

\(a\) Actor–Criticxtx\_\{t\}Vθ​\(t,x\)V\_\{\\theta\}\(t,x\)V^t\\widehat\{V\}\_\{t\}backprop∇xtVθ\\nabla\_\{x\_\{t\}\}V\_\{\\theta\}learned critic networkO⁡\(1\)O\(1\)biasedneeds critic\(b\) Multi\-step rolloutxtx\_\{t\}r⁡\(x1\)r\(x\_\{1\}\)backprop×T\\times TTTnetwork callsunbiasedO⁡\(T\)O\(T\)rolloutO⁡\(T\)O\(T\)memory\(c\) WTF \(ours\)Xt,τu¯X^\{\\bar\{u\}\}\_\{t,\\tau\}Xτ,1u¯X^\{\\bar\{u\}\}\_\{\\tau,1\}xtx\_\{t\}u¯τ\\bar\{u\}\_\{\\tau\}xτx\_\{\\tau\}r⁡\(x1\)r\(x\_\{1\}\)single Jacobian backprop22flow map callsunbiasedO⁡\(1\)O\(1\)no critic

Figure 3:Flow map value estimation\.Common alternatives use either a learned critic, as in actor–critic methods, or a multi\-step rollout that backpropagates throughTTsteps\. WTF instead reachesx1x\_\{1\}fromxtx\_\{t\}in two flow map evaluations and a single Jacobian backprop, which isO⁡\(1\)O\(1\), unbiased, and does not require a critic\.

### Understanding the regularizer

Combining𝒯b\\mathcal\{T\}\_\{b\}with the reward functionrrgives the regularized reward maximization problem

ρ~1=argmaxν∈𝒫⁡\(ℝd\)\{λ​𝔼x∼ν​\[r⁡\(x\)\]−𝒯b​\(ρ0,ν\)\},\\tilde\{\\rho\}\_\{1\}=\\argmax\_\{\\nu\\in\\mathcal\{P\}\(\\mathbb\{R\}^\{d\}\)\}\\ \\left\\\{\\ \\lambda\\,\\mathbb\{E\}\_\{x\\sim\\nu\}\[r\(x\)\]\-\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu\)\\ \\right\\\},\(12\)whose optimizerρ~1\\tilde\{\\rho\}\_\{1\}defines a*Wasserstein\-tilted flow*\(WTF\)\. To understand the nature of this problem, we examine how the distribution\-level optimization acts on individual source samples\. The following result shows that the optimum can be understood as transporting each source sample toward a high\-reward endpoint, subject to the cost of deviating from the pre\-trained flow\.

###### Proposition 3\.1\.

Under the regularity conditions of[SectionD\.1](https://arxiv.org/html/2609.27033#A4.SS1),

supν\{λ​𝔼x∼ν​\[r⁡\(x\)\]−𝒯b​\(ρ0,ν\)\}=𝔼x∼ρ0​\[supy∈ℝd\{λ​r​\(y\)−cb​\(x,y\)\}\],\\sup\_\{\\nu\}\\ \\left\\\{\\ \\lambda\\,\\mathbb\{E\}\_\{x\\sim\\nu\}\[r\(x\)\]\-\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu\)\\ \\right\\\}\\;=\\;\\mathbb\{E\}\_\{x\\sim\\rho\_\{0\}\}\\\!\\left\[\\,\\sup\_\{y\\in\\mathbb\{R\}^\{d\}\}\\ \\bigl\\\{\\,\\lambda\\,r\(y\)\-c\_\{b\}\(x,y\)\\,\\bigr\\\}\\,\\right\],\(13\)and the optimum is attained by choosingy∗​\(x\)=argmaxy\{λ​r​\(y\)−cb​\(x,y\)\}y^\{\*\}\(x\)=\\argmax\_\{y\}\\\{\\,\\lambda\\,r\(y\)\-c\_\{b\}\(x,y\)\\,\\\}for eachx∼ρ0x\\sim\\rho\_\{0\}\.

The proof, given in[SectionE\.1](https://arxiv.org/html/2609.27033#A5.SS1), follows from Kantorovich duality\[[40](https://arxiv.org/html/2609.27033#bib.bib16),[45](https://arxiv.org/html/2609.27033#bib.bib15)\]\. The result gives a concrete particle\-level interpretation of the objective\. The optimal solution transports each source sample to the destination that gives its best reward\-cost tradeoff, rather than selecting the terminal law only through global reweighting\.

#### Comparison with KL reward tilting\.

The sample\-level distinction above reflects a broader structural difference between the two objectives\. KL reward tilting reweights the base distribution according to reward, whereas WTF transports samples under a cost induced by the pre\-trained drift\. As a consequence, reward reweighting preserves the base conditional distribution within each reward level set, while WTF can redistribute mass within a level set\.[AppendixI](https://arxiv.org/html/2609.27033#A9)gives exact examples of this distinction\.

While our proposed regularizer[Equation11](https://arxiv.org/html/2609.27033#S3.E11)is natural for the problem under study, and while[Proposition3\.1](https://arxiv.org/html/2609.27033#S3.Thmtheorem1)demonstrates a clean interpretation of what it does at the sample level, it is still unclear how to implement it algorithmically\. The following result shows that[Equation12](https://arxiv.org/html/2609.27033#S3.E12)can be reformulated as a deterministic optimal control problem with terminal reward, which will enable us to develop scalable algorithms\.

###### Proposition 3\.2\.

Under standard regularity conditions onrr,bb, andρ0\\rho\_\{0\}, the optimal value of[Equation12](https://arxiv.org/html/2609.27033#S3.E12)equals the optimal value of the deterministic optimal control problem

supu𝔼x0∼ρ0​\[λ​r​\(x1u\)−12​∫01∥ut​\(xtu\)∥2​dt\],\\displaystyle\\sup\_\{u\}\\ \\mathbb\{E\}\_\{x\_\{0\}\\sim\\rho\_\{0\}\}\\\!\\left\[\\,\\lambda\\,r\(x\_\{1\}^\{u\}\)\-\\frac\{1\}\{2\}\\int\_\{0\}^\{1\}\\lVert u\_\{t\}\(x\_\{t\}^\{u\}\)\\rVert^\{2\}\\,dt\\,\\right\],\(14\)subject tox˙tu=bt​\(xtu\)\+ut​\(xtu\),x0u=x0\.\\displaystyle\\text\{subject to\}\\quad\\dot\{x\}\_\{t\}^\{u\}=b\_\{t\}\(x\_\{t\}^\{u\}\)\+u\_\{t\}\(x\_\{t\}^\{u\}\),\\quad x\_\{0\}^\{u\}=x\_\{0\}\.

Whenever an optimizeru∗:\[0,1\]×ℝd→ℝdu^\{\*\}:\[0,1\]\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}exists, the terminal lawρ1u∗=Law⁡\(x1u∗\)\\rho^\{u^\{\*\}\}\_\{1\}=\\mathrm\{Law\}\(x\_\{1\}^\{u^\{\*\}\}\)solves[Equation12](https://arxiv.org/html/2609.27033#S3.E12)\.

The proof is given in[SectionE\.2](https://arxiv.org/html/2609.27033#A5.SS2), which proceeds via a Benamou\-Brenier\-style dynamic reformulation of the static regularizer[Equation11](https://arxiv.org/html/2609.27033#S3.E11)established by[Chen et al\. \[42\]](https://arxiv.org/html/2609.27033#bib.bib49),[Chen et al\. \[43\]](https://arxiv.org/html/2609.27033#bib.bib48),[Chen et al\. \[44\]](https://arxiv.org/html/2609.27033#bib.bib25)\. A direct consequence is that the fine\-tuned model is again a single driftv~t=bt\+ut∗\\tilde\{v\}\_\{t\}=b\_\{t\}\+u^\{\*\}\_\{t\}that replacesbtb\_\{t\}at inference\. We further show in[SectionD\.2](https://arxiv.org/html/2609.27033#A4.SS2)that the problem[Equation14](https://arxiv.org/html/2609.27033#S3.E14)is the zero\-noise limit of a Schrödinger\-bridge problem\[[46](https://arxiv.org/html/2609.27033#bib.bib24),[43](https://arxiv.org/html/2609.27033#bib.bib48),[42](https://arxiv.org/html/2609.27033#bib.bib49),[44](https://arxiv.org/html/2609.27033#bib.bib25)\], connecting WTF to the stochastic optimal control framework used by Adjoint Matching and related methods\[[16](https://arxiv.org/html/2609.27033#bib.bib37),[17](https://arxiv.org/html/2609.27033#bib.bib44),[18](https://arxiv.org/html/2609.27033#bib.bib46),[38](https://arxiv.org/html/2609.27033#bib.bib31)\]\.[Table8](https://arxiv.org/html/2609.27033#A4.T8)situates the corresponding regularizers and algorithms\.

## Solving the optimal control problem

We now develop a practical algorithm to solve the optimal control problem[Equation14](https://arxiv.org/html/2609.27033#S3.E14)using a pre\-trained flow mapXs,t​\(x\)=x\+\(t−s\)​vs,t​\(x\)X\_\{s,t\}\(x\)=x\+\(t\-s\)\\,v\_\{s,t\}\(x\)for the uncontrolled dynamicsx˙t=bt​\(xt\)\\dot\{x\}\_\{t\}=b\_\{t\}\(x\_\{t\}\)withbt=vt,tb\_\{t\}=v\_\{t,t\}\. In practice, we find that direct backpropagation through a sampled version of[Equation14](https://arxiv.org/html/2609.27033#S3.E14)exhibits a reward\-hacking failure mode, where the terminal map can increase reward while the diagonal velocity fails to learn the corresponding controlled dynamics; see[SectionD\.3](https://arxiv.org/html/2609.27033#A4.SS3)for details\. To avoid this pathology, we instead regress the diagonal control onto the value gradient of a frozen reference control, giving a policy\-improvement update\. The flow map’s long\-range transitions let us construct the on\-policy states and value estimates required by this update without numerically simulating every intermediate state, yielding simulation\-free rollouts and an end\-to\-end fine\-tuning algorithm for flow maps\. For a controluu, define the value function as the reward\-to\-go from statexxat timettminus the remaining control cost,

Vtu​\(x\)=λ​r​\(Xt,1u​\(x\)\)−12​∫t1∥uτ​\(Xt,τu​\(x\)\)∥2​𝑑τ,V^\{u\}\_\{t\}\(x\)=\\lambda\\,r\(X^\{u\}\_\{t,1\}\(x\)\)\-\\frac\{1\}\{2\}\\int\_\{t\}^\{1\}\\lVert u\_\{\\tau\}\(X^\{u\}\_\{t,\\tau\}\(x\)\)\\rVert^\{2\}\\,d\\tau,\(15\)so that maximizing𝔼x0​\[Vu​\(0,x0\)\]\\mathbb\{E\}\_\{x\_\{0\}\}\[V^\{u\}\(0,x\_\{0\}\)\]recovers[Equation14](https://arxiv.org/html/2609.27033#S3.E14)\. ThusVtu​\(x\)V^\{u\}\_\{t\}\(x\)is the future objective for a trajectory starting fromxxat timett;[AppendixC](https://arxiv.org/html/2609.27033#A3)provides a self\-contained introduction to deterministic optimal control and the role of the value function\. A standard result in optimal control theory\[[47](https://arxiv.org/html/2609.27033#bib.bib58)\]states that the optimal control for[Equation14](https://arxiv.org/html/2609.27033#S3.E14)is the gradient of the optimal value function,

ut∗​\(x\)=∇xVt∗​\(x\),Vt∗​\(x\)=supuVtu​\(x\)\.u^\{\*\}\_\{t\}\(x\)=\\nabla\_\{x\}V^\{\*\}\_\{t\}\(x\),\\quad V\_\{t\}^\{\*\}\(x\)=\\sup\_\{u\}V\_\{t\}^\{u\}\(x\)\.\(16\)The following result shows that fitting our trainable control to the current estimate of the value gradient in an iterative fashion is guaranteed to improve the objective at the population level\.

###### Proposition 4\.1\(Performance difference and policy improvement\)\.

Letu¯\\bar\{u\}be a reference control with value functionVtu¯V^\{\\bar\{u\}\}\_\{t\}and value gradientg¯t​\(x\)=∇xVtu¯​\(x\)\\bar\{g\}\_\{t\}\(x\)=\\nabla\_\{x\}V^\{\\bar\{u\}\}\_\{t\}\(x\)\. For any candidate controlwwand any initial pair\(s,x\)\(s,x\)with controlled trajectoryxtw=Xs,tw​\(x\)x^\{w\}\_\{t\}=X^\{w\}\_\{s,t\}\(x\)fort∈\[s,1\]t\\in\[s,1\],

Vsw​\(x\)−Vsu¯​\(x\)=12​∫s1\[∥u¯t​\(xtw\)−g¯t​\(xtw\)∥2−∥wt​\(xtw\)−g¯t​\(xtw\)∥2\]​𝑑t\.V^\{w\}\_\{s\}\(x\)\-V^\{\\bar\{u\}\}\_\{s\}\(x\)=\\frac\{1\}\{2\}\\int\_\{s\}^\{1\}\\\!\\left\[\\,\\lVert\\bar\{u\}\_\{t\}\(x^\{w\}\_\{t\}\)\-\\bar\{g\}\_\{t\}\(x^\{w\}\_\{t\}\)\\rVert^\{2\}\-\\lVert w\_\{t\}\(x^\{w\}\_\{t\}\)\-\\bar\{g\}\_\{t\}\(x^\{w\}\_\{t\}\)\\rVert^\{2\}\\,\\right\]dt\.\(17\)In particular, settingw=g¯w=\\bar\{g\}makes the second norm vanish and givesVsg¯​\(x\)≥Vsu¯​\(x\)V^\{\\bar\{g\}\}\_\{s\}\(x\)\\geq V^\{\\bar\{u\}\}\_\{s\}\(x\)for all\(s,x\)\(s,x\), with equality if and only ifu¯t​\(xtg¯\)=g¯t​\(xtg¯\)\\bar\{u\}\_\{t\}\(x^\{\\bar\{g\}\}\_\{t\}\)=\\bar\{g\}\_\{t\}\(x^\{\\bar\{g\}\}\_\{t\}\); in that caseu¯\\bar\{u\}is the optimal control of[Equation14](https://arxiv.org/html/2609.27033#S3.E14)\.

The proof is given in[SectionE\.3](https://arxiv.org/html/2609.27033#A5.SS3)\. The identity[Equation17](https://arxiv.org/html/2609.27033#S4.E17)implies thatw=g¯w=\\bar\{g\}, the gradient of the current value function, maximizes the improvement term\. Starting from a reference controlu¯\\bar\{u\}, we perform policy iteration by fitting the trainable controlwwto the corresponding value gradientg¯\\bar\{g\}and then usingwwas the reference control for the next iteration\.

By[Proposition4\.1](https://arxiv.org/html/2609.27033#S4.Thmtheorem1), each iterate strictly improves the objective\. The iteration continues to improve until it reaches the fixed pointu¯=g¯\\bar\{u\}=\\bar\{g\}, which is the global optimum\.

#### Simulation\-free value estimation\.

Computing∇xVu¯\\nabla\_\{x\}V^\{\\bar\{u\}\}for a frozen referenceu¯\\bar\{u\}requires integrating along and then backpropagating through the entire future trajectory, both of which are computationally expensive\. To avoid this, we observe that the integrated control cost can be written as an expectation,∫t1∥u¯τ​\(Xt,τu¯​\(x\)\)∥2​𝑑τ=\(1−t\)​𝔼τ∼Unif⁡\[t,1\]​\[∥u¯τ​\(Xt,τu¯​\(x\)\)∥2\]\\int\_\{t\}^\{1\}\\lVert\\bar\{u\}\_\{\\tau\}\(X^\{\\bar\{u\}\}\_\{t,\\tau\}\(x\)\)\\rVert^\{2\}\\,d\\tau=\(1\-t\)\\,\\mathbb\{E\}\_\{\\tau\\sim\\mathrm\{Unif\}\[t,1\]\}\\\!\\left\[\\lVert\\bar\{u\}\_\{\\tau\}\(X^\{\\bar\{u\}\}\_\{t,\\tau\}\(x\)\)\\rVert^\{2\}\\right\]\. Definingxτ=Xt,τu¯​\(x\)x\_\{\\tau\}=X^\{\\bar\{u\}\}\_\{t,\\tau\}\(x\)andx1=Xτ,1u¯​\(xτ\)x\_\{1\}=X^\{\\bar\{u\}\}\_\{\\tau,1\}\(x\_\{\\tau\}\), this gives an unbiased single\-sample Monte Carlo estimate of[Equation15](https://arxiv.org/html/2609.27033#S4.E15),

V^​\(t,x\)=λ​r​\(x1\)−1−t2​∥u¯τ​\(xτ\)∥2\.\\widehat\{V\}\(t,x\)=\\lambda\\,r\(x\_\{1\}\)\-\\frac\{1\-t\}\{2\}\\lVert\\bar\{u\}\_\{\\tau\}\(x\_\{\\tau\}\)\\rVert^\{2\}\.\(18\)

The flow map gives the endpointx1x\_\{1\}without numerically simulating the full future trajectory, so the only remaining Monte Carlo approximation is the scalar time integral\. We emphasize that the only variance in[Equation18](https://arxiv.org/html/2609.27033#S4.E18)comes from our Monte Carlo estimate of a one\-dimensional time integral, in contrast to the high\-dimensional integrals that need to be computed for stochastic value functions arising from reward tilting\. This formula gives an unbiased regression targetg^​\(t,x\)=∇xV^​\(t,x\)\\widehat\{g\}\(t,x\)=\\nabla\_\{x\}\\widehat\{V\}\(t,x\)via automatic differentiation\. The WTF training objective combines value gradient regression on the diagonals=ts=twith an off\-diagonal flow map self\-distillation loss:

ℒWTF​\(w^\)\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{WTF\}\}\(\\hat\{w\}\)=𝔼x0,t,τ​\[∥w^t,t​\(x¯t\)−bt​\(x¯t\)−g^​\(t,x¯t\)∥2\]\+β​ℒdist​\(w^\),\\displaystyle=\\mathbb\{E\}\_\{x\_\{0\},t,\\tau\}\\\!\\left\[\\,\\lVert\\hat\{w\}\_\{t,t\}\(\\bar\{x\}\_\{t\}\)\-b\_\{t\}\(\\bar\{x\}\_\{t\}\)\-\\widehat\{g\}\(t,\\bar\{x\}\_\{t\}\)\\rVert^\{2\}\\,\\right\]\+\\beta\\,\\mathcal\{L\}\_\{\\mathrm\{dist\}\}\(\\hat\{w\}\),\(19\)w¯\\displaystyle\\bar\{w\}=sg\(w^\),x¯t=Xw¯0,t\(x0\),\\displaystyle=\\mathrm\{sg\}\\left\(\\hat\{w\}\\right\),\\quad\\bar\{x\}\_\{t\}=X^\{\\bar\{w\}\}\_\{0,t\}\(x\_\{0\}\),

wheresg⁡\(⋅\)\\mathrm\{sg\}\\left\(\\cdot\\right\)denotes the stop\-gradient operator,g^​\(t,x\)=∇xV^​\(t,x\)\\widehat\{g\}\(t,x\)=\\nabla\_\{x\}\\widehat\{V\}\(t,x\)is built from the Monte Carlo estimator[Equation18](https://arxiv.org/html/2609.27033#S4.E18), andℒdist\\mathcal\{L\}\_\{\\mathrm\{dist\}\}is any of the self\-distillation losses[Equation25](https://arxiv.org/html/2609.27033#A1.E25)reviewed in[AppendixA](https://arxiv.org/html/2609.27033#A1)\. Herew¯=sg⁡\(w^\)\\bar\{w\}=\\mathrm\{sg\}\\left\(\\hat\{w\}\\right\)is a frozen reference used to form the on\-policy states and value\-gradient targets\. In addition to ensuring that the output of our method is a fine\-tuned flow map, continual self\-distillation ensures that we always have access toamortized on\-policy rolloutsfor value\-function estimation\. At the population level, we callw^\\hat\{w\}a fixed point when[Equation19](https://arxiv.org/html/2609.27033#S4.E19)is stationary inw^\\hat\{w\}withw¯\\bar\{w\}and the stopped targets held fixed, andw¯=w^\\bar\{w\}=\\hat\{w\}\.

###### Proposition 4\.2\(Fixed\-point optimality of WTF\)\.

Under the regularity conditions of[SectionD\.1](https://arxiv.org/html/2609.27033#A4.SS1)and assumingρ0\\rho\_\{0\}has full support onℝd\\mathbb\{R\}^\{d\}, any population fixed point of[Equation19](https://arxiv.org/html/2609.27033#S4.E19)satisfiesw^t,t=bt\+ut∗\\hat\{w\}\_\{t,t\}=b\_\{t\}\+u^\{\*\}\_\{t\}\. Moreover,Xs,tw^X^\{\\hat\{w\}\}\_\{s,t\}is the flow map ofbt\+ut∗b\_\{t\}\+u^\{\*\}\_\{t\}\.

The proof, given in[SectionE\.4](https://arxiv.org/html/2609.27033#A5.SS4), factorizes the fixed\-point condition\. The diagonal regression givesw^t,t=bt\+∇xVtw¯\\hat\{w\}\_\{t,t\}=b\_\{t\}\+\\nabla\_\{x\}V^\{\\bar\{w\}\}\_\{t\}on the on\-policy support, while the off\-diagonal distillation makesXw^X^\{\\hat\{w\}\}the flow map ofw^t,t\\hat\{w\}\_\{t,t\}\. Atw¯=w^\\bar\{w\}=\\hat\{w\},[Proposition4\.1](https://arxiv.org/html/2609.27033#S4.Thmtheorem1)then gives the optimal control\.

[Algorithm1](https://arxiv.org/html/2609.27033#algorithm1)summarizes the training procedure and shows the two choices of reward gradient discussed in[Section6](https://arxiv.org/html/2609.27033#S6)\.

The synthetic experiments use[Algorithm1](https://arxiv.org/html/2609.27033#algorithm1)as written, while the large\-scale image experiments use the Euclidean reward\-gradient variant described in[AppendixF](https://arxiv.org/html/2609.27033#A6)\. Both loss terms are evaluated on every iteration;[SectionG\.1](https://arxiv.org/html/2609.27033#A7.SS1)records the practical choice of frozen referencew¯\\bar\{w\}, which is an exponential moving average ofw^\\hat\{w\}on ImageNet and the stop\-gradient copy of the live weights on text\-to\-image\.

Algorithm 1WTF: Wasserstein\-tilted fine\-tuning of flow mapsInput:Pre\-trained flow map

Xs,t​\(x\)=x\+\(t−s\)​vs,t​\(x\)X\_\{s,t\}\(x\)=x\+\(t\-s\)\\,v\_\{s,t\}\(x\); reward

rr; scale

λ\\lambda; distillation weight

β\\beta; off\-diagonal sampler

ps,tp\_\{s,t\}; learning rate

η\\eta; reward\-gradient type

∈\{exact,Euclidean\}\\in\\\{\\mathrm\{exact\},\\mathrm\{Euclidean\}\\\}
Output:Fine\-tuned flow map

Xs,tw^​\(x\)=x\+\(t−s\)​w^s,t​\(x\)X^\{\\hat\{w\}\}\_\{s,t\}\(x\)=x\+\(t\-s\)\\,\\hat\{w\}\_\{s,t\}\(x\)
Initialize

w^\\hat\{w\}from

vv; set frozen reference

w¯=sg⁡\(w^\)\\bar\{w\}=\\mathrm\{sg\}\\left\(\\hat\{w\}\\right\);//implementation:w¯\\bar\{w\}may be an EMA ofw^\\hat\{w\}, see[SectionG\.1](https://arxiv.org/html/2609.27033#A7.SS1)

repeat

Sample

x0∼ρ0x\_\{0\}\\sim\\rho\_\{0\},

t∼Unif⁡\[0,1\]t\\sim\\mathrm\{Unif\}\[0,1\],

τ∼Unif⁡\[t,1\]\\tau\\sim\\mathrm\{Unif\}\[t,1\];

Forward \(frozen

w¯\\bar\{w\}\):

x¯t=X0,tw¯​\(x0\)\\bar\{x\}\_\{t\}=X^\{\\bar\{w\}\}\_\{0,t\}\(x\_\{0\}\),

x¯τ=Xt,τw¯​\(x¯t\)\\bar\{x\}\_\{\\tau\}=X^\{\\bar\{w\}\}\_\{t,\\tau\}\(\\bar\{x\}\_\{t\}\),

x¯1=Xτ,1w¯​\(x¯τ\)\\bar\{x\}\_\{1\}=X^\{\\bar\{w\}\}\_\{\\tau,1\}\(\\bar\{x\}\_\{\\tau\}\);

Monte Carlo value:

V^=λ​r​\(x¯1\)−1−t2​∥w¯τ,τ​\(x¯τ\)−bτ​\(x¯τ\)∥2\\widehat\{V\}=\\lambda\\,r\(\\bar\{x\}\_\{1\}\)\-\\tfrac\{1\-t\}\{2\}\\,\\lVert\\bar\{w\}\_\{\\tau,\\tau\}\(\\bar\{x\}\_\{\\tau\}\)\-b\_\{\\tau\}\(\\bar\{x\}\_\{\\tau\}\)\\rVert^\{2\};

Value gradient:

g^=∇x¯tV^\\widehat\{g\}=\\nabla\_\{\\bar\{x\}\_\{t\}\}\\widehat\{V\};//autograd through frozenw¯\\bar\{w\}

if*reward\-gradient type==exact*then

g^r←∇Xt,1w¯\(x¯t\)⊤∇r\(x¯1\)\\widehat\{g\}\_\{r\}\\leftarrow\\nabla X^\{\\bar\{w\}\}\_\{t,1\}\(\\bar\{x\}\_\{t\}\)^\{\\top\}\\nabla r\(\\bar\{x\}\_\{1\}\);

else

g^r←∇r​\(x¯1\)\\widehat\{g\}\_\{r\}\\leftarrow\\nabla r\(\\bar\{x\}\_\{1\}\);

Diagonal loss:

ℒval=∥w^t,t​\(x¯t\)−sg⁡\(bt​\(x¯t\)\+g^\)∥2\\mathcal\{L\}\_\{\\mathrm\{val\}\}=\\lVert\\hat\{w\}\_\{t,t\}\(\\bar\{x\}\_\{t\}\)\-\\mathrm\{sg\}\\left\(\\,b\_\{t\}\(\\bar\{x\}\_\{t\}\)\+\\widehat\{g\}\\,\\right\)\\rVert^\{2\};

Sample

\(s′,t′\)∼ps,t\(s^\{\\prime\},t^\{\\prime\}\)\\sim p\_\{s,t\}over the upper triangle; compute

ℒdist\\mathcal\{L\}\_\{\\mathrm\{dist\}\}via[Equation84](https://arxiv.org/html/2609.27033#A6.E84);

Update:

θ←θ−η​∇θ\(ℒval\+β​ℒdist\)\\theta\\leftarrow\\theta\-\\eta\\,\\nabla\_\{\\theta\}\\\!\\left\(\\mathcal\{L\}\_\{\\mathrm\{val\}\}\+\\beta\\,\\mathcal\{L\}\_\{\\mathrm\{dist\}\}\\right\);

until*converged*;

return

Xs,tw^​\(x\)=x\+\(t−s\)​w^s,t​\(x\)X^\{\\hat\{w\}\}\_\{s,t\}\(x\)=x\+\(t\-s\)\\,\\hat\{w\}\_\{s,t\}\(x\);

## Related work

#### Generative modeling via dynamical transport\.

Diffusion and flow models generate samples by transporting a simple prior to the data distribution through learned stochastic or deterministic dynamics\[[28](https://arxiv.org/html/2609.27033#bib.bib20),[29](https://arxiv.org/html/2609.27033#bib.bib19),[30](https://arxiv.org/html/2609.27033#bib.bib7),[31](https://arxiv.org/html/2609.27033#bib.bib6)\]\. Flow matching learns the deterministic dynamics directly, but sampling still requires numerical integration and repeated network evaluations\. Flow maps amortize this computation by learning the solution operator between pairs of times\[[21](https://arxiv.org/html/2609.27033#bib.bib9),[22](https://arxiv.org/html/2609.27033#bib.bib45),[23](https://arxiv.org/html/2609.27033#bib.bib5),[24](https://arxiv.org/html/2609.27033#bib.bib8),[35](https://arxiv.org/html/2609.27033#bib.bib32),[25](https://arxiv.org/html/2609.27033#bib.bib3)\], enabling generation in one or a few evaluations\. Although developed primarily for accelerated inference, flow maps have recently been used for reward alignment\[[17](https://arxiv.org/html/2609.27033#bib.bib44),[18](https://arxiv.org/html/2609.27033#bib.bib46),[48](https://arxiv.org/html/2609.27033#bib.bib54)\]\. We study direct reward fine\-tuning of the flow map, so the fine\-tuned model retains few\-step generation\.

#### Reward fine\-tuning\.

Reward fine\-tuning adapts a pre\-trained generative model to improve a downstream reward while limiting deviation from the base model\. A common formulation is KL\-regularized reward maximization, whose optimizer is the exponential reward tilt\. Adjoint Matching\[[16](https://arxiv.org/html/2609.27033#bib.bib37)\]casts this target as memoryless stochastic optimal control, while RAM\[[49](https://arxiv.org/html/2609.27033#bib.bib71)\]derives a regression objective for the same KL\-regularized target\. Reinforcement\-learning methods including DDPO\[[50](https://arxiv.org/html/2609.27033#bib.bib61)\], DPOK\[[51](https://arxiv.org/html/2609.27033#bib.bib62)\], and Flow\-GRPO\[[52](https://arxiv.org/html/2609.27033#bib.bib68)\]instead optimize reward through the generative sampler\. VGG\-Flow\[[53](https://arxiv.org/html/2609.27033#bib.bib53)\]also derives flow fine\-tuning from deterministic optimal control, but learns a value\-gradient critic and updates the velocity field\. WTF instead regularizes with the transport cost induced by the pre\-trained deterministic drift\. This yields a deterministic control problem in which flow map transitions provide simulation\-free value estimates without a learned critic\. Optimal transport has also appeared in reinforcement learning for generative policies\[[54](https://arxiv.org/html/2609.27033#bib.bib70)\], where transport enters through a critic\-based policy objective\. In WTF, the reward\-independent transport cost is built from the pre\-trained dynamics and serves directly as the regularizer\.

#### Flow maps for reward alignment\.

Flow maps have also begun to play a direct role in reward\-alignment methods\. MFM\[[17](https://arxiv.org/html/2609.27033#bib.bib44)\]uses stochastic flow maps for efficient conditional endpoint sampling and value estimation, but its fine\-tuning procedure returns a velocity field and therefore requires distillation for few\-step deployment\. VFM\[[48](https://arxiv.org/html/2609.27033#bib.bib54)\]learns a noise adapter together with a flow map for one\-step conditional generation, while score\-distillation approaches\[[55](https://arxiv.org/html/2609.27033#bib.bib56),[56](https://arxiv.org/html/2609.27033#bib.bib57)\]regularize one\-step generators toward a base distribution\. These methods differ in how the flow map enters the alignment procedure and in the model retained after fine\-tuning\. WTF directly fine\-tunes the deterministic flow map across time pairs, preserving inference across NFE budgets and compatibility with subsequent flow map alignment\.

#### Inference\-time alignment\.

Complementary to fine\-tuning, inference\-time methods keep model parameters fixed and modify generation through reward guidance, particle\-based sampling, or test\-time optimization\[[19](https://arxiv.org/html/2609.27033#bib.bib39),[57](https://arxiv.org/html/2609.27033#bib.bib34),[58](https://arxiv.org/html/2609.27033#bib.bib36),[59](https://arxiv.org/html/2609.27033#bib.bib11),[60](https://arxiv.org/html/2609.27033#bib.bib12),[61](https://arxiv.org/html/2609.27033#bib.bib1),[62](https://arxiv.org/html/2609.27033#bib.bib4),[63](https://arxiv.org/html/2609.27033#bib.bib13),[64](https://arxiv.org/html/2609.27033#bib.bib2)\]\. FMRG\[[26](https://arxiv.org/html/2609.27033#bib.bib55)\]formulates few\-step, single\-trajectory guidance of flow maps as deterministic optimal control\. Diamond Maps\[[18](https://arxiv.org/html/2609.27033#bib.bib46)\]instead uses stochastic flow maps for value estimation and supports guidance, search, and sequential Monte Carlo\. Because WTF retains a flow map after fine\-tuning, methods such as FMRG can be applied directly for additional inference\-time alignment; velocity\-field fine\-tuning methods require flow map distillation first\.

## Experiments

Figure 4:Comparison of WTF and KL regularization in one dimension\.Each panel places the reward at a different distanceKKfrom the mean of the base law and shows the resulting terminal distributions\. Both terminal laws are computed from their exact population optima\. As the reward moves into lower\-density regions, KL tilting reweights the existing mass, whereas WTF transports mass toward the reward\.To gain intuition for how the Wasserstein\-tilted and KL\-tilted objectives differ, we first compare their population optima in one dimension, where both can be evaluated without training\. We then evaluate trained flow maps on ImageNet\-256 and text\-to\-image generation\. The exact reward gradient differentiates the terminal reward through the flow map transition, giving the pullback∇Xt,1\(x¯t\)⊤∇r\(x¯1\)\\nabla X\_\{t,1\}\(\\bar\{x\}\_\{t\}\)^\{\\top\}\\nabla r\(\\bar\{x\}\_\{1\}\)\. We use this gradient in the synthetic experiments\. For our large\-scale image experiments, we instead use∇r​\(x¯1\)\\nabla r\(\\bar\{x\}\_\{1\}\)directly while leaving the forward endpoint unchanged\. This removes the flow map Jacobian from the reward gradient and computes the update in the direction that increases reward at the generated endpoint\. Following the terminology of FMRG\[[26](https://arxiv.org/html/2609.27033#bib.bib55)\], we refer to the resulting update as the Euclidean reward gradient\. We find that this choice improves both reward and diversity, although the exact gradient also achieves strong results and converges rapidly\. We highlight these choices in[Algorithm1](https://arxiv.org/html/2609.27033#algorithm1)\.[AppendixF](https://arxiv.org/html/2609.27033#A6)gives implementation details, and[SectionH\.2](https://arxiv.org/html/2609.27033#A8.SS2)compares the two empirically\. For ImageNet\-256 we fine\-tune DMF XL/2\[[65](https://arxiv.org/html/2609.27033#bib.bib35)\], and for text\-to\-image we fine\-tune TiM\-T2I\[[27](https://arxiv.org/html/2609.27033#bib.bib67)\]\. We use HPSv2\[[1](https://arxiv.org/html/2609.27033#bib.bib41)\], a learned human\-preference score, as the training reward in both settings\. We report PickScore\[[66](https://arxiv.org/html/2609.27033#bib.bib43)\]and ImageReward\[[10](https://arxiv.org/html/2609.27033#bib.bib42)\]as additional reward metrics\. We measure diversity by the mean pairwise squared distance in DreamSim and CLIP embedding space, which collapses to zero when all samples coincide\. On ImageNet, we compare against Adjoint Matching\[[16](https://arxiv.org/html/2609.27033#bib.bib37)\]and other flow map\-based fine\-tuning methods, MFM\[[17](https://arxiv.org/html/2609.27033#bib.bib44)\]and VFM\[[48](https://arxiv.org/html/2609.27033#bib.bib54)\]\. On text\-to\-image, we compare against Adjoint Matching, a first\-order method, and Flow\-GRPO\[[52](https://arxiv.org/html/2609.27033#bib.bib68)\], a zeroth\-order method\. Because WTF fine\-tunes the full flow map across time pairs, the same checkpoint can be evaluated at different NFE budgets without post\-hoc distillation\. Full training and evaluation details are in[AppendicesF](https://arxiv.org/html/2609.27033#A6)and[G](https://arxiv.org/html/2609.27033#A7)\.

### Synthetic experiments

We first isolate the effect of the regularizer by comparing the WTF and KL population optima in one dimension\. The KL optimumρ~∝ρ1​eλ​r\\tilde\{\\rho\}\\propto\\rho\_\{1\}e^\{\\lambda r\}reweights the base law, while[Proposition3\.1](https://arxiv.org/html/2609.27033#S3.Thmtheorem1)gives the WTF optimum by transporting each sample to its reward\-cost optimum\. We takeρ1=𝒩⁡\(0,1\)\\rho\_\{1\}=\\mathcal\{N\}\(0,1\)and use the Gaussian reward bumprK\(x\)=exp\(−\(x−K\)2/2w2\)r\_\{K\}\(x\)=\\exp\(\-\(x\-K\)^\{2\}/2w^\{2\}\)centered atKK, with widthw=2\.25w=2\.25and reward scaleλ=7\\lambda=7\. Both terminal laws are evaluated without training\. We compute the KL tilt by numerical quadrature\. For WTF, the prior\-action cost is available in closed form, after which we solve the pointwise optimization in[Proposition3\.2](https://arxiv.org/html/2609.27033#S3.Thmtheorem2)and obtain the terminal density by change of variables\. This gives an exact comparison of the two population optima without sampling\.[Figure4](https://arxiv.org/html/2609.27033#S6.F4)shows that WTF achieves higher reward than KL tilting, with the gap increasing as the reward moves into lower\-density regions of the base distribution\. This behavior follows from the different regularizers\. KL tilting can only reweight mass already present under the base law, whereas WTF can transport mass toward high\-reward regions\. Transport is therefore most advantageous when high\-reward regions carry little mass under the base distribution

The two objectives also differ in which terminal laws they can reach\. A tiltν∝ρ1​eλ​r\\nu\\propto\\rho\_\{1\}e^\{\\lambda r\}cannot change the relative density of two points with equal reward: ifr⁡\(xa\)=r⁡\(xb\)r\(x\_\{a\}\)=r\(x\_\{b\}\), both are multiplied by the same factoreλ​r​\(xa\)e^\{\\lambda r\(x\_\{a\}\)\}, so

ν⁡\(xa\)ν⁡\(xb\)=ρ1​\(xa\)ρ1​\(xb\)for every​λ\.\\frac\{\\nu\(x\_\{a\}\)\}\{\\nu\(x\_\{b\}\)\}=\\frac\{\\rho\_\{1\}\(x\_\{a\}\)\}\{\\rho\_\{1\}\(x\_\{b\}\)\}\\qquad\\text\{for every \}\\lambda\.Transport has no such constraint and can move mass between such points, so WTF can reach terminal laws that no tilt ofρ1\\rho\_\{1\}byrrproduces\.

### ImageNet\-256 main results

We fine\-tune DMF XL/2 with HPSv2 as the reward\. WTF results are means over three matched seeds, and every method was given the same fine\-tuning budget of roughly five hours on one8×H1008\\times\\mathrm\{H100\}node\.

Table 5:ImageNet results\.WTF is evaluated at11and250250NFE from the same fine\-tuned flow map\. Results are means over three matched seeds; the base model is DMF XL/2\.Figure 6:Reward against cumulative training compute\.HPSv2 versus GPU\-hours at5050NFE, all methods fine\-tuned from the same flow map\. WTF reaches the peak rewards of Adjoint Matching and Flow\-GRPO with47×47\\timesand280×280\\timesless compute\. Shading is±1\\pm 1standard error; protocol in[SectionG\.4](https://arxiv.org/html/2609.27033#A7.SS4)\.[Table5](https://arxiv.org/html/2609.27033#S6.T5)compares WTF with Adjoint Matching, MFM, and VFM\. WTF achieves higher reward while maintaining competitive diversity\. The same checkpoint also remains effective from one to 250 NFE\. MFM fine\-tunes the diagonal velocity field, while VFM targets the one\-step sampler, so neither directly provides the same arbitrary\-budget flow map\.[SectionH\.1](https://arxiv.org/html/2609.27033#A8.SS1)shows matched\-latent samples\.

### Text\-to\-image main results

We fine\-tune the TiM\-T2I flow map\[[27](https://arxiv.org/html/2609.27033#bib.bib67)\]against HPSv2\. Every method was fine\-tuned under the same budget of roughly4848GPU\-hours on one8×H1008\\times\\mathrm\{H100\}node\.[SectionG\.1](https://arxiv.org/html/2609.27033#A7.SS1)gives the full setup\.[Table7](https://arxiv.org/html/2609.27033#S6.T7)evaluates each released baseline checkpoint in our harness\.

Table 7:Text\-to\-image results\.Baseline checkpoints are evaluated in the same harness, and WTF is evaluated across inference budgets from a single fine\-tuned flow map\. At 50 NFE, WTF attains the highest HPSv2 \(the training reward\), PickScore and ImageReward of any method in the table\.[Table7](https://arxiv.org/html/2609.27033#S6.T7)compares WTF with released Adjoint Matching and Flow\-GRPO checkpoints evaluated in the same harness\. As on ImageNet, WTF achieves higher reward while maintaining comparable diversity, and the same fine\-tuned flow map remains effective in the few\-step regime\.[Figure6](https://arxiv.org/html/2609.27033#S6.F6)shows that WTF also converges substantially faster, reaching the peak rewards attained by Flow\-GRPO and Adjoint Matching with up to280×280\\timesand47×47\\timesless training compute, respectively\.[SectionH\.3](https://arxiv.org/html/2609.27033#A8.SS3)variesλ\\lambdato characterize the reward\-diversity tradeoff, and[SectionG\.4](https://arxiv.org/html/2609.27033#A7.SS4)gives the compute comparison\.

## Conclusion

We introduced a transport\-regularized formulation of reward fine\-tuning for deterministic flows\. The objective transports samples toward higher reward rather than reweighting the base distribution and is equivalent to adeterministic optimal control problem\. A pre\-trained flow map makes value estimation simulation\-free and critic\-free, yielding direct end\-to\-end fine\-tuning of the flow map in which the same few\-step sampler is used both at training and at deployment\. On ImageNet\-256 and text\-to\-image generation, WTF matches or surpasses baselines on reward, reaching the peak rewards of the comparison methods with substantially less training compute\. The fine\-tuned model remains compatible with flow map inference\-time alignment without an intervening distillation stage\. We hope that our framework spurs broader interest in the use of flow maps as a foundational primitive for accelerated post\-training of generative models\.

#### Limitations\.

[Proposition4\.2](https://arxiv.org/html/2609.27033#S4.Thmtheorem2)is a population\-level fixed\-point result\. The large\-scale experiments use the Euclidean reward gradient described in[AppendixF](https://arxiv.org/html/2609.27033#A6)\. Increasing reward can reduce sample diversity, with the reward scale controlling this tradeoff\. WTF also assumes access to a pre\-trained flow map\. Starting from a velocity model requires a preceding flow map distillation stage, whose cost is not included in our fine\-tuning comparison\.

## Acknowledgements

AM is supported by the Clarendon Fund Scholarship, University of Oxford\. We gratefully acknowledge fal for providing the computational resources that enabled this work\. We also thank Modal for additional compute support\. The authors acknowledge the use of resources provided by the Isambard\-AI National AI Research Resource \(AIRR\)\[[67](https://arxiv.org/html/2609.27033#bib.bib40)\]\. Isambard\-AI is operated by the University of Bristol and is funded by the UK Government’s Department for Science, Innovation and Technology \(DSIT\) via UK Research and Innovation; and the Science and Technology Facilities Council \[ST/AIRR/I\-A\-I/1023\]\.

## References

- \[1\]X\. Wu, Y\. Hao, K\. Sun, Y\. Chen, F\. Zhu, R\. Zhao, and H\. Li\(2023\)Human preference score v2: a solid benchmark for evaluating human preferences of text\-to\-image synthesis\.External Links:2306\.09341Cited by:[Figure 1](https://arxiv.org/html/2609.27033#S0.F1),[Figure 1](https://arxiv.org/html/2609.27033#S0.F1.5.1),[§1](https://arxiv.org/html/2609.27033#S1.p1.1),[§6](https://arxiv.org/html/2609.27033#S6.p1.1)\.
- \[2\]R\. Rombach, A\. Blattmann, D\. Lorenz, P\. Esser, and B\. Ommer\(2022\)High\-resolution image synthesis with latent diffusion models\.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp\. 10684–10695\.Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1)\.
- \[3\]P\. Esser, S\. Kulal, A\. Blattmann, R\. Entezari, J\. Müller, H\. Saini, Y\. Levi, D\. Lorenz, A\. Sauer, F\. Boesel, D\. Podell, T\. Dockhorn, Z\. English, K\. Lacey, A\. Goodwin, Y\. Marek, and R\. Rombach\(2024\)Scaling rectified flow transformers for high\-resolution image synthesis\.InInternational Conference on Machine Learning,External Links:2403\.03206Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1)\.
- \[4\]A\. Blattmann, T\. Dockhorn, S\. Kulal, D\. Mendelevitch, M\. Kilian, D\. Lorenz, Y\. Levi, Z\. English, V\. Voleti, A\. Letts,et al\.\(2023\)Stable video diffusion: scaling latent video diffusion models to large datasets\.arXiv preprint arXiv:2311\.15127\.Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1)\.
- \[5\]J\. Abramson, J\. Adler, J\. Dunger, R\. Evans, T\. Green, A\. Pritzel, O\. Ronneberger, L\. Willmore, A\. J\. Ballard, J\. Bambrick, S\. W\. Bodenstein, D\. A\. Evans, C\. Hung, M\. O’Neill, D\. Reiman, K\. Tunyasuvunakool, Z\. Wu, A\. Žemgulytė, E\. Arvaniti, C\. Beattie, O\. Bertolli, A\. Bridgland, A\. Cherepanov, M\. Congreve, A\. I\. Cowen\-Rivers, A\. Cowie, M\. Figurnov, F\. B\. Fuchs, H\. Gladman, R\. Jain, Y\. A\. Khan, C\. M\. R\. Low, K\. Perlin, A\. Potapenko, P\. Savy, S\. Singh, A\. Stecula, A\. Thillaisundaram, C\. Tong, S\. Yakneen, E\. D\. Zhong, M\. Zielinski, A\. Žídek, V\. Bapst, P\. Kohli, M\. Jaderberg, D\. Hassabis, and J\. M\. Jumper\(2024\)Accurate structure prediction of biomolecular interactions with alphafold 3\.Nature630\(8016\),pp\. 493–500\.Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1)\.
- \[6\]C\. Zeni, R\. Pinsler, D\. Zügner, A\. Fowler, M\. Horton, X\. Fu, Z\. Wang, A\. Shysheya, J\. Crabbé, S\. Ueda, R\. Sordillo, L\. Sun, J\. Smith, B\. Nguyen, H\. Schulz, S\. Lewis, C\. Huang, Z\. Lu, Y\. Zhou, H\. Yang, H\. Hao, J\. Li, C\. Yang, W\. Li, R\. Tomioka, and T\. Xie\(2025\)A generative model for inorganic materials design\.Nature639\(8055\),pp\. 624–632\.Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1)\.
- \[7\]C\. Lee, J\. Yoo, M\. Agarwal, S\. Shah, J\. Huang, A\. Raghunathan, S\. Hong, N\. M\. Boffi, and J\. Kim\(2026\)Flow map language models: one\-step language modeling via continuous denoising\.arXiv preprint arXiv:2602\.16813\.Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1)\.
- \[8\]D\. Roos, O\. Davis, F\. Eijkelboom, M\. Bronstein, M\. Welling, İ\. İ\. Ceylan, L\. Ambrogioni, and J\. van de Meent\(2026\)Categorical flow maps\.arXiv preprint arXiv:2602\.12233\.Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1)\.
- \[9\]P\. Potaptchik, J\. Yim, A\. Saravanan, P\. Holderrieth, E\. Vanden\-Eijnden, and M\. S\. Albergo\(2026\)Discrete flow maps\.arXiv preprint arXiv:2604\.09784\.Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1)\.
- \[10\]J\. Xu, X\. Liu, Y\. Wu, Y\. Tong, Q\. Li, M\. Ding, J\. Tang, and Y\. Dong\(2023\)ImageReward: learning and evaluating human preferences for text\-to\-image generation\.InAdvances in Neural Information Processing Systems,External Links:2304\.05977Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1),[§6](https://arxiv.org/html/2609.27033#S6.p1.1)\.
- \[11\]L\. Ouyanget al\.\(2022\)Training language models to follow instructions with human feedback\.Advances in Neural Information Processing Systems35\.External Links:2203\.02155Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1)\.
- \[12\]L\. Gao, J\. Schulman, and J\. Hilton\(2023\)Scaling laws for reward model overoptimization\.InInternational Conference on Machine Learning,pp\. 10835–10866\.External Links:2210\.10760Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1)\.
- \[13\]J\. L\. Watson, D\. Juergens, N\. R\. Bennett, B\. L\. Trippe, J\. Yim, H\. E\. Eisenach, W\. Ahern, A\. J\. Borst, R\. J\. Ragotte, L\. F\. Milles, B\. I\. M\. Wicky, N\. Hanikel, S\. J\. Pellock, A\. Courbet, W\. Sheffler, J\. Wang, P\. Venkatesh, I\. Sappington, S\. V\. Torres, A\. Lauko, V\. De Bortoli, E\. Mathieu, S\. Ovchinnikov, R\. Barzilay, T\. S\. Jaakkola, F\. DiMaio, M\. Baek, and D\. Baker\(2023\)De novo design of protein structure and function with rfdiffusion\.Nature620\(7976\),pp\. 1089–1100\.Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1)\.
- \[14\]G\. Corso, H\. Stärk, B\. Jing, R\. Barzilay, and T\. Jaakkola\(2023\)DiffDock: diffusion steps, twists, and turns for molecular docking\.InInternational Conference on Learning Representations,External Links:2210\.01776Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p1.1)\.
- \[15\]K\. Clark, P\. Vicol, K\. Swersky, and D\. J\. Fleet\(2024\)Directly fine\-tuning diffusion models on differentiable rewards\.InInternational Conference on Learning Representations,External Links:2309\.17400Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p2.1),[§2\.2](https://arxiv.org/html/2609.27033#S2.SS2.p1.2)\.
- \[16\]C\. Domingo\-Enrich, M\. Drozdzal, B\. Karrer, and R\. T\. Q\. Chen\(2025\)Adjoint matching: fine\-tuning flow and diffusion generative models with memoryless stochastic optimal control\.InInternational Conference on Learning Representations,External Links:2409\.08861Cited by:[Appendix B](https://arxiv.org/html/2609.27033#A2.SS0.SSS0.Px2.p1.6),[Appendix B](https://arxiv.org/html/2609.27033#A2.p1.1),[§G\.3](https://arxiv.org/html/2609.27033#A7.SS3.p3.1),[§1](https://arxiv.org/html/2609.27033#S1.p2.1),[§2\.2](https://arxiv.org/html/2609.27033#S2.SS2.p1.2),[§2\.2](https://arxiv.org/html/2609.27033#S2.SS2.p2.3),[§3\.1](https://arxiv.org/html/2609.27033#S3.SS1.p1.2),[§3\.2](https://arxiv.org/html/2609.27033#S3.SS2.SSS0.Px1.p3.1),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px2.p1.1),[§6](https://arxiv.org/html/2609.27033#S6.p1.1)\.
- \[17\]P\. Potaptchik, A\. Saravanan, A\. Mammadov, A\. Prat, M\. S\. Albergo, and Y\. W\. Teh\(2026\)Meta flow maps enable scalable reward alignment\.External Links:2601\.14430Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p2.1),[§2\.2](https://arxiv.org/html/2609.27033#S2.SS2.p1.2),[§3\.2](https://arxiv.org/html/2609.27033#S3.SS2.SSS0.Px1.p3.1),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2609.27033#S6.p1.1)\.
- \[18\]P\. Holderrieth, D\. Chen, L\. Eyring, I\. Shah, G\. Anantharaman, Y\. He, Z\. Akata, T\. Jaakkola, N\. M\. Boffi, and M\. Simchowitz\(2026\)Diamond maps: efficient reward alignment via stochastic flow maps\.External Links:2602\.05993Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p2.1),[§2\.2](https://arxiv.org/html/2609.27033#S2.SS2.p1.2),[§3\.2](https://arxiv.org/html/2609.27033#S3.SS2.SSS0.Px1.p3.1),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px4.p1.1)\.
- \[19\]M\. Uehara, Y\. Zhao, C\. Wang, X\. Li, A\. Regev, S\. Levine, and T\. Biancalani\(2025\)Inference\-time alignment in diffusion models with reward\-guided generation: tutorial and review\.External Links:2501\.09685Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p2.1),[§2\.2](https://arxiv.org/html/2609.27033#S2.SS2.p1.2),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px4.p1.1)\.
- \[20\]I\. Karatzas and S\. Shreve\(2014\)Brownian motion and stochastic calculus\.Vol\.113,springer\.Cited by:[Appendix B](https://arxiv.org/html/2609.27033#A2.SS0.SSS0.Px1.p1.4),[§1](https://arxiv.org/html/2609.27033#S1.p2.1),[§2\.2](https://arxiv.org/html/2609.27033#S2.SS2.p2.1)\.
- \[21\]N\. M\. Boffi, M\. S\. Albergo, and E\. Vanden\-Eijnden\(2024\)Flow Map Matching with Stochastic Interpolants: a mathematical framework for consistency models\.arXiv:2406\.07507\.Cited by:[Appendix F](https://arxiv.org/html/2609.27033#A6.SS0.SSS0.Px3.p1.1),[§1](https://arxiv.org/html/2609.27033#S1.p3.1),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p2.1),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p2.2),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1)\.
- \[22\]N\. M\. Boffi, M\. S\. Albergo, and E\. Vanden\-Eijnden\(2025\)How to build a consistency model: learning flow maps via self\-distillation\.Advances in Neural Information Processing Systems38\.Cited by:[Appendix F](https://arxiv.org/html/2609.27033#A6.SS0.SSS0.Px3.p1.1),[§1](https://arxiv.org/html/2609.27033#S1.p3.1),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p2.1),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p2.2),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1)\.
- \[23\]Y\. Song, P\. Dhariwal, M\. Chen, and I\. Sutskever\(2023\)Consistency models\.InInternational Conference on Machine Learning,External Links:2303\.01469Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p3.1),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p2.1),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1)\.
- \[24\]D\. Kim, C\. Lai, W\. Liao, N\. Murata, Y\. Takida, T\. Uesaka, Y\. He, Y\. Mitsufuji, and S\. Ermon\(2024\)Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion\.InInternational Conference on Learning Representations,External Links:2310\.02279Cited by:[§1](https://arxiv.org/html/2609.27033#S1.p3.1),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p2.1),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1)\.
- \[25\]Z\. Geng, M\. Deng, X\. Bai, J\. Z\. Kolter, and K\. He\(2025\)Mean flows for one\-step generative modeling\.External Links:2505\.13447Cited by:[Appendix F](https://arxiv.org/html/2609.27033#A6.SS0.SSS0.Px3.p1.1),[§1](https://arxiv.org/html/2609.27033#S1.p3.1),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p2.1),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p2.2),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1)\.
- \[26\]J\. Y\. Huang, J\. Lin, S\. Shah, K\. Nair, and N\. M\. Boffi\(2026\)How to guide your flow: few\-step alignment via flow map reward guidance\.External Links:2604\.27147Cited by:[Appendix F](https://arxiv.org/html/2609.27033#A6.SS0.SSS0.Px1.p1.1),[§1](https://arxiv.org/html/2609.27033#S1.p3.1),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px4.p1.1),[§6](https://arxiv.org/html/2609.27033#S6.p1.1)\.
- \[27\]Z\. Wang, Y\. Zhang, X\. Yue, X\. Yue, Y\. Li, W\. Ouyang, and L\. Bai\(2025\)Transition models: rethinking the generative learning objective\.arXiv preprint arXiv:2509\.04394\.Cited by:[§G\.1](https://arxiv.org/html/2609.27033#A7.SS1.SSS0.Px2.p1.1),[item 4](https://arxiv.org/html/2609.27033#S1.I1.i4.p1.1),[§6\.3](https://arxiv.org/html/2609.27033#S6.SS3.p1.1),[§6](https://arxiv.org/html/2609.27033#S6.p1.1)\.
- \[28\]Y\. Lipman, R\. T\. Chen, H\. Ben\-Hamu, M\. Nickel, and M\. Le\(2023\)Flow matching for generative modeling\.InThe Eleventh International Conference on Learning Representations,External Links:2210\.02747Cited by:[Appendix A](https://arxiv.org/html/2609.27033#A1.SS0.SSS0.Px1.p1.2),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p1.2),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1)\.
- \[29\]X\. Liu, C\. Gong, and Q\. Liu\(2023\)Flow straight and fast: learning to generate and transfer data with rectified flow\.InThe Eleventh International Conference on Learning Representations,External Links:2209\.03003Cited by:[Appendix A](https://arxiv.org/html/2609.27033#A1.SS0.SSS0.Px1.p1.2),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p1.2),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1)\.
- \[30\]M\. S\. Albergo, N\. M\. Boffi, and E\. Vanden\-Eijnden\(2023\)Stochastic interpolants: a unifying framework for flows and diffusions\.arXiv preprint arXiv:2303\.08797\.Cited by:[Appendix A](https://arxiv.org/html/2609.27033#A1.SS0.SSS0.Px1.p1.1),[Appendix A](https://arxiv.org/html/2609.27033#A1.SS0.SSS0.Px1.p1.2),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p1.2),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1)\.
- \[31\]M\. S\. Albergo and E\. Vanden\-Eijnden\(2023\)Building normalizing flows with stochastic interpolants\.InThe Eleventh International Conference on Learning Representations,External Links:2209\.15571Cited by:[Appendix A](https://arxiv.org/html/2609.27033#A1.SS0.SSS0.Px1.p1.1),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p1.2),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1)\.
- \[32\]Y\. Song, J\. Sohl\-Dickstein, D\. P\. Kingma, A\. Kumar, S\. Ermon, and B\. Poole\(2020\)Score\-based generative modeling through stochastic differential equations\.arXiv:2011\.13456\.Cited by:[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p1.2)\.
- \[33\]N\. M\. Boffi and E\. Vanden\-Eijnden\(2023\)Probability flow solution of the Fokker–Planck equation\.Machine Learning: Science and Technology4\(3\),pp\. 035012\.Cited by:[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p1.2)\.
- \[34\]T\. Karras, M\. Aittala, T\. Aila, and S\. Laine\(2022\)Elucidating the Design Space of Diffusion\-Based Generative Models\.arXiv:2206\.00364\.Cited by:[§F\.1](https://arxiv.org/html/2609.27033#A6.SS1.p1.1),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p1.2)\.
- \[35\]K\. Frans, D\. Hafner, S\. Levine, and P\. Abbeel\(2025\)One step diffusion via shortcut models\.InInternational Conference on Learning Representations,External Links:2410\.12557Cited by:[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p2.1),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1)\.
- \[36\]A\. Sabour, S\. Fidler, and K\. Kreis\(2025\)Align your flow: scaling continuous\-time flow map distillation\.External Links:2506\.14603Cited by:[Appendix F](https://arxiv.org/html/2609.27033#A6.SS0.SSS0.Px3.p1.1),[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p2.2)\.
- \[37\]L\. Zhou, M\. Parger, A\. Haque, and J\. Song\(2026\)Terminal velocity matching\.InInternational Conference on Learning Representations,External Links:2511\.19797Cited by:[§2\.1](https://arxiv.org/html/2609.27033#S2.SS1.p2.2)\.
- \[38\]P\. Holderrieth, U\. Singer, T\. Jaakkola, R\. T\. Q\. Chen, Y\. Lipman, and B\. Karrer\(2025\)GLASS flows: transition sampling for alignment of flow and diffusion models\.External Links:2509\.25170Cited by:[§2\.2](https://arxiv.org/html/2609.27033#S2.SS2.p1.2),[§3\.2](https://arxiv.org/html/2609.27033#S3.SS2.SSS0.Px1.p3.1)\.
- \[39\]J\. L\. Doob\(1984\)Classical potential theory and its probabilistic counterpart\.Grundlehren der mathematischen Wissenschaften ; 262,Springer,New York ;\.Cited by:[§2\.2](https://arxiv.org/html/2609.27033#S2.SS2.p2.1)\.
- \[40\]C\. Villani\(2009\)Optimal transport: old and new\.Vol\.338,Springer\.Cited by:[§D\.1](https://arxiv.org/html/2609.27033#A4.SS1.p2.1),[§3\.1](https://arxiv.org/html/2609.27033#S3.SS1.SSS0.Px1.p1.1),[§3\.1](https://arxiv.org/html/2609.27033#S3.SS1.p3.1),[§3\.2](https://arxiv.org/html/2609.27033#S3.SS2.p2.1)\.
- \[41\]J\. Benamou and Y\. Brenier\(2000\)A computational fluid mechanics solution to the monge\-kantorovich mass transfer problem\.Numerische Mathematik84\(3\),pp\. 375–393\.Cited by:[§3\.1](https://arxiv.org/html/2609.27033#S3.SS1.SSS0.Px1.p4.1)\.
- \[42\]Y\. Chen, T\. T\. Georgiou, and M\. Pavon\(2017\)Optimal transport over a linear dynamical system\.IEEE Transactions on Automatic Control62\(5\),pp\. 2137–2152\.Note:arXiv:1502\.01265Cited by:[§E\.2](https://arxiv.org/html/2609.27033#A5.SS2.p1.1),[§3\.1](https://arxiv.org/html/2609.27033#S3.SS1.SSS0.Px1.p4.1),[§3\.2](https://arxiv.org/html/2609.27033#S3.SS2.SSS0.Px1.p3.1)\.
- \[43\]Y\. Chen, T\. T\. Georgiou, and M\. Pavon\(2016\)On the relation between optimal transport and Schrödinger bridges: a stochastic control viewpoint\.Journal of Optimization Theory and Applications169\(2\),pp\. 671–691\.Note:arXiv:1412\.4430Cited by:[§D\.2](https://arxiv.org/html/2609.27033#A4.SS2.p5.1),[§E\.2](https://arxiv.org/html/2609.27033#A5.SS2.p1.1),[§3\.1](https://arxiv.org/html/2609.27033#S3.SS1.SSS0.Px1.p4.1),[§3\.2](https://arxiv.org/html/2609.27033#S3.SS2.SSS0.Px1.p3.1)\.
- \[44\]Y\. Chen, T\. T\. Georgiou, and M\. Pavon\(2021\)Stochastic control liaisons: richard sinkhorn meets gaspard monge on a schrodinger bridge\.Siam Review63\(2\),pp\. 249–313\.Cited by:[§D\.2](https://arxiv.org/html/2609.27033#A4.SS2.p5.1),[§E\.2](https://arxiv.org/html/2609.27033#A5.SS2.p1.1),[§3\.1](https://arxiv.org/html/2609.27033#S3.SS1.SSS0.Px1.p4.1),[§3\.2](https://arxiv.org/html/2609.27033#S3.SS2.SSS0.Px1.p3.1)\.
- \[45\]F\. Santambrogio\(2015\)Optimal transport for applied mathematicians\.Birkäuser, NY55\(58\-63\),pp\. 94\.Cited by:[§D\.1](https://arxiv.org/html/2609.27033#A4.SS1.p2.1),[§3\.2](https://arxiv.org/html/2609.27033#S3.SS2.p2.1)\.
- \[46\]C\. Léonard\(2014\)A survey of the schrödinger problem and some of its connections with optimal transport\.Discrete and Continuous Dynamical Systems\-Series A34\(4\),pp\. 1533–1574\.Cited by:[§3\.2](https://arxiv.org/html/2609.27033#S3.SS2.SSS0.Px1.p3.1)\.
- \[47\]D\. Liberzon\(2011\)Calculus of variations and optimal control theory: a concise introduction\.Princeton University Press\.Cited by:[Appendix C](https://arxiv.org/html/2609.27033#A3.p1.1),[§D\.1](https://arxiv.org/html/2609.27033#A4.SS1.p2.1),[§4](https://arxiv.org/html/2609.27033#S4.p1.2)\.
- \[48\]A\. Mammadov, S\. Takao, B\. Chen, R\. Baptista, M\. Mardani, Y\. W\. Teh, and J\. Berner\(2026\)Variational flow maps: make some noise for one\-step conditional generation\.External Links:2603\.07276Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px1.p1.1),[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px3.p1.1),[§6](https://arxiv.org/html/2609.27033#S6.p1.1)\.
- \[49\]A\. Bergmeister, S\. Jegelka, N\. Nüsken, C\. Domingo\-Enrich, and J\. Pidstrigach\(2026\)Reinforce adjoint matching: scaling rl post\-training of diffusion and flow\-matching models\.arXiv preprint arXiv:2605\.10759\.Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px2.p1.1)\.
- \[50\]K\. Black, M\. Janner, Y\. Du, I\. Kostrikov, and S\. Levine\(2024\)Training diffusion models with reinforcement learning\.InInternational Conference on Learning Representations,External Links:2305\.13301Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px2.p1.1)\.
- \[51\]Y\. Fan, O\. Watkins, Y\. Du, H\. Liu, M\. Ryu, C\. Boutilier, P\. Abbeel, M\. Ghavamzadeh, K\. Lee, and K\. Lee\(2023\)DPOK: reinforcement learning for fine\-tuning text\-to\-image diffusion models\.InAdvances in Neural Information Processing Systems,External Links:2305\.16381Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px2.p1.1)\.
- \[52\]J\. Liu, G\. Liu, J\. Liang, Y\. Li, J\. Liu, X\. Wang, P\. Wan, D\. Zhang, and W\. Ouyang\(2025\)Flow\-grpo: training flow matching models via online rl\.arXiv preprint arXiv:2505\.05470\.Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px2.p1.1),[§6](https://arxiv.org/html/2609.27033#S6.p1.1)\.
- \[53\]Z\. Liu, T\. Z\. Xiao, C\. Domingo\-Enrich, W\. Liu, and D\. Zhang\(2025\)Value gradient guidance for flow matching alignment\.Note:VGG\-FlowCited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px2.p1.1)\.
- \[54\]M\. Sun, P\. Ding, W\. Zhang, and D\. Wang\(2025\)Score\-based diffusion policy compatible with reinforcement learning via optimal transport\.InProceedings of the 42nd International Conference on Machine Learning \(ICML\),Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px2.p1.1)\.
- \[55\]N\. Kumari, S\. Wang, N\. Zhao, Y\. Nitzan, Y\. Li, K\. K\. Singh, R\. Zhang, E\. Shechtman, J\. Zhu, and X\. Huang\(2026\)Learning an image editing model without image editing pairs\.InICLR,Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px3.p1.1)\.
- \[56\]G\. Chen, S\. Huang, K\. Liu, J\. Zhu, X\. Qu, P\. Chen, Y\. Cheng, and Y\. Sun\(2025\)Flash\-dmd: towards high\-fidelity few\-step image generation with efficient distillation and joint reinforcement learning\.External Links:2511\.20549Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px3.p1.1)\.
- \[57\]P\. Dhariwal and A\. Nichol\(2021\)Diffusion models beat gans on image synthesis\.InAdvances in Neural Information Processing Systems,External Links:2105\.05233Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px4.p1.1)\.
- \[58\]J\. Ho and T\. Salimans\(2022\)Classifier\-free diffusion guidance\.arXiv preprint arXiv:2207\.12598\.Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px4.p1.1)\.
- \[59\]H\. Ye, H\. Lin, J\. Han, M\. Xu, S\. Liu, Y\. Liang, J\. Ma, J\. Zou, and S\. Ermon\(2024\)TFG: unified training\-free guidance for diffusion models\.InAdvances in Neural Information Processing Systems,External Links:2409\.15761Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px4.p1.1)\.
- \[60\]J\. Yu, Y\. Wang, C\. Zhao, B\. Ghanem, and J\. Zhang\(2023\)FreeDoM: training\-free energy\-guided conditional diffusion model\.InProceedings of the IEEE/CVF International Conference on Computer Vision,External Links:2303\.09833Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px4.p1.1)\.
- \[61\]H\. Chung, J\. Kim, M\. T\. Mccann, M\. L\. Klasky, and J\. C\. Ye\(2023\)Diffusion posterior sampling for general noisy inverse problems\.InInternational Conference on Learning Representations,External Links:2209\.14687Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px4.p1.1)\.
- \[62\]J\. Kim, B\. S\. Kim, and J\. C\. Ye\(2025\)FlowDPS: flow\-driven posterior sampling for inverse problems\.arXiv preprint arXiv:2503\.08136\.Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px4.p1.1)\.
- \[63\]M\. Skreta, T\. Akhound\-Sadegh, V\. Ohanesian, R\. Bondesan, A\. Aspuru\-Guzik, A\. Doucet, R\. Brekelmans, A\. Tong, and K\. Neklyudov\(2025\)Feynman\-kac correctors in diffusion: annealing, guidance, and product of experts\.InInternational Conference on Machine Learning,External Links:2503\.02819Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px4.p1.1)\.
- \[64\]R\. Singhal, Z\. Horvitz, R\. Teehan, M\. Ren, Z\. Yu, K\. McKeown, and R\. Ranganath\(2025\)A general framework for inference\-time scaling and steering of diffusion models\.InInternational Conference on Machine Learning,External Links:2501\.06848Cited by:[§5](https://arxiv.org/html/2609.27033#S5.SS0.SSS0.Px4.p1.1)\.
- \[65\]K\. Lee, S\. Yu, and J\. Shin\(2025\)Decoupled meanflow: turning flow models into flow maps for accelerated sampling\.arXiv preprint arXiv:2510\.24474\.Cited by:[§G\.1](https://arxiv.org/html/2609.27033#A7.SS1.SSS0.Px1.p1.1),[§6](https://arxiv.org/html/2609.27033#S6.p1.1)\.
- \[66\]Y\. Kirstain, A\. Polyak, U\. Singer, S\. Matiana, J\. Penna, and O\. Levy\(2023\)Pick\-a\-pic: an open dataset of user preferences for text\-to\-image generation\.InAdvances in Neural Information Processing Systems,External Links:2305\.01569Cited by:[§6](https://arxiv.org/html/2609.27033#S6.p1.1)\.
- \[67\]S\. McIntosh\-Smith, S\. R\. Alam, and C\. Woods\(2024\)Isambard\-ai: a leadership class supercomputer optimised specifically for artificial intelligence\.Cited by:[Acknowledgements](https://arxiv.org/html/2609.27033#Sx1.p1.1)\.
- \[68\]W\. H\. Fleming and R\. W\. Rishel\(1975\)Deterministic and stochastic optimal control\.Springer,New York, NY\.Cited by:[Appendix C](https://arxiv.org/html/2609.27033#A3.p1.1),[§D\.1](https://arxiv.org/html/2609.27033#A4.SS1.p2.1)\.
- \[69\]M\. Bardi and I\. Capuzzo\-Dolcetta\(1997\)Optimal control and viscosity solutions of Hamilton–Jacobi–Bellman equations\.Birkhäuser,Boston, MA\.Cited by:[Appendix C](https://arxiv.org/html/2609.27033#A3.p1.1),[§D\.1](https://arxiv.org/html/2609.27033#A4.SS1.p2.1)\.
- \[70\]L\. C\. Evans\(2010\)Partial differential equations\.2 edition,Graduate Studies in Mathematics, Vol\.19,American Mathematical Society,Providence, RI\.Cited by:[§D\.1](https://arxiv.org/html/2609.27033#A4.SS1.p2.1)\.
- \[71\]X\. Wu, Y\. Hao, M\. Zhang, K\. Sun, Z\. Huang, G\. Song, Y\. Liu, and H\. Li\(2024\)Deep reward supervisions for tuning text\-to\-image diffusion models\.InEuropean Conference on Computer Vision,Cited by:[Appendix F](https://arxiv.org/html/2609.27033#A6.SS0.SSS0.Px1.p1.2)\.
- \[72\]J\. Chen, H\. Cai, J\. Chen, E\. Xie, S\. Yang, H\. Tang, M\. Li, Y\. Lu, and S\. Han\(2024\)Deep compression autoencoder for efficient high\-resolution diffusion models\.arXiv preprint arXiv:2410\.10733\.Cited by:[§G\.1](https://arxiv.org/html/2609.27033#A7.SS1.SSS0.Px2.p1.1)\.
- \[73\]Gemma Team\(2025\)Gemma 3 technical report\.arXiv preprint arXiv:2503\.19786\.Cited by:[§G\.1](https://arxiv.org/html/2609.27033#A7.SS1.SSS0.Px2.p1.1)\.
- \[74\]G\. Ilharco, M\. Wortsman, R\. Wightman, C\. Gordon, N\. Carlini, R\. Taori, A\. Dave, V\. Shankar, H\. Namkoong, J\. Miller, H\. Hajishirzi, A\. Farhadi, and L\. Schmidt\(2021\)OpenCLIP\.InZenodo,External Links:[Document](https://dx.doi.org/10.5281/zenodo.5143773)Cited by:[§G\.1](https://arxiv.org/html/2609.27033#A7.SS1.SSS0.Px4.p1.1)\.

## Appendix ABackground on flow\-based generative models

In this section, we provide a self\-contained introduction to flow matching via stochastic interpolants and their self\-distillation into flow maps\.

#### Stochastic interpolants and flow matching\.

Given a dataset of samples\{x1i\}i=1n∼ρ1\\\{x\_\{1\}^\{i\}\\\}\_\{i=1\}^\{n\}\\sim\\rho\_\{1\}and a base distributionρ0=𝒩⁡\(0,I\)\\rho\_\{0\}=\\mathcal\{N\}\(0,I\), the stochastic\-interpolant framework\[[31](https://arxiv.org/html/2609.27033#bib.bib6),[30](https://arxiv.org/html/2609.27033#bib.bib7)\]introduces the time\-dependent random variable

It=αt​x0\+βt​x1,x0∼ρ0,x1∼ρ1,I\_\{t\}=\\alpha\_\{t\}\\,x\_\{0\}\+\\beta\_\{t\}\\,x\_\{1\},\\qquad x\_\{0\}\\sim\\rho\_\{0\},\\ x\_\{1\}\\sim\\rho\_\{1\},\(20\)withα,β:\[0,1\]→\[0,1\]\\alpha,\\beta:\[0,1\]\\to\[0,1\]satisfyingα0=β1=1\\alpha\_\{0\}=\\beta\_\{1\}=1andα1=β0=0\\alpha\_\{1\}=\\beta\_\{0\}=0\. A standard result is thatLaw⁡\(It\)=Law⁡\(xt\)\\mathrm\{Law\}\(I\_\{t\}\)=\\mathrm\{Law\}\(x\_\{t\}\)wherextx\_\{t\}solves the probability flow[Equation1](https://arxiv.org/html/2609.27033#S2.E1)with velocitybt​\(x\)=𝔼⁡\[I˙t∣It=x\]b\_\{t\}\(x\)=\\mathbb\{E\}\[\\dot\{I\}\_\{t\}\\mid I\_\{t\}=x\]\. This conditional expectation is learned by minimizing the flow\-matching objective\[[28](https://arxiv.org/html/2609.27033#bib.bib20),[29](https://arxiv.org/html/2609.27033#bib.bib19),[30](https://arxiv.org/html/2609.27033#bib.bib7)\]

ℒb​\(b^\)=𝔼t,x0,x1​\[∥b^t​\(It\)−I˙t∥2\]\\mathcal\{L\}\_\{b\}\(\\hat\{b\}\)=\\mathbb\{E\}\_\{t,x\_\{0\},x\_\{1\}\}\\\!\\left\[\\lVert\\hat\{b\}\_\{t\}\(I\_\{t\}\)\-\\dot\{I\}\_\{t\}\\rVert^\{2\}\\right\]\(21\)over a class of neural networks, converting generative modeling into a regression problem\.

#### Flow map parameterization\.

Throughout the paper we use the flow map parameterization

Xs,t​\(x\)=x\+\(t−s\)​vs,t​\(x\),X\_\{s,t\}\(x\)=x\+\(t\-s\)\\,v\_\{s,t\}\(x\),\(22\)wherev:\[0,1\]2×ℝd→ℝdv:\[0,1\]^\{2\}\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}is the function to be learned\. On the diagonals=ts=t,Xt,tX\_\{t,t\}reduces to the identity, and thetangent condition

vt,t​\(x\)=bt​\(x\)v\_\{t,t\}\(x\)=b\_\{t\}\(x\)\(23\)identifies the diagonal velocity of the flow map with the drift of the underlying probability flow\. We refer tovt,tv\_\{t,t\}as theimplicit velocityofXs,tX\_\{s,t\}; this property allows the flow map to serve both as an accelerated sampler*and*as the source of the velocity field that our fine\-tuning algorithm operates on\.

#### Three characterizations and self\-distillation losses\.

The flow map of the deterministic dynamicsx˙t=bt​\(xt\)\\dot\{x\}\_\{t\}=b\_\{t\}\(x\_\{t\}\)admits three equivalent characterizations,

∂tXs,t​\(x\)\\displaystyle\\partial\_\{t\}X\_\{s,t\}\(x\)=bt​\(Xs,t​\(x\)\)\\displaystyle=b\_\{t\}\(X\_\{s,t\}\(x\)\)\(Lagrangian\),\\displaystyle\(\\text\{Lagrangian\}\),\(24\)∂sXs,t​\(x\)\+∇Xs,t​\(x\)​bs​\(x\)\\displaystyle\\partial\_\{s\}X\_\{s,t\}\(x\)\+\\nabla X\_\{s,t\}\(x\)\\,b\_\{s\}\(x\)=0\\displaystyle=0\(Eulerian\),\\displaystyle\(\\text\{Eulerian\}\),Xs,t​\(x\)\\displaystyle X\_\{s,t\}\(x\)=Xu,t​\(Xs,u​\(x\)\)\\displaystyle=X\_\{u,t\}\(X\_\{s,u\}\(x\)\)\(semigroup, for​s≤u≤t\),\\displaystyle\(\\text\{semigroup, for \}s\\leq u\\leq t\),each of which gives rise to a different self\-distillation training objective by squaring the corresponding residual and replacingbtb\_\{t\}with the \(stop\-gradient\) diagonal velocityvt,tv\_\{t,t\}via the tangent condition[Equation23](https://arxiv.org/html/2609.27033#A1.E23):

ℒLSD​\(X\)\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{LSD\}\}\(X\)=𝔼⁡\[∥∂tXs,t​\(xs\)−sg⁡\(vt,t​\(Xs,t​\(xs\)\)\)∥2\],\\displaystyle=\\mathbb\{E\}\\\!\\left\[\\,\\lVert\\partial\_\{t\}X\_\{s,t\}\(x\_\{s\}\)\-\\mathrm\{sg\}\\left\(v\_\{t,t\}\(X\_\{s,t\}\(x\_\{s\}\)\)\\right\)\\rVert^\{2\}\\,\\right\],\(25\)ℒESD​\(X\)\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{ESD\}\}\(X\)=𝔼⁡\[∥∂sXs,t​\(xs\)\+sg⁡\(∇Xs,t​\(xs\)​vs,s​\(xs\)\)∥2\],\\displaystyle=\\mathbb\{E\}\\\!\\left\[\\,\\lVert\\partial\_\{s\}X\_\{s,t\}\(x\_\{s\}\)\+\\mathrm\{sg\}\\left\(\\nabla X\_\{s,t\}\(x\_\{s\}\)\\,v\_\{s,s\}\(x\_\{s\}\)\\right\)\\rVert^\{2\}\\,\\right\],ℒPSD​\(X\)\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{PSD\}\}\(X\)=𝔼⁡\[∥Xs,t​\(xs\)−sg⁡\(Xu,t​\(Xs,u​\(xs\)\)\)∥2\]\.\\displaystyle=\\mathbb\{E\}\\\!\\left\[\\,\\lVert X\_\{s,t\}\(x\_\{s\}\)\-\\mathrm\{sg\}\\left\(X\_\{u,t\}\(X\_\{s,u\}\(x\_\{s\}\)\)\\right\)\\rVert^\{2\}\\,\\right\]\.Each of these objectives can be augmented with the flow\-matching objective[Equation21](https://arxiv.org/html/2609.27033#A1.E21)to anchor the diagonal velocityvt,tv\_\{t,t\}to the data\-derived driftbtb\_\{t\}\. In the main text we instead use a value gradient matching objective on the diagonal that pullsw^t,t\\hat\{w\}\_\{t,t\}toward the optimal control direction\.

## Appendix BBackground on the stochastic optimal control formulation of fine\-tuning

In this section, we provide the standard stochastic optimal control formulation of fine\-tuning, following the convention of[Domingo\-Enrich et al\. \[16\]](https://arxiv.org/html/2609.27033#bib.bib37), with notation aligned to ours\.

#### Setup\.

Letbt:ℝd→ℝdb\_\{t\}:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}be the pre\-trained probability\-flow drift, letρt\\rho\_\{t\}denote its marginals, and letr:ℝd→ℝr:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}be a terminal reward\. For a diffusion scaleσt\\sigma\_\{t\}, define the score\-corrected stochastic drift

at\(x\)=bt\(x\)\+12σt2∇logρt\(x\)\.a\_\{t\}\(x\)=b\_\{t\}\(x\)\+\\tfrac\{1\}\{2\}\\,\\sigma\_\{t\}^\{2\}\\,\\nabla\\log\\rho\_\{t\}\(x\)\.\(26\)The reference process

d​Xt=at​\(Xt\)​d​t\+σt​d​Wt,X0∼ρ0,dX\_\{t\}=a\_\{t\}\(X\_\{t\}\)\\,dt\+\\sigma\_\{t\}\\,dW\_\{t\},\\qquad X\_\{0\}\\sim\\rho\_\{0\},\(27\)has the same one\-time marginalsρt\\rho\_\{t\}as the deterministic probability flow, withσt\>0\\sigma\_\{t\}\>0a time\-dependent diffusion coefficient andWtW\_\{t\}a standard Brownian motion\. A controlled processXtuX^\{u\}\_\{t\}evolves according to

d​Xtu=\(at​\(Xtu\)\+σt​ut​\(Xtu\)\)​d​t\+σt​d​Wt,X0u∼ρ0,dX^\{u\}\_\{t\}=\\big\(a\_\{t\}\(X^\{u\}\_\{t\}\)\+\\sigma\_\{t\}\\,u\_\{t\}\(X^\{u\}\_\{t\}\)\\big\)\\,dt\+\\sigma\_\{t\}\\,dW\_\{t\},\\qquad X^\{u\}\_\{0\}\\sim\\rho\_\{0\},\(28\)for a square\-integrable feedback controlu:\[0,1\]×ℝd→ℝdu:\[0,1\]\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}\. Girsanov’s theorem says that the path\-space KL betweenPuP^\{u\}and the referenceRRis given by\[[20](https://arxiv.org/html/2609.27033#bib.bib26)\]

KL\(Pu∥R\)=12𝔼Pu\[∫01∥ut\(Xtu\)∥2dt\]\.\\mathrm\{KL\}\\left\(P^\{u\}\\\|R\\right\)=\\frac\{1\}\{2\}\\mathbb\{E\}\_\{P^\{u\}\}\\\!\\left\[\\int\_\{0\}^\{1\}\\lVert u\_\{t\}\(X^\{u\}\_\{t\}\)\\rVert^\{2\}\\,dt\\right\]\.\(29\)The reward\-tilted distribution

ρ~1​\(x\)∝eλ​r​\(x\)​ρ1​\(x\)\\tilde\{\\rho\}\_\{1\}\(x\)\\propto e^\{\\lambda\\,r\(x\)\}\\,\\rho\_\{1\}\(x\)\(30\)arises as the maximizer ofλ𝔼\[r\(X1\)\]−KL\(ρ~1∥ρ1\)\\lambda\\,\\mathbb\{E\}\[r\(X\_\{1\}\)\]\-\\mathrm\{KL\}\\left\(\\tilde\{\\rho\}\_\{1\}\\\|\\rho\_\{1\}\\right\), whereρ1=Law⁡\(X1\)\\rho\_\{1\}=\\mathrm\{Law\}\(X\_\{1\}\)under the reference[Equation27](https://arxiv.org/html/2609.27033#A2.E27)\.

#### Lifting the reward\-tilted problem to path space\.

Combining the reward with the path\-space KL[Equation29](https://arxiv.org/html/2609.27033#A2.E29)gives the regularized reward\-maximization problem

supu𝔼Pu​\[λ​r​\(X1u\)−12​∫01∥ut​\(Xtu\)∥2​dt\],\\displaystyle\\sup\_\{u\}\\ \\mathbb\{E\}\_\{P^\{u\}\}\\\!\\left\[\\,\\lambda\\,r\(X^\{u\}\_\{1\}\)\-\\tfrac\{1\}\{2\}\\int\_\{0\}^\{1\}\\lVert u\_\{t\}\(X^\{u\}\_\{t\}\)\\rVert^\{2\}\\,dt\\,\\right\],\(31\)subject tod​Xtu=\(at​\(Xtu\)\+σt​ut​\(Xtu\)\)​d​t\+σt​d​Wt,X0u∼ρ0,\\displaystyle\\text\{subject to\}\\quad dX^\{u\}\_\{t\}=\\big\(a\_\{t\}\(X^\{u\}\_\{t\}\)\+\\sigma\_\{t\}\\,u\_\{t\}\(X^\{u\}\_\{t\}\)\\big\)\\,dt\+\\sigma\_\{t\}\\,dW\_\{t\},\\quad X^\{u\}\_\{0\}\\sim\\rho\_\{0\},a stochastic optimal control problem\. The value function under a candidate controluumeasures the expected cost\-to\-go from statexxat timett,

Vtu​\(x\)=𝔼Pu​\[λ​r​\(X1u\)−12​∫t1∥us​\(Xsu\)∥2​𝑑s\|Xtu=x\],V^\{u\}\_\{t\}\(x\)=\\mathbb\{E\}\_\{P^\{u\}\}\\\!\\left\[\\,\\lambda\\,r\(X^\{u\}\_\{1\}\)\-\\tfrac\{1\}\{2\}\\int\_\{t\}^\{1\}\\lVert u\_\{s\}\(X^\{u\}\_\{s\}\)\\rVert^\{2\}\\,ds\\;\\Big\|\\;X^\{u\}\_\{t\}=x\\,\\right\],\(32\)and the optimal value functionVt∗​\(x\)=supuVtu​\(x\)V^\{\*\}\_\{t\}\(x\)=\\sup\_\{u\}V^\{u\}\_\{t\}\(x\)admits the closed\-form Cole–Hopf representation

Vt∗​\(x\)=log⁡𝔼R​\[eλ​r​\(X1\)∣Xt=x\],ut∗​\(x\)=σt​∇xVt∗​\(x\),V^\{\*\}\_\{t\}\(x\)=\\log\\mathbb\{E\}\_\{R\}\\\!\\left\[\\,e^\{\\lambda\\,r\(X\_\{1\}\)\}\\mid X\_\{t\}=x\\,\\right\],\\qquad u^\{\*\}\_\{t\}\(x\)=\\sigma\_\{t\}\\,\\nabla\_\{x\}V^\{\*\}\_\{t\}\(x\),\(33\)where the expectation is taken under the reference process[Equation27](https://arxiv.org/html/2609.27033#A2.E27)rather than underP∗P^\{\*\}\. The optimally\-controlled processP∗P^\{\*\}is the Doobhh\-transform ofRRwithht​\(x\)=eVt∗​\(x\)=𝔼R​\[eλ​r​\(X1\)∣Xt=x\]h\_\{t\}\(x\)=e^\{V^\{\*\}\_\{t\}\(x\)\}=\\mathbb\{E\}\_\{R\}\[e^\{\\lambda\\,r\(X\_\{1\}\)\}\\mid X\_\{t\}=x\], whose joint density on the endpoints is

P∗​\(x0,x1\)=ρ0​\(x0\)​ρ1\|0​\(x1∣x0\)​eλ​r​\(x1\)h0​\(x0\),P^\{\*\}\(x\_\{0\},x\_\{1\}\)=\\rho\_\{0\}\(x\_\{0\}\)\\,\\rho\_\{1\\mid 0\}\(x\_\{1\}\\mid x\_\{0\}\)\\,\\frac\{e^\{\\lambda\\,r\(x\_\{1\}\)\}\}\{h\_\{0\}\(x\_\{0\}\)\},\(34\)whereρ1\|0​\(x1∣x0\)\\rho\_\{1\\mid 0\}\(x\_\{1\}\\mid x\_\{0\}\)is the conditional density ofX1X\_\{1\}givenX0=x0X\_\{0\}=x\_\{0\}under the reference process\. In general, the implicit terminal marginal obtained by integrating[Equation34](https://arxiv.org/html/2609.27033#A2.E34)overx0x\_\{0\}is not the target tiltρ~1\\tilde\{\\rho\}\_\{1\}, because the1/h0​\(x0\)1/h\_\{0\}\(x\_\{0\}\)factor couples the endpoints through the initial\-time value function\.[Domingo\-Enrich et al\. \[16\]](https://arxiv.org/html/2609.27033#bib.bib37)identify the so\-called*memoryless*noise schedule

σtml=2​ηt,ηt=αt​\(β˙tβt​αt−α˙t\),\\sigma^\{\\mathrm\{ml\}\}\_\{t\}=\\sqrt\{2\\,\\eta\_\{t\}\},\\qquad\\eta\_\{t\}=\\alpha\_\{t\}\\\!\\left\(\\frac\{\\dot\{\\beta\}\_\{t\}\}\{\\beta\_\{t\}\}\\,\\alpha\_\{t\}\-\\dot\{\\alpha\}\_\{t\}\\right\),\(35\)which for the standard linear interpolantαt=1−t,βt=t\\alpha\_\{t\}=1\-t,\\ \\beta\_\{t\}=treduces toσtml=2​\(1−t\)/t\\sigma^\{\\mathrm\{ml\}\}\_\{t\}=\\sqrt\{2\(1\-t\)/t\}\. This schedule ensures thatX0X\_\{0\}is independent ofX1X\_\{1\}under the reference process, so thatρ1\|0​\(x1∣x0\)=ρ1​\(x1\)\\rho\_\{1\\mid 0\}\(x\_\{1\}\\mid x\_\{0\}\)=\\rho\_\{1\}\(x\_\{1\}\)and[Equation34](https://arxiv.org/html/2609.27033#A2.E34)factors and the terminal marginal is precisely the reward tilt,P∗​\(x1\)∝eλ​r​\(x1\)​ρ1​\(x1\)=ρ~1​\(x1\)P^\{\*\}\(x\_\{1\}\)\\propto e^\{\\lambda\\,r\(x\_\{1\}\)\}\\,\\rho\_\{1\}\(x\_\{1\}\)=\\tilde\{\\rho\}\_\{1\}\(x\_\{1\}\)\.

In[SectionD\.2](https://arxiv.org/html/2609.27033#A4.SS2)we connect this SOC recipe to the abstract terminal\-distribution framing of[Equation8](https://arxiv.org/html/2609.27033#S3.E8), where we show that the implicit regularizer behind the SOC problem is the entropic Schrödinger\-bridge cost betweenρ0\\rho\_\{0\}and the candidate terminalν\\nu\([PropositionD\.6](https://arxiv.org/html/2609.27033#A4.Thmtheorem6)\)\. This cost has our static optimal transport regularizer𝒯b\\mathcal\{T\}\_\{b\}as its zero\-noise limit \([PropositionD\.7](https://arxiv.org/html/2609.27033#A4.Thmtheorem7)\)\.

## Appendix CBackground on deterministic optimal control

In this section, we collect several standard results from deterministic optimal control theory that we use to design our algorithm in[Section4](https://arxiv.org/html/2609.27033#S4)\. For textbook treatments, we recommend[Fleming and Rishel \[68\]](https://arxiv.org/html/2609.27033#bib.bib59),[Bardi and Capuzzo\-Dolcetta \[69\]](https://arxiv.org/html/2609.27033#bib.bib60),[Liberzon \[47\]](https://arxiv.org/html/2609.27033#bib.bib58)\.

#### Setup\.

We study the deterministic optimal control problem[Equation14](https://arxiv.org/html/2609.27033#S3.E14),

supu𝔼x0∼ρ0​\[λ​r​\(x1u\)−12​∫01∥ut​\(xtu\)∥2​dt\],\\displaystyle\\sup\_\{u\}\\ \\mathbb\{E\}\_\{x\_\{0\}\\sim\\rho\_\{0\}\}\\\!\\left\[\\,\\lambda\\,r\(x\_\{1\}^\{u\}\)\-\\frac\{1\}\{2\}\\int\_\{0\}^\{1\}\\lVert u\_\{t\}\(x\_\{t\}^\{u\}\)\\rVert^\{2\}\\,dt\\,\\right\],\(36\)subject tox˙tu=bt​\(xtu\)\+ut​\(xtu\),x0u=x0,\\displaystyle\\text\{subject to\}\\quad\\dot\{x\}\_\{t\}^\{u\}=b\_\{t\}\(x\_\{t\}^\{u\}\)\+u\_\{t\}\(x\_\{t\}^\{u\}\),\\quad x\_\{0\}^\{u\}=x\_\{0\},through a feedback controlu:\[0,1\]×ℝd→ℝdu:\[0,1\]\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}\. The value function under a candidate controluu[Equation15](https://arxiv.org/html/2609.27033#S4.E15),

Vtu​\(x\)=λ​r​\(Xt,1u​\(x\)\)−12​∫t1∥uτ​\(Xt,τu​\(x\)\)∥2​𝑑τ,V^\{u\}\_\{t\}\(x\)=\\lambda\\,r\(X^\{u\}\_\{t,1\}\(x\)\)\-\\frac\{1\}\{2\}\\int\_\{t\}^\{1\}\\lVert u\_\{\\tau\}\(X^\{u\}\_\{t,\\tau\}\(x\)\)\\rVert^\{2\}\\,d\\tau,\(37\)measures the reward\-to\-go from statexxat timett, and the optimal value function isVt∗​\(x\)=supuVtu​\(x\)V^\{\*\}\_\{t\}\(x\)=\\sup\_\{u\}V^\{u\}\_\{t\}\(x\)\. This definition mirrors its stochastic counterpart[Equation32](https://arxiv.org/html/2609.27033#A2.E32), but in the deterministic settingV∗V^\{\*\}admits no Cole–Hopf representation as an expectation over the reference process: the deterministic reference flow has no randomness to integrate against, and the closed form[Equation33](https://arxiv.org/html/2609.27033#A2.E33)collapses to a pointwise evaluation that returns no information about the optimal control\. We instead characterizeV∗V^\{\*\}via the Hamilton–Jacobi–Bellman equation, derived next\.

#### Hamilton–Jacobi–Bellman equation\.

A standard dynamic\-programming argument shows that the optimal value functionV∗V^\{\*\}solves the Hamilton–Jacobi–Bellman equation

∂tV∗t\(x\)\+supu∈ℝd\{\(bt\(x\)\+u\)⋅∇xV∗t\(x\)−12∥u∥2\}=0,V∗1\(x\)=λr\(x\)\.\\displaystyle\\partial\_\{t\}V^\{\*\}\_\{t\}\(x\)\+\\sup\_\{u\\in\\mathbb\{R\}^\{d\}\}\\\!\\left\\\{\\,\(b\_\{t\}\(x\)\+u\)\\cdot\\nabla\_\{x\}V^\{\*\}\_\{t\}\(x\)\-\\tfrac\{1\}\{2\}\\lVert u\\rVert^\{2\}\\,\\right\\\}=0,\\quad V^\{\*\}\_\{1\}\(x\)=\\lambda\\,r\(x\)\.\(38\)The pointwise supremum is attained at the optimal control

ut∗​\(x\)=∇xVt∗​\(x\),u^\{\*\}\_\{t\}\(x\)=\\nabla\_\{x\}V^\{\*\}\_\{t\}\(x\),\(39\)and substituting back into[Equation38](https://arxiv.org/html/2609.27033#A3.E38)reduces to

∂tVt∗​\(x\)\+bt​\(x\)⋅∇xVt∗​\(x\)\+12​∥∇xVt∗​\(x\)∥2=0\.\\partial\_\{t\}V^\{\*\}\_\{t\}\(x\)\+b\_\{t\}\(x\)\\cdot\\nabla\_\{x\}V^\{\*\}\_\{t\}\(x\)\+\\tfrac\{1\}\{2\}\\lVert\\nabla\_\{x\}V^\{\*\}\_\{t\}\(x\)\\rVert^\{2\}=0\.\(40\)The closed\-form solution[Equation39](https://arxiv.org/html/2609.27033#A3.E39)is what motivates the value gradient regression target in[Section4](https://arxiv.org/html/2609.27033#S4): knowingV∗V^\{\*\}pointwise gives the optimal control directly via differentiation inxx\.

#### Policy\-evaluation identity\.

For any fixed reference controlu¯\\bar\{u\}\(not necessarily optimal\), the value functionVu¯V^\{\\bar\{u\}\}satisfies a linear transport equation that we record as a lemma for later reuse\. This identity is the workhorse behind the proof of the performance\-difference proposition[Proposition4\.1](https://arxiv.org/html/2609.27033#S4.Thmtheorem1)given in[SectionE\.3](https://arxiv.org/html/2609.27033#A5.SS3)\.

###### Lemma C\.1\(Policy evaluation\)\.

Under the regularity assumptions of[SectionD\.1](https://arxiv.org/html/2609.27033#A4.SS1), the value functionVu¯V^\{\\bar\{u\}\}of any feedback controlu¯∈𝒰\\bar\{u\}\\in\\mathcal\{U\}satisfies the linear transport equation

∂tVtu¯​\(x\)\+\(bt​\(x\)\+u¯t​\(x\)\)⋅∇xVtu¯​\(x\)=12​∥u¯t​\(x\)∥2\.\\partial\_\{t\}V^\{\\bar\{u\}\}\_\{t\}\(x\)\+\(b\_\{t\}\(x\)\+\\bar\{u\}\_\{t\}\(x\)\)\\cdot\\nabla\_\{x\}V^\{\\bar\{u\}\}\_\{t\}\(x\)=\\tfrac\{1\}\{2\}\\lVert\\bar\{u\}\_\{t\}\(x\)\\rVert^\{2\}\.\(41\)

###### Proof\.

Fixu¯∈𝒰\\bar\{u\}\\in\\mathcal\{U\}and letxtu¯x\_\{t\}^\{\\bar\{u\}\}denote the controlled trajectoryx˙tu¯=bt​\(xtu¯\)\+u¯t​\(xtu¯\)\\dot\{x\}\_\{t\}^\{\\bar\{u\}\}=b\_\{t\}\(x\_\{t\}^\{\\bar\{u\}\}\)\+\\bar\{u\}\_\{t\}\(x\_\{t\}^\{\\bar\{u\}\}\)\. From the integral definition[Equation37](https://arxiv.org/html/2609.27033#A3.E37)ofVu¯V^\{\\bar\{u\}\}, the value along the controlled trajectory satisfies

Vtu¯​\(xtu¯\)=λ​r​\(x1u¯\)−∫t112​∥u¯τ​\(xτu¯\)∥2​𝑑τ,V^\{\\bar\{u\}\}\_\{t\}\(x\_\{t\}^\{\\bar\{u\}\}\)=\\lambda\\,r\(x\_\{1\}^\{\\bar\{u\}\}\)\-\\int\_\{t\}^\{1\}\\tfrac\{1\}\{2\}\\lVert\\bar\{u\}\_\{\\tau\}\(x\_\{\\tau\}^\{\\bar\{u\}\}\)\\rVert^\{2\}\\,d\\tau,\(42\)whose right\-hand side depends onttonly through the lower limit of the integral\. Differentiating inttgives the trajectory\-level identity

dd​t​Vtu¯​\(xtu¯\)=12​∥u¯t​\(xtu¯\)∥2\.\\frac\{d\}\{dt\}V^\{\\bar\{u\}\}\_\{t\}\(x\_\{t\}^\{\\bar\{u\}\}\)=\\tfrac\{1\}\{2\}\\lVert\\bar\{u\}\_\{t\}\(x\_\{t\}^\{\\bar\{u\}\}\)\\rVert^\{2\}\.\(43\)By the chain rule, the same derivative equals

dd​t​Vtu¯​\(xtu¯\)=∂tVtu¯​\(xtu¯\)\+x˙tu¯⋅∇xVtu¯​\(xtu¯\)=∂tVtu¯​\(xtu¯\)\+\(bt\+u¯t\)​\(xtu¯\)⋅∇xVtu¯​\(xtu¯\)\.\\frac\{d\}\{dt\}V^\{\\bar\{u\}\}\_\{t\}\(x\_\{t\}^\{\\bar\{u\}\}\)=\\partial\_\{t\}V^\{\\bar\{u\}\}\_\{t\}\(x\_\{t\}^\{\\bar\{u\}\}\)\+\\dot\{x\}\_\{t\}^\{\\bar\{u\}\}\\cdot\\nabla\_\{x\}V^\{\\bar\{u\}\}\_\{t\}\(x\_\{t\}^\{\\bar\{u\}\}\)=\\partial\_\{t\}V^\{\\bar\{u\}\}\_\{t\}\(x\_\{t\}^\{\\bar\{u\}\}\)\+\(b\_\{t\}\+\\bar\{u\}\_\{t\}\)\(x\_\{t\}^\{\\bar\{u\}\}\)\\cdot\\nabla\_\{x\}V^\{\\bar\{u\}\}\_\{t\}\(x\_\{t\}^\{\\bar\{u\}\}\)\.\(44\)Equating[Equation43](https://arxiv.org/html/2609.27033#A3.E43)and[Equation44](https://arxiv.org/html/2609.27033#A3.E44)and noting that varying the initial conditionx0x\_\{0\}makesxtu¯x\_\{t\}^\{\\bar\{u\}\}range overℝd\\mathbb\{R\}^\{d\}at eachttyields[Equation41](https://arxiv.org/html/2609.27033#A3.E41)pointwise\. ∎

At the optimal controlu¯=u∗=∇xV∗\\bar\{u\}=u^\{\*\}=\\nabla\_\{x\}V^\{\*\}, comparing[Equation41](https://arxiv.org/html/2609.27033#A3.E41)with the reduced HJB[Equation40](https://arxiv.org/html/2609.27033#A3.E40)confirms the consistencyVu∗=V∗V^\{u^\{\*\}\}=V^\{\*\}\.

## Appendix DAdditional derivations

In this section, we collect additional derivations and structural results that support the development in the main text\.

### Regularity assumptions

In the following, we assume the below standard regularity assumptions\.

###### Assumption D\.1\(Drift regularity\)\.

The pre\-trained driftb:\[0,1\]×ℝd→ℝdb:\[0,1\]\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}is jointly measurable, locally Lipschitz inxxuniformly intt, and has at most linear growth at infinity, so the ODEx˙t=bt​\(xt\)\\dot\{x\}\_\{t\}=b\_\{t\}\(x\_\{t\}\)generates a global flowX:\[0,1\]2×ℝd→ℝdX:\[0,1\]^\{2\}\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}onℝd\\mathbb\{R\}^\{d\}\.

###### Assumption D\.2\(Reward regularity\)\.

The rewardr:ℝd→ℝr:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}is bounded above and upper semicontinuous; for results that involve∇r\\nabla rwe further assumer∈C1​\(ℝd\)r\\in C^\{1\}\(\\mathbb\{R\}^\{d\}\)with bounded gradient\.

###### Assumption D\.3\(Marginals\)\.

The base distributionρ0\\rho\_\{0\}and the candidate terminal distributionsν\\nuhave finite second moment, with𝒯b​\(ρ0,ν\)<∞\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu\)<\\inftyfor the relevantν\\nu\.

###### Assumption D\.4\(Trainable flow map class\)\.

The trainable fieldw^\\hat\{w\}ranges over a convex class𝒲\\mathcal\{W\}of two\-time velocity fieldsw^:\{0≤s≤t≤1\}×ℝd→ℝd\\hat\{w\}:\\\{0\\leq s\\leq t\\leq 1\\\}\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}that are continuously differentiable inxxand inss, and locally Lipschitz inxxuniformly in\(s,t\)\(s,t\)\. We assume𝒲\\mathcal\{W\}contains, for every continuous fieldϕ\\phion the upper triangle, the averaged field\(s,t,x\)↦1t−s​∫stϕ⁡\(σ,t,x\)​𝑑σ\(s,t,x\)\\mapsto\\tfrac\{1\}\{t\-s\}\\int\_\{s\}^\{t\}\\phi\(\\sigma,t,x\)\\,d\\sigma, and thatw^t,t−bt\\hat\{w\}\_\{t,t\}\-b\_\{t\}is an admissible control for everyw^∈𝒲\\hat\{w\}\\in\\mathcal\{W\}\.

###### Assumption D\.5\(Control class\)\.

The admissible control class𝒰\\mathcal\{U\}consists of locally Lipschitz feedback controlsu:\[0,1\]×ℝd→ℝdu:\[0,1\]\\times\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}satisfying𝔼x0∼ρ0​\[∫01∥ut​\(xtu\)∥2​𝑑t\]<∞\\mathbb\{E\}\_\{x\_\{0\}\\sim\\rho\_\{0\}\}\\\!\\left\[\\int\_\{0\}^\{1\}\\lVert u\_\{t\}\(x\_\{t\}^\{u\}\)\\rVert^\{2\}dt\\right\]<\\infty\.

These assumptions are standard in the theory of dynamical systems and optimal control to ensure uniqueness of solutions\[[68](https://arxiv.org/html/2609.27033#bib.bib59),[69](https://arxiv.org/html/2609.27033#bib.bib60),[47](https://arxiv.org/html/2609.27033#bib.bib58),[70](https://arxiv.org/html/2609.27033#bib.bib66)\], and in optimal transport theory to ensure existence of solutions to the corresponding optimization problems\[[40](https://arxiv.org/html/2609.27033#bib.bib16),[45](https://arxiv.org/html/2609.27033#bib.bib15)\]\. In particular,[AssumptionsD\.1](https://arxiv.org/html/2609.27033#A4.Thmtheorem1),[D\.2](https://arxiv.org/html/2609.27033#A4.Thmtheorem2)and[D\.3](https://arxiv.org/html/2609.27033#A4.Thmtheorem3)together imply that the costcbc\_\{b\}is lower semicontinuous and coercive in its endpoints, so that the static infimum𝒯b​\(ρ0,ν\)\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu\)is attained\.

### Schrödinger bridges and the zero\-noise limit

The SOC approach in[AppendixB](https://arxiv.org/html/2609.27033#A2)optimizes over path measuresPP, but the reward depends only on the terminalX1∼P1X\_\{1\}\\sim P\_\{1\}\. We first show that the implicit terminal\-level regularizer behind this recipe, in the sense of[Equation8](https://arxiv.org/html/2609.27033#S3.E8), is the entropic Schrödinger\-bridge cost \([PropositionD\.6](https://arxiv.org/html/2609.27033#A4.Thmtheorem6)\), then take the zero\-noise limit of this cost to recover our static prior\-action transport regularizer𝒯b\\mathcal\{T\}\_\{b\}\([PropositionD\.7](https://arxiv.org/html/2609.27033#A4.Thmtheorem7)\)\.

For a noise scaleε\>0\\varepsilon\>0, letRεR^\{\\varepsilon\}denote the reference path measure onC⁡\(\[0,1\],ℝd\)C\(\[0,1\];\\mathbb\{R\}^\{d\}\)associated with the reference SDE

d​Xt=at​\(Xt\)​d​t\+ε​d​Wt,X0∼ρ0,dX\_\{t\}=a\_\{t\}\(X\_\{t\}\)\\,dt\+\\sqrt\{\\varepsilon\}\\,dW\_\{t\},\\qquad X\_\{0\}\\sim\\rho\_\{0\},\(45\)which is the constant\-noise specialization of[AppendixB](https://arxiv.org/html/2609.27033#A2)withσt=ε\\sigma\_\{t\}=\\sqrt\{\\varepsilon\}\. It preserves the marginalsρt\\rho\_\{t\}and converges to the deterministic probability flow asε→0\\varepsilon\\to 0\. The entropic Schrödinger\-bridge cost betweenρ0\\rho\_\{0\}and a candidate terminalν\\nuis the infimum of the path\-space KL over all path measures with prescribed marginals,

SBε\(ρ0,ν\)=infP:P0=ρ0,P1=νKL\(P∥Rε\)\.\\mathrm\{SB\}\_\{\\varepsilon\}\(\\rho\_\{0\},\\nu\)=\\inf\_\{P\\,:\\,P\_\{0\}=\\rho\_\{0\},\\,P\_\{1\}=\\nu\}\\,\\mathrm\{KL\}\\left\(P\\\|R^\{\\varepsilon\}\\right\)\.\(46\)
###### Proposition D\.6\(Path\-KL collapses to entropic SB at the terminal level\)\.

For anyλ\>0\\lambda\>0andε\>0\\varepsilon\>0, the path\-space fine\-tuning problem

supP:P0=ρ0\{λ𝔼P\[r\(X1\)\]−KL\(P∥Rε\)\}\\sup\_\{P\\,:\\,P\_\{0\}=\\rho\_\{0\}\}\\ \\left\\\{\\ \\lambda\\,\\mathbb\{E\}\_\{P\}\[r\(X\_\{1\}\)\]\-\\mathrm\{KL\}\\left\(P\\\|R^\{\\varepsilon\}\\right\)\\ \\right\\\}\(47\)has the same value as the terminal\-space problem

supν\{λ​𝔼x∼ν​\[r⁡\(x\)\]−SBε​\(ρ0,ν\)\}\.\\sup\_\{\\nu\}\\ \\left\\\{\\ \\lambda\\,\\mathbb\{E\}\_\{x\\sim\\nu\}\[r\(x\)\]\-\\mathrm\{SB\}\_\{\\varepsilon\}\(\\rho\_\{0\},\\nu\)\\ \\right\\\}\.\(48\)A path measureP∗P^\{\*\}achieves the supremum in[Equation47](https://arxiv.org/html/2609.27033#A4.E47)if and only if its terminal lawν∗=P1∗\\nu^\{\*\}=P^\{\*\}\_\{1\}achieves the supremum in[Equation48](https://arxiv.org/html/2609.27033#A4.E48)andP∗P^\{\*\}achieves the infimum[Equation46](https://arxiv.org/html/2609.27033#A4.E46)definingSBε​\(ρ0,ν∗\)\\mathrm\{SB\}\_\{\\varepsilon\}\(\\rho\_\{0\},\\nu^\{\*\}\)\.

###### Proof\.

The reward𝔼P​\[r⁡\(X1\)\]=𝔼x∼P1​\[r⁡\(x\)\]\\mathbb\{E\}\_\{P\}\[r\(X\_\{1\}\)\]=\\mathbb\{E\}\_\{x\\sim P\_\{1\}\}\[r\(x\)\]depends onPPonly through its terminal marginal\. Splitting the supremum,

supP:P0=ρ0\{λ𝔼P\[r\(X1\)\]−KL\(P∥Rε\)\}\\displaystyle\\sup\_\{P\\,:\\,P\_\{0\}=\\rho\_\{0\}\}\\ \\big\\\{\\lambda\\,\\mathbb\{E\}\_\{P\}\[r\(X\_\{1\}\)\]\-\\mathrm\{KL\}\\left\(P\\\|R^\{\\varepsilon\}\\right\)\\big\\\}=supνsupP:P0=ρ0,P1=ν\{λ𝔼ν\[r\]−KL\(P∥Rε\)\}\\displaystyle=\\sup\_\{\\nu\}\\ \\sup\_\{P\\,:\\,P\_\{0\}=\\rho\_\{0\},\\,P\_\{1\}=\\nu\}\\ \\big\\\{\\lambda\\,\\mathbb\{E\}\_\{\\nu\}\[r\]\-\\mathrm\{KL\}\\left\(P\\\|R^\{\\varepsilon\}\\right\)\\big\\\}\(49\)=supν\{λ𝔼ν\[r\]−infP:P0=ρ0,P1=νKL\(P∥Rε\)\},\\displaystyle=\\sup\_\{\\nu\}\\ \\big\\\{\\lambda\\,\\mathbb\{E\}\_\{\\nu\}\[r\]\-\\inf\_\{P\\,:\\,P\_\{0\}=\\rho\_\{0\},\\,P\_\{1\}=\\nu\}\\mathrm\{KL\}\\left\(P\\\|R^\{\\varepsilon\}\\right\)\\big\\\},\(50\)where the second equality moves theinf\\infinside because the reward term does not depend onPP\. The inner infimum isSBε​\(ρ0,ν\)\\mathrm\{SB\}\_\{\\varepsilon\}\(\\rho\_\{0\},\\nu\)by definition, giving the value equality\. The optimizer characterization follows from the same splitting, as a maximizerP∗P^\{\*\}of the joint supremum must achieve both the outer supremum \(inν∗\\nu^\{\*\}\) and the inner infimum \(for thatν∗\\nu^\{\*\}\), and conversely\. ∎

The following result clarifies the relationship between Adjoint Matching and the WTF formulation of[Section3](https://arxiv.org/html/2609.27033#S3), showing that WTF can be viewed as a suitable zero\-noise limit of a Schrödinger bridge problem, while Adjoint Matching is a specific Schrödinger bridge problem with the memoryless noise schedule\.

###### Proposition D\.7\(Zero\-noise limit\)\.

Under the standing regularity assumptions of[SectionD\.1](https://arxiv.org/html/2609.27033#A4.SS1), the rescaled entropic Schrödinger\-bridge cost converges to the static prior\-action transport cost,

ε​SBε​\(ρ0,ν\)⟶𝒯b​\(ρ0,ν\)as​ε→0\+,\\varepsilon\\,\\mathrm\{SB\}\_\{\\varepsilon\}\(\\rho\_\{0\},\\nu\)\\;\\longrightarrow\\;\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu\)\\quad\\text\{as \}\\varepsilon\\to 0^\{\+\},\(51\)and the corresponding stochastic optimal control formulation ofSBε\\mathrm\{SB\}\_\{\\varepsilon\}reduces in the same limit to the deterministic OC problem[Equation14](https://arxiv.org/html/2609.27033#S3.E14)\.

The proof combines Freidlin–Wentzell large\-deviations andΓ\\Gamma\-convergence with the dynamic\-form representation of𝒯b\\mathcal\{T\}\_\{b\}from[LemmaE\.1](https://arxiv.org/html/2609.27033#A5.Thmtheorem1)\. The key steps are carried out by[Chen et al\. \[43\]](https://arxiv.org/html/2609.27033#bib.bib48),[Chen et al\. \[44\]](https://arxiv.org/html/2609.27033#bib.bib25); we refer the reader there for the proof\.[PropositionD\.7](https://arxiv.org/html/2609.27033#A4.Thmtheorem7)formalizes the statement that WTF is the deterministic limit of the path\-KL fine\-tuning recipe of[AppendixB](https://arxiv.org/html/2609.27033#A2)\.[Table8](https://arxiv.org/html/2609.27033#A4.T8)situates WTF alongside the KL\-based and SB\-based recipes by what each regularizer constrains, what algorithm it induces, and the rollout cost it pays at training time\.

Table 8:Comparison of methods\.Regularizers for fine\-tuning generative models, organized by what they constrain\. WTF \(last row\) is defined directly in terms of the deterministic dynamics of the pre\-trained flow, which lets the resulting algorithm exploit the flow map for constant\-number flow map evaluations at training time\.
### Failure of the naive control objective

The most direct algorithmic translation of the optimal control problem[Equation14](https://arxiv.org/html/2609.27033#S3.E14)uses the fine\-tuned flow map[Equation83](https://arxiv.org/html/2609.27033#A6.E83)Xs,tw^​\(x\)=x\+\(t−s\)​w^s,t​\(x\)X^\{\\hat\{w\}\}\_\{s,t\}\(x\)=x\+\(t\-s\)\\,\\hat\{w\}\_\{s,t\}\(x\)from the main text and writes the residual control asut=w^t,t−btu\_\{t\}=\\hat\{w\}\_\{t,t\}\-b\_\{t\}, leading to the sampled objective

ℒ⁡\(w^\)=𝔼x0,t​\[12​∥w^t,t​\(X0,tw^​\(x0\)\)−bt​\(X0,tw^​\(x0\)\)∥22−λ​r​\(X0,1w^​\(x0\)\)\]\+β​ℒdist​\(w^\),\\mathcal\{L\}\(\\hat\{w\}\)=\\mathbb\{E\}\_\{x\_\{0\},t\}\\\!\\left\[\\,\\tfrac\{1\}\{2\}\\lVert\\hat\{w\}\_\{t,t\}\(X^\{\\hat\{w\}\}\_\{0,t\}\(x\_\{0\}\)\)\-b\_\{t\}\(X^\{\\hat\{w\}\}\_\{0,t\}\(x\_\{0\}\)\)\\rVert\_\{2\}^\{2\}\-\\lambda\\,r\(X^\{\\hat\{w\}\}\_\{0,1\}\(x\_\{0\}\)\)\\,\\right\]\+\\beta\\,\\mathcal\{L\}\_\{\\mathrm\{dist\}\}\(\\hat\{w\}\),\(52\)whereℒdist\\mathcal\{L\}\_\{\\mathrm\{dist\}\}is an off\-diagonal self\-distillation loss\. In practice we find that[Equation52](https://arxiv.org/html/2609.27033#A4.E52)fails to provide a sufficiently informative learning signal for the diagonal velocityw^t,t\\hat\{w\}\_\{t,t\}\. The reward acts on the terminal mapX0,1w^X^\{\\hat\{w\}\}\_\{0,1\}, which depends on the off\-diagonalw^s,t\\hat\{w\}\_\{s,t\}for\(s,t\)\(s,t\)all the way up to\(0,1\)\(0,1\)\. The control cost in contrast only seesw^t,t\\hat\{w\}\_\{t,t\}on the diagonal\. The model can therefore change the terminal map in directions that increase reward while leaving the sampled diagonal velocity nearly unchanged\. This creates a map\-dynamics inconsistency: the off\-diagonal map adapts without learning the corresponding diagonal dynamics\. Early experiments with this objective motivated the value\-gradient approach in the main text\.

## Appendix EOmitted proofs

In this section, we restate and prove the mathematical results from the main text\.

### Proof of[Proposition3\.1](https://arxiv.org/html/2609.27033#S3.Thmtheorem1)

See[3\.1](https://arxiv.org/html/2609.27033#S3.Thmtheorem1)

###### Proof\.

Substituting the Kantorovich definition[Equation11](https://arxiv.org/html/2609.27033#S3.E11)of𝒯b\\mathcal\{T\}\_\{b\}into the static reward\-regularized problem,

supν\{λ​𝔼y∼ν​\[r⁡\(y\)\]−𝒯b​\(ρ0,ν\)\}\\displaystyle\\sup\_\{\\nu\}\\big\\\{\\lambda\\,\\mathbb\{E\}\_\{y\\sim\\nu\}\[r\(y\)\]\-\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu\)\\big\\\}=supνsupπ∈Π⁡\(ρ0,ν\)∫\[λ​r​\(y\)−cb​\(x,y\)\]​π​\(d​x​dy\)\\displaystyle=\\sup\_\{\\nu\}\\,\\sup\_\{\\pi\\in\\Pi\(\\rho\_\{0\},\\,\\nu\)\}\\ \\int\\big\[\\,\\lambda\\,r\(y\)\-c\_\{b\}\(x,y\)\\,\\big\]\\,\\pi\(dx\\,dy\)\(53\)=supπ:π0=ρ0∫\[λr\(y\)−cb\(x,y\)\]π\(dxdy\),\\displaystyle=\\sup\_\{\\pi\\,:\\,\\pi\_\{0\}=\\rho\_\{0\}\}\\ \\int\\big\[\\,\\lambda\\,r\(y\)\-c\_\{b\}\(x,y\)\\,\\big\]\\,\\pi\(dx\\,dy\),where the second equality uses that the joint sup over\(ν,π\)\(\\nu,\\pi\)withπ∈Π⁡\(ρ0,ν\)\\pi\\in\\Pi\(\\rho\_\{0\},\\nu\)reduces to a sup over couplingsπ\\piwith first marginalρ0\\rho\_\{0\}and arbitrary second marginal\. The constraintπ0=ρ0\\pi\_\{0\}=\\rho\_\{0\}lets us disintegrate any feasible coupling asπ⁡\(d​x​d​y\)=ρ0​\(d​x\)​π1\|0​\(d​y∣x\)\\pi\(dx\\,dy\)=\\rho\_\{0\}\(dx\)\\,\\pi\_\{1\\mid 0\}\(dy\\mid x\), where for eachxxthe conditionalπ1\|0\(⋅∣x\)\\pi\_\{1\\mid 0\}\(\\cdot\\mid x\)is a probability measure onℝd\\mathbb\{R\}^\{d\}and is otherwise unconstrained\. Substituting this disintegration gives

supπ:π0=ρ0∫\[λr\(y\)−cb\(x,y\)\]π\(dxdy\)\\displaystyle\\sup\_\{\\pi\\,:\\,\\pi\_\{0\}=\\rho\_\{0\}\}\\ \\int\\big\[\\lambda\\,r\(y\)\-c\_\{b\}\(x,y\)\\big\]\\,\\pi\(dx\\,dy\)\(54\)=sup\{π1\|0\(⋅∣x\)\}x∫ρ0\(dx\)∫\[λr\(y\)−cb\(x,y\)\]π1\|0\(dy∣x\)\\displaystyle=\\sup\_\{\\\{\\pi\_\{1\\mid 0\}\(\\cdot\\mid x\)\\\}\_\{x\}\}\\ \\int\\rho\_\{0\}\(dx\)\\,\\int\\big\[\\lambda\\,r\(y\)\-c\_\{b\}\(x,y\)\\big\]\\,\\pi\_\{1\\mid 0\}\(dy\\mid x\)=∫ρ0\(dx\)supπ1\|0\(⋅∣x\)∫\[λr\(y\)−cb\(x,y\)\]π1\|0\(dy∣x\)\\displaystyle=\\int\\rho\_\{0\}\(dx\)\\,\\sup\_\{\\pi\_\{1\\mid 0\}\(\\cdot\\mid x\)\}\\,\\int\\big\[\\lambda\\,r\(y\)\-c\_\{b\}\(x,y\)\\big\]\\,\\pi\_\{1\\mid 0\}\(dy\\mid x\)=𝔼x∼ρ0​\[supy∈ℝd\{λ​r​\(y\)−cb​\(x,y\)\}\],\\displaystyle=\\mathbb\{E\}\_\{x\\sim\\rho\_\{0\}\}\\\!\\left\[\\,\\sup\_\{y\\in\\mathbb\{R\}^\{d\}\}\\ \\big\\\{\\,\\lambda\\,r\(y\)\-c\_\{b\}\(x,y\)\\,\\big\\\}\\,\\right\],which is[Equation13](https://arxiv.org/html/2609.27033#S3.E13)\. The first equality substitutes the disintegration, the second exchanges the sup with the outer integral, and the third observes that the inner sup over probability measures onℝd\\mathbb\{R\}^\{d\}of the linear functionalμ↦∫f⁡\(y\)​μ​\(𝑑y\)\\mu\\mapsto\\int f\(y\)\\,\\mu\(dy\)withf⁡\(y\)=λ​r​\(y\)−cb​\(x,y\)f\(y\)=\\lambda\\,r\(y\)\-c\_\{b\}\(x,y\)equalssupyf⁡\(y\)\\sup\_\{y\}f\(y\), attained by a Dirac at any maximizery∗​\(x\)∈argmaxy\{λ​r​\(y\)−cb​\(x,y\)\}y^\{\*\}\(x\)\\in\\argmax\_\{y\}\\\{\\lambda\\,r\(y\)\-c\_\{b\}\(x,y\)\\\}\. ∎

### Proof of[Proposition3\.2](https://arxiv.org/html/2609.27033#S3.Thmtheorem2)

The proof relies on the dynamic formulation of optimal transport with prior dynamics established by[Chen et al\. \[42\]](https://arxiv.org/html/2609.27033#bib.bib49),[Chen et al\. \[43\]](https://arxiv.org/html/2609.27033#bib.bib48),[Chen et al\. \[44\]](https://arxiv.org/html/2609.27033#bib.bib25), due to[Chen et al\. \[42\]](https://arxiv.org/html/2609.27033#bib.bib49)for linear priors and to[Chen et al\. \[43\]](https://arxiv.org/html/2609.27033#bib.bib48),[Chen et al\. \[44\]](https://arxiv.org/html/2609.27033#bib.bib25)for general drifts\. We state the result in our notation for completeness; it holds under the regularity assumptions of[SectionD\.1](https://arxiv.org/html/2609.27033#A4.SS1), and we refer the reader to those works for its proof\.

###### Lemma E\.1\(Benamou–Brenier with prior dynamics\)\.

Under the regularity assumptions of[SectionD\.1](https://arxiv.org/html/2609.27033#A4.SS1), for everyν\\nusuch that𝒯b​\(ρ0,ν\)<∞\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu\)<\\infty,

𝒯b\(ρ0,ν\)=inf\(ρ,v\)\\displaystyle\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu\)=\\inf\_\{\(\\rho,v\)\}∫01∫12​∥vt​\(x\)−bt​\(x\)∥2​ρt​\(x\)​dx​dt,\\displaystyle\\int\_\{0\}^\{1\}\\\!\\\!\\int\\tfrac\{1\}\{2\}\\lVert v\_\{t\}\(x\)\-b\_\{t\}\(x\)\\rVert^\{2\}\\,\\rho\_\{t\}\(x\)\\,dx\\,dt,\(55\)subject to\\displaystyle\\text\{subject to\}∂tρt\+∇⋅\(vtρt\)=0,ρ\|t=0=ρ0,ρ\|t=1=ν\.\\displaystyle\\partial\_\{t\}\\rho\_\{t\}\+\\nabla\\\!\\cdot\\\!\(v\_\{t\}\\,\\rho\_\{t\}\)=0,\\quad\\rho\|\_\{t=0\}=\\rho\_\{0\},\\quad\\rho\|\_\{t=1\}=\\nu\.

We now use[LemmaE\.1](https://arxiv.org/html/2609.27033#A5.Thmtheorem1)to prove[Proposition3\.2](https://arxiv.org/html/2609.27033#S3.Thmtheorem2)\. See[3\.2](https://arxiv.org/html/2609.27033#S3.Thmtheorem2)

###### Proof\.

For brevity we write

JOC​\(u\)\\displaystyle J\_\{\\mathrm\{OC\}\}\(u\)=𝔼x0∼ρ0​\[λ​r​\(x1u\)−12​∫01∥ut​\(xtu\)∥2​dt\],\\displaystyle=\\mathbb\{E\}\_\{x\_\{0\}\\sim\\rho\_\{0\}\}\\\!\\left\[\\lambda\\,r\(x\_\{1\}^\{u\}\)\-\\tfrac\{1\}\{2\}\\int\_\{0\}^\{1\}\\lVert u\_\{t\}\(x\_\{t\}^\{u\}\)\\rVert^\{2\}\\,dt\\right\],\(56\)Jstat​\(ν\)\\displaystyle J\_\{\\mathrm\{stat\}\}\(\\nu\)=λ​𝔼x∼ν​\[r⁡\(x\)\]−𝒯b​\(ρ0,ν\),\\displaystyle=\\lambda\\,\\mathbb\{E\}\_\{x\\sim\\nu\}\[r\(x\)\]\-\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu\),for the OC and static objectives, respectively\. We prove the value equality

supu∈𝒰JOC​\(u\)=supνJstat​\(ν\)\\sup\_\{u\\in\\mathcal\{U\}\}J\_\{\\mathrm\{OC\}\}\(u\)=\\sup\_\{\\nu\}J\_\{\\mathrm\{stat\}\}\(\\nu\)\(57\)via matching upper and lower bounds\. The lower\-bound construction also yields the optimizer\-realization claim under the regularity hypotheses\.

We first prove the upper boundsupuJOC​\(u\)≤supνJstat​\(ν\)\\sup\_\{u\}J\_\{\\mathrm\{OC\}\}\(u\)\\leq\\sup\_\{\\nu\}J\_\{\\mathrm\{stat\}\}\(\\nu\)\. Fix anyu∈𝒰u\\in\\mathcal\{U\}and letρtu=Law⁡\(xtu\)\\rho\_\{t\}^\{u\}=\\mathrm\{Law\}\(x\_\{t\}^\{u\}\)denote the law of the controlled trajectory, with corresponding velocityvtu=bt\+utv\_\{t\}^\{u\}=b\_\{t\}\+u\_\{t\}\. By the standard correspondence between trajectories and densities,\(ρu,vu\)\(\\rho^\{u\},v^\{u\}\)satisfies the continuity equation

∂tρtu\+∇⋅\(vtu​ρtu\)=0,ρu\|t=0=ρ0,ρu\|t=1=Law⁡\(x1u\),\\partial\_\{t\}\\rho\_\{t\}^\{u\}\+\\nabla\\cdot\(v\_\{t\}^\{u\}\\,\\rho\_\{t\}^\{u\}\)=0,\\qquad\\rho^\{u\}\|\_\{t=0\}=\\rho\_\{0\},\\qquad\\rho^\{u\}\|\_\{t=1\}=\\mathrm\{Law\}\(x\_\{1\}^\{u\}\),\(58\)so\(ρu,vu\)\(\\rho^\{u\},v^\{u\}\)is a feasible candidate for the dynamic prior\-action transport problem[Equation55](https://arxiv.org/html/2609.27033#A5.E55)fromρ0\\rho\_\{0\}toLaw⁡\(x1u\)\\mathrm\{Law\}\(x\_\{1\}^\{u\}\)\. Pulling the squared deviation12​∥vtu−bt∥2=12​∥ut∥2\\tfrac\{1\}\{2\}\\lVert v\_\{t\}^\{u\}\-b\_\{t\}\\rVert^\{2\}=\\tfrac\{1\}\{2\}\\lVert u\_\{t\}\\rVert^\{2\}back to trajectories underρtu\\rho\_\{t\}^\{u\}gives the action identity

∫01∫12​∥vtu​\(x\)−bt​\(x\)∥2​ρtu​\(x\)​𝑑x​𝑑t=12​𝔼x0∼ρ0​\[∫01∥ut​\(xtu\)∥2​𝑑t\],\\int\_\{0\}^\{1\}\\\!\\\!\\int\\tfrac\{1\}\{2\}\\lVert v\_\{t\}^\{u\}\(x\)\-b\_\{t\}\(x\)\\rVert^\{2\}\\,\\rho\_\{t\}^\{u\}\(x\)\\,dx\\,dt=\\tfrac\{1\}\{2\}\\,\\mathbb\{E\}\_\{x\_\{0\}\\sim\\rho\_\{0\}\}\\\!\\left\[\\int\_\{0\}^\{1\}\\lVert u\_\{t\}\(x\_\{t\}^\{u\}\)\\rVert^\{2\}\\,dt\\right\],\(59\)which is the change of variables from densities to trajectories, usingLaw⁡\(xtu\)=ρtu\\mathrm\{Law\}\(x^\{u\}\_\{t\}\)=\\rho^\{u\}\_\{t\}\. Applying[LemmaE\.1](https://arxiv.org/html/2609.27033#A5.Thmtheorem1)to the candidate\(ρu,vu\)\(\\rho^\{u\},v^\{u\}\), the right\-hand side is at least𝒯b​\(ρ0,Law⁡\(x1u\)\)\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\mathrm\{Law\}\(x\_\{1\}^\{u\}\)\), and so

JOC​\(u\)\\displaystyle J\_\{\\mathrm\{OC\}\}\(u\)=λ​𝔼x∼Law⁡\(x1u\)​\[r⁡\(x\)\]−12​𝔼x0∼ρ0​\[∫01∥ut​\(xtu\)∥2​dt\]\\displaystyle=\\;\\lambda\\,\\mathbb\{E\}\_\{x\\sim\\mathrm\{Law\}\(x\_\{1\}^\{u\}\)\}\[r\(x\)\]\-\\tfrac\{1\}\{2\}\\,\\mathbb\{E\}\_\{x\_\{0\}\\sim\\rho\_\{0\}\}\\\!\\left\[\\int\_\{0\}^\{1\}\\lVert u\_\{t\}\(x\_\{t\}^\{u\}\)\\rVert^\{2\}\\,dt\\right\]\(60\)≤λ​𝔼x∼Law⁡\(x1u\)​\[r⁡\(x\)\]−𝒯b​\(ρ0,Law⁡\(x1u\)\)=Jstat​\(Law⁡\(x1u\)\)\.\\displaystyle\\leq\\;\\lambda\\,\\mathbb\{E\}\_\{x\\sim\\mathrm\{Law\}\(x\_\{1\}^\{u\}\)\}\[r\(x\)\]\-\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\mathrm\{Law\}\(x\_\{1\}^\{u\}\)\)\\;=\\;J\_\{\\mathrm\{stat\}\}\(\\mathrm\{Law\}\(x\_\{1\}^\{u\}\)\)\.The first equality is the change of variables𝔼x0∼ρ0​\[r⁡\(x1u​\(x0\)\)\]=𝔼x∼Law⁡\(x1u\)​\[r⁡\(x\)\]\\mathbb\{E\}\_\{x\_\{0\}\\sim\\rho\_\{0\}\}\[r\(x\_\{1\}^\{u\}\(x\_\{0\}\)\)\]=\\mathbb\{E\}\_\{x\\sim\\mathrm\{Law\}\(x\_\{1\}^\{u\}\)\}\[r\(x\)\], the inequality is[LemmaE\.1](https://arxiv.org/html/2609.27033#A5.Thmtheorem1)applied to\(ρu,vu\)\(\\rho^\{u\},v^\{u\}\)via[Equation59](https://arxiv.org/html/2609.27033#A5.E59), and the final equality is the definition ofJstatJ\_\{\\mathrm\{stat\}\}\. Taking the sup overu∈𝒰u\\in\\mathcal\{U\}on both sides and using that the set of attainable terminal laws\{Law⁡\(x1u\):u∈𝒰\}\\\{\\mathrm\{Law\}\(x\_\{1\}^\{u\}\):u\\in\\mathcal\{U\}\\\}is contained in the admissible set forsupνJstat​\(ν\)\\sup\_\{\\nu\}J\_\{\\mathrm\{stat\}\}\(\\nu\)yieldssupuJOC​\(u\)≤supνJstat​\(ν\)\\sup\_\{u\}J\_\{\\mathrm\{OC\}\}\(u\)\\leq\\sup\_\{\\nu\}J\_\{\\mathrm\{stat\}\}\(\\nu\)\.

We now prove the matching lower boundsupuJOC​\(u\)≥supνJstat​\(ν\)\\sup\_\{u\}J\_\{\\mathrm\{OC\}\}\(u\)\\geq\\sup\_\{\\nu\}J\_\{\\mathrm\{stat\}\}\(\\nu\)\. Letν∗\\nu^\{\*\}achieve the supremum ofJstatJ\_\{\\mathrm\{stat\}\}, whose existence follows from the assumptions of[SectionD\.1](https://arxiv.org/html/2609.27033#A4.SS1)\. By[LemmaE\.1](https://arxiv.org/html/2609.27033#A5.Thmtheorem1), there exists a pair\(ρ∗,v∗\)\(\\rho^\{\*\},v^\{\*\}\)achieving the dynamic prior\-action cost,

∫01∫12​∥vt∗​\(x\)−bt​\(x\)∥2​ρt∗​\(x\)​𝑑x​𝑑t=𝒯b​\(ρ0,ν∗\),ρ∗\|t=0=ρ0,ρ∗\|t=1=ν∗\.\\int\_\{0\}^\{1\}\\\!\\\!\\int\\tfrac\{1\}\{2\}\\lVert v\_\{t\}^\{\*\}\(x\)\-b\_\{t\}\(x\)\\rVert^\{2\}\\,\\rho\_\{t\}^\{\*\}\(x\)\\,dx\\,dt=\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu^\{\*\}\),\\qquad\\rho^\{\*\}\|\_\{t=0\}=\\rho\_\{0\},\\qquad\\rho^\{\*\}\|\_\{t=1\}=\\nu^\{\*\}\.\(61\)Setut∗=vt∗−btu^\{\*\}\_\{t\}=v^\{\*\}\_\{t\}\-b\_\{t\}\. Under the regularity assumed in the theorem statement,u∗∈𝒰u^\{\*\}\\in\\mathcal\{U\}and the controlled flowx˙tu∗=bt​\(xtu∗\)\+ut∗​\(xtu∗\)\\dot\{x\}^\{u^\{\*\}\}\_\{t\}=b\_\{t\}\(x^\{u^\{\*\}\}\_\{t\}\)\+u^\{\*\}\_\{t\}\(x^\{u^\{\*\}\}\_\{t\}\)generates the velocity fieldv∗v^\{\*\}, so thatLaw⁡\(xtu∗\)=ρt∗\\mathrm\{Law\}\(x\_\{t\}^\{u^\{\*\}\}\)=\\rho\_\{t\}^\{\*\}and in particularLaw⁡\(x1u∗\)=ν∗\\mathrm\{Law\}\(x\_\{1\}^\{u^\{\*\}\}\)=\\nu^\{\*\}\. Pulling the cost back to trajectories as in[Equation59](https://arxiv.org/html/2609.27033#A5.E59)and substituting[Equation61](https://arxiv.org/html/2609.27033#A5.E61)yields

JOC​\(u∗\)\\displaystyle J\_\{\\mathrm\{OC\}\}\(u^\{\*\}\)=λ​𝔼x∼ν∗​\[r⁡\(x\)\]−12​𝔼x0∼ρ0​\[∫01∥ut∗​\(xtu∗\)∥2​dt\]\\displaystyle=\\;\\lambda\\,\\mathbb\{E\}\_\{x\\sim\\nu^\{\*\}\}\[r\(x\)\]\-\\tfrac\{1\}\{2\}\\,\\mathbb\{E\}\_\{x\_\{0\}\\sim\\rho\_\{0\}\}\\\!\\left\[\\int\_\{0\}^\{1\}\\lVert u^\{\*\}\_\{t\}\(x\_\{t\}^\{u^\{\*\}\}\)\\rVert^\{2\}\\,dt\\right\]\(62\)=λ​𝔼x∼ν∗​\[r⁡\(x\)\]−𝒯b​\(ρ0,ν∗\)=Jstat​\(ν∗\),\\displaystyle=\\;\\lambda\\,\\mathbb\{E\}\_\{x\\sim\\nu^\{\*\}\}\[r\(x\)\]\-\\mathcal\{T\}\_\{b\}\(\\rho\_\{0\},\\nu^\{\*\}\)\\;=\\;J\_\{\\mathrm\{stat\}\}\(\\nu^\{\*\}\),where the first equality usesLaw⁡\(x1u∗\)=ν∗\\mathrm\{Law\}\(x\_\{1\}^\{u^\{\*\}\}\)=\\nu^\{\*\}, the second substitutes[Equation61](https://arxiv.org/html/2609.27033#A5.E61)via[Equation59](https://arxiv.org/html/2609.27033#A5.E59), and the third is the definition ofJstatJ\_\{\\mathrm\{stat\}\}\. Sinceu∗u^\{\*\}is feasible,supuJOC​\(u\)≥JOC​\(u∗\)=Jstat​\(ν∗\)=supνJstat​\(ν\)\\sup\_\{u\}J\_\{\\mathrm\{OC\}\}\(u\)\\geq J\_\{\\mathrm\{OC\}\}\(u^\{\*\}\)=J\_\{\\mathrm\{stat\}\}\(\\nu^\{\*\}\)=\\sup\_\{\\nu\}J\_\{\\mathrm\{stat\}\}\(\\nu\)\.

Combining the upper and lower bounds gives the value equality[Equation57](https://arxiv.org/html/2609.27033#A5.E57)\. The lower\-bound construction additionally yields the optimizer\-realization claim: the specific feasible controlu∗=v∗−bu^\{\*\}=v^\{\*\}\-bachievesJOC​\(u∗\)=Jstat​\(ν∗\)J\_\{\\mathrm\{OC\}\}\(u^\{\*\}\)=J\_\{\\mathrm\{stat\}\}\(\\nu^\{\*\}\)by[Equation62](https://arxiv.org/html/2609.27033#A5.E62), this in turn equalssupνJstat​\(ν\)\\sup\_\{\\nu\}J\_\{\\mathrm\{stat\}\}\(\\nu\)by the choice ofν∗\\nu^\{\*\}, and the value equality[Equation57](https://arxiv.org/html/2609.27033#A5.E57)identifies the right\-hand side withsupuJOC​\(u\)\\sup\_\{u\}J\_\{\\mathrm\{OC\}\}\(u\)\. HenceJOC​\(u∗\)=supuJOC​\(u\)J\_\{\\mathrm\{OC\}\}\(u^\{\*\}\)=\\sup\_\{u\}J\_\{\\mathrm\{OC\}\}\(u\)withu∗∈𝒰u^\{\*\}\\in\\mathcal\{U\}, sou∗u^\{\*\}attains the OC supremum, and its terminal lawLaw⁡\(x1u∗\)=ν∗\\mathrm\{Law\}\(x\_\{1\}^\{u^\{\*\}\}\)=\\nu^\{\*\}solves[Equation12](https://arxiv.org/html/2609.27033#S3.E12)\. ∎

### Proof of[Proposition4\.1](https://arxiv.org/html/2609.27033#S4.Thmtheorem1)

See[4\.1](https://arxiv.org/html/2609.27033#S4.Thmtheorem1)

###### Proof\.

We first establish the performance\-difference identity[Equation17](https://arxiv.org/html/2609.27033#S4.E17)and then specialize tow=g¯w=\\bar\{g\}\. Fix a candidate controlwwand an initial pair\(s,x\)\(s,x\), and letxtw=Xs,tw​\(x\)x^\{w\}\_\{t\}=X^\{w\}\_\{s,t\}\(x\)denote the controlled trajectory ofx˙tw=bt​\(xtw\)\+wt​\(xtw\)\\dot\{x\}^\{w\}\_\{t\}=b\_\{t\}\(x^\{w\}\_\{t\}\)\+w\_\{t\}\(x^\{w\}\_\{t\}\)fort∈\[s,1\]t\\in\[s,1\]\.

DifferentiatingVtu¯​\(xtw\)V^\{\\bar\{u\}\}\_\{t\}\(x^\{w\}\_\{t\}\)inttalong the trajectory, applying[LemmaC\.1](https://arxiv.org/html/2609.27033#A3.Thmtheorem1)at\(t,xtw\)\(t,x^\{w\}\_\{t\}\), and usingg¯t​\(xtw\)=∇xVtu¯​\(xtw\)\\bar\{g\}\_\{t\}\(x^\{w\}\_\{t\}\)=\\nabla\_\{x\}V^\{\\bar\{u\}\}\_\{t\}\(x^\{w\}\_\{t\}\)yields

dd​t​Vtu¯​\(xtw\)=12​∥u¯t​\(xtw\)∥2\+\(wt​\(xtw\)−u¯t​\(xtw\)\)⋅g¯t​\(xtw\)\.\\frac\{d\}\{dt\}V^\{\\bar\{u\}\}\_\{t\}\(x^\{w\}\_\{t\}\)=\\tfrac\{1\}\{2\}\\lVert\\bar\{u\}\_\{t\}\(x^\{w\}\_\{t\}\)\\rVert^\{2\}\+\\big\(w\_\{t\}\(x^\{w\}\_\{t\}\)\-\\bar\{u\}\_\{t\}\(x^\{w\}\_\{t\}\)\\big\)\\cdot\\bar\{g\}\_\{t\}\(x^\{w\}\_\{t\}\)\.\(63\)Completing the square inwt​\(xtw\)w\_\{t\}\(x^\{w\}\_\{t\}\),

dd​t​Vtu¯​\(xtw\)=12​∥u¯t​\(xtw\)−g¯t​\(xtw\)∥2−12​∥wt​\(xtw\)−g¯t​\(xtw\)∥2\+12​∥wt​\(xtw\)∥2\.\\frac\{d\}\{dt\}V^\{\\bar\{u\}\}\_\{t\}\(x^\{w\}\_\{t\}\)=\\tfrac\{1\}\{2\}\\lVert\\bar\{u\}\_\{t\}\(x^\{w\}\_\{t\}\)\-\\bar\{g\}\_\{t\}\(x^\{w\}\_\{t\}\)\\rVert^\{2\}\-\\tfrac\{1\}\{2\}\\lVert w\_\{t\}\(x^\{w\}\_\{t\}\)\-\\bar\{g\}\_\{t\}\(x^\{w\}\_\{t\}\)\\rVert^\{2\}\+\\tfrac\{1\}\{2\}\\lVert w\_\{t\}\(x^\{w\}\_\{t\}\)\\rVert^\{2\}\.\(64\)Integrating fromssto11along the trajectory and using the terminal conditionV1u¯​\(x\)=λ​r​\(x\)V^\{\\bar\{u\}\}\_\{1\}\(x\)=\\lambda\\,r\(x\),

λ​r​\(x1w\)−Vsu¯​\(x\)=∫s1\[12​∥u¯t​\(xtw\)−g¯t​\(xtw\)∥2−12​∥wt​\(xtw\)−g¯t​\(xtw\)∥2\+12​∥wt​\(xtw\)∥2\]​𝑑t\.\\lambda\\,r\(x^\{w\}\_\{1\}\)\-V^\{\\bar\{u\}\}\_\{s\}\(x\)=\\int\_\{s\}^\{1\}\\\!\\left\[\\,\\tfrac\{1\}\{2\}\\lVert\\bar\{u\}\_\{t\}\(x^\{w\}\_\{t\}\)\-\\bar\{g\}\_\{t\}\(x^\{w\}\_\{t\}\)\\rVert^\{2\}\-\\tfrac\{1\}\{2\}\\lVert w\_\{t\}\(x^\{w\}\_\{t\}\)\-\\bar\{g\}\_\{t\}\(x^\{w\}\_\{t\}\)\\rVert^\{2\}\+\\tfrac\{1\}\{2\}\\lVert w\_\{t\}\(x^\{w\}\_\{t\}\)\\rVert^\{2\}\\,\\right\]dt\.\(65\)From the definition of the value function,Vsw​\(x\)=λ​r​\(x1w\)−∫s112​∥wt​\(xtw\)∥2​𝑑tV^\{w\}\_\{s\}\(x\)=\\lambda\\,r\(x^\{w\}\_\{1\}\)\-\\int\_\{s\}^\{1\}\\tfrac\{1\}\{2\}\\lVert w\_\{t\}\(x^\{w\}\_\{t\}\)\\rVert^\{2\}\\,dt, so subtracting the integral of12​∥wt​\(xtw\)∥2\\tfrac\{1\}\{2\}\\lVert w\_\{t\}\(x^\{w\}\_\{t\}\)\\rVert^\{2\}from both sides of[Equation65](https://arxiv.org/html/2609.27033#A5.E65)yields the performance\-difference identity[Equation17](https://arxiv.org/html/2609.27033#S4.E17)for any candidate controlww\.

Specializing tow=g¯w=\\bar\{g\}makes the second norm in[Equation17](https://arxiv.org/html/2609.27033#S4.E17)vanish, giving

Vsg¯​\(x\)−Vsu¯​\(x\)=12​∫s1∥u¯t​\(xtg¯\)−g¯t​\(xtg¯\)∥2​𝑑t≥0,V^\{\\bar\{g\}\}\_\{s\}\(x\)\-V^\{\\bar\{u\}\}\_\{s\}\(x\)=\\frac\{1\}\{2\}\\int\_\{s\}^\{1\}\\lVert\\bar\{u\}\_\{t\}\(x^\{\\bar\{g\}\}\_\{t\}\)\-\\bar\{g\}\_\{t\}\(x^\{\\bar\{g\}\}\_\{t\}\)\\rVert^\{2\}\\,dt\\;\\geq\\;0,\(66\)with equality if and only ifu¯t​\(xtg¯\)=g¯t​\(xtg¯\)\\bar\{u\}\_\{t\}\(x^\{\\bar\{g\}\}\_\{t\}\)=\\bar\{g\}\_\{t\}\(x^\{\\bar\{g\}\}\_\{t\}\)for almost everyt∈\[s,1\]t\\in\[s,1\]\. At equality,u¯=g¯=∇xVu¯\\bar\{u\}=\\bar\{g\}=\\nabla\_\{x\}V^\{\\bar\{u\}\}, and substituting this into the policy\-evaluation identity[Equation41](https://arxiv.org/html/2609.27033#A3.E41)of[LemmaC\.1](https://arxiv.org/html/2609.27033#A3.Thmtheorem1)reduces it to

∂tVtu¯​\(x\)\+bt​\(x\)⋅∇xVtu¯​\(x\)\+12​∥∇xVtu¯​\(x\)∥2=0,V1u¯​\(x\)=λ​r​\(x\),\\partial\_\{t\}V^\{\\bar\{u\}\}\_\{t\}\(x\)\+b\_\{t\}\(x\)\\cdot\\nabla\_\{x\}V^\{\\bar\{u\}\}\_\{t\}\(x\)\+\\tfrac\{1\}\{2\}\\lVert\\nabla\_\{x\}V^\{\\bar\{u\}\}\_\{t\}\(x\)\\rVert^\{2\}=0,\\qquad V^\{\\bar\{u\}\}\_\{1\}\(x\)=\\lambda\\,r\(x\),\(67\)which is exactly the reduced HJB equation[Equation40](https://arxiv.org/html/2609.27033#A3.E40)whose unique solution is the optimal value functionV∗V^\{\*\}\. HenceVu¯=V∗V^\{\\bar\{u\}\}=V^\{\*\}andu¯=∇xVu¯=∇xV∗=u∗\\bar\{u\}=\\nabla\_\{x\}V^\{\\bar\{u\}\}=\\nabla\_\{x\}V^\{\*\}=u^\{\*\}by the optimal control formula[Equation39](https://arxiv.org/html/2609.27033#A3.E39)\. ∎

### Proof of[Proposition4\.2](https://arxiv.org/html/2609.27033#S4.Thmtheorem2)

See[4\.2](https://arxiv.org/html/2609.27033#S4.Thmtheorem2)

###### Proof\.

The proof proceeds in three steps\. Step 1 analyzes the off\-diagonal distillation critical condition; with the stop\-gradient target frozen the loss is a convex quadratic functional whose minimum is zero and attained, so every critical point forcesXs,tw^X^\{\\hat\{w\}\}\_\{s,t\}to be the flow map of the velocityw^t,t\\hat\{w\}\_\{t,t\}on the support of the sampling distribution over the upper triangle of the\(s,t\)\(s,t\)plane\. Step 2 analyzes the diagonal regression critical condition; using Step 1’s flow map property it establishes unbiasedness of the Monte Carlo value estimator, and forcesw^t,t=bt\+∇xVtw¯\\hat\{w\}\_\{t,t\}=b\_\{t\}\+\\nabla\_\{x\}V^\{\\bar\{w\}\}\_\{t\}on the on\-policy support\. Step 3 stitches these together via the stop\-gradient consistencyw¯=sg⁡\(w^\)\\bar\{w\}=\\mathrm\{sg\}\\left\(\\hat\{w\}\\right\)and the policy\-iteration fixed\-point characterization of[Proposition4\.1](https://arxiv.org/html/2609.27033#S4.Thmtheorem1)\.

For the proof, we decomposeℒWTF​\(w^\)=ℒdiag​\(w^,w¯\)\+β​ℒdist​\(w^,w¯\)\\mathcal\{L\}\_\{\\mathrm\{WTF\}\}\(\\hat\{w\}\)=\\mathcal\{L\}\_\{\\mathrm\{diag\}\}\(\\hat\{w\};\\bar\{w\}\)\+\\beta\\,\\mathcal\{L\}\_\{\\mathrm\{dist\}\}\(\\hat\{w\};\\bar\{w\}\)withw¯=sg⁡\(w^\)\\bar\{w\}=\\mathrm\{sg\}\\left\(\\hat\{w\}\\right\)\. The variational analysis below uses this factorization throughout, treating every stop\-gradient quantity as held fixed under variations ofw^\\hat\{w\}, including those depending onw¯\\bar\{w\}and the distillation targetsg\(∇Xt,τw^w^t,t\)\\mathrm\{sg\}\\left\(\\nabla X^\{\\hat\{w\}\}\_\{t,\\tau\}\\,\\hat\{w\}\_\{t,t\}\\right\), which is a function ofw^\\hat\{w\}itself\. We work with the Eulerian self\-distillation loss[Equation84](https://arxiv.org/html/2609.27033#A6.E84), but the Lagrangian and progressive variants follow by analogous arguments using the corresponding flow map characterization in[Equation24](https://arxiv.org/html/2609.27033#A1.E24)\. The full\-support hypothesis onps,tp\_\{s,t\}is satisfied by the uniform\-on\-upper\-triangle sampler used in our algorithm, andw^\\hat\{w\}varies over the class𝒲\\mathcal\{W\}of[AssumptionD\.4](https://arxiv.org/html/2609.27033#A4.Thmtheorem4)\. Sinceρ0\\rho\_\{0\}has full support andX0,tw¯X^\{\\bar\{w\}\}\_\{0,t\}is a homeomorphism, the law ofx¯t\\bar\{x\}\_\{t\}has full support onℝd\\mathbb\{R\}^\{d\}for everytt; combined with the full support ofps,tp\_\{s,t\}this makessupp\(ρt,τ,xw¯\)=\{0≤t≤τ≤1\}×ℝd\\mathrm\{supp\}\(\\rho^\{\\bar\{w\}\}\_\{t,\\tau,x\}\)=\\\{0\\leq t\\leq\\tau\\leq 1\\\}\\times\\mathbb\{R\}^\{d\}\.

#### Step 1: Off\-diagonal distillation\.

The Eulerian self\-distillation loss[Equation84](https://arxiv.org/html/2609.27033#A6.E84)reads

ℒdist​\(w^,w¯\)=𝔼x0,t,τ​\[∥∂tXt,τw^​\(x¯t\)\+sg⁡\(∇Xt,τw^​\(x¯t\)​w^t,t​\(x¯t\)\)∥2\],\\mathcal\{L\}\_\{\\mathrm\{dist\}\}\(\\hat\{w\};\\bar\{w\}\)=\\mathbb\{E\}\_\{x\_\{0\},t,\\tau\}\\\!\\left\[\\,\\lVert\\partial\_\{t\}X^\{\\hat\{w\}\}\_\{t,\\tau\}\(\\bar\{x\}\_\{t\}\)\+\\mathrm\{sg\}\\left\(\\,\\nabla X^\{\\hat\{w\}\}\_\{t,\\tau\}\(\\bar\{x\}\_\{t\}\)\\,\\hat\{w\}\_\{t,t\}\(\\bar\{x\}\_\{t\}\)\\,\\right\)\\rVert^\{2\}\\,\\right\],\(68\)with\(t,τ\)\(t,\\tau\)sampled from the off\-diagonal samplerps,tp\_\{s,t\}on the upper triangle andx¯t=X0,tw¯​\(x0\)\\bar\{x\}\_\{t\}=X^\{\\bar\{w\}\}\_\{0,t\}\(x\_\{0\}\)\. Here and below∂t\\partial\_\{t\}differentiates the*first*time argument of the flow map, as in the Eulerian characterization of[Equation24](https://arxiv.org/html/2609.27033#A1.E24)\. Writeρt,τ,xw¯\\rho^\{\\bar\{w\}\}\_\{t,\\tau,x\}for the joint density of\(t,τ,x¯t\)\(t,\\tau,\\bar\{x\}\_\{t\}\)under the upper\-triangle sampler and the on\-policy trajectory, and

c⁡\(t,τ,x\)≔∇Xt,τw^​\(x\)​w^t,t​\(x\)c\(t,\\tau,x\)\\coloneqq\\nabla X^\{\\hat\{w\}\}\_\{t,\\tau\}\(x\)\\,\\hat\{w\}\_\{t,t\}\(x\)\(69\)for the stop\-gradient target, which the semi\-gradient convention holds fixed under variations ofw^\\hat\{w\}\. The off\-diagonal velocitiesw^t,τ\\hat\{w\}\_\{t,\\tau\},t<τt<\\tau, therefore enter[Equation68](https://arxiv.org/html/2609.27033#A5.E68)only through the un\-stopped factor∂tXt,τw^​\(x¯t\)\\partial\_\{t\}X^\{\\hat\{w\}\}\_\{t,\\tau\}\(\\bar\{x\}\_\{t\}\)\. Under the flow map parameterizationXt,τw^​\(x\)=x\+\(τ−t\)​w^t,τ​\(x\)X^\{\\hat\{w\}\}\_\{t,\\tau\}\(x\)=x\+\(\\tau\-t\)\\,\\hat\{w\}\_\{t,\\tau\}\(x\),

∂tXt,τw^​\(x\)=−w^t,τ​\(x\)\+\(τ−t\)​∂tw^t,τ​\(x\),\\partial\_\{t\}X^\{\\hat\{w\}\}\_\{t,\\tau\}\(x\)=\-\\hat\{w\}\_\{t,\\tau\}\(x\)\+\(\\tau\-t\)\\,\\partial\_\{t\}\\hat\{w\}\_\{t,\\tau\}\(x\),\(70\)which is*linear*inw^\\hat\{w\}\. Withccfrozen, the residual in[Equation68](https://arxiv.org/html/2609.27033#A5.E68)is thus affine in the off\-diagonal velocities and

w^⟼ℒdist​\(w^,w¯\)=𝔼x0,t,τ​\[∥∂tXt,τw^​\(x¯t\)\+c⁡\(t,τ,x¯t\)∥2\]\\hat\{w\}\\longmapsto\\mathcal\{L\}\_\{\\mathrm\{dist\}\}\(\\hat\{w\};\\bar\{w\}\)=\\mathbb\{E\}\_\{x\_\{0\},t,\\tau\}\\\!\\left\[\\,\\lVert\\partial\_\{t\}X^\{\\hat\{w\}\}\_\{t,\\tau\}\(\\bar\{x\}\_\{t\}\)\+c\(t,\\tau,\\bar\{x\}\_\{t\}\)\\rVert^\{2\}\\,\\right\]\(71\)is a convex quadratic functional of them\. Every stationary point of a convex quadratic functional is a global minimizer, so it is enough to identify the global minimum\.

That minimum is zero, and it is attained within the parameterized class\. The field

w^t,τ⋆​\(x\)=1τ−t​∫tτc⁡\(σ,τ,x\)​𝑑σ,t<τ,\\hat\{w\}^\{\\star\}\_\{t,\\tau\}\(x\)=\\frac\{1\}\{\\tau\-t\}\\int\_\{t\}^\{\\tau\}c\(\\sigma,\\tau,x\)\\,d\\sigma,\\qquad t<\\tau,\(72\)leaves the diagonal untouched – sinceXτ,τw^=idX^\{\\hat\{w\}\}\_\{\\tau,\\tau\}=\\mathrm\{id\}gives∇Xτ,τw^=I\\nabla X^\{\\hat\{w\}\}\_\{\\tau,\\tau\}=Iand hencec⁡\(τ,τ,x\)=w^τ,τ​\(x\)c\(\\tau,\\tau,x\)=\\hat\{w\}\_\{\\tau,\\tau\}\(x\), so[Equation72](https://arxiv.org/html/2609.27033#A5.E72)extends continuously tow^τ,τ⋆=w^τ,τ\\hat\{w\}^\{\\star\}\_\{\\tau,\\tau\}=\\hat\{w\}\_\{\\tau,\\tau\}– and lies in the same class of two\-time velocity fields asw^\\hat\{w\}, withw^⋆\\hat\{w\}^\{\\star\}inheriting the regularity ofcc\. It satisfiesXt,τw^⋆​\(x\)=x\+∫tτc⁡\(σ,τ,x\)​𝑑σX^\{\\hat\{w\}^\{\\star\}\}\_\{t,\\tau\}\(x\)=x\+\\int\_\{t\}^\{\\tau\}c\(\\sigma,\\tau,x\)\\,d\\sigmaand therefore∂tXt,τw^⋆​\(x\)=−c⁡\(t,τ,x\)\\partial\_\{t\}X^\{\\hat\{w\}^\{\\star\}\}\_\{t,\\tau\}\(x\)=\-c\(t,\\tau,x\), so its residual vanishes identically andℒdist​\(w^⋆,w¯\)=0\\mathcal\{L\}\_\{\\mathrm\{dist\}\}\(\\hat\{w\}^\{\\star\};\\bar\{w\}\)=0\.

This is where the argument differs from a variational\-derivative calculation: rather than dividing out the Jacobian factor∂\(∂tXt,τw^\)/∂w^t,τ\\partial\(\\partial\_\{t\}X^\{\\hat\{w\}\}\_\{t,\\tau\}\)/\\partial\\hat\{w\}\_\{t,\\tau\}and arguing that it is non\-degenerate, convexity of the frozen objective together with realizability of the frozen target gives the conclusion with nothing to invert\. Since a critical point attains the global minimum00of[Equation71](https://arxiv.org/html/2609.27033#A5.E71), the residual vanishesρt,τ,xw¯\\rho^\{\\bar\{w\}\}\_\{t,\\tau,x\}\-almost everywhere:

∂tXt,τw^​\(x\)=−∇Xt,τw^​\(x\)​w^t,t​\(x\),\(t,τ,x\)∈supp⁡\(ρt,τ,xw¯\)\.\\partial\_\{t\}X^\{\\hat\{w\}\}\_\{t,\\tau\}\(x\)=\-\\nabla X^\{\\hat\{w\}\}\_\{t,\\tau\}\(x\)\\,\\hat\{w\}\_\{t,t\}\(x\),\\qquad\(t,\\tau,x\)\\in\\mathrm\{supp\}\\big\(\\rho^\{\\bar\{w\}\}\_\{t,\\tau,x\}\\big\)\.\(73\)Equation[Equation73](https://arxiv.org/html/2609.27033#A5.E73)is the Eulerian transport equation[Equation24](https://arxiv.org/html/2609.27033#A1.E24)characterizingXs,tw^X^\{\\hat\{w\}\}\_\{s,t\}as the flow map generated by the velocity fieldw^t,t\\hat\{w\}\_\{t,t\}\. Because it holds for everyx∈ℝdx\\in\\mathbb\{R\}^\{d\}, it holds in particular along the integral curves ofw^t,t\\hat\{w\}\_\{t,t\}\. Combined with the diagonal initial conditionXt,tw^​\(x\)=xX^\{\\hat\{w\}\}\_\{t,t\}\(x\)=x– which holds for free under the flow map parameterizationXs,tw^​\(x\)=x\+\(t−s\)​w^s,t​\(x\)X^\{\\hat\{w\}\}\_\{s,t\}\(x\)=x\+\(t\-s\)\\,\\hat\{w\}\_\{s,t\}\(x\)– standard uniqueness results for the method of characteristics imply

∂tXs,tw^​\(x\)=w^t,t​\(Xs,tw^​\(x\)\),Xs,tw^=Xσ,tw^∘Xs,σw^,0≤s≤σ≤t≤1,\\partial\_\{t\}X^\{\\hat\{w\}\}\_\{s,t\}\(x\)=\\hat\{w\}\_\{t,t\}\\big\(X^\{\\hat\{w\}\}\_\{s,t\}\(x\)\\big\),\\qquad X^\{\\hat\{w\}\}\_\{s,t\}=X^\{\\hat\{w\}\}\_\{\\sigma,t\}\\circ X^\{\\hat\{w\}\}\_\{s,\\sigma\},\\qquad 0\\leq s\\leq\\sigma\\leq t\\leq 1,\(74\)that is,Xw^X^\{\\hat\{w\}\}is the flow map generated by its own diagonal velocity\. In particular,Xw^X^\{\\hat\{w\}\}– and henceXw¯X^\{\\bar\{w\}\}at the stop\-gradient consistencyw¯=sg⁡\(w^\)\\bar\{w\}=\\mathrm\{sg\}\\left\(\\hat\{w\}\\right\)– satisfies the semigroup property\. We reserve the symboluufor the residual controlw^t,t−bt\\hat\{w\}\_\{t,t\}\-b\_\{t\}introduced in Step 3, so thatXuX^\{u\}keeps its usual meaning as the flow ofbt\+utb\_\{t\}\+u\_\{t\}\.

#### Step 2: Diagonal regression\.

The diagonal regression term reads

ℒdiag​\(w^,w¯\)=𝔼x0,t,τ​\[∥w^t,t​\(x¯t\)−bt​\(x¯t\)−g^​\(t,x¯t,τ\)∥2\],\\mathcal\{L\}\_\{\\mathrm\{diag\}\}\(\\hat\{w\};\\bar\{w\}\)=\\mathbb\{E\}\_\{x\_\{0\},t,\\tau\}\\\!\\left\[\\,\\lVert\\hat\{w\}\_\{t,t\}\(\\bar\{x\}\_\{t\}\)\-b\_\{t\}\(\\bar\{x\}\_\{t\}\)\-\\widehat\{g\}\(t,\\bar\{x\}\_\{t\};\\tau\)\\rVert^\{2\}\\,\\right\],\(75\)witht∼Unif⁡\[0,1\]t\\sim\\mathrm\{Unif\}\[0,1\],τ∼Unif⁡\[t,1\]\\tau\\sim\\mathrm\{Unif\}\[t,1\],x¯t=X0,tw¯​\(x0\)\\bar\{x\}\_\{t\}=X^\{\\bar\{w\}\}\_\{0,t\}\(x\_\{0\}\), andg^​\(t,x,τ\)=∇xV^​\(t,x,τ\)\\widehat\{g\}\(t,x;\\tau\)=\\nabla\_\{x\}\\widehat\{V\}\(t,x;\\tau\)the value\-gradient estimator obtained by differentiating the Monte Carlo estimator[Equation18](https://arxiv.org/html/2609.27033#S4.E18)at the auxiliary sampleτ\\tau\.

The unbiasedness ofV^\\widehat\{V\}as an estimator ofVtw¯​\(x\)V^\{\\bar\{w\}\}\_\{t\}\(x\)relies on the semigroup property ofXw¯X^\{\\bar\{w\}\}established in Step 1\. The construction of[Equation18](https://arxiv.org/html/2609.27033#S4.E18)writesx1=Xτ,1w¯​\(Xt,τw¯​\(x\)\)x\_\{1\}=X^\{\\bar\{w\}\}\_\{\\tau,1\}\(X^\{\\bar\{w\}\}\_\{t,\\tau\}\(x\)\), and the semigroup property collapses this toXt,1w¯​\(x\)X^\{\\bar\{w\}\}\_\{t,1\}\(x\)independently ofτ\\tau\. Combined with the substitution\(1−t\)​𝔼τ​\[∥⋅∥2\]=∫t1∥⋅∥2​𝑑τ\(1\-t\)\\,\\mathbb\{E\}\_\{\\tau\}\[\\lVert\\cdot\\rVert^\{2\}\]=\\int\_\{t\}^\{1\}\\lVert\\cdot\\rVert^\{2\}\\,d\\taufor the control\-cost term, this yields

𝔼τ∼Unif⁡\[t,1\]​\[V^​\(t,x,τ\)\]\\displaystyle\\mathbb\{E\}\_\{\\tau\\sim\\mathrm\{Unif\}\[t,1\]\}\\\!\\left\[\\widehat\{V\}\(t,x;\\tau\)\\right\]=λ​r​\(Xt,1w¯​\(x\)\)−12​∫t1∥w¯τ,τ​\(Xt,τw¯​\(x\)\)−bτ​\(Xt,τw¯​\(x\)\)∥2​𝑑τ\\displaystyle=\\lambda\\,r\\\!\\left\(X^\{\\bar\{w\}\}\_\{t,1\}\(x\)\\right\)\-\\frac\{1\}\{2\}\\int\_\{t\}^\{1\}\\lVert\\bar\{w\}\_\{\\tau,\\tau\}\(X^\{\\bar\{w\}\}\_\{t,\\tau\}\(x\)\)\-b\_\{\\tau\}\(X^\{\\bar\{w\}\}\_\{t,\\tau\}\(x\)\)\\rVert^\{2\}\\,d\\tau\(76\)=Vtw¯​\(x\),\\displaystyle=V^\{\\bar\{w\}\}\_\{t\}\(x\),where the second equality is the definition[Equation15](https://arxiv.org/html/2609.27033#S4.E15)ofVtw¯V^\{\\bar\{w\}\}\_\{t\}with controluτ=w¯τ,τ−bτu\_\{\\tau\}=\\bar\{w\}\_\{\\tau,\\tau\}\-b\_\{\\tau\}\. Differentiating inxxand interchanging differentiation with expectation under the standard regularity conditions yields

𝔼τ​\[g^​\(t,x,τ\)\]=∇xVtw¯​\(x\)\.\\mathbb\{E\}\_\{\\tau\}\\\!\\left\[\\widehat\{g\}\(t,x;\\tau\)\\right\]=\\nabla\_\{x\}V^\{\\bar\{w\}\}\_\{t\}\(x\)\.\(77\)
Now we compute the variational derivative ofℒdiag\\mathcal\{L\}\_\{\\mathrm\{diag\}\}with respect tow^t,t\\hat\{w\}\_\{t,t\}at a test pointxx\. Since the regression targetbt​\(x¯t\)\+g^​\(t,x¯t,τ\)b\_\{t\}\(\\bar\{x\}\_\{t\}\)\+\\widehat\{g\}\(t,\\bar\{x\}\_\{t\};\\tau\)is a function ofw¯\\bar\{w\}only, it is held fixed under variations ofw^t,t\\hat\{w\}\_\{t,t\}\. With its target frozen, the distillation term depends onw^\\hat\{w\}only through its off\-diagonal components, so it contributes nothing to a diagonal variation and the critical\-point condition forℒWTF\\mathcal\{L\}\_\{\\mathrm\{WTF\}\}reduces to that forℒdiag\\mathcal\{L\}\_\{\\mathrm\{diag\}\}\. The variational derivative is

δ​ℒdiagδ​w^t,t​\(x\)=2​𝔼τ​\[w^t,t​\(x\)−bt​\(x\)−g^​\(t,x,τ\)\]​ρw¯​\(t,x\),\\frac\{\\delta\\mathcal\{L\}\_\{\\mathrm\{diag\}\}\}\{\\delta\\hat\{w\}\_\{t,t\}\(x\)\}=2\\,\\mathbb\{E\}\_\{\\tau\}\\\!\\left\[\\hat\{w\}\_\{t,t\}\(x\)\-b\_\{t\}\(x\)\-\\widehat\{g\}\(t,x;\\tau\)\\right\]\\rho^\{\\bar\{w\}\}\(t,x\),\(78\)whereρw¯​\(t,x\)\\rho^\{\\bar\{w\}\}\(t,x\)denotes the joint density of\(t,x¯t\)\(t,\\bar\{x\}\_\{t\}\)undert∼Unif⁡\[0,1\]t\\sim\\mathrm\{Unif\}\[0,1\],x0∼ρ0x\_\{0\}\\sim\\rho\_\{0\}\. By[Equation77](https://arxiv.org/html/2609.27033#A5.E77), the bracketed expectation simplifies tow^t,t​\(x\)−bt​\(x\)−∇xVtw¯​\(x\)\\hat\{w\}\_\{t,t\}\(x\)\-b\_\{t\}\(x\)\-\\nabla\_\{x\}V^\{\\bar\{w\}\}\_\{t\}\(x\)\.

Setting the variational derivative to zero on the support ofρw¯\\rho^\{\\bar\{w\}\}yields

w^t,t​\(x\)=bt​\(x\)\+∇xVtw¯​\(x\),\(t,x\)∈supp⁡\(ρw¯\)\.\\hat\{w\}\_\{t,t\}\(x\)=b\_\{t\}\(x\)\+\\nabla\_\{x\}V^\{\\bar\{w\}\}\_\{t\}\(x\),\\qquad\(t,x\)\\in\\mathrm\{supp\}\(\\rho^\{\\bar\{w\}\}\)\.\(79\)Becausew¯=sg⁡\(w^\)\\bar\{w\}=\\mathrm\{sg\}\\left\(\\hat\{w\}\\right\), the right\-hand side is the value gradient underw^\\hat\{w\}itself, and the on\-policy support coincides with that ofρw^\\rho^\{\\hat\{w\}\}:

w^t,t​\(x\)=bt​\(x\)\+∇xVtw^​\(x\),\(t,x\)∈supp⁡\(ρw^\)\.\\hat\{w\}\_\{t,t\}\(x\)=b\_\{t\}\(x\)\+\\nabla\_\{x\}V^\{\\hat\{w\}\}\_\{t\}\(x\),\\qquad\(t,x\)\\in\\mathrm\{supp\}\(\\rho^\{\\hat\{w\}\}\)\.\(80\)

#### Step 3: Joint consistency closes the loop\.

From[Equation80](https://arxiv.org/html/2609.27033#A5.E80),

w^t,t​\(x\)=bt​\(x\)\+∇xVtw^​\(x\)on​supp​\(ρw^\)\.\\hat\{w\}\_\{t,t\}\(x\)=b\_\{t\}\(x\)\+\\nabla\_\{x\}V^\{\\hat\{w\}\}\_\{t\}\(x\)\\qquad\\text\{on \}\\mathrm\{supp\}\(\\rho^\{\\hat\{w\}\}\)\.\(81\)Hereρw^\\rho^\{\\hat\{w\}\}denotes the joint law of\(t,X0,tw^​\(x0\)\)\(t,X^\{\\hat\{w\}\}\_\{0,t\}\(x\_\{0\}\)\)undert∼Unif⁡\[0,1\]t\\sim\\mathrm\{Unif\}\[0,1\]andx0∼ρ0x\_\{0\}\\sim\\rho\_\{0\}, which is the lawρw¯\\rho^\{\\bar\{w\}\}of[Equation78](https://arxiv.org/html/2609.27033#A5.E78)evaluated at the stop\-gradient consistencyw¯=sg⁡\(w^\)\\bar\{w\}=\\mathrm\{sg\}\\left\(\\hat\{w\}\\right\)\. From[Equation74](https://arxiv.org/html/2609.27033#A5.E74),Xs,tw^X^\{\\hat\{w\}\}\_\{s,t\}is the flow map generated by the velocityw^t,t\\hat\{w\}\_\{t,t\}\. Define the residual controlut≔w^t,t−btu\_\{t\}\\coloneqq\\hat\{w\}\_\{t,t\}\-b\_\{t\}, so thatw^t,t=bt\+ut\\hat\{w\}\_\{t,t\}=b\_\{t\}\+u\_\{t\}is the controlled drift andXs,tw^=Xs,tuX^\{\\hat\{w\}\}\_\{s,t\}=X^\{u\}\_\{s,t\}; accordingly we writeρu=ρw^\\rho^\{u\}=\\rho^\{\\hat\{w\}\}for the on\-policy law\.

Substitutingw^t,t=bt\+ut\\hat\{w\}\_\{t,t\}=b\_\{t\}\+u\_\{t\}into[Equation81](https://arxiv.org/html/2609.27033#A5.E81)gives the on\-policy fixed\-point condition

ut​\(x\)=∇xVtu​\(x\),\(t,x\)∈supp⁡\(ρu\),u\_\{t\}\(x\)=\\nabla\_\{x\}V^\{u\}\_\{t\}\(x\),\\qquad\(t,x\)\\in\\mathrm\{supp\}\(\\rho^\{u\}\),\(82\)whereVtu≡Vtw^V^\{u\}\_\{t\}\\equiv V^\{\\hat\{w\}\}\_\{t\}since the value function depends only on the controlled driftbt\+utb\_\{t\}\+u\_\{t\}\.

Sinceρ0\\rho\_\{0\}has full support onℝd\\mathbb\{R\}^\{d\}and the controlled flowX0,tuX^\{u\}\_\{0,t\}is a diffeomorphism under standard regularity conditions, the on\-policy densityρu\\rho^\{u\}has full support\. Hence[Equation82](https://arxiv.org/html/2609.27033#A5.E82)givesut=∇xVtuu\_\{t\}=\\nabla\_\{x\}V^\{u\}\_\{t\}everywhere, and the equality clause of[Proposition4\.1](https://arxiv.org/html/2609.27033#S4.Thmtheorem1)\(withu¯=u\\bar\{u\}=uandg¯=∇xVtu\\bar\{g\}=\\nabla\_\{x\}V^\{u\}\_\{t\}\) yieldsu=u∗u=u^\{\*\}, the optimal control of[Equation14](https://arxiv.org/html/2609.27033#S3.E14)\.

Combining the above arguments, we conclude thatw^t,t=bt\+ut∗\\hat\{w\}\_\{t,t\}=b\_\{t\}\+u^\{\*\}\_\{t\}and thatXs,tw^X^\{\\hat\{w\}\}\_\{s,t\}is the flow map ofbt\+ut∗b\_\{t\}\+u^\{\*\}\_\{t\}\. ∎

## Appendix FAlgorithmic aspects

We next describe the implementation of the WTF objective[Equation19](https://arxiv.org/html/2609.27033#S4.E19), summarized in[Algorithm1](https://arxiv.org/html/2609.27033#algorithm1)\.

#### Euclidean reward gradient\.

The exact reward contribution to the value gradient is∇Xt,1\(xt\)⊤∇r\(x1\)\\nabla X\_\{t,1\}\(x\_\{t\}\)^\{\\top\}\\nabla r\(x\_\{1\}\)\. For our image experiments, we instead use∇r​\(x1\)\\nabla r\(x\_\{1\}\)while leaving the endpointx1=Xt,1​\(xt\)x\_\{1\}=X\_\{t,1\}\(x\_\{t\}\)unchanged\. This removes the flow map Jacobian from the reward gradient and computes the update in the direction that increases reward at the generated endpoint\. Following the terminology of[Huang et al\. \[26\]](https://arxiv.org/html/2609.27033#bib.bib55), we refer to this as the*Euclidean*reward gradient\. For the residual parameterizationXt,1​\(x\)=x\+\(1−t\)​wt,1​\(x\)X\_\{t,1\}\(x\)=x\+\(1\-t\)\\,w\_\{t,1\}\(x\), we implement it as

x1=xt\+sg⁡\(Xt,1​\(xt\)−xt\)\.x\_\{1\}=x\_\{t\}\+\\mathrm\{sg\}\\left\(X\_\{t,1\}\(x\_\{t\}\)\-x\_\{t\}\\right\)\.This leaves the endpoint unchanged while omitting the flow map Jacobian from the backward pass\. Related updates have been used for diffusion reward fine\-tuning, including DRTune\[[71](https://arxiv.org/html/2609.27033#bib.bib69)\]\.[Algorithm1](https://arxiv.org/html/2609.27033#algorithm1)shows both gradient choices, and[SectionH\.2](https://arxiv.org/html/2609.27033#A8.SS2)compares them empirically\.

#### Parameterization\.

The derivation writes the control as a residualuuon the base driftbb\. We fine\-tune a single flow map

Xs,tw^​\(x\)=x\+\(t−s\)​w^s,t​\(x\),X^\{\\hat\{w\}\}\_\{s,t\}\(x\)=x\+\(t\-s\)\\,\\hat\{w\}\_\{s,t\}\(x\),\(83\)The network is initialized from the pre\-trained mean velocityvs,tv\_\{s,t\}\. The residual is recovered asus,t=w^s,t−vs,tu\_\{s,t\}=\\hat\{w\}\_\{s,t\}\-v\_\{s,t\}and is used only in the objective\.[SectionF\.1](https://arxiv.org/html/2609.27033#A6.SS1)gives the corresponding diffusion\-time convention\.

#### Off\-diagonal self\-distillation\.

The off\-diagonal regularizerℒdist\\mathcal\{L\}\_\{\\mathrm\{dist\}\}in[Equation19](https://arxiv.org/html/2609.27033#S4.E19)can be any of the standard self\-distillation losses for flow maps; we use the Eulerian \(mean\-flow\) variant\[[21](https://arxiv.org/html/2609.27033#bib.bib9),[22](https://arxiv.org/html/2609.27033#bib.bib45),[36](https://arxiv.org/html/2609.27033#bib.bib10),[25](https://arxiv.org/html/2609.27033#bib.bib3)\]

ℒdist​\(w^\)=𝔼⁡\[∥∂tXt,τw^​\(xt\)\+sg⁡\(∇Xt,τw^​\(xt\)​w^t,t​\(xt\)\)∥2\],\\mathcal\{L\}\_\{\\mathrm\{dist\}\}\(\\hat\{w\}\)=\\mathbb\{E\}\\\!\\left\[\\lVert\\partial\_\{t\}X^\{\\hat\{w\}\}\_\{t,\\tau\}\(x\_\{t\}\)\+\\mathrm\{sg\}\\left\(\\nabla X^\{\\hat\{w\}\}\_\{t,\\tau\}\(x\_\{t\}\)\\,\\hat\{w\}\_\{t,t\}\(x\_\{t\}\)\\right\)\\rVert^\{2\}\\right\],\(84\)wherext=sg⁡\(X0,tw¯​\(x0\)\)x\_\{t\}=\\mathrm\{sg\}\\left\(X^\{\\bar\{w\}\}\_\{0,t\}\(x\_\{0\}\)\\right\)is sampled on\-policy along the current flow andsg⁡\(⋅\)\\mathrm\{sg\}\\left\(\\cdot\\right\)denotes a stop\-gradient\. We use the Eulerian self\-distillation objective, matching the objective used to pre\-train the base models\.

#### Reward\-gradient scaling\.

Because the base velocity and reward gradient can have very different numerical scales, we normalize the reward gradient by their norm ratio and useλeff=λ​κ\\lambda\_\{\\mathrm\{eff\}\}=\\lambda\\kappain[Equation19](https://arxiv.org/html/2609.27033#S4.E19)\.[SectionF\.2](https://arxiv.org/html/2609.27033#A6.SS2)gives the definition, and[Table9](https://arxiv.org/html/2609.27033#A7.T9)reports the raw values ofλ\\lambda\.

### Diffusion\-time convention

Many pre\-trained checkpoints, including score\-based, variance\-preserving, and EDM\-style models\[[34](https://arxiv.org/html/2609.27033#bib.bib18)\], use timet=1t=1for the Gaussian base distribution andt=0t=0for the data distribution\. The reward is evaluated att=0t=0, and the signs and integration limits differ from the convention in the main text\. We state the resulting objective and value\-gradient target below\.

Letbtb\_\{t\}denote a pre\-trained velocity for which integratingx˙t=bt​\(xt\)\\dot\{x\}\_\{t\}=b\_\{t\}\(x\_\{t\}\)fromt=1t=1tot=0t=0transports a Gaussian samplex1∼𝒩⁡\(0,I\)x\_\{1\}\\sim\\mathcal\{N\}\(0,I\)into a data samplex0∼ρ∗x\_\{0\}\\sim\\rho^\{\*\}, and letXs,tu​\(x\)X\_\{s,t\}^\{u\}\(x\)denote the state at timettunder the controlled ODEx˙τu=bτ​\(xτu\)\+uτ​\(xτu\)\\dot\{x\}\_\{\\tau\}^\{u\}=b\_\{\\tau\}\(x\_\{\\tau\}^\{u\}\)\+u\_\{\\tau\}\(x\_\{\\tau\}^\{u\}\)started fromxxat timess\. In this convention, one typically hast≤st\\leq s\. The OC problem[Equation14](https://arxiv.org/html/2609.27033#S3.E14)becomes

supu𝔼x1∼𝒩⁡\(0,I\)​\[λ​r​\(x0u\)−12​∫01∥ut​\(xtu\)∥2​dt\],\\displaystyle\\sup\_\{u\}\\mathbb\{E\}\_\{x\_\{1\}\\sim\\mathcal\{N\}\(0,I\)\}\\\!\\left\[\\,\\lambda\\,r\(x\_\{0\}^\{u\}\)\-\\frac\{1\}\{2\}\\int\_\{0\}^\{1\}\\lVert u\_\{t\}\(x\_\{t\}^\{u\}\)\\rVert^\{2\}\\,dt\\,\\right\],\(85\)subject tox˙tu=bt​\(xtu\)\+ut​\(xtu\),x1u=x1\.\\displaystyle\\text\{subject to\}\\quad\\dot\{x\}\_\{t\}^\{u\}=b\_\{t\}\(x\_\{t\}^\{u\}\)\+u\_\{t\}\(x\_\{t\}^\{u\}\),\\quad x\_\{1\}^\{u\}=x\_\{1\}\.Only the terminal\-time indexing changes relative to[Equation14](https://arxiv.org/html/2609.27033#S3.E14), since the control\-cost integrand is orientation\-free\. We note that the integral of the control runs opposite to the direction of integration of the flow, which ensures the control cost remains positive\.

The value function under controluuis

Vtu​\(x\):=λ​r​\(Xt,0u​\(x\)\)−12​∫0t∥uτ​\(Xt,τu​\(x\)\)∥22​𝑑τ,V^\{u\}\_\{t\}\(x\):=\\lambda\\,r\\\!\\left\(X\_\{t,0\}^\{u\}\(x\)\\right\)\-\\frac\{1\}\{2\}\\int\_\{0\}^\{t\}\\lVert u\_\{\\tau\}\(X\_\{t,\\tau\}^\{u\}\(x\)\)\\rVert\_\{2\}^\{2\}\\,d\\tau,\(86\)where again the integration domain is the unsigned physical interval\[0,t\]\[0,t\]so that the control cost is accumulated as a positive quantity\. Applying dynamic programming with infinitesimal stept↦t−ϵt\\mapsto t\-\\epsilontoward the data gives the Hamilton–Jacobi–Bellman equation

∂tVt∗​\(x\)\+bt​\(x\)⋅∇xVt∗​\(x\)−12​∥∇xVt∗​\(x\)∥22=0,V0∗​\(x\)=λ​r​\(x\),\\partial\_\{t\}V^\{\*\}\_\{t\}\(x\)\+b\_\{t\}\(x\)\\cdot\\nabla\_\{x\}V^\{\*\}\_\{t\}\(x\)\-\\frac\{1\}\{2\}\\lVert\\nabla\_\{x\}V^\{\*\}\_\{t\}\(x\)\\rVert\_\{2\}^\{2\}=0,\\qquad V^\{\*\}\_\{0\}\(x\)=\\lambda\\,r\(x\),\(87\)with the optimal control

ut∗​\(x\)=−∇xVt∗​\(x\)\.u\_\{t\}^\{\*\}\(x\)=\-\\nabla\_\{x\}V^\{\*\}\_\{t\}\(x\)\.\(88\)The minus sign in[Equation88](https://arxiv.org/html/2609.27033#A6.E88), in contrast tout∗=\+∇xVt∗u^\{\*\}\_\{t\}=\+\\nabla\_\{x\}V^\{\*\}\_\{t\}in the data\-terminal convention, is the consequence of time reversal\. Because integration of the generative process flows backwards in time, this negative sign implements gradientascentat inference, matching the forward\-time result\.

The single\-sample Monte Carlo estimator of[Equation86](https://arxiv.org/html/2609.27033#A6.E86)also changes accordingly\. Fix a frozen referenceu¯\\bar\{u\}, sampleτ∼Unif⁡\[0,t\]\\tau\\sim\\mathrm\{Unif\}\[0,t\], and computexτ=Xt,τu¯​\(x\)x\_\{\\tau\}=X\_\{t,\\tau\}^\{\\bar\{u\}\}\(x\)andx0=Xτ,0u¯​\(xτ\)x\_\{0\}=X\_\{\\tau,0\}^\{\\bar\{u\}\}\(x\_\{\\tau\}\)to form

V^t​\(x\)=λ​r​\(x0\)−t2​∥u¯τ​\(xτ\)∥22,\\widehat\{V\}\_\{t\}\(x\)=\\lambda\\,r\(x\_\{0\}\)\-\\frac\{t\}\{2\}\\lVert\\bar\{u\}\_\{\\tau\}\(x\_\{\\tau\}\)\\rVert\_\{2\}^\{2\},\(89\)where the factorttis the length of the remaining interval\[0,t\]\[0,t\], replacing the factor1−t1\-tfrom[Equation18](https://arxiv.org/html/2609.27033#S4.E18)\. Differentiation yields the gradient targetg^t​\(x\)=∇xV^t​\(x\)\\widehat\{g\}\_\{t\}\(x\)=\\nabla\_\{x\}\\widehat\{V\}\_\{t\}\(x\), and the diagonal value gradient regression of[Equation19](https://arxiv.org/html/2609.27033#S4.E19)fitsw^t,t​\(x¯t\)→bt​\(x¯t\)−g^t​\(x¯t\)\\hat\{w\}\_\{t,t\}\(\\bar\{x\}\_\{t\}\)\\to b\_\{t\}\(\\bar\{x\}\_\{t\}\)\-\\widehat\{g\}\_\{t\}\(\\bar\{x\}\_\{t\}\)along the reference trajectoryx¯t=X1,tu¯​\(x1\)\\bar\{x\}\_\{t\}=X\_\{1,t\}^\{\\bar\{u\}\}\(x\_\{1\}\), with the sign flip ong^t\\widehat\{g\}\_\{t\}matching[Equation88](https://arxiv.org/html/2609.27033#A6.E88)\.

### Reward\-gradient scaling

The reward gradient and the base flow velocity gradient typically live on different numerical scales\. The base flowbtb\_\{t\}has been trained to pushρ0\\rho\_\{0\}all the way toρ1\\rho\_\{1\}in unit time, so the magnitude ofbtb\_\{t\}is set by the dataset and the training schedule\. Hand\-coded or learned rewards typically have gradient magnitudes set by an unrelated convention\. Without rescaling,λ\\lambdain the value gradient[Equation18](https://arxiv.org/html/2609.27033#S4.E18)must be tuned by orders of magnitude per reward to balanceλ∇r\\lambda\\nabla ragainst the implicit scale of the base velocity field\. To normalize this in standardized units, we rescale the reward gradient by the ratio of the base flow velocity norm to the reward gradient norm:

κ=𝔼⁡\[∥bt​\(x¯t\)∥2∥∇xr​\(x¯1\)∥2\],λeff=λ⋅κ,\\kappa=\\mathbb\{E\}\\left\[\\frac\{\\lVert b\_\{t\}\(\\bar\{x\}\_\{t\}\)\\rVert\_\{2\}\}\{\\lVert\\nabla\_\{x\}r\(\\bar\{x\}\_\{1\}\)\\rVert\_\{2\}\}\\right\],\\qquad\\lambda\_\{\\mathrm\{eff\}\}=\\lambda\\cdot\\kappa,\(90\)and useλeff\\lambda\_\{\\mathrm\{eff\}\}in place ofλ\\lambdainside[Equation18](https://arxiv.org/html/2609.27033#S4.E18)\. We track the two norms with exponential moving averages and do not freeze them after warmup\. This requires one reward backward pass and one evaluation ofbtb\_\{t\}per iteration\.[SectionH\.3](https://arxiv.org/html/2609.27033#A8.SS3)sweepsλ\\lambdaon text\-to\-image\.

## Appendix GImplementation details

### Training setup

#### ImageNet\.

We fine\-tune DMF XL/2\[[65](https://arxiv.org/html/2609.27033#bib.bib35)\]end to end using HPSv2 as the reward, starting from the pre\-traineddmf\_xl\_2\_256checkpoint at256×256256\\times 256resolution over the full10001000\-class conditioning\. We update the full network without an adapter\. A single parameterization is used forw^t,t\\hat\{w\}\_\{t,t\}andw^s,t\\hat\{w\}\_\{s,t\}\. Training usesbf16mixed precision for12,50012\{,\}500optimizer steps, with the EMA weights as the frozen referencew¯\\bar\{w\}in[Equation19](https://arxiv.org/html/2609.27033#S4.E19)\. The text\-to\-image experiments instead use a stop\-gradient copy of the current weights\.

#### Text\-to\-image\.

We fine\-tune the pre\-trained TiM\-T2I model\[[27](https://arxiv.org/html/2609.27033#bib.bib67)\]at512×512512\\times 512resolution using thetim\_xl\_p1\_t2iconfiguration\. Images are represented on a16×1616\\times 16latent grid with3232channels from a frozen deep compression autoencoder\[[72](https://arxiv.org/html/2609.27033#bib.bib72)\]\(mit\-han\-lab/dc\-ae\-f32c32\-sana\-1\.1\-diffusers\), which downsamples by3232spatially to3232latent channels\. Captions are encoded by a frozen Gemma 3 1B instruction\-tuned text encoder\[[73](https://arxiv.org/html/2609.27033#bib.bib73)\]\(google/gemma\-3\-1b\-it\) with a maximum sequence length of256256\. We train LoRA adapters on all attention\(Q,K,V,out\)\(Q,K,V,\\text\{out\}\)and MLP\(fc1,fc2\)\(\\text\{fc\}\_\{1\},\\text\{fc\}\_\{2\}\)projections in every transformer block\. Prompts are drawn from the photo and painting categories of HPDv2\.

#### Fine\-tuning\.

A single LoRA adapter of rank1616with scalingα=16\\alpha=16is used for the diagonal update and all off\-diagonal flow map evaluations\. We train for1,4601\{,\}460optimizer steps and use the Euclidean reward gradient described in[AppendixF](https://arxiv.org/html/2609.27033#A6)\.

#### Reward model\.

Reward gradients use HPSv2 v2\.1 with the fine\-tunedxswu/HPSv2weights on an OpenCLIP\[[74](https://arxiv.org/html/2609.27033#bib.bib74)\]CLIP ViT\-H\-14 backbone\. Images in\[0,1\]\[0,1\]are processed at224224pixels with the HPSv2 mask\-aware normalization and resize pipeline\.

#### Optimization\.

Both benchmarks use AdamW\. ImageNet uses\(β1,β2\)=\(0\.9,0\.95\)\(\\beta\_\{1\},\\beta\_\{2\}\)=\(0\.9,0\.95\)with weight decay00, so its update coincides with Adam; text\-to\-image uses\(β1,β2\)=\(0\.9,0\.999\)\(\\beta\_\{1\},\\beta\_\{2\}\)=\(0\.9,0\.999\)with weight decay10−210^\{\-2\}\. The learning rate is constant\.[Table9](https://arxiv.org/html/2609.27033#A7.T9)quotes the global batch, together with the per\-GPU count and accumulation factor that produce it across the88GPUs\. The EMA decayμ\\muaverages the trainable parameters of the fine\-tuned model, updated once per optimizer step\. For text\-to\-image, we clip each per\-sample reward gradient to the0\.80\.8quantile of the batch gradient norms\. For both benchmarks, we apply a global parameter\-gradientL2L\_\{2\}clip of11\. The reward scaleλ\\lambdain[Table9](https://arxiv.org/html/2609.27033#A7.T9)is the raw coefficient before the normalization in[SectionF\.2](https://arxiv.org/html/2609.27033#A6.SS2)\. Every sample contributes to both the diagonal value\-gradient term and the regularizer;pdiagp\_\{\\mathrm\{diag\}\}mixes the times at which the regularizer is applied rather than routing samples between losses\. The text\-to\-image runs use a different regularizer and carry no such mixing probability\.[Table9](https://arxiv.org/html/2609.27033#A7.T9)gives the configuration for both benchmarks\.

### Sampling and evaluation

#### ImageNet\.

We evaluate at NFE∈\{1,250\}\\in\\\{1,250\\\}with CFG scale1\.01\.0\. Using the first3232ImageNet classes, we generate1616samples per class, conditioned on the class label\. We use a fixed evaluation seed and share initial latents across methods\. HPSv2, PickScore, and ImageReward are averaged over the same512512images\. Diversity is the class\-averaged mean pairwise squared distance in DreamSim and CLIP embedding space\. WTF results are means over three matched training seeds\. We use the first3232ImageNet classes, fixed across all methods\.

#### Text\-to\-image\.

We evaluate at NFE∈\{1,2,4,8,50\}\\in\\\{1,2,4,8,50\\\}with CFG scale2\.52\.5and share starting latents across methods\. The evaluation pool contains100100fixed prompts from the photo and painting categories of HPDv2\. Reward metrics use1,0001\{,\}000scored generations drawn from this pool under a fixed evaluation seed, so every method is scored on the identical prompt sequence\. Diversity is computed from1616samples for each of a fixed subset of3232prompts using the same DreamSim and CLIP pairwise distances\. Each sampler transition uses one network evaluation\. A joint conditional and unconditional forward pass counts as one NFE\. ImageNet uses CFG scale1\.01\.0without an unconditional branch\.

Table 9:Implementation details for both benchmarks\.Reward scales are the raw coefficientλ\\lambda; The applied weight isλeff\\lambda\_\{\\mathrm\{eff\}\}from[SectionF\.2](https://arxiv.org/html/2609.27033#A6.SS2)\.

### Baseline implementations

Both baselines use the sampling and evaluation protocol of[SectionG\.2](https://arxiv.org/html/2609.27033#A7.SS2)\. Training stops when the evaluation reward does not improve for six consecutive checkpoints\.

Flow\-GRPO uses group size2424with one prompt per rank,192192samples per iteration,1010SDE steps, LoRA rank3232withα=32\\alpha=32, learning rate10−410^\{\-4\}, and weight decay10−210^\{\-2\}\. Advantages use a per\-prompt mean and a globally gathered reward standard deviation, clipped to\[−5,5\]\[\-5,5\]\. The KL penalty is the analytic same\-variance Gaussian transition KL against the frozen backbone, with weight0\.010\.01\. Each rollout batch is followed by one optimizer update over all1010timesteps, and applies no EMA to the adapter weights\.

Adjoint Matching\[[16](https://arxiv.org/html/2609.27033#bib.bib37)\]uses reward scale1\.2×1051\.2\\times 10^\{5\},N=40N=40rollout steps,K=20K=20timesteps per update, LoRA rank88, learning rate2×10−52\\times 10^\{\-5\}, and effective batch2828\. Half of theKKtimesteps are drawn uniformly from the first three quarters of the trajectory; the remaining1010are the final1010timesteps\. Adapter parameters are float32\. Loss terms whose norm exceeds an exponential moving average of the globally gathered0\.90\.9quantile are masked\.

### Training cost and compute comparison

We compare fine\-tuning compute on the same8×H1008\\times\\mathrm\{H100\}node\.[Table10](https://arxiv.org/html/2609.27033#A7.T10)reports GPU\-hours to each displayed checkpoint and excludes pre\-training, evaluation, and post\-hoc flow map distillation\. The text\-to\-image totals are approximate and use the median time between consecutive training logs\.[Table10](https://arxiv.org/html/2609.27033#A7.T10)reports the cost to each displayed checkpoint, whereas[Figure6](https://arxiv.org/html/2609.27033#S6.F6)reports the compute WTF requires to reach each baseline’s peak reward\. These comparisons use different endpoints, so their ratios are not directly commensurable\.

Table 10:Fine\-tuning compute\.GPU\-hours to the reported checkpoint on the same8×H1008\\times\\mathrm\{H100\}node\. Post\-hoc flow map distillation is excluded\.

## Appendix HAdditional experimental results

### Qualitative comparison on ImageNet

[Figure11](https://arxiv.org/html/2609.27033#A8.F11)compares samples against the base model and adjoint matching at a fixed initial latent\.

![Refer to caption](https://arxiv.org/html/2609.27033v1/figures/overview_imagenet_v2.jpg)Figure 11:Qualitative results on ImageNet\.Samples at 250 NFE from the base model, adjoint matching, and WTF, with the initial latent fixed down each column\. Adjoint matching stays close to the base sample\. WTF changes composition and color while keeping the class\.
### Exact versus Euclidean reward gradient

The exact reward contribution uses the flow map pullbackJ⊤Xt,1∇rJ^\{\\top\}\_\{X\_\{t,1\}\}\\nabla r, while the Euclidean estimator replaces it with∇r\\nabla r\.[Table12](https://arxiv.org/html/2609.27033#A8.T12)compares the two estimators on text\-to\-image at matched iterations\. The Euclidean estimator attains higher HPSv2 and higher DreamSim and CLIP diversity\.

Table 12:Exact versus Euclidean reward gradients on text\-to\-image\.Both rows are evaluated at800800optimizer steps under one protocol \(5050NFE, guidance2\.52\.5,1,0001\{,\}000quality prompts,32×1632\\times 16diversity prompts\)\. A matched comparison over three seeds at600600iterations andλ=25\\lambda=25moves HPSv2 from0\.3590\.359to0\.3760\.376and DreamSim from0\.1540\.154to0\.2170\.217, in the same direction and of comparable size\.
### Reward scale and reward\-diversity tradeoff

We retrain the text\-to\-image model over a4×4\\timesrange of reward scales at matched training steps and evaluation protocol\. Increasingλ\\lambdaincreases HPSv2 while decreasing DreamSim and CLIP diversity\. Increasingλ\\lambdafrom1010to4040raises HPSv2 by0\.0230\.023and lowers DreamSim and CLIP diversity by0\.0220\.022and0\.0130\.013, respectively\.

## Appendix ISynthetic experiments: transport versus reweighting

We compare the WTF and KL population optima on settings where both laws can be computed exactly\. The experiments isolate how transport and reward reweighting produce different terminal laws\.

### Gaussian benchmark

The base isρ0=ρ1=𝒩⁡\(0,1\)\\rho\_\{0\}=\\rho\_\{1\}=\\mathcal\{N\}\(0,1\)in one dimension, and the reward is a bounded bumprK​\(y\)=exp⁡\(−12​\(\(y−K\)/w\)2\)r\_\{K\}\(y\)=\\exp\(\-\\tfrac\{1\}\{2\}\(\(y\-K\)/w\)^\{2\}\)of widthw=2\.25w=2\.25centered atKKbase standard deviations from the mean\. The width is chosen so that the reward overlaps the base appreciably rather than sitting in its tail: atK=4\.75K=4\.75the base already attains𝔼ρ1​\[r\]=0\.142\\mathbb\{E\}\_\{\\rho\_\{1\}\}\[r\]=0\.142, and15\.8%15\.8\\%of the base mass hasr\>0\.25r\>0\.25\.

Both laws are evaluated exactly rather than sampled\. We compute the KL tilt by numerical quadrature\. For the affine base drift, the prior\-action cost has a closed form, so the WTF optimum is obtained from the pointwise maximization in[Proposition3\.1](https://arxiv.org/html/2609.27033#S3.Thmtheorem1)and its terminal density by change of variables\.[Figure4](https://arxiv.org/html/2609.27033#S6.F4)and[Table14](https://arxiv.org/html/2609.27033#A9.T14)report the resulting population quantities\.[Figure13](https://arxiv.org/html/2609.27033#A9.F13)quantifies the same comparison across the reward centreKKand the reward scaleλ\\lambda\.

Figure 13:Quantitative sweeps for the one\-dimensional comparison\.Left: expected reward against the reward centreKKat fixed reward scale\. Right: expected reward against the reward scaleλ\\lambdaat fixed reward centre\. Both quantify the trend that[Figure4](https://arxiv.org/html/2609.27033#S6.F4)shows qualitatively\.Atλ=7\\lambda=7the difference between the two laws widens monotonically as the reward moves outward, from0\.0680\.068atK=2K=2to0\.2560\.256atK=5K=5\.

Table 14:Expected reward as the reward moves outward\.The reward bump of widthw=2\.25w=2\.25is centeredKKbase standard deviations from the mean\. Both fine\-tuned laws are population optima, the tilt by quadrature and WTF in closed form, so no seed variation enters\.
### Qualitative comparison on text\-to\-image

The figures below show text\-to\-image samples across inference budgets, with the prompts listed in each caption\.

#### Prompts and classes in[Figure1](https://arxiv.org/html/2609.27033#S0.F1)\.

The text\-to\-image rows of[Figure1](https://arxiv.org/html/2609.27033#S0.F1)use, from top to bottom, “The image depicts the god dreaming at the end of time\.”, “An otherworldly world depicted with vivid colors by Fuco Ueda\.”, and “A skull\-shaped island with rocks and vegetation, painted by Ghibli with strong light and shadow\.” The ImageNet rows are class\-conditional rather than prompted, and show the classes*platypus*,*Shetland sheepdog*and*volcano*\.

![Refer to caption](https://arxiv.org/html/2609.27033v1/figures/appx_t2i_1.jpg)Figure 15:Reward\-aligned text\-to\-image samples across inference budgets \(first set\)\.Each column is one prompt and each row an inference budget, with the pre\-trained model in the top row\. Within a column every WTF row uses the same initial noise, so the sequence shows one sample refining as the budget grows rather than independent draws\. Prompts, left to right: \(a\) “A zentangle pizza illustration with colorful ink\. \(b\) “Cross section of an apple in a limited neutral palette with a beautiful graphic design and a painterly style\. \(c\) “Scary African voodoo paintings by Jean\-Michel Basquiat\. \(d\) “A digital painting of the legendary water city of Atlantis, featuring a Greek temple, statues, and a red flag\.” \(e\) “A pencil sketch of Danny Devito by Milt Kahl\. \(f\) “The image features a surreal fox and skulls in highly detailed, liquid oilpaint style\. \(g\) “An art piece by Wojciech Siudmak depicting an individual gazing at the vast cosmos\. \(h\) “A cobblestone street with a tree over the sea at sunset, illuminated by sun rays\.![Refer to caption](https://arxiv.org/html/2609.27033v1/figures/appx_t2i_2.jpg)Figure 16:Reward\-aligned text\-to\-image samples across inference budgets \(second set\)\.Each column is one prompt and each row an inference budget, with the pre\-trained model in the top row\. Within a column every WTF row uses the same initial noise, so the sequence shows one sample refining as the budget grows rather than independent draws\. Prompts, left to right: \(a\) “A painting of a firefall cascading over a high cliff\. \(b\) “An image depicting the concept of yin and yang\. \(c\) “Psytrance artwork by Lee Madgwick\. \(d\) “Artwork depicting a futuristic car, created by Ed Roth\. \(e\) “A night scene of a lavender field with a town and church in the background, reminiscent of Van Gogh\. \(f\) “An image depicting the concept of yin and yang\. \(g\) “Portrait of a creature with bat ears, a wolf snout and eagle features, wearing a poncho and helmet\. \(h\) “The image is a drawing of a skeletal, frail figure driving a chariot pulled by two skeletal hounds\.

相似文章

Flow-Map GRPO:基于锚定随机组合的少步流图生成器强化学习

arXiv cs.LG

提出了Flow-Map GRPO,一种用于确定性少步流图生成器的在线RL后训练框架,引入了锚定随机流图组合(ASFMC)以在不改变原始模型参数化的情况下实现随机优化。在基于FLUX的MeanFlow和sCM上的实验表明,在基于奖励的、感知的和任务级别的指标上均有改进。

基于价值梯度流的强化学习

Hugging Face Daily Papers

价值梯度流(VGF)提出了一种可扩展的行为正则化强化学习方法,将其构建为通过离散梯度流求解的最优传输问题,在离线强化学习和大型语言模型强化学习基准测试中取得了最先进的成果。该方法消除了显式的策略参数化,同时通过控制传输预算实现了自适应的测试时缩放。

通过非梯度向量流的流图学习

arXiv cs.LG

本文介绍了SGFlow,一种为扩散模型学习流图的方法,该方法避免了可逆性约束和通过模型迭代的反向传播,在CIFAR上取得了有竞争力的FID分数,并具有经过证明的稳定点保证。

MeanFlowNFT: 将前向过程强化学习引入平均速度生成器

Hugging Face Daily Papers

MeanFlowNFT为平均速度生成器引入了一种前向过程强化学习方法,能够在保持快速少步采样的同时,高效地与人类偏好对齐。实验表明,它在大多数指标上优于先前经过强化学习调优的少步生成器,甚至超越了多步调优的扩散模型。