Meta Flow Maps enable scalable reward alignment
Peter Potaptchik, Adhi Saravanan, Abbas Mammadov, Alvaro Prat, Michael Albergo, Yee-Whye Teh
Abstract
Controlling generative models is computationally expensive. This is because optimal alignment with a reward function-whether via inference-time steering or fine-tuning-requires estimating the value function. This task demands access to the conditional posterior p 1|t (x1|xt), the distribution of clean data x1 consistent with an intermediate state xt, a requirement that typically compels methods to resort to costly trajectory simulations. To address this bottleneck, we introduce Meta Flow Maps (MFMs), a framework extending consistency models and flow maps into the stochastic regime. MFMs are trained to perform stochastic one-step posterior sampling, generating arbitrarily many i.i.d. draws of clean data x1 from any intermediate state. Crucially, these samples provide a differentiable reparametrisation that unlocks efficient value function estimation. We leverage this capability to solve bottlenecks in both paradigms: enabling inference-time steering without inner rollouts, and facilitating unbiased, off-policy fine-tuning to general rewards. Empirically, our single-particle steered-MFM sampler outperforms a Best-of-1000 baseline on ImageNet across multiple rewards at a fraction of the compute.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers8
- Diamond Maps: Efficient Reward Alignment via Stochastic Flow MapsPeter Holderrieth, Douglas Chen, Luca Eyring, Ishin Shah et al.ICML 2026 · 18 citations
- Calibrated Test-Time Guidance for Bayesian InferenceDaniel Geyfman, Felix Draxler, Jan Groeneveld, Hyunsoo Lee et al.ICML 2026 · 4 citations
- Flow Sampling : Learning to Sample from Unnormalized Densities via Denoising Conditional ProcessesAaron Havens, Brian Karrer, Neta ShaulICML 2026 · 2 citations
- DAISI: Data Assimilation with Inverse Sampling using Stochastic InterpolantsMartin Andrae, Erik Larsson, So Takao, Tomas Landelius et al.ICML 2026 · 2 citations
- Variational Flow Maps: Make Some Noise for One-Step Conditional GenerationAbbas Mammadov, So Takao, Bohan Chen, Ricardo Baptista et al.ICML 2026 · 2 citations
Builds on36
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Diffusion Models Beat GANs on Image SynthesisPrafulla Dhariwal, Alexander Quinn NicholNeurIPS 2021 · 13,211 citations
- High-Resolution Image Synthesis with Latent Diffusion ModelsRobin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser et al.CVPR 2022 · 13,123 citations
- Scalable Diffusion Models with TransformersWilliam Peebles, Saining XieICCV 2023 · 5,568 citations
- Elucidating the Design Space of Diffusion-Based Generative ModelsTero Karras, Miika Aittala, Timo Aila, Samuli LaineNeurIPS 2022 · 3,959 citations
Related papers
- GLASS Flows: Efficient Inference for Reward Alignment of Flow and Diffusion ModelsPeter Holderrieth, Uriel Singer, Tommi Jaakkola, Ricky T. Q. Chen et al.ICLR 2026 · 7 citations
- How to Guide Your Flow: Few-Step Alignment via Flow Map Reward GuidanceJerry Huang, Justin Lin, Sheel Shah, Kartik Nair et al.ICML 2026
- Diffusion Tree Sampling: Scalable inference‑time alignment of diffusion modelsVineet Jain, Kusha Sareen, Mohammad Pedramfar, Siamak RavanbakhshNeurIPS 2025 · 41 citations
- CMT: Mid-Training for Efficient Learning of Consistency, Mean Flow, and Flow-Map ModelsZheyuan Hu, Chieh-Hsin Lai, Yuki Mitsufuji, Stefano ErmonICLR 2026 · 21 citations
- Multisample Flow Matching: Straightening Flows with Minibatch CouplingsAram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos et al.ICML 2023 · 243 citations
