PaperlayerView exact revision on arXiv

Source-linked AI summary

Diamond Maps: Efficient Reward Alignment via Stochastic Flow Maps

Peter Holderrieth, Douglas Chen, Luca Eyring, Ishin Shah, Giri Anantharaman, Yutong He, Zeynep Akata, Tommi Jaakkola, Nicholas Matthew Boffi, Max Simchowitz

arXiv:2602.05993v3cs.LGcs.AI

TL;DR

Reward alignment remains costly and brittle because existing approaches require reward-specific fine-tuning or approximate inference-time adaptation. Diamond Maps redesign flow-based generators as stochastic flow maps that amortize simulation while supporting accurate value-function estimation, with reported fast and high-quality reward alignment and scalable search and SMC.

  • Problem

    Existing reward-alignment approaches are costly or limited in accuracy, motivating generative models that support fast and accurate inference-time alignment to arbitrary rewards.

  • Method

    Diamond Maps combine flow-map distillation with stochastic transitions, using Posterior and Weighted designs for efficient value-function estimation.

  • Results

    Diamond Maps provide fast and high-quality reward alignment, with Posterior Diamond Maps enabling scalable proposals for search and SMC.

  • Takeaways & Limitations

    Diamond Maps provide a framework for adapting flow and diffusion models at inference time through stochastic flow-map estimators.

  • Takeaways & Limitations

    Importance-sampling estimator accuracy depends on effective sample size, while Posterior Diamond Maps require GLASS Flow distillation for training.

Abstract

from arXiv · show

Flow and diffusion models produce high-quality samples, but adapting them to user preferences or constraints post-training remains costly and brittle, a challenge commonly called reward alignment. We argue that efficient reward alignment should be a property of the generative model itself, not an afterthought, and redesign the model for adaptability. We propose "Diamond Maps", stochastic flow map models that enable efficient and accurate alignment to arbitrary rewards at inference time. Diamond Maps amortize many simulation steps into a single-step sampler, like flow maps, while preserving the stochasticity required for optimal reward alignment. This design makes search, Sequential Monte Carlo, and guidance scalable by enabling efficient and consistent estimation of the value function. Our experiments show that Diamond Maps can be learned efficiently via distillation from GLASS Flows, achieve stronger reward alignment performance, and scale better than existing methods. Our results point toward a practical route to generative models that can be rapidly adapted to arbitrary preferences and constraints at inference time.

1. Introduction

Existing reward-alignment methods either require costly reward-specific fine-tuning or modify sampling with limited accuracy. Diamond Maps redesign generative models around inference-time adaptability by combining stochasticity with flow-map efficiency.

  • Reward fine-tuning is complex, costly, and must be repeated for each new reward, while inference-time methods leave the pretrained model unchanged.
  • Value-function estimation is central to exact guidance but is considered intractable because standard models require many SDE or ODE simulation steps.
  • Diamond Maps use flow maps to amortize many inference steps into one neural-network evaluation for efficient value-function estimation and exact guidance.
  • Diamond Maps are stochastic flow maps that support one-step look-aheads for reward evaluation, enabling exploration, search, and guidance.
  • Posterior Diamond Maps distill GLASS Flows into one-step posterior samplers with simple value-function estimation.
  • Weighted Diamond Maps make standard flow maps stochastic through renoising, enabling consistent value-function estimation with off-the-shelf distilled models.
  • Posterior Diamond Maps also act as one-step samplers of the time-reversal SDE, enabling scalable proposals for search and SMC.
  • Experiments show that Diamond Maps provide fast and high-quality reward alignment.

2. Background

