Meta Flow Maps enable scalable reward alignment

TL;DR

Meta Flow Maps (MFMs) provide a scalable, efficient method for reward alignment in generative models by enabling stochastic posterior sampling, reducing computational costs in steering and fine-tuning tasks.

Contribution

We introduce Meta Flow Maps, a novel framework extending flow models to stochastic regimes, allowing efficient posterior sampling and improved reward alignment in generative models.

Findings

Steered-MFM outperforms baseline on ImageNet across multiple rewards.

MFMs enable inference-time steering without inner rollouts.

MFMs facilitate unbiased, off-policy fine-tuning for general rewards.

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}(x_1|x_t)$, the distribution of clean data $x_1$ consistent with an intermediate state $x_t$, 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 $x_1$ from any intermediate state. Crucially, these samples provide a differentiable reparametrization 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.