跳转至

Inference-Time Scaling of Diffusion Models via Progressive Seed Pruning

会议: ECCV 2026
论文: ECCV 原文
代码: https://github.com/rogerioagjr/psp
领域: 图像生成
关键词: 扩散模型, 流匹配, 推理时扩展, 种子剪枝, 奖励引导生成

一句话总结

本文提出渐进式种子剪枝(PSP),通过在去噪早期前置探索大量候选噪声种子、依据中间去噪外推估计值与黑盒奖励进行分阶段剪枝,在固定计算预算下以确定性采样大幅提升扩散与流匹配模型的提示词对齐能力。

研究背景与动机

条件图像生成已全面步入以扩散模型与流匹配模型为主导的时代。借鉴大语言模型中通过测试时搜索、采样与验证实现推理时扩展(Inference-Time Scaling)的成功经验,视觉生成领域也迫切需要探索如何有效分配推理算力。然而,传统单纯增加去噪步数的做法边际效益极低;大量实证研究表明,初始随机噪声种子对生成画面的宏观布局与语义质量具有决定性影响。因而,基于黑盒奖励模型引导的种子搜索或重采样逐渐成为测试时扩展的主流途径。

现有方法在分配推理算力时存在根本性的机制与系统局限。经典的 Best-of-N(BoN)对所有候选种子做全流程去噪,大量算力被浪费在最终注定落败的劣质种子上;而近期基于重要性采样的粒子滤波(如 FK-Steering)或树搜索方法(如 DSearch),均强行假设整个推理阶段必须维持恒定的并行粒子数(即恒定显存占用)。为了生成多样化候选,它们极度依赖随机采样注入扰动,这在采用直线 ODE 解算器的现代流匹配模型(如 SD 3.5)中会严重破坏生成质量。更关键的是,许多生成任务(如提示词对齐)的核心宏观语义在去噪早期就已基本固化,恒定显存假设硬生生切断了在早期广域探索更多种子的可能。

本文的切入点是打破这一不必要的恒定显存约束:现代云端与多 GPU 推理集群天然支持弹性内存与动态批处理。去噪模型在每一步预测中本身就隐含了当前时刻对清晰样本的单步外推估计,无需额外的生成网络前向开销即可提前获取信息量充足的中间奖励评分。核心 idea:提出渐进式种子剪枝(PSP),在去噪初期前置算力并发探索超大初始种子池,随后在预设时间节点利用单步去噪外推的中间奖励评分激进剔除低潜力候选,将剩余算力集中于少数优质轨迹,在固定总计算步数下实现确定性且高效的测试时扩展。

方法详解

整体框架

PSP 将测试时算力前置到去噪初期的种子探索阶段。整体流程包含四个阶段:超大初始种子并发、零额外网络开销的中间清晰图像外推、基于黑盒奖励的 Top-K 剪枝、以及对收敛候选集的最终精细去噪。在给定总模型评估步数(有效算力倍率 \(\bar{N} = C/T\))的前提下,PSP 在去噪关键节点按预设离散调度阶段性缩小粒子集合,既保证了全流程计算预算恒定,又契合多 GPU 分布式推理的系统调度。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入条件 c 与超大初始种子池<br/>采样 k_T 个高斯噪声种子"] --> B["前置探索阶段<br/>并发推演多条去噪轨迹"]
    B --> C["单步无开销清晰图像外推<br/>利用网络内在预测计算 x_0 估计值"]
    C --> D["黑盒奖励评估与阶段剪枝<br/>计算奖励评分并执行 Top-K 截断"]
    D --> E["幸存子集继续去噪<br/>算力集中于收敛候选"]
    E --> F["输出最终高质量生成图像"]

关键设计

1. 算力前置探索:以非恒定显存换取超大初始种子池覆盖 以往重要性采样与树搜索方法将整个推理过程的活跃粒子数固定为常量,导致初始探索的种子数严重受限。PSP 针对图像宏观构图与语义对齐在去噪早期即已定型的物理特性,打破了恒定显存假设。在起始阶段构建规模达到 \(k_T = 2\bar{N}\) 的种子池,利用去噪前期的快速演化覆盖更广阔的模式空间。这种前置算力分配策略使模型以同等计算预算接触到两倍于 BoN 的种子多样性,显著降低了因起始噪声不佳而陷入局部低质解的概率。

2. 单步无开销清晰图像外推:利用生成网络输出构建中间评估代理 在去噪中间状态评估图像质量的最大难点在于中间隐变量充满噪声,若调用多步采样解码将引入灾难性的额外算力。PSP 巧妙利用了扩散模型与流匹配模型的内在数学特性:在扩散模型中,网络预测的噪声 \(\epsilon_\theta\) 可直接导出对清晰图像的条件期望估计 \(\hat{x}_0(x_t, t, c) = \mathbb{E}[x_0 \mid x_t, t, c]\);在流匹配模型中,当前速度场预测 \(v_\theta(x_t, t, c)\) 亦可通过一阶外推直接得到无噪声端点估计 \(\hat{x}_0(x_t, t, c) = x_t - \epsilon_t v_\theta(x_t, t, c)\)。中间外推图像经轻量级 VAE 解码后,即可直接交由黑盒奖励模型计算中间评分 \(s_t = r(\hat{x}_0, c)\),整个过程不需要引入任何额外的生成网络前向传播。

