ScaleLong : Vers un entraînement plus stable des modèles de diffusion via la mise à l'échelle des connexions longues du réseau

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

October 20, 2023

Auteurs: Zhongzhan Huang, Pan Zhou, Shuicheng Yan, Liang Lin

cs.AI

Résumé

Dans les modèles de diffusion, l'UNet est l'architecture de réseau la plus couramment utilisée, car ses connexions longues (LSCs) reliant des blocs distants du réseau permettent d'agréger des informations à longue distance et d'atténuer le problème du gradient qui disparaît. Malheureusement, l'UNet souffre souvent d'une instabilité lors de l'entraînement dans les modèles de diffusion, ce qui peut être atténué en réduisant les coefficients de ses LSCs. Cependant, une compréhension théorique de l'instabilité de l'UNet dans les modèles de diffusion, ainsi que de l'amélioration des performances apportée par la mise à l'échelle des LSCs, fait encore défaut. Pour résoudre ce problème, nous démontrons théoriquement que les coefficients des LSCs dans l'UNet ont un impact significatif sur la stabilité de la propagation avant et arrière, ainsi que sur la robustesse de l'UNet. Plus précisément, les caractéristiques cachées et le gradient de l'UNet à n'importe quelle couche peuvent osciller, et leurs amplitudes d'oscillation sont en réalité importantes, ce qui explique l'instabilité de l'entraînement de l'UNet. De plus, l'UNet est également théoriquement sensible aux entrées perturbées et prédit une sortie éloignée de la sortie souhaitée, entraînant une perte oscillatoire et donc un gradient oscillatoire. Par ailleurs, nous observons également les avantages théoriques de la mise à l'échelle des coefficients des LSCs de l'UNet en termes de stabilité des caractéristiques cachées et du gradient, ainsi que de robustesse. Enfin, inspirés par notre théorie, nous proposons un cadre efficace de mise à l'échelle des coefficients, appelé ScaleLong, qui ajuste les coefficients des LSCs dans l'UNet et améliore la stabilité de l'entraînement de l'UNet. Les résultats expérimentaux sur quatre ensembles de données célèbres montrent que nos méthodes sont supérieures pour stabiliser l'entraînement et permettent une accélération de l'entraînement d'environ 1,5 fois sur différents modèles de diffusion utilisant des architectures UNet ou UViT. Code : https://github.com/sail-sg/ScaleLong

English

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