Test-time scaling of diffusions with flow maps (2511.22688v1)

Abstract: A common recipe to improve diffusion models at test-time so that samples score highly against a user-specified reward is to introduce the gradient of the reward into the dynamics of the diffusion itself. This procedure is often ill posed, as user-specified rewards are usually only well defined on the data distribution at the end of generation. While common workarounds to this problem are to use a denoiser to estimate what a sample would have been at the end of generation, we propose a simple solution to this problem by working directly with a flow map. By exploiting a relationship between the flow map and velocity field governing the instantaneous transport, we construct an algorithm, Flow Map Trajectory Tilting (FMTT), which provably performs better ascent on the reward than standard test-time methods involving the gradient of the reward. The approach can be used to either perform exact sampling via importance weighting or principled search that identifies local maximizers of the reward-tilted distribution. We demonstrate the efficacy of our approach against other look-ahead techniques, and show how the flow map enables engagement with complicated reward functions that make possible new forms of image editing, e.g. by interfacing with vision LLMs.

• The paper introduces FMTT, a method using learned flow maps to integrate terminal reward gradients during test-time in diffusion models.
• The paper demonstrates that FMTT outperforms traditional denoiser guidance in tasks such as MNIST generation and text-to-image synthesis, ensuring unbiased sampling.
• The paper establishes a robust theoretical framework with flexible tilting strategies, enabling advanced reward-conditioned generative modeling.

Test-Time Scaling of Diffusions with Flow Maps: A Technical Summary

Introduction

This work presents Flow Map Trajectory Tilting (FMTT), a methodology for adapting diffusion and flow-based generative models at inference time to better satisfy complex, often terminal, reward functions. The core insight is that model guidance via reward gradients—routinely employed for improved sample alignment—suffers from practical and theoretical deficiencies when rewards are only well-defined at the data manifold (i.e., terminal states). FMTT introduces a generic solution leveraging the predictive capabilities of neural flow maps, which directly learn the time-integrated solution operator of a generative process, thereby providing a superior mechanism for test-time reward incorporation compared to heuristic denoiser-based approaches.

Figure 1: Test-time search can overcome model biases and reliably sample from regions of the distribution (e.g., precise clock times) that baselines fail to capture.

Flow Maps and Inference-Time Guidance

Classical diffusion and flow models encode generation as the integration of an ODE or SDE from a simple base distribution to the data manifold. Guidance techniques inject reward gradients into these dynamics; however, since user rewards ($r(x)$) are typically defined only at $x \sim \rho_1$ (terminal samples), approximating the signal at intermediate steps is necessary. Using a score-based denoiser for a single-step look-ahead is a common workaround, but this provides only limited reward signal, especially in early, high-noise regime.

FMTT capitalizes on the structure of learned flow maps, which allow arbitrary-step look-ahead via directly modeling $X_{s, t}(x_s) = x_t$. This enables precise evaluation of terminal-state rewards at any intermediate point, greatly enhancing the informativeness of reward gradients throughout the generative trajectory.

Figure 2: Schematic overview of test-time adaptation of diffusions with flow map tilting. The flow map enables reward incorporation through look-ahead states, yielding more targeted trajectory guidance.

Theoretical Framework

FMTT grounds test-time guidance in a well-principled framework. The method samples from the reward-tilted distribution

$\hat{\rho}_1(x) = \rho_1(x) e^{r(x) + \hat{F}}$

by evolving a modified SDE that incorporates both the original drift and the gradient of the terminal reward evaluated through the flow map:

$$d\tilde{x}_t = v_{t,t}(\tilde{x}_t) dt + \epsilon_t [s_t(\tilde{x}_t) + t \nabla_{\tilde{x}_t} r(X_{t,1}(\tilde{x}_t))] dt + \sqrt{2 \epsilon_t} dW_t$$

Augmentation with importance weights, derived via a Jarzynski estimator, yields unbiased sampling from the reward-tilted distribution. Notably, when the time-dependent reward is defined via the flow map, the weight accumulation simplifies to a trajectory integral of the reward:

$A_t = \int_0^t r(X_{s,1}(\tilde{x}_s)) \, ds$

