ScaleLong: Towards More Stable Training of Diffusion Model via Scaling Network Long Skip Connection
Authors: Zhongzhan Huang
Pan Zhou
Shuicheng YAN
Liang Lin
Venue: NeurIPS 2023 poster (2023)
Keywords: Diffusion Model
Stable Training
Network architectures

Abstract: 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 improve the  training stability of UNet. Experimental results on CIFAR10, CelebA, ImageNet and COCO 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.

ScaleLong: Towards More Stable Training of Diffusion Model via Scaling Network Long Skip Connection
Zhongzhan Huang 1.2
Pan Zhou Shuicheng Yan Liang Lin'*
Sun Yat-Sen University, Sea Al Lab
huangzhah230ma112..ygg sysu. edu. cn; zhoupan@sea. com; shui chuicheng-ro yan@gmail com; linliangdieee.oo org
Abstract
In diffusion models, UNet 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 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 UNet the stableness 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 improve 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 different diffusion models with UNet or UVIT backbones. Click here for Code.
Introduction
Recently, diffusion models (DMs) [1-9] have become the most successful generative models because of their superiority modeling real- istic data distributions. The methodology of DMs includes forward diffusion process and corresponding reverse diffusion process. For for- ward process, it progressively injects Gaussian noise into realistic sample until the sample becomes Gaussian noise, while the reverse process trains a neural network to denoise the injected noise at each step for gradually map- ping the Gaussian noise into the vanilla sample. By decomposing the complex generative task into sequential application of denoising small noise, DMs achieve much better synthesis performance than other generative models, e.g., generative adversarial networks [10] and variational autoencoders []]], on image [12-15]. 3D [16, 17] and video [18-21] data and beyond.
X00000
Depth
X X10000 Depth
7.0
Org
2.0
feature
1/2
1/22-C3
I.S (ours)
5.0
LS (ours)
CS (ours)
CS (ours)
L.0.
0.5
20000
20000
30000
6
10000
20000
30000
lterations
Iterations
Figure 1: The feature oscillation in UNet.
"Coresponning author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Motivation, Since the reverse diffusion process in DMs essentially addresses a denoising task, most DMs follow previous image denoising and restoration works [22-26] to employ UNet as their de-facto backbone. This is because UNet uses long skip connects (LSCss) to connect the distant and symmetrical network blocks residual network, which helps long-distant information aggregation and alleviates the vanishing gradient issue, yielding superior image denoising and restoration performance. However, in DMs, its reverse process uses a shared UNet to predict the injected noise at any step, even though the input noisy data have time-varied distribution due to the progressively injected noise. This inevitably leads to the training difficulty and instability of UNet in DMs. Indeed, as shown in Fig. 1, the output of hidden layers in UNet, especially for deep UNet, oscillates rapidly along with training iterations. This feature oscillation also indicates that there are some layers whose parameters suffer from oscillations which often impairs the parameter learning speed. Moreover, the unstable hidden features also result in an unstable input for subsequent layers and greatly increase their learning difficulty. Some works [27-31] I] empirically find that 1/ 1/v2- constant scaling (1/v2-CS), that scales the coefficients of LSC (i.e. {x,} } in Fig. 2) from one used in UNet to 1/v2, alleviates the oscillations to some extent as shown in Fig. 1, which, however, lacks intuitive and theoretical explanation and still suffers from feature and gradient oscillations. So it is unclear yet 1) why UNet in DMs is unstable and also 2) why scaling the coefficients of LSC in UNet helps stabilize training, which hinders the designing new and more advanced DMs in principle way.
Contribution. In this work, we address the above two funda- mental questions and contribute to deriving some new results and alternatives for UNet. Particularly, we theoretically analyze the above training instability of UÑet in DMs, and identify the key role of the LSC coefficients in UNet for its unstable training, interpreting the stabilizing effect and limitations of the 1/ 1//2-scaling technique. Inspired by our theory, we propose the framework ScaleLong including two novel-yet-effectiv coefficient scaling methods which scale the coefficients of LSC in UNet and better improve the training stability of UNet. Our main contributions are highlighted below.
Input
Output
(1
02
b2
KN
by
fx
Long skip connection
Forward
Block
Scaling
Our first contribution is proving that the coefficients of LSCs in UNet have big effects on the stableness of the forward and backward propagation and robustness of UNet. Formally, for forward propagation, we show that the norm of the hidden feature at any layer can be lower- and also upper-bounded, and the oscillation range between lower and upper bounds is of the order (mm) En Ci=1- m.), where {N²} denotes the scaling coefficients of N LSCs and the input dimension m of UNet is often of hundreds. Since standard UNet sets K2- (Vi), the above oscillation range becomes O (m.V and is large, which partly explains the oscillation and instability of hidden features as observed in Fig. 1. The 1/V2-CS technique can slightly reduce the oscillation range, and thus helps stabilize UNet. Similarly, for backward propagation, the gradient magnitude in UNet is upper bounded by 0(m E In standard UNet, this bound becomes a large bound (mNN) and yields the possible big incorrect parameter update, impairing the parameter learning speed during training. This also explains the big oscillation of hidden features, and reveals the stabilizing effects of the 1/V2-C3 technique. Furthermore, for robustness that measures the prediction change when adding a small perturbation into an input, we show robustness bound of UNet is O Ki MM) with a model-related constant Mo 1. This result also implies a big robustness bound of standard UNet, and thus shows the sensitivity of UNet to the noise which can give a big incorrect gradient and explain the unstable training. It also shows that 1/V2-C5 technique improves the robustness.
coefficients
Figure 2: The diagram of UNet.
Inspired by our theory, we further propose novel framework ScaleLong including two scaling methods to adjust the coefficients of the LSCs in UNet for stabilizing the training: 1) constant scaling method (CS for short) and 2) Icarnable scaling method (LS). CS sets the coefficients {K,} serir of aaponentially-dccayee
constants, i..., K; = k' with a contant K e (0, 1]. As a result, CS greatly stabilizes the UNettraining by largely reducing the robustness error bound from (M,)) to (Mo(()))) N.)). which is also (K N-2) times smaller than the bound O(xM,)) of the universal scaling method, i.c., K; = k. like 1/ 1/V2-CS technique. Meanwhile, CS can ensure the information transmission of LSC without degrading into feedforward network. Similarly, the oscillation range

---

of hidden features and also gradient magnitude can also be controlled by our CS. For LS, it uses tiny shared network for all LSCs to predict the scaling coefficients for each LSC. In this way, LS is more flexible and adaptive than the CS method, it can learn the scaling coefficients according to the training data, and can also adjust the scaling coefficients along with training iterations which could benefit the training in terms of stability and convergence speed. Fig. 1 shows the effects of CS and LS in stabilizing the UNet training in DMs.
Extensive experiments on CIFARIO [32], CelebA [33], ImageNet [34], and COCO [35], show the effectiveness of our CS and LS methods in enhancing training stability and thus accelerating the learning speed by at least 1.5x in most cases across several DMs, e.g. UNet and UViT networks under the unconditional [3-5], class-conditional [36, 37] and text-to-image [38 42] settings.
2 Preliminaries and other related works
Diffusion Model (DM). DDPM-alike DMs [1-6, 8, 9, 43] generates a sequence of noisy samples
{x by repeatedly adding Gaussian noise to a sample xo until attatining xT ~ (0, 1). This noise injection process, a.k.a. the forward process, can be formalized as a Markov chain g (x17) (xo) (xxllaax B,I), a and Be depend on and satisfy ag + 6g- = I. By using the properties of Gaussian distribution, the forward process can be written as
II q (x+ -1),,,, where (xxe(xx-1)
(x xo) = (xx;;; X, (1 -1 ,,).),
where o₂ -I ag. Next, one can sample x X + v1 â where e, ~N I). Then DM adopts neural network éo (-, t) to invert the forward process and predict the noise added at each time step. This process to recover data Xg from Gaussian noise by minimizing the loss
C Cimpple ()) Elle &o(vāx + t
(2)
UNet-alike Network. Since the network Es(., t) in DMs predicts the noise to denoise, it plays a similar role of UNet-alike network widely in image restoration tasks, e.g. image de-raining, image denoising [22-26, 44, 45]. This inspires DMs to use UNet-alike network (3) that uses LSCs to connect distant parts of the network for long-range information transmission and aggregation
UNet(2) (z) = fo(x), fx()) = bi+1 o [8:+1 a++1 x + fi+1 (a'+1 o x)], 0 O N (3)
where E I.Th denotes the input, a, and by (i 2 1) are the trainable parameter of the i-th block, K; > 0 (i 2 1) are the scaling coefficients and are set to standard UNet. JN is the middle block of UNet. For the vector operation it can be designed to implement different networks.
W.l.o.g, in this paper. we consider as matrix multiplication, and set the i-th block as a stacked network [46, 47] which is implemented as a; = WT Φ( (W= (W(+))) with ReLU activation function & and learnable matrices W E Rmxm (j 2 1). Moreover, let f' also have the same architecture, i.c. (x) w of (W/ (.....(((f Following [36], the above UNet has absorbed the fusion operation of the feature 0g x and Si (a, o x) into UNet backbone. So for simplicity, we analyze the UNet in (3). Moreover, as there are many variants of UNet, e.g. transformer UNet, it is hard to provide a unified analysis for all of them. Thus, here we only analyze the most widely used UNet in (3), and leave the analysis of other UNet variants in future work.
Other related works. Previous works [46, 48-55] mainly studied classification task and observed performance improvement of scaling block output instead of skip connections in ResNet and Trans- former, i.e. X,+1 = X; + fx,,),, where fx()) is the i-th block. In contrast, we study UNet for DMs which uses long skip connections (LSC) instead of short skip connections, and stabilize UNet training via scaling LSC via Eq. (3) which is superior to scaling block output as shown in Section 5.
3 Stability Analysis on UNet
As shown in Section 1, standard U-Net (K; 1, Vi) often suffers from the unstable training issue in diffusion models (DMs), while scaling the long skip connection (LSC) coefficients kę(vi) in UNet can help to stabilize the training. However, theoretical understandings and also the effects of scaling LSC coeffeicents remain absent yet, hindering the development of new and more advanced DMs principle way. To address this problem, in this section, we will theoretically and comprehensively analyze 1) why UNet in DMs is unstable and also 2) why scaling
the instability of U-Net
3

---

