生成式人工智能/大语言模型

重新思考如何训练 Diffusion 模型

在探索了扩散模型采样、参数化和训练的基础知识之后,我们的团队开始研究这些网络架构的内部结构。请参考 生成式 AI 研究聚焦:揭开基于扩散的模型的神秘面纱 了解更多详情。

结果证明这是一项令人沮丧的练习。任何直接改进这些模型的尝试都会使结果更加糟糕。它们似乎处于微妙、微调、高性能的状态,任何更改都会破坏平衡。虽然通过彻底重新调整超参数可以实现好处,但下一组改进将需要重新经历整个过程。

如果您熟悉这种繁琐的开发循环,但您不直接使用扩散,请继续阅读。我们的研究结果针对大多数神经网络及其训练背后的普遍问题和组件。

我们决定打破这个循环,回顾一下基础知识。为什么架构如此易碎?网络中是否存在破坏训练进程的未知现象?我们如何使其更加稳健?归根结底:由于这些问题,我们目前还剩下多少性能?

我们最近的论文 分析和改进扩散模型的训练动力学 中报告了我们的研究结果和细节。我们仔细分析和重新思考训练动态,这是许多旗舰级图像生成器模型的基础。我们报告了在训练速度和生成质量方面的出色性能,以实现降噪扩散。图 1 显示,我们的模型的质量与先前的工作相当,只是模型复杂性和训练时间的一小部分,并且在使用更大的模型时明显优于先前的工作。

A scatter graph plotting FID score on vertical axis and model complexity on horizontal axis.
图 1.生成的图像质量随模型复杂性和最新先进方法的变化而变化。较低的 FID 分数 (垂直轴) 表示质量更好

我们得出了简化的网络架构和训练方法,我们称之为 EDM2,这是一种可靠、清晰的方法,可隔离 ADM 的强大核心,同时摆脱历史包和束缚。

此外,我们还介绍了一个理解不足但至关重要的网络权重指数移动平均值程序,并大大简化了此超参数的调优。

这篇文章介绍这项研究的主要发现。要查看参考实现代码,请访问在 GitHub 上发布的 NVlabs/edm2

训练动态是什么?为何如此重要?

假设基准 ADM 架构如图 2 所示。

A diagram of connected layers constituting the ADM network.
图 2.ADM 网络架构

我们可以通过从输出端发送的梯度反馈信号来成功训练如此深度和复杂的网络,这真是一个奇迹。我们如何确保它以健康和平衡的方式真正到达网络的所有部分,并训练每个层充分发挥其潜力?

从历史上看,这一点很不明确,并且由于梯度爆炸和消失等问题,网络将受到性能不佳的影响。在 2010 年代早期引入现代优化器以及归一化和初始化程序后,情况迅速改善。训练动态已经解锁,到目前为止,我们还没有看到模型复杂性和数据的扩展。

随意散布批量或组归一化层等组件通常可以解决最明显的问题,而这只能解决。但是,现代网络是真正运行良好的学习机器,还是它们在日益复杂的架构和训练任务中作响?扩散损失非常杂乱和复杂,通过仔细检查,训练几乎无法将所有内容保持在一起。

控制重量和激活量

第一个提示是:ADM 训练并非一切都正确。跟踪网络中传播的激活函数中看到的值的大小以及存储在层中的学习权重张量,可以显示整个训练过程中的稳步增加 (线对应于选择的不同层)。

Line graph of Activations and Weights with magnitude on vertical axis, and training time as horizontal axis. The magnitudes grow steadily over time.
图 3.作为时间函数绘制的不同层的激活量和权重大小

这种行为很常见 (不仅仅是在扩散中),通常会被忽略为另一种奇怪的好奇心,因为它似乎并不阻止网络学习。但是,我们发现这是一种不健康的现象,会严重影响训练的速度、可靠性和可预见性,并影响最终结果的质量。

为什么会出现这个问题?是什么导致了重量和活化增长,以及它们有哪些有害影响?这些问题带来了许多问题,每个问题都在复杂的交互网络中相互交织。

首先,我们发现 **权重** 只是一个数值伪影,而不是训练“有意”为了达到某些有益的目的而寻求的东西。这种现象是由梯度更新中存在的噪声和数值近似值错误所引起的统计趋势。尽管存在这种附带性质,但由于两个不同的原因,它存在问题:…

  1. 训练会随着时间的推移而饱和,因为当现有权重已经很大时,权重的更新会产生较小的相对影响。这就像将一成分投入罐子里一样。过了一段时间,再添加一不会有太大的变化。训练速度会降低到爬行 (每层的速度不同),即使有更多的东西需要学习。换言之,训练会受到每层学习率下降的控制。
  2. 权重乘法作用于 **激活** 通过卷积和矩阵乘积。当权重增加时,该层的激活也会增加。因此,权重增长会导致激活增长。

那么,激活增长又如何呢?为什么会出现问题呢?

