PaperlayerView exact revision on arXiv

Source-linked AI summary

Meta Flow Maps enable scalable reward alignment

Peter Potaptchik, Adhi Saravanan, Abbas Mammadov, Alvaro Prat, Michael S. Albergo, Yee Whye Teh

arXiv:2601.14430v2stat.MLcs.LG

TL;DR

Reward alignment in generative models requires conditional-posterior samples for value-function estimation, but exact rollouts are expensive and deterministic maps cannot preserve posterior diversity. Meta Flow Maps amortise posterior-targeting transport into stochastic one-step maps, enabling differentiable samples for efficient steering and unbiased off-policy fine-tuning. The paper reports competitive sample quality, improved controllability, and reduced computation across reward settings.

  • Problem

    Reward alignment requires conditional-posterior samples for value-function estimation, while exact rollouts are expensive and deterministic flow maps cannot represent the full posterior.

  • Method

    Meta Flow Maps train a single amortised stochastic flow-map model that generates arbitrarily many differentiable one-step posterior samples from intermediate states.

  • Results

    MFMs provide efficient inference-time steering and unbiased off-policy fine-tuning across diverse rewards, with competitive sample quality, improved controllability, and reduced computation.

  • Takeaways & Limitations

    MFMs eliminate explicit inner rollouts for the supported steering and fine-tuning settings while retaining stochastic posterior sampling.

  • Takeaways & Limitations

    Existing posterior samplers face a scope boundary: exact methods are slow, whereas deterministic accelerated methods cannot capture the stochasticity required for posterior sampling.

Abstract

from arXiv · show

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.

1 Introduction

Reward alignment requires estimating a value function from conditional clean-data posteriors, but exact sampling is costly and deterministic accelerations lose posterior diversity. Meta Flow Maps address this bottleneck with stochastic one-step posterior sampling for efficient steering and fine-tuning.

  • Reward alignment: Inference-time steering and fine-tuning both depend on the value function associated with the reward-tilted target distribution.The value-function gradient determines the optimal drift correction for transient steering or permanent fine-tuning.
  • Computational bottleneck: Exact posterior sampling requires repeated ODE or SDE trajectory integration, making steering slow and reward-specific fine-tuning impractical.Heuristic posterior approximations are efficient but biased.
  • Limits of deterministic maps: Deterministic flow maps efficiently predict endpoints but collapse each intermediate state to one output instead of representing its full conditional posterior.The required posterior can contain multiple valid clean endpoints for a fixed noisy state.
  • Stochastic flow maps: Stochastic flow maps use exogenous noise and an intermediate state to generate arbitrarily many i.i.d. samples from the conditional posterior.These samples provide a differentiable reparametrization for estimating the value function and its gradient.
  • Meta Flow Maps: Meta Flow Maps train one amortised model over an infinite family of posterior-targeting flow maps.The model extends consistency models and flow maps into the stochastic regime.
  • Empirical scope: MFMs support efficient steering and off-policy fine-tuning across diverse rewards, with competitive sample quality, improved controllability, and reduced computation.The contribution list specifically identifies asymptotically exact inference-time steering and unbiased fine-tuning objectives.

2 Dynamical Measure Transport

Dynamical transport learns time-dependent dynamics that move a reference distribution to data, while flow maps compress trajectory integration into direct state-to-state maps. Training combines local drift matching with consistency constraints to obtain valid global flow maps.

  • ODE Transport: ODE transport evolves samples from a reference distribution p0 to a data distribution p1 using a time-dependent drift.The drift is learned so the terminal state follows p1.
  • Training: Stochastic interpolants define intermediate distributions between prior and data samples, with a valid drift given by the conditional expectation of the interpolant derivative.The drift can be parameterised by a neural network and trained with mean-squared regression.
  • Efficient sampling: Numerically integrating ODE trajectories is expensive because sampling typically requires many neural-network evaluations.Consistency models and flow maps seek to compress this integration into one or a small number of steps.
  • Flow maps: A flow map directly predicts the trajectory state at time u from its state at time s without explicitly simulating infinitesimal dynamics.The operator Xs,u maps xs to xu for any pair of times.
  • Flow-map training: Flow-map training uses diagonal drift matching to enforce local correctness and consistency objectives to propagate it into a valid global trajectory.Consistency rules constrain the average drift for arbitrary time triples rather than only the diagonal s = u.