Flow matching interpolates Gaussian noise and data through a learned vector field, but sampling requires costly step-wise ODE simulation. Flow maps address this cost by amortizing multiple simulation steps into one network evaluation.

  • Flow matching applies to score-based diffusion models and represents data as vectors z in R^d along a probability path from Gaussian noise to the data distribution.
  • The marginal vector field is obtained by averaging conditional vector fields, and ODE simulation from Gaussian noise produces trajectories with the desired marginals.
  • At time 1, the simulated state X1 is distributed according to pdata, yielding a sample from the target data distribution.
  • Flow matching requires computationally expensive step-wise ODE simulation.
  • Flow maps amortize many ODE simulation steps into a single neural-network evaluation by mapping a state from time t to t′ along an ODE trajectory.
  • Flow maps can be learned through distillation from existing flow or diffusion models or through self-distillation.

3. Value function estimation via Stochastic Flow Maps

Reward alignment targets a reward-tilted version of the data distribution, making accurate value-function estimation essential for guidance, search, and SMC. Diamond Maps seek efficient, statistically consistent estimators by restoring stochasticity to flow-map transitions.

  • Reward alignment samples from a distribution proportional to pdata(z) exp(r(z)), where the pretrained model regularizes the search space.
  • The value function is a central object for inference-time alignment because it evaluates noisy states using expected future reward.
  • The log-expectation form of the value function arises from denoising a Gaussian kernel in flow and diffusion models.
  • The noised reward-tilted distribution measures how much more likely a noisy state is under the reward-tilted distribution than under the original data distribution.
  • Guidance: Accurate guidance is equivalent to accurately estimating the gradient of the value function.
  • Search and SMC: Accurate and efficient value-function estimation enables optimal particle filtering and search over the support of the prior.
  • Denoiser approximation: The value function is difficult to compute because posterior sampling requires expensive SDE or ODE trajectory simulation.
  • Denoiser approximation: The denoiser approximation can be highly biased, producing the Jensen gap and limiting inference-time adaptation methods to simple rewards.

4. Posterior Diamond Maps

Posterior Diamond Maps are stochastic one-step posterior samplers designed for efficient value-function estimation and inference-time reward alignment. They can also support iterative guidance, search, and SMC through Diamond DDPM sampling, which avoids the error accumulation of iterative denoising and noising.

  • Posterior Diamond Maps: Posterior Diamond Maps sample from p1|t(·|xt) with a stochastic flow map over an inner time axis while holding xt fixed.The inner axis evolves Gaussian noise into a posterior sample, making the model adaptable to rewards at inference time.
  • Value-function estimation: Consistent value-function and gradient estimators enable guidance using Monte Carlo posterior samples.Algorithm 1 samples from the posterior map and averages reward gradients to estimate the guidance signal.
  • Training: Posterior Diamond Maps are trained by distilling GLASS Flows, whose conditional flow evolves an inner state from Gaussian noise to p1|t(·|xt).The GLASS velocity field is obtained from a pretrained flow-matching model and then distilled into a flow map.
  • Iterative sampling: Diamond Early Stop DDPM sampling directly transitions from xt to xt′ by stopping the inner flow at a selected s* and re-transforming the state.The selected inner time satisfies t*(s*,t)=t′, avoiding a full denoise-to-one and renoise-back procedure.
  • Iterative sampling: Diamond DDPM sampling is contained in the Posterior Diamond Map and provides more efficient proposal distributions for search and SMC.The paper reports significantly improved performance over iterative denoising and noising.
  • Guidance: Naive reward-gradient guidance can collapse images toward the reward target, whereas the corrected value function preserves realism through added regularization.The blueness-reward illustration contrasts collapse under unweighted guidance with more realistic samples under the corrected estimator.

5. Weighted Diamond Maps

