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
cs.LG, cs.CV, stat.ML
2026-09-23
An OT regularizer built from the pretrained drift lets flow maps be fine-tuned directly: top HPSv2 at 4 NFE on text-to-image, up to 280× less compute than Flow-GRPO.
Reward fine-tuning is a standard post-training step for generative models: take a learned scorer and push the pretrained generator toward higher scores. The dominant recipe writes this as KL-regularized reward maximization, whose optimum is a tilted distribution, the base law reweighted by e^(λr). Two frictions have persisted in that framing. First, KL tilting can only redistribute probability mass over the support the base model already has: two points with equal reward keep their relative density no matter how large λ gets, and moving mass into high-reward regions where the base model is thin is mechanically out of reach. Second, tractable training objectives lean on Girsanov's theorem, a tool for stochastic diffusions. Teams using deterministic flow matching end up detouring: treat the flow as a diffusion, fine-tune, convert back, and roll out samples along the way to estimate value functions at high variance. Distilled one-step flow maps do not fit into this pipeline at all.
WTF (Wasserstein-Tilted Flow Maps) replaces the KL regularizer with an optimal-transport cost built from the pretrained drift: a candidate velocity pays for the time-integrated squared deviation from the base flow, so only the departure from the base model is billed, and moving farther costs more.
Three theoretical steps turn that into a runnable algorithm:
The flow map itself needs a sentence. A flow matching model learns the velocity field of an ODE, and sampling integrates it step by step, tens to hundreds of network calls. A flow map learns the ODE's solution operator: one call jumps across any time interval, and multi-step sampling chains calls through the semigroup property. WTF fine-tunes the flow map directly, so the output is few-step by construction, with no post-hoc distillation stage. At scale the reward gradient takes the Euclidean form, the plain gradient at the endpoint, which the authors found more stable than the exact Jacobian pullback.
ImageNet-256, base model DMF XL/2, HPSv2 as training reward, three seeds, every method given the same budget of roughly five hours on one 8×H100 node:
| Method | NFE | HPSv2 | DreamSim diversity |
| Base | 1 | 0.214 | 0.790 |
| Adjoint Matching | 250 | 0.243 | 0.453 |
| VFM | 1 | 0.307 | 0.482 |
| WTF | 1 | 0.319 | 0.482 |
Text-to-image, base model TiM-T2I, roughly 48 GPU-hours of fine-tuning:
| Method | NFE | HPSv2 | PickScore | ImageReward |
| Base | 1 | 0.228 | 20.36 | 0.182 |
| Flow-GRPO | 50 | 0.303 | 22.06 | 1.210 |
| Adjoint Matching | 50 | 0.357 | 22.75 | 1.380 |
| WTF | 4 | 0.401 | 22.97 | 1.383 |
On compute, measured as the budget needed to reach a baseline's peak reward, WTF uses 47× less than Adjoint Matching and 280× less than Flow-GRPO. At one NFE, WTF's 0.319 already exceeds both the base model at 250 steps (0.218) and Adjoint Matching at 250 steps (0.243).
RL post-training is becoming a line item for generative models, and whether rollouts run at 4 steps or 50 steps changes the cost by an order of magnitude, which is where the 280× comes from. The methodological point is broader: KL tilting is not the only legitimate regularizer for reward fine-tuning, and for deterministic flows a transport-based one is native. Few-step accelerated samplers stop being an inference-side trick and become infrastructure that post-training can target directly.
Four are stated by the authors: the fixed-point result holds at the population level with no finite-sample guarantee; large-scale experiments use the Euclidean reward gradient rather than the exact one analyzed; higher reward comes with lower diversity, traded off through λ; and the method assumes a pretrained flow map is already in hand, so starting from a velocity model requires a distillation stage whose cost is excluded from the 47× and 280× comparisons. Two more stand out on a close read. HPSv2 is also the training objective, and the paper does not disentangle how much of the PickScore and ImageReward gains come from correlation among reward models. Each benchmark uses a single base model, so reproduction on other backbones is untested. Diversity also needs an honest read: WTF's 0.482 against the base model's 0.841 at 250 steps is a clear drop, merely better than peer fine-tuning methods such as MFM at 0.379.