(a)
The distribution of ||x,1|
(b)
The distribution ||c/|///1 x
Extra noise d3 0.01
Frequency Frequency
200
60
distribution Xg
CIFARIO
Frequency
30.
CIFARIO
FID-50K
C3 (uun) Org UNT
(ous)
1.6x
100
20
1.0
0.
20.
10.
2.49%
1.5 x
1.4
300
distribution Xg
CelebA
Frequency
30.
CelebA
100
200-
20
20
0.93%
wib
Extra noise 0g 0.02
noise
3.1x
15-
LS
Frequency
300
200
distribation
ImageNet
Frequency
20-
ImageNot
FID 50k]
10-
CS 2.9x
UVi
2.6x
20
100-
20
10-
2.17%
0.2m
0.4m
9.6m
0.8m
1.0m
Values
Values
wio extra noise
W/ extra noise
Figure 3: Distribution of Xo.
||x
12 and
1x and effect of extra noise to performance. ||<.1///
the coeffeicents of LSC in UNet helps stablize training. To this end, in the following, we will analyze the stableness of the forward and backward propagation of UNet, and investigate the robustness of UNet to the noisy input. All these analyses not only deepen the understanding of the instability of UNet, but also inspire an effective solution for more stable UNet as introduced in Section 4.
3.1 Stability of Forward Propagation
Here we study the stableness of the hidden features UNet. This is very important, since for certain layer, if its features vary significantly along with the training iterations, it means there is at least layer whose parameters suffer from oscillations which often impairs the parameter learning speed. Moreover, the unstable hidden features also result in unstable input for subsequent layers and greatly increase their learning difficulty, since a network cannot easily learn an unstable distribution, e.g. Internal Covariate Shift [56, 57], which is also shown in many works [58-60]. In the following. we theoretically characterize the hidden feature output by bounding its norm.
Theorem 3.1. Assume that all learnable matrices of UNet in Eq. (3) are independently initialized
as Kaiming's 's initialization, ie.., (0, 2). Then for any p (0, 1). by minimizing the training loss Eq. (2) of DMs, with probability at least O(N) expl exp[- (SQmmp²) we have
N R3 (1 1-D))(cilxx1 Σn x3+ oz S 141l (1+ s1++
N Σ,-
c [cl1xx118
00
(4)
VVixo + VI- Eees is the where the hidden feature h; is the output of fi defined in Eq. (3), =;= input of UNet; ci and C2 are two constants; K; is the scaling coefficient of the i-th LSC in UNet.
Theorem 3.1 reveals that with high probability, the norm of hidden features h; of UNet can be both
lower- and also hpper-boundedd and its biggest oscillation range is of the order [1]xe]] 11 Σ j=i x3)
which is decided by the scaling coefficients {K;} and the UNet input x So when ||x; [|2 or Ei=1 K? large, the above oscillation range is large which allows hidden feature h2 to oscillate largely as shown in Fig. 1, leading to the instability of the forward processing in UNet.
In fact, the following Lemma 3.2 shows that ||x+|] 1I3 is around feature dimension m in DDPM.
Lemma 3.2. For x; Vā,xo vI agE; defined in Eq. (1) as input of UNet, e, (O, if Xo follow the uniform distribution U U[- 1, 1]. then we have
E x 12 = (1 2E, ///33mmmmmmmmmmmm3 = (()).
(5)
Indeed, the conditions and conclusions of Lemma 3.2 hold universally in commonly used datasets, such CIFAR10, CelebA, and ImageNet, which can be empirically valid in Fig. 3 (a). In this way, in standard UNet whose scaling coefficients satisfy K; = 1, the oscillation range of hidden feature ||ball becomes a very large bound O(Nm), where N denotes the number of LSCCs. Accordingly, UNet suffers from the unstable forward process caused by the possible big oscillation of the hidden feature h which accords with the observations in Fig. 1. To relieve the instability issue, one possible

---

solution is to scale the coefficients } of LSCs, which could reduce the oscillation range of hidden feature| ||h,112 and yield more stable hidden features h; to benefit the subsequent layers. This is also one reason why 1/ /2sscaling technique [27-31] can help stabilize the UNet. However, it is also important to note that these coefficients cannot be too small, as they could degenerate a network into a feedforward network and negatively affect the network representation learning ability.
3.2 Stability of Backward Propagation
As mentioned in Section 3.1, the reason behind the instability of UNet forward process in DMs is the parameter oscillation. Here we go step further to study why these UNet parameters oscillate. Since parameters are updated by the gradient given in the backward propagation process, we analyze the gradient of UNet in DMs. A stable gradient with proper magnitude can update model parameters smoothly and also effectively, and yields more stable training. On the contrary, if gradients oscillate largely as shown in Fig. 1, they greatly affect the parameter learning and lead the unstable hidden features which further increases the learning difficulty of subsequent layers. In the following, we analyze the gradients of UNet, and investigate the effects of LSC coefficient scaling on gradients.
For brivity, we use W to denote all the learnable matrices of the i-th block (see Eq. (3)), i... W = [WT], Wi--...., WT W Wh whi). Assume the training loss is 1 Ex= (W), where Cg (W) denotes the training loss of the s-th sample among the n training samples. Next, we analyze the gradient W of each training sample, and summarize the main results in Theorem 3.3, in which we omit the subseript & of V&,(W) for brevity. See the proof in Appendix.
Theorem 3.3. Assume that all learnable matrices of UNet in Eq. (3) are independently initialized as Kaiming's 's initialization, i.e., NN, ). Then for any E (0, 1]. with probability at least 1 (nNN) expl- (m)m ), for a sample Xt, in training set, we have
(((W)
N ", ca)
IV (W) S
E (1,2, {ai, be);; (6)
where x, X0 + v v denotes the noisy sample of the s-th sample, G. ~ N(O,I). N the number of LSCs, C3 is a small constant.
Theorem 3.3 reveals that for any training sample in the training dataset, the gradient of any trainable xÿ) Since Lemma 3.2 shows Xx 112 is at the order of (m) in expectation, the bound of gradient norm becomes (m,, (W) Σj= xÿ). Consider the feature dimension m is often hundreds and the loss ((W) would be large at the beginning of the training phase, if the number N of LSCs is relatively large, the gradient bound in standard UNet (K; = 1 vi) would become O ( (W)) and is large. This means that UNet has the risk of instability gradient during training which is actually observed in Fig. 1. To address this issue, similar to the analysis of Theorem 3.3, we can appropriately scale the coefficients {F²} ensure that the model has enough representational power while preventing too big a gradient and avoid unstable UNet parameter updates. This also explains why the 1 /2-scaling technique can improve the stability of UNet in practice.
parameter W p in UNet can be bounded by (e(W)- ||x E ,, (up to the constant term).
3.3 Robustness On Noisy Input
In Section 3.1 and 3.2, we have revealed the reasons of instability of UNet and also the benign effects of scaling coefficients of LSCCs in UNet during DM training. Here we will further discuss the stability of UNet in terms of the robustness on noisy input. A network is said to be robust if its output does not change much when adding a small noise perturbation to its input. This robustness is important for stable training of UNet in DMs. As shown in Section 2, DMs aim at using UNet to predict the noise & ~ N (0, I) in the noisy sample for denoising at each time step t. However, in practice, additional unknown noise caused by random minibatch, data collection, and preprocessing strategies is inevitably introduced into noise 6, during the training process, and greatly affects the training. is because if a network is sensitive to noisy input, its output would be distant from the desired output, and yields oscillatory loss which leads to oscillatory gradient and parameters. In contrast, a robust network would have stable output, loss and gradient which boosts the parameter learning speed.
Here we use a practical example to show the importance of the robustness of UNet in DMs. According to Eq. (2), DM aims to predict noise e, from the UNet input x, (xxx / Vārxo, (1 ā))
5

---

Assume an extra Gaussian noise Es ~ o30 is injected into x, which yields a new input x ~ N (xx: Va,Xo, [11 - .it) + =.]l). Accordingly, if the variance og of extra noise is large, it can dominate the noise et in x,, and thus hinders predicting the desired noise et. In practice, the variance of this extra noise could vary along training iterations, further exacerbating the instability of UNet training. Indeed, Fig. 3 (c) shows that the extra noise Eς with different o can significantly affect the performance of standard UNet for DMs, and the scaled LSCs can alleviate the impact of extra noise to some extent. We then present the following theorem to quantitatively analyze these observations.
Theorem 3.4. For UNet in Eq. (3), assume Mo = max{llb; aill2: S i S ] and fx is Lo- Lipschitz continuous. co a constant related to Mo and Lo- Suppose xi" an perturbated input of the vanilla input X₁ with small perturbation co = ||x; xel12. Then we have
N
[ UNet(x))l2 S c8
K,M +
where x, = VoyXo + VI- 4 ~ (O, N is the number of the long skip connections.
Theorem 3.4 shows that for a perturbation magnitude cs, the robustness error bound of UNet in DMs is O(es(E= k,M))). For standard UNet (x; 1vi), this bound becomes very large bound D(NM))) This implies that standard UNet is sensitive extra noise in the input, especially when LSC number N is large (Mo 2 in practice, see Appendix). Hence, is necessary to control the coefficients Ką of LSCs which accords with the observations in Fig. 3 (c). Therefore, setting appropriate scaling coefficients { can not only control the oscillation of hidden features and gradients during the forward and backward processes as described in Section 3.1 and Section 3.2, but also enhance the robustness of the model to input perturbations, thereby improving better stability for DMs as discussed the beginning of Section 3.3.
4 Theorr-inspired Scaling Methods For Long Skip Connections
In Section 3, we have theoretically shown that scaling the coefficients of LSCs in UNet can help to improve the training stability of UNet on DM generation tasks by analyzing the stability of hidden feature output and gradient, and the robustness of UNet on noisy input. These theoretical results can directly be used to explain the effectiveness of 1/v 2-s-calnng technique in [27-31] which scales all LSC coefficients K; from one in UNet to a constant, and is called 1/v2 (1/ 12//2 for short). Here will use these theoretical results to design novel framework ScaleLongg including two more effective scaling methods for DM stable training. 1) constant scaling method (CS for short) and 2) learnable scaling method (LS for short).
4.1 Constant Scaling Method
In Theorem 3.1 3.4, we already show that adjusting the scaling coefficients {Ne} to smaller ones can effectively decrease the oscillation of hidden features, avoids too large gradient and also improves the robustness of U-Net to noisy input. Motivated by this, we propose an effective constant scaling method (CS) that exponentially scales the coefficients {N²} of LSCs in UNet:
CS:
Kg ki-1
(i= 1,2,,.... N, VK E (0)
(8)
When i = 1, the scaling coefficient of the first long skip connection is and thus ensures that at least one Kg is not too small to cause network degeneration.
This exponential scaling in CS method is more effective than the universal scaling methods which universally scales Kü = e.g. K 1/v2 in the 1/V2-CS method, in terms of reducing the training oscillation of UNet. We first look at the robustness error bound Theorem 3.4. For our CS, its bound is at the order of O (Mo(sxM)))- the universal scaling method. This shows the superiority of our CS method. Moreover, from the analysis in Theorem 3.1. CS method can effectively compress the oscillation range of the hidden feature h4 of the i-th layer to (n) Compared with the oscillation range O (NN of standard UNet, CS also greatly alleviates the oscillation since 2i < N. A similar analysis on Theorem 3.3 also shows that CS can effectively control the gradient magnitude, helping stabilize UNet training.
which is N22 times smaller than the bound (KMM of
Now we discuss how to set the value of in Eq. (8). To this end, we use the widely used DDPM framework on the benchmarks (CIFAR1O, CelebA and ImageNet) to estimate k. Firstly, Theorem

---

3.1 reveals that the norm of UNet output is of order O[|x; Σ Σド= Rÿ). Next, from the loss function of DMs, i..., Eq. (2), the target output is noise e, -N 1). Therefore, expect that O[|x, 11 E n23) = |e.112. and we have Σ n. = O[Ie, 12/1|/ Moreover, Fig. 3 (b) shows the distribution of ||<=|///|>xe|/ for the dataset considered, which is long-tailed and more than 97% of the values are smaller than 5. For this distribution, during the training, these long-tail samples result
in a larger value of l/e.1l/ ||x:122 They may cause the output order O[ Σソ= x3) to deviate more from the ||x:1ž. leading to unstable loss and impacting training. Therefore, it is advisable to appropriately increase سj=1* x2 to address the influence of long-tail samples. [61-65]. So the value of Σ}= n3 should be between the mean values of ||ea | l12. i... about 1.22, and rough upper bound of 5. Combining Eq. (8) and the settings of N in UViT and UNet, we can estimate [0.5, 0.95]. In practice, the hyperparameter x can be adjusted around this range.
N
4.2 Learnable Scaling Method
In Section 4. an effective constant scaling (CS) method is derived from theory, and can improve the training stability of UNet in practice as shown in Section 5. Here, we provide an alternative solution, namely, the learnable scaling method (LS), which uses a tiny network to predict the scaling coefficients for
Calibration
Feature Flement-wisr multiplicatinn
each long skip connection. Accordingly, LS is more flexible and adaptive than the constant scaling method, since it can learn the coefficients according to the training data, and can also adjust the coefficients along with training iterations which could benefit the training in terms of stability and convergence speed, which will be empirically demonstrated in Section 5.
Avg token
Figure The diagram of L.S.
As shown in Fig. 4, LS designs a calibration network usually shared by all long skip connections to predict scaling coefficients }., Specifically, for the i-th long skip connection, its input feature is I; E IR Bx where B denotes the batch size. For convolution-baser UNet, N and D respectively denote the channel number and spatial dimension (H x W) of the feature map; for transformer-basse UVIT [36], N and D respectively represent the token number and token dimension. See more discussions input features of other network architectures in the Appendix. Accordingly. LS feeds Ie E Rx into the tiny calibration network Ço parameterized by o for prediction:
LS: K; = (So (GA((z)))) E BBxNx1 i N,
(9)
where GAP denotes the global average pooling that averages 2; along the last dimension to obtain feature map of size B x N x 1; a is a sigmoid activation function. After learning {x²). J. LS uses is to scale the input F; via element- wise multiplication. In practice, network Ço is very small, and has only about 0.01M parameters in our all experiments which brings an ignorable cost compared with the cost of UNet but greatly improves the training stability and generation performance of UNet. In fact, since the parameter count of Ça itself is not substantial, to make LS more versatile for different network architecture, we can individually set a scaling module into each long skip connection. This configuration typically does not lead to performance degradation and can even result in improved performance, while maintaining the same inference speed as the original LS.
Experiment
In this section, we evaluate our methods by using UNet [3, 4] and also UViT [36, 66, 67] under the unconditional [3-5], class-conditional [36, 37] and text-to-image [38-42] settings on several commonly used datasets, including CIFAR10 [32], CelebA [33], ImageNet [34], and MS-COCO [35]. We follow the settings of [36] and defer the specific implementation details to the Appendix.
Training Stability. Fig. 1 shows that our CS and L.S methods can significantly alleviate the oscillation in UNet training under different DM settings, which is consistent with the analysis for controlling hidden features and gradients in Section 3. Though the 1/V2-CS method can stabilize training to some extent, there relatively large oscillations yet during training, particularly for deep networks. Additionally, Fig. 3 (c) shows that CS and LS can resist extra noise interference to some extent, further demonstrating the effectiveness of our proposed methods in improving training stability.
Convergence Speed. Fig. shows the training efficiency of our CS and LS methods under different settings which is a direct result of stabilizing effects of our methods. Specifically, during training, for
7

---

Architecture
Batch size (b2)
Model depth
Dataset
Org
Org
Org
Org
LS (ours)
20
LS (ours)
LS (ours)
(ours)
FID. 10Kg
CS (ours)
CS (ouns)
CS (ours)
(ours)
15
2.5%
1.7%
LO
1.7x
1,4%
2.0%
2.0%
UVIT
1.7%
by-64
64
1.7x
depth=16
MS-COCO
2/0
300
400
200
200
300
100
300
500
700
Training Heertion (k)
Training Ueriiio (8)
Training Aerution (k)
Training bematnn (k)
26
Org
Org
Org
Org
LS (ours)
35-
LS (ours)
(ours)
(ours)
FID-10K
CS (ours)
CS (ours)
(ouns)
102
(ours)
75.
18
2,0x
63
3
10
1.5%
1,5x
UNet
depth=20
CelebA
1.5x
100
140
180
200
200
300
400
130
150
200
Training iteration
Training iterstion (k)
Training Herraion (k)
Training itestion
Figure 5: Training curve comparison under different training settings in which all experiments here adopt 12-layered UVIT and batch size 128 CIFARIO unless specified in the figures.
Table 1: Synthesis performance comparison under different diffusion model settings.
Model (CIFARIO)
Type
#Param
FID.
Model (ImageNet 64)
Type
#Param
FID.
EDM [68]
Uncondition
56M
1.97
Glow [69]
3.81
IDDPM [3]
Uncondition
53M
2,90
IDDPM (small) [3]
Class-condition
100M
6.92
DDPM+ cont. [27)
Uncondition
62M
2.55
IDDPM (large) [3]
Class-condition
270M
2.92
GenViT [70]
Uncondition
11M
20.20
ADM [71]
Class-condition
296M
2.07
UNet [4]
Uncondition
36M
3.17
EDM [68]
Class-condition
296M
1.36
UVTT-S/2 [36]
Uncondition
44M
3.11
UVIT-M/4 [36]
Class-condition
131M
5.85
UNet+1/v22CS
Uncondition
36M
3.14
UVIT-L/4 [36]
Class-condition
287M
4.26
UVIT-S/2-I/V2-CS
Uncondition
44M
3.11
UWIT-M/4-I/V--CS
Class-condition
131M
5.80
UNe++CS (ouns)
Uncondition
36M
3.05
UVIT-L/4-1V/-CCS
Class-condition
287M
4.20
UNet+LS (ours)
Uncondition
36M
3.07
UNet [3]
Class-condition
144M
3.77
UVIT-S/2+CS (ours)
Uncondition
44M
2.98
UNet-1/v2-CS
Class-condition
144M
3.66
UVIT-S/2+LS (ours)
Uncondition
44M
3.01
UNet+CS (ours)
Class-condition
144M
3.48
UNc++LS (ours)
Class-condition
144M
3.39
Model (MS-COCO)
Type
#Param
FID1
UVIT-M/4+CS (ours)
Class-condition
131M
5.68
VQ-Diffusion [41]
Text-to-Image
370M
19.75
UVIT-M/4+LS (ours)
Class-condition
131M
5.75
Friro [72]
pext-to-Image
766M
8.97
UVIT-I,4+CS (ours)
Class-condition
287M
3.83
UNet [36]
Text-to-Iman Image
260M
7.32
UVIT-L/4+LS (ours)
Class-condition
287M
4.08
UNet+1/V2-CS
Text-to- Text-to-Image
260M
7.19
UVIT-//+11//2-CS
Text-to- -mage
252M
5.88
Model (CelebA)
Type
#Param
FID.
UVIT-S/2 [36]
252M
5,95
DDIM [73]
Uncondition
79M
3.26
UVTT-S/2 (Deep) [36]
Text-to-muuge
265M
5,48
Soft Truncation [74]
Uncondition
62M
1.90
UNec+CS (ours)
Text-to-Imuge
260M
7.04
UV [36]
Uncondition
44M
2,87
UN (ours)
Text-to-Imuge
260M
6.89
UVIT-S/-1/V2-CC
Uncondition
44M
3.15
UVIT-S/2+CS (ours)
Text-to-Iman Image
252M
5.77
UVIT-S/2+CS (ours)
Uncondition
44M
2,78
UVIT-S/2+LS (ours)
Text-to-Image
252M
5.60
UVIT-S/2+LS (ours)
Uncondition
44M
2.73
most cases, they consistently accelerate the learning speed by least 1.5x faster in terms of training steps across different datasets, batch sizes, network depths, and architectures. This greatly saves the actual training time of DMs. For example, when using 32-sized batch, I.S trained for 6.7 hours (200 k steps) achieves superior performance than standard UVIT trained for 15.6 hours (500 k steps). All these results reveal the effectiveness of the long skip connection scaling for faster training of DMs.
Performance Comparison. Tables show that our CS and LS methods can consistently enhance the widely used UNet and UVIT backbones under the unconditional [3-5], class-conditional [36, 37] and text-to-image [38-42] settings with almost no additional computational cost. Moreover, on both UNet and UVTT, our CS and IS also achieve much better performance than 1/v2-CS. All these results are consistent with the analysis in Section 4.1.
8

---

bz- 12
b2-64/depth-12
b2-32/deptth- 12
b2-64td
by-64/depthh20
Ong UVIT
3.11
5.36
13.89
8.31
6.62
v2 -CS
3.11
4.72
12.67
5.67
4.95
(ours)
3.01
3.58
5.90
4.08
3.56
CS (ours)
2.98
3.95
8.51
5.36
4.14
2.5
3.05
315
3.50
1.50
5.50
6.0
10.0
13.0
1.0
5.0
6.
10.0
4.3
6.0
3.0
FID-50K]
FID-50k]
FID-S0k]
FID-50k]
FID-50k]
Figure 6: Synthesis performance of proposed methods under small batch size and deep architecture.
Robustness to Batch Size and Network Depth. UViT is often sensitive to batch size and network depth [36], since small batch sizes yield more noisy stochastic gradients that further aggravate gradient instability. Additionally, Section 3 shows that increasing network depth impairs the stability of both forward and backward propagation and its robustness to noisy inputs. So in Fig. 6, we evaluate our methods on standard UVIT with different batch sizes and network depths. The results show that our CS and LS methods can greatly help UViT on CIFAR 10 by showing big improvements on standard UVIT and 1/ //22-SS-base UVTT under different training batch sizes and network depth.
Robustness of LS to Network Design. For the calibration module in LS, here we inves- tigate the robustness of LS to the module designs. Table 2 shows that using advanced modules, e.g. IE [77, 80] and SE [79]. can improve LS performance compared with simply treating K; as learnable coefficients. So we use SE module [79, 81, 82) in our experiments (see Appendix). Furthermore, we observe that LS is robust to different activation functions and compression ratios r in the SE module. Hence, our LS does not require elaborated model crafting and adopts default 16 and ReLU throughout this work.
Table 2: Robustness of LS to network design
Ratio
FIDJ
Activation
FIDL
Modules
FID
3.06
ELU [75]
3.08
Leamable Ka
3.46
3.10
RcLL [76] (ours)
3.01
1b: [77] 3.09
16 (ours)
3.01
Mish [78)
3.08
SE [79] (ours)
3.01
Other scaling methods. Our LS and CS methods scale LSC, unlike previous scaling methods [46, 48-53] that primarily scale block output in classification tasks. Table 3 reveals that these block-output scaling methods have worse performance than LS (FID: 3.01) and CS (FID: 2.98) under diffusion model setting with UVIT as backbone.
Table 3: Other scaling
Mchtod (CIFARIO)
FIDI
stable 0.1 [46]
3.77
stable r 0.3 [46]
3.75
ReZero [48]
3.46
Discussion
Relationship between LS and CS. Taking UVIT on CIFAR1O, MS-
1.0
ImageNet MS-COCo
COCO, ImageNet and CelebA as baselines, we randomly sample 30 Gaussian noises from (0, I) and measure the scaling coefficients, i.e.
0.8
CelebA
K; 1, learned by LS for each LSC. Fig. 7 shows that LS and CS share similar characteristics. Firstly, the predicted coefficients of LS fall within the orange area which is the estimated coefficient range given by CS in Section 4. Moreover, these learned coefficients K2 share the almost the same exponential curve with Eq. (8) in CS. These shared characteristics have the potential to serve as crucial starting points for further optimization of DM network architectures in the future. Moreover, we try to preliminary analyze these shared characteristics.
CIFARIO
0.6-
0.24
4-th
skip connecion
Figure 7: The K; from LS.
First, for CS, how is the direction of applying the exponentially decayed scale determined? In fact, if we use the reverse direction, namely, K; = N (x < 1), the stability bound in Theorem 3.4 is extremely large. The main term of stability bound in Theorem 3.4 can be written S, = EN Hi, Mg = ':N Mo + M3 + ... + K.MYN and could be very large, since My > 1 is large when N is large and scaling it by factor k could not sufficiently control its magnitude. In contrast, our default setting K; = - of CS can well control the main term in stability bound: S S Ra = Mn + K. MB + K NNN where the larger terms Mt+1 are weighted by smaller coefficients s². In this way, S is much smaller than S,, which shows the advantages of our default setting. Besides, the following Table 4 also compares the above two settings by using UVIT on Cifar (batch size = 64), and shows that our default setting exhibits significant advantages.
Next, for LS, there are two possible reasons for why LS discovers decaying scaling curve similar to the CS. On the one hand, from a theoretical view, as discussed for the direction of scaling, for the i-th long skip connection (1 < i S N). the learable K; should be smaller to better control the
9