Weighted Diamond Maps make standard deterministic flow maps stochastic at inference time through renoising and reward correction. They reduce sampling cost while providing value-function estimates, but their efficiency relative to Posterior Diamond Maps remains theoretically unresolved.

  • Stochastic flow maps: A renoising map simulates the diffusion forward process backward in time, allowing a deterministic flow map to produce stochastic clean-data samples.The resulting distribution locally explores around xt but is not generally the true posterior.
  • Value-function correction: Renoised flow-map samples alone do not yield a correct value-function estimator because q1|t(·|xt) differs from the posterior p1|t(·|xt).The paper introduces correction terms to address this mismatch.
  • Value-function correction: The corrected estimator adds recovery-reward and score-based terms, which the paper identifies as crucial for exact guidance estimation.The score term is tractable through the relationship between the score function and the pretrained velocity field.
  • Inference cost: The corrected estimator requires O(2KN) function evaluations, with further reduction possible by approximating the score through Tweedie’s formula and a flow-map denoiser.K is the number of Monte Carlo samples and N is the number of ODE simulation steps.
  • Inference cost: Using a critical guidance window reduces total evaluations to O(K ∗Nguidance + Nno-guidance), while Weighted Diamond Maps can also estimate the value function itself.Only selected guidance steps require the Monte Carlo reward-gradient computation.
  • Posterior versus Weighted Maps: Weighted Diamond Maps can use off-the-shelf distilled models, whereas Posterior Diamond Maps require GLASS Flow distillation; their relative inference efficiency remains an open theoretical question.The comparison depends on importance-weight effective sample size and training quality.

6. Related Work

The paper situates Diamond Maps among flow-map and consistency-model methods, introducing stochastic flow maps to preserve randomness beyond deterministic ODE distillation. It also distinguishes the method from approximate posterior-learning approaches by claiming an exact construction.

  • Extensions: The paper additionally proposes Weighted Diamond Maps that reuse existing flow maps, alongside Diamond DDPM sampling.These extensions are presented as related-work distinctions and are discussed in later sections.
  • Posterior sampling: Related posterior methods learn approximations beyond the denoiser mean, while the paper states that Diamond Maps are exact rather than approximate.The comparison concerns posterior sampling and value-function estimation.
  • Stochastic flow maps: Prior flow-map and consistency-model methods distill the same marginal-vector-field ODE trajectory, whereas Diamond Maps construct stochastic flow maps not fully determined by the current state.This stochasticity differentiates the proposed model design from deterministic trajectory distillation.

7. Experiments

Experiments evaluate Posterior and Weighted Diamond Maps across inverse problems, prompt alignment, and high-resolution text-to-image generation. They report one-step posterior distillation, improved prompt adherence, and favorable compute–quality scaling against inference-time baselines.

  • Posterior Diamond Maps: Posterior Diamond Maps are distilled from pretrained CIFAR10, CelebA, and ImageNet1k models to evaluate one-step posterior sampling.The experiments include unconditional CIFAR10 and CelebA models and a class-conditional ImageNet1k model.
  • Posterior Diamond Maps: Posterior Diamond Maps successfully distill GLASS Flows into one-step posterior samplers, while Diamond Early Stop DDPM sampling outperforms iterative denoising and noising.The paper reports additional few-step and posterior-sampling results in Tables 3–5.
  • Posterior Diamond Maps: For inverse problems, Posterior Diamond Maps achieve a better Pareto frontier across reward scales and NFEs, with stochasticity improving robustness at high reward scales on CelebA-64.For ImageNet, increasing Monte Carlo samples becomes more useful after base drift-step scaling reaches diminishing returns.
  • Posterior Diamond Maps: Posterior Diamond Maps extend a class-conditional ImageNet1k model to simple text conditioning and generate out-of-distribution images following prompts.The evaluation uses ImageReward-based guidance and reports qualitative examples.
  • Weighted Diamond Maps: Diamond Maps improve compositional prompt adherence over base models and Flow Map Guidance, capturing attributes and spatial relationships missed at earlier stages.Comparisons use shared initial noise seeds per row to isolate guidance effects.
  • Weighted Diamond Maps: For FLUX, GenEval ablations identify SNR=1.5 and g=10 as the best compute–quality trade-off, while increasing Monte Carlo samples yields consistent gains.Excessive noise, especially SNR≥3.0 with many guidance steps, degrades quality as perturbation errors accumulate.
  • Weighted Diamond Maps: Weighted Diamond Maps use Monte Carlo particles to scale reward-tilted velocity estimation, producing a more favorable compute–quality frontier than Best-of-N and Flow Map Guidance.This scaling axis differs from independent-sample scaling and additional ODE-step refinement.
  • Weighted Diamond Maps: Weighted Diamond Maps outperform competing approaches with fewer NFEs and scale more favorably across both velocity-parameterized FLUX and direct-prediction SANA-Sprint flow maps.The comparison includes Best-of-N, Reward-based Noise Optimization, and Prompt Optimization.

