ScaleLong: Naar een stabielere training van diffusiemodellen via schaling van netwerk-long-skip-verbindingen

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

Samenvatting

In diffusiemodellen is UNet de meest populaire netwerkbackbone, omdat de lange sprongverbindingen (LSCs) die verre netwerkblokken verbinden, informatie over lange afstand kunnen samenvoegen en het verdwijnen van gradiënten kunnen verminderen. Helaas heeft UNet vaak last van onstabiele training in diffusiemodellen, wat kan worden verlicht door de LSC-coëfficiënten kleiner te schalen. Echter, theoretische inzichten in de instabiliteit van UNet in diffusiemodellen en ook de prestatieverbetering door LSC-schaling ontbreken nog. Om dit probleem op te lossen, tonen we theoretisch aan dat de coëfficiënten van LSCs in UNet een grote invloed hebben op de stabiliteit van de voorwaartse en achterwaartse propagatie en de robuustheid van UNet. Specifiek kunnen de verborgen kenmerken en gradiënt van UNet in elke laag oscilleren en hun oscillatiebereiken zijn daadwerkelijk groot, wat de instabiliteit van UNet-training verklaart. Bovendien is UNet ook bewezen gevoelig voor verstoorde invoer en voorspelt het een uitvoer die ver verwijderd is van de gewenste uitvoer, wat resulteert in oscillerend verlies en dus oscillerende gradiënt. Daarnaast observeren we ook de theoretische voordelen van de LSC-coëfficiëntschaling van UNet in de stabiliteit van verborgen kenmerken en gradiënt en ook robuustheid. Ten slotte, geïnspireerd door onze theorie, stellen we een effectief coëfficiëntschalingsframework ScaleLong voor dat de coëfficiënten van LSC in UNet schaalt en de trainingsstabiliteit van UNet beter verbetert. Experimentele resultaten op vier bekende datasets tonen aan dat onze methoden superieur zijn in het stabiliseren van training en ongeveer 1,5x trainingsversnelling opleveren op verschillende diffusiemodellen met UNet- of UViT-backbones. 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