PaperlayerView exact revision on arXiv

Source-linked AI summary

How to Guide Your Flow: Few-Step Alignment via Flow Map Reward Guidance

Jerry Y. Huang, Justin Lin, Sheel Shah, Kartik Nair, Nicholas M. Boffi

arXiv:2604.27147v3cs.LGcs.AI

TL;DR

Reward guidance seeks samples that satisfy user-specified criteria, but existing methods can require expensive multi-particle, many-step schemes or rely on poorly understood approximations. This paper casts guidance as deterministic optimal control and uses the flow map for single-trajectory guidance. FMRG matches or surpasses baselines across inverse problems and reward-guided generation with as few as 3 NFEs.

  • Problem

    Reward-tilted sampling is difficult, while multi-particle methods are costly and single-particle heuristic approximations leave the resulting distribution poorly characterized.

    Source: §1 · p. 2

  • Method

    The paper formulates guidance as deterministic optimal control, with the flow map providing feedback for integrating and guiding a single trajectory.

    Source: §1 · p. 3, §2.3 · p. 4

  • Results

    FMRG matches or surpasses baselines across latent-space inverse problems and reward-guided image generation with as few as 3 NFEs.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The flow-map perspective provides a way to understand existing heuristic guidance approximations whose accuracy relative to reward-tilted sampling has been questioned.

    Source: §6 · p. 16

    FMRG requires a pretrained flow map, and performance depends on its quality; FMRG-J uses approximately 48 GB VRAM at 512px resolution.

    Source: §6 · p. 16

Abstract

from arXiv · show

In generative modeling, we often wish to produce samples that maximize a user-specified reward such as aesthetic quality or alignment with human preferences, a problem known as \textit{guidance}. Despite their widespread use, existing guidance methods either require expensive multi-particle, many-step schemes or rely on poorly understood approximations. We reformulate guidance as a \textit{deterministic optimal control problem}, yielding a hierarchy of algorithms that subsumes existing approaches at the coarsest level. We show that the \textit{flow map}, an object of significant recent interest for its role in fast inference, arises naturally in the optimal solution. Based on this observation, we propose \textbf{Flow Map Reward Guidance (FMRG)}: a training-free, \textit{single-trajectory} framework that uses the flow map to both integrate and guide the flow. At text-to-image scale, FMRG matches or surpasses baselines across inverse problems and reward-guided generation with \textbf{as few as 3 NFEs}, giving at least an order-of-magnitude speedup in comparison to prior state of the art. Code is available at https://github.com/jrrhuang/fmrg.

1 Introduction

Reward guidance seeks samples that meet user criteria, but reward-tilted sampling is difficult and existing approaches either require costly multi-particle schemes or rely on poorly characterized approximations. The paper formulates guidance as deterministic optimal control and uses the flow map to build a single-trajectory method for few-step generation.

1 Introduction

  • Reward guidance targets samples satisfying criteria such as aesthetic quality, measurement agreement, physical plausibility, or user intent.Guidance steers sampling toward high-reward samples without additional training.Source: §1 · p. 2
  • Reward-tilted sampling is computationally intractable even for simple rewards, while multi-particle methods generically require prohibitive particle counts and integration steps.Single-particle approaches such as DPS use heuristic approximations whose produced distribution is poorly characterized.Source: §1 · p. 2
  • Deterministic probability flows and accelerated flow maps have become common, but a principled guidance method native to deterministic, few-step, single-trajectory generation is missing.Recent approaches that simulate stochastic transitions within an ODE sampler retain the structure of the stochastic formulation.Source: §1 · pp. 2–3
  • The paper formulates guidance as deterministic optimal control, where the flow map appears in the optimal feedback and supports integration and guidance along one trajectory.The framework is presented as native to few-step, single-trajectory generation.Source: §1 · p. 3
  • The framework also analyzes interventions such as early stopping and reward-gradient choice, and subsumes single-trajectory methods such as DPS and seed-optimization methods as special cases.The paper states that this unifying characterization is developed in Appendix D.Source: §1 · p. 3
  • FMRG matches or surpasses baselines across latent-space inverse problems and reward-guided image generation with as few as 3 NFEs.The paper reports at least an order-of-magnitude speedup over competing approaches.Source: §1 · p. 3

2 Theoretical framework

