Test-time scaling of diffusions with flow maps¶

会议: NeurIPS 2026
arXiv: 2511.22688
代码: https://flow-map-trajectory-tilting.github.io/
领域: VLM 推理
关键词: 扩散模型, 流映射, 测试时缩放, 奖励引导, 视觉语言模型

一句话总结¶

本文提出流映射轨迹倾斜（Flow Map Trajectory Tilting, FMTT），利用流映射具备任意时刻瞬时向前预测（look-ahead）终点的性质计算下游复杂奖励梯度，结合简化的 Jarzynski 重要性权重与序列蒙特卡洛/贪婪搜索，实现了扩散模型在测试时的高效可扩展推理与多模态 VLM 语义对齐。

研究背景与动机¶

基于扩散模型（Diffusion Models）和流匹配（Flow Matching）的连续生成模型在计算机视觉与跨学科科学生成中取得了巨大成功，其核心生成机制是将随机噪声通过数值求解常微分方程（ODE）或随机微分方程（SDE）传输至目标数据分布。在实际应用中，用户往往希望生成样本能够精确对齐特定任务的奖励函数（例如生成显示精确时间的表盘、严格对称的几何结构、或满足多模态复杂逻辑判别的图像）。然而，标准的引导采样（guidance）方法通常需要将测试时奖励的梯度注入到生成动力学中；而绝大多数用户自定义的奖励函数（如预训练分类器、美学评估模型或 VLM 评判模型）仅在生成轨迹终点 $t=1$ 的无噪真实图像流形上有良好定义，在中间高噪声状态下毫无物理意义或缺乏有效信号。

这种困境引发了核心的技术矛盾：如果每一步扩散生成都通过数值求解 ODE 积分到终点来评估奖励，计算开销将随步数呈二次方激增而无法承受；如果像现有方案那样采用单步去噪器（1-step denoiser）作为启发式前瞻代理，在轨迹前期的强噪声区域，去噪器重构出的图像严重模糊且严重偏离数据流形，导致反向传播计算的奖励梯度极其嘈杂甚至完全误导生成轨迹。此外，朴素地在扩散漂移中加入奖励梯度会导致轨迹分布发生不可控的滞后（lag），无法从理论上精确采样目标倾斜分布（reward-tilted distribution）。

针对上述瓶颈，本文从流映射（flow map）理论切入。流映射直接学习流 ODE 的解算子，具备在任意时间步直接前瞻终点状态的解析能力。本文的核心 idea 是：利用流映射与瞬时速度场之间的切线与欧拉恒等式，构建流映射轨迹倾斜（FMTT）机制，不仅通过流映射前瞻在整条轨迹上提供高保真终点预测与稳定奖励梯度，还将 Jarzynski 恒等式下的重要性权重微分方程化简为沿前瞻轨迹累积奖励的极简形式，从而在采样与搜索两端实现严格且高效的测试时计算缩放。

方法详解¶

整体框架¶

FMTT 的核心目标是在无需微调模型参数的前提下，使生成过程能够精确采样受奖励倾斜的目标分布 $\hat{\rho}_1(x) \propto \rho_1(x) e^{r(x)}$，或在测试时通过贪婪前瞻搜索高奖励候选样本。系统由两大核心支柱构成：一是基于流映射 $X_{s,t}$ 的任意步前瞻算子，通过复合目标奖励定义连续时间依赖奖励 $r_t(x) = t \cdot r(X_{t,1}(x))$；二是基于粒子动力学与重要性权重演化的序列蒙特卡洛（SMC）/束搜索框架，通过注入奖励梯度引导粒子前行，并利用权重演化进行重采样或 top-$n$ 筛选。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["初始高斯噪声粒子 X_0 ~ N(0, I)"] --> B["前瞻状态预测与奖励评估<br/>X_t,1 前瞻终点 + 计算梯度 ∇r_t"]
    B --> C["动力学漂移与扩散更新<br/>复合速度场 + 引导梯度 + 随机扰动"]
    C --> D["重要性对数权重演化<br/>解析化简 dA_t/dt = r(X_t,1)"]
    D -->|采样模式 Sampling| E["SMC 权重重采样<br/>Softmax 概率加权克隆与淘汰"]
    D -->|搜索模式 Searching| F["Top-n 贪婪粒子筛选<br/>保留高奖励粒子并重置克隆"]
    E --> G["循环步进至 t=1 终点样本"]
    F --> G