---

Table 4: The FID-10k], result of different direction of scaling.
Training step
5k
10k
15k
20k
25k
30k
35k
40k
45k
K N-1+1
67.26
33.93
22.78
16.91
15.01
14.01
12.96
12.34
12.26
Ri = = (ours)
85.19
23.74
15.36
11.38
10.02
8.61
7.92
7.27
6.65
magnitude of Mj so that the stability bound, e-g. in Theorem 3.4, is small. This directly yields the decaying scaling strategy which is also learnt by the scaling network. On the other hand, we can also analyze this observation a more intuitive manner. Specifically, considering the UNet architecture, the gradient that travels through the i-th long skip connection during the backpropagation process influences the updates of both the first blocks in the encoder and the last blocks in the UNet decoder. As result, to ensure stable network training, it's advisable to slightly reduce the gradients on the long skip connections involving more blocks (i.c., those with larger values) to prevent any potential issues with gradient explosion.
Limitations. CS does not require any additional parameters or extra computational cost. But it can only estimate a rough range of the scaling coefficient ς as shown in Section 4.1. This means that one still needs to manually select к from the range. In the future, we may be able to design better methods to assist DMs in selecting better values of K. or to have better estimates of the range of k. Another effective solution is to use our LS method which can automatically learn qualified scaling coefficient for each LSC. But LS inevitably introduces additional parameters and computational costs. Luckily, the prediction network in LS is very tiny and has only about 0.01 M parameters, working very well in all our experiments.
Moreover, our CS and LS mainly focus on stabilizing the training of DMs, and indeed well achieve their target: greatly improving the stability of UNet and UVIT whose effects is learning speed boosting of them by more than 1.5 x as shown in Fig. 5. But our CS and LS do not bring a very big improvement in terms of the final FID performance in Table 1, although they consistently improve model performance. Nonetheless, considering their simplicity, versatility, stabilizing ability to DMs, and faster learning speed, LS and CS should be greatly valuable in DMs.
Conclusion. We theoretically analyze the instability risks from widely used standard UNet for diffusion models (DMs). These risks are about the stability of forward and backward propagation, as well as the robustness of the network to extra noisy inputs. Based on these theoretical results, we propose a novel framework ScaleLong including two effective scaling methods, namely LS and CS, which can stabilize the training of DMs in different settings and lead to faster training speed and better generation performance.
Acknowledgments, This work was supported in part by National Key R&D Program of China under Grant No.2021ZDO0 1 1601, National Natural Science Foundation of China (NSFC) under Grant No.61836012, 62325605, U21A20470, GuangDong Basic and Applied Basic Research Foundation under Grant No. 2023A151501 1374, GuangDong Province Key Laboratory of Information Security Technology.
References
[1] Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 11461-11471, 2022.
[2] Tiankai Hang, Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, and Baining Guo. Efficient diffusion training via min-snr weighting strategy. arXiv preprint arXiv:2303.0953 095556, 2023.
[3] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffuusio probabilistic models. In International Conference on Machine Learning, pages 8162-8171. PMLR, 2021.
[4] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Demoising diffusion probabilistic models. Advances in Neural Information Processing Systems. 33:6840-6851, 2020.
[5] Jacob Austin, Daniel D Johnson. Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17981-17993 2021.
[6] Florinel-Alin Croitoru, Vlad Hondru, Radu Tudor lonescu, and Mubarak Shah. Diffusion models in vision: A survey, IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
10

---

[7] Nithin Gopalakrishnan Nair. Kangfu Mei, and Vishal M. Patel. At-ddpm: Restoring faces degraded by atmospheric turbulence using denoising diffusion probabilistic models. 2023 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), pages 3423-3432, 2022.
[8] Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. Ilvr: Conditioning method for denoising diffusion probabilistic models. 2021 IEEF/CVF International Conference on Computer Vision (ICCV). pages 14347-14333 2021.
[9] Julian Wyatt, Adam Leach, Sebastian M. Schmon, and Chris G. Willcocks. Anoddpm: Anomaly detection with denoising diffusion probabilistic models using simplex noise. 2022 IEEE/CVF Conference on Computer Vision and Puttern Recognition Workshops (CVPRW), pages 649-655, 2022.
[10] Ian Goodfellow, Jean Pouget-Abadic Mehdi Mirza, Bing Xu, David Warde-Farles Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11):139-
144, 2020.
[11] Shengjia Zhao, Jiaming Song, and Stefano Ermon. Infovae: Balancing learning and inference in variational autoencoders. Proceedings of the aaai conference on artificial intelligence, volume 33, pages 9885-5892, 2019.
[12] Gwanghyun Kim, Taesung Kwon, and Jong Chul Ye, Diffusionclip: Text-guided diffusion models for robust image manipulation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 2426-2435, 2022.
[13] Omri Avrahami, Dani I Lischinski, and Ohad Fried. Blended diffusion for text-driven editing of natural images. In Proceedings of the IEEF/CVF Conference Computer Vision and Pattern Recognition, pages 12208-1818,, 2022.
[14] Alexander Quinn Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob Megrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. In International Conference on Machine Learning, pages 16784-16804. PMI.R, 2022.
[15] Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is strong image synthesizer. ArXiv, 2023.
[16] Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall, Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint ar 220: /4988, 2022.
[17] Tengfei Wang, Bo Zhang. Ting Zhang, Shuyang Gu, Jianmin Bao, Tadas Baltrusaitis, Jingjing Shen, Dong Chen, Fang Wen, Qifeng Chen, et al. Rodin: A generative model for sculpting 3d digital avatars using diffusion. arXiv preprint arXiv:2212.06115, 2022.
[18] Jonathan Ho, Tim Salimans, Alexey Gritsenko. William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. arXiv preprint arXiv:2204.0 033458. 2022.
[19] Eyal Molad, Eliahu Horwitz, Dani Valevski, Alex Rav Acha, Yossi Matias, Yael Pritch, Yaniv Leviathan, and Yedid Hoshen. Dreamix: Video diffusion models are general video editors. arXiv preprint 2xXvv2330.003329 202..
[20] Patrick Esser, Johnathan Chiu, Parmida Atighchchian, Jonathan Granskog, and Anastasis Germanidis. Structure and content- geieent-ggdiddddd video synthesis with diffusion models. arXiv preprint arXiv:2302.0301. 2023.
[21] Andreas Blattmann. Robin Rombach, Huan Ling. Tim Dockhorn, Seung Wook Kim, Sanja Fidler, and Karsten Kreis. Align your latents: High-resolution video synthesis with latent diffussin models. arXiv preprint arXiv:2304..88: 088888. 2023.
[22] Yi Zhang, Dasong Li, Xiaoyu Shi, Dailan He, Kangning Song, Xiaogang Wang. Hongwei Qin, and Hongsheng Li. Kbnet: Kernel basis network for image restoration. arXiv preprint arXiv:2303.0228. 2023.
[23] Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, Ming-Hsuan Yang, and Ling Shao. Multi-stage progressive image restoration. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 14821-14831, 2021.
[24] Zhendong Wang. Xiaodong Cun, Jianmin Bao, Wengang Zhou, Jianzhuang Liu, and Houqiang Li. Uformer:
A general shapee transformer for image restoration. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 17683 17683-17693, 2022.
11