通过网络中间某个地方的单个层的眼睛查看世界,在任何训练迭代中,它收到的输入激活量的增长略大于过去看到的增长。虽然其权重已适应以前较小的输入,因此已经过时,但它使用这些权重。后来,它收到一个梯度信号,告诉它如何解决这一差异。但是反馈来得太迟,计算已经有点错误。在下一次迭代中,输入激活量再次增加,并且重复相同的过程。

通过诱导,在不断增长的激活下,所有层在网络输出始终存在略有错误。无论训练持续多久,优化永远无法完全达到。这在扩散降噪等高精度任务中非常重要。

补救办法

认识到问题后,如何解决这些问题?为了简化任务,我们首先严格检查了网络中的每个移动部件,并消除了不必要的松。例如,我们删除了所有可学习的偏差 (令人惊讶的是,没有任何不良反应),这使我们更容易推理大小,因为它为我们留下了更少的可学习参数类型。论文详细介绍了许多其他改进,这些改进共同使网络的行为更加可预测和稳定。

清晰明了之后,我们针对之前详述的问题进行了核心改进。以下各节详细介绍了关键步骤:

  1. 为所有应用统一的大小保存原则。层
  2. 控制有效学习率
  3. 删除现已被取代 (且存在问题的) 组归一化层

消除激活和重量增长

我们采用一个简单的总体原则:每个层都必须根据统计预期保留其输入激活的大小。由于没有层主动改变激活大小,它们会随着时间的推移和网络各部分之间保持大致不变。