关键设计¶

1. 基于流映射的精确前瞻与奖励梯度构建

传统扩散模型前瞻方法依赖经验去噪器 $\mathsf{D}_t(x) = \mathbb{E}[x_1 | I_t = x]$，在生成早期 $t \approx 0$ 时，去噪图像仅是整个流形的均值模糊图像，梯度的余弦相似度极低且梯度方向杂乱。而双时间流映射 $X_{s,t}(x)$ 满足流方程的积分算子关系 $X_{s,t}(x_s) = x_t$，并遵循欧拉方程 $\partial_s X_{s,t}(x) + b_s(x) \cdot \nabla X_{s,t}(x) = 0$。FMTT 直接将测试时奖励定义为 $r_t(x) = t \cdot r(X_{t,1}(x))$。借助流映射的切线恒等式 $\lim_{s \to t} \partial_t X_{s,t}(x) = v_{t,t}(x) = b_t(x)$，模型无需额外的嵌套 ODE 求解即可在单步前瞻中直接得到清晰、高保真的终点图像预测，从而计算出高质量的奖励梯度 $\nabla_{\tilde{x}_t} r_t(\tilde{x}_t) = t \nabla X_{t,1}(\tilde{x}_t)^\top \nabla r(X_{t,1}(\tilde{x}_t))$，在整条生成轨迹上提供稳定方向指导。

2. 基于非平衡统计物理的权重化简与无偏采样

朴素向 SDE 注入奖励梯度导致粒子的实际演化概率密度 $\tilde{\rho}_t$ 偏离理想的倾斜概率密度 $\hat{\rho}_t \propto \rho_t e^{r_t + \hat{F}_t}$。为了修正这一分布滞后，FMTT 基于扩展的 Jarzynski 恒等式引入轨迹权重补偿项 $A_t$。一般情况下，$A_t$ 的演化包含复杂的偏导数与散度项，但在流映射前瞻形式 $r_t(x) = t \cdot r(X_{t,1}(x))$ 下，结合流映射的欧拉恒等式 $\partial_t X_{t,1}(x) + \nabla X_{t,1}(x) b_t(x) = 0$，对数权重的演化微分方程产生了惊人的抵消与简化： $$\frac{\mathrm{d}A_t}{\mathrm{d}t} = r(X_{t,1}(\tilde{x}_t))$$ 该结论（论文 Proposition 2.2）证明了对数权重的瞬时增长率恰好等于终点前瞻状态下的原始奖励值。通过在 SMC 流程中计算归一化增量权重，算法能够以无偏估计对目标分布进行精确粒子重采样，并最小化热力学长度（Thermodynamic Length），实现极高采样效率。

3. 奖励修正漂移与测试时搜索泛化

为了在复杂高维空间（如高分辨率文生图）中进一步增强奖励上升速度，FMTT 进一步将奖励梯度扩展到漂移速度场本身，构建通用的广义 SDE： $$\mathrm{d}\tilde{x}_t = \left[ v_{t,t}(\tilde{x}_t) + \chi_t \nabla r_t(\tilde{x}_t) + \epsilon_t (s_t(\tilde{x}_t) + \nabla r_t(\tilde{x}_t)) \right] \mathrm{d}t + \sqrt{2\epsilon_t} \mathrm{d}W_t$$ 其中 $\chi_t$ 为漂移调控参数。在严格采样场景下，取 $\chi_t = 0$ 可使权重计算最为廉价且热力学长度最小；而在高维图像生成与离散搜索场景下，选取倾斜分数动力学参数 $\chi_t = \eta_t = \frac{1.05(1-t)}{t+0.05}$ 能够最大化奖励驱动效果。结合 Top-$n$ 剪枝与粒子克隆机制，FMTT 成功将理论采样框架泛化为一种高强度测试时搜索算法，能够突破基础扩散模型的分布先验，探索原本无法生成的长尾甚至分布外（OOD）图像特征。

一个完整示例¶

