Flow Map Trajectory Tilting (FMTT)

• Flow Map Trajectory Tilting (FMTT) is a mathematically principled test-time adaptation method that leverages the exact flow-map look-ahead in diffusion models to incorporate terminal rewards.
• It integrates the underlying ODE/SDE dynamics with direct reward evaluation to achieve unbiased sample generation and improved efficiency over denoiser-based approximations.
• FMTT offers provable guarantees for reward ascent and reduced sampling variance, making it highly applicable for complex tasks like image editing and semantic control.

Flow Map Trajectory Tilting (FMTT) is a mathematically principled test-time adaptation technique for diffusion models, introduced to address the challenge of maximizing user-specified reward functions—such as classifier log-likelihoods or vision–LLM (VLM) scores—that are only well defined at the endpoint of the generation process. Leveraging the flow map associated with the deterministic or stochastic ODE/SDE underlying the diffusion process, FMTT enables exact look-ahead to final samples, yielding both unbiased sampling via exact importance weighting and efficient search for reward-maximizing samples. It stands in contrast to prior methods that rely on myopic approximations, such as denoiser-based look-ahead, and provides provable guarantees for reward ascent and sample efficiency (Sabour et al., 27 Nov 2025).

1. Mathematical Background and Flow Maps

A diffusion model generates a path of densities $\{\rho_t(x)\}_{t\in[0,1]}$ interpolating from a simple noise prior $\rho_0(x)$ (e.g., $N(0,I)$) to a data density $\rho_1(x)$, typically via:

• The stochastic process (SDE):

$$d x_t = [b_t(x_t) + \epsilon_t s_t(x_t)]\,dt + \sqrt{2\epsilon_t}\,dW_t$$

• Or the deterministic ODE (“probability flow”):

$d x_t = b_t(x_t)\,dt$

Here, $b_t(x)$ is the drift, $s_t(x) = \nabla_x \log \rho_t(x)$ the score, and $\epsilon_t \geq 0$ a user-chosen noise schedule.

The instantaneous velocity field $v(t, x)$ is defined as $dx_t = v(t, x_t)dt$, typically with $v(t, x) = b_t(x)$ in the ODE setting.

The two-time flow map $\Phi_{s \to t}(x)$ (also denoted $X_{s, t}(x)$) integrates the ODE from $s$ to $t$: the marginal law $X_{0,1\sharp}\rho_0 = \rho_1$ implies one-shot sampling if $X_{0,1}$ is known. Key identities include the Eulerian equation $$\partial_s X_{s, t}(x) + b_s(x) \cdot \nabla_x X_{s, t}(x) = 0$$ and the tangent identity $\lim_{s\to t} \partial_t X_{s, t}(x) = b_t(x)$ (Sabour et al., 27 Nov 2025).

2. Incorporating Terminal Rewards through Flow-Map Look-Ahead

FMTT aims to generate samples from a "tilted" target $p_R(x) \propto \rho_1(x)\,e^{R(x)}$ for a terminal reward $R(x)$, which may only be computable at $t = 1$. Standard practice augments the SDE drift with an estimated $\nabla R$, but this is generally ill-posed since $R$ is undefined off the endpoint distribution. A typical workaround uses an approximate denoiser $D_t(x) \approx \mathbb{E}[x_1 | x_t]$, with $r_t(x) = t R(D_t(x))$, but this estimate is inaccurate at early $t$.

FMTT instead utilizes exact flow-map look-ahead: $r_t(x) = t R(\Phi_{t\to 1}(x))$, so each intermediate state is evaluated by "looking ahead" to its unique deterministic endpoint under the learned flow map. The resulting continuous-time SDE is:

$$d x_t = v_{t,t}(x_t)\,dt + \epsilon_t [s_t(x_t) + t \nabla_x R(\Phi_{t\to 1}(x_t))]\,dt + \sqrt{2\epsilon_t}\,dW_t$$

where $v_{t,t} = b_t$.

3. Importance Weighting and Sampling Algorithm

Sampling from the reward-tilted density $\hat \rho_t(x) = \rho_t(x) e^{r_t(x) + F_t}$, with $F_t = -\log \int \rho_t(x) e^{r_t(x)} dx$, requires exact trajectory weighting. Jarzynski’s equality yields the required importance weights: for each trajectory,

$$A_t = \int_0^t \left[ b_s \cdot \nabla r_s + \partial_s r_s \right](x_s) ds$$

When $r_t(x) = t R(\Phi_{t \to 1}(x))$, the log-weight simplifies to the accumulated terminal reward along the look-ahead path:

