ChatPaper.aiChatPaper

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
PDF41December 15, 2024