PaperlayerView exact revision on arXiv

Source-linked AI summary

WTF?! Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps

Abbas Mammadov, Jerry Y. Huang, Justin Lin, Partha Kaushik, Sheel Shah, Kartik Nair, Yee Whye Teh, Nicholas M. Boffi

arXiv:2609.27033v1cs.LGcs.CVstat.ML

TL;DR

Reward fine-tuning seeks to improve sample rewards while remaining close to a pre-trained flow, but existing approaches center on reward tilting and diffusion-based procedures. WTF introduces a drift-based optimal-transport regularizer and an equivalent deterministic optimal-control formulation for direct flow-map fine-tuning. On ImageNet-256 and text-to-image, it matches or surpasses baseline rewards, with substantially reduced text-to-image training compute.

  • Problem

    Existing reward fine-tuning commonly targets a reward-tilted distribution through diffusion-based procedures, despite the use of deterministic flow maps and few-step samplers.

    Source: §1 · p. 2, §3 · p. 4

  • Method

    WTF uses an optimal-transport regularizer built from the pre-trained drift, equivalent to deterministic optimal control, to fine-tune a flow map directly.

    Source: §1 · p. 3

  • Results

    Across ImageNet-256 and text-to-image, WTF matches or surpasses baselines on reward and reaches matched text-to-image reward with up to 280× and 47× less compute than zeroth-order and first-order baselines, respectively.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The fine-tuned model retains few-step flow-map inference without post-hoc distillation.

    Source: §7 · p. 12

    WTF assumes access to a pre-trained flow map; starting from a velocity model requires a preceding distillation stage whose cost is excluded from the fine-tuning comparison.

    Source: §2.1 · p. 3, §7 · p. 12

Abstract

from arXiv · show

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.

1 Introduction

Reward fine-tuning seeks higher-reward samples while staying close to a pre-trained model, but existing approaches rely on reward tilting and diffusion-based procedures rather than directly fine-tuning deterministic flow maps. WTF uses a drift-based optimal-transport regularizer to move samples toward higher reward, enabling simulation-free flow-map fine-tuning with strong benchmark results and reduced compute.

1 Introduction

  • Existing reward-fine-tuning methods commonly require expensive rollouts to score terminal samples, making rollout cost a major part of the post-training budget.Source: §1 · p. 2
  • Existing procedures remain diffusion-based despite growing use of deterministic flows and few-step flow maps, motivating direct end-to-end fine-tuning of flow maps.Source: §1 · p. 2
  • WTF builds an optimal-transport regularizer from the pre-trained drift, moving individual samples toward higher reward rather than reweighting the base distribution.The approach is presented as a direct fine-tuning framework for flow maps.Source: §1 · p. 3
  • The regularized fine-tuning problem is equivalent to deterministic optimal control with terminal reward and is the rescaled small-noise limit of stochastic optimal-control algorithms for reward tilting.Source: §1 · p. 3
  • WTF provides end-to-end flow-map fine-tuning with unbiased, critic-free value estimates requiring O(1) rather than O(T) flow-map evaluations for a T-step rollout.The fine-tuned model retains the flow map’s few-step inference at deployment.Source: §1 · p. 3
  • On ImageNet-256 and text-to-image modeling, WTF matches or surpasses baselines on reward; at text-to-image scale, it reaches matched reward with up to 280× and 47× less compute than zeroth- and first-order baselines, respectively.Source: §1 · p. 3

2 Background

Flow maps learn the solution operator for a pre-trained flow, allowing samples to be generated in one evaluation or with a few steps. Existing reward-fine-tuning methods instead target a reward-tilted distribution through stochastic processes and convert the fine-tuned diffusion back into a flow.

2.1 Flow and flow map-based generative models

  • A learned flow map is the solution operator of the flow dynamics, enabling one-evaluation sampling or higher-accuracy sampling over multiple steps.Recent applications obtain samples in four or fewer steps with performance matching many-step flows.Source: §2.1 · pp. 3–4