3 Reward Alignment

Reward alignment requires estimating the value-function gradient from the conditional posterior of clean endpoints, but existing exact samplers are too costly and deterministic flow maps cannot represent posterior diversity. The section develops the posterior-sampling formulation that motivates Meta Flow Maps.

  • Unified objective: Reward alignment targets a reward-tilted terminal distribution through inference-time steering or fine-tuning.Both paradigms use the same controlled-dynamics objective, differing in whether the correction is applied during sampling or distilled into model parameters.
  • Unified objective: The value function measures expected future reward under p1|t(·|x), and its gradient supplies the optimal drift correction for both control paradigms.The correction can be applied transiently during inference or permanently through fine-tuning.
  • Why MFMs: Differentiable posterior samples are also required for reparameterized gradient estimators, whereas differentiating through exact ODE or SDE solvers is prohibitively expensive.This computational burden has often forced reliance on coarse approximations.
  • Posterior estimation: Exact posterior sampling can use inner SDE rollouts, but nesting simulations inside steering or fine-tuning makes the approach prohibitively costly.Independent noise produces posterior samples, yet each sample batch requires a full forward simulation from the current state.
  • Posterior estimation: Conditional ODE flows can target p1|t(·|x), but sampling still requires separate initial conditions and trajectory integration for every Monte Carlo sample.Even analytically accessible conditional drifts remain expensive to unroll.
  • Why MFMs: Standard flow maps are efficient marginal transports but deterministic, so they cannot represent the full conditional posterior when multiple endpoints are valid for one intermediate state.This creates a trade-off between exact but slow stochastic sampling and efficient but insufficient deterministic acceleration.

4 Meta Flow Maps

Meta Flow Maps extend one-step flow maps into the stochastic regime, producing diverse conditional-posterior samples while amortising a family of posterior-targeting transports. Their training combines diagonal velocity supervision with consistency objectives, yielding one-shot posterior sampling across contexts.

  • Stochastic Flow Maps: Stochastic Flow Maps map exogenous noise and an intermediate state directly to conditional-posterior samples, enabling arbitrarily many distinct draws for one context.The noise input lets the operator cover the full support of each target posterior.
  • Meta Flow Maps: Meta Flow Maps amortise a family of context-dependent flow maps, with (t, x) selecting the transport map targeting p1|t(·|x).The model learns an infinite collection of posterior-targeting flow maps simultaneously rather than one transport map.
  • Meta Flow Maps: MFM auxiliary trajectories need not intersect the conditioning state x; only their endpoint is constrained to follow p1|t(·|x).Intermediate auxiliary states generally lack a direct interpretation in the original generative process.
  • Training: MFM training combines diagonal velocity supervision with consistency losses so each fixed-context map integrates its conditional velocity field correctly.The framework supports self-distillation from data and teacher-distillation from a pretrained flow-matching model.
  • Extensions: The framework extends beyond linear interpolants to arbitrary context sets and general stochastic processes.The stated examples include processes such as video frames or weather observations.
  • Sampling: A trained MFM produces one-shot conditional-posterior samples, and at t=0 its endpoint distribution equals the marginal data distribution p1.The sampling construction uses independent base-noise draws for the same context.

5 MFMs for Reward Alignment

