【AIGC】2024-CVPR-分析和改进扩散模型的训练动态

分析和改进扩散模型的训练动态

作者：Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine
单位：NVIDIA, Aalto University
论文地址：https://openaccess.thecvf.com/content/CVPR2024/html/Karras_Analyzing_and_Improving_the_Training_Dynamics_of_Diffusion_Models_CVPR_2024_paper.html

摘要

 扩散模型目前以其无与伦比的大数据集扩展能力在数据驱动图像合成领域占据主导地位。在本文中，我们在不改变其高级结构的情况下，确定并纠正了流行的 ADM 扩散模型架构中训练不均匀和无效的几个原因。在训练过程中观察到网络激活和权重中不受控制的幅度变化和不平衡，我们重新设计了网络层以保持激活、权重和更新幅度的预期。我们发现，系统地应用这种理念可以消除观察到的漂移和不平衡，从而在同等计算复杂度下产生更好的网络。我们的修改将 ImageNet-512 合成中之前的 FID 记录从 2.41 提高到 1.81，这是使用快速确定性采样实现的。

 作为一项独立贡献，我们提出了一种事后设置（post-hoc）指数移动平均（exponential moving average，EMA）参数的方法，即在完成训练运行后。这允许精确调整 EMA 长度而无需进行多次训练，并揭示其与网络架构、训练时间和指导的令人惊讶的交互。

1. 引言

 由于去噪扩散模型的进步，基于文本提示、示例图像或其他形式输入的高质量图像合成已变得广受欢迎 [21, 46, 62–65, 70]。基于扩散的方法可以生成高质量的图像，同时提供多功能控件 [8, 17, 20, 44, 75] 和引入新主题的便捷方式 [12, 56]，并且它们还扩展到其他模态，如音频 [38, 51]、视频 [5, 22, 24] 和 3D 形状 [42, 50, 52, 61]。Yang et al. [71] 给出了方法和应用的最新概述。

 从高层次上讲，扩散模型通过反复应用图像去噪将纯噪声图像转换为新的生成图像。从数学上讲，每个去噪步骤都可以通过分数匹配 [26] 的视角来理解，并且通常使用配备自注意力 [69] 层的 U-Net [21, 55] 来实现。由于我们对扩散模型背后的理论没有贡献，因此我们建议感兴趣的读者参考 Sohl-Dickstein et al. [62]、Song 和 Ermon [64] 和 Ho et al. [21] 的开创性著作，以及 Karras et al. [33]，他们在共同的背景下构建了各种数学框架。

 尽管似乎可以毫无阻力地扩展到非常大的数据集和模型，但由于高度随机的损失函数，扩散模型的训练动态仍然具有挑战性。最终的图像质量由整个采样链中预测的微弱图像细节决定，中间步骤的小错误可能会在后续迭代中产生滚雪球效应。网络必须准确估计各种噪声水平、高斯噪声实现和条件输入中的平均清晰图像。考虑到在所有这些方面都是随机的混沌训练信号，学习这样做很困难。

 为了在如此嘈杂的训练环境中有效学习，网络理想情况下应该对参数更新具有可预测且均匀的响应。我们认为，目前最先进的设计并未满足这一理想条件，这会损害模型的质量，并且由于超参数、网络设计和训练设置之间的复杂相互作用，很难改进模型。

 我们的首要目标是了解有时微妙的、由于意外现象而导致评分网络训练动态失衡的方式，并逐一消除这些影响。我们方法的核心是权重、激活、梯度和权重更新的预期量级，所有这些都已被确定为先前工作中的重要因素（例如，[1, 3, 6, 7, 9, 37, 39-41, 59, 73, 74]）。粗略地说，我们的方法是通过一组清晰的设计选择来标准化所有量级，以统一的方式解决它们的相互依赖关系。

 具体来说，我们在不改变其整体结构的情况下对 ADM [11] U-Net 架构进行了一系列修改，并在此过程中展示了显著的质量改进（第 2 节）。最终的网络是 ADM 的替代品。它在有指导和无指导的情况下分别创下了 ImageNet-512 图像合成的新纪录 FID 1.81 和 1.91，而之前最先进的 FID 分别为 2.41 和 2.99。它在模型复杂度方面表现尤为出色（图 1），并且使用快速确定性采样而不是以前方法中使用的慢得多的随机采样来实现这些结果。

 作为一项独立贡献，我们提出了一种事后设置指数移动平均（EMA）参数的方法，即在训练运行完成后设置 (第 3 节)。模型平均 [27, 49, 57, 68, 72] 是所有高质量图像合成方法中不可或缺的技术 [2, 11, 23, 29, 31, 33, 46, 48, 54, 60, 63, 65]。不幸的是，EMA 衰减常数是一个难以调整的超参数，因为只有当训练接近收敛时，微小变化的影响才会显现出来。我们的事后 EMA 允许基于训练期间存储的预集成权重快照准确高效地重建具有任意 EMA 配置文件的网络。它还支持许多以前在计算上不可行的探索 (第 3.3 节)。

 我们的实现和预训练模型可在 https://github.com/NVlabs/edm2 上找到。

