Meta Flow Maps enable scalable reward alignment

TL;DR

This work tackles the computational bottleneck of reward alignment for generative models by introducing Meta Flow Maps (MFMs), stochastic flow-map operators that generate arbitrarily many one-shot samples from the conditional posterior $p_{1|t}(\cdotig| x)$ for any intermediate state. By providing differentiable posterior samples, MFMs enable asymptotically exact estimation of the value-function gradient $ abla V_t(x)$, which is central to both inference-time steering and off-policy fine-tuning toward general rewards. The authors present a dual estimator framework (MFM-GF and MFM-G) for $ abla V_t$, derive convergence guarantees for steered samplers, and develop an unbiased fine-tuning objective (MFM-FT). Empirically, MFMs demonstrate competitive posterior sampling quality, improved steering efficiency, and substantial compute savings on ImageNet compared to Best-of-N baselines, while also enabling scalable reward alignment through training-time fine-tuning.

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.

Figures (24)

• Figure 1: Samples from a Meta Flow Map (MFM) trained on ImageNet ($256 \times 256$). (Left) 4-step samples from a base MFM. (Right) Inference-time steering with MFMs and HPSv2 using the prompts shown. The base MFM generates images using only class labels and so all prompt alignment comes from the MFM steering with HPSv2.
• Figure 2: An MFM $X$ conditions on an intermediate time–state pair $(t, x)$ (visualized as noisy images) along the stochastic interpolant and learns a shared conditional flow $X_{s,u}(\cdot \,; t, x)$ that maps base noise $\epsilon \sim p_0$ to endpoint samples $x_1 \sim p_{1|t}(\cdot |x)$ (visualized as clean images) via $X_{0,1}(\epsilon; t, x)$. For a given $(t, x)$ pair, varying the initial noise $\epsilon' \sim p_0$ yields multiple samples from the same posterior $p_{1|t}(\cdot | x)$. Conversely, for the same initial noise $\epsilon$, conditioning on two different time–state pairs $(t_a, x_a)$ and $(t_b, x_b)$ yields one sample from each of two different posteriors $p_{1|t_a}(\cdot | x_a)$ and $p_{1|t_b}(\cdot | x_b)$.
• Figure 3: Conditional endpoint samples on ImageNet. In each block, the first column shows a ground-truth ImageNet image, the second column shows a corrupted version $x_t$ at the indicated noise level, and the remaining four columns are four independent one-shot samples from the Meta Flow Map, $\hat{x}^{(i)}_1 = X_{0,1}(\epsilon^{(i)}; t, x_t)$, targeting $p_{1|t}(\cdot \mid x_t)$, for independent $\epsilon^{(i)}$. The noise variables $\epsilon^{(i)}$ for the posterior-sample columns are coupled across the left and right sub-figures, showing how the same $\epsilon^{(i)}$ yields different endpoints as $(t,x_t)$ changes.
• Figure 4: Comparison for GMM inverse problems. Left plot shows the prior, analytic posterior and density maps of posterior samples (ODE). Center graphs show the Sliced-Wasserstein, $\mathcal{S}$-$\mathcal{W}_2$, between 4096 samples from the true posterior and our inference-time steering schemes. For SMC, we use $K=4096$ particles, and report the mean $\mathcal{S}$-$\mathcal{W}_2$ over 20 random seeds. Right table reports $\mathcal{S}$-$\mathcal{W}_2$ and MMD for select inference-time steering set-ups and fine-tuning (ODE).
• Figure 5: Comparison of the drift estimators through the empirical probability mass function (PMF) over the classes in the samples (from 4096 samples) and the ground-truth PMF of the target posterior (defined by the weight vector $\mathbf{w}$). Left graph shows the $\mathcal{L}_2$ distance between the empirical PMF and ground-truth PMF for increasing number of MC samples. Right graph plots the ground-truth PMF, alongside the empirical PMF obtained through different drift estimators.

Theorems & Definitions (8)

• Proposition 5.1: Convergence Rates
• Proposition C.1
• proof
• Proposition C.2: Formal Convergence Guarantees
• Remark C.3
• proof
• Proposition C.4: Drift Estimator Moments
• proof