ScaleLong: Hin zu stabilerem Training von Diffusionsmodellen durch Skalierung der langen Sprungverbindungen im Netzwerk
ScaleLong: Towards More Stable Training of Diffusion Model via Scaling Network Long Skip Connection
October 20, 2023
Autoren: Zhongzhan Huang, Pan Zhou, Shuicheng Yan, Liang Lin
cs.AI
Zusammenfassung
In Diffusionsmodellen ist UNet das beliebteste Netzwerk-Backbone, da seine langen Skip-Connections (LSCs), die entfernte Netzwerkblöcke verbinden, weit entfernte Informationen aggregieren und das Verschwinden von Gradienten mildern können. Leider leidet UNet oft unter instabilem Training in Diffusionsmodellen, was durch die Skalierung seiner LSC-Koeffizienten verringert werden kann. Allerdings fehlen bisher theoretische Erkenntnisse über die Instabilität von UNet in Diffusionsmodellen und auch über die Leistungsverbesserung durch die LSC-Skalierung. Um dieses Problem zu lösen, zeigen wir theoretisch, dass die Koeffizienten der LSCs in UNet einen großen Einfluss auf die Stabilität der Vorwärts- und Rückwärtspropagation sowie auf die Robustheit von UNet haben. Insbesondere können die verborgenen Merkmale und der Gradient von UNet in jeder Schicht oszillieren, und ihre Oszillationsbereiche sind tatsächlich groß, was die Instabilität des UNet-Trainings erklärt. Darüber hinaus ist UNet nachweislich auch empfindlich gegenüber gestörten Eingaben und sagt eine Ausgabe voraus, die weit von der gewünschten Ausgabe entfernt ist, was zu oszillierenden Verlusten und somit oszillierenden Gradienten führt. Zusätzlich beobachten wir die theoretischen Vorteile der LSC-Koeffizienten-Skalierung von UNet in Bezug auf die Stabilität der verborgenen Merkmale und des Gradienten sowie auf die Robustheit. Schließlich schlagen wir, inspiriert durch unsere Theorie, ein effektives Koeffizienten-Skalierungsframework namens ScaleLong vor, das die Koeffizienten der LSCs in UNet skaliert und die Trainingsstabilität von UNet besser verbessert. Experimentelle Ergebnisse auf vier bekannten Datensätzen zeigen, dass unsere Methoden überlegen sind, um das Training zu stabilisieren, und etwa eine 1,5-fache Trainingsbeschleunigung bei verschiedenen Diffusionsmodellen mit UNet- oder UViT-Backbones erzielen. 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