2. 改善训练动态

 现在让我们继续研究和消除与评分网络训练动态中各种不平衡相关的影响。作为我们的基线，我们采用在 EDM [33] 框架中实现的 ADM [11] 网络。该架构将 U-Net [55] 与自注意 [69] 层相结合（图 2a、b），其变体已广泛应用于大规模扩散模型，包括 Imagen [58]、稳定扩散 [54]、eDiff-I [2]、DALL-E 2 [47, 53] 和 DALL-E 3 [4]。我们的训练和采样设置基于 EDM 公式，具有恒定学习率和 32 个确定性二阶采样步骤。

 我们使用类条件 ImageNet [10] 512×512 数据集进行评估，并且像大多数高分辨率扩散模型一样，在执行 8 倍空间上采样的预训练解码器 [54] 的潜在空间中运行。因此，我们的输出在解码之前是 64×64×4。在探索过程中，我们使用一个适度大小的网络配置，大约有 3 亿个可训练参数，第 4 节稍后将介绍扩大网络的结果。对 2147M（$=2^{31}$）张图像进行训练，每批 2048 张，这足以让这些模型达到最佳 FID。

 我们将分几个步骤构建改进的架构和训练程序。我们的阐述重点介绍基本原理和与网络相关的变化。有关每个架构步骤的详细信息以及相关方程式，请参阅附录 B。

 Improved baseline (CONFIG B)。由于原始 EDM 配置针对的是 RGB 图像，我们将输出通道数增加到 4，并用 ImageNet-512 图像的 64×64×4 潜在表示替换训练数据集，全局标准化为零均值和标准差 ${\sigma }_{data}=0.5$。在此设置中，我们获得了 8.00 的基线 FID（见表 1）。

 然后，我们解决了原始 EDM 训练设置中的一个缺点：虽然 EDM 中的损失权重在初始化时将所有噪声水平的损失幅度标准化为 1.0，但随着训练的进行，这种情况不再成立。然后，梯度反馈的幅度在噪声水平之间变化，以不受控制的方式重新加权它们的相对贡献。

 为了抵消这种影响，我们采用了 Kendall et al. [34] 提出的多任务损失的连续泛化。实际上，我们跟踪原始损失值作为噪声水平的函数，并根据其倒数缩放训练损失。有关更多详细信息和推理，请参阅附录 B.2。这些变化共同将 FID 从 8.00 降低到 7.24。

 Architectural streamlining (CONFIG C)。为了便于分析训练动态，我们继续精简和稳定架构。为了避免处理多种不同类型的可训练参数，我们从所有卷积层和线性层以及调节路径中删除了附加偏差。为了恢复网络偏移数据的能力，我们将一个常数 1 的附加通道连接到网络的输入。我们进一步使用 He 的均匀初始化 [14] 统一所有权重的初始化，从 ADM 的原始位置编码方案切换到更标准的傅立叶特征 [67]，并通过删除其均值减法和学习缩放来简化组规范化层。

 最后，我们观察到，由于训练过程中关键和查询向量的幅度增长，注意力图通常会以脆弱和尖锐的配置结束。我们通过切换到余弦注意力 [13, 43, 45] 来纠正这个问题，在计算点积之前对向量进行规范化。实际好处是，这允许在整个网络中使用 16 位浮点数学，从而提高效率。这些变化共同将 FID 从 7.24 降低到 6.96。

2.2. 标准化激活量级

 随着架构的简化，我们现在开始解决训练动态中的第一个问题：激活量级。如图 3 第一行所示，尽管在每个块内使用了组归一化，但随着训练的进行，CONFIG C 中的激活量级仍会不受控制地增长。值得注意的是，增长没有显示出在训练运行结束时逐渐减弱或稳定的迹象。

 查看图 2b 中的架构，增长可能并不太令人惊讶：由于编码器、解码器和自注意力块的残差结构，ADM 网络包含长信号路径而没有任何归一化。这些路径积累了残差分支的贡献，并可以通过重复卷积放大它们的激活。我们假设，这种持续的激活量级增长不利于训练，因为它会使网络始终处于不收敛和非最佳状态。

 我们也尝试在主路径中引入组归一化层，但这导致结果质量显著下降。这可能与之前关于 StyleGAN [31] 的发现有关，过度标准化会削弱网络的功能，以至于各层学会了通过人为的图像伪影来绕过它。受到 StyleGAN2 [32] 中采用的解决方案以及其他寻求显式标准化替代方案的作品 [1, 6, 37] 的启发，我们选择修改网络，使各个层和路径保持预期的激活幅度，目标是消除或至少减少对数据依赖性标准化的需求。

 Magnitude-preserving learned layers (CONFIG D)。为了保持预期的激活幅度，我们将每层的输出除以该层引起的激活幅度的预期缩放，而不考虑激活本身。我们首先将其应用于模型每个部分的所有学习层（卷积和全连接）。

 鉴于我们寻求一种与输入激活的实际内容无关的方案，我们必须对它们做出一些统计假设。为简单起见，我们假设像素和特征图相互不相关，且具有相同的标准差 ${\sigma }_{act}$。全连接层和卷积层都可以看作由堆叠单元组成，每个输出通道一个单元。每个单元实际上将权重向量 ${\mathbf{w}}_i\in {\mathbb{R}}^n$ 的点积应用于输入激活的某个子集，以产生每个输出元素。在我们的假设下，第 $i$ 个通道的输出特征的标准差变为 ${\Vert {\mathbf{w}}_i\Vert }_2{\sigma }_{\mathrm{act}}$。为了恢复输入激活幅度，我们按通道除以 ${\Vert {\mathbf{w}}_i\Vert }_2$。¹

 我们同样可以将标量除法视为应用于 ${\mathbf{w}}_i$ 本身。只要梯度通过计算范数传播，该方案就等同于没有学习到输出尺度的权重正则化 [59]；以后我们将使用此术语。由于整体权重幅度不再对激活产生影响，我们通过从单位高斯分布中抽取来初始化所有权重。

 此修改消除了网络学习改变整体激活幅度的任何直接方法，如图 3（CONFIG D）所示，幅度漂移被成功消除。FID 也显着提高，从 6.96 提高到 3.75。

