引言
近年来，大型语言模型(LLMs)在推理过程中的扩展性得到了充分的验证。诸如o1、o3、DeepSeek R1、QwQ和Step Reasoner mini等模型都表明，在推理过程中增加计算量可以显著提高性能。那么，这种原理是否也适用于扩散模型呢？纽约大学谢赛宁领导的团队对此进行了深入研究。通过使用通用搜索框架进行的系统调查，他们发现推理时的扩展对于扩散模型确实有效。在推理过程中增加计算量可以显著提高生成样本的质量。此外，图像的复杂性使得框架内的组件可以进行不同的组合，以适应各种应用场景。
主要发现
- 推理时扩展对扩散模型有效： 在推理过程中分配更多的计算资源可以产生更高质量的样本。
- 组件组合的灵活性： 该框架允许不同的组件配置，以适应各种应用。
- 超越去噪步骤： 研究表明，在采样过程中搜索更好的噪声是扩展NFE的另一个维度，而不仅仅是增加去噪步骤。
- 两个设计轴： 该框架侧重于两个关键的设计轴：
- 验证器(Verifiers)： 在搜索过程中提供反馈。
- 算法(Algorithms)： 寻找更好的噪声候选。
研究方法
该团队探索了三种不同的验证器场景，模拟了各种用例：
- 可以使用关于最终评估的特权信息的情况。
- 可以使用条件信息来指导生成的情况。
- 没有额外信息可用的情况。
对于算法，他们研究了：
- 随机搜索(Random Search)： 从一组固定的候选项中选择最佳的。
- 零阶搜索(Zero-Order Search)： 使用验证器反馈迭代改进噪声候选项。
- 路径搜索(Path Search)： 使用验证器反馈迭代改进扩散采样轨迹。
该研究最初在一个相对简单的ImageNet类条件生成设置中探索了这些设计。随后，他们将这些设计应用于更大规模的文本条件生成，并评估了他们提出的框架。
推理时间扩展
该论文提出了一个在扩散模型中扩展推理时间的框架，将挑战定义为寻找最佳采样噪声。该过程涉及两个核心组件：
- 验证器(Verifiers)： 这些是预训练模型，用于评估生成样本的质量。它们接收生成的样本以及可选的条件，并输出一个标量分数。
- 算法(Algorithms)： 这些算法使用验证器分数来寻找更好的候选样本。该函数接收一个验证器(V)、一个预训练的扩散模型(_θ)以及一组生成的样本和条件，并输出最佳的初始噪声。
总推理预算通过函数评估总数(NFE)来衡量，包括去噪步骤和搜索成本。
搜索验证器
研究人员首先使用Oracle验证器，该验证器具有关于所选样本最终评估的完整信息。对于ImageNet，这包括诸如FID和IS之类的指标。然后，他们探索了更易于访问的预训练模型作为监督验证器，例如CLIP和DINO。这些模型用于对样本进行分类，选择与类标签对应的最高logit的样本。
然而，他们观察到，这些逐点操作的分类器仅部分地与FID分数的目的相符。这导致随着计算量的增加，样本方差减少和模式崩溃。这种被称为“验证器攻击”的现象因随机搜索算法的无约束搜索空间而加速。有趣的是，研究发现验证器不一定需要条件信息来有效地指导搜索。他们观察到DINO/CLIP分类器的logits与低噪声水平下的x预测和最终干净样本之间的特征空间余弦相似度之间存在很强的相关性。这导致了自监督验证器的使用，这些验证器不需要额外的条件信息，并且仍然表现出有效的扩展行为。
搜索算法
为了缓解验证器攻击，研究人员探索了更精细的搜索算法，这些算法可以逐步优化候选样本。这包括一种零阶搜索方法：
- 从随机高斯噪声作为支点开始。
- 在支点的邻域中找到N个候选。
- 运行通过ODE求解器的候选以获得样本和验证器分数。
- 使用最佳候选更新支点，并重复步骤1-3。
他们还研究了一种路径搜索算法，该算法探索了沿采样轨迹搜索的可能性：
- 采样N个初始噪声样本，并将ODE求解器运行到噪声水平σ。
- 向每个样本添加噪声，并模拟前向噪声过程。
- 在每个噪声样本上运行ODE求解器，并根据验证器分数保留前N个候选，重复直到ODE求解器达到σ=0。
- 随机搜索剩余的N个样本，并保留最佳样本。
与随机搜索相比，零阶搜索和路径搜索算法都保持了很强的局部性。
文本到图像场景中的扩展
该团队研究了搜索框架在更大规模的文本到图像任务中的扩展能力。他们使用DrawBench和T2I-CompBench数据集进行评估，并以FLUX.1-dev模型作为骨干。他们还扩展了监督验证器的选择，包括Aesthetic Score Predictor、CLIPScore和ImageReward。此外，他们还通过组合这三个验证器创建了一个验证器集成。
分析：验证器-任务对齐
该研究比较了不同数据集上各种验证器-算法组合的结果。在DrawBench上，他们发现使用所有验证器通常可以提高样本质量。但是，他们观察到，单独使用Aesthetic和CLIP验证器可能会导致过度拟合其偏差，从而对彼此产生负面影响。这源于它们的评估重点不匹配：Aesthetic Score侧重于视觉质量，通常偏爱高度风格化的图像，而CLIP优先考虑视觉-文本对齐，有时会牺牲视觉质量。他们指出，某些验证器更适合特定任务，并且验证器的有效性取决于其与任务要求的对齐。
算法性能
三种搜索算法（随机、零阶和路径）都有效地提高了DrawBench上的采样质量。但是，由于其他两种方法的局部性，随机搜索在某些方面表现更好。随机搜索更快地收敛到验证器偏差，而其他两种算法则需要改进次优候选。
与微调的兼容性
该团队研究了他们的搜索方法与微调模型的兼容性。他们使用了DPO微调的Stable Diffusion XL模型，发现该搜索方法可以推广到不同的模型，并提高已经对齐的模型的性能。
不同推理计算维度的影响
该研究探讨了推理计算的不同方面如何影响结果：
- 搜索迭代次数： 增加迭代次数会使噪声更接近最优值。
- 每次搜索迭代的计算量： 调整每次迭代的去噪步骤数会显示不同的计算最优区域。
- 最终生成计算： 该团队使用了最终去噪步骤的最佳设置，以确保最高的最终样本质量。
计算投入的有效性
研究人员探讨了在较小的扩散模型上进行推理时扩展的有效性。他们发现，对于ImageNet，扩展较小的模型可能非常有效。在某些情况下，在较小的模型上进行搜索可以胜过没有搜索的较大模型。但是，有效性取决于较小模型的基线性能。在基于文本的设置中，PixArt-Σ仅使用一小部分计算量就优于FLUX-1.dev。这些结果表明，在训练过程中花费的大量计算资源可以通过在生成过程中少量计算来弥补，从而更有效地产生更高质量的样本。