The theoretical framework casts guidance as deterministic control and derives flow-map-based approximations for single-trajectory sampling. It also characterizes how greedy guidance affects reward and diversity, and how early stopping can recover variance.

2.1 Flow-based generative models

  • A learned flow map accelerates inference by replacing numerical integration with a single function evaluation from the base sample to the target.Numerical integration typically requires 50–100 network evaluations.Source: §2.1 · pp. 3–4

2.2 The guidance problem

  • Reward-tilt sampling requires integrating over stochastic trajectories, making it intractable in high dimensions; finite-particle gradient estimates can also be biased and high variance.High variance can arise from rare samples with high reward.Source: §2.2 · p. 4

2.3 An optimal control formulation

  • The control objective balances terminal reward against deviation from the base flow, with larger λ allowing more aggressive guidance but greater departure from the learned distribution.The exact guidance term is circular because it requires the optimally guided flow-map endpoint during inference.Source: §2.3 · pp. 4–5

2.4 A tractable approximation

  • In the small-λ limit, the HJB equation yields a tractable guidance signal from a pre-trained flow map; the same signal is optimal for a greedy short-window control as the window shrinks.The authors justify using the signal beyond the small-λ regime because both analyses produce it.Source: §2.4 · p. 5
  • DPS is a coarsest-level approximation in the guidance hierarchy, replacing the exact flow-map endpoint and Jacobian with one-step Euler approximations.Other methods, including FlowDPS, FlowChef, and MPGD, also fall within the hierarchy.Source: §2.4 · p. 6

2.5 On-manifold guidance

  • When the data distribution lies on a smooth manifold, the flow-map Jacobian removes reward-gradient components orthogonal to it, keeping the controlled trajectory on-manifold.FMRG-J retains the Jacobian-projected signal, whereas FMRG-E drops the Jacobian for memory efficiency.Source: §2.5 · p. 6

2.6 Terminal distribution and early stopping

  • In the Gaussian setting, greedy guidance achieves the highest expected reward for a given λ but contracts variance more aggressively; early stopping can mitigate mode collapse.At t_stop = 0.3, early stopping recovers variance comparable to reward tilting.Source: §2.6 · p. 7, Figure 4 · p. 7
  • Guiding only until t_stop and then following the uncontrolled flow can control mode collapse; in the Gaussian setting, a suitable stopping time recovers exact control’s polynomial variance scaling.Smaller t_stop ends guidance earlier, while t_stop = 1 guides the entire trajectory.Source: §2.6 · p. 7

3 Algorithmic framework

FMRG implements deterministic reward guidance by combining flow-map integration with a guidance update, with design choices for the gradient, optimization steps, and NFE budget. Its Jacobian and Euclidean variants trade off on-manifold behavior against direct reward optimization.

3 Algorithmic framework

  • Operator splitting integrates the base dynamics with the pre-trained flow map and applies the guidance term separately.The framework organizes methods by endpoint approximation, base-dynamics integration, and guidance weight.Source: §3 · p. 8, Table 8 · p. 9
  • Multiple reward-gradient steps can be taken within each time interval, and they may also be used before the first flow step for seed optimization.The number of gradient steps per interval is denoted nopt.Source: §3 · p. 9
  • FMRG-J backpropagates through the flow-map Jacobian to project reward gradients onto the data-manifold tangent space, while FMRG-E uses the Euclidean gradient directly.FMRG-E avoids backpropagation through the flow map but may leave the data manifold; FMRG-J keeps trajectories on-manifold and more robustly preserves features, whereas FMRG-E can obtain higher reward with artifacts.Source: Figure 7 · p. 8, §3 · p. 9
  • FMRG-E is effective when the reward landscape aligns with the data manifold, while FMRG-J tends to work better for complex neural-network-based rewards.The Euclidean variant has a better NFE–performance trade-off for ℓ2 reconstruction, whose optima lie close to the manifold.Source: §3 · p. 9
  • Linearizing intermediate flow advancement reduces FMRG-E to one NFE per step while retaining a full flow-map lookahead for reward estimation.With early stopping, generation can finish with a single flow-map step.Source: §3 · p. 9
  • A small number of independent reinitializations can address dependence on the initial noise sample; at most four suffice even for complex neural-network rewards.Each reinitialization starts a fresh single-trajectory guided run, unlike SMC's evolving particle population.Source: §3 · p. 10