¹ 主要目标是切断权重和激活幅度之间的直接联系；为此，统计假设不需要完全成立。

2.3. 标准化权重和更新

 在标准化激活后，我们将注意力转向网络权重和学习率。如图 3 所示，CONFIG D 中的网络权重明显呈增长趋势，甚至比 CONFIG C 中的增长趋势还要快。造成这种情况的机制众所周知 [59]：使用前对权重进行归一化会迫使损失梯度垂直于权重向量，沿着这个方向迈出一步总是会落在离原点更远的点上。即使使用 Adam 优化器标准化了梯度幅度，净效应也是有效学习率（即网络权重更新的相对大小）会随着训练的进行而衰减。

 虽然有人认为有效学习率的衰减是一种理想的效果 [59]，但我们主张对其进行明确控制，而不是让它在各层之间不受控制地、不均匀地漂移。因此，我们将其视为我们寻求补救的另一种训练动态不平衡。请注意，将所有权重初始化为单位高斯可确保初始化时有效学习率统一，但之后则不然。

 Controlling effective learning rate (CONFIG E)。我们建议通过强制权重归一化来解决权重增长问题，即在每个训练步骤之前将每个权重向量 $w_i$ 明确归一化为单位方差。重要的是，我们仍然在训练期间在此基础上应用 “标准” 权重归一化，即在使用时对权重向量进行归一化。这会将训练梯度投影到 $w_i$ 所在的现在单位幅值超球面的切平面上（有关推导，请参阅附录 B.4）。这确保 Adam 的方差估计是根据实际切平面步骤计算的，并且不会被梯度向量的即将被擦除的法向分量所破坏。现在，网络中的权重和梯度幅值都相等，我们也统一了有效学习率。假设权重和梯度之间没有相关性，现在每个 Adam 步骤都会用梯度替换大约固定比例的权重。一些优化器 [3, 39, 73] 通过基于数据的梯度重新缩放明确实现了类似的效果。

 我们现在可以直接控制有效学习率。恒定的学习率不再导致收敛，因此我们引入了 Kingma 和 Ba [36] 所倡导的平方根倒数学习率衰减计划。具体来说，我们定义 $\alpha \left(t\right)={\alpha }_{\mathrm{ref}}/\sqrt{\max \left(t/t_{\mathrm{ref}},1\right)}$，其中 $t$ 是当前训练迭代，αref 和 tref 是超参数（有关实现细节，请参阅附录 D）。如图 3 所示，生成的 CONFIG E 在整个训练过程中成功保留了激活和权重大小。结果，FID 从 3.75 提高到 3.02。

2.4. 移除组归一化（CONFIG F）

 在控制了激活、权重和更新幅度后，我们现在可以移除跨像素操作的数据相关组归一化层，这些层可能会产生有害结果 [32]。尽管网络在没有任何归一化层的情况下也能成功训练，但我们发现，在编码器主路径中引入更弱的像素归一化 [30] 层仍然会带来一些好处。我们的假设是，像素归一化有助于抵消违反 CONFIG D 中标准化工作背后的统计假设的相关性。因此，我们移除所有组归一化层，并用 1/4 的像素归一化层替换它们。我们还从嵌入网络中删除了第二个线性层，并从网络输出中删除了非线性，并将残差块中的重采样操作合并到主路径上。FID 从 3.02 提高到 2.71。

