WTF?! Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps
Not on the map. Why?
Abstract (arXiv)
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× 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.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 60 of its 74 references that are in the map. First posted on arXiv 22 Sep 2026.
Builds on
60 of 74 references.
References · 60
- 2021High-Resolution Image Synthesis with Latent Diffusion Modelscited by 18,694
- 2022Training language models to follow instructions with human feedbackcited by 15,712
- 2020Score-Based Generative Modeling through Stochastic Differential Equationscited by 8,964
- 2021Diffusion Models Beat GANs on Image Synthesiscited by 8,185
- 2022Classifier-Free Diffusion Guidancecited by 5,717
- 2022Flow Matching for Generative Modelingcited by 5,426
- 2024Scaling Rectified Flow Transformers for High-Resolution Image Synthesiscited by 4,094
- 2022Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flowcited by 3,550
- 2022Elucidating the Design Space of Diffusion-Based Generative Modelscited by 2,994
- 2023Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasetscited by 2,526
- 2023Consistency Modelscited by 1,778
- 2025Gemma 3 Technical Reportcited by 1,686
38 more
- 2022Diffusion Posterior Sampling for General Noisy Inverse Problemscited by 1,375
- 2022Scaling Laws for Reward Model Overoptimizationcited by 1,056
- 2023ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generationcited by 1,010
- 2022Building Normalizing Flows with Stochastic Interpolantscited by 964
- 2023Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generationcited by 950
- 2023Training Diffusion Models with Reinforcement Learningcited by 939
- 2023Stochastic Interpolants: A Unifying Framework for Flows and Diffusionscited by 897
- 2023Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesiscited by 847
- 2025Flow-GRPO: Training Flow Matching Models via Online RLcited by 605
- 2025Mean Flows for One-step Generative Modelingcited by 524
- 2023Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusioncited by 448
- 2023DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Modelscited by 431
- 2023Directly Fine-Tuning Diffusion Models on Differentiable Rewardscited by 431
- 2024One Step Diffusion via Shortcut Modelscited by 373
- 2022DiffDock: Diffusion Steps, Twists, and Turns for Molecular Dockingcited by 307
- 2023FreeDoM: Training-Free Energy-Guided Conditional Diffusion Modelcited by 290
- 2013A survey of the Schr"odinger problem and some of its connections with optimal transportcited by 286
- 2024Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Controlcited by 222
- 2024Deep Compression Autoencoder for Efficient High-Resolution Diffusion Modelscited by 209
- 2025A General Framework for Inference-time Scaling and Steering of Diffusion Modelscited by 201
- 2024TFG: Unified Training-Free Guidance for Diffusion Modelscited by 103
- 2014On the relation between optimal transport and Schr"odinger bridges: A stochastic control viewpointcited by 99
- 2024Flow map matching with stochastic interpolants: A mathematical framework for consistency modelscited by 99
- 2025How to build a consistency model: Learning flow maps via self-distillationcited by 96
- 2025Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Expertscited by 85
- 2025Inference-Time Alignment in Diffusion Models with Reward-Guided Generation: Tutorial and Reviewcited by 82
- 2025Align Your Flow: Scaling Continuous-Time Flow Map Distillationcited by 81
- 2025FlowDPS: Flow-Driven Posterior Sampling for Inverse Problemscited by 65
- 2024Deep Reward Supervisions for Tuning Text-to-Image Diffusion Modelscited by 59
- 2024Isambard-AI: a leadership class supercomputer optimised specifically for Artificial Intelligencecited by 59
- 2026Flow Map Language Models: One-step Language Modeling via Continuous Denoisingcited by 49
- 2022Probability flow solution of the Fokker-Planck equationcited by 46
- 2026Categorical Flow Mapscited by 37
- 2025Transition Models: Rethinking the Generative Learning Objectivecited by 32
- 2026Meta Flow Maps enable scalable reward alignmentcited by 32
- 2026Diamond Maps: Efficient Reward Alignment via Stochastic Flow Mapscited by 26
- 2026Discrete Flow Mapscited by 23
- 2025Decoupled MeanFlow: Turning Flow Models into Flow Maps for Accelerated Samplingcited by 19
50 most cited of 60 All 60 (JSON)
Built on by
Not cited yet.