MFMs use differentiable posterior samples to estimate value-function gradients for inference-time steering and to construct reward-alignment objectives for fine-tuning. The resulting samplers have explicit discretisation and Monte Carlo errors, while fine-tuning can be performed off-policy.

  • Inference-Time Steering: MFMs estimate the value-function gradient by drawing direct conditional-posterior samples and using them in Monte Carlo steering estimators.The MFM endpoint map X0,1(ϵ; t, x) supplies differentiable samples for the required posterior expectations.
  • Inference-Time Steering: The extracted marginal drift and score, together with the estimated value gradient, provide the components needed for optimally steered SDE or ODE dynamics.The marginal drift is recovered from the unconditional flow, and a Gaussian base enables score extraction by reparametrisation.
  • Inference-Time Steering: MFM steering errors arise from both time discretisation and Monte Carlo estimation of the optimal drift.The convergence analysis quantifies these errors for the SDE sampler using Euler–Maruyama steps and independent Monte Carlo samples.
  • Inference-Time Steering: The convergence constant in the steering guarantee depends exponentially on the dimension d, while the analogous ODE result is stated in W2 distance.The KL bound is specific to the stochastic setting and relies on Girsanov’s theorem.
  • Training-Time Fine-Tuning: MFMs support fine-tuning toward the reward-tilted distribution through an explicitly off-policy objective defined pointwise over all states and times.Training can sample states from the interpolant without simulating the current model.
  • Training-Time Fine-Tuning: The reward-aligned drift can be distilled into a one-step tilted flow map or corresponding tilted MFM for permanent model alignment.This extends MFMs beyond transient inference-time steering.
  • Training-Time Fine-Tuning: The fine-tuning objective avoids the finite-sample bias of self-normalised Monte Carlo drift estimators by enforcing an equivalent implicit optimality condition.A surrogate regression loss is constructed whose expected gradient vanishes at the optimal drift.

6 Related Work

Prior work addresses reward alignment through approximate steering, exact particle methods, rollout-based search, and fine-tuning, but these approaches trade accuracy, compute, or robustness. MFMs differ by learning cheap differentiable posterior samples for tilted-drift estimation.

  • Inference-Time Alignment: Inference-time alignment methods either approximate the exact tilted dynamics or use particle-based resampling and search.The first class targets the reward-tilted distribution through approximate posterior or drift calculations.
  • Inference-Time Alignment: Gaussian posterior approximations are efficient but can be biased and fail in multimodal or nonlinear settings where the mean is not a valid data sample.LGD is described as approximating the posterior with a Gaussian centered at the posterior mean.
  • Inference-Time Alignment: Sequential Monte Carlo can target the tilted distribution unbiasedly in principle, but requires many particles to avoid weight degeneracy and collapse.Search-based alternatives estimate intermediate rewards through explicit rollouts, which remains costly per sampling step.
  • Few-Step Samplers: Unlike standard few-step models that primarily accelerate sampling, MFMs provide one-step differentiable samples from p1|t(·|x) for tilted-drift estimation.This requirement distinguishes MFM training from prior consistency-model and flow-map approaches.
  • Posterior Sampling: Existing posterior-sampling approaches often rely on SDE trajectory rollouts, while GLASS still requires expensive ODE solving during inference.MFMs instead eliminate iterative integration by learning the transitions during training.
  • Generative Fine-Tuning: Reward-maximisation fine-tuning can cause mode collapse and overfitting by concentrating on a single high-reward mode rather than the true posterior.The related-work discussion contrasts this with distribution-matching approaches.

7 Experiments

Experiments evaluate MFM steering, posterior sampling, value estimation, and fine-tuning across synthetic, MNIST, and ImageNet settings. MFMs improve target-posterior fidelity and reward-alignment efficiency, including performance exceeding Best-of-N with substantially fewer function evaluations.

  • 7.1.1 Gaussian Mixture Models: GMM steering schemes outperform DPS and TDS with as few as N = 2 or 4 Monte Carlo samples.DPS significantly over-represents the largest mode, while the analytic posterior enables direct sampling-fidelity assessment.
  • 7.1.1 Gaussian Mixture Models: Increasing Monte Carlo samples improves steering performance for the drift estimators.
  • 7.1.2 MNIST: On MNIST, MFM-GF and MFM-G approach correct target class ratios as Monte Carlo samples increase, while MFM-G significantly outperforms MFM-GF.DPS substantially over-represents the dominant mode.
  • 7.2.1 MFM Performance: MFM-XL/2 achieves FID 1.97 in 4 steps while also providing stochastic one-step posterior samples for reward alignment.The ImageNet benchmark uses 50,000 generated and reference images for FID computation.
  • 7.2.1 MFM Performance: MFMs outperform explicit ODE rollouts in posterior recovery and value-function estimation across conditioning times and NFEs.The improvement is greatest at one function evaluation, where differentiating through rollouts is prohibitively expensive.
  • 7.2.2 Inference-Time Steering: MFM-G outperforms Best-of-N and DPS across reward models, with N = 1 exceeding Best-of-N at NBoN = 1000 while using over 100× fewer NFEs.Steering with different reward models does not degrade the remaining evaluated metrics and often improves them.
  • 7.2.3 Fine-Tuning: Fine-tuning on HPSv2 produces stable, consistent increases in scores from all three reward models while preserving the base samples’ semantic content.