2.5. 保持幅度的固定函数层（CONFIG G）

 为了完整起见，我们注意到网络仍然有不保持激活幅度的层。首先，傅里叶特征的正弦和余弦函数没有单位方差，我们通过将其放大 $\sqrt{2}$ 来纠正。其次，除非对此进行补偿，否则 SiLU [16] 非线性会衰减预期的激活单位方差分布。因此，我们修改它们以将输出除以 ${\mathbb{E}}_{x\sim \mathcal{N}\left(0,1\right)}{\left[\mathrm{silu}\left(\mathrm{x}\right)\right]}^{1/2}\approx 0.596$。

 第三，我们考虑两个网络分支通过加法或连接而连接的情况。在以前的配置中，每个分支对输出的贡献取决于不受控制的激活幅度。现在我们可以预期这些是标准化的，因此分支之间的平衡被揭示为一个有意义的可控参数 [7]。我们将加法运算切换为加权和，并通过实验观察到，在编码器和解码器块中，固定残差路径权重为 30% 效果最佳，在嵌入中为 50%。我们将输出除以该加权和的预期标准偏差。

 解码器中 U-Net 跳跃的连接已经保持了幅度，因为我们可以预期两个分支的幅度相似。然而，两个输入在后续层中的相对贡献与它们各自的通道数成正比，我们认为这是编码器和解码器超参数之间不必要且不直观的依赖关系。我们通过缩放输入来消除这种依赖性，使得连接结果的整体幅度保持不变，但输入的贡献变得相等。

 标准化完成后，我们确定了两个仍然需要按学习量缩放激活的特定位置。首先，我们在网络的最末端添加一个学习到的、零初始化的标量增益（即缩放），因为我们不能期望所需的输出始终具有单位方差。其次，我们将类似的学习增益应用于每个残差块内的调节信号，以便在初始化时禁用调节，并且其在每个编码器/解码器块中的强度成为学习参数。此时，我们可以在训练期间禁用 dropout [19, 66] 而不会产生任何不良影响，这在以前是不可能的。

 图 2c 展示了我们的最终设计，它比基线简单得多，也更容易推理。最终的 FID 为 2.56，与当前最先进的技术相比具有很强的竞争力，尤其是考虑到我们的探索架构的计算复杂性适中。

3. 事后 EMA

 众所周知，模型权重的指数移动平均（EMA）在生成图像合成中起着重要作用 [46, 65]，其衰减参数的选择对结果有显著影响 [29, 46]。

 尽管众所周知其重要性，但人们对衰减参数与训练和采样的其他方面之间的关系知之甚少。为了分析这些问题，我们开发了一种事后选择 EMA 配置文件的方法，即无需在训练前指定它。这使我们能够密集地采样 EMA 的长度并绘制其对质量的影响，揭示与网络架构、训练时间和无分类器指导的有趣相互作用。

 本节中方程和方法的更多细节、推导和讨论包含在附录 C 中。

3.1. 幂函数 EMA 曲线

 传统 EMA 维护网络参数的运行加权平均值 ${\hat{\theta }}_{\beta }$ 以及正在训练的参数 $\theta$。在每个训练步骤中，平均值通过 ${\hat{\theta }}_{\beta }\left(t\right)=\beta {\hat{\theta }}_{\beta }\left(t-1\right)+\left(1-\beta \right)\theta \left(t\right)$ 更新，其中 $t$ 表示当前训练步骤，从而产生早期训练步骤贡献的指数衰减曲线。衰减率由常数 $\beta$ 决定，该常数通常接近于 1。

 出于两个原因，我们建议使用基于幂函数而不是指数衰减的略微改变的平均曲线。首先，我们的架构修改倾向于更长的平均值；然而，非常长的指数 EMA 会在网络参数大多是随机的训练初始阶段赋予不可忽略的权重。其次，我们观察到一个明显的趋势，即较长的训练运行受益于较长的 EMA 衰减，因此平均曲线理想情况下应该随着训练时间自动扩展。

 上述两个要求都是通过幂函数来满足的。我们将时间 $t$ 的平均参数定义为

 其中常数 $\gamma$ 控制轮廓的锐度。使用此公式，${\theta }_{t=0}$ 的权重始终为零。这是可取的，因为随机初始化对平均值没有影响。得到的平均轮廓也是与尺度无关的：将训练时间加倍会自动将轮廓拉伸相同的倍数。

 为了在实践中计算 $\hat{\theta }\left(t\right)$，我们在每个训练步骤后执行增量更新，如下所示：

 因此，更新与传统 EMA 类似，但 $\beta$ 取决于当前训练时间。²

 最后，虽然参数 $\gamma$ 在数学上很简单，但它对平均曲线的形状有某种不直观的影响。因此，我们更喜欢通过其相对标准偏差 ${\sigma }_{\mathrm{rel}}$ 来参数化曲线，即其峰值相对于训练时间的 “宽度”：${\sigma }_{\mathrm{rel}}={\left(\gamma +1\right)}^{1/2}{\left(\gamma +2\right)}^{-1}{\left(\gamma +3\right)}^{-1/2}$。因此，当报告 EMA 长度为 10% 时，我们指的是 ${\sigma }_{\mathrm{rel}}=0.10$ 的曲线（相当于 $\gamma \approx 6.94$）。

² 从技术角度来说，将其称为 “EMA 曲线” 是一种错误的说法，因为权重衰减不是指数衰减。然而，考虑到它与传统 EMA 的用途相同，我们认为在这里创造一个新术语会产生误导。