8. Discussion

Diamond Maps enable efficient inference-time reward alignment through stochastic flow maps, with Posterior and Weighted variants supporting value estimation and one-step posterior or time-reversal sampling. The framework remains limited by reward-distribution mismatch and imperfections in distilled flow maps.

  • Diamond Maps enable efficient and accurate inference-time alignment using flow maps.
  • Posterior Diamond Maps distill GLASS Flows into one-step samplers for the posterior itself.
  • Weighted Diamond Maps turn standard flow maps into consistent value-function estimators through renoising.
  • When pdata is far from the reward-tilted distribution, estimator variance can become very large, making alignment difficult in high dimensions.This reflects a fundamental difficulty of inference-time alignment rather than a limitation specific to Diamond Maps.
  • Practical estimator accuracy depends on underlying flow-map quality, so distillation imperfections can contribute to remaining alignment error.
  • The framework is presented as applicable beyond images, including molecules, robotics, video, and audio.

A.3. Proof of Theorem 4.1

The proof establishes that Diamond Map constructions provide consistent value-function estimation and one-step stochastic transitions suitable for SMC and search. These properties yield an O(N) speed-up over standard flow models for comparable estimators.

  • The value-function estimator is consistent when constructed from posterior samples generated by the Diamond Map.
  • Posterior flow constructions reproduce the relevant SDE or DDPM transition distributions through matching joint distributions.
  • Monte Carlo samples provide the estimator construction, while gradient estimates follow by differentiating the resulting value expression.
  • Diamond Maps support SMC by evolving particles with one-step DDPM transitions and resampling using potentials that favor value-function increases.
  • O(N) speed-up: Diamond Maps obtain both pDDPM transitions and p1|t samples in one step compared with standard flow models.
  • The same construction applies to search, which collapses the particle population into a single particle with maximum potential.

D.1.1. CIFAR10/CELEBA-64

The experiments evaluate Posterior Diamond Maps across CIFAR-10, CelebA-64, and ImageNet settings, comparing their architecture, posterior recovery, sampling quality, and scalable Monte Carlo behavior. The reported setup emphasizes distillation from flow-based teachers and additional conditioning inputs.

  • Posterior Diamond Maps use additional time and conditioning inputs beyond the corresponding flow-map architectures.
  • Table 3 compares FID across step counts for CIFAR-10, CelebA-64, and ImageNet, with Posterior Diamond Map flow maps performing almost on par with standard flow maps.
  • Posterior recovery is evaluated by comparing FID for GLASS and Posterior Diamond Maps at different timesteps and step counts.
  • Posterior Diamond Maps are evaluated for super-resolution guidance and qualitative CLIPScore and ImageReward alignment.
  • Increasing Monte Carlo samples improves Sequential Monte Carlo performance even when batch size decreases under an approximately fixed NFE budget.
  • ImageNet Posterior Diamond Maps are initialized from pretrained flow maps and trained by distilling the GLASS velocity field.

D.1.3. GUIDANCE