8 Conclusion

The conclusion presents MFMs as amortized stochastic flow maps for sampling conditional posteriors efficiently. This supports asymptotically exact inference-time steering and unbiased off-policy fine-tuning to general rewards, while leaving training-objective design and broader stochastic-process extensions open.

  • 8 Conclusion: MFMs amortize an infinite family of conditional transport problems into one few-step model that samples from conditional posteriors.
  • 8 Conclusion: The resulting samples eliminate explicit inner rollouts for efficient, asymptotically exact inference-time steering.
  • 8 Conclusion: MFMs also enable efficient, unbiased off-policy fine-tuning that permanently aligns generative models with general rewards.
  • 8 Conclusion: The framework is agnostic to the training objective and extends beyond fixed-endpoint generation to arbitrary intermediate-time prediction and more general stochastic processes.

A Methodology

MFMs provide stochastic one-step posterior maps that generate differentiable samples for value estimation, supporting efficient steering and fine-tuning. The framework also extends to arbitrary intermediate-time prediction and gradient-free reward search.

  • Stochastic posterior sampling: A trained MFM can be reparametrized to sample posteriors from interpolants with different coefficient paths.The construction uses a corresponding time t* and preserves differentiable one-shot sampling.
  • Stochastic posterior sampling: MFMs reparametrize posterior sampling so varying exogenous noise produces arbitrarily many i.i.d. one-step samples from p1|t(·|x).These samples provide a differentiable reparametrization for asymptotically exact value-function and gradient estimation.
  • Extended prediction: The construction generalizes to prediction at arbitrary intermediate times for stochastic processes over a general index set T.It can condition on multiple time points and is not restricted to interpolants or flow-matching processes.
  • Extended prediction: MFMs define conditional transport bridges whose endpoint law matches the target posterior pr|t(·|x), rather than reproducing the physical process evolution.The auxiliary flow transports a base distribution q to the conditional posterior.
  • Gradient-free reward search: MFM-Search uses efficient posterior samples to explore high-reward regions when rewards are non-differentiable, but it has no asymptotic guarantee of sampling the true tilted distribution.The method is presented as one heuristic within a broader search and optimization design space.

B.2 Architecture

MFM architectures augment flow maps with conditioning on an outer time and intermediate state while retaining standard generative-model backbones. Training combines diagonal and consistency objectives, with initialization and adaptive-loss choices supporting practical implementation.

  • Architecture: MFMs extend standard flow-map architectures by conditioning on the outer time t and state x.For DiT, this uses an additional time embedding and a patch embedder for x, modulated through AdaLN-Zero.
  • Initialization: Flow-map initialization can preserve the original map by zero-initializing the new time embedding and state-conditioning modulation.This retains the flow map at initialization while enabling the added MFM inputs to be learned.

C.2 Convergence Guarantees

The convergence analysis bounds errors from Monte Carlo drift estimation and time discretization under regularity assumptions. It establishes convergence of MFM steering to the reward-tilted target in Wasserstein-2 and KL divergence, with a boundary caveat at t = 1.

  • Formal guarantees: Under the stated regularity assumptions, MFM steering converges to the target distribution with a bound uniform in the discretization and Monte Carlo parameters.The theorem applies to the terminal distribution generated with K time steps and N independent samples per step.
  • Boundary limitation: The guarantees are formally applied before t = 1 when the target distribution lies on a low-dimensional manifold and the score becomes singular.The Wasserstein bound includes a smoothing error term W2(p1−ε, preward).
  • Error decomposition: Fresh independent posterior samples at each time interval make the sampling noise a martingale contribution in the Wasserstein analysis.The resulting error decomposition separates Monte Carlo variance, estimator bias, and Euler discretization.
  • Formal guarantees: The convergence proof controls Wasserstein-2 and KL errors by combining Monte Carlo estimation error with discretization error.The KL analysis uses path-measure data processing and Girsanov’s theorem under a positive lower diffusion bound.

