[论文解读] ScaleLong: Towards More Stable Training of Diffusion Model via Scaling Network Long Skip Connection

Zhongzhan Huang, Pan Zhou+2 more|arXiv (Cornell University)|Oct 20, 2023

Advanced Neuroimaging Techniques and Applications被引用 5

一句话总结

本文提出 ScaleLong，一种通过缩放 U-Net 的长跳跃连接（LSC）系数来稳定扩散模型训练的框架。理论上，它表明较大的 LSC 系数会导致特征和梯度振荡；实证上，它证明系数缩放可减少振荡，并在多个数据集和架构上将训练速度提升最多 1.5 倍。

ABSTRACT

In diffusion models, UNet is the most popular network backbone, since its long skip connects (LSCs) to connect distant network blocks can aggregate long-distant information and alleviate vanishing gradient. Unfortunately, UNet often suffers from unstable training in diffusion models which can be alleviated by scaling its LSC coefficients smaller. However, theoretical understandings of the instability of UNet in diffusion models and also the performance improvement of LSC scaling remain absent yet. To solve this issue, we theoretically show that the coefficients of LSCs in UNet have big effects on the stableness of the forward and backward propagation and robustness of UNet. Specifically, the hidden feature and gradient of UNet at any layer can oscillate and their oscillation ranges are actually large which explains the instability of UNet training. Moreover, UNet is also provably sensitive to perturbed input, and predicts an output distant from the desired output, yielding oscillatory loss and thus oscillatory gradient. Besides, we also observe the theoretical benefits of the LSC coefficient scaling of UNet in the stableness of hidden features and gradient and also robustness. Finally, inspired by our theory, we propose an effective coefficient scaling framework ScaleLong that scales the coefficients of LSC in UNet and better improves the training stability of UNet. Experimental results on four famous datasets show that our methods are superior to stabilize training and yield about 1.5x training acceleration on different diffusion models with UNet or UViT backbones. Code: https://github.com/sail-sg/ScaleLong

研究动机与目标

• 理解为何尽管存在长跳跃连接，U-Net 在扩散模型中的训练仍不稳定。
• 从理论上解释 LSC 系数缩放（特别是 $1/\sqrt{2}$-缩放技术）的稳定效应。
• 开发一个原则性框架，以提升扩散模型中的训练稳定性和收敛速度。
• 通过实证验证，缩放 LSC 系数可减少特征和梯度振荡，并加速训练。

提出的方法

• 理论上推导了 U-Net 中隐藏特征和梯度范数的边界，表明振荡范围与 $\mathcal{O}(m\sum_{i=1}^{N}\kappa_i^2)$ 成正比，其中 $\kappa_i$ 为 LSC 系数。
• 证明了 U-Net 对输入扰动的鲁棒性受 $\mathcal{O}(\sum_{i=1}^{N}\kappa_i M_0^i)$ 限制，其中 $M_0 > 1$ 表明存在不稳定性风险。
• 提出 ScaleLong，一种对 LSC 应用自适应系数缩放的框架，以减小振荡范围并提升训练稳定性。
• 通过在 CIFAR10、CelebA 和 ImageNet64 上的特征可视化与梯度追踪，实证验证了理论边界。
• 在 ScaleLong 内引入两种新颖且有效的系数缩放方法，其性能优于 $1/\sqrt{2}$-缩放基线。
• 通过消融研究表明，$\kappa > 1$ 会加剧不稳定性，尤其在小批量训练时，而 $\kappa \in (0,1)$ 可确保稳定性。

实验结果

研究问题

• RQ1 为何 U-Net 在扩散模型中尽管具有长跳跃连接，仍表现出训练不稳定？
• RQ2 长跳跃连接的系数如何影响 U-Net 中前向和反向传播的稳定性？
• RQ3 对 LSC 系数采用 $1/\sqrt{2}$-缩放具有稳定效应的理论基础是什么？
• RQ4 能否设计一个原则性系数缩放框架，使训练稳定性和收敛速度超越经验启发式方法？
• RQ5 U-Net 对输入扰动的鲁棒性如何依赖于 LSC 系数？

主要发现

• 当 $\kappa_i = 1$ 时，U-Net 中隐藏特征和梯度的振荡范围为 $\mathcal{O}(mN)$，这解释了训练过程中观察到的不稳定性。
• 将 LSC 系数缩放至 $1/\sqrt{2}$ 可减小振荡范围并提升训练稳定性，但无法完全消除不稳定性。
• 理论边界表明，梯度幅值受 $\mathcal{O}(m\sum_{i=1}^{N}\kappa_i^2)$ 上限约束，当 $\kappa_i = 1$ 时该值较大，导致参数更新不稳定。
• U-Net 对输入扰动的鲁棒性受 $\mathcal{O}(\sum_{i=1}^{N}\kappa_i M_0^i)$ 限制，其中 $M_0 > 1$，表明对微小输入变化敏感。
• ScaleLong 减少了特征和梯度振荡，使多种具有 UNet 或 UViT 主干的扩散模型训练速度提升最多 1.5 倍。
• 实验结果证实，$\kappa > 1$ 会导致训练快速崩溃，尤其在小批量设置下，验证了理论中对 $\kappa \in (0,1)$ 的限制域。