ScaleLong: Towards More Stable Training of Diffusion Model via Scaling
  Network Long Skip Connection
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
6535f8519ded17e619eb721e
Zhongzhan Huang
6535f8519ded17e619eb721f
Pan Zhou
6535f8519ded17e619eb7220
Shuicheng Yan
extracted_pending
2025-03-31T14:14:20.059Z
67eaa070b9fa8908e151fd7d
/avatars/1fe2fd678d2e71099a83a9bcb9ab517e.svg
shuicheng yan
user
shuicheng
6535f8519ded17e619eb7221
Liang Lin
ScaleLong: Towards More Stable Training of Diffusion Model via Scaling
  Network Long Skip Connection
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
The paper analytically examines the instability of UNet in diffusion models due to long skip connections and proposes a coefficient scaling framework to enhance training stability and speed.
diffusion models UNet long skip connects vanishing gradient forward and backward propagation robustness hidden feature gradient oscillation perturbed input UViT ScaleLong training stability