4 Related work

Prior guidance methods either maintain costly particle populations or apply heuristics to single trajectories, while reward fine-tuning requires additional training. FMRG instead uses a pre-trained flow map for single-trajectory guidance, and the cited inverse-problem results show strong performance at low NFEs.

4 Related work

  • FMRG uses the flow map for both integration and guidance along a single deterministic trajectory, unlike multi-particle flow-map tilting.The paper reports up to a 70× reduction in NFE at matched quality versus flow-map trajectory tilting.Source: §4 · p. 10
  • FMRG provides an optimal-control foundation for heuristic single-trajectory methods, recovering DPS and related approaches through endpoint approximation rather than a small-variance assumption.Replacing the exact flow-map endpoint with one Euler step recovers DPS-type methods.Source: §4 · p. 10
  • FMRG achieves reward alignment at inference without retraining or modifying the underlying flow, unlike reward fine-tuning and approaches requiring new models or auxiliary training.This allows different rewards to be applied to a single checkpoint.Source: §4 · pp. 10–11
  • FMRG outperforms baselines on super-resolution, motion deblurring, and inpainting at low NFEs, and FMRG-E performs notably better in the low-NFE AFHQ trade-off.Figure 10 shows the inverse-problem comparisons and the LPIPS-versus-FID trade-off.Source: Figure 10 · p. 11

5 Experiments

FMRG delivers competitive or superior reward-guided generation across inverse problems, style transfer, human-preference rewards, and VLM rewards, with large NFE reductions. Results and design-choice analyses show how performance varies with reward type and guidance configuration.

5 Experiments

  • Across inverse problems and reward-guided generation, FMRG achieves competitive or superior sample quality with up to 10× fewer NFEs on inverse problems and 70× fewer on reward-guided generation.Experiments include latent-space reconstruction losses, human-preference rewards, and VLM rewards.Source: §5 · p. 11

5.1 Latent-space inverse problems

  • FMRG outperforms all baselines on AFHQ-Dog and FFHQ, with FMRG-E competitive at 3 and 6 NFEs and both FMRG variants strongest at higher NFEs.The advantage is especially pronounced in unconditional generation, where FlowDPS and FlowChef degrade significantly.Source: §5.1 · p. 12

5.2 Style guidance

  • FMRG-J better preserves image quality while incorporating target style, whereas FMRG-E optimizes style more directly; DPS fails to incorporate style and FlowChef produces artifacts.The comparison is consistent with the proposed approximation hierarchy.Source: Figure 11 · p. 12, §5.2 · p. 13

5.3 Reward-guided generation

5.4 VLM guidance

  • FMRG-J steers generations toward detailed, compositional prompts that unguided FLUX misses, including object attributes, spatial relationships, and text rendering.The guidance uses a 7B-parameter VLM reward for prompt-image alignment.Source: §5.4 · p. 14, Figure 15 · p. 15

5.5 Analysis of design choices

  • Early stopping helps most for preference rewards susceptible to adversarial maximization, while Jacobian guidance matters more when reward optima may lie far from the data manifold.Without early stopping, preference-guided generation degrades from over-optimization artifacts; FMRG-J leads FMRG-E by 3.4 points at matched NFE for human-preference rewards.Source: §5.5 · p. 15

6 Conclusion and limitations

The paper presents FMRG as an inference-time guidance framework grounded in deterministic optimal control and flow maps. It also identifies dependence on pre-trained flow-map quality, hardware costs for FMRG-J, and over-optimization without early stopping as limitations.

6 Conclusion and limitations

  • FMRG frames inference-time guidance as deterministic optimal control, using a flow-map-centric approach for efficient sampling and strong practical performance.The paper positions this as an alternative to approximate reward-tilted sampling.Source: §6 · p. 16
  • FMRG depends on a pre-trained flow map and its quality; at 512px, FMRG-J uses about 48 GB VRAM, while FMRG-E runs on one L40S GPU.Without early stopping, greedy guidance can over-optimize rewards and cause mode collapse and artifacts.Source: §6 · p. 16

A Background on flow maps

  • A flow map X_s,t maps a state at time s to its state at time t along the probability flow, and its parameterized velocity recovers the drift on the diagonal.The paper calls this relation the tangent condition.Source: A Background on flow maps · p. 21