$A_t = \int_0^t R(\Phi_{s \to 1}(x_s)) ds$

Algorithmic Steps

Time is discretized ($t_0 = 0 < t_1 < \dots < t_K = 1$) and $N$ particles are evolved. For each step $k \to k+1$ and particle $n$:

• Propagate $x_k^n \to x_{k+1}^n$ using Euler–Maruyama update of the FMTT SDE.
• Update accumulated reward $$A_{k+1}^n \leftarrow A_k^n + (t_{k+1} - t_k) R(\Phi_{t_k \to 1}(x_k^n))$$.
• (Optional) Resample if the effective sample size (ESS) drops below a threshold.

After $K$ steps, weights $w^n = \exp(A_K^n)$ yield unbiased estimators for expectations under $\hat \rho_1$. Alternatively, keeping top-M particles at each resampling for high $R(\Phi_{t_k \to 1}(x_k))$ implements greedy search for reward maximizers (Sabour et al., 27 Nov 2025).

Step	Input	Update Operation
Propagation	$x_k^n$	Euler–Maruyama for FMTT SDE
Weight increment	$A_k^n$	$$A_{k+1}^n = A_k^n + \Delta t\,R(\Phi_{t_k \to 1}(x_k^n))$$
Resampling	All particles	If $\mathrm{ESS} < \tau$, reweight

4. Theoretical Guarantees

FMTT’s use of exact flow-map look-ahead yields provably more faithful alignment of drift with the true $\nabla R$ at generation endpoint: for any $t$,

$$\mathbb{E}[R(\Phi_{t\to 1}(x_t))] - \mathbb{E}[R(\Phi_{t\to 1}(\tilde x_t))]$$

is smaller (in first order) under the FMTT drift than standard gradient steering.

In the sequential Monte Carlo (SMC) context, the variance of the importance sampling normalizer is governed by the “incremental discrepancy” $$D(t, t+\Delta t) = \log[1 + \text{Var}_{x_t}[G_{t, t+\Delta t}]]$$ and the “thermodynamic length” $\Lambda = \sum_k \sqrt{D(t_{k-1}, t_k)}$. FMTT guarantees strictly lower $D_{tot}$ and $\Lambda$ than naïve gradient control, resulting in lower SMC variance and greater sampling efficiency.

5. Practical Implementation and Empirical Results

The learned flow map $\Phi_{t \to 1}(x_t)$ can be evaluated efficiently, in 1–4 network queries via "any-step" consistency models. This direct look-ahead enables the use of complex black-box rewards, such as those provided by VLMs, for which prior denoiser-approximate methods are ineffective.

Examples demonstrated include precise clock editing, geometric constraints (e.g., symmetry, anti-symmetry), and masked-region inpainting. FMTT enables text–image alignment with VLM rewards (e.g., Qwen2.5-VL, Skywork-VL), a regime where denoiser look-ahead fails. On GenEval (550+ prompts, human-rewarded), FMTT with beam search supports a mean object-alignment score of 0.79, compared to 0.75 for FLUX and 0.76 for multi–Best-of-N baselines. On UniGenBench++ (600 prompts, VLM-evaluated), FMTT at $\sim$2000 NFEs achieves $\sim$75% overall, outperforming Best-of-N ($\sim$73%).

Critically, one-step denoiser look-ahead shows no improvement over Best-of-N, affirming the necessity of the true flow-map signal for nontrivial reward maximization (Sabour et al., 27 Nov 2025).

6. Significance and Implications

FMTT establishes a general, theoretically sound protocol for test-time adaptation of flow-based and diffusion models toward sample selection or generation tasks defined by arbitrary or black-box reward functions. It addresses the challenge posed by rewards that are ill-defined away from terminal data distributions through explicit integration of the flow map, bypassing the limitations of denoiser-based surrogates.

A plausible implication is that the approach can serve as a foundation for downstream editing and control tasks requiring sample-efficient, unbiased, and reward-tailored sample generation—particularly in domains interfacing with complex, non-differentiable evaluators such as VLMs or multimodal classifiers.

7. Relation to Prior Look-Ahead and Gradient Steering Methods

Conventional test-time strategies inject reward gradients into SDEs, but in the presence of terminal-only rewards, this is ill-defined or yields poor alignment between drift and desired distribution. FMTT’s use of the exact flow map ensures that the look-ahead signal is both mathematically accurate and efficiently computable, yielding provable improvements in reward ascent and SMC variance.

Empirical comparisons demonstrate the inadequacy of one-step denoiser-based look-ahead for complex rewards and highlight FMTT’s superiority in practical tasks requiring intricate, semantic, or structural constraints at generation time.