3.2. 训练后合成新的 EMA 曲线

 我们的目标是允许在训练后自由选择 $\gamma$ 或等效的 ${\sigma }_{\mathrm{rel}}$。为了实现这一点，我们在训练期间维护两个平均参数向量 ${\hat{\theta }}_{{\gamma }_1}$ 和 ${\hat{\theta }}_{{\gamma }_2}$，其中常数 ${\gamma }_1=16.97$ 和 ${\gamma }_2=6.94$，分别对应于 ${\sigma }_{\mathrm{rel}}$ $0.05$ 和 $0.10$。这些平均参数向量会定期存储在训练运行期间保存的快照中。在我们所有的实验中，我们每约 8 百万张训练图像存储一次快照，即每 4096 个训练步骤存储一次快照，批处理大小为 2048。

 为了在训练期间或训练后的任何时间点重建对应于任意 EMA 轮廓的近似 $\hat{\theta }$，我们找到存储的 ${\hat{\theta }}_{{\gamma }_i}$ 的 EMA 轮廓与所需 EMA 轮廓之间的最小二乘最优拟合，并取存储的 ${\hat{\theta }}_{{\gamma }_i}$ 的相应线性组合。参见图 4 的说明。

 我们注意到，事后 EMA 重建不仅限于幂函数平均轮廓，也不仅限于对快照和重建使用相同类型的轮廓。此外，它甚至可以从每个快照的单个存储 $\hat{\theta }$ 完成，尽管其准确度远低于两个存储 $\hat{\theta }$。只要有足够数量的训练快照，就可以重新审视以前没有考虑事后 EMA 的训练运行，并尝试新的平均轮廓。

3.3. 分析

 利用事后 EMA 技术，我们现在可以分析不同设置中不同 EMA 长度的影响。

 图 5a 显示了表 1 中配置 B-G 中 FID 如何根据 EMA 长度而变化。我们可以看到，最佳 EMA 长度在配置之间有很大差异。此外，当我们接近最终配置 G 时，最佳值变得更窄，这最初可能看起来令人担忧。

 然而，如图 5b 所示，最佳值的狭窄似乎可以通过模型变得更加统一来解释，就每个权重张量 “首选” 的 EMA 长度而言。在此测试中，我们首先从网络的不同部分选择权重张量的子集。然后，对于每个选定的张量，我们分别执行扫描，其中仅更改选定张量的 EMA，而所有其他张量保持全局最优。结果显示为每个张量一行，揭示了对 FID 的惊人影响。有趣的是，虽然一个权重张量与其他张量不同步显然是有害的，但我们观察到在 CONFIG B 中，FID 可以提高 10%，从 7.24 降至 ∼6.5。在一个例子中，这是使用非常短的每个张量 EMA 实现的，在另一个例子中，则是使用非常长的 EMA。我们假设这些不同的偏好意味着任何全局选择都是一个难以妥协的妥协。对于我们最终的 CONFIG G，这种影响消失了，最优值更加明显：看不到 FID 的显著改善，张量现在对最佳 EMA 达成了一致。虽然事后 EMA 允许根据每个张量选择 EMA 长度，但我们尚未在本次实验之外探索这个机会。

 最后，图 5c 说明了最佳 EMA 长度在训练过程中的演变。尽管我们对 EMA 长度的定义已经与训练长度相关，但我们观察到，随着训练的进行，最佳值慢慢转向相对较长的 EMA。