B Background on reward tilting and stochastic optimal control

  • Under the memoryless condition, the optimally controlled process has a reward-tilted terminal marginal, whose drift correction is the gradient of the stochastic-control value function.The cited background associates this formulation with auxiliary stochastic trajectories.Source: B Background on reward tilting and… · p. 22

C.1 Setup and regularity assumptions

  • The omitted proofs establish regularity-based bounds for flow-map Jacobians and Hessians that support analysis of the greedy approximation.The setup assumes a domain containing relevant trajectories and standard regularity conditions.Source: §C.1 · pp. 23–24

C.2 Proof of Proposition 2.1

  • The proof characterizes optimal guidance through the adjoint dynamics and derives the optimal control using the flow-map Jacobian.The adjoint equation is linked to the flow-map Jacobian evolution.Source: §C.2 · pp. 24–25

C.3 HJB characterization and small-λ expansion

  • The small-λ analysis identifies greedy guidance as the λ = 0 limit and states that it is accurate to O(λ^2).Source: §C.3 · p. 25

C.4 Proof of Proposition 2.3

The proof derives the optimal guidance signal as the limit of a short control window. It then uses a scalar Gaussian setting with a quadratic reward to compare reward tilting and optimal control.

C.4 Proof of Proposition 2.3

  • As the control window shrinks to zero, the restricted problem’s optimal guidance signal converges to the stated guidance expression.The proof relates controlled and uncontrolled flow maps to compute first-order perturbations.Source: §C.4 · p. 27

C.5 Gaussian case

  • The Gaussian analysis uses a standard Gaussian base, a Gaussian target, and a quadratic reward, with the scalar case chosen for simpler equations.The analysis is stated to extend to multivariate cases with dense covariance matrices.Source: §C.5 · p. 28

C.5.1 Reward-tilted distribution

  • For the Gaussian target and quadratic reward, the reward-tilted distribution remains Gaussian and shifts its mean toward the reward maximum.Source: §C.5.1 · pp. 28–29

C.5.2 Greedy guidance

The Gaussian analysis derives closed-form behavior for greedy guidance, including the effect of early stopping on terminal variance. It also characterizes exact optimal control as an analytical LQR solution in this setting.

C.5.2 Greedy guidance

  • In the Gaussian setting, greedy guidance yields a Gaussian terminal distribution, and early stopping can recover the rational-squared variance scaling of exact optimal control.A unique stopping time matches the exact terminal variance and depends continuously on λ.Source: §C.5.2 · pp. 29–30

C.5.3 Exact optimal control

  • For a Gaussian base and target with quadratic reward, the dynamics and costs form an LQR problem with an analytical optimal-control solution.Source: §C.5.3 · p. 31
  • The exact and greedy controls contract toward the same target trajectory, differing only in gain.The exact solution has a quadratic value function and a closed-form terminal Gaussian distribution.Source: §C.5.3 · pp. 31–32

C.5.4 Comparison of guidance schemes

The comparison characterizes how reward tilting, greedy guidance, and exact optimal control trade off reward concentration and variance. It also identifies early stopping as a way to tune reward optimization against diversity.

Controls

  • Both trajectory-based schemes use linear-in-x controls that drive trajectories toward a common reference path ending at reward maximum a.The greedy and exact controls differ in their gain.Source: §C.5.4 · p. 33

Variances

  • As λ →∞, reward tilting concentrates at O(1/λ), exact optimal control at O(1/λ²), and greedy guidance contracts exponentially.Greedy guidance applies an instantaneous contraction without accounting for future control effort; the exact gain is softened by the Riccati equation.Source: §C.5.4 · p. 33

Expected reward

  • Greedy guidance achieves exponentially higher expected reward than exact optimal control and reward tilting, at the cost of more aggressive variance reduction.The comparison quantifies a reward–diversity trade-off.Source: §C.5.4 · pp. 32–33

Early stopping

  • Stopping greedy guidance at tstop=0.3 preserves variance comparable to the target while shifting the mean toward the reward maximum.Early stopping truncates the contraction integral and provides a tunable reward–diversity trade-off.Source: §C.5.4 · p. 33

C.6 Proof of Proposition 2.5