我们首先用“非学习”层 (例如非线性 (我们认为这些层与学习层分开)、连接和加法来说明这一想法。例如,SiLU 非线性会降低其输入激活的大小,将负值拉向零。使用大量随机数据馈送此层,平均而言,其大小将变化 0.596 倍。保持大小的 SiLU 版本将输出激活除以这个统计预期常量,近似于取消大小的变化。

如何将此原则应用于包含可学习权重的卷积和全连接层?激活大小的预期变化与权重大小成正比。为了消除这一点,我们重新调整了权重,使其始终保持在单位大小 (带有一些细微之处 – 请参阅论文附录或代码版本以了解详情)。

至关重要的是,这可以消除权重增长会导致激活次数增加的直接因果关系,并消除永远过时的层的问题。激活数量的标准化还使网络更具可预见性。例如,当组合来自两个不同分支的激活时,我们可以根据选择的比例混合它们的贡献,而不会因为不受控制的数量而意外超过一个分支。

请注意,这如何让两鸟获得一块石头:激活过程保持稳定,并且明确禁止权重增长 (图 4)。

A line graph of magnitude-preserving layers showing flat lines over time
图 4.使用大小保持层时,激活和权重变得稳定

控制学习率衰减

消除权重增长还可以消除不受控制的学习率的下降。这是一个很好的基准情况,因为现在每次 Adam 优化器权重更新都会对每个权重产生大致一致的影响,而不管训练进展有多远或位于网络中的位置。

相比之下,原始网络层根据其不受控制的数量具有“有效学习率”:一些层学习速度太慢,另一些则速度太快,不稳定;一些在完成之前就停止了学习,而另一些则从未安顿下来。

但是,学习率衰减是必要的,因为它允许网络在学习时逐渐专注于数据的更精细的细节。我们希望根据我们的条件来执行此操作,并且在归一化权重下,我们可以相信所有层都遵循指定的速率。在实践中,简单的受控衰减调度在此设置中效果很好。

删除组归一化层

此时,ADM 网络中的众多组归一化层已成为冗余,因为幅度保存方案已经控制了激活幅度。有趣的是,删除这些层可以进一步改善结果。这表明直接应用激活统计信息可能会产生不利的副作用。

这激发了我们对原始 StyleGAN 架构中所发现的几个有趣的发现的回忆,详情如下:适用于生成式对抗网络的基于风格的生成器架构。事实证明,生成器网络已经学会了通过在其内部激活贴图中增加一个大规模的本地化 Blob(在输出图像中留下一个难看的标记)来“欺诈”。这将通过破坏其计算的平均值来绕过激活量级归一化。由于网络为了逃脱而跳的圆圈,严格的归一化肯定会严重限制网络的期望。

我们认为,除以给定激活张量中的平均大小会在所有像素和特征之间引入意外的远程依赖关系和复杂的平衡要求。权重归一化在这方面较为柔和。一个角出现的明亮像素不会使图像的其余部分变暗。

指数移动平均线

模型平均是深度学习中广泛使用的一种实践。我们的想法是在训练期间跟踪“最近”权重的运行指数移动平均 (EMA).这些权重不在训练期间使用,而是在推理时使用。我们的想法是,原始训练权重很杂,对最近显示的训练样本影响太大。

扩散模型训练的一个令人费解的方面 (通常隐藏在研究论文附录中的不透明超参数表或代码库中的默认参数中) 是需要应用很长的平均值来获得良好结果,通常是整个训练长度的几个百分比。使用最新的训练权重 (无 EMA) 会产生灾难性的糟糕结果。另一方面,过多的平均值会过度强调早期阶段的权重,包括随机初始化。这并不方便,因为对于平均速率的逻辑似乎并不多,并且需要在训练之前确定数量。还是这样?

后期重建

在 分析和改进扩散模型的训练动力学 中,我们改变了 EMA 配置文件曲线的形状,使其随着训练长度“舒展”,并提出了一种在训练后重建具有不同 EMA 长度的网络的后期方法。我们的想法是存储具有固定较短 EMA 长度的中间训练状态的定期快照,然后我们可以通过合适的线性混合从中重建各种更长的 EMA 配置文件。与重新运行整个训练相比,组合几十个快照的效率要高出几个数量级;确定最佳 EMA 长度是一个简单的机械过程。

图 5 展示了如何在权重历史中将单个快照合并为更长的平均值的示例。着色区域表示每个快照在不同时间收集的权重。

Two-part graph, with top part showing a sequence of short EMA profiles and the bottom part showing a combined long EMA profile.
图 5.在 训练期间保存的每个快照都会随着时间的推移收集一段较短的 EMA 片段。在 Pose – hoc 重建中,这些快照会组合起来以生成更长的 EMA 配置文件

此方法支持将性能指标绘制为 EMA 长度的函数,使我们能够对其行为进行前所未有的深入了解。密集测量作为 EMA 长度函数的生成图像的质量 (FID 分数越低越好) (占完整训练运行长度的百分比),会得出如图 6 所示的图形。它比较不同网络配置的性能,这些配置用字母表示。

A plot with FID on vertical axis and EMA length on horizontal axis, showing approximately parabola-shaped lines corresponding to different models. The bottoms of the graphs are highlighted as the optimal choices of EMA.
图 6.作为 EMA 长度函数的一系列不同模型的质量 (FID,越低越好)

最佳值非常锐利,当偏离最佳值时,质量会迅速下降。如果选择错误的 EMA,一个好的模型可能会表现出任意较差的性能。盲目使用次优 EMA (例如,8%或 18%) 会错误地表明获胜的配置 G 低于其他一些竞争对手。

使用此新工具可以发现一些令人惊讶的现象。例如,最佳 EMA 长度对使用无分类器指导非常敏感 (图 7)。

Similar to Figure 6, this plot shows the performance plots of a single network, but using a range of different choices for classifier-free guidance parameter.
图 7.启用指导后,最佳 EMA 长度会显著缩短

基于这些知识,扫描 EMA 值是未来模型调整和评估的重要组成部分。

结果和结论

我们的团队使用隐扩散彻底评估了常见 ImageNet-512 设置的改进,并在此广泛使用的基准测试中达到创纪录的 1.81 FID.但是,仅仅查看行数可能会产生误解。重要的是随大小进行缩放。

为此,我们训练了六个规模越来越大的模型,分别为 EDM2-XS、S、M、L、XL 和 XXL.图 1 显示了每个模型的 FID 分数 (质量) 作为模型复杂性 (每个网络评估的千兆浮点运算) 的函数,以及最近的先进竞争对手,如扩散 Transformer 和 VDM++.对于每个模型,我们还绘制了在采样期间启用无分类器指导的变体,因为这是改进当前模型扩散结果的关键工具。

因此,我们的模型在给定的预算范围内一致地达到了最好的降价效果,并达到了与之前最先进的模型相似的质量,模型缩小了 5 倍。在参数计数、推理时间采样费用和训练成本方面进行的比较也说明了类似的情况 (图 8)。

Similar to Figure 1, but plotted against different quantities on the horizontal axis.
图 8.在模型容量、采样成本和训练成本方面衡量的性能

当然,生成的图像也可以而且应该被查看 (图 9)。但该领域严重依赖指标是有原因的,因为很难直观判断来自少量图像的训练数据的保真度。

 A grid of images showing various subjects, such as animals and food.
图 9.从 ImageNet-512 数据集上的先进模型生成的样本

我们还评估了使用 FD 的性能DINOv2指标,该指标最近被提出来改进 FID 的一些盲点。我们看到我们的方法也有类似的优势。

除了更大限度地提高性能外,一个关键目标是简化架构,使其更易于调整和探索。根据我们的实践经验,最终架构的行为似乎更可预测,我们发现它在变化方面更加可靠。

您可能已经注意到,大多数研究结果都是通用的,与 ADM 架构无关,对此,甚至没有与扩散有关。这确实是未来研究的一个有趣的途径。这些原则也可以直接应用于,例如,扩散 Transformer 训练。例如,在先前的研究中,图像分类器已从归一化方案的相关重新设计中受益。强制权重归一化和大小保留层等关键组件是自包含的,可以轻松应用到其他网络。

 

Tags