扩散模型也能推理时Scaling，谢赛宁团队重磅研究可能带来文生图新范式

来源 | 机器之心

对于 LLM，推理时 scaling 是有效的！这一点已经被近期的许多推理大模型证明：o1、o3、DeepSeek R1、QwQ、Step Reasoner mini…… 

但这个说法也适用于扩散模型吗？

近日，纽约大学谢赛宁领导的一个团队对这一方向进行了探索。具体来说，他们借助通用搜索框架系统性地探索了扩散模型的推理时 scaling 情况。他们发现，推理时 scaling 对扩散模型是有效的 —— 增加推理时间计算可以显著提高扩散模型生成的样本的质量，并且由于图像的复杂性，可以针对不同的应用场景，对框架中的组件进行不同形式的组合。

这篇论文的署名部分还透露出了一个有趣的小细节：谢赛宁的所属机构标记成了谷歌并且通讯作者邮箱也使用了其谷歌域名的邮箱。但我们目前还并不清楚这是否意味着这位著名 AI 研究者已经正式加入谷歌，还是在谷歌兼职（谢赛宁此前已经是谷歌研究院访问学者）。

• 论文标题：Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps

• 论文链接：https://arxiv.org/pdf/2501.09732

有趣的是，就在前几天，来自纽约大学和哥伦比亚大学的研究者也发布了一篇类似方向的论文。看来，scaling 推理将是改进扩散模型的一个非常有前途的方向。

论文概览

OpenAI o1 等模型的问世已经证明，在推理阶段增加计算量可以让 LLM 的性能进一步提升。但对于扩散模型，如何有效 scaling 推理时间计算以进一步提升性能还不够明确。

扩散模型经过训练以去除数据中的噪声，是一类在连续数据领域占主导地位的生成模型，如图像、音频和视频。为了生成单个样本，它们的生成过程通常从纯噪声开始，需要经过训练模型的多次前向传递来去噪并获得干净的数据。这些前向传递因此被称为去噪步骤。由于去噪步骤的数量可以调整，以在样本质量和计算成本之间进行权衡，扩散模型的生成过程自然提供了在推理时分配计算预算的灵活性。

对于生成模型，这种计算预算通常用函数评估次数（NFE）来衡量，以确保与其他使用迭代采样过程但没有去噪能力的模型系列进行合理比较。

经验观察表明，仅通过将计算投入到去噪步骤中，性能提升在达到某个 NFE 后往往会趋于平稳，限制了推理过程中增加计算的收益。因此，以前关于扩散模型的工作长期以来一直专注于在保持高性能的同时尽可能减少推理时的 NFE 以提高效率。

但是，谷歌 DeepMind 和谢赛宁等人新论文的研究方向与此相反。

与 LLM 相比，扩散模型处理的是作为初始样本注入的噪声或在采样过程中注入的噪声的显式随机性。已有研究表明这些噪声并非等价，即某些噪声会带来更好的生成结果。这一观察为 scaling NFE 提供了除增加去噪步骤之外的另一个维度 —— 在采样中搜索更好的噪声。

新论文研究了在推理过程中通过搜索有效利用计算资源的方法，而不是仅将 NFE 分配给去噪步骤（这往往会很快导致性能不再增长），从而在推理时改进扩散模型的性能和可扩展性（图 1）。

作者主要考虑搜索框架中的两个设计轴：用于在搜索中提供反馈的验证器，以及用于寻找更好噪声候选项的算法，遵循 LLM 中使用的术语。

对于验证器，他们考虑了三种不同的设置，这些设置旨在模拟三种不同的使用场景：

1. 掌握最终评估如何进行的特权信息的场景；

2. 掌握用于指导生成的条件信息的场景；

3. 没有额外信息可用的场景。

对于算法，他们研究了：

1. 随机搜索，它只是从固定的候选集中选择最佳项；

2. 零阶搜索，它利用验证器反馈来迭代改进噪声候选项；

3. 路径搜索，它利用验证器反馈来迭代改进扩散采样轨迹。

作者首先在 ImageNet 类别条件生成这个相对简单的设置中探讨这些设计选择并展示它们的有效性，为新框架提供一个具体实例。然后他们将这些设计选择应用到更大规模的文本条件生成设置中，并评估他们提出的框架。

由于图像的复杂性质和文本条件包含的丰富信息，需要对生成质量进行更全面的评估。因此，他们采用多个验证器来进行搜索中的推理时计算 scaling。这也使他们能够探究每个验证器具有的「偏见」，以及它们与生成任务的匹配程度。为了缓解对单个验证器的过拟合，他们还尝试了验证器集成，并展示了它在不同基准测试中的良好泛化能力。

新论文的贡献总结如下：

