【AI论文】Diffusion模型的推理时缩放:超越降噪步骤的缩放策略
摘要:生成模型在各个领域产生了重大影响,这主要得益于它们能够在训练过程中通过增加数据、计算资源和模型规模来进行扩展,这一现象被称为扩展定律。最近的研究已经开始探索大型语言模型(LLMs)在推理时间的扩展行为,揭示了如何通过增加推理过程中的计算来进一步提升性能。与LLMs不同,Difussion(扩散)模型天生具有通过调整去噪步数来灵活调整推理时间计算的能力,尽管性能提升通常在几十个步数之后趋于平稳。在本文中,我们探索了扩散模型在推理时间的扩展行为,超越了仅增加去噪步数的限制,并研究了如何通过增加计算来进一步提升生成性能。具体来说,我们考虑了一个搜索问题,旨在识别扩散采样过程中更好的噪声。我们沿着两个轴构建了设计空间:用于提供反馈的验证器,以及用于寻找更好噪声候选的算法。
Huggingface链接:Paper page ,论文链接:2501.09732
1. 引言
生成模型在各个领域都取得了显著的影响,这主要得益于它们在训练期间通过增加数据、计算资源和模型规模进行扩展的能力,这一现象被称为扩展定律。最近的研究开始探索大型语言模型(LLMs)在推理时的扩展行为,揭示了当额外的计算资源在推理时可用时,模型性能如何进一步提升。与LLMs不同,扩散模型天生就具有通过调整降噪步骤的数量来灵活调整推理时计算的能力,尽管性能提升通常在几十个降噪步骤后就趋于平稳。本文探讨了扩散模型在推理时的扩展行为,不仅限于增加降噪步骤,还研究了如何通过增加计算来进一步提高生成性能。
2. 背景与动机
扩散模型概述:
扩散模型是一类生成模型,通过反向学习一个参考的“加噪”过程来工作。它们从一个纯噪声开始,通过多次前向传递训练好的模型来逐步去噪并获得干净的数据。这些前向传递因此被称为降噪步骤。由于降噪步骤的数量可以调整以权衡样本质量和计算成本,扩散模型的生成过程可以根据计算预算进行定制。
推理时扩展的挑战:
尽管增加降噪步骤可以提高性能,但这种提升在达到一定数量的步骤后会趋于平稳。本文的研究动机是探索推理时计算的有效利用方法,以超越仅通过增加降噪步骤来提高性能的限制。
3. 推理时缩放方法
本文提出了一种框架,将推理时缩放问题表述为一个在采样噪声上的搜索问题。具体来说,就是如何知道哪些采样噪声是好的,以及如何搜索这些噪声。框架主要包含两个设计维度:验证器和搜索算法。
验证器:
验证器用于评估候选样本的好坏。它们通常是一些预训练模型,能够提供反馈。具体来说,验证器是函数V:,其中输入是生成的样本(可选地包括相应的条件),输出是每个生成样本的分数。
- Oracle验证器:假设具有关于最终评估程序的全部特权信息,如FID和IS等评估指标。
- 监督验证器:使用预训练的模型,如CLIP和DINO,通过分类视角来评估样本。
- 自监督验证器:利用模型在低噪声水平下的特征预测与最终生成样本之间的相似度作为评估标准,不需要额外的条件信息。
搜索算法:
搜索算法用于找到更好的噪声候选。它们通常是函数 ,其中输入是验证器、预训练的扩散模型以及N对生成样本和相应条件,输出是根据噪声与样本之间的确定性映射找到的最佳初始噪声。
- 随机搜索:从固定数量的候选噪声中随机采样,选择验证器分数最高的样本。
- 零阶搜索:在初始噪声的邻域内采样候选噪声,通过验证器反馈迭代地细化噪声候选。
- 路径搜索:沿着采样轨迹上的噪声注入进行搜索,通过验证器反馈迭代地细化扩散采样轨迹。
4. 实验与结果
ImageNet上的类条件生成:
本文首先在ImageNet上的类条件生成任务中测试了所提出的框架。通过随机搜索、零阶搜索和路径搜索算法,以及不同的验证器组合,研究了推理时计算预算对生成性能的影响。实验结果表明,通过增加推理时计算,可以显著提高样本质量,并且不同的验证器和搜索算法组合在不同的任务中表现各异。
文本到图像的生成:
在更大规模的文本到图像生成任务中,本文进一步评估了所提出框架的推理时扩展能力。使用了DrawBench和T2I-CompBench两个数据集,以及FLUX.1-dev模型作为骨干模型。实验结果表明,通过搜索可以找到与不同生成任务更好对齐的验证器,从而进一步提高样本质量。特别是,对于需要广泛评估方面的生成任务,如DrawBench,使用包含更多细微评估方面的验证器组合(如Verifier Ensemble)能够显著提高样本质量。
5. 推理时计算投资的分析
本文还分析了推理时计算投资的多个维度,包括搜索迭代次数、每次搜索迭代的计算量以及最终生成的计算量。实验结果表明,通过调整这些维度,可以在不同的计算预算下实现性能的最优化。特别是,对于较小的扩散模型,在有限的推理时计算预算下,通过搜索可以显著提高性能,甚至超过没有搜索的较大模型。
6. 与微调方法的兼容性
本文还探讨了搜索方法与微调方法的兼容性。实验结果表明,搜索方法可以泛化到不同的模型,并可以进一步提高已经通过微调对齐的模型的性能。这对于缓解微调模型与奖励模型之间可能存在的不一致性问题,以及提高模型在其他指标上的泛化能力具有重要意义。
7. 相关工作
本文讨论了与推理时扩展、扩散模型微调以及扩散模型中样本选择与优化相关的工作。与现有工作相比,本文提出了一个更全面的框架来探索扩散模型在推理时的扩展行为,并强调了验证器选择和搜索算法设计的重要性。
8. 结论
本文提出了一种扩散模型在推理时的缩放框架,通过搜索找到更好的采样噪声来提高生成性能。实验结果表明,通过增加推理时计算,可以显著提高扩散模型的生成性能,并且不同的验证器和搜索算法组合在不同的任务中表现各异。此外,本文还分析了推理时计算投资的多个维度,并探讨了搜索方法与微调方法的兼容性。未来的工作可以进一步探索更多定制的验证器设计和更高效的搜索算法。
具体细节与案例分析
ImageNet上的实验细节:
- 使用SiT-XL模型在ImageNet上进行类条件生成。
- 通过随机搜索算法和不同的验证器(如FID、IS、CLIP、DINO)组合,研究了推理时计算预算对FID和IS的影响。
- 实验结果表明,随着推理时计算的增加,FID和IS均得到显著改善。
文本到图像生成的实验细节:
- 使用FLUX.1-dev模型在DrawBench和T2I-CompBench上进行文本到图像生成。
- 通过随机搜索算法和Verifier Ensemble验证器,研究了推理时计算预算对样本质量的影响。
- 实验结果表明,Verifier Ensemble验证器在需要广泛评估方面的生成任务中表现优异。
案例分析:
- 在ImageNet上,使用FID作为验证器时,随着推理时计算的增加,样本质量显著提高,但增加到一定程度后提升趋于平稳。
- 在DrawBench上,使用Verifier Ensemble验证器时,通过搜索可以找到与文本提示更好对齐的样本,从而显著提高样本质量。
这些具体细节和案例分析进一步支持了本文所提出的推理时缩放框架的有效性和实用性。