4. 结果

 我们使用 512×512 分辨率的 ImageNet [10] 作为主要数据集。表 2 总结了使用我们的方法以及几种早期技术的各种模型大小的 FID。

 让我们首先考虑没有指导的 FID [20]，其中最好的先前方法是 VDM++ [35]，FID 为 2.99。即使是我们在第 2 节中用于架构探索的小模型 EDM2-S 也以 2.56 的 FID 击败了它。进一步扩大我们的模型将 FID 提高到 1.91，大大超过了之前的记录。如图 1 所示，我们的结果在模型复杂度方面更加显著。我们发现，在出现过度拟合的情况下，即当训练损失继续减少但验证损失和 FID 开始增加时，启用 dropout [19, 66] 可以改善我们的结果。因此，我们在出现过拟合迹象的较大配置 $\left(M–XXL\right)$ 中启用 dropout，而在有害的较小配置 $\left(\mathrm{XS},\mathrm{S}\right)$ 中禁用 dropout。

 附录 A 中给出了本节的其他定量结果、示例图像和详细比较。

 指导。值得注意的是，几种早期方法 [11, 48] 仅在使用无分类器指导 [20] 时才报告具有竞争力的结果。虽然指导仍然是控制单个结果图像的感知质量与生成分布的覆盖范围之间平衡的宝贵工具，但当目标是简单地匹配图像分布时，它应该不是必需的。

 图 6 绘制了使用各种指导强度作为 EMA 长度函数的小型模型 $\left(EDM2-\mathrm{S}\right)$ 的 FID。令人惊讶的是，最佳 EMA 长度在很大程度上取决于指导强度。如果没有事后 EMA，这类研究的成本非常高，因此我们假设某些现有技术中 vanilla 结果和指导结果之间的巨大差异可能部分是使用非最佳 EMA 参数的结果。利用我们最大的模型，适量的指导（1.2）进一步将 ImageNet-512 FID 从 1.91 提高到 1.81，创下了该数据集的新纪录。

 低成本指导。实现无分类器指导的标准方法是训练一个模型来支持条件和无条件生成 [20]。虽然概念上很简单，但这隐含地假设这两个任务都需要一个同样复杂的模型。然而，事实似乎并非如此：在我们的测试中，发现最小的 $\left(\mathrm{XS}\right)$ 无条件模型足以指导甚至最大的 $\left(\mathrm{XXL}\right)$ 条件模型——使用更大的无条件模型根本没有改善结果。

 表 2 中的结果是我们在所有配置中使用 $\mathrm{XS}$ 大小的无条件模型计算得出的。使用小型无条件模型可以大大减少指导的典型 2 倍计算开销。

 ImageNet-64。为了证明我们的方法不仅限于潜在扩散，我们提供了 ImageNet-64 中 RGB 空间扩散的结果。表 3 显示我们的结果优于使用确定性采样的早期方法。EDM [33] 创下的先前记录 FID 2.22 在类似模型复杂度下提高到 1.58，并通过缩放进一步提高到 1.33。L 尺寸模型能够饱和此数据集。

 此结果接近 RIN 使用随机采样实现的记录 FID 1.23。随机采样可以纠正去噪网络的不准确性，但这需要相当大的调整工作量和计算成本（例如，1000 与 63 NFE），使得随机采样对于大规模系统没有吸引力。使用随机采样很可能可以进一步改善我们的结果，但我们将其留待将来研究。

 事后 EMA 观察。除了前面几节讨论的相互作用之外，我们还得出了与 EMA 长度相关的两个初步发现。我们在这里将它们作为轶事介绍，并将详细研究留待将来的工作。

 首先，我们观察到，当学习率增加时，最佳 EMA 长度会下降，反之亦然，大致根据 ${\sigma }_{\mathrm{rel}}\propto 1/\left({\alpha }_{\mathrm{ref}}^2{t}_{\mathrm{ref}}\right)$。由此产生的 FID 也在 tref 的 2 倍范围内保持相对稳定。在实践中，将 ${\alpha }_{\mathrm{ref}}$ 和 $t_{\mathrm{ref}}$ 设置在正确的范围内似乎就足够了，这减少了仔细调整这些超参数的需要。

 其次，我们观察到，当模型容量增加时，最佳 EMA 长度趋于下降，并且当数据集的复杂性降低时也是如此。这似乎意味着更简单的问题需要更短的 EMA。

5. 讨论和未来工作

 我们改进的降噪器架构旨在替代广泛使用的 ADM 网络，因此我们希望它能够在大型图像生成器中得到广泛应用。现在，训练的各个方面都不再那么纠结，因此更容易对架构进行局部修改而不会破坏其他部分。这应该可以进一步研究 U-Net 的结构和平衡等。

 一个有趣的问题是，类似的方法是否同样适用于其他扩散架构，例如 RIN [28] 和 DiT [48]，以及扩散模型以外的其他应用领域。似乎这种以幅度为重点的工作在 ImageNet 分类器 [6, 7] 的特定主题之外引起了相对较少的关注。

 我们相信，事后 EMA 将使一系列以前无法实现的有趣研究成为可能。如果没有它，我们的一些图表将需要一千个 GPU 年才能制作出来；现在它们只需要一个月的 GPU。我们希望，廉价生产的 EMA 数据将使我们能够取得新的突破，以理解 EMA 在扩散模型中的确切作用，并找到设置 EMA 长度的原则性方法——可能基于每个层或每个参数。

 致谢。我们感谢 Eric Chan、Qinsheng Zhang、Erik Härkönen、Tuomas Kynkäänniemi、Arash Vahdat、Ming-Yu Liu 和 David Luebke 的讨论和评论，感谢 Tero Kuosmanen 和 Samuel Klenberg 维护我们的计算基础设施。

参考文献

