ScaleLong: К более стабильному обучению моделей диффузии через масштабирование длинных пропущенных соединений в сети

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

October 20, 2023

Авторы: Zhongzhan Huang, Pan Zhou, Shuicheng Yan, Liang Lin

cs.AI

Аннотация

В диффузионных моделях UNet является наиболее популярной архитектурой сети, поскольку её длинные пропускные соединения (LSCs), связывающие удалённые блоки сети, позволяют агрегировать информацию на больших расстояниях и смягчать проблему исчезающего градиента. К сожалению, UNet часто сталкивается с нестабильностью обучения в диффузионных моделях, что можно частично устранить, уменьшая коэффициенты LSCs. Однако теоретическое понимание причин нестабильности UNet в диффузионных моделях, а также улучшения производительности за счёт масштабирования LSCs до сих пор отсутствует. Чтобы решить эту проблему, мы теоретически показываем, что коэффициенты LSCs в UNet оказывают значительное влияние на стабильность прямого и обратного распространения, а также на устойчивость сети. В частности, скрытые признаки и градиенты UNet на любом слое могут колебаться, причём диапазоны этих колебаний оказываются значительными, что объясняет нестабильность обучения UNet. Более того, UNet также теоретически чувствительна к возмущённым входным данным и предсказывает выход, далёкий от желаемого, что приводит к колебаниям функции потерь и, следовательно, градиентов. Кроме того, мы также наблюдаем теоретические преимущества масштабирования коэффициентов LSCs в UNet для стабильности скрытых признаков, градиентов и устойчивости сети. Вдохновлённые нашей теорией, мы предлагаем эффективную структуру масштабирования коэффициентов ScaleLong, которая масштабирует коэффициенты LSCs в UNet и значительно улучшает стабильность обучения. Экспериментальные результаты на четырёх известных наборах данных показывают, что наш метод превосходит другие подходы в стабилизации обучения и обеспечивает ускорение обучения примерно в 1,5 раза на различных диффузионных моделях с архитектурами UNet или UViT. Код: 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