分析和改进扩散模型的训练动力学.
0. TL; DR
本文针对当前流行的扩散模型（如ADM）在训练过程中存在的不均衡和低效问题进行了深入分析，并提出了一系列改进方法。通过重新设计网络层以保持激活、权重和更新的预期大小，显著提升了模型性能，同时降低了计算复杂度。改进后的模型在ImageNet-512图像合成任务中取得了1.81的FID，刷新了记录。此外，文章还提出了一种后处理指数移动平均（EMA）参数的方法，允许在训练完成后精确调整EMA长度，揭示了其与网络架构、训练时间和引导的复杂交互。
1. 背景介绍
扩散模型在数据驱动的图像合成领域占据主导地位，尤其是在处理大规模数据集时表现出色。然而，其训练动力学仍然具有挑战性，主要是因为损失函数的高度随机性。在训练过程中，网络需要准确估计在不同噪声水平、高斯噪声实现和条件输入下的平均干净图像，这使得训练过程复杂且容易出错。此外，网络必须在极其嘈杂的训练信号中高效学习，而当前的扩散模型架构（如ADM）在训练动力学方面存在一些问题，例如网络激活和权重的无控制变化以及训练过程中的不平衡。
2. 改进扩散模型的动力学
本文的核心目标是通过重新设计网络层来解决训练动态中的不平衡问题。具体来说，作者提出了以下改进措施：
（1）Baseline模型
Baseline模型由 U-Net 和 Self-Attention 混合组成，在$512\times 512$的生成任务中输入噪声维度是$64\times 64\times 4$，Baseline 的 FID 为 8.00。
（2）改进Baseline
作者首先调整超参数 (学习率、EMA 长度、训练噪声水平分布等) 来优化 Baseline 模型的性能，并禁用了 $32\times 32$ 分辨率的 Self-Attention。
原始扩散损失权重在初始化时将所有噪声水平的损失幅度标准化为$1.0$，作者采用了 Uncertainty Loss 将原始损失值跟踪为噪声水平的函数，并通过其倒数缩放训练损失。
上述变化将 FID 从 8.00 降低到 7.24。
（3）结构流水线化 (Architectural streamlining)
为了避免处理多种不同类型的可训练参数，作者从所有卷积层和线性层以及调节路径中去除加性 bias。为了恢复网络偏移数据的能力，作者将常数 $1$ 的附加通道连接到网络的输入。作者把原始位置编码方案切换到更标准的傅里叶特征，并简化 Group Normalization 层。
在训练过程中，由于 Key 和 Query 向量的大小增长，注意力图通常表现出尖刺。作者使用余弦注意力机制在计算点积之前对向量进行归一化。这允许在整个网络中使用 $16$ 位浮点数，提高了整体的效率。
上述变化将 FID 从 7.24 降低到 6.96。
（4）幅值保持的学习层 (Magnitude-preserving learned layers)
作者观察到随着训练的进行，激活幅值会出现不可控的增长。这是由于网络中的残差结构包含较长的信号路径，且没有任何归一化。这些路径从残差分支累积，并且可以通过重复的卷积放大激活值。这种激活值幅度的增长会将整个模型置于非最佳状态下训练。
为了保留预期的激活幅度，作者将每一层的输出除以该层引起的激活幅度的预期缩放。为了恢复输入激活幅值，作者将每层的权重逐通道地L2归一化。由于整体权重大小不再对激活产生影响，因此使用单位高斯分布初始化所有权重。
上述变化将 FID 从 6.96 降低到 3.75。
（5）有效的学习率控制 (Controlling effective learning rate )
随着训练过程的进行，网络权重出现了明显的增长趋势。有效学习率 (即权重更新量的相对大小) 仍然随着训练的进行而衰减。文作者认为应该显式地控制它，而不是让它在层之间不可控和不均匀地漂移。
作者提出了一种 Forced Weight Normalization 技术，在每个训练步骤之前显式地将每个权重向量归一化为单位方差。同时在训练期间仍然在此之上应用 “标准” 的权重归一化。
作者还引入了逆平方根的学习率衰减策略：
\[\alpha(t) = \frac{\alpha_{ref}}{\sqrt{\max(t/t_{ref},1)}}\]上述变化将 FID 从 3.75 降低到 3.02。
（6）移除组归一化 (Removing group normalizations)
作者去除了 Group Normalization。尽管网络在没有任何归一化层的情况下可以成功训练，作者发现在 Encoder 的主路径引入更弱的 Pixel Normalization 层仍然有好处。作者还从 Embedding 网络中删除了第2个线性层和网络输出的非线性，并将残差块中的重采样操作合入主路径中。
上述变化将 FID 从 3.02 降低到 2.71。
（7）幅度保持的固定函数层 (Magnitude-preserving fixed-function layers)
傅里叶特征的正弦函数和余弦函数没有单位方差，作者通过将它们放大$\sqrt{2}$倍来纠正。SiLU非线性激活函数衰减了方差，因此作者将将输出除以$E_{x\sim N(0,1)}[silu(x)^2]^{1/2}\approx 0.596$来补偿。
此外每个分支之间可以通过可控制的参数来取得平衡，作者把加法操作换成加权求和。作者在整个模型的结尾增加了一个可学习的、零初始化的标量的增益，并在每个残差块内的 Condition 信号应用类似的增益。
上述变化将 FID 从 2.71 降低到 2.56。
（8）事后EMA (Post-hoc EMA)
指数移动平均 (Exponential Moving Average, EMA) 对网络权重采用$\theta_{t} = \beta \theta_{t-1} + (1-\beta)\theta_{t}$的更新方式， 使得早期训练步骤的贡献呈现出指数级衰减。
非常长的指数 EMA 对网络参数的初始阶段施加了不可忽略的权重，而初始阶段的参数常常是随机的。其次更长的训练运行受益于更长的 EMA 衰减。因此作者提出使用基于幂函数而不是指数函数的衰减:
\[\theta_{t} = \beta_\gamma(t) \theta_{t-1} + (1-\beta_\gamma(t))\theta_{t} \\ \beta_\gamma(t) = (1-1/t)^{\gamma+1}\]事后 EMA 的目标是为了在训练之后，可以灵活选择$\gamma$。为此，作者在训练的过程中维持了两套平均参数向量，对应的$\gamma$分别是 0.05 和 0.1。这些平均参数向量定期存储在训练期间保存的 Snapshots 中。在所有的实验中，作者每大约 800 万个训练图像存储一次快照，即每 4096 个训练步骤，Batch Size 为 2048。
为了得到在训练期间或之后任意一点的任何 EMA profile 的参数，作者找到了存储的参数和需求的 EMA profile 之间的最小二乘最优拟合，并进行相应线性组合。
图a显示了模型B-G版本对应的最优EMA长度变化，配置之间的最佳 EMA 长度差异很大，并且最终配置 G 的最优值变得更窄。图b显示了网络权重对EMA长度的敏感程度，作者从网络的不同部分选择权重张量的子集，每次只有所选张量的 EMA 发生了变化，而其他所有张量都保持在全局最优值，结果表明最终配置 G 的最优EMA长度更加稳定。图c显示了训练过程中最佳 EMA 长度的演变，随着训练的进行，最优值在缓慢向着相对较长的 EMA 长度移动。
3. 实验分析
上述一系列扩散模型动力学的优化技巧能够有效提高模型的性能：
作者使用 ImageNet 512×512 图像生成任务作为主要实验。之前最好的结果是 VDM++，FID 的值为 2.99。即使是使用小模型 EDM2-S 也取得了 2.56 的 FID，缩放模型尺寸之后可以进一步将 FID 提高到 1.91，大大超过了之前的记录。适度的 Guidance 可以进一步将 FID 提高到 1.81。
为了证明所提方法不局限于在隐空间中的扩散模型，作者提供了在 ImageNet-64 中的 RGB 空间的结果。如本文的结果优于一些使用确定性采样的早期方法。之前的记录是通过 EDM实现的，FID 为 2.22。本文在相似的计算复杂度下将其提升到了 1.58，进一步缩放之后达到 1.33。