• 提出了一个用于扩散模型推理时 scaling 的基础框架。论文表明，通过搜索去 scaling NFE 可以在各种生成任务和模型规模上带来实质性改进，超越了仅增加去噪步骤的方法。此外，作者对推理时计算预算如何影响 scaling 性能进行了全面的实证分析。

• 在提出的搜索框架中确定了两个关键设计轴：提供反馈的验证器和寻找更好噪声候选项的算法。作者研究了不同验证器 - 算法组合在各种任务中的表现，他们的发现表明没有一种配置是普遍最优的；每个任务反而需要一个独特的搜索设置才能实现最佳的 scaling 性能。

• 对验证器与不同生成任务之间的匹配度进行了广泛分析。论文的结果揭示了不同验证器中嵌入的偏见，以及在每个不同的视觉生成任务中需要专门设计验证器的必要性。

如何在推理时间 scale

本文提出了一个框架，用于扩散模型的推理时间 scaling 。本文将这一挑战表述为对采样噪声的搜索问题；具体来说，本文需要知道哪些采样噪声是好的，以及如何搜索这些噪声？

在高层次上，本文考虑了两个设计轴：

一是验证器，用于评估候选样本的质量。这些验证器通常是能够提供反馈的预训练模型；验证器函数形式如下：

这些函数接收生成的样本以及可选的对应条件，并输出一个标量值作为每个生成样本的评分。

二是算法，用于根据验证器的评分找到更好的候选样本。函数形式如下：  

它接收验证器 V、预训练扩散模型  𝐷_θ ，以及 N 对生成的样本及其相应的条件，并根据噪声和样本之间的确定性映射输出最佳的初始噪声。在整个搜索过程中，f 通常会多次通过 𝐷_θ 进行前向传递。本文将这些额外的前向传递称为搜索成本，用 NFE 来衡量。

举例来说，本文使用在 ImageNet 上预训练的 SiT-XL 模型，分辨率为 256×256，并使用二阶 Heun 采样器进行采样，即除了初始噪声外，采样过程中没有其他随机性来源。此外，本文还通过去噪步骤和搜索成本的总 NFE 来衡量推理计算预算。

本文从最简单的搜索算法开始，即随机采样高斯噪声，使用 ODE 生成样本，并选择与最佳验证器评分对应的样本（图 2）。本文将这种算法称为随机搜索，它本质上是一种在所有噪声候选上应用一次的 Best-of-N 策略。

对于验证器，本文从最佳验证器开始，即 Oracle 验证器，他们假设其拥有关于选定样本最终评估的全部特权信息。对于 ImageNet，由于 FID 和 IS 通常被用作评估指标，因而直接将它们作为 Oracle 验证器。

搜索验证器

虽然通过搜索 scalingNFE 在使用 Oracle 验证器时表现出令人印象深刻的性能（如图 3 所示），但关键问题在于这种有效性是否可以推广到使用更易获取的预训练模型的监督验证器上。

为了研究这一点，本文选择了两个模型：CLIP 和 DINO 。对于 CLIP，本文遵循 Radford 等人的方法，使用通过提示工程生成的嵌入权重作为零样本分类器。对于 DINO，本文直接使用预训练的线性分类头。在搜索过程中，将样本输入分类器，并选择与生成中使用的类别标签对应的 logits 最高的样本。

如图 4 所示，与单纯通过增加去噪步骤 scaling NFE 相比（图 1），这种策略也有效提升了模型在 IS 上的性能。然而，本文注意到，由于这些分类器是逐点操作的，它们仅部分与 FID 评分的目标对齐。具体来说，它们生成的 logits 仅关注单个样本的质量，而没有考虑总体多样性，这导致样本方差显著减少，并最终表现为随着计算量增加而出现的模式崩溃。随机搜索算法也难辞其咎，由于其无约束的搜索空间，加速了搜索向验证器偏见的收敛。这种现象类似于强化学习中的奖励黑客行为，因此将其称为验证器黑客。

尽管条件信息在现实世界的生成任务中至关重要，但本文发现验证器并不一定需要条件信息来有效指导搜索。如图 5 所示，他们发现 DINO/CLIP 分类器输出的 logits 与模型在低噪声水平（σ=0.4）下的 x 预测值与最终生成的干净样本（σ=0）之间的特征空间（分别由 DINO/CLIP 提取）余弦相似度之间存在强相关性。

因此，本文使用这种相似度评分作为分类 logits 的替代，并将这类验证器称为自监督验证器，因为它们不需要额外的条件信息。图 5 再次观察到了有效的 scaling 行为。

这一结果对于条件信息不可用或难以获取的用例来说（如医学图像生成任务）是令人鼓舞的。由于这些限制在现实场景中并不常见，本文将自监督验证器的进一步研究留待未来工作。

搜索算法