The proof shows that the flow-map Jacobian maps perturbations into the endpoint manifold’s tangent space. Consequently, transposed-Jacobian reward gradients depend only on their tangent component.

Proof of Proposition 2.5

  • The image of ∇X_t,1(x_t) lies in the tangent space T_x1M because perturbing x_t traces a curve on M through the endpoint.The proof applies the Jacobian to an arbitrary direction and identifies the resulting vector as the curve’s tangent at the endpoint.Source: §C.6 · pp. 33–34
  • For any reward-gradient vector v, ∇X_t,1(x_t)^T v equals ∇X_t,1(x_t)^T v∥, where v∥ is its projection onto T_x1M.The component orthogonal to the tangent space has zero inner product with every Jacobian image.Source: §C.6 · pp. 34–35

D Existing methods as special cases of FMRG

The appendix interprets several guidance methods as special cases or approximations within FMRG. Their relationships largely turn on whether they use the full flow map or a one-step velocity approximation, and on their guidance weights.

D Existing methods as special cases of FMRG

  • FMRG provides a unifying framework that recovers existing guidance methods as special cases or approximations.The appendix extends the DPS connection to FlowDPS and FlowChef, which reduce to FMRG-E with different effective guidance weights.Source: D Existing methods as special cases of… · p. 35

D.1 DPS as a one-step approximation

  • DPS-style guidance is a coarse approximation to greedy guidance: it backpropagates reward through a posterior-mean endpoint estimate rather than the full flow map.For the linear interpolant, the posterior mean is a single Euler step; the full flow map is the limit of infinitely many Euler steps.Source: §D.1 · p. 36

D.2 FlowDPS as an approximation of FMRG-E

  • FlowDPS approximates FMRG-E by replacing both flow-map velocities in its operator-splitting update with the instantaneous velocity.Its denoiser-based data prediction is the one-step Euler estimate of the flow-map endpoint.Source: §D.2 · pp. 36–37

D.3 FlowChef as an approximation of FMRG-E

  • FlowChef makes the same two velocity substitutions as FlowDPS and discards the Jacobian factor in its endpoint reward-gradient update.Its gradient correction uses the terminal reward gradient at the endpoint estimate.Source: §D.3 · p. 37

D.4 MPGD as an approximation of FMRG-E

  • MPGD without projection updates the Tweedie clean estimate directly, while MPGD-AE additionally applies an autoencoder projection.In the flow setting, MPGD without projection has the same update structure as FlowDPS and FlowChef, with a different effective guidance weight.Source: §D.4 · p. 38

D.5 Summary of approximations

  • FlowDPS, FlowChef, and MPGD replace off-diagonal flow-map velocities with the instantaneous velocity, while FMRG uses the flow map for integration and on-manifold endpoint predictions.The passage reports effective weights of λ·Δt for FMRG-E, constant s′ for FlowChef, and t+·ct for MPGD; FMRG enables guidance in as few as 3 NFEs.Source: §D.5 · p. 39

D.6 ReNO and D-Flow

  • ReNO and D-Flow restrict FMRG-J to t=0, optimizing the initial noise seed rather than correcting the trajectory during generation.ReNO approximates the flow map with a one-step generator, whereas D-Flow backpropagates through a multi-step ODE solver.Source: §D.6 · p. 39

E.1 Flow map model

  • The experiments use flow maps distilled from FLUX.1-Dev via Lagrangian distillation and parameterized as LoRA adapters.The inverse-problem and reward-guided experiments use rank-32 and rank-64 adapters, respectively.Source: §E.1 · p. 39

E.2 Practical implementation

The implementation configures guidance duration, per-step flow-map evaluations, and guidance strength. The choice to reuse an endpoint trades accuracy at higher NFE budgets for more guided steps at low budgets.

Algorithm parameters

  • The early stopping time tstop sets how long guidance is applied before the remaining trajectory is completed with one unguided flow-map step.A smaller tstop means earlier stopping.Source: §E.2 · p. 39

Flow map stepping

  • Reusing the flow-map endpoint for both reward gradients and stepping performs better at low NFEs, while separate evaluations perform better at higher NFEs.The implementation uses one NFE per step for low-NFE configurations and two for high-NFE configurations.Source: §E.2 · pp. 39–40