以生成显示精确时刻“4点45分”的挂钟图像为例： 1. 初始化：在 $t=0$ 采样 $N=32$ 个高斯白噪声粒子，并进行 $C=2$ 倍克隆扩展为 64 个粒子。 2. 前瞻与梯度反传：在时间步 $t_k$，流映射网络以 1 步前向直接输出每个粒子对应的终点预测图像 $X_{t_k, 1}(x_{ij}^k)$；将前瞻图像输入预训练视觉语言模型（VLM）评估由 3 个原子问题构成的组合奖励（“时针是否指向 4 与 5 之间”、“分针是否指向 9”、“秒针是否指向 12”），并反向传播获得梯度 $\nabla r_{t_k}$。 3. 粒子更新与演化：粒子依据加入 $\nabla r_{t_k}$ 的 SDE 更新一步，沿奖励梯度方向朝“4:45 表盘”流形弯曲。 4. 测试时筛选与重平衡：在 $k = 100$ 步时（轨迹中间），根据当前各粒子的前瞻奖励 $r_{t_k}$ 进行贪婪筛选，仅保留表现最好的 top-32 个粒子并重新克隆，淘汰表盘指针错位或模糊的轨迹分支。最终在 $t=1$ 成功生成精确对齐复杂视觉约束的高清图像。

损失函数 / 训练策略¶

FMTT 本身为纯推断阶段（training-free）的测试时自适应算法，底座模型采用基于 Align Your Flow 蒸馏训练的 FLUX.1-dev 4-step 流映射模型，其流映射权重在推断时保持冻结。在测试时粒子演化中，Euler-Maruyama 离散化步数设为 $K$（MNIST 采样 $K=200$，FLUX 图像生成 $K=4$），粒子数 $N \in [8, 256]$，扩散系数设为 $\epsilon_t = 1-t$。VLM 奖励采用 Skywork-VL Reward 或 Qwen2.5-VL，通过二分类 logits 差值的 sigmoid 构造标量目标。

实验关键数据¶

主实验¶

在文生图对象对齐基准 GenEval（550 个复杂提示词，评估物体计数、属性绑定、空间关系等）上的主实验结果如下表所示。FMTT 在综合性能与属性绑定等多维度上显著优于无梯度搜索和基于初始潜码优化的 ReNO 方法。

方法 (Method)	平均分 (Mean ↑)	单物体 (Single Obj.)	双物体 (Two Obj.)	计数 (Counting)	颜色 (Colors)	位置 (Position)	属性绑定 (Attr. Binding)	NFE (函数评估数)
FLUX.1 [dev] (Base)	0.65	0.99	0.78	0.70	0.78	0.18	0.45	180
Flow Map (4-step)	0.62	0.99	0.72	0.63	0.80	0.19	0.39	16
FLUX.1 [dev] + Best-of-NN	0.75	0.99	0.94	0.83	0.86	0.26	0.57	1440
Flow Map + Best-of-NN	0.73	1.00	0.88	0.82	0.85	0.25	0.59	128
Flow Map + Multi-best-of-NN	0.76	1.00	0.95	0.84	0.85	0.26	0.69	1280
Flow Map + Beam Search	0.75	1.00	0.92	0.86	0.85	0.29	0.58	1200
Flow Map + ReNO	0.71	0.98	0.89	0.79	0.89	0.20	0.57	1280
FMTT (本文完整方法)	0.79	1.00	0.97	0.90	0.91	0.30	0.64	1400
FMTT (1步去噪器前瞻)	0.75	0.99	0.90	0.87	0.87	0.26	0.59	350
FMTT (4步扩散采样前瞻)	0.75	0.99	0.93	0.86	0.89	0.27	0.57	1400

消融实验¶

在 UniGenBench++ 统一语义评估基准（600 个多维度提示词，使用 Qwen2.5-VL-72B 评判）上，对比不同前瞻机制与搜索策略随计算预算（粒子数 $N$ 及 NFE）的测试时扩展性能：