前文的探索主要集中于简单的随机搜索设置，他们发现这种方法可能导致验证器黑客行为。基于这一认知，作者开始研究更细致的搜索算法，这些算法利用验证器的反馈逐步优化候选样本，每次只进行小幅调整，从而减轻过拟合风险。

具体来说，本文考虑了一种零阶搜索方法：

1. 从一个随机高斯噪声 n 作为支点开始。

2. 在支点的邻域中找到 N  个候选项。

3. 通过 ODE 求解器运行候选项以获得样本及其对应的验证器得分。

4. 找到最佳候选项，将其更新为支点，并重复步骤 1-3。

与零阶优化类似，零阶搜索不涉及昂贵的梯度计算；相反，它通过在邻域内进行多次前向函数评估来近似梯度方向。

本文还注意到，由于许多验证器是可微的，理论上可以使用真实梯度进行一阶搜索，并且在实际中已有应用。然而，这需要在整个采样过程中进行反向传播，通常在时间和空间复杂度上都极为昂贵，尤其是在扩展大模型时。在实践中，本文发现一阶搜索在 ImageNet 上并没有表现出比零阶搜索更明显的优势，尽管它的成本更高。

扩散采样过程的迭代特性为设计局部搜索算法提供了其他可能性，沿着注入噪声的采样轨迹进行搜索是可行的。本文提出了路径搜索来探索其中一种可能性。具体来说，

1. 采样 N 个初始独立同分布噪声，并运行 ODE 求解器直到某个噪声水平 σ。噪声样本 x_σ 作为搜索起点。

2. 为每个噪声样本采样 M 个独立同分布噪声，并模拟从 σ 到 σ+Δf 的前向加噪过程，生成大小为 M 的 {x_(σ+Δf)}。

3. 在每个 x_(σ+Δf) 上运行 ODE 求解器到噪声水平 σ + Δf - Δb，并获得 x_(σ+Δf-Δb)。在这些样本上运行验证器并保留前 N 个候选样本。重复步骤 2-3，直到 ODE 求解器达到 σ=0。

4. 对剩余的 N  个样本通过随机搜索并保留最佳的一个。

为了确保迭代会终止，该团队严格要求 Δb > Δf。另外，由于验证器通常不适应有噪声输入，因此他们在步骤 3 中执行一个额外的去噪步骤，并使用干净的 x 预测与验证器交互。这里，主要的 scaling 轴是步骤 2 中添加的噪声数量 M，在实践中，他们研究了不同初始噪声数量 N 的 scaling 行为。

因此该算法被命名为 Paths-N。这两种算法见图 2，从中可以看到，与随机搜索相比，零阶搜索和路径搜索都保持了很强的局部性：前者在初始噪声的邻域中运作，后者则是在采样过程的中间步骤中搜索。

这些算法的性能见图 6。由于这两种算法的局部性质，它们都在一定程度上缓解了 FID 的多样性问题，同时保持了 Inception Score 的 scaling。对于零阶搜索，可以注意到增大 N 的效果很小，N = 4 似乎已经是对局部最优的良好估计。至于路径搜索，可以看到不同的 N 值会导致不同的 scaling 行为，小 N 值在小生成预算下计算效率高，而大 N 值在增加更多计算时更具优势。

在文生图场景中的推理时 scaling

在更大规模文生图任务中，该团队研究了搜索框架的推理时 scaling 能力，并研究了验证器与特定图像生成任务之间的对齐性能。

数据集方面，该团队采用了 DrawBench 和 T2I-CompBench 以实现更全面的评估。骨干模型则是 FLUX.1-dev。验证器方面，考虑到文生图任务本质上的复杂性质，该团队还扩展了监督式验证器的选择：Aesthetic Score Predictor、CLIPScore、ImageReward。其中 Aesthetic Score Predictor 可以预测人类对合成图像的视觉质量的评分，CLIPScore 则是通过 4 亿对人工标注的图像 - 文本数据来对齐视觉和文本特征，ImageReward 则更能体现更一般化的偏好。

此外，该团队还创新地将这三个验证器组合到一起，得到了第四个验证器，称为 Verifier Ensemble。可用于进一步扩展验证器在不同评估方面的能力。

更多实验设置详情请参看原论文。

分析结果：验证器 - 任务对齐

下面比较在不同数据集上验证器和算法组合的结果。

DrawBench。这是一个高度通用的数据集，包含来自不同类别的文本提示词。图 8 展示了在该数据集上的一些结果。

如图 8 所示，根据 LLM Grader 指标，可知使用所有验证器进行搜索通常都会提高样本质量，而具体的改进行为在不同设置下有所不同。这表明，可以根据不同的应用场景选择专门的搜索设置。