F.1 Gaussian Mixture Model (GMM)

The GMM experiment evaluates MFM steering on a three-component two-dimensional mixture and a noisy linear inverse-problem reward. It compares posterior samples with analytic-posterior samples using SW2 and MMD, while ImageNet experiments use adapted DiT models.

  • GMM setup: The synthetic prior is a 2D GMM with three components centered at (−3, −3), (0, 0), and (3, 3), each with covariance 0.5I2×2.A small MLP is trained with the semigroup MFM loss to sample the prior.
  • GMM setup: The reward is the likelihood of a noisy linear inverse problem with a = [1.2, −0.8]⊤, σ = 0.2, and yobs = −1.0.Steering targets the posterior p(x|yobs = −1).
  • Evaluation: Evaluation generates 4096 posterior samples and compares them with analytic-posterior samples using sliced Wasserstein-2 and MMD.The methods use Euler or Euler–Maruyama with N = 1000 discretization steps; SMC uses K = 4096 particles.
  • Evaluation: The reported figures examine MMD against true posterior samples, density maps relative to the analytic posterior, and empirical class ratios against the target ratio.These correspond to Figures 11–13.
  • ImageNet implementation: For ImageNet at 256 × 256, MFM adaptations of DiT increase parameter counts from 131 →134M and 675 →684M for B/2 and XL/2 models.The models are initialized from pretrained flow or flow-map checkpoints and trained with data or distillation objectives.

F.3.1 Ablations.

The ablations compare FID across increasing numbers of function evaluations and identify the strongest adaptive-loss configuration for 2- and 4-step XL/2 generation.

  • Ablations: The best configuration, MFM-XL/2, is selected from comparisons across model scales and training objectives.The main-body and Table 2 designation corresponds to the best configuration in Table 4.
  • Ablations: The default adaptive-loss coefficients are p_diag = 1.0 and p_cons = 1.0, while p_diag = 0.5 and p_cons = 1.0 was marginally more effective for 2- and 4-step XL/2 generation.The alternative was observed specifically in the XL/2 experiments.
  • Ablations: FID is reported for increasing NFE, with lower values indicating better generation quality.Bracketed values specify the adaptive-loss parameters for diagonal and consistency terms, while c remains fixed at 0.01.

F.3.2 Inference-Time Steering

The inference-time steering ablations compare MFM-based estimators and baselines under compute-normalized evaluation, alternative reward models, varying Monte Carlo samples, and practical ImageNet implementation constraints.

  • Compute-normalized comparisons: MFM-GF performs far worse than both MFM-G and MFM-Search in the compute-normalized comparison.This comparison is retained from the main-body plot for completeness.
  • Robustness checks: Alternative reward-model evaluations test whether steering gains transfer beyond the reward model used for steering.The steering reward appears in the first bolded subplot, with other reward models shown in subsequent subplots.
  • Estimator costs: MFM-GF uses K base-drift evaluations plus N posterior samples and N reward evaluations per step, giving NFE = K + 2NK.The estimator samples posterior endpoints at each discretization step to estimate the value function.
  • Estimator costs: MFM-G estimates the value function with N posterior samples and N reward evaluations per step, while its gradient requires an additional backward pass.Its cost is summarized as NFE = K + 4NK under the stated 2× backward-pass assumption.
  • Baselines: Best-of-N generates NBoN samples with K discretization steps and evaluates each final sample for selection, requiring NFE = KNBoN + NBoN.The baseline is evaluated for NBoN from 1 to 1000 using a pool of 128,000 samples.
  • Implementation constraints: ImageNet steering gradients can greatly exceed the unconditional drift, making finer time discretization necessary to avoid excessive discretization error.The experiments instead rescale the steering gradient, introducing bias only in drift magnitude rather than direction.
Loading 2601.14430v2…