2.2 Reward fine-tuning

  • Existing reward-fine-tuning methods approximate sampling from a reward-tilted distribution using a stochastic value function, then convert the fine-tuned diffusion process back into a flow.The value function uses a conditional expectation over a memoryless stochastic process.Source: §2.2 · p. 4

3 Wasserstein-tilted flow maps

Existing reward-tilting procedures rely on diffusion rollouts, while WTF defines deviation from a pre-trained deterministic flow through an optimal-transport cost. The resulting objective has a particle-level reward-cost interpretation and an equivalent deterministic optimal-control formulation.

3 Wasserstein-tilted flow maps

  • Diffusion-based value-gradient rollouts can be slow and high variance, so WTF uses a flow map to amortize deterministic value-gradient rollouts.Figure 3 contrasts critic-based or T-step rollout alternatives with WTF’s two flow-map evaluations and one Jacobian backprop, which are O(1), unbiased, and critic-free.Source: §3 · p. 4, Figure 3 · p. 6

3.1 An optimal transport regularizer

  • WTF regularizes with the optimal-transport cost T_cb(ρ0, ν), measuring residual control needed to steer the base flow from the same prior to a candidate terminal law.Unlike ordinary W2 costs against ρ1 or ρ0, this cost incorporates the pre-trained drift.Source: §3.1 · p. 5

3.2 Understanding the regularizer

  • The regularized objective transports each source sample to its best reward-cost endpoint, rather than choosing the terminal law only through global reward reweighting.WTF can redistribute mass within a reward level set, whereas KL reward tilting preserves the base conditional distribution there.Source: §3.2 · p. 6
  • The objective is equivalent to deterministic optimal control with terminal reward, and its fine-tuned model is a single drift replacing the base drift at inference.The problem is also identified as the zero-noise limit of a Schrödinger-bridge problem.Source: §3.2 · pp. 6–7

4 Solving the optimal control problem

The section formulates reward fine-tuning as deterministic optimal control and develops policy iteration with simulation-free value-gradient estimates. The resulting training procedure combines diagonal value-gradient regression with off-diagonal flow-map self-distillation.

4 Solving the optimal control problem

  • Direct backpropagation through sampled controlled trajectories can reward-hack the terminal map, so WTF instead regresses the diagonal control to a frozen reference’s value gradient.This policy-improvement update uses flow-map transitions to construct on-policy states and value estimates without numerically simulating every intermediate state.Source: §4 · p. 7
  • The value function represents future reward minus remaining control cost, and its gradient supplies the optimal control for the fine-tuning objective.Fitting a candidate control to the reference value gradient is guaranteed to improve the objective at the population level.Source: §4 · p. 7
  • Policy iteration fits a trainable control to the frozen reference’s value gradient, then reuses that control as the next reference.Each iterate strictly improves the objective until the global-optimum fixed point is reached.Source: §4 · p. 7
  • Flow-map endpoints avoid simulating the full future trajectory, leaving an unbiased Monte Carlo estimate of a scalar time integral for value-gradient regression.The estimator’s variance comes only from that one-dimensional integral, rather than high-dimensional integrals in stochastic value functions for reward tilting.Source: §4 · p. 8
  • The training objective combines on-diagonal value-gradient regression with off-diagonal flow-map self-distillation using frozen-reference states and targets.At a population fixed point, the learned diagonal control and flow map satisfy the optimal-control solution under the stated regularity and full-support conditions.Source: §4 · p. 8

5 Related work

Prior work uses flow maps to accelerate generation and has explored reward alignment through stochastic maps, adapters, or inference-time guidance. WTF instead fine-tunes the deterministic flow map with a transport regularizer and simulation-free value estimates.

5 Related work

  • Flow maps amortize numerical integration by learning solution operators between times, enabling generation in one or a few network evaluations.Prior alignment work includes stochastic flow maps, while this paper studies direct reward fine-tuning that retains few-step generation.Source: §5 · p. 8