[1] Devansh Arpit, Yingbo Zhou, Bhargava Kota, and Venu Govindaraju. Normalization propagation: A parametric tech- nique for removing internal covariate shift in deep networks. In Proc. ICML, 2016. 2, 3
[2] YogeshBalaji,SeungjunNah,XunHuang,ArashVahdat,Ji- aming Song, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, Bryan Catanzaro, Tero Karras, and Ming-Yu Liu. eDiff-I: Text-to-image diffusion models with ensemble of expert denoisers. CoRR, abs/2211.01324, 2022. 2
[3] JeremyBernstein,ArashVahdat,YisongYue,andMing-Yu Liu. On the distance between two neural networks and the stability of learning. In Proc. NIPS, 2020. 2, 4
[4] James Betker, Gabriel Goh, Li Jing, Tim Brooks, Jianfeng Wang, Linjie Li, Long Ouyang, Juntang Zhuang, Joyce Lee, Yufei Guo, Wesam Manassra, Prafulla Dhariwal, Casey Chu, Yunxin Jiao, and Aditya Ramesh. Improving image generation with better captions. Technical report, OpenAI, 2023. 2
[5] AndreasBlattmann,RobinRombach,HuanLing,TimDock- horn, Seung Wook Kim, Sanja Fidler, and Karsten Kreis. Align your latents: High-resolution video synthesis with la- tent diffusion models. In Proc. CVPR, 2023. 1
[6] Andrew Brock, Soham De, and Samuel L. Smith. Charac- terizing signal propagation to close the performance gap in unnormalized ResNets. In Proc. ICLR, 2021. 2, 3, 8
[7] Andrew Brock, Soham De, Samuel L. Smith, and Karen Simonyan. High-performance large-scale image recognition without normalization. In Proc. ICML, 2021. 2, 5, 8
[8] Tim Brooks, Aleksander Holynski, and Alexei A. Efros. In- structPix2Pix: Learning to follow image editing instructions. In Proc. CVPR, 2023. 1
[9] Minhyung Cho and Jaehyung Lee. Riemannian approach to batch normalization. In Proc. NIPS, 2017. 2
[10] JiaDeng,WeiDong,RichardSocher,Li-JiaLi,KaiLi,andLi Fei-Fei. ImageNet: A large-scale hierarchical image database. In Proc. CVPR, 2009. 2, 7
[11] Prafulla Dhariwal and Alex Nichol. Diffusion models beat GANs on image synthesis. In Proc. NeurIPS, 2021. 2, 3, 7, 8
[12] Rinon Gal, Yuval Alaluf, Yuval Atzmon, Or Patashnik, Amit Haim Bermano, Gal Chechik, and Daniel Cohen-Or. An image is worth one word: Personalizing text-to-image generation using textual inversion. In Proc. ICLR, 2023. 1
[13] Spyros Gidaris and Nikos Komodakis. Dynamic few-shot visual learning without forgetting. In Proc. CVPR, 2018. 3
[14] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level perfor- mance on ImageNet classification. In Proc. ICCV, 2015. 3
[15] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In Proc. ECCV, 2016. 3
[16] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (GELUs). CoRR, abs/1606.08415, 2016. 5
[17] Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or. Prompt-to-prompt image editing with cross attention control. In Proc. ICLR, 2023. 1
[18] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bern- hard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to a local Nash equilibrium. In Proc. NIPS, 2017. 2
[19] GeoffreyE.Hinton,NitishSrivastava,AlexKrizhevsky,Ilya Sutskever, and Ruslan R. Salakhutdinov. Improving neural networks by preventing co-adaptation of feature detectors. CoRR, abs/1207.0580, 2012. 5, 7
[20] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021. 1, 7, 8
[21] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffu- sion probabilistic models. In Proc. NeurIPS, 2020. 1, 2
[22] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P. Kingma, Ben Poole, Mohammad Norouzi, David J. Fleet, and Tim Salimans. Ima- gen Video: High definition video generation with diffusion models. CoRR, abs/2210.02303, 2022. 1
[23] JonathanHo,ChitwanSaharia,WilliamChan,DavidJ.Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. JMLR, 23(1), 2022. 2
[24] Jonathan Ho, Tim Salimans, Alexey A. Gritsenko, William Chan, Mohammad Norouzi, and David J. Fleet. Video diffu- sion models. In Proc. ICLR Workshop on Deep Generative Models for Highly Structured Data, 2022. 1
[25] Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. Sim- ple diffusion: End-to-end diffusion for high resolution images. In Proc. ICML, 2023. 2, 7
[26] Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. JMLR, 6(24), 2005. 1
[27] PavelIzmailov,DmitriiPodoprikhin,TimurGaripov,Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In Proc. Uncertainty in Artificial Intelligence, 2018. 2
[28] Allan Jabri, David J. Fleet, and Ting Chen. Scalable adaptive computation for iterative generation. In Proc. ICML, 2023. 7, 8
[29] MingukKang,Jun-YanZhu,RichardZhang,JaesikPark,Eli Shechtman, Sylvain Paris, and Taesung Park. Scaling up GANs for text-to-image synthesis. In Proc. CVPR, 2023. 2, 5
[30] TeroKarras,TimoAila,SamuliLaine,andJaakkoLehtinen. Progressive growing of GANs for improved quality, stability, and variation. In Proc. ICLR, 2018. 5
[31] Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proc. CVPR, 2019. 2, 3
[32] Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of StyleGAN. In Proc. CVPR, 2020. 3, 5
[33] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In proc. NeurIPS, 2022. 1, 2, 8
[34] Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In Proc. CVPR, 2018. 3
[35] Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the ELBO with data augmentation. In Proc. NeurIPS, 2023. 7, 8
[36] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proc. ICLR, 2015. 4
[54] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proc. CVPR, 2022.
[37] Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. Self-normalizing neural networks. In Proc. 2 NIPS, 2017. 2, 3
[38] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. DiffWave: A versatile diffusion model for audio synthesis. In Proc. ICLR, 2021. 1
[39] Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay balances learning across
neural networks. CoRR, abs/2305.17212, 2023. 2, 4
[40] Twan van Laarhoven. L2 regularization versus batch and weight normalization. CoRR, abs/1706.05350, 2017.
[41] Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning. In Proc. ICLR, 2020. 2
[42] Chen-HsuanLin,JunGao,LumingTang,TowakiTakikawa, Xiaohui Zeng, Xun Huang, Karsten Kreis, Sanja Fidler, Ming- Yu Liu, and Tsung-Yi Lin. Magic3D: High-resolution text-to-3D content creation. In Proc. CVPR, 2023. 1
[43] Chunjie Luo, Jianfeng Zhan, Xiaohe Xue, Lei Wang, Rui Ren, and Qiang Yang. Cosine normalization: Using cosine
similarity instead of dot product in neural networks. In Proc. ICANN, 2018. 3
[44] Ron Mokady, Amir Hertz, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or. NULL-text inversion for editing real im- ages using guided diffusion models. In Proc. CVPR, 2023. 1
[45] Quang-Huy Nguyen, Cuong Q. Nguyen, Dung D. Le, and Hieu H. Pham. Enhancing few-shot image classification with cosine transformer. IEEE Access, 11, 2023. 3
[46] Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In Proc. ICML, pages 8162– 8171, 2021. 1, 2, 5
[47] Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. GLIDE: Towards photorealistic image genera- tion and editing with text-guided diffusion models. In Proc. ICML, 2022. 2
[48] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proc. ICCV, 2023. 2, 7, 8
[49] Boris Polyak and Anatoli Juditsky. Acceleration of stochastic approximation by averaging. SIAM Journal on Control and Optimization, 30(4), 1992. 2
[50] BenPoole,AjayJain,JonathanT.Barron,andBenMildenhall. DreamFusion: Text-to-3D using 2D diffusion. In Proc. ICLR, 2023. 1
[51] Vadim Popov, Ivan Vovk, Vladimir Gogoryan, Tasnima Sadekova, and Mikhail Kudinov. Grad-TTS: A diffusion probabilistic model for text-to-speech. In Proc. ICML, 2021. 1
[52] Amit Raj, Srinivas Kaza, Ben Poole, Michael Niemeyer, Ben Mildenhall, Nataniel Ruiz, Shiran Zada, Kfir Aberman, Michael Rubenstein, Jonathan Barron, Yuanzhen Li, and Varun Jampani. DreamBooth3D: Subject-driven text-to-3D generation. In Proc. ICCV, 2023. 1
[55] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional networks for biomedical image segmentation. In Proc. MICCAI, 2015. 1, 2, 3
[56] Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. DreamBooth: Fine tuning text-to-image diffusion models for subject-driven gen- eration. In Proc. CVPR, 2023. 1
[57] David Ruppert. Efficient estimations from a slowly con- vergent Robbins–Monro process. Technical report, Cornell University – Operations Research and Industrial Engineering, 1988. 2
[58] ChitwanSaharia,WilliamChan,SaurabhSaxena,LalaLi,Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S. Sara Mahdavi, Rapha Gontijo Lopes, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. In Proc. NeurIPS, 2022. 2
[59] TimSalimansandDiederikP.Kingma.Weightnormalization: A simple reparameterization to accelerate training of deep neural networks. In Proc. NIPS, 2016. 2, 4
[60] Axel Sauer, Katja Schwarz, and Andreas Geiger. StyleGAN- XL: Scaling StyleGAN to large diverse datasets. In Proc. SIGGRAPH, 2022. 2, 7
[61] J. Ryan Shue, Eric Ryan Chan, Ryan Po, Zachary Ankner, Jiajun Wu, and Gordon Wetzstein. 3D neural field generation using triplane diffusion. In Proc. CVPR, 2023. 1
[62] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Proc. ICML, 2015. 1
[63] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Proc. ICLR, 2021. 2
[64] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Proc. NeurIPS, 2019. 1
[65] YangSong,JaschaSohl-Dickstein,DiederikP.Kingma,Ab- hishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Proc. ICLR, 2021. 1, 2, 5
[66] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. JMLR, 15 (56), 2014. 5, 7
[67] MatthewTancik,PratulP.Srinivasan,BenMildenhall,Sara Fridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ra- mamoorthi, Jonathan T. Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimen- sional domains. In Proc. NeurIPS, 2020. 3
[68] Antti Tarvainen and Harri Valpola. Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results. In Proc. NIPS, 2017.
[53] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, 2 and Mark Chen. Hierarchical text-conditional image genera- tion with CLIP latents. CoRR, abs/2204.06125, 2022. 2
[69] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkor- eit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proc. NIPS, 2017. 1, 2
[70] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation, 23(7):1661–1674, 2011. 1
[71] LingYang,ZhilongZhang,YangSong,ShendaHong,Run-sheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming- Hsuan Yang. Diffusion models: A comprehensive survey of methods and applications. ACM Comput. Surv., 56(4), 2023. 1
[72] Yasin Yazıcı, Chuan-Sheng Foo, Stefan Winkler, Kim-Hui Yap, Georgios Piliouras, and Vijay Chandrasekhar. The un- usual effectiveness of averaging in GAN training. In Proc. ICLR, 2019. 2
[73] Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. CoRR, abs/1708.03888, 2017. 2, 4
[74] GuodongZhang,ChaoqiWang,BowenXu,andRogerGrosse. Three mechanisms of weight decay regularization. In Proc. ICLR, 2019. 2
[75] Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. In Proc. ICCV, 2023. 1