USF++：扩散概率模型求解器搜索的统一采样框架——实现极低步数下的高质量生成

《IEEE Transactions on Pattern Analysis and Machine Intelligence》：USF++: A Unified Sampling Framework for Solver Searching of Diffusion Probabilistic Models

本文针对扩散概率模型(DPM)采样速度慢的问题，提出统一采样框架USF++，通过进化搜索自动寻找最优求解器调度策略。该研究在CIFAR-10、ImageNet等数据集上显著提升极低函数评估次数(NFE≤5)下的生成质量，在Stable-Diffusion上实现2倍加速，为快速高质量生成提供新思路。

　　

在人工智能生成内容领域，扩散概率模型(Diffusion Probabilistic Models, DPMs)近年来展现出强大的图像生成能力。然而，其采样过程需要大量函数评估次数(Number of Function Evaluations, NFE)，通常需数百甚至上千步计算，严重限制了实际应用效率。尽管已有研究通过高阶数值方法将采样步数缩减至20步左右，但在极低NFE（少于10步）设置下，样本质量仍不尽如人意。

现有加速方法主要分为两类：一类需重新训练神经网络，成本高昂；另一类致力于设计高效的微分方程(Differential Equations, DE)求解器。后者虽无需重新训练，但其求解策略往往在整个采样过程中保持固定，未能充分考虑不同时间步的特性差异，导致在有限步数下性能提升受限。

针对这一挑战，清华大学电子工程系团队在《IEEE Transactions on Pattern Analysis and Machine Intelligence》上发表论文，提出了一种名为USF++的统一采样框架。该框架基于扩散常微分方程(Ordinary Differential Equation, ODE)的指数积分形式，将单步求解过程分解为七个可独立决策的组件：时间步离散化方案、神经网络预测类型、泰勒展开阶数、导数估计方法、ODE校正器使用、解析第一步(Analytical First Step, AFS)技术以及U-net结构中的跳跃连接系数。

研究团队通过理论分析和实验验证发现，不同时间步适合的求解策略存在显著差异。这一关键观察表明，精心设计的求解器调度策略（即分配给每个时间步的求解策略组合）对样本质量有重要影响。然而，由于决策空间巨大，手动设计最优调度策略极具挑战。

为此，研究人员采用进化搜索方法自动寻找高性能的求解器调度。该方法以FID(Fréchet Inception Distance)作为评估指标，在固定NFE预算下，通过初始化、变异、交叉和选择等操作，迭代优化求解器调度策略。

关键技术方法包括：基于指数积分的统一求解框架构建，支持噪声预测、数据预测和g_θ预测三种网络输出类型；泰勒展开阶数和导数估计方法的灵活配置；解析第一步技术节省NFE；U-net跳跃连接系数的优化；以及针对不同NFE预算的进化搜索策略。实验在CIFAR-10、ImageNet-256、LSUN-Bedroom等数据集上进行，使用预训练模型验证。

主要结果

实验结果表明，USF++在多个数据集上均显著优于现有最优方法。在CIFAR-10数据集上，仅用5步NFE即可达到3.89的FID值，远超基线方法的12.45；在LSUN-Bedroom数据集上，3步NFE即可实现8.62的FID值，而基线方法DPM-Solver-v3的FID为81.33。此外，将USF++应用于Stable-Diffusion模型，在文本到图像生成任务上实现了2倍加速（从10步减少到5步），且不牺牲生成质量。

消融研究

通过消融实验验证了各组件的重要性。当移除跳跃连接系数后，性能显著下降，表明其是必要的可搜索组件。同时，研究比较了使用不同数量图像计算FID对搜索效果的影响，发现1000张图像在估计准确性和搜索广度间取得了最佳平衡。

讨论分析

研究深入探讨了USF++的有效性原因。传统方法依赖局部误差阶分析，假设高阶方法必然带来更好性能。然而实验发现，盲目使用高阶更新（如对所有步使用3阶泰勒展开）会导致性能急剧下降（FID从3.93升至459.91）。这表明局部误差阶分析忽略了误差累积和模型本身不准确性等现实因素，而USF++通过端到端的质量驱动搜索，隐式考虑了这些复杂因素。

结论与展望

USF++为扩散概率模型的采样策略研究提供了系统性的分析框架，揭示了时间步异质性的重要性，并通过自动化搜索实现了在极低NFE下的高质量生成。该方法无需重新训练模型，为扩散模型的实用化部署提供了有效解决方案。

未来研究方向包括：系统探索当前搜索空间之外的其他策略，开发加速评估方法以降低计算开销，以及研究更高效的搜索算法进一步优化调度发现过程。这些扩展将增强框架的实际适用性，推动扩散模型在更多场景下的应用。