ScaleLong: Hacia un Entrenamiento Más Estable de Modelos de Difusión mediante la Escalación de Conexiones de Salto Largo en la Red

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

Resumen

En los modelos de difusión, UNet es la arquitectura de red más popular, ya que sus conexiones de salto largas (LSCs, por sus siglas en inglés) que conectan bloques distantes de la red pueden agregar información de largo alcance y mitigar el problema del gradiente que desaparece. Desafortunadamente, UNet suele presentar inestabilidad durante el entrenamiento en modelos de difusión, lo cual puede aliviarse escalando los coeficientes de sus LSCs a valores más pequeños. Sin embargo, aún faltan comprensiones teóricas sobre la inestabilidad de UNet en modelos de difusión y también sobre la mejora de rendimiento que ofrece el escalado de LSCs. Para abordar este problema, demostramos teóricamente que los coeficientes de las LSCs en UNet tienen un gran impacto en la estabilidad de la propagación hacia adelante y hacia atrás, así como en la robustez de UNet. Específicamente, las características ocultas y el gradiente de UNet en cualquier capa pueden oscilar, y sus rangos de oscilación son en realidad amplios, lo que explica la inestabilidad del entrenamiento de UNet. Además, UNet también es teóricamente sensible a entradas perturbadas y predice una salida distante de la deseada, lo que genera una pérdida oscilatoria y, por ende, un gradiente oscilatorio. Asimismo, observamos los beneficios teóricos del escalado de coeficientes de LSCs en UNet en términos de la estabilidad de las características ocultas, el gradiente y también la robustez. Finalmente, inspirados por nuestra teoría, proponemos un marco efectivo de escalado de coeficientes llamado ScaleLong, que ajusta los coeficientes de las LSCs en UNet y mejora significativamente la estabilidad del entrenamiento de UNet. Los resultados experimentales en cuatro conjuntos de datos famosos muestran que nuestros métodos son superiores para estabilizar el entrenamiento y logran una aceleración del entrenamiento de aproximadamente 1.5x en diferentes modelos de difusión con arquitecturas UNet o 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