ScaleLong: Towards More Stable Training of Diffusion Model via Scaling Network Long Skip Connection

要約

拡散モデルでは、UNet が最も人気のあるネットワーク バックボーンです。これは、遠方のネットワーク ブロックを接続するロング スキップ接続 (LSC) により、長距離の情報を集約し、勾配消失を軽減できるためです。
残念ながら、UNet は拡散モデルのトレーニングが不安定になることがよくありますが、LSC 係数を小さくすることで緩和できます。
しかし、拡散モデルにおける UNet の不安定性や LSC スケーリングのパフォーマンス向上についての理論的理解はまだ得られていません。
この問題を解決するために、UNet の LSC の係数が順方向および逆方向伝播の安定性と UNet の堅牢性に大きな影響を与えることを理論的に示します。
具体的には、UNet の隠れた特徴と勾配はどの層でも振動する可能性があり、その振動範囲は実際には大きいため、UNet トレーニングの不安定性が説明されます。
さらに、UNet は摂動された入力にも敏感であることが証明されており、目的の出力から離れた出力を予測するため、振動損失が発生し、したがって振動勾配が発生します。
さらに、隠れた特徴と勾配の安定性、および堅牢性において、UNet の LSC 係数スケーリングの理論上の利点も観察されています。
最後に、私たちの理論に触発されて、UNet の LSC の係数をスケーリングし、UNet のトレーニングの安定性を向上させる効果的な係数スケーリング フレームワーク ScaleLong を提案します。
4 つの有名なデータセットに関する実験結果は、私たちの方法がトレーニングを安定させるのに優れており、UNet または UViT バックボーンを備えたさまざまな拡散モデルで約 1.5 倍のトレーニング加速をもたらすことを示しています。
コード: https://github.com/sail-sg/ScaleLong

要約(オリジナル)

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

arxiv情報

提供元, 利用サービス

arxiv.jp, Google