ScaleLong: Verso un Addestramento Più Stabile dei Modelli di Diffusione Tramite il Ridimensionamento delle Connessioni Long Skip della Rete

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

October 20, 2023

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

cs.AI

Abstract

Nei modelli di diffusione, UNet è l'architettura di rete più popolare, poiché le sue connessioni a salto lungo (LSC) che collegano blocchi di rete distanti possono aggregare informazioni a lunga distanza e alleviare il problema del gradiente scomparso. Sfortunatamente, UNet spesso soffre di instabilità durante l'addestramento nei modelli di diffusione, che può essere mitigata riducendo i coefficienti delle LSC. Tuttavia, una comprensione teorica dell'instabilità di UNet nei modelli di diffusione e del miglioramento delle prestazioni ottenuto ridimensionando le LSC è ancora assente. Per risolvere questo problema, dimostriamo teoricamente che i coefficienti delle LSC in UNet hanno un grande impatto sulla stabilità della propagazione in avanti e indietro e sulla robustezza di UNet. Nello specifico, la feature nascosta e il gradiente di UNet in qualsiasi livello possono oscillare, e i loro intervalli di oscillazione sono effettivamente ampi, il che spiega l'instabilità dell'addestramento di UNet. Inoltre, UNet è anche dimostrabilmente sensibile a input perturbati e predice un output lontano da quello desiderato, generando una perdita oscillatoria e quindi un gradiente oscillatorio. Oltre a ciò, osserviamo anche i benefici teorici del ridimensionamento dei coefficienti delle LSC di UNet nella stabilità delle feature nascoste e del gradiente, nonché nella robustezza. Infine, ispirati dalla nostra teoria, proponiamo un framework efficace di ridimensionamento dei coefficienti chiamato ScaleLong, che ridimensiona i coefficienti delle LSC in UNet e migliora ulteriormente la stabilità dell'addestramento di UNet. I risultati sperimentali su quattro famosi dataset mostrano che i nostri metodi sono superiori nel stabilizzare l'addestramento e producono un'accelerazione dell'addestramento di circa 1,5x su diversi modelli di diffusione con architetture UNet o UViT. Codice: 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