---

[25] Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, and Ming-Hsuan Yang. Restormer: Efficient transformer for high-resolution image restoration. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 5728-5739, 2022.
[26] Liangyu Chen, Xiaojie Chu, Xiangyu Zhang, and Jian Sun. Simple baselines for image restoration. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part VII, pages 17-33. Springer, 2022.
[27] Yang Song, Jascha Narain Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations, ArXiv, abs/201 1.3455, 2020.
[28] Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45:4713 44266, 2021.
[29] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Seyed Kam- yar Seyed Ghasemipour, Burcu Karagol Ayan, Seyedch Sara Mahdavi, Raphael Gontijo Lopes, Tim Salimans, Jonathan Ho, David J. Flect, and Mohammad Norouzi. Photorcalistic text-to-image diffusion models with deep language understanding. ArXiv, abs/2205. 11487, 2022.
[30] Jonathan Ho, Ajay Jain, and P. Abheel. Denoising diffusion probabilistic models. ArXiv, ahs/2006.111111 11239, 2020.
[31] Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 4396-4405, 2018.
[32] Alex Krizhevsky, Learning multiple layers of features from tiny images. 2009.
[33] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 201.5.
[34] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Scan Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael S. Bernstein, Alexander C. Berg, and Li Fei-Fei. Imagenet large scale visual recognition challenge. International Journal of Computer Vision, 11::211-222, 2014.
[35] Tsung- Yi Lin, Michael Maire. Serge J. Belongie, James Hays. Pietro Perona, Deva Ramanan, Piotr Dollár. and C. Lawrence Zitnick. Microsoft coco: Common objects in context. In European Conference on Computer Vision, 2014.
[36] Fan Bao, Chongxuan Li, Yue Cao, and Jun Zhu, All worth words: vit backbone for score-based diffusion models. arXiv preprint (XXi:22209112555 2022.
[37] Jonathan Ho. Classifier- free diffusion guidance. ArXiv, abs/2207. 12598, 2022.
[38] Shanshan Zhong, Zhongzhan Huang, Wushao Wen, Jinghui Qin, and Liang Lin. Sur-adapter: Enhancing text-to-image pre-trained diffusion models with large language models. arXiv preprint 285.52300..05
2023.
[39] Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. Dream- booth: Fine tuning text-to-image diffusion models for subject-driven generation. ArXiv, 112//2288.12222,, 2022.
[40] Nupur Kumari, Bin Zhang, Richard Zhang, Eli Shechtman, and Jun- Yan Zhu. Multi-concept customization of text-to-imme iot-to- diffusion. ArXiv, abs/2212.04442 2022.
[41] Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for text-to-image synthesis. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10686-10696, 2021.
[42] Jiale Xu, Xintao Wang, Weihao Cheng, Yan-Pei Cao, Ying Shan, Xiaohu Qie, and Shenghua Gao. 3ream3d: Zero-shot text-to-3d synthesis using 3d shape prior and text-to-image diffusion models. ArXiv, 147044 2022.
[43] Firas Khader, Gustav Mueller-Franzes, Soroosh Tayebi Arasteh, Tianyu Han, Christoph Haarburger. Maximilian Franz Schulze- Hagen, Philipp Schad, Sandy Engelhardt, Bettina Baessler, Sebastian Foersch, J. Stegmaier. Christiane Kuhl, Sven Nebelung, Jakob Nikolas Kather, and Daniel Truhn. Denoising diffusion probabilistic models for 3d medical image generation. Scientific Reports, 13, 2022.
12

---

[44] Liangyu Chen, Xin Lu, J. Zhang, Xiaojie Chu, and Chengpeng Chen. Hinet: Half instance normalization network for image restoration. 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops (CVPRW), pages 182-192, 2021.
[45] Ying Tai, Jian Yang, Xiaoming Liu, and Chunyan Xu. Memnet: A persistent memory network for image restoration. 2017 IEEE International Conference on Computer Vision (ICCV), pages 4549-4557, 2017.
[46] Huishuai Zhang, Da Yu, Mingyang Yi, Wei Chen, and Tie-Yan Liu. Stabilize deep resnet with sharp scaling factor tau. 2019.
[47] Huishuai Zhang, Da Yu, Wei Chen, and Tin-Yan Liu. On the stability of multi-branch network, 2021.
[48] Thomas C. Bachlechner, Bodhisattwa Prasad Majumder. Huanru Henry Mao, G. Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence large depth. In Conference on Uncertainty in Artificial Intelligence, 2020.
[49] Hongyi Zhang. Yann Dauphin, and Tengyu Ma. Fixup initializaatin: Residual Icarning without normaliza- tion. ArXiv, 0tt//9901.932//1 2019.
[50] Soham De and Samuel I.. Smith. Batch normalization biases residual blocks towards the identity function in deep networks. arXiv: Learning, 2020.
[51] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Yan Liu. On layer normalization in the transformer architecture, 2020.
[52] Bobby He, James Martens, Guodong Zhang, Aleksandar Botev, Andy Brock, Samuel L. Smith, and Yee Whye Tch. Deep transformers without shortcuts; Modifying self-attention for faithful signal propaga- tion. ArXiv, abs/2302. 10322, 2023.
[53] Soufiane Hayou, Fugenio Clerico, Bo He, George Deligiannidis, A. Doucet, and Judith Rousseau. Stable resnet. In International Conference on Artificial Intelligence and Statistics, 2020.
[54] Boris Hanin and David Rolnick, How to start training: The effect of initialization and architecture. Advances in Neural Information Processing Systems, 31, 2018.
[55] David Balduzzi. Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma. and Brian McWiliiams. The shattered gradients problem: If resnets are the answer, then what is the question? In International Conference on Machine Learning, pages 342-350. PMLR, 2017.
[56] Muhammad Awais. Md. Tauhid Bin lqbal, and Sung- Ho Bae. Revisiting internal covariate shift for batch normalization. IEEE Transactions on Neural Networks and Learning Systems, 32:5082-5092, 2020.
[57] You Huang and Yuanlong Yu. An internal covariate shift bounding algorithm for deep neural networks by unitizing layers' outputs. 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 8462-882 8470, 2020.
[58] Sergey Loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, 2015.
[59] Shibani Santurkar, Dimitris Tsipras, Andrew Ilyas, and Aleksander Madry. How does batch normalization help optimization? (no, is not about internal covariate shift), ArXiv, abs/1805. 1604, 2018
[60] Tim Cooijmans, Nicolas Ballas, César Laurent, and Aaron C. Courville. Recurrent batch normalization. ArXiv, abs/ 690.0905, 2016
[61] Tong Wu, Ziwei Liu, Qingqiu Huang, Yu Wang, and Dahua Lin. Adversarial robustness under long-tailed distribution. 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 8655 8664, 2021.
[62] Zhuo Chen, Jinglong Chen, Yong Feng, Shen Liu, Tianci Zhang, Kaiyu Zhang, and Wenrong Xiao. Imbalance fault diagnosis under long-tailed distribution: Challenges, solutions and prospects. Knowl. Based Syst., 258: 10008, 2022.
[63] Xiyu He and Xiang Qian. A real-time surface defect detection system for industrial products with long- tailed distribution. 2021 IEEE Industrial Electronics and Applications Conference (IEACon), pages 313-317, 2021.
[64] Xudong Wang, Long Lian, Zhongqi Miao, Ziwei Liu, and Stella X. Yu. Long-tailed recognition by routing diverse distribution- aware experts. ArXiv, abs/2010.01809, 2020.
13

---

[65] Shaden Alshammari. Yu-Xiong Wang, Deva Ramanan, and Shu Kong. Long- tailed recognition via weight balancing. 2022 IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), pages 6887-6897, 2022.
[66] Fan Bao, Shen Nie, Kaiwen Xue, Chongxuan Li, Shiliang Pu, Yaole Wang. Gang Yue, Yue Cao. Hang Su, and Jun Zhu. One transformer fits all distributions in multi-modal diffusion at scale. ArXiv, abs/2303.06555, 2023.
[67] Zehin You, Yong Zhong, Fan Bao, Jiacheng Sun, Chongxuan Li, and Jun Zhu. Diffusion models and semi-supervised learners benefit mutually with few labels. ArXiv, abs/2302. 10586, 2023.
[68] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. ArXiv, abs/2206. 003644 2022.
[69] Diederik Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible Ixl convolutions. ArXiv, abs/1 807.007.0000 007.03039, 2018.
[70] Xiulong Yang, Sheng-Min Shih, Yinlin Fu, Xiaoting Zhao, and Shihao Ji. Your vit is secretly hybrid discriminnaive-geerraaie diffusion model. arXiv preprint arXiv:2208.0779) 2022.
[71] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advunces in Neural Information Processing Systems, 34:8780-8794, 2021.
[72] Wanshu Fan, Yen-Chur Chen, Dongdong Chen, Yu Cheng, Lu Yuan, and Yu- Yu-Chiang Frank Wang. Frido: Feature pyramid diffusion for complex scene image synthesis. ArXiv, abs/2208. 13753, 2022.
[73] Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint vrXiv:2010.02.502 2020.
[74] Dongjun Kim, Scung-Jac Shin, Kyungwoo Song, Wanmo Kang, and Il-Chul Moon. Soft truncation:
A universal training technique of score-based diffusion model for high precision score estimation. In International Conference Machine Learning, 2021.
[75] Amn Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv: Learning, 2015.
[76] Abien Fred Agarap. Deep learning using rectified linear units (relu). ArXiv, 0bs//880308333 2018.
[77] Senwei Liang, Zhongzhan Huang, Mingfu Liang, and Haizhao Yang. Instance enhancement batch normalization: an adaptive regulator of batch noise. In AAAI Conference on Artificial Intelligence, 2019.
[78] Diganta Misra. Mish: self regularized non-monotonic activation function. British Machine Vision Conference, 2020.
[79] Hu, Shen, Samuel Albanie, Gang Sun, and Enhua Wu. Squeeze-and-excitttie networks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 42:2011 10221, 2017.
[80] Shan Zhong, Wushao Wen, and Jinghui Qin. Switchable self-attention module, ArXiv, a552220.0060000 2022.
[81] Sanghyun Woo, Jongchan Park, Joon-Young Lee, and In-So Kweon. Cham: Convolutional block attention module. Furppean Conference on Computer Vision, 2018
[82] Zhongzhan Huang, Senwei Liang, Mingfu Liang, and Haizhao Yang. Dianet: Dense-and-implici attention network. In AAAI Conference on Artificial Intelligence, 2019.
[83] Gavin E Crooks. Survey of simple, continuous, univariate probability distributions. Technical report, Technical report, Lawrence Berkeley National Lab, 2013., 2012.
[84] Robert B Davies. Algorithm as 155: The distribution of a linear combination of x 2 random variables. Applied Statistics, pages 323-333, 1980.
[85] Robert B Davies. Numerical inversion of a characteristic function. Biometrika, 60(2)415--17, 1973.
[86] Zhongzhan Huang, Wenqi Shao, Xinjiang Wang, Liang Lin, and Ping Luo. Rethinking the pruning criteria for convolutional neural network. Advances in Neural Information Processing Syxtems. 14:16305-16318, 2021.
[87] Zeyuan Allen-Zhu Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference Machine Learning, pages 242-252. PMLR, 2019.
14

---

The appendix is structured as follows. In Appendix A, we provide the training details and network architecture design for all the experiments conducted in this paper. Additionally, we present the specific core code used in our study. Appendix B presents the proofs for all the theorems and lemmas introduced this paper. Finally, in Appendix C, we showcase additional examples of feature oscillation to support our argument about the training instability of existing UNet models when training diffusion models. Furthermore, we investigate the case of scaling coefficient A greater than in CS, and the experimental results are consistent with our theoretical analysis.
A Implementation details
ImageNet 64x 64
CelcbA 64 64
CIFARIO 32 32
MS-COCO
Figure 8: The examples of the generated images by UVIT with our proposed LS.
A.1 Experimental setup
In this paper, we consider four commonly used datasets for our experiments: CIFAR1O [32], CelebA [33], ImageNet [34], and MS-COCO [35]. For the unconditional setting, we use CIFARIO 10, which consists of 50,000 training images, and CelebA 64, which contains 162,770 training images of human faces. For the class-conditional settings, we utilize the well-known ImageNet dataset at 64 x 64 resolutions, which contains 1, 1,281,165 training images from 1,000 different classes. For the text-to-image setting, we consider the MS-COCO dataset 256 x 256 resolution. Unless otherwise specified, all of our experiments are conducted on A100 GPU and follow the experimental settings in [36].
Table 5: The summary of the datasets for the experiments.
Dataset
#class
#training
#type
#Image size
CIFARIO
10
50,000
Uncondition Uncondition
32 x 32
CelebA
100
162,770
64 64
ImageNet 64
1000
1,281,167
Class-condition
64 x 64
MS-COCO
82,783
Text-to-Image
256 256
A.2 Network architecture details.
The scaling strategies we proposed on long skip connections, namely CS and LS, are straightforward to implement and can be easily transferred to neural network architectures of different diffusion models. Fig. and Fig. 10 compare the original UVIT with the code incorporating CS L.S, demonstrating that proposed methods require minimal code modifications.
15