Reward fine-tuning.

  • KL-regularized reward fine-tuning targets an exponential reward tilt; WTF instead uses transport cost induced by the pre-trained deterministic drift and simulation-free value estimates.Adjoint Matching and RAM address the KL target, while VGG-Flow uses a learned value-gradient critic and updates the velocity field.Source: Reward fine-tuning · p. 9
  • MFM returns a velocity field requiring distillation for few-step deployment, whereas WTF fine-tunes the deterministic flow map across time pairs and preserves inference across NFE budgets.VFM learns a noise adapter with a flow map for one-step conditional generation.Source: Reward fine-tuning · p. 9
  • Inference-time methods leave parameters fixed; because WTF retains a flow map, FMRG can provide additional alignment without first distilling a velocity-field model.In the one-dimensional comparison, KL tilting reweights existing mass while WTF transports mass toward the reward as it moves into lower-density regions.Source: Reward fine-tuning · p. 9, Figure 5 · p. 10

6 Experiments

Experiments compare exact one-dimensional population optima and trained flow maps on ImageNet-256 and text-to-image. WTF achieves higher reward with competitive diversity and retains few-step inference; on text-to-image it reaches baseline peak rewards with substantially less training compute.

6 Experiments

  • The experiments compare exact and Euclidean reward gradients, evaluate HPSv2 and additional reward metrics, and assess diversity using DreamSim and CLIP embedding distances.The exact gradient is used in synthetic experiments; large-scale image experiments use the Euclidean reward gradient.Source: §6 · p. 10

6.1 Synthetic experiments

  • Exact population optima show WTF achieving higher reward than KL tilting, with a widening gap as the reward moves into lower-density regions of the base distribution.KL reweights existing mass, whereas WTF transports mass toward high-reward regions and can reach terminal laws unavailable to reward tilting.Source: §6.1 · pp. 10–11

6.2 ImageNet-256 main results

  • WTF achieves higher reward than Adjoint Matching, MFM, and VFM while maintaining competitive diversity, and the same checkpoint remains effective from 1 to 250 NFE.Table 6 reports ImageNet results for DMF XL/2, with WTF means over three matched seeds.Source: §6.2 · p. 11, Table 6 · p. 11

6.3 Text-to-image main results

  • At 50 NFE, in the same harness, WTF attains the table’s highest HPSv2 (the training reward), PickScore, and ImageReward while maintaining comparable diversity, reaching Flow-GRPO’s and Adjoint Matching’s peak rewards with 280× and 47× less training compute, respectively.The same fine-tuned flow map remains effective in the few-step regime.Source: §6.3 · pp. 11–12, Figure 7 · p. 11, Table 8 · p. 12

7 Conclusion

The conclusion presents a transport-regularized formulation that moves samples toward higher reward and is equivalent to deterministic optimal control. It reports reward performance on ImageNet-256 and text-to-image, alongside limitations on scope and initialization.

7 Conclusion

  • WTF transports samples toward higher reward rather than reweighting the base distribution, and its objective is equivalent to deterministic optimal control.The conclusion describes direct, end-to-end fine-tuning of the flow map.Source: §7 · p. 12
  • A pre-trained flow map enables simulation-free, critic-free value estimation and use of the same few-step sampler during training and deployment.The fine-tuned model remains compatible with inference-time alignment without an intervening distillation stage.Source: §7 · p. 12
  • On ImageNet-256 and text-to-image, WTF matches or surpasses baselines on reward and reaches their peak rewards with substantially less training compute.The conclusion also reports compatibility with flow map inference-time alignment without post-hoc distillation.Source: §7 · p. 12
  • The large-scale experiments use a Euclidean reward gradient, increasing reward can reduce diversity, and starting from a velocity model requires distillation whose cost is excluded.WTF assumes access to a pre-trained flow map; Proposition 4.2 is a population-level fixed-point result.Source: §7 · p. 12

D.2 Schrödinger bridges and the zero-noise limit

The section relates WTF to Schrödinger-bridge fine-tuning and its zero-noise limit, and explains a failure mode of a naive control objective. It also states pointwise and optimal-control equivalence results underlying the framework.