3. 预设离散几何剪枝调度:兼顾排序稳定性与工业级系统部署 连续动态调配粒子数在工程部署中极难做批处理优化。PSP 采用预设的几何减半(Successive Halving)修剪调度。以默认的四倍算力(\(\bar{N}=4\))为例,系统起始启动 \(8\) 个候选种子;在去噪进度达到 \(25\%\) 时计算各轨迹的 \(\hat{x}_0\) 奖励并截断保留 Top-4;在去噪进度达到 \(50\%\) 时再次截断保留 Top-2;最后将剩余 \(50\%\) 的去噪步数完全分配给这 \(2\) 条最具潜力的轨迹,最终挑选奖励最高者输出。固定时间点与固定批大小使得显存峰值和推理耗时完全可预测,极易与现代弹性推理框架契合。

4. 纯确定性解算流:消除随机扰动并解锁离线调度搜索 不同于 FK-Steering 等重要性采样方法必须依赖随机采样器产生多样化子代粒子,PSP 坚持采用完全确定性求解器(扩散模型采用 \(\eta=0\) 的 DDIM,流匹配模型采用 Euler Discrete)。这彻底规避了随机噪声对流匹配整流轨迹的破坏。更为关键的是,确定性使得所有候选轨迹与中间奖励在初始种子确定后完全可复现。算法设计者仅需在小规模 Benchmark 上离线运行一次种子轨迹生成并缓存中间奖励,即可零生成成本快速网格搜索针对特定下游任务的最优剪枝时间点与保留比例。

损失函数 / 训练策略

PSP 属于完全无需训练(Training-free)与无需梯度(Gradient-free)的测试时扩展方案,不修改底座扩散模型或流匹配模型的任何权重。引导评估采用轻量级黑盒评分模型 ImageReward,仅在预设的少数剪枝节点对单步外推图进行评分。中间批次通过 VAE 切片(slicing)技术进行串行或小批次解码,在几乎不增加运行耗时的前提下有效抑制了瞬时显存开销。

实验关键数据

主实验

在 GenEval 提示词基准上,固定相同总计算量(\(\bar{N}=4\)),评估不同测试时扩展方法在 SD v1.5、SDXL(扩散模型)与 SD 3.5(流匹配模型)上的引导奖励(ImageReward)、人类偏好分(HPSv2)、自动化提示词对齐率(GenEval)以及众包人工评估对齐率(Human)。

模型底座 采样方法 步数 \(T\) 算力倍率 \(\bar{N}\) ImageReward \(\uparrow\) HPSv2 \(\uparrow\) GenEval \(\uparrow\) 人工评估 \(\uparrow\)
SD v1.5 Standard 64 1 -0.159 0.257 0.434 0.431
SD v1.5 Best-of-N 64 4 0.655 0.273 0.542 0.594
SD v1.5 FK-Steering 64 4 0.640 0.261 0.531 0.569
SD v1.5 DSearch 64 4 0.783 0.275 0.506 0.514
SD v1.5 BFS 64 4 0.820 0.263 0.564 -
SD v1.5 PSP(本文) 64 4 0.827 0.278 0.574 0.624
SDXL Standard 64 1 0.431 0.275 0.529 0.531
SDXL Best-of-N 64 4 1.098 0.290 0.629 0.682
SDXL FK-Steering 64 4 1.189 0.284 0.627 0.676
SDXL DSearch 64 4 1.186 0.301 0.589 0.649
SDXL BFS 64 4 1.247 0.285 0.636 -
SDXL PSP(本文) 64 4 1.224 0.294 0.645 0.713
SD 3.5 Standard 32 1 1.045 0.297 0.713 0.787
SD 3.5 Best-of-N 32 4 1.336 0.304 0.747 0.831
SD 3.5 FK-Steering 32 4 1.294 0.284 0.742 0.837
SD 3.5 DSearch 32 4 1.144 0.297 0.704 0.824
SD 3.5 BFS 32 4 1.343 0.285 0.747 -
SD 3.5 PSP(本文) 32 4 1.380 0.306 0.747 0.841

消融实验

在单一 NVIDIA H200 GPU 上对计算开销、显存开销进行定量基准测试,并对比与 DPO 偏好微调模型的正交互补性。

1. 推理耗时与峰值显存对比(H200 测试)

模型 方法配置 算力倍率 \(\bar{N}\) 单图推理耗时 (s) 峰值显存 (GiB) 说明
SD v1.5 BoN 4 2.65 4.29 等算力基线
SD v1.5 PSP 4 2.99 5.30 探索 8 种子,增加耗时仅 12.8%
SD v1.5 BoN 8 4.70 5.31 全探索基线,PSP 耗时大幅领先
SDXL BoN 4 12.48 8.28 等算力基线
SDXL PSP 4 13.74 13.27 探索 8 种子,增加耗时仅 10.1%
SDXL BoN 8 23.90 10.79 PSP 比全量 BoN 提速约 1.74×
SD 3.5 BoN 4 35.97 29.85 等算力基线
SD 3.5 PSP 4 37.59 34.85 探索 8 种子,增加耗时仅 4.5%
SD 3.5 BoN 8 71.71 33.86 PSP 逼近 2× 速度优势且对齐分更高

