Test-time scaling of diffusions with flow maps

TL;DR

Flow Map Trajectory Tilting (FMTT) is a principled reward-guided sampler for diffusion / flow-matching models that uses a flow map (à la Align Your Flow: Scaling Continuous-Time Flow Map Distillation) as the look-ahead operator inside reward-gradient guidance. Standard reward-guided diffusion has to evaluate a reward that is only defined at the clean endpoint, so common practice substitutes a one-step denoiser — which is blurry and off-manifold at high noise. FMTT replaces that denoiser with the few-step flow map, derives a remarkably simple Jarzynski/SMC importance weight via the flow map’s Eulerian identity, and shows that this both samples the exact tilted distribution and outperforms denoiser look-ahead empirically. As a headline application, it is the first sampler to make pretrained VLMs viable as natural-language reward functions for test-time scaling of text-to-image models, beating gradient-free best-of-N and gradient-based ReNO on GenEval and UniGenBench++.

Key claims

Reward-gradient guidance on a diffusion SDE is ill-posed because user rewards are defined on the data distribution at $t=1$ , not along the noisy trajectory; common workarounds use a one-step denoiser as a look-ahead, but at early noise levels the denoiser prediction is too blurry to give meaningful reward gradients [§2.2 “Naïve look-ahead” / “Denoiser look-ahead”, Fig. 2].
A learned two-time flow map $f(x_t, t, s)$ provides an any-step look-ahead that, by exploiting the Eulerian identity $v(x,t) = \partial_s f(x,t,s)|_{s=t}$ , accesses both the instantaneous velocity used in the SDE and a clean-image estimate in 1 or a few NFE per step [§2.1, Eq. (11)].
The reward-tilted SDE driven by $\nabla_x r(f(x_t,t,1))$ does not sample the tilted distribution exactly; a Jarzynski-style log-weight ODE corrects the bias, and with the flow-map look-ahead the correction reduces to simply integrating $\partial_t r(f(x_t,t,1))$ along the trajectory (Proposition 2.2) [§2.3, Eqs. (17)–(19)].
The flow map’s Eulerian equation $\partial_t f + (\partial_x f) v = 0$ is what collapses the otherwise ugly derivative chain in the importance weight to that single reward time-derivative term — i.e. flow-map structure is load-bearing, not cosmetic [Proof in §A.2].
Thermodynamic length of the tilting process — defined via the excess variance of the cumulative log-weights — is a resampling-free proxy for SMC efficiency; FMTT minimizes it relative to no-look-ahead and denoiser-look-ahead baselines on a controlled MNIST class-tilt task [§2.3 “Characterizing effectiveness”, §3.1, Fig. 9].
The same importance-weight machinery supports both exact sampling (resampling like SMC) and search (top-N replacement instead of resampling) — i.e. FMTT exposes a single dial between faithful tilted-distribution sampling and greedy reward maximization [§2.3 “From sampling to search”].
On GenEval over a distilled Align-Your-Flow 4-step flow map of FLUX.1-dev, FMTT reaches mean 0.79 — beating FLUX.1-dev itself (0.65 at 180 NFE), gradient-free best-of-N (0.75, 1440 NFE), multi-best-of-N (0.76, 1280 NFE), beam search (0.75, 1200 NFE), and gradient-based ReNO on the noise latent (0.71, 1280 NFE) [Table 1].
The flow-map look-ahead is causally responsible for the gain: swapping it for a 1-step denoiser drops the mean to 0.75, and swapping for a 4-step diffusion sampler (same NFE budget) drops to 0.75 [Table 1, FMTT ablation rows].
For geometric rewards (symmetry, anti-symmetry, rotation invariance, masking) that are deep in the tail of the base model, best-of-N and ReNO fail to satisfy the constraint while FMTT’s trajectory-tilting drives the generation onto the constraint manifold; the flow-map look-ahead is necessary to fully maximize the masked reward, denoiser look-ahead only partially does [§3.2 “Geometric Transformation Rewards”, Figs. 5, 7].
FMTT enables, for the first time as a test-time-scaling sampler, pretrained VLMs (Skywork-VL-Reward, Qwen2.5-VL-7B-Instruct) as reward functions specified entirely in natural language via yes/no questions; on UniGenBench++ the compute-vs-score curve shows the flow-map look-ahead scaling efficiently while a 1-step-denoiser look-ahead barely beats best-of-N [§3.2 “VLMs as a Judge”, Fig. 8].
The denoiser-look-ahead curve flattens because its high-noise predictions are blurred off the data manifold and so produce unhelpful reward gradients from a VLM that has only seen natural images — explaining why VLM-as-reward only becomes practical once a flow map is available [§3.2 paragraph after Fig. 8].

Method

FMTT operates on top of a pretrained two-time flow map $f_\theta(x_t, t, s)$ (the paper uses a 4-step distillation of FLUX.1-dev from Align Your Flow). For a user-specified terminal reward $r$ , the time-dependent reward along the trajectory is defined by composition with the flow map: $\tilde r(x,t) := r(f_\theta(x, t, 1))$ . The sampler runs