D.2 Schrödinger bridges and the zero-noise limit

  • In the zero-noise limit, the rescaled entropic Schrödinger-bridge cost converges to the static prior-action transport cost, and its stochastic control problem reduces to deterministic optimal control.The section identifies WTF as a deterministic limit of the path-KL fine-tuning recipe.Source: §D.2 · pp. 22–23

D.3 Failure of the naive control objective

  • The naive control objective can change the terminal map through off-diagonal velocities while leaving the diagonal control cost nearly unchanged, creating map-dynamics inconsistency.This issue motivated the value-gradient approach in the main text.Source: §D.3 · p. 23

E.1 Proof of Proposition 3.1

  • For each base sample, the static objective's optimum chooses a terminal sample maximizing reward minus transport cost.The optimizer is selected separately for each sample drawn from the base distribution.Source: §E.1 · p. 23

E.2 Proof of Proposition 3.2

  • Under standard regularity conditions, the static fine-tuning objective and deterministic optimal-control problem have equal optimal values, with an optimizer's terminal law solving the static problem.The proof's construction also realizes a static optimizer by an optimal control.Source: §E.2 · pp. 24–25

E.3 Proof of Proposition 4.1

Proposition 4.1 derives a performance-difference identity for any candidate control and shows that fitting the current value gradient yields an optimal control when equality holds.

E.3 Proof of Proposition 4.1

  • V^g improves on the reference value, with equality only when the reference control matches the value gradient along the trajectory induced by g.At equality, the reference control is optimal.Source: §E.3 · p. 26
  • The proof establishes a performance-difference identity for any candidate control by integrating along its controlled trajectory and subtracting its control cost.The identity is then specialized to the value-gradient control.Source: §E.3 · p. 26
  • At equality, the reference control equals the value gradient, reducing policy evaluation to the HJB equation whose unique solution is the optimal value function.This identifies the reference control with the optimal control.Source: §E.3 · p. 26

E.4 Proof of Proposition 4.2

The proof characterizes population fixed points of WTF as the optimal controlled drift, subject to regularity and full-support assumptions. It also describes implementation choices, time conventions, reward-gradient scaling, and training setups.

E.4 Proof of Proposition 4.2

  • Under regularity conditions and a full-support prior, every population fixed point of the WTF objective has diagonal velocity equal to the base drift plus the optimal control.The corresponding flow map is generated by that controlled drift.Source: §E.4 · pp. 26–27
  • The proof combines off-diagonal distillation, which enforces flow-map consistency, with diagonal regression, which fits the value gradient, and closes the loop using stop-gradient consistency.Full support lets the on-policy fixed-point condition imply the optimal control.Source: §E.4 · pp. 27–29

F Algorithmic aspects

  • For image experiments, the Euclidean reward gradient omits the flow-map Jacobian while leaving the generated endpoint unchanged.This computes the update in the direction that increases reward at the endpoint.Source: F Algorithmic aspects · p. 29

F.1 Diffusion-time convention

  • With data at t = 0 and Gaussian noise at t = 1, time reversal changes the signs and integration limits while retaining a positive control cost.The resulting negative sign implements gradient ascent at inference.Source: §F.1 · p. 30

Reward-gradient scaling.

  • Reward-gradient scaling normalizes the reward gradient by the ratio of base-velocity and reward-gradient norms, with the norms tracked by exponential moving averages.The authors note that without rescaling, the reward coefficient may need reward-specific tuning across orders of magnitude.Source: Reward-gradient scaling · p. 29, §F.2 · pp. 30–31

G.1 Training setup

  • ImageNet fine-tunes the full DMF XL/2 network, while text-to-image trains rank-16 LoRA adapters; both use HPSv2 reward gradients.ImageNet training runs 12,500 optimizer steps; text-to-image runs 1,460 steps.Source: §G.1 · p. 31
  • The benchmarks use AdamW with different optimizer settings, and text-to-image clips per-sample reward gradients to the batch’s 0.8 quantile.Both apply a global parameter-gradient L2 clip of 1.Source: §G.1 · p. 31