Guidance strength λt

  • FMRG-E sets λt=η·t·(1−t+), downweighting updates near the endpoints.η is the single tuned constant in each case.Source: §E.2 · p. 40
  • FMRG-J normalizes the guidance signal to the flow-map velocity magnitude and sets λt=η·Δt, making η invariant to reward-model gradient scale.The passage notes this matters because human-preference and ℓ2-loss rewards have vastly different gradient scales.Source: §E.2 · p. 40

E.3 Latent-space inverse problems

The inverse-problem experiments compare FMRG with baselines on AFHQ-Dog and FFHQ, including prompt-removal and early-stopping ablations. In style transfer, FMRG-J achieves the best reported style similarity while remaining competitive on text alignment.

E.3 Latent-space inverse problems

  • The inverse-problem evaluation uses 1,000 256 × 256 images each from AFHQ-Dog and FFHQ, with 4× average-pooling, motion blur, or masked-region inpainting degradations.Hyperparameters are selected on 50 held-out images and evaluated on 1,000 separate images.Source: §E.3 · p. 40
  • Early stopping helps FMRG-E on super-resolution at low NFEs, but its benefits diminish at higher NFEs.The evaluated reward is negative ℓ2 reconstruction loss, and observations include added Gaussian noise with σ = 0.03.Source: §E.3 · p. 40, Table E.3 · p. 41
  • When prompts are varied or removed, FMRG degrades minimally while FlowDPS and FlowChef degrade noticeably; the gap widens in unconditional generation.The prompt-removal comparison covers AFHQ and FFHQ.Source: §E.3 · p. 41, Figure E.6 · p. 44, Figure E.7 · p. 44
  • Across AFHQ-Dog and FFHQ, FMRG-E and FMRG-J dominate distortion and perception metrics at every NFE budget, with the largest gap at low NFEs.The full metric plots cover text-conditional generation.Source: Figure E.4 · p. 42, Figure E.5 · p. 43
  • Qualitative comparisons show sharper detail from FMRG on 4× super-resolution, observation-consistent motion deblurring, and context-consistent box inpainting.Examples include both AFHQ and FFHQ images.Source: Figure E.8 · p. 45, Figure E.9 · p. 46, Figure E.10 · p. 47

E.4 Style transfer

  • For style transfer, FMRG-J achieves the best CLIP-I style similarity and competitive CLIP-T text alignment.The style loss is the Frobenius norm between Gram matrices of CLIP ViT-B/16 layer-2 features from reference and generated images.Source: §E.4 · p. 52

E.5 Reward-guided generation

For reward-guided generation, the experiments examine GenEval early stopping, aesthetic enhancement, and warmup particle selection. They also describe stochastic renoising for VLM guidance and report wall-clock time and peak VRAM for the evaluated methods.

E.5 Reward-guided generation

  • A warmup with K = 3 independent noise particles selects the highest-reward predicted endpoint before continuing guidance on that single trajectory.The paper reports that K = 3 restarts is sufficient to consistently improve performance.Source: §E.5 · p. 54
  • On GenEval, tstop = 0.25 consistently outperforms other tested levels, while removing early stopping produces over-optimization artifacts.The ablation tests FMRG-J at NFE 30 and 100; the qualitative examples show improving image quality as tstop decreases from 0.75 to 0.25.Source: §E.5 · p. 54, Table E.18 · p. 55, Figure E.19 · p. 55
  • Additional aesthetic comparisons show FMRG producing more dramatic enhancements than ReNO at matched low NFE budgets.The comparisons are shown in Figures E.20 and E.21.Source: §E.5 · p. 54, Figure E.20 · p. 56, Figure E.21 · p. 57

E.6 VLM guidance

  • For VLM-guided generation, injecting a small amount of stochastic noise late in generation is reported to reduce minor artifacts and sharpen image quality.The procedure adds diversity without significantly altering the guided trajectory.Source: §E.6 · p. 58

E.7 Wall-clock time and VRAM

  • The paper reports wall-clock time per image and peak VRAM for reward-guided generation at 512px on H100 and inverse problems at 256px on L40S.The measurements are presented in Tables E.22 and E.23.Source: §E.7 · p. 58, Table E.22 · p. 58, Table E.23 · p. 58
Loading 2604.27147v3…