2. 与 DPO 偏好微调模型的互补性实验

基础模型 是否经过 DPO 微调 采样策略 步数 \(T\) 算力倍率 \(\bar{N}\) ImageReward \(\uparrow\) GenEval \(\uparrow\)
SD v1.5 Standard 64 1 -0.159 0.434
SD v1.5 Standard 64 1 -0.022 0.454
SD v1.5 PSP 64 4 0.827 0.574
SD v1.5 BoN 64 4 0.751 0.562
SD v1.5 PSP 64 4 0.907 0.593
SDXL Standard 64 1 0.431 0.529
SDXL Standard 64 1 0.894 0.581
SDXL PSP 64 4 1.224 0.645
SDXL BoN 64 4 1.300 0.657
SDXL PSP 64 4 1.365 0.667

关键发现

  • 探索范围完胜中途随机重采样:在所有三款主流底座上,PSP 在 GenEval 和人类真实盲测评价中均取得了最优表现。FK-Steering 和 DSearch 虽然在某些设定下优化了中间奖励,但频繁引入随机扰动容易引发奖励作弊(Reward Hacking),生成的图像破坏了整体结构;而 PSP 将算力全部用于扩大确定性种子探索面,泛化对齐表现最稳健。
  • 算力扩展与失真衰减规律:当算力倍率从 \(\bar{N}=1\) 扩展至 \(\bar{N}=16\) 时,PSP 的奖励收益呈现严格的单调递增,且相对全量 \(2\bar{N}\) 真实最佳种子的遗憾值(Regret)随算力增加显著收敛。这表明在更大算力下,早期中间评分的鲁棒性足以安全滤除绝大部分失败轨迹。
  • 与微调正交互补:在未微调的基础模型上运行 PSP(\(\bar{N}=4\)),其文本对齐指标显著超越了经过 DPO 训练的模型在标准采样下的表现;而将 PSP 应用于 DPO 模型后,各项指标进一步跃升,证明测试时扩展与训练时对齐具有强互补性。
  • 外延至 Prompt Rewrite 离散空间:将候选池从噪声种子替换为由大型语言模型(ChatGPT)生成的多样化提示词重写文本,PSP 同样能在固定预算下显著优于 BoN,展现出通用的测试时候选空间修剪能力。

亮点与洞察

  • 敏锐指出现有测试时扩展对“恒定显存”假定的不必要依赖,将算力前置至去噪早期的种子探索,以轻微且可控的峰值显存代价换取了翻倍的种子候选空间。
  • 单步外推估计 \(\hat{x}_0\) 零计算代价复用去噪网络当期输出,打破了中间状态评估必须额外前向传播的固有模式。
  • 完全确定性的解算机制不仅与流匹配模型的直轨迹天然契合,更使得离线仿真成为可能,可直接读取中间特征缓存秒级搜索特定任务的最佳剪枝时机。

局限与展望

  • 单一标量奖励局限:PSP 依赖单标量评分进行排序筛选,在需要细粒度空间几何约束(如 ControlNet 级空间边缘匹配)或复杂多目标权衡时难以表达。
  • 对晚期微观属性的增益相对有限:种子的决定性主要作用于全局布局与高层语义;对于画风纹理、微观细节等在去噪末期才沉淀的属性,前置剪枝带来的边际增益收窄。
  • 后续优化方向:探索多奖励协同的帕累托前沿动态剪枝,以及自适应置信度驱动的弹性剪枝步长。

相关工作与启发

  • vs Best-of-N (BoN):BoN 独立走完所有种子的全程去噪,产生大量无效开销;PSP 在去噪中途按几何减半策略淘汰低分种子,在相同总计算评估次数下将有效探索的种子总数翻倍。
  • vs FK-Steering (Feynman-Kac 重要性采样):FK-Steering 全程维持固定粒子数并引入随机重采样,但在流匹配模型中随机噪声会严重劣化直线积分精度;PSP 坚持确定性轨迹,通过前期多开、中期剪除的方式规避了重采样噪声。
  • vs DSearch / BFS (树搜索方法):树搜索方法在中后期构建分叉树,不仅对齐收益易陷入局部过拟合,且动态分支严重破坏了 GPU 批量调度的稳定性;PSP 的剪枝计划完全预定,系统吞吐与部署确定性极高。

评分

  • 新颖性: ⭐⭐⭐⭐☆ 突破恒定显存教条,提出面向去噪特性的前置种子剪枝范式。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖扩散与流匹配主流底座,结合了自动化检测与超两万次人工双盲评估。
  • 写作质量: ⭐⭐⭐⭐⭐ 动机阐述清晰透彻,系统级权衡与实验剖析扎实全面。
  • 价值: ⭐⭐⭐⭐⭐ 即插即用、无需训练且确定性友好,是文生图生产环境极具落地价值的推理扩展方案。