ScaleLong：通过缩放网络长跳跃连接实现扩散模型更稳定的训练
计算机视觉与模式识别
2023-10-23 v1 人工智能
摘要
在扩散模型中，UNet 是最流行的网络骨干，因其长跳跃连接（LSCs）可连接远端网络块，聚合远距离信息并缓解梯度消失。遗憾的是，UNet 在扩散模型中常遭遇训练不稳定，可通过将其 LSC 系数缩小来缓解这个问题。然而，关于 UNet 在扩散模型中不稳定的理论解释以及 LSC 缩放带来的性能提升，目前仍属空白。为解决该问题，我们从理论上证明 UNet 中 LSCs 的系数对前向与反向传播的稳定性和 UNet 的鲁棒性有重大影响。具体而言，UNet 任意层的隐藏特征与梯度可能发生振荡，且其振荡范围实际上很大，这解释了 UNet 训练的不稳定性。此外，UNet 被证明对扰动输入敏感，并预测出远离期望输出的结果，从而产生振荡损失进而振荡梯度。另外，我们也观察到 UNet 的 LSC 系数缩放在隐藏特征与梯度稳定性以及鲁棒性方面的理论益处。最后，受我们的理论启发，我们提出一种有效的系数缩放框架 ScaleLong，其对 UNet 中 LSC 的系数进行缩放，更好地改善了 UNet 的训练稳定性。在四个著名数据集上的实验结果表明，我们的方法在稳定训练方面更优，并在采用 UNet 或 UViT 骨干的不同扩散模型上带来约 1.5 倍的训练加速。代码：https://github.com/sail-sg/ScaleLong
引用
@article{arxiv.2310.13545,
title = {ScaleLong: Towards More Stable Training of Diffusion Model via Scaling Network Long Skip Connection},
author = {Zhongzhan Huang and Pan Zhou and Shuicheng Yan and Liang Lin},
journal= {arXiv preprint arXiv:2310.13545},
year = {2023}
}
备注
accepted by NeurIPS 2023