方法配置 (Config)	粒子数 $N$	综合均分 (Mean)	实体布局 (Entity Layout)	逻辑推理 (Logical Reas.)	文本渲染 (Text Gen.)	NFE
FLUX.1 [dev] (Baseline)	-	61.78 ± 3.58	67.16 ± 4.32	24.08 ± 3.63	30.17 ± 5.77	180
Best-of-NN	N=16	67.57 ± 1.53	72.95 ± 2.26	31.19 ± 2.34	36.49 ± 4.01	720
Best-of-NN	N=64	70.93 ± 1.38	77.99 ± 1.71	38.30 ± 4.12	42.82 ± 3.29	2880
Best-of-NN	N=256	73.75 ± 1.06	80.97 ± 1.54	45.18 ± 3.14	44.83 ± 3.54	11520
FMTT (1-step denoiser)	N=64	70.89 ± 1.32	78.36 ± 3.03	37.61 ± 4.05	43.97 ± 2.21	2880
FMTT (1-step denoiser)	N=256	72.90 ± 0.89	81.53 ± 1.70	42.89 ± 4.12	39.94 ± 1.88	11520
FMTT (1-step flow map)	N=64	71.13 ± 0.62	77.05 ± 2.66	37.39 ± 3.39	45.11 ± 3.39	960
FMTT (1-step flow map)	N=256	74.86 ± 0.50	81.72 ± 2.14	46.33 ± 1.38	46.55 ± 4.10	3840
FMTT (2-step flow map)	N=256	75.55 ± 0.53	81.53 ± 1.10	46.79 ± 1.45	47.70 ± 3.09	6720
FMTT (4-step flow map)	N=256	75.14 ± 0.80	80.41 ± 1.10	45.87 ± 2.25	50.86 ± 2.62	12480

关键发现¶

去噪器前瞻的失效根源：1步去噪器前瞻在 UniGenBench++ 扩展实验中甚至无法击败无梯度的 Best-of-NN 基线。其本质原因在于扩散初期噪声极高，去噪器输出严重失真，其产生的梯度未能提供有效指引，反而白白耗费了反向传播计算算力。
流映射的高效扩展曲线：基于 1 步流映射前瞻的 FMTT 在仅消耗 3840 NFE（$N=256$）时便达到了 74.86% 的均分，超越了消耗 11520 NFE 的 Best-of-NN（73.75%），以不到三分之一的计算成本实现了更高的精度。
热力学长度理论验证：在 MNIST 严格采样验证中，基于流映射的 FMTT（$\chi_t = 0$）的总偏差与热力学长度在所有方法中最低，且平均奖励与类别熵严格落在理论真值置信区间内，证明了理论无偏修正的正确性。

亮点与洞察¶

流映射与欧拉恒等式的优雅对偶：将微分流方程与积分流映射结合，利用 $\partial_t X_{t,1} + \nabla X_{t,1} b_t = 0$ 这一欧拉不变性，将极其繁琐的非平衡态重要性权重积分奇迹般化简为终点奖励的原函数积分，实现了理论美感与工程极简性的统一。
首次验证 VLM 作为测试时通用奖励函数：将多模态大语言模型（如 Skywork-VL、Qwen2.5-VL）的二分类 logits 差值直接转化为生成轨迹的连续梯度指引，使得图像空间生成能够直接遵循“自然语言指令与视觉提问”驱动的测试时优化。
突破模型内生偏置的 OOD 探索：在旋转对称、遮罩定位等几何约束任务中，基座模型因先验偏差无法生成完全对称的目标，而 FMTT 的动力学梯度倾斜能够强行引导粒子跳出预训练支持域，成功生成长尾且满足严格几何规律的图像。

局限与展望¶

VLM 奖励黑客（Reward Hacking）风险：作者坦承，将 VLM 作为连续优化目标时，粒子搜索容易利用模型漏洞“作弊”（例如要求生成 4:45 表盘时，模型可能会在表盘上手绘出文本数字“4:45”而非正确画出指针）。缓解方案需将复杂逻辑拆解为多道严格的细粒度子问题求和。
反向传播的内存与延迟开销：通过复杂的 VLM 或多步流映射反向传播需要维护计算图，显存开销大于无梯度的拒绝采样，目前在超大规模批量生成时对硬件要求较高。
未来方向：探索无梯度流映射引导或一阶轻量化近似，以及将 FMTT 推广至 3D 几何生成、分子动力学模拟与视频长程规划等连续时间物理建模任务。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将流映射全域解算子引入扩散模型测试时扩展，导出极简形式的重要性权重理论。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖理论玩具模型（MNIST 严格采样）、标准文生图基准（GenEval）、语义全维度基准（UniGenBench++）及复杂几何约束与 VLM 问答。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽，理论直觉与实验分析阐述清晰。
价值: ⭐⭐⭐⭐⭐ 为生成式模型的 Test-time Scaling 提供了坚实的理论基石与前沿实用方案。