原创 小鹿 2024-10-24 13:17 北京
作者 | 小鹿
家人们，OpenAI 又上新了！
这次他们推出了全新的生成式模型sCM(Simplifying Continuous-Time Consistency Models)，支持视频、图像、三维模型和音频的生成。
这款模型的性能非常惊人，与传统的扩散模型相比，生成同等质量内容的速度提高了惊人的50倍！最大号的 1.5B 参数模型，在单卡 A100 上生成一张图片只需 0.11 秒！而且效果与Diffusion模型相近。
如下图所示，几乎同样的时间，扩散模型还是一团马赛克，sCM已经开始初见雏形了！
OpenAI同时发布了研究论文，两位华人作者全都毕业于清华。
官方博客：
https://openai.com/index/simplifying-stabilizing-and-scaling-continuous-time-consistency-models/
研究动机
尽管在扩散模型为生成式人工智能带来了巨大的突破，使得生成逼真的图像、3D模型、音频和视频成为可能，但其低效率仍是一个待解决的问题。
扩散模型的采样过程通常需要数十甚至数百步才能生成一个样本，这极大地限制了其在实际应用中的效率和可扩展性。而为了提高扩散模型生成的效率，研究人员们提出了各种蒸馏技术加速采样，但这些技术可能导致计算成本过高、训练过程复杂和样本质量下降等新问题。而连续模型（Consistency Models, CMs）又引入了额外的超参数，容易产生离散误差。
因此，sCM的研究团队提出改进扩散模型的训练模式，学习一个函数，将带噪声的图像映射到其在PF-ODE轨迹上的下一个时间步的更清晰版本。这种方法不是一步到位地去除所有噪声，而是根据PF-ODE的方向，逐步将图像向更清晰的方向推进。在两步采样的情况下，sCM会进行两次这样的映射，最终得到一个相对更清晰的图像。
关键改进
sCM的核心理念在于一致性，旨在使模型在连续时间步的输出保持一致性,通过学习 PF-ODE 的单步解析，sCM 能够直接将噪声转化为清晰图像，实现一步到位的转换。主要改进点包括：
时间条件策略改进（Identity Time Transformation）：采用 替代 ，解决了当 t 接近 T 时的数值不稳定性问题。
位置时间嵌入（Positional Time Embeddings）：使用位置嵌入替换傅里叶嵌入，以避免由傅里叶嵌入引起的不稳定性。
自适应双重归一化（Adaptive Double Normalization）：保留了 AdaGN 在扩散训练中的表现力，但消除了其在 CM 训练中的不稳定性。
自适应权重（Adaptive Weighting）：之前的权重设计是手动的，对于不同的数据分布和网络架构可能不是最优的。根据EDM2提出了数据分布和网络结构自动调整训练目标的权重，减少了手动调参的需求。
切线归一化/裁剪（Tangent Normalization/Clipping）：控制梯度方差，进一步增强训练的稳定性。
JVP 重新排列和 Flash Attention 的 JVP 计算：在Flash Attention风格的单次前向传递中高效地计算softmax自注意力及其JVP，显著减少了注意力层中JVP计算的GPU内存使用，提高了大规模模型训练的数值精度和效率。
渐进式退火：使训练过程更加稳定，易于扩展到大规模模型。
扩散微调和切线预热：通过对预训练的扩散模型进行微调和逐步预热切线函数的第二项，加速收敛并增强稳定性。
实验分析
研究团队在多个图像数据集进行了实验分析，包括CIFAR-10、ImageNet 64×64和ImageNet 512×512等。
在可视化分析中，路径的展示非常清晰：蓝线代表扩散模型逐步进行的采样过程，而红线则显示了一致性模型更加直接和快速的采样方法。通过采用一致性训练或一致性蒸馏技术，sCM能够有效地训练出一致性模型，显著减少了生成高质量样本所需的步骤数量。
目前研究团队训练的最大的sCM模型拥有15亿个参数。在一台未经过推理优化的A100 GPU上，生成一个样本的时间仅需0.11秒。通过对硬件和系统进行定制化优化，可以进一步提高生成速度，使得在图像、音频、视频等多个领域的实时生成成为可能。
sCM模型在标准的FID（Fréchet Inception Distance）评分上表现出色，分数越低表示质量越高。此外，sCM模型在有效采样计算量方面也表现优异，即生成每个样本所需的总计算成本。如下图所示，使用2步sCM生成的样本在质量上与当前最佳方法相当，而所需的有效采样计算量不到10%，极大地提高了采样效率。
结语
sCM通过改进的架构和训练目标，简化并稳定了连续时间一致性模型的训练过程，使其能够有效地扩展到拥有15亿参数的ImageNet 512×512规模。这些改进结合起来，不仅提高了在不同数据集和模型规模上的性能，而且在大规模应用中表现出了优于其他少步采样方法的可预测的扩展性。此外，研究团队也指出，在未来的研究中，根据特定应用的需求，评估sCM的质量可能需要采用不同的方法。