G.2 Sampling and evaluation

The evaluation protocols fix or share prompts and starting latents across methods and measure reward and diversity at specified inference budgets. Baseline and compute comparisons use defined training protocols, with compute endpoints that are not directly commensurable.

G.2 Sampling and evaluation

  • ImageNet evaluations use 1 or 250 NFE on 512 shared-seed images, while text-to-image uses 100 fixed prompts and evaluates at 1, 2, 4, 8, or 50 NFE.Diversity is measured with DreamSim and CLIP pairwise distances; text-to-image reward metrics score 1,000 generations.Source: §G.2 · pp. 31–32

G.3 Baseline implementations

  • Flow-GRPO and Adjoint Matching use distinct rollout, update, and adapter configurations, and baseline training stops after six consecutive checkpoints without reward improvement.Both baselines follow the sampling and evaluation protocol used for the paper’s comparisons.Source: §G.3 · p. 32

G.4 Training cost and compute comparison

  • Compute is compared on the same 8 × H100 node, but Table 11 checkpoint costs and Figure 7 costs to baseline peak reward use different endpoints.The endpoint ratios are therefore not directly commensurable; pre-training, evaluation, and post-hoc distillation are excluded.Source: §G.4 · p. 33, Table 11 · p. 33

H Additional experimental results

Additional experiments compare sample changes, reward-gradient estimators, reward-scale tradeoffs, and transport versus reward reweighting. The reported results include higher reward and diversity for the Euclidean estimator, alongside a reward-diversity tradeoff as reward scale increases.

H.2 Exact versus Euclidean reward gradient

  • In text-to-image comparisons at matched iterations, the Euclidean estimator achieves higher HPSv2, DreamSim, and CLIP diversity than the other estimator.Source: §H.2 · p. 33

H.3 Reward scale and reward-diversity tradeoff

  • Increasing λ from 10 to 40 raises HPSv2 by 0.023 and lowers DreamSim and CLIP diversity by 0.022 and 0.013, respectively.The text-to-image sweep retrains models over a 4× range of reward scales at matched training steps and evaluation protocol.Source: §H.3 · p. 33

H.1 Qualitative comparison on ImageNet

  • At 250 NFE with each column sharing an initial latent, WTF changes ImageNet sample composition and color while keeping the class, whereas adjoint matching stays close to the base sample.The comparison includes samples from the base model, adjoint matching, and WTF.Source: §H.1 · p. 33, Figure 12 · p. 34

I Synthetic experiments: transport versus reweighting

  • The synthetic experiments compare exactly computable WTF and KL population optima to isolate how transport and reward reweighting produce different terminal laws.The supplied passages state the experimental aim but not its outcome.Source: I Synthetic experiments: transport… · p. 34

I.1 Gaussian benchmark

The Gaussian benchmark compares exact population optima for a KL reward tilt and WTF as a reward bump moves outward; the text-to-image appendix shows samples across inference budgets.

I.1 Gaussian benchmark

  • The benchmark uses a N(0,1) base and prior with a width-2.25 bounded reward bump; the tilt is computed by quadrature and WTF in closed form.The bump is chosen to overlap appreciably with the base: at K = 4.75, base expected reward is 0.142 and 15.8% of base mass has r > 0.25.Source: §I.1 · p. 34
  • Figure 14 sweeps expected reward over reward centre K at fixed scale and reward scale λ at fixed centre, quantifying Figure 5’s qualitative trend.Source: Figure 14 · p. 35
  • At λ = 7, the difference between the two laws widens monotonically from 0.068 at K = 2 to 0.256 at K = 5.Table 15 reports expected reward as the reward moves outward; both laws are population optima.Source: §I.1 · p. 34, Table 15 · p. 35

I.2 Qualitative comparison on text-to-image

  • Figures 16–17 place prompts in columns and inference budgets in rows; each WTF sequence holds initial noise fixed, showing one sample refining rather than independent draws.Source: Figure 16 · p. 36, Figure 17 · p. 37

Related papers

Loading 2609.27033v1…