mlsanityData to 7 Oct 2026

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

arXiv 2609.27033 ↗First posted 22 Sep 2026cs.LG cs.CV stat.MLcited by 0

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

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