---

Table 6: The training settings for different dataset.
Dataset
CIFAR10 64x64
CelehA 64x64
ImageNet 64x64
MS-COCO
Latent shape
32x32x4
Batch size
128/64/32
128
1024
256
Training iterations
500k
500k
300k
1M
Warm-up steps
2.5k
4k
5k
5k
Optimizer
Adam
Adam W
Adam W 3. 3.00E-04 0.03 (0.99.0.99)
AdamW
Learning rate
2.00E-04
2.00E-04
2.00E-04
Weight decay
0.03
0.03
0.03
Beta
(0.99,0.999)
(0.99.0.99)
(0.9,0.9)
Origin
CS (ours)
skips [
skips [] for bik in self-in_blocks:
selF.in, 2a1.in__bae iiFock_koock::
b1k(x)
(xpeend()))
(xppen))
self. www.fmddbbbleee block(xx)
self self.nid_block@
layer eere_cttt__
bik self blocks:
for bik in self.out blocks:
bik(x, pop)))
skips..pop) ke
depth layer_cr
layer_cnt cnt +2
= sel1..nor(x)
(x)mm)
self.decoder pred(x)
x self. dec.dddedeeeneeee decoder_pred(xx)
assert x1size(1) selF.extras
assert x.size(1) self.extras
x(;, self self.extrasi,
x = x.:, self.extras:, :]
unpateniff(x,, self.in_chan chans)
unpatchify(x, aaff.in_chrns) x self.final_layerr
self.final_aae layer_)
return
Figure 9: The comparison between the original UVIT and our proposed CS.
CS requires no additional parameters and can rapidly adapt to any UNet-alike architecture and dataset. In our experiments, have chosen values of к 0.5, 0.5, 0.7, and 0.8 for the four datasets illustrated in Table 5, respectively. In various training environments, we recommend setting x within the range of 0.5 to 0.95. I.S designs calibration network shared by all long skip connections to predict scaling coefficients }. Specifically, for the i-th long skip connection, its input feature is 2; E BR where B denotes the batch size. For convolution- channel number and spatial dimension (H x W) of the feature map; for transformer-based UViT [36], N and D respectively represent the token number and token dimension.
UNet, N and D respectively denote the
In the main text, we define the LS as
= (Co GAP((2)) BxNxI 1PXXNX,1141 1SiSN,
(10)
For Co. we employ a neural network structure based on SENet [79], given by:
(x(z) W, 35 W2 z),
(11)
where & E R., denotes matrix multiplication, and vi represents the ReLU activation function. The matrices W, E N/NN// and WI E R where r is the compression ratio, are set to r = 16 in this paper. In the ablation study of the main text, have demonstrated that the choice of has minimal impact on the performance of the diffusion model. It is worth noting that LS is applied to long skip connections where the feature dimensions remain unchanged, such in UViT, then the calibration network is only shared (z(z).
However, when LS is applied to neural networks where the dimension of features in certain long skip connections may change, such as in UNet with downsample and upsample operations, we need to provide set of encoders and decoders in the calibration network to ensure dimension alignment for different feature dimensions.
Assuming there are d different feature dimensions in the long skip connections, where the dimension of the i-th feature (B N, D), 1 S i d, the input dimension of (a(x)) is set Nuin min N, Additionally, set up d sets of encoders and decoders, where features with the same dimension
16

---

Origin
LS (ours)
skips []
skips []
for blk self ii_bblccks::
for bik in self.in_blocks: blk(x)
x blk(x)
(xpped)
akips.append(x)
x self.mid mid_blocc(xx)
self.mid block(x)
for blk self.out bllf.out_blockk: blocks:
for bik in self.our otff.ouutbllckk:
x blk(x, skips.pop())
skip skips.pop()
blk(x, skip self.15(skiip
x self.norm(x)
x self.norm(x)
x self.decoder (xcooder_pred(x))
x self.decoder_pnee
assert x.size(1) self.extras
assert x.size(1) self.extras
x[:, self.extras;,
x x(:, self.extrast, :]
x unpatchify(x, self.in_chars channss_
x unpatchify(x, self.in relf.in_chans
x selyer))ol_laee(x))
x self.fimel_1at www.fllllllayr(( (x)
return
Figure 10: The comparison between the original UViT and our proposed L.S.
share encoder and decoder. The i-th set of encoder (() and decoder See also adopt the SENet structure, given by:
(i)
W Nuin N,/r, 2 W N./r,N,, oz), x ZER
(12)
where T is the compression ratio, typically set to [N./4], and
de = W WN,xN./r o 16 W Ne/r,xNaan о x), I EER Im
(13)
In fact, since the parameter count of (x(r) itself is not substantial, to make LS more versatile, we can also individually set scaling module into each long skip connection. This configuration typically does not lead to sigmificant performance degradation and can even result in improved performance, while maintaining the same inference speed as the original LS.
17

---

B
The proof of the Theorems and Lemma
B.1
Preliminaries
Diffusion Model (DM). DDPM-alike DMs [1-6, 8, 9, 43] generates sequence of noisy samples {x:]__ by repeatedly adding Gaussian noise sample Xo until attatining Xy ~ (0, I). This noise injection process, a.k.a. the forward process, can be formalized as Markov chain g (x17 I q (x, (x11)))) where g(x{|x{-1) Vaxx B,T), 0g and 33 depend and satisfy ag + 3. = 1. By using the properties of Gaussian distribution, the forward process can be written as
g(x, xn) = (x,; vā, (1 -1-68 ā,)
(14)
Vā,X0 vI a where 4 (O, I). Then DM adopts neural network éØ () to invert the forward process and predict the noise added at each time step. This process is recover data Xo from Gaussian noise by minimizing the loss
where 02 IT Li=1 Og. Next, can sample
C impll (0) Eli. - co(vārxo + v1 t)
(15)
UNet-alike Network. Since the network Co(:; t) in DMs predicts the noise to denoise, it plays similar role of UNet-alike network widely in image restoration tasks, e.g. image de-raining, image denoising [22-26, 44, 45]. This inspires DMs to use UNet-alike network (16) that uses LSCs to connect distant parts of the network for long-range information transmission and aggregation
UNet(2) fo(x), fi(x) = d;+1 o [Re+1 0,+1 x + fit1 (4:+1 c x)), 0 0 N (16)
where x E I.Th denotes the input, a; and by (i 2 1) are the trainable parameter of the i-th block, N2 (i 1) are the scaling coefficients and are set to in standard UNet. fN is the middle block of UNet. For the vector operation it can be designed to implement different networks.
W.L.o.g, in this paper, consider matrix multiplication, and set the block a; and b, the stacked networks [46, 47] which is implemented as
W oi
6, w (( ( 60
(17)
with ReLU activation function & and learnable matrices W W E RRXm (j 2 1). Moreover, let fN also have the same architecture, i.e. fx(x) w o W (W
Proposition B.1 (Amoroso distribution). The Amoroso distribution a four parameter, continuous, univariate, unimodal probability density, with semi-infinite range [83]. And its probability density function is
A Amoroso(X
1
3
0,
exp
X x-a a -(-
(18)
T(a)
0
0
for x. a, O, at 3 E ER, > and range x 2 0 0, x S a if0 < The mean and variance of Amoroso distribution are
T(a + } T(a)
(19)
and
+(a+ T(a))
T(a + +1) T(a)²
Var Varx (mox-AAmmrooxX 0²
(20)
Proposition B.2 (Stirling's formula). For big enough x and x. E RRt, we have an approximation of Gamma function:
z z
I(x + 1) V 2
=
(21)
18

---

and Z4, = 1...., k are lo independent, normally distributed random variables with mean 0 and standard deviation The statistic (x) = Ei x follows Scaled Chi distribution [83]. Moreover, f2(X) also fol- lows Amoroso(z)0) 22, 2... By Eq. (19) and Eq. (20), the mean and variance of Scaled Chi distribution are
Proposition B.3 (Scaled Chi distribution). Let X
Ex Ex 20/303 0.. I(
(22)
r($)
and
[T 1)
r( )
12x = 20²
(23)
I(
r(
Lemma B.4. For big enough x and R+, we have
2+1
lim +
2
(24)
[(2)
2
And
T( +1)
)(
lim
(25)
++++0
7(5)
r(5)
4
Proof.
T(
2
2m
1
2
26
1
lim
lim
from Proposition. B.2
['(2)
2m
2-2
2
2e
lim
1 -
2
R-
2 1
1
2
a-
26
2e
2-2
1
1
= lim
1 2
0-+00
2
2e
on the other hand, we have
11)+
T
1
z-2
lim
lim x-+o
1+
T($)
((#)
2
2
2 2
2e
1
lim
c-
(1
2e
x
(-e)
2
Lemma B.5. 184, 85] Let x (I1, ..., NN), where 1 < i S N, are independently and identically distributed random variables with Di ((,,0²) this case, the matheraatcaal expectation E 2?) can be calculated as N2N N + N2N.
19

---

Lemma B.6. [86] Let F,y E IR NN be the j-th filter of the i-th convolutional layer. If all filters are initialized as Kaiming's 's initialization, i.e., NO, N,xkxk), in i-th layer, Fij (j = 1, ..... N411) are ii.d and approximately follow such distribution during training:
N,0, - -Nxokok)-
(26)
Lemma B.7. [47, 87] For the feedforward neural network Wi (1 with initialization
(O, )) and W,, 1 < S L, mmxm vp (0, 1), the following inequality holds with at least
1 O()) expl- /t]]:
(1 (1-P) -) S [ l2 I (1 hol a E (1,2,
1),,
(27)
where ha is @W (Whh-1) for 1,2, .... 1 with ho
Lemma B.8. For a matrix W E Rm and its elements are independently follow the distribution N(0,c)) )²), we have the follwing estimation while m large enough
[Wsl2 = ((1ss|1)),
(28)
where 8 R.m.
Proof. Let U4, 1 < i S m be the i-th column of matrix W, and we have 22 ~ c' We first prove the following two facts while m - oo:
(1) ||wl ||,,Ile + /2c- r((m+1)/2) T(m/2)
Si<js m;
(2) (m,, 03) + 0, < <i<j
For the fact (1), since Chebyshev inequality, for S i S m and given M, we have
P P(lllaala B([|u]|aa) 2 / /MVar(|o:lla)) <
1
(29)
M
from Proposition B.3 and Lemma. B.4, we can rewrite Eq. (29) when m - 0o:
P IIvilla V2c. I((m + 1)/2)
{
M e, 22c.
I((m 1)/2)
M e
+
21
(30)
T(m/2)
T(m/2)
M
For a small enough
> 0, let M = 1/e. Note that V c == c/v 2e is constant. When m + oc.
/2c-
1((m+1)/2) ((m/2)
M z
Hence, for any i E [1, m] and any small enough e, we have
P Plbbila x V2c- I( + 1)/2) 2-c T(m/2)
(31)
So the fact (1) holds. Moreover, consider the (e,, 0z), 1 < <
Let U (Wa,, ..., Usm) and dy = (U/1, Uy2, ... Ujm). So (u), 13)) and by are independent, hence
Ep Uipllpp- Note that, 0
(((ipUjp) = 0,
(32)
Var(Uip"jp) Var(cip) Var(r;p) + [B(()))) Var(u;p)) + ( c4,
(33)
since central limit theorem, we have
Σ Up - -0) (0,
Vm. c2.
(34)
According to Eq. (22), Lemma B.4 and Eq. (34), when m - oo, we have
(1), 1z) 11u
1
(v,
vz)
N(,00 ((,
(35)
Vm
Vm
20

---

Therefore, the fact (2) holds and we have (v,, 2y) - 0, S 0,,,,,4 sisje & m.
Next, consider [W 8||2 for any IRm
W s = 8 [m1, u2,... Umm] [11, 02, -2,-.., vmn
vn 12
of 02
oi Um
02 11
10221
(36)
From the facts (1) and (2), when m is large enough, we have 2." + 0, i + j. and
| ~ 102212 ~
pm Cg. Therefore, we have
[W =8 T c. = 0(1|s112)
(37)
Let = 1 in Lemma B.6, and the convolution kernel degenerates into a learnable matrix. In this
case, observe that if W E R" is initialized using Kaiming initialization, W(0, 2/m). during the training process, the elements of W will approximately follow a Gaussian distribution (0, (0mmm))) In this case,
easily rewrite Lemma B.7 the following Corollary B.9:
Corollary B.9. For the feedforward neural network W tialization during training.
(eW((((W2))) (x)) with kaiming's 's ini- (0,1), the following inequality holds with at least Vp
(I)) (mpp
O((1 - O(1--))lkol1> & |/hll2 S [(1 + pholll)
aE (1,2, 1},
(38)
where ha is of hhh for a 1,2, I..,..--. with ho x.
Moreover, considering Lemma B.6 and Lemma B.8, we can conclude that the matrix Wi in Corollary B.9 will not change the magnitude its corresponding input ((Wi-I..&(WW (W)))).. In other words, we have the following relationship:
O((1 - p)]-o)) & [Wihh-llle -
(39)
Therefore, for a; and b, defined in Eq. (17), we have following corollary:
Corollary B.10. For o; and b, defined in Eq. (17), 1 S sis N, and their corresponding input v. and Ib. Vp E (0, 1), the following each inequality holds with at least 1 -O() ((mpp?
||a Fall2 O(11xall2), ||b, moll2 ((1xxa122)
(40)
Lemma B.11. [47] Let D be diagonal matrix representing the ReLU activation layer. Assume the training loss is ""= (W),, where (,W denotes the training loss of the s-th sample among the training samples. For the Wi defined in Eq. (17) and (x,) defined in Eq. (16), where p E (1,2,...,1.00 {ai, bi), the following inequalities hold with least [xxp[))))m[
(D ... .... D(W..)T ....) I vll2 S 0(10012),
(41)
and
TT. T... .T.W") ...") < O([[v112,
(42)
where u E R
B.2
The proof of Theorem 3.1.
21

---

Theorem 3.1. Assume that all learnable matrices of UNet in Eq. (16) are independently initialized as Kaiming's s initialization, i.e., WO, ). Then for any p (0, 1), by minimizing the
training loss Eq. (15) of DMs, with probability at least (NN) expl-
- [[mpp)) we have
N Σ j=++1
N
(1 p1²
cillx
cil1x
02
S 11 1 p1+
cil1xe112
Σ j=++1
Ky
+
(43)
where the hidden feature h2 the output of defined in Eq. (16), X Var vI O3er is the input of UNet; ej and 02 are two constants; Ki is the scaling coefficient of the i-th LSC in UNet.
Proof. Let s represent the imput of f. in Eq. (16), i... 8 = ag o a--1 ...... 01 o x,, and let h; = f()). By expanding Eq. (16), can observe that h4 f((s) is mainly composed of the following N terms in the long skip connection:
Ki+1 bi41 04+1 si
K442 b44 li42 o a,f2 0;+1 9;;;
(44)
KN bb4 o b,42 ..42.-000
o by ax ............................... a,+1 8
When N is sufficiently large, h; is predominantly dominated by N long skip connection terms, while the information on the residual term is relatively minor. Hence, we can approximate Eq. (16) as follows:
b J((8) 23
Σ f=++1
Kj TH O 0s, 0SiSN-1,
(45)
where ग.े = bi+1 b142..... bj. Tt == aj o a;-1
a++1.
In practice, the depth I of block og or by is a universal constant, therefore, from Corollary B.9, pp E (0,1) and j fi++1,i N}. with probability least O(N) expl (mmpp)))m)))))) } we have
| 112 O((1 + p112))
(46)
where
hNg o (. (W w r)), r b,42 o by o T/. os,
(47)
s.t. h, fi(s) 2 Ej K2 W o buy. From Lemma B.6, the elements of w will approxi- mately follow a Gaussian distribution (00, (mm))), therefore, EE E(Wh++ hng) O. Moreover, from Lemma B.8,
Var Whit hwr)) = hT, 0( Imxmh hxy = O(m 1)(/k 112 Imxm
(48)
Therefore [47], h4 R Ejmit Wh o hay can be approximately regarded as a Gaussian variable
with zero mean and covariance matrix En (²00m=>)]l,, 12 Imxm= According to Lemma B.5, we have
14 2 x E
Ry whit
hR2)
Σ (=0mm-))lla j-i+1
Σ j=1+1
11. (49)
j=it1
From Eq. (46) and Eq. (49), we have
N +
call c111x1112 p)
Σ #=
K + C2
S 11
calxx112
Σ -=
02
(50)
where Ci denote ((1) in Eq. (49) and the constant 02 compensate for the estimation loss caused by all the approximation calculations in the derivation process.
22

---

B.3 The proof of Lemma 3.2.
Lemma 3.2. For X; = vā:xa + / Er defined in Eq. (14) as a input UNet, é, ~ (O if xo follow the uniform distribution U| U 1, 11, then we have
E E1 (1 2R/00/33)m ((m).
(51)
Proof. For /0,X0 + VI 04 we have
Xx 12 Ex (āt
x R))m)
(52)
where No (j) represents the j-th element of Xp- The equality in Eq. (52) follows from the fact that,
from Eq. (14) and Lemma B.5, we know that x, follows A (xxx; v Vo, (1 (x).). Therefore, ||x, l1 follows chi-square distribution, whose mathematical expectation is Ex (ār Σ5= x6(i + (1 m.)m).
Moreover, since X follow the uniform distribution U[-1, 1], we have
77
X 12
E,
Σ
x²P 2=-U(-1,](a)ddr E2 a.)m
(53)
j=1
E, m/3 3 + E,(1 a²)m = (1 -- 3)m m= (m)..
where Pa~U1- 1.11 is density function of the elements in Xg. The last equality in Eq. (53) follows the fact that, according to the setup of the diffusion model, is generally observed that ān monotonically decreases to as f increases, with ā S 1. Consequently, we can easily deduce that E, is of
the order (1)) For instance, in the setting of DDPM, we have E, = E, (II V 0.27 = 0(1)
0.02i ~
1000
Theorem 3.3. Assume that all learnable matrices of UNet in Eq. (16) are independently initialized as Kaiming's initialization, i.e., NO,, z.). Then for any p E (0, 1], with probability at least OnNN) (1m)), for sample x; in training set, we have
I & 0 (W) ||x 113 Σ tra) (p (1,2,...,744 E (a,, b,)),
(54)
where X = /a + /1 ācde denotes the noisy sample of the s-th sample, ~ N(0.1). N is the number of LSCs, C3 is small constant. n is the size of the training set.
Proof. From Eq. (16) and Eq. (44), we can approximate the output of UNet fo(x;) by weight sum of N forward path g², S < N, i.e.
N Ki g¹
fo(x.)
(55)
1=1
where  denote by b2 o... o aj 0...0 o Xx, and S S N.
Moreover, we use 9f.9 represent the feature map after W9 in forward path g i.e., o6. Wihh4g and h:j o).g.)). Next, we first estimate we: la W 2
23

---

A ww (w)
Σ de, (W)/ 0g' /0g/ 0g) Wo
|0L (W)0g'' 0g/ /ōgb 0gp /ow
DI D(W,)) oe. || 1
(56)
From Lemma B.1 1 and Corollary B.9, with probability at least expl we have
N O(1Io, - ||x+112
VI w.: (W)I
1=1
(O(see((W)) )(1) |
(57)
= ( | Σ
where the constant Os can compensate for the estimation loss caused by all the approximation calculations in Eq. (55). Similarly, by Eq. (42), we can prove the situation while q 6, i.e.,
|| la (3) $ O (W) ||x+ 112 Σe R3 + C3)
B.4 The proof of Theorem 3.4.
Theorem 3.4. For UNet in Eq. (16), assume M, = max o a4/12, S i S N) and fN is Lo-Lipschitz continuous. co is constant related to Mo and Lu- Suppose x is an perturbated input of the vanilla input X, with small perturbation ||*; xx|l2. Then we have
ΣE
[UNet(x;)) UNet(x:)ll S
Ki + co
(58)
where x, Vauxa + VI diec, ~ INO,I), N is the number of the long skip connections.
Proof.
From Eq. (16), we have
UNet(x:) = fo(x,) = b, o [K1 ax + f1 (a1 (a xx)]],
(59)
and
UNet(x(²) = -o(x)')= = ( a + + (a) o xx,))
(60)
Using Taylor expansion, have
UNet(x(") UNet(x;) b, o [M1 ai (xi' xx) + fi (a1 o xx') f1 (al ox)]
= b1 o oa [M (x x1) + (xe* Tv (2) [==0 (61)
Therefore,
[[UNet(x(') - UNet(x:) 1/2 = bbl oar o [K] (xi' - xx) + (xf² - V (z) lxmajox)] 1|2
S cs bb oa1l2 (R1 + IV (2) l==a,ox,11) l12)
(62)
Next, we estimate the || V (2) 1x= 1|2, From Eq. (16), have
24

---

==fi(=) = dab2 (K2 d2 02-2- -2 (az o 2)]
(63)
b2 02
b2 a2 Vu/2 (u) lu=goz:
Therefore, we have
VV (z) lamajox 22 = [1(x2 b2 o d2 o I by
Vaf2 (u) lamaa=x
|| b2 02 o II b2 o 02 o Vaf2 (u) lx= 122
(64)
< ||0 o 020222222222222 (k2 + VV.f22 (u) ox 122..
From Eq. (62) and Eq. (64), let [UNet(x UNet(x;)l3 we have
A es
bb| o
a
|| (m
VV (z)
=x l12)
S cs
bb 01
[K1 + |6 o 02ll2
(K2 + V (u) (xx 112]]
es
(M1 1b|
oai
||2
K2
||bg 02l12 llb ailll || a2l12- V/V (xmoooooo
s es Σ
(K:
1b3 001112)
|bb o aalll Lo]
S es (K1 M3
M
co),
(65)
where co = Mb Lo. and hence we have [NNe(x;;) UNet(x:)ll S es |Σ RgM + co for UNet Eq. (16) and x; = vāixo + v1 OfEt, ~ (0,,1), with small perturbation Es-
Next, we demonstrate through empirical experiments that Mg Theorem 3.4 generally greater than 1. As shown in Fig. 11, we directly computed ||b azll2 for the Eq. (16) on the CIFAR10 and CelebA datasets. We observe that [íb; asl/2 exhibits monotonically increasing trend with training iterations and is always greater than at any iteration. Since Mo max |b, a, a;l12, 1 i S N,
conclude that Mo 2
60
CIFAR 10
CelebA
60
Norm
40
Norm
40
20
20
Norm 1.0
Norm 1.0
0
10
20
30
0
10
20
30
Training iteration 5k
Training iteration 5k
Figure 11: The norm of bb, o aill2 on different dataset.
25

---

C Other experimental observations
Fig. 12 is a supplement to Fig. the main text. We investigate the case of UNet with depth of 12 on the CIFAR10 dataset, where each image represents the visualization results obtained with different random seeds. We observe that the original UNet consistently exhibits feature oscillation, while our
proposed method effectively reduces the issue of oscillation, thereby stabilizing the training of the diffusion model.
(9/10)
10000
10000
Org.
Org
3.01
Org
Org
LS (ours)
2.0
(ours)
(ours)
(ours)
Norm of feanure
CS (ours)
Norm of feature
(ours)
Norn of teature
CS (ours)
Norm of feeature
(ours)
Zo
Lo
t.
20000 lierations
30000
leerations
Iterations
Iterations
Figure 12: More examples for feature oscillation under different random seeds (UNet's 's depth 12).
In the main text, we specify that the domain of x. in CS should be (0,1). This ensures stable forward and backward propagation of the model and enhances the model's robustness input noise, as discussed in Section 3. Otherwise, it would exacerbate the instability of the model training.
In Fig. 13, illustrate the impact of having 6 1 on the performance of the diffusion model. We observe that generally, having K > 1 leads to a certain degree of performance degradation, especially in unstable training settings, such as when the batch size small. For example, in the case of ImageNet64, when the batch size is reduced to 256, the model's training rapidly collapses when 1. These experimental results regarding x align with the analysis presented in Section 3. ImagoNet64-bz-22
CIFARIO-be-123 128
CIFARIO-bZ-66
ImageNet64-bz- 1024
K2 -1.
kF=1.5
K2 115
600
K2 =1.
K2 1121
20-
Kg -1.2 (ours)
CS (ours)
(ours)
400
CS (ours)
FID 10k]
10.
200
10
100
200
300
400
500
100
700
900
400
500
100
150
200
750
503
100
150 Training
Training Heration
Training Eieration
Training iteraiin
Figure 13: The situation that the x;
under different settings.
26


Reviews:
[{"review_id": "oLiKcX0e0x", "reviewer": "Reviewer_KfUJ", "summary": "The paper presents a study and algorithm for scaling UNet's long-range connections such that convergence and stability can be improved. The results are strong in the setup of training diffusion models.", "questions": "Minor:\n- L86: There should be a space between Unet and training.\n- Fig.3 (a): What does the unit m mean on the x-axis? I'm also missing the point this figure is trying to make, as referred to by L161-L163.\n", "limitations": "Yes", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is very well written. The text and graphs are polished and it's easy to follow the idea and contributions.\n- The proposed modification of using an exponential scale over depth is simple and effective, which could be directly applicable to any UNet architecture. \n- The experiments are thorough and the results are strong. The proposed CS and LS algorithms are shown to be better than the vanilla UNet and heuristic $1/\\sqrt(2)$ scaling rule in terms of stability and performance.", "weaknesses": "- Explanation of the algorithm: Although the CS and LS algorithms are shown effective, an intuitive explanation on why they work is missing.\n    - For CS, how is the direction of applying the exponentially decayed scale determined? For instance, can we do $\\kappa_i=\\kappa^{D-i-1}$? In the explanation L251-259, I don't see anything specific to the direction or order. \n    - For LS,  I'm assuming there is no additional supervision to the calibration network. In that case, is there any explanation of why it can discover a reasonable scaling rule (Fig.7)? Since both the weights and the scale contribute to the final output of each block, I think decoupling these would be challenging without additional regularization and supervision. Is there any explanation why LS discovers a decaying scaling curve similar to the CS heuristics, but not an increasing scale?\n- Experiments\n    - For most of the experiments, the results are reported for \"Org\" (no scaling) and \"CS-$1/\\sqrt(2)$\". However, Fig. 5 does not contain results for \"CS-$1/\\sqrt(2)$\". How does the proposed method compare to it in terms of convergence?\n- Significance\n    - Although the paper primarily focuses on training diffusion models, many of the analyses seem general enough to be applicable to other scenarios that use UNets, e.g., depth prediction and segmentations. I wonder if this is true and if there is a reason for only showing results on the generation tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful and very positive comments! In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which may further help solve the current issues.\n\n**`(1)` For the direction of  the exponentially decayed scale in CS**, it is derived by our theory. Specifically,  if we use the reverse direction, namely,  $\\kappa_i = \\kappa^{N-i+1} \\ (\\kappa<1)$, the stability bound in Theorem 3 is extremely large.  The main  term of stability bound in Theorem 3  can be written as \n$\\mathcal{S}_{\\text{r}} =\\sum_{i=1}^N\\kappa_i M_0^i =  \\kappa^NM_0 + \\kappa^{N-1}M_0^2+...+ \\kappa M_0^N$, and could be very large,  since $M_0^N$ is large when $N$ is large and scaling it by a factor $\\kappa$ could not sufficiently control its magnitude (here $M_0>1$, please see Line 230 in manuscript).\n\nIn contrast,  our default setting $\\kappa_i=\\kappa^{i-1}$ of CS  can well control the main  term in stability bound: \n$\\mathcal{S} = \\sum_{i=1}^N\\kappa_i M_0^i =  M_0 + \\kappa^{1}M_0^2+...+ \\kappa ^{N-1}M_0^N$, \nwhere the larger  terms $M_0^{i+1}$ are weighted by smaller coefficients $\\kappa^{i}$. In this way, $\\mathcal{S} $ is much smaller than $\\mathcal{S}_{\\text{r}}$,  which shows the advantages of our default setting. \n\nBesides,  the following Table 1  also  compares the above two settings by using  UViT on Cifar10  (batch size = 64), and shows that our default setting exhibits significant advantages.\n\n| Training step                  | 5k    | 10k   | 15k   | 20k   | 25k   | 30k   | 35k   | 40k   | 45k   |\n|--------------------------------|-------|-------|-------|-------|-------|-------|-------|-------|-------|\n| $\\kappa_i = \\kappa^{N-i+1}$    | 67.26 | 33.93 | 22.78 | 16.91 | 15.01 | 14.01 | 12.96 | 12.34 | 12.26 |\n| $\\kappa_i=\\kappa^{i-1}$ (ours) | 85.19 | 23.74 | 15.36 | 11.38 | 10.02 | 8.61  | 7.92  | 7.27  | 6.65  |\n\n\n**`(2)` There are two possible reasons  for why  LS discovers a decaying scaling curve similar to the CS.**\n- On one hand, from a theoretical view, as discussed in our reply (1), for the $i$-th long skip connection $(1\\leq i \\leq N)$, the learnable $\\kappa_i$ should be smaller to control the magnitude of $M_0^i$ better so that the stability bound, e.g. in Theorem 3, is small. This directly yields the decaying scaling strategy which is also learnt by the scaling network.  \n\n- On the other hand, we can also analyze this observation in a more intuitive manner. Specifically, considering the UNet architecture, the gradient that travels through the $i$-th long skip connection during the backpropagation process influences the updates of both the first $i$ blocks in the encoder and the last $i$ blocks in the UNet decoder. As a result, to ensure stable network training, it's advisable to slightly scale the gradients on the long skip connections involving more blocks (i.e., those with larger $i$ values) to prevent any potential issues with gradient explosion.\n\n\n\n**`(3)` For the convergence curve of $1/\\sqrt{2}$-CS**, we do not include it in Fig. 5 for more clear comparison, since too many curves may lead to unclear comparison and possible misunderstanding. But in Figure 1 of the vanilla manuscript, we have compared with $1/\\sqrt{2}$-CS, and find our methods can better stabilize the UNet training. Moreover, in Table 1, our methods also show better synthesis performance than $1/\\sqrt{2}$-CS. Per your suggestion, we have compared with $1/\\sqrt{2}$-CS in terms of convergence.   Taking the setting about UViT on CIFAR10 as examples, Fig.1 (a) and (b) in the rebuttal PDF  show that our methods reveal much faster convergence behaviors than $1/\\sqrt{2}$-CS.\n\n\n**`(4)` Our analysis are based on diffusion models, and may not be directly applicable to other applications**. For example, our mathematical derivation relies on the particularity of the diffusion model, such as the approximately normal distribution of the network's predictions (i.e., noise). For other typical scenarios, such as image segmentation and depth prediction, the neural network's output consists of segmentation masks or depth maps which do not fully satisfy our analysis needs. While our theories cannot be directly applied to other tasks, there is potential to bridge these gaps through more calibrated analyses. This could be an important direction for future work.\n\n\n**`(5)` For the value of $m$ in Fig.3 (a)**, it denotes the feature dimension of $\\mathbf{x}_t$ (Line 158). The small subfigure in Fig.3(a) is to experimentally verify the distribution of $\\mathbf{x}_0$ required in Lemma 3.2. For the larger subfigure, it is used to validate the conclusion of Lemma 3.2 that $\\mathbb{E}(||\\mathbf{x}_t||_2^2)$ is of the order of $\\mathcal{O}(m)$. In conclusion, the theoretical analysis and findings in Lemma 3.2 are consistent with the real statistical results on three datasets (CIFAR, CelebA, and ImageNet) in Fig.3.\n\n\n**`(6)` For the typos**,  per your suggestions, we will fix them in revision. \n "}, {"review_id": "uv9ufjx4ld", "reviewer": "Reviewer_7ePs", "summary": "The paper discusses the stability issues observed while training UNet in diffusion models, and theorizes on the role of Long Skip Connections (LSCs) in causing this instability. \n\nDiffusion models (DMs), lauded for their ability to model realistic data distributions, involve a forward and a reverse diffusion process. \n\nMost DMs use UNet as their backbone due to its use of LSCs which facilitate long-distance information aggregation and prevent vanishing gradient issues. \n\nHowever, despite the use of shared UNet for predicting injected noise at each step, instability is noticed during training.\n\nThe paper's main contributions revolve around investigating this instability and deriving effective (and efficient) methods to address it. \n\nSpecifically, it is theoretically proven that the coefficients of LSCs in UNet significantly affect the stability of forward and backward propagation as well as the robustness of UNet. \n\nThe paper also proposes two coefficient scaling methods, Constant Scaling (CS) and Learnable Scaling (LS), designed to adjust the coefficients of LSCs in UNet for training stability.\n\nCS involves setting the coefficients as a series of exponentially-decaying constants while LS uses a small shared network to predict the scaling coefficients for each LSC.", "questions": "- ", "limitations": "- ", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- One of the key strengths of the paper is its rigorous theoretical analysis. It provides a comprehensive understanding of the instability of UNet in diffusion models by focusing on the significant impact of the coefficients of Long Skip Connections (LSCs). It's an important insight that broadens the understanding of UNet's performance in diffusion models.\n\n\n- The paper proposes two novel scaling methods, Constant Scaling (CS) and Learnable Scaling (LS), designed to enhance the stability of the UNet training process. These simple methods address the identified problem and provide tangible ways to improve training stability, indicating a proactive approach to problem-solving.\n\n", "weaknesses": "- The paper primarily focuses on the role of Long Skip Connections (LSCs) in causing training instability in UNet. While LSCs are a significant part of the model, the authors might have explored other potential contributing factors to this instability as well. Broadening the scope of their investigation could potentially have led to a more comprehensive understanding of the issue.\n\n\n- The placement of Figure 1 doesn't align well with its corresponding description, causing some disconnection. Ideally, figures on the first page should offer readers a quick, intuitive understanding of the paper's content. In its present state, Figure 1 seems to lack sufficient information to accomplish this. If this positioning is due to space constraints, please consider relocating it to a more suitable page.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful and positive comments.  In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which may further help solve the current issues.\n\n\n**`(1)` In addition to LSCs, we  have also explored other potential contributing factors.** We found that the decoder of UNet is also a significant part of the model. We have extensively explored how to apply scaling strategy  to the output of each block in the decoder, hoping to stabilize model training. However, we find that this approach cannot achieve satisfactory performance for diffusion models, and are exploring the reasons behind. \n\n\n \n**`(2)` For the layout of the figures and the corresponding text**,  per your suggestion, we will arrange them in the revision so that the figures and the corresponding text are close,  improving the readability."}, {"review_id": "bwdhdCYisn", "reviewer": "Reviewer_Ttzr", "summary": "In this paper, the authors focus on the challenge of the instability arising from the commonly adapted U-Net architecture for diffusion models. In particular, the authors start by theoretically analyzing the influence of the coefficients of long skip connects in U-Net-based diffusion models, specifically on the stability of the forward and backward propagation, along with the robustness of the network. Motivated by the theoretical discussions, the authors propose two corresponding scaling methods that adapt the coefficients of the long skip connections for more stable training, including a constant scaling method and a learnable scaling method, which is straightforward and ideally similar to that in dynamic neural networks and meta learning, but backed by solid theoretical analysis. Experiments on four datasets validate the effectiveness of the proposed two forms of scaling manners, with key codes provided in the supplement.", "questions": "Please kindly refer to the weakness section for more details.", "limitations": "Limitations are discussed in Sect. 6 of the paper.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed constant and learnable scaling methods are very straightforward, but with solid theoretical analysis and guarantee. The results shown in Fig. 1 are also promising that successfully alleviate the oscillations.\n- Extensive ablation studies including the robustness of LS to network architecture are provided.\n- Key codes are provided in the supplementary material for handy reproducibility.\n", "weaknesses": "- Method: The second part of learnable scaling method in this paper is algorithmically related to meta learning and dynamic neural networks. However, there is a lack of discussions on related works. Also, it would be great if the authors could discuss some alternative ways for learnable scaling by borrowing the ideas from meta learning and dynamic neural networks.\n- Experiments: The authors are encouraged to provide some qualitative comparative results with existing methods, as a complement to the quantitative results reported at, for example, Tab. 1.\n\nThe instability issue in diffusion models with U-Net is not my primary research area. As such, I may not be capable of correctly evaluating the novelty of this paper. I will refer to the comments of other reviewers for my final justification.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful and very positive comments! In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which may further help solve the current issues.\n\n**`(1)` Per your suggestion, here we discuss some alternative ways for learnable scaling** by borrowing the ideas from meta-learning and dynamic neural networks.\n\n- For dynamic neural networks, the modules like LS can be viewed as attention modules. The work [1] suggests that such modules can be considered as adaptive feature-regulating dynamic neural networks. So by replacing or designing improved attention modules, we can effectively borrow ideas from dynamic neural networks to enhance the performance of diffusion models. Indeed, in our manuscript, we have already explored different attention modules. For example, Table 2 in the main text also presents the performance of some other alternative modules. \n\n- For meta-learning. We can loosely follow the steps inspired by MAML [2]. First, We can regard the scaling module $\\mathcal{M}$ as a meta-learner in meta-learning, and regard each sample $\\mathcal{T}_ i$  as a task. For the inner loop, given a task (sample) $\\mathcal{T}_ i$, the vanilla diffusion loss can be used to tune $\\mathcal{M}$  via a step of gradient descent to adapt $\\mathcal{M}$  to the task $\\mathcal{T} _ i$  and obtain a new task-specific scaling module $\\mathcal{M}_ {T_ i}$.  In the outer loop, we feed the sample $\\mathcal{T}_ i$ into the diffusion model with a task-specific scaling module $\\mathcal{M}_ {T_ i}$ again, and update the diffusion model and the scaling module $\\mathcal{M}$ via gradient. The benefits of this meta-learning is that we can learn a meta model $\\mathcal{M}$ as the scaling module, and  adapt it to a specific sample for generating  LSC weights. Though this idea is promising,this meta-learning method requires to compute the Hessian matrix, and thus suffers from high computational cost which is not very suitable for diffusion models whose model size is often large. \n      \n   In the revision, we will supplement the above discussion, and try to provide further validation experiments and analysis. \n\n[1] Han Y et al. Dynamic neural networks: A survey[J]. TPAMI, 2021\n\n[2] Finn C et al. Model-agnostic meta-learning for fast adaptation of deep networks[C]. ICML 2017.\n\n**`(2)` For qualitative comparative results,**  we have already provided some visualizations of our method (in Tab. 1) in Appendix for qualitative analysis. From the visual results (namely, the generated images), our methods can generate highly qualified images. Per your suggestion, we will include more comparisons with other methods in the revision."}, {"review_id": "LOsgXzU3n1", "reviewer": "Reviewer_xZoY", "summary": "This paper proposes to scale the skip connection in diffusion model Unets by an exponential factor. The authors should that the feature norms of vanilla Unets oscillate across batches, and that their proposed method results in much smaller feature output oscillations. They conjecture that this method stabilizes training, and prove bounds on the oscillation which their method decreases. The proposed method shows faster convergence and shows improvements across multiple models, datasets, model sizes.", "questions": "1. The authors derive bounds for oscillations (eq 4), and they show that their methods decreases this bound, but how tight is this bound? Comparing the minimum and maximum norms predicted by this formula, with that in Figure 1., is important to show if this bound is indeed useful, or perhaps their proposed method has some other side effect that results in the performance improvements.\n1. Direct evidence of parameter oscillation (mentioned in line 47, 172 can be shown), perhaps by comparing a per-parameter Adam style \"m\" and \"v\" across the baseline method and the author's method. If oscillation of parameters is indeed smaller, one would perhaps expect the norm of \"m/v\" to be larger in the author's case. Is this indeed the case?\n1. Given that the exponential change in scaling with increasing total number of layers, does the performance of the method suffer compared to $\\frac{1}{\\sqrt{2}}-CS$ as number of layers increases to more depths (as large depth will cause the skip connection to essentially be scaled to zero)?\n\nMinor presentation issues (authors are not expected to respond to these, as they are minor issues that can be fixed later):\n1. Line 85 \"serir\" should perhaps be \"series\"?\n1. In table 1, comparisons to author's method \"x+LS/CS\" should perhaps be placed directly below \"x\", to make it easier to compare scores.\n1. The values for CS for table 1 are mentioned in the supplementary, but perhaps they should be mentioned at line 272 or in 315.\n1. Quotes around 'CS.py' in supplementary above Figure 2 are inverted.\n1. $\\frac{1}{\\sqrt{2}}-CS$ reads as \"1 by square root 2 minus CS\", and for a moment I was trying to think what $CS$ are the authors subtracting from  $\\frac{1}{\\sqrt{2}}$. Perhaps a better notation for this may be something along the lines of $(CS)\\frac{1}{\\sqrt{2}}$?", "limitations": "-", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors method follows prior work on scaling block output and/or the skip connection, where such scaling methods have been shown to increase stability/robustness and result in Lipschitz continuous models.\n1. The learnt scales in LS method seems to somewhat mirror the author's theoretical exponential CS method, providing further validation.\n1. The authors method results in faster convergence and shows improvements across multiple models, datasets, model sizes, and compares favourably to some other scaling methods.\n1. The author's proposed method can be very easily adapted to existing models.", "weaknesses": "1. The authors show that feature norms oscillate across samples, and they say this implies parameters must also be oscillating (line 47, line 172). While the authors demonstrate that feature norms do oscillate (and I am willing to agree the same may be true of gradient norms), this may not directly cause the parameters/updates to oscillate.\n1. The authors works (and proposed exponential CS solution) is remarkably similar to [1], where a scaling of $b^l$ is proposed for the block output (compared to this paper's scaling of the skip connection by a similar value). An experimental comparison with this method should be added to Table 3, or perhaps the similarities/differences discussed.\n1. Discussion of prior related works should be expanded. For example, [2] proposes scaling of the skip connection and/or the block output, and has similar exponential bounds.\n1. The constants in theorem 3.1 are ignored (assumed to be 1) in line 270 to derive the values for the scaling parameter. If this constant is for example 10 or 0.1, that will dramatically change the proposed values for CS. Further investigation into this constant is required.\n\n\n[1] Hanin, B., & Rolnick, D. (2018). How to Start Training: The Effect of Initialization and Architecture. Advances in Neural Information Processing Systems, 31, 571–581.\n[2] Balduzzi, D., Frean, M., Leary, L., Lewis, J.P., Ma, K.W. &amp; McWilliams, B.. (2017). The Shattered Gradients Problem: If resnets are the answer, then what is the question?. <i>Proceedings of the 34th International Conference on Machine Learning</i>, in <i>Proceedings of Machine Learning Research</i> 70:342-350 Available from https://proceedings.mlr.press/v70/balduzzi17b.html.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful and positive comments! In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which may further help solve the current issues.\n\n**`(1)` For parameter oscillation,**  we indeed observe it as discussion below, and our approach effectively mitigates this oscillation. \n\n As shown in Fig.2 in the attached PDF, taking UViT on CIFAR-10  as an example, we select several parameters from both the encoder and decoder's final layers. We then computed the average sliding variance (using a window size of 5) of these parameters during training for the three methods (org, LS, CS). This approach provides more direct observation of parameter oscillation compared to monitoring the norm of \"m/v.\" From the visualization results, one can observe that parameter oscillation does occur, and our method effectively mitigates this oscillation to stabilize model training.\n\n**`(2)` For the scaling method  on the block output in the work [1],** we compare it with ours in  Table 2 in the attached PDF, where the scaling method in [1] is denoted as $b^l$. We follow three settings from  the work [1]  by respectively setting $b=0.90, 0.75, 0.50$, and find that this kind of scaling method on the block output cannot achieve good  performance improvement.\n Therefore, for stable training of the diffusion model, we believe that applying scaling on long skip connections, similar to our proposed CS, would be more effective. We will include this comparison into our revision. \n\n**`(3)` In the revision, we will discuss more related works.** For the work [2], our work mainly differs from it in two notable differences. \n\t\n- (**Architecture**) Our method focuses on the UNet architecture, particularly the long skip connections, which are not considered in [2]. Most of the discussions and analyses about ResNet (Kaiming He 2016) from [2] cannot be directly applied to UNet. \n\t\n- (**Scaling location**) The work [2] also primarily focuses on scaling at the block output, rather than on scaling at skip connections. As mentioned in section 3.3 of paper [2], scaling at skip connections may not work well in practice for ResNet (Kaiming He 2016). On the contrary, in the scenario of UNet and diffusion models, we notice that scaling at long skip connections is more effective than scaling at the block output.\n\t\nOur work's analysis and conclusions contribute to an expanded understanding of scaling strategies in neural network architectures within the community. In the revision, we will include the above discussion to further strengthen our work. \n\n**`(4)` For the constant $c_1$ in theorem 3.1,**\nas mentioned in Line 270-272, we use a rough estimate for constant $c_1 \\approx 1$ which is valid in many experiments. \n Moreover, per your suggestion,  we conduct statistical analysis on different datasets such as CIFAR10, CelebA, and ImageNet64, and find that the real range of the constant $c_1$ is  $(0.752, 1.060)$ for analysis in Line 270-272, which is close to our rough estimate $c_1\\approx 1$ and will not dramatically change the proposed values for CS.  In the revision, we will modify this section to enhance our writing.\n\n**`(5)` For the oscillation bound in Eq (4),**\ndue to the use of some approximations and consideration of worst-case scenarios in Eq (4), our bound may not be very tight. And it's also not easy to be visualized in Figure 1 since some parameters are hard to measure during training directly and quantitatively, e.g., $c_2$. In this paper, Eq(4) serves as a qualitative estimation of feature oscillations, providing inspiration and explanation of our proposed scaling method from a feature map perspective. This indicates that the estimation of Eq(4) is meaningful and insightful. \nMoreover, we are the first to conduct such an analysis on the diffusion model of the UNet neural network architecture. In future work, we will attempt more sophisticated analytical approaches to obtain tighter bounds.\n\n\n**`(6)` The comparison between    $\\frac{1}{\\sqrt{2}}$-CS and ours on very deep network.**  If the number of layers is large enough, indeed, our scaling will tend to zero. However, this situation is relatively rare in the practical diffusion model tasks using UNet. On the most popular image or text generation tasks, a UNet usually doesn't have a large number of long skip connections. Therefore, in this practical real-world setting,  our scaling factor will not be scaled to zero, and our methods are more effective than $\\frac{1}{\\sqrt{2}}$-CS.\n\t\nAs technology advances, our proposed scaling strategy will need further refinement to handle potential cases where diffusion models may be employed in extremely deep networks. Additionally, if a long skip connection approaches zero but the performance can be maintained, it may not necessarily be a bad thing, as these connections can be almost removed to save memory and improve model inference speed without sacrificing performance. \n\n\n**`(7)` For the minor presentation issues**,  per your suggestions, we will fix them in revision. \n\t\n[1] Hanin, B. et al. How to Start Training: The Effect of Initialization and Architecture. NeurIPS 2018.\n\t\n[2] Balduzzi, D. et al. The Shattered Gradients Problem: If resnets are the answer, then what is the question? ICML 2017."}, {"review_id": "KbDFwJBVro", "reviewer": "Reviewer_VLJa", "summary": "U-Net is the most popular neural network backbone for diffusion models. In U-Net, the long skip connection (LSC) links the long-distant information near to the input and the intermediate network outputs. However, they suffers from unstable training, which is resolved by scaling down the LSC coefficients. This paper addresses the theoretical aspects of stabilizing LSC, robustifies the training process, and accelerates this. Both hand-crafted and learning-based parameterization are proposed.", "questions": "* This work focuses on scaling the skip connection coefficients U-Net-based diffusion models. Does this technique also hold for other U-Net based methods rather than diffusion models?\n* Can you add some figures or tables that addresses the optimal value of $\\kappa$ in CS cases, like the $\\kappa$ from LS in Figure 7?\n* Please write the full name of IE and SE in line 330. \n* Please address the number of sampling steps in each methods.\n* Does parameterizing LC to some additional module (Figure 4) and training the network end-to-end adequately learn $\\kappa_i$? \n\n===========\n\nCorrections\n* Line 85: serir --> series\n* Line 86: UNettraining --> UNet training", "limitations": "The authors adequately addressed the limitations and introduced further directions.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The organization of the paper is clear, and the figures of motivating experiments are well demonstrated.\n* The method and simple yet efficient, improving both the training speed and sampling performance while requiring only some hand-crafted hyperparameters (CS) or lightweight learning module (LS). ", "weaknesses": "* This work does not deal with more recent diffusion model baselines such as EDMs, which is much more powerful than the methods that are compared in this work.\n* Prior, related works are not introduced, especially for those who are not familiar with the scaling of skip connections in other network architectures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful and positive comments! In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which may further help solve the current issues.\n\n**`(1)` About the baseline EDMs.**  Thank you for your suggestion. Since our method focuses on the neural network architecture for the diffusion model,  it could be applicable to EDMs which also uses UNet for training. However, this needs to be thoughtfully tested. Since EDMs require  much longer training time, such as 16 GPU days for CIFAR10 and over 300 GPU days for ImageNet64,  it is really challenging for us to finish the experiments within the rebuttal phase. In the revision, we will try our best to evaluate our method on EDMs and report the results in revision.\n\n**`(2)` About the related works.** Thank you for your suggestion.  At present, due to the limited space (9 pages), we have to briefly introduce the related works in Line 22-26 and Line 101-128, and spend much space to introduce three theoretical analysis from three different aspects, and our proposed methods. We very much appreciate your suggestions about   introducing more related works, including the  previous scaling methods of skip connections, and will try to discuss more in the revision since  the final version often allows us to use an extra page.\n\n**`(3)` Our technique may not directly improve other UNet based tasks.** Our analysis is based on diffusion models. For example, our mathematical derivation relies on the particularity of the diffusion model, such as the approximately normal distribution of the network's predictions (i.e., noise). For other typical UNet-based scenarios, such as image segmentation and depth prediction, the neural network's output consists of segmentation masks or depth maps which may not fully satisfy our analysis needs. \n\nWhile our theories cannot be directly applied to other tasks, our technology, especially the experimental setup based on our theories, may not be directly transferable to other UNet-based scenarios. In the future, we will test our proposed LS and CS in some common UNet-based scenarios, such as image segmentation, to further enhance the versatility of our approach.\n\n**`(4)` The optimal value of $\\kappa$ in CS cases** please refer to rebuttal PDF Fig.1 (c). The values of $\\kappa$ in CS are 0.5 (for CIFAR10), 0.5 (for CelebA), 0.8 (for MS-COCO) and 0.95 (for ImageNet64).\n\n\n**`(5)` About the full name of IE and SE**. Thanks for your suggestion. We will add the full name of IE (Instance enhancement batch normalization [1]) and SE (Squeeze-and-excitation networks [2]) in the revision.\n\n\n**`(6)` About the sampling steps in each methods**. We follow the default and official sampling settings in UViT. Specifically, for CIFAR10 and CelebA64, the experiments used the \"Euler-maruyama-sde\" method with 1000 sampling steps for sampling. On the other hand, for MS-COCO and ImageNet64, the experiments used the \"DPM solver\" method with 50 sampling steps.\n\n**`(7)` About the LC parameterizing**. Apologies, we may not fully understand your question. Let me try to clarify:\n\t\n- If you're asking whether LS can be used to end-to-end adaptively learn $\\kappa_i$, then according to the explanation in Section 4.2, our LS is originally designed for adaptive learning of $\\kappa_i$, rather than being a fixed value like CS.\n\t\n- If you're inquiring about the feasibility of applying LS with additional (other) modules for adaptive learning of $\\kappa_i$, Table 2 explores the use of alternative modules for learning $\\kappa_i$. However, the results indicate that the current structure already achieves satisfactory performance.\n \nI hope my response addresses your issue. If not, please let me know. \n\n\n**`(8)` For the typos**,  per your suggestions, we will fix them in revision. \n\t\n[1] Senwei et al. Instance enhancement batch normalization: an adaptive regulator of batch noise. AAAI 2020.\n\n[2] Jie Hu et al. Squeeze-and-excitation networks. TPAMI 2017.\n\t\n\t"}, {"review_id": "JRqg30DiwV", "reviewer": "Reviewer_Ehww", "summary": "They state that diffusion models using Unet suffer from unstable training\nand oscillations of features and gradients. They also state that while this is sensitive to coefficients related to scaling the skip connections of Unet. They set out to provide an explanation and more robust scaling methods for the skip connections to address this. Their methods are constant scaling based exponentially on the depth of the skip connection, and learnable scaling. The methods seem to consistently reduce FID versus the same model with regular scaling.", "questions": "N/A", "limitations": "N/A", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I believe this paper’s contribution is simple and meaningful, and written\nclearly. The experiments seem thorough to changes made and the questions you would have reading it. I cannot fully speak either way to novelty in this area, but it seems like a sound paper.\n", "weaknesses": "I did not see meaningful weaknesses in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**`(1)` Thank you for your encouragement and taking the time to review our article.** We will further improve our paper based on the suggestions of other reviewers. "}]

Author rebuttal:
We provide the necessary charts for the rebuttal stage in the attached PDF. Reviewers are kindly requested to refer to them.

Meta-review: This paper proposes two scaling methods to stabilize the training of UNets in Diffusion Models, and provides both theoretical and empirical support. The paper received unanimous accept recommendations (2 accepts, 3 weak accepts, and 1 borderline accept) from reviewers. The paper's strengths include a simple and effective approach, rigorous theoretical analysis, and thorough experimental analysis and good results. Most of the initial weaknesses surrounding discussion of related work, questions regarding the approach, and experiments were addressed by the rebuttal. The ACs agree with the reviewers that this paper can make a good contribution to the machine learning and computer vision communities; in particular, the proposed approach is simple and effective, and has potential to be applicable to various UNet architectures. The ACs recommend acceptance.
Decision: Accept (poster)