[論文レビュー] 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

ひとこと要約

本稿では、U-Netの長距離スキップ接続（LSC）の係数をスケーリングすることで、拡散モデルにおける学習の安定化を図るScaleLongというフレームワークを提案する。理論的には、大きな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$は不安定性リスクを示唆する。
• 振動範囲の低減と学習安定性の向上を目的として、LSCに適応的係数スケーリングを適用するScaleLongというフレームワークを提案した。
• CIFAR10、CelebA、ImageNet64の複数のデータセットで、特徴量可視化および勾配追跡を用いて理論的境界の実証的妥当性を検証した。
• ScaleLong内に、$1/\sqrt{2}$-スケーリングベースラインを上回る2つの新規かつ効果的な係数スケーリング手法を導入した。
• アブレーションスタディにより、$\kappa > 1$は特に小規模バッチサイズ下で不安定性を悪化させることを示し、理論的領域制限$\kappa \in (0,1)$の妥当性を裏付けた。

実験結果

リサーチクエスチョン

• RQ1 なぜU-Netは長距離スキップ接続を備えながらも、拡散モデルにおいて学習が不安定になるのか？
• RQ2 長距離スキップ接続の係数が、U-Netにおける前向きおよび後向き伝搬の安定性にどのように影響するか？
• RQ3 $1/\sqrt{2}$-スケーリングがLSC係数に与える安定化効果の理論的根拠は何か？
• RQ4 経験的ヒューリスティクスを凌駕する、原理的係数スケーリングフレームワークが学習の安定性と収束速度を向上させられるか？
• RQ5 U-Netの入力摂動に対するロバストネスは、LSC係数にどのように依存するか？

主な発見

• 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)$の妥当性が裏付けられた。