dX_t = \big[b(X_t,t) + \tfrac{1}{2}(\epsilon_t^2 + \sigma_t^2)\,\nabla_x \tilde r(X_t,t)\big]\,dt + \sigma_t\,dW_t

with $b$ recovered from the flow map via the Eulerian identity $v = \partial_s f|_{s=t}$ , and the score implicit in the stochastic-interpolant parameterization. In parallel, each particle carries a log-weight $A_t$ that satisfies

dA_t = \partial_t \tilde r(X_t, t)\,dt

— this is the simplification Proposition 2.2 buys, where without the flow map the importance weight would involve $\partial_t r + (\nabla r)\cdot b + \tfrac{1}{2}\Delta r$ . Standard SMC resampling with weights $\exp(A_t)/\sum \exp(A_t)$ gives an unbiased sampler of the tilted distribution; substituting top-N selection for resampling turns FMTT into a gradient-based search algorithm that still benefits from the flow-map look-ahead. A variant (“reward-modified drift”) additionally adds the reward gradient to the velocity itself, with a compensating term in the weight ODE that keeps the marginal correct [§2.3 “Reward-modified drift”, Eqs. (20)–(21)].

For VLM-as-reward, the reward at clean image $x_1$ is the softmax-probability the VLM assigns to “yes” on a yes/no caption-correctness question: $r(x_1) = p_{\text{VLM}}(\text{yes} \mid x_1, q)$ — making the entire reward channel definable in natural language, and supporting multi-image queries via VLMs that accept image lists.

Results

GenEval (Table 1, base = AYF-distilled FLUX.1-dev 4-step flow map): FMTT 0.79 mean (Single Obj. 1.00 / Two Obj. 0.97 / Counting 0.90 / Colors 0.91 / Position 0.30 / Attr. Binding 0.64) at 1400 NFE vs. base flow map 0.62 (16 NFE), FLUX.1-dev 0.65 (180 NFE), best-of-N 0.73–0.75 (128–1440 NFE), multi-best-of-N 0.76 (1280 NFE), beam search 0.75 (1200 NFE), ReNO 0.71 (1280 NFE).
GenEval look-ahead ablation (Table 1): replacing the flow map with a 1-step denoiser drops FMTT to 0.75 mean (at 350 NFE); with a 4-step diffusion look-ahead, 0.75 (at 1400 NFE).
UniGenBench++ with Skywork-VL reward: the compute-vs-score scaling plot (Fig. 8) shows FMTT-with-flow-map dominates both best-of-N and FMTT-with-1-step-denoiser across the compute axis; the denoiser look-ahead barely improves over best-of-N, attributed to blurred high-noise denoiser predictions [§3.2 “VLMs as a Judge”].
Geometric rewards (Figs. 5, 7): on rotational, symmetry, and masking constraints, FMTT-with-flow-map satisfies the constraint while best-of-N, ReNO, multi-best-of-N produce either blurry or constraint-violating outputs; FMTT-with-denoiser maximizes masked-region darkness but fails to relocate content out of the masked region.
MNIST exact-sampling diagnostic (Fig. 9): all three look-ahead variants are unbiased (their reward and class entropy match the ground-truth tilted-distribution values within standard error), but the flow-map look-ahead has the smallest total discrepancy and thermodynamic length — i.e. more efficient sampling per particle budget.

Why it’s interesting

This is the first paper on the wiki that frames a flow map as the lookahead operator inside reward guidance rather than as a fast sampler, and it brings together three threads that the wiki has been tracking separately. First, it directly extends Align Your Flow: Scaling Continuous-Time Flow Map Distillation — same authors, same NVIDIA distillation of FLUX.1-dev — into a test-time-scaling regime: AYF was the SOTA few-step distillation, FMTT is what to do with that 4-step flow map at inference once you also want reward signal. The flow map stops being a curiosity that beats consistency models on FID and becomes the load-bearing piece of a guidance algorithm. Second, it complements Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models on the Diffusion Distillation reward-post-training thread: where HyperNoise amortizes the reward into a trained noise modulator (one-time training, cheap inference), FMTT spends test-time compute to integrate the reward gradient — both target distilled few-step models, but the inference-time / training-time tradeoff is opposite, and the recipes look composable (train a HyperNoise once, then run FMTT at inference for the long tail). Third, the VLM-as-reward result is the first sharp instance the wiki has filed where the VLM-as-Evaluator thread becomes actionable as a training-free reward channel — and the failure mode of denoiser look-ahead (blurry predictions are off-manifold for VLMs trained only on natural images) predicts that any “VLM-judges-the-noisy-trajectory” recipe needs a look-ahead at least as good as a few-step flow map. The reward-hacking caveat — yes/no questions need to be detailed enough to prevent the search loop exploiting loopholes — is the same brittleness the Classifier-Free Guidance variants cluster sees in other inference-only reward stacks.