从图 8 的左两列中，可以看到使用 Aesthetic 和 CLIP 验证器进行搜索会过度拟合它们固有的偏见，对彼此产生负面影响。该团队推测这两个验证器在评估上存在重大不匹配：Aesthetic Score 仅关注视觉质量，经常偏向于与文本提示词偏离的高度风格化图像，而 CLIP 优先考虑视觉 - 文本对齐，牺牲了视觉质量。因此，如果在搜索过程中使用一个验证器（例如 Aesthetic Score），其偏见便会降低由另一个验证器（例如 CLIP）评估的指标。

另外，与 ImageNet 设置类似，该团队在这里也观察到了随着搜索预算增长而表现出的 scaling 行为。

通过观察验证器在 DrawBench 和 T2I-CompBench 上的行为，可以发现：验证器的有效性取决于其标准与任务特定需求的匹配程度，某些验证器比其他验证器更适合特定任务。

算法。表 2 展示了三种搜索算法在 DrawBench 上的表现。

可以看到，这三种方法全都能有效提高采样质量，其中随机搜索在某些方面优于其他两种方法。同样，该团队认为这种行为的原因是零阶搜索和路径搜索的局部性质（图 2）。由于这里展示的所有验证器和指标都是基于每个样本进行评估的，随机搜索将大大加速向验证器偏见的收敛，而其他两种算法需要对次优候选项进行改进。

搜索与微调是兼容的

搜索和微调的目标都是将最终样本与显式奖励模型或人类偏好对齐。前者将样本模式向特定验证器的偏见移动，后者是直接修改模型的分布以与奖励对齐。于是问题来了：在模型分布被修改后，是否仍然可以根据验证器移动样本模式？

在这里研究过的所有微调方法中，作为 RLHF 方法的一个更高效和更简单的替代方案，DiffusionDPO 已被广泛用于对齐大规模文生图模型。为了回答这个问题，该团队采用论文《Diffusion model alignment using direct preference optimization》中的 DPO 微调的 Stable Diffusion XL 模型，并在 DrawBench 数据集上进行搜索。由于该模型是在 Pick-a-Pic 数据集上微调的，于是他们也用 PickScore 评估器替换 ImageReward。结果见表 3。

可以看到，搜索方法可以泛化到不同的模型，并可以改进已对齐模型的性能。看起来，这个工具很有用，可以缓解微调模型与奖励模型不一致的情况，并提高它们在其他指标上的泛化能力。

投入推理计算上不同维度的不同效果

由于扩散模型本质上的迭代采样性质，因此通过搜索来 scaling NFE 的维度有很多。这篇论文也对此展开了一些讨论。

搜索迭代的次数。直观地说，增加搜索迭代次数可以使选定的噪声相对于验证器更接近最优集，从而可以显著提高性能。

每次搜索迭代的计算量。在每次搜索迭代中，去噪步骤数量都可以调整。这里将模型采用的去噪步骤次数记为 NFEs/iter。尽管仅增加去噪步骤时模型性能很快达到平稳（图 1），但该团队观察到在搜索过程中，调整 NFEs/iter 可以揭示不同的计算最优区域，如图 10 所示。在搜索过程中，较小的 NFEs/iter 能实现高效收敛，但最终性能较低。相反，较大的 NFEs/iter 会导致收敛较慢但产生更好的性能。

最终生成的计算量。尽管可以自由调整最终生成的去噪步骤，但该团队的做法是始终使用最优设置以获得最佳的最终样本质量。在 ImageNet 中，他们将去噪预算固定为 250 NFEs，在文本到图像设置中使用 30 步采样器，因为进一步 scaling 时的性能将很快趋平。

投资计算的有效性

该团队探索了在较小的扩散模型上进行推理时计算 scaling 的有效性，并强调了相对于它们不进行搜索的大型模型的性能效率。结果见图 11。

可以看到，在 ImageNet 上，对小型模型进行推理时计算 scaling 会非常有效。在固定的推理计算预算下，对小型模型进行搜索可以超过不进行搜索的大型模型。举个例子，在推理计算量有限时，SiT-L 相对于 SiT-XL 存在优势。但是，将 SiT-B 与其他两个模型进行比较表明，这种有效性取决于小型模型相对强大的基线性能。当小型模型的基线性能明显落后时，scaling 的收益是有限的，会导致结果次优。

在基于文本的设置中也出现这种现象，如表 4 所示。

可以看到，仅使用十分之一的计算量，PixArt-Σ 就超过了不进行搜索的 FLUX-1.dev。而如果使用大约两倍的计算量，PixArt-Σ 还能大幅超过不进行搜索的 FLUX-1.dev。这些结果具有重要的实践意义：在训练中投入的大量计算资源可以通过生成期间的一小部分计算来抵消，从而能更高效地获得更高质量的样本。