ScaleLong: Rumo a um Treinamento Mais Estável de Modelos de Difusão por meio da Escalonagem de Conexões de Salto Longo na Rede

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

October 20, 2023

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

cs.AI

Resumo

Nos modelos de difusão, a UNet é a arquitetura de rede mais popular, pois suas conexões de salto longo (LSCs, do inglês Long Skip Connects) que ligam blocos distantes da rede podem agregar informações de longa distância e aliviar o problema do gradiente que desaparece. Infelizmente, a UNet frequentemente sofre com instabilidade durante o treinamento em modelos de difusão, o que pode ser mitigado ao reduzir os coeficientes das LSCs. No entanto, ainda faltam entendimentos teóricos sobre a instabilidade da UNet em modelos de difusão e também sobre a melhoria de desempenho proporcionada pelo ajuste dos coeficientes das LSCs. Para resolver esse problema, demonstramos teoricamente que os coeficientes das LSCs na UNet têm um grande impacto na estabilidade da propagação direta e reversa, bem como na robustez da UNet. Especificamente, as características ocultas e o gradiente da UNet em qualquer camada podem oscilar, e suas faixas de oscilação são, na verdade, amplas, o que explica a instabilidade do treinamento da UNet. Além disso, a UNet também é comprovadamente sensível a entradas perturbadas e prevê uma saída distante da saída desejada, resultando em perda oscilatória e, consequentemente, em gradiente oscilatório. Além disso, também observamos os benefícios teóricos do ajuste dos coeficientes das LSCs da UNet na estabilidade das características ocultas, do gradiente e também na robustez. Por fim, inspirados por nossa teoria, propomos uma estrutura eficaz de ajuste de coeficientes chamada ScaleLong, que ajusta os coeficientes das LSCs na UNet e melhora a estabilidade do treinamento da UNet. Resultados experimentais em quatro conjuntos de dados famosos mostram que nossos métodos são superiores para estabilizar o treinamento e proporcionam uma aceleração de cerca de 1,5x no treinamento em diferentes modelos de difusão com backbones UNet ou UViT. Código: 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