These theoretical results generalize to different "tilting" strategies, allowing for flexible search and sampling protocols, including greedy maximization, SMC-driven sampling, and hybrid approaches.

Empirical Evaluation

MNIST: Thermodynamic Discrepancy and Weight Efficiency

On class-conditional generation with MNIST, FMTT delivers lower total discrepancy and thermodynamic length—proxies for sampling efficiency—than conventional denoiser-based or naive guidance. Only FMTT achieves unbiased estimation of the desired target distribution with error bars overlapping the ground truth, highlighting both its effectiveness and its theoretical optimality.

Figure 3: Comparison of MNIST tilted sampling to generate digits that would be classified as zeros. In FMTT, flow map look-ahead yields superior reward targeting and sample diversity.

High-Dimensional Text-to-Image: Human, Geometric, and VLM Rewards

Experiments on challenging open-text-to-image tasks (building on distilled flow maps from "Align Your Flow") demonstrate that FMTT is robust to reward function complexity:

• Human preference rewards: Benchmarking on GenEval shows FMTT marginally surpasses competitive baselines in image-object alignment, even when the base model is already highly rewarded-aligned.
• Geometric invariance rewards: FMTT significantly outperforms both best-of-$N$ and gradient-based denoiser look-ahead in enforcing symmetry, anti-symmetry, and rotation invariance, producing sharper, more compliant samples.

  Figure 4: Qualitative comparison on three basic geometric rewards. Flow map-based gradient guidance sharpens outputs and enforces structural constraints more reliably than prior approaches.

• Masked and localized rewards: Only FMTT with flow map look-ahead can reliably localize content to unmasked regions, accomplishing penalization or incentivization in constrained image areas.

  Figure 5: Only flow map-based FMTT concentrates content reliably in masked reward tasks, outperforming denoiser and model-only baselines.

• Vision-LLM (VLM) Rewards: Using VLMs as evaluators, with rewards specified via natural language, FMTT achieves stronger scaling on the UniGenBench++ benchmark compared to best-of-$N$ and denoiser-based methods. Denoiser look-ahead fails to reliably improve VLM scores, likely due to the mismatch between intermediate blurred states and terminal evaluation criteria.

  Figure 6: Scaling results on UniGenBench++. FMTT with flow map look-ahead consistently attains higher reward scores per compute than competing strategies.

Qualitative Analysis and Reward Hacking

FMTT consistently produces samples more closely adhering to specified rewards—whether human-alignment, geometric, or natural language-defined—than other search paradigms.

Figure 7: Prompts where the base model fails are corrected by FMTT, with flow map look-ahead yielding the most robust improvements.

However, maximizing complex, non-robust rewards (e.g., VLM-based) can uncover failure modes such as reward hacking, where the search discovers undesirable yet high-reward configurations (e.g., cheating via written numerals).

Figure 8: VLM reward hacking—flow map-based search exploits reward model weaknesses by embedding cheat cues to maximize reward.

Implications and Future Directions

The introduction of FMTT offers a principled, theoretically justified, and highly versatile approach for inference-time scaling of diffusion models. Unlike approaches reliant on score-based denoiser look-ahead, FMTT leverages the function approximation capabilities of learned flow maps to provide efficient, trajectory-wise reward incorporation. This leads to strong empirical improvements on reward-aligned generation—most notably under non-stationary, highly compositional, and VLM-defined reward conditions—without retraining or data augmentation.

Given the generality of the theory and flexibility of empirical instantiations, FMTT is likely to become foundational for reward-aligned generative modeling in image, video, and broader domains, further bridging the gap between conditional generation capability and user-controllable model behavior. Future research will likely extend FMTT to joint distribution modeling, hierarchical flow control, and reinforcement learning settings, where the reward signal is itself dynamic and multi-modal.

Conclusion

FMTT demonstrates that flow map-based guidance is both mathematically superior and more practically robust for test-time scaling of diffusion and flow-based models. The methodology resolves a key limitation of reward gradient guidance by enabling consistent and efficient inference-time optimization, particularly when models are required to achieve nuanced or out-of-distribution objectives, such as those provided by VLMs or complex constraint sets.