The guidance experiments apply Posterior Diamond Maps to inverse problems and prompt alignment, varying guidance strength and Monte Carlo schedules. Prompt-alignment evaluation uses optimized and held-out rewards to assess reward hacking, while the method’s NFE accounting includes base, Monte Carlo, and posterior-sampling costs.

  • Monte Carlo schedules are ablated, including front-loaded schedules, while retaining at least one sample at every guidance timestep.
  • NFE accounting includes the costs of base flow steps, Monte Carlo samples, and posterior samples.
  • Guidance experiments evaluate Posterior Diamond Maps on noisy super-resolution inverse problems for CelebA-64 and ImageNet.
  • Prompt alignment on ImageNet optimizes either ImageReward or CLIPScore while evaluating the other as a held-out reward.
  • Weighted Diamond Maps provide a guidance algorithm using a pretrained flow map, Monte Carlo samples, and a differentiable reward.
  • Figure 16 reports less reward hacking for Posterior Diamond Maps on a held-out reward while optimized-reward performance remains similar.

D.1.4. SMC

The SMC and guidance setups combine Monte Carlo sampling, flow-map transitions, and reward-gradient normalization to make reward alignment computationally tractable. The experiments examine prompt alignment, hyperparameter trade-offs, and guidance behavior under fixed inference budgets.

  • SMC: SMC varies Monte Carlo samples K, softmax temperature τ, and ImageNet guidance cutoff tmax, while resampling at every outer step.The setup uses an ESS threshold of 1.0, shared noise across particles, and a constant number of MC samples across time.
  • SMC: Prompt-alignment experiments evaluate CLIPScore on CelebA-64 and ImageReward on ImageNet, comparing particle counts under matched total NFE budgets.ImageNet evaluation uses 50 prompts and reports both mean and best terminal ImageReward across particles.
  • Guidance normalization: Analytical guidance weights can overwhelm the reward gradient because score and likelihood gradients scale with data dimensionality d.The method instead unit-normalizes each gradient contribution and scales their combination with a shared λ.
  • Guidance normalization: For SANA-Sprint, λ = 20 is fixed across text-to-image experiments, while substantially larger values can cause visual artifacts by dominating the base velocity.The same implementation is applied to the SANA-Sprint and FLUX Flow Map settings.
  • Flow-map inference: The FLUX Flow Map uses 25 ODE steps and requires two inference steps per Monte Carlo sample and guidance step.One step advances to the next time point, while another predicts the clean endpoint for reward evaluation.
  • Guidance hyperparameters: For GenEval, SNR=1.5 and g = 10 with n = 25 provides the best reported compute trade-off, while moderate SNR values are most robust.The reported robust SNR range is [1.1, 1.5].

D.2.3. SANA-SPRINT ABLATIONS

The SANA-Sprint ablations show how endpoint-derived velocity enables Weighted Diamond Maps without retraining, while particle count, guidance steps, and time horizon shape performance. Guidance improves with several particles and more steps, but both dimensions exhibit diminishing returns or bounded useful ranges.

  • SANA-Sprint as a Flow Map: SANA-Sprint requires no retraining or architectural changes because Tweedie’s formula derives velocity from its existing clean-endpoint prediction.The resulting Weighted Diamond Map procedure costs n + g ∗mc NFE.
  • SNR and Particle Scaling: Using more than two particles gives significant improvements, but further particle scaling adds no additional improvement under the tested setup.The ablation keeps five guidance steps fixed over the [0.05, 0.25] time frame.
  • Number of guidance steps: Increasing guidance steps improves performance, with diminishing returns after N > 10.The experiments recommend keeping particles in the range [4, 8] while adjusting guidance steps to available compute.
  • Number of guidance steps: The [0.05, 0.25] guidance time frame performs better than longer horizons with the same number of steps, including [0.05, 0.45].The comparison is reported in the guidance-strategy ablation at SNR 20.0 and 4 particles.
  • Prompt alignment: Posterior Diamond Maps extend class-conditional ImageNet models to simple text prompts for ImageReward alignment.The qualitative examples are presented as guidance with Posterior Diamond Maps.
Loading 2602.05993v3…