Title: ScaleLong: Towards More Stable Training of Diffusion Model via Scaling Network Long Skip Connection
Venue: NeurIPS 2023
Accepted: yes
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

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 four famous datasets show that our methods are superior to stabilize training, and yield about 1.5 × training acceleration on different diffusion models with UNet or UViT backbones. Click here for Code.

1 Introduction

Recently, diffusion models (DMs) [1-9] have become the most successful generative models because of their superiority on modeling realistic data distributions. The methodology of DMs includes a forward diffusion process and a corresponding reverse diffusion process. For forward process, it progressively injects Gaussian noise into a realistic sample until the sample becomes a Gaussian noise, while the reverse process trains a neural network to denoise the injected noise at each step for gradually map-

LS (ours)

CS (ours)

Org

-CS

LS (ours)

CS (ours)

0

30000

20000

10000

Iterations

0

30000

20000

10000

Iterations

Norm of feature

×10000

×10000

0.5

1.0

0

1.5

2.0

7.0

5.0

3.0

1.0

Depth = 12

Depth = 16

Org

-CS

ping the Gaussian noise into the vanilla sample. By decomposing the complex generative task into a 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 [11], on image [12-15], 3D [16, 17] and video [18-21] data and beyond.

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 (LSCs) to connect the distant and symmetrical network blocks in a 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] empirically find that 1 / √ 2 -constant scaling ( 1 / √ 2 -CS), that scales the coefficients of LSC (i.e. { κ i } in Fig. 2) from one used in UNet to 1 / √ 2 , 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 a principle way.

Contribution. In this work, we address the above two fundamental questions and contribute to deriving some new results and alternatives for UNet. Particularly, we theoretically analyze the above training instability of UNet 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 / √ 2 -scaling technique. Inspired by our theory, we propose the framework ScaleLong including two novel-yet-effective 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.

Our first contribution is proving that the coefficients of LSCs in UNet have big effects on the stableness of the forward and

Long skip

connection

Forward

Block

Input

Output

Scaling

coefficients

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 O ( m ∑ N i =1 κ 2 j ) , where { κ i } denotes the scaling coefficients of N LSCs and the input dimension m of UNet is often of hundreds. Since standard UNet sets κ i = 1 ( ∀ i ) , the above oscillation range becomes O ( mN ) and is large, which partly explains the oscillation and instability of hidden features as observed in Fig. 1. The 1 / √ 2 -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 O ( m ∑ N i =1 κ 2 i ) . In standard UNet, this bound becomes a large bound O ( mN ) , 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 / √ 2 -CS 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 ( ∑ N i =1 κ i M i 0 ) with a model-related constant M 0 > 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 / √ 2 -CS technique improves the robustness.

Inspired by our theory, we further propose a 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) learnable scaling method ( LS ). CS sets the coefficients { κ i } as a serir of exponentially-decaying constants, i.e., κ i = κ i -1 with a contant κ ∈ (0 , 1] . As a result, CS greatly stabilizes the UNettraining by largely reducing the robustness error bound from O ( M N 0 ) to O ( M 0 ( κM 0 ) N -1 ) , which is also O ( κ N -2 ) times smaller than the bound O ( κM N 0 ) of the universal scaling method, i.e., κ i = κ , like 1 / √ 2 -CS technique. Meanwhile, CS can ensure the information transmission of LSC without degrading into a feedforward network. Similarly, the oscillation range

of hidden features and also gradient magnitude can also be controlled by our CS. For LS, it uses a 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, as 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 CIFAR10 [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.5 × 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 i } T i =1 by repeatedly adding Gaussian noise to a sample x 0 until attatining x T ∼ N ( 0 , I ) . This noise injection process, a.k.a. the forward process, can be formalized as a Markov chain q ( x 1: T | x 0 ) = ∏ T t =1 q ( x t | x t -1 ) , where q ( x t | x t -1 ) = N ( x t | √ α t x t -1 , β t I ) , α t and β t depend on t and satisfy α t + β t = 1 . By using the properties of Gaussian distribution, the forward process can be written as

where ¯ α t = ∏ t i =1 α i . Next, one can sample x t = √ ¯ α t x 0 + √ 1 -¯ α t ϵ t , where ϵ t ∼ N ( 0 , I ) . Then DMadopts a neural network ˆ ϵ θ ( · , t ) to invert the forward process and predict the noise ϵ t added at each time step. This process is to recover data x 0 from a Gaussian noise by minimizing the loss

UNet-alike Network . Since the network ˆ ϵ θ ( · , 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 in (3) that uses LSCs to connect distant parts of the network for long-range information transmission and aggregation

where x ∈ R m denotes the input, a i and b i ( i ≥ 1) are the trainable parameter of the i -th block, κ i > 0 ( i ≥ 1) are the scaling coefficients and are set to 1 in standard UNet. f N 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 i ◦ x = W a i l ϕ ( W a i l -1 ...ϕ ( W a i 1 x )) with a ReLU activation function ϕ and learnable matrices W a i j ∈ R m × m ( j ≥ 1) . Moreover, let f N also have the same architecture, i.e. f N ( x ) = W f i l ϕ ( W f i l -1 ...ϕ ( W f i 1 x )) . Following [36], the above UNet has absorbed the fusion operation of the feature a i ◦ x and f i ( a i ◦ 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 Transformer, i.e. x i +1 = x i + κ i f i ( x i ) where f i ( x i ) 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 ( κ i = 1 , ∀ i ) often suffers from the unstable training issue in diffusion models (DMs), while scaling the long skip connection (LSC) coefficients κ i ( ∀ i ) in UNet can help to stabilize the training. However, theoretical understandings of the instability of U-Net and also the effects of scaling LSC coeffeicents remain absent yet, hindering the development of new and more advanced DMs in a 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

1

5

9

w/ extra noise

w/o extra noise

FID-50k↓

Extra noise       = 0.01

0

5

10

15

w/ extra noise

w/o extra noise

FID-50k↓

1.6×

1.4×

1.5×

LS (ours)

CS (ours)

Org UViT

LS (ours)

CS (ours)

Org UViT

Extra noise       = 0.02

1

7

5

9

3

Values

Frequency

10

20

30

10

20

30

10

20

30

Frequency

Frequency

The distribution of

2.49%

0.93%

2.17%

(a)

(b)

(c)

100

200

300

100

200

300

100

200

300

Frequency

Frequency

Frequency

0

0.6m

0.4m

1.0m

0.2m

Values

0.8m

The distribution of

0.0

1.0

-1.0

60

20

distribution

CIFAR10

CelebA

ImageNet

60

20

60

20

distribution

distribution

CIFAR10

CelebA

ImageNet

13

3.1×

2.9×

2.6×

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 a more stable UNet as introduced in Section 4.

3.1 Stability of Forward Propagation

Here we study the stableness of the hidden features of UNet. This is very important, since for a certain layer, if its features vary significantly along with the training iterations, it means there is at least a layer 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, 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 initialization, i.e., N (0 , 2 m ) . Then for any ρ ∈ (0 , 1] , by minimizing the training loss Eq. (2) of DMs, with probability at least 1 -O ( N ) exp[ -Ω( mρ 2 )] , we have

where the hidden feature h i is the output of f i defined in Eq. (3), x t = √ ¯ α t x 0 + √ 1 -¯ α t ϵ t is the input of UNet; c 1 and c 2 are two constants; κ i 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 i of UNet can be both lower- and also upper-bounded, and its biggest oscillation range is of the order O ( ∥ x t ∥ 2 2 ∑ N j = i κ 2 j ) which is decided by the scaling coefficients { κ i } and the UNet input x t . So when ∥ x t ∥ 2 2 or ∑ N i =1 κ 2 i is large, the above oscillation range is large which allows hidden feature h i 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 t ∥ 2 2 is around feature dimension m in DDPM.

Lemma 3.2. For x t = √ ¯ α t x 0 + √ 1 -¯ α t ϵ t defined in Eq. (1) as a input of UNet, ϵ t ∼ N (0 , I ) , if x 0 follow the uniform distribution U [ -1 , 1] , then we have

Indeed, the conditions and conclusions of Lemma 3.2 hold universally in commonly used datasets, such as CIFAR10, CelebA, and ImageNet, which can be empirically valid in Fig. 3 (a). In this way, in standard UNet whose scaling coefficients satisfy κ i = 1 , the oscillation range of hidden feature ∥ h i ∥ 2 2 becomes a very large bound O ( Nm ) , where N denotes the number of LSCs. Accordingly, UNet suffers from the unstable forward process caused by the possible big oscillation of the hidden feature h i which accords with the observations in Fig. 1. To relieve the instability issue, one possible

solution is to scale the coefficients { κ i } of LSCs, which could reduce the oscillation range of hidden feature ∥ h i ∥ 2 2 and yield more stable hidden features h i to benefit the subsequent layers. This is also one reason why 1 / √ 2 -scaling 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 a 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 to 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.e., W = [ W a 1 l , W a 1 l -1 , ..., W a 1 1 , W a 2 l ..., W b 1 2 , W b 1 1 ] . Assume the training loss is 1 n ∑ n s =1 ℓ s ( W ) , where ℓ s ( 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 subscript s of ∇ ℓ s ( 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 initialization, i.e., N (0 , 2 m ) . Then for any ρ ∈ (0 , 1] , with probability at least 1 -O ( nN ) exp[ -Ω( m )] , for a sample x t in training set, we have

where x t = √ ¯ α t x 0 + √ 1 -¯ α t ϵ t denotes the noisy sample of the s -th sample, ϵ t ∼ N ( 0 , I ) , N is the number of LSCs, c 3 is a small constant.

Theorem 3.3 reveals that for any training sample in the training dataset, the gradient of any trainable parameter W q p in UNet can be bounded by O ( ℓ ( W ) · ∥ x t ∥ 2 2 ∑ N j = i κ 2 j ) (up to the constant term). Since Lemma 3.2 shows ∥ x t ∥ 2 2 is at the order of O ( m ) in expectation, the bound of gradient norm becomes O ( mℓ s ( W ) ∑ N j = i κ 2 j ) . 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 ( κ i = 1 ∀ i ) would become O ( mNℓ s ( 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 { κ i } to 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.

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 LSCs 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 ϵ t ∼ 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 ϵ t during the training process, and greatly affects the training. It 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 ϵ t from the UNet input x t ∼ N ( x t ; √ ¯ α t x 0 , (1 -¯ α t ) I ) .

Assume an extra Gaussian noise ϵ δ ∼ N (0 , σ 2 δ I ) is injected into x t which yields a new input x ′ t ∼ N ( x t ; √ ¯ α t x 0 , [(1 -¯ α t ) + σ 2 δ ] I ) . Accordingly, if the variance σ 2 δ of extra noise is large, it can dominate the noise ϵ t in x ′ t , and thus hinders predicting the desired noise ϵ t . 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 ϵ δ with different σ 2 δ 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 M 0 = max {∥ b i ◦ a i ∥ 2 , 1 ≤ i ≤ N } and f N is L 0 -Lipschitz continuous. c 0 is a constant related to M 0 and L 0 . Suppose x ϵ δ t is an perturbated input of the vanilla input x t with a small perturbation ϵ δ = ∥ x ϵ δ t -x t ∥ 2 . Then we have

Theorem 3.4 shows that for a perturbation magnitude ϵ δ , the robustness error bound of UNet in DMs is O ( ϵ δ ( ∑ N i =1 κ i M i 0 )) . For standard UNet ( κ i = 1 ∀ i ), this bound becomes a very large bound O ( NM N 0 ) . This implies that standard UNet is sensitive to extra noise in the input, especially when LSC number N is large ( M 0 ≥ 1 in practice, see Appendix). Hence, it is necessary to control the coefficients κ i of LSCs which accords with the observations in Fig. 3 (c). Therefore, setting appropriate scaling coefficients { κ i } 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 at the beginning of Section 3.3.

4 Theory-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 / √ 2 -scaling technique in [27-31] which scales all LSC coefficients κ i from one in UNet to a constant, and is called 1 / √ 2 ( 1 / √ 2 -CS for short). Here we will use these theoretical results to design a novel framework ScaleLong 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 { κ i } 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 { κ i } of LSCs in UNet:

When i = 1 , the scaling coefficient of the first long skip connection is 1, and thus ensures that at least one κ i 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 κ i = κ , e.g. κ = 1 / √ 2 in the 1 / √ 2 -CS method, in terms of reducing the training oscillation of UNet. We first look at the robustness error bound in Theorem 3.4. For our CS, its bound is at the order of O ( M 0 ( κM 0 ) N -1 ) which is O ( κ N -2 ) times smaller than the bound O ( κM N 0 ) of the universal scaling method. This shows the superiority of our CS method. Moreover, from the analysis in Theorem 3.1 , our CS method can effectively compress the oscillation range of the hidden feature h i of the i -th layer to O ( κ 2 i m ) . Compared with the oscillation range O ( Nm ) of standard UNet, CS also greatly alleviates the oscillation since κ 2 i ≪ N . A similar analysis on Theorem 3.3 also shows that CS can effectively control the gradient magnitude, helping stabilize UNet training.

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 (CIFAR10, CelebA and ImageNet) to estimate κ . Firstly, Theorem

3.1 reveals that the norm of UNet output is of order O ( ∥ x t ∥ 2 2 ∑ N j =1 κ 2 j ) . Next, from the loss function of DMs, i.e., Eq. (2), the target output is a noise ϵ t ∼ N (0 , I ) . Therefore, we expect that O ( ∥ x t ∥ 2 2 ∑ N j =1 κ 2 j ) ≈ ∥ ϵ t ∥ 2 2 , and we have ∑ N j =1 κ 2 j ≈ O ( ∥ ϵ t ∥ 2 2 / ∥ x t ∥ 2 2 ) . Moreover, Fig. 3 (b) shows the distribution of ∥ ϵ t ∥ 2 2 / ∥ x t ∥ 2 2 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 ∥ ϵ t ∥ 2 2 / ∥ x t ∥ 2 2 . They may cause the output order O ( ∥ x t ∥ 2 2 ∑ N j =1 κ 2 j ) to deviate more from the ∥ ϵ t ∥ 2 2 , leading to unstable loss and impacting training. Therefore, it is advisable to appropriately increase ∑ N j =1 κ 2 j to address the influence of long-tail samples. [61-65]. So the value of ∑ N j =1 κ 2 j should be between the mean values of ∥ ϵ t ∥ 2 2 / ∥ x t ∥ 2 2 , i.e., about 1.22, and a 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 κ can be adjusted around this range.

4.2 Learnable Scaling Method

In Section 4.1, 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 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

Avg token

Calibration

Feature

Element-wise

multiplication

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.

As shown in Fig. 4, LS designs a calibration network usually shared by all long skip connections to predict scaling coefficients { κ i } . Specifically, for the i -th long skip connection, its input feature is x i ∈ R B × N × D , where B denotes the batch size. For convolution-based UNet, N and D respectively denote the channel number and spatial dimension ( H × W ) of the feature map; for transformer-based UViT [36], N and D respectively represent the token number and token dimension. See more discussions on input features of other network architectures in the Appendix. Accordingly, LS feeds x i ∈ R B × N × D into the tiny calibration network ζ ϕ parameterized by ϕ for prediction:

where GAP denotes the global average pooling that averages x i along the last dimension to obtain a feature map of size B × N × 1 ; σ is a sigmoid activation function. After learning { κ i } , LS uses κ i to scale the input x i via element-wise multiplication. In practice, network ζ ϕ 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 ζ ϕ 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.

5 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 LS 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 / √ 2 -CS method can stabilize training to some extent, there are 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. 5 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

100

400

300

500

200

Training iteration (k)

6

10

FID-10k↓

LS (ours)

CS (ours)

Org

100

400

300

500

200

Training iteration (k)

6

10

20

LS (ours)

CS (ours)

Org

400

300

500

200

Training iteration (k)

10

15

20

LS (ours)

CS (ours)

Org

100

700

500

300

Training iteration (k)

8

10

12

LS (ours)

CS (ours)

Org

FID-10k↓

LS (ours)

CS (ours)

Org

100

400

300

500

200

Training iteration (k)

15

25

35

LS (ours)

CS (ours)

Org

Training iteration (k)

Architecture

Batch size (bz)

Model depth

Dataset

UViT

UNet

bz=64

bz=32

depth=16

depth=20

MS-COCO

CelebA

2.5×

1.7×

1.4×

1.7×

1.7×

1.7×

2.0×

2.0×

1.5×

2.0×

3.1×

2.0×

1.5×

1.5×

100

400

300

200

10

20

30

LS (ours)

CS (ours)

Org

1.5×

1.5×

20

60

100

80

200

120

Training iteration (k)

160

LS (ours)

CS (ours)

Org

14

18

22

26

100

180

140

Training iteration (k)

most cases, they consistently accelerate the learning speed by at least 1.5 × 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 a 32-sized batch, LS 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 1 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 UViT, our CS and LS also achieve much better performance than 1 / √ 2 -CS. All these results are consistent with the analysis in Section 4.1.

2.95

3.05

FID-50k↓

3.15

3.11

3.11

3.01

2.98

4.50

5.50

3.50

5.36

4.72

3.58

3.95

6.0

10.0

14.0

6.0

8.0

10.0

4.0

13.89

12.67

5.90

8.51

8.31

5.67

4.08

5.36

LS (ours)

CS (ours)

Org UViT

-CS

bz=128/depth=12

bz=64/depth=12

bz=32/depth=12

bz=64/depth=16

bz=64/depth=20

FID-50k↓

FID-50k↓

FID-50k↓

FID-50k↓

6.0

8.0

4.0

6.62

3.56

4.14

4.95

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 CIFAR10 by showing big improvements on standard UViT and 1 / √ 2 -CS-based UViT under different training batch sizes and network depth.

Robustness of LS to Network Design . For the calibration module in LS, here we investigate 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 κ i 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 r = 16 and ReLU throughout this work.

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.

6 Discussion

Relationship between LS and CS . Taking UViT on CIFAR10, MSCOCO, ImageNet and CelebA as baselines, we randomly sample 30 Gaussian noises from N (0 , I ) and measure the scaling coefficients, i.e., κ i ≤ 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.1. Moreover, these learned coefficients κ i 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.

-th Long skip connecion

1

2

5

4

3

6

1.0

0.8

0.6

0.4

0.2

0.0

ImageNet

MS-COCO

CelebA

CIFAR10

Figure 7: The

κ

i

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, κ i = κ N -i +1 ( κ < 1) , the stability bound in Theorem 3.4 is extremely large. The main term of stability bound in Theorem 3.4 can be written as S r = ∑ N i =1 κ i M i 0 = κ N M 0 + κ N -1 M 2 0 + ... + κM N 0 , and could be very large, since M N 0 > 1 is large when N is large and scaling it by a factor κ could not sufficiently control its magnitude. In contrast, our default setting κ i = κ i -1 of CS can well control the main term in stability bound: S = ∑ N i =1 κ i M i 0 = M 0 + κ 1 M 2 0 + ... + κ N -1 M N 0 , where the larger terms M i +1 0 are weighted by smaller coefficients κ i . In this way, S is much smaller than S r, which shows the advantages of our default setting. Besides, the following Table 4 also compares the above two settings by using UViT on Cifar10 (batch size = 64), and shows that our default setting exhibits significant advantages.

Next, for LS, there are two possible reasons for why LS discovers a 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 ≤ N ) , the learnable κ i should be smaller to better control the

magnitude of M i 0 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 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 reduce 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.

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 κ , or to have better estimates of the range of κ . Another effective solution is to use our LS method which can automatically learn qualified scaling coefficient κ i 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.01M 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 × 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.2021ZD0111601, National Natural Science Foundation of China (NSFC) under Grant No.61836012, 62325605, U21A20470, GuangDong Basic and Applied Basic Research Foundation under Grant No. 2023A1515011374, GuangDong Province Key Laboratory of Information Security Technology.

References

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 IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11461-11471, 2022.

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.09556 , 2023.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , pages 8162-8171. PMLR, 2021.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

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.

Florinel-Alin Croitoru, Vlad Hondru, Radu Tudor Ionescu, and Mubarak Shah. Diffusion models in vision: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2023.

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.

Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. Ilvr: Conditioning method for denoising diffusion probabilistic models. 2021 IEEE/CVF International Conference on Computer Vision (ICCV) , pages 14347-14356, 2021.

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 Pattern Recognition Workshops (CVPRW) , pages 649-655, 2022.

Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM , 63(11):139144, 2020.

Shengjia Zhao, Jiaming Song, and Stefano Ermon. Infovae: Balancing learning and inference in variational autoencoders. In Proceedings of the aaai conference on artificial intelligence , volume 33, pages 5885-5892, 2019.

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 Pattern Recognition , pages 2426-2435, 2022.

Omri Avrahami, Dani Lischinski, and Ohad Fried. Blended diffusion for text-driven editing of natural images. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 18208-18218, 2022.

Alexander Quinn Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob Mcgrew, 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. PMLR, 2022.

Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is a strong image synthesizer. ArXiv , 2023.

Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209.14988 , 2022.

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.06135 , 2022.

Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. arXiv preprint arXiv:2204.03458 , 2022.

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 arXiv:2302.01329 , 2023.

Patrick Esser, Johnathan Chiu, Parmida Atighehchian, Jonathan Granskog, and Anastasis Germanidis. Structure and content-guided video synthesis with diffusion models. arXiv preprint arXiv:2302.03011 , 2023.

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 diffusion models. arXiv preprint arXiv:2304.08818 , 2023.

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.02881 , 2023.

Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, Ming-Hsuan Yang, and Ling Shao. Multi-stage progressive image restoration. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 14821-14831, 2021.

Zhendong Wang, Xiaodong Cun, Jianmin Bao, Wengang Zhou, Jianzhuang Liu, and Houqiang Li. Uformer: A general u-shaped transformer for image restoration. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 17683-17693, 2022.

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 Pattern Recognition , pages 5728-5739, 2022.

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.

Yang Song, Jascha Narain Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. ArXiv , abs/2011.13456, 2020.

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-4726, 2021.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, Seyedeh Sara Mahdavi, Raphael Gontijo Lopes, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. ArXiv , abs/2205.11487, 2022.

Jonathan Ho, Ajay Jain, and P. Abbeel. Denoising diffusion probabilistic models. ArXiv , abs/2006.11239, 2020.

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.

Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.

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 2015.

Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean 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 , 115:211-252, 2014.

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.

Fan Bao, Chongxuan Li, Yue Cao, and Jun Zhu. All are worth words: a vit backbone for score-based diffusion models. arXiv preprint arXiv:2209.12152 , 2022.

Jonathan Ho. Classifier-free diffusion guidance. ArXiv , abs/2207.12598, 2022.

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 arXiv:2305.05189 , 2023.

Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation. ArXiv , abs/2208.12242, 2022.

Nupur Kumari, Bin Zhang, Richard Zhang, Eli Shechtman, and Jun-Yan Zhu. Multi-concept customization of text-to-image diffusion. ArXiv , abs/2212.04488, 2022.

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.

Jiale Xu, Xintao Wang, Weihao Cheng, Yan-Pei Cao, Ying Shan, Xiaohu Qie, and Shenghua Gao. Dream3d: Zero-shot text-to-3d synthesis using 3d shape prior and text-to-image diffusion models. ArXiv , abs/2212.14704, 2022.

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.

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.

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.

Huishuai Zhang, Da Yu, Mingyang Yi, Wei Chen, and Tie-Yan Liu. Stabilize deep resnet with a sharp scaling factor tau. 2019.

Huishuai Zhang, Da Yu, Wei Chen, and Tie-Yan Liu. On the stability of multi-branch network, 2021.

Thomas C. Bachlechner, Bodhisattwa Prasad Majumder, Huanru Henry Mao, G. Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence at large depth. In Conference on Uncertainty in Artificial Intelligence , 2020.

Hongyi Zhang, Yann Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. ArXiv , abs/1901.09321, 2019.

Soham De and Samuel L. Smith. Batch normalization biases residual blocks towards the identity function in deep networks. arXiv: Learning , 2020.

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.

Bobby He, James Martens, Guodong Zhang, Aleksandar Botev, Andy Brock, Samuel L. Smith, and Yee Whye Teh. Deep transformers without shortcuts: Modifying self-attention for faithful signal propagation. ArXiv , abs/2302.10322, 2023.

Soufiane Hayou, Eugenio Clerico, Bo He, George Deligiannidis, A. Doucet, and Judith Rousseau. Stable resnet. In International Conference on Artificial Intelligence and Statistics , 2020.

Boris Hanin and David Rolnick. How to start training: The effect of initialization and architecture. Advances in Neural Information Processing Systems , 31, 2018.

David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams. 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.

Muhammad Awais, Md. Tauhid Bin Iqbal, and Sung-Ho Bae. Revisiting internal covariate shift for batch normalization. IEEE Transactions on Neural Networks and Learning Systems , 32:5082-5092, 2020.

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-8470, 2020.

Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning , 2015.

Shibani Santurkar, Dimitris Tsipras, Andrew Ilyas, and Aleksander Madry. How does batch normalization help optimization? (no, it is not about internal covariate shift). ArXiv , abs/1805.11604, 2018.

Tim Cooijmans, Nicolas Ballas, César Laurent, and Aaron C. Courville. Recurrent batch normalization. ArXiv , abs/1603.09025, 2016.

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.

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:110008, 2022.

Xiyu He and Xiang Qian. A real-time surface defect detection system for industrial products with longtailed distribution. 2021 IEEE Industrial Electronics and Applications Conference (IEACon) , pages 313-317, 2021.

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.

Shaden Alshammari, Yu-Xiong Wang, Deva Ramanan, and Shu Kong. Long- tailed recognition via weight balancing. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 6887-6897, 2022.

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.

Zebin 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.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. ArXiv , abs/2206.00364, 2022.

Diederik P. Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. ArXiv , abs/1807.03039, 2018.

Xiulong Yang, Sheng-Min Shih, Yinlin Fu, Xiaoting Zhao, and Shihao Ji. Your vit is secretly a hybrid discriminative-generative diffusion model. arXiv preprint arXiv:2208.07791 , 2022.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems , 34:8780-8794, 2021.

Wanshu Fan, Yen-Chun Chen, Dongdong Chen, Yu Cheng, Lu Yuan, and Yu-Chiang Frank Wang. Frido: Feature pyramid diffusion for complex scene image synthesis. ArXiv , abs/2208.13753, 2022.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020.

Dongjun Kim, Seung-Jae 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 on Machine Learning , 2021.

Djork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv: Learning , 2015.

Abien Fred Agarap. Deep learning using rectified linear units (relu). ArXiv , abs/1803.08375, 2018.

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.

Diganta Misra. Mish: A self regularized non-monotonic activation function. In British Machine Vision Conference , 2020.

Jie Hu, Li Shen, Samuel Albanie, Gang Sun, and Enhua Wu. Squeeze-and-excitation networks. IEEE Transactions on Pattern Analysis and Machine Intelligence , 42:2011-2023, 2017.

Shan Zhong, Wushao Wen, and Jinghui Qin. Switchable self-attention module. ArXiv , abs/2209.05680, 2022.

Sanghyun Woo, Jongchan Park, Joon-Young Lee, and In-So Kweon. Cbam: Convolutional block attention module. In European Conference on Computer Vision , 2018.

Zhongzhan Huang, Senwei Liang, Mingfu Liang, and Haizhao Yang. Dianet: Dense-and-implicit attention network. In AAAI Conference on Artificial Intelligence , 2019.

Gavin E Crooks. Survey of simple, continuous, univariate probability distributions. Technical report, Technical report, Lawrence Berkeley National Lab, 2013., 2012.

Robert B Davies. Algorithm as 155: The distribution of a linear combination of χ 2 random variables. Applied Statistics , pages 323-333, 1980.

Robert B Davies. Numerical inversion of a characteristic function. Biometrika , 60(2):415-417, 1973.

Zhongzhan Huang, Wenqi Shao, Xinjiang Wang, Liang Lin, and Ping Luo. Rethinking the pruning criteria for convolutional neural network. Advances in Neural Information Processing Systems , 34:16305-16318, 2021.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via overparameterization. In International Conference on Machine Learning , pages 242-252. PMLR, 2019.

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 in 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 κ greater than 1 in CS, and the experimental results are consistent with our theoretical analysis.

A Implementation details

ImageNet 64x 64

CelebA 64x 64

CIFAR1032×32

MS-COCO

A.1 Experimental setup

In this paper, we consider four commonly used datasets for our experiments: CIFAR10 [32], CelebA [33], ImageNet [34], and MS-COCO [35]. For the unconditional setting, we use CIFAR10, 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 × 64 resolutions, which contains 1,281,167 training images from 1,000 different classes. For the text-to-image setting, we consider the MS-COCO dataset at 256 × 256 resolution. Unless otherwise specified, all of our experiments are conducted on A100 GPU and follow the experimental settings in [36].

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. 9 and Fig. 10 compare the original UViT with the code incorporating CS or LS, demonstrating that our proposed methods require minimal code modifications.

Origin

CS (ours)

skips = []

skips = []

for blk in self.in_blocks:

for blk in self.in_blocks:

x = blk(x)

x = blk(x)

skips.append(x)

skips.append(x)

X = self.mid_block(x)

X = self.mid_block(x)

layer_cnt = 0

for blk in self.out_blocks:

for blk in self.out_blocks:

x = blk(x, skips.pop())

x =blk(x,skips.pop()* k**

(self.depth - 1 - layer_cnt))

layer_cnt += 1

x = self.norm(x)

X = self.norm(x)

x = self.decoder_pred(x)

x = self.decoder_pred(x)

x=x[:,self.extras:,:]

x=x[:, self.extras:,:]

X = unpatchify(x, self.in_chans)

X = unpatchify(x, self.in_chans)

x = self.final_layer(x)

x = self.final_layer(x)

return x

return x

CS requires no additional parameters and can rapidly adapt to any UNet-alike architecture and dataset. In our experiments, we have chosen values of κ as 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 κ within the range of 0.5 to 0.95. LS designs a calibration network shared by all long skip connections to predict scaling coefficients { κ i } . Specifically, for the i -th long skip connection, its input feature is x i ∈ R B × N × D , where B denotes the batch size. For convolution-based UNet, N and D respectively denote the channel number and spatial dimension ( H × W ) of the feature map; for transformer-based UViT [36], N and D respectively represent the token number and token dimension.

In the main text, we define the LS as

For ζ ϕ , we employ a neural network structure based on SENet [79], given by:

where x ∈ R N , ◦ denotes matrix multiplication, and ψ represents the ReLU activation function. The matrices W 1 ∈ R N × N/r and W 1 ∈ R N/r × N , where r is the compression ratio, are set to r = 16 in this paper. In the ablation study of the main text, we have demonstrated that the choice of r has a minimal impact on the performance of the diffusion model. It is worth noting that if LS is applied to long skip connections where the feature dimensions remain unchanged, such as in UViT, then the calibration network is only a shared ζ ϕ ( x ) .

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 a 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 is ( B × N i × D ), 1 ≤ i ≤ d , the input dimension of ζ ϕ ( x ) is set to N min = min N i . Additionally, we set up d sets of encoders and decoders, where features with the same dimension

Origin

LS (ours)

skips = []

skips = []

for blk in self.in_blocks:

for blk in self.in_blocks:

x = blk(x)

x = blk(x)

skips.append(x)

skips.append(x)

x = self.mid_block(x)

X = self.mid_block(x)

for blk in self.out_blocks:

for blk in self.out_blocks:

x = blk(x, skips.pop())

skip = skips.pop()

x = blk(x, skip * self.LS(skip) )

x = self.norm(x)

X = self.norm(x)

X = self.decoder_pred(x)

x = self.decoder_pred(x)

x = x[:, self.extras:, :]

X=x[:, self.extras:,:]

X = unpatchify(x, self.in_chans)

X = unpatchify(x, self.in_chans)

x = 5elf.final_layer(x)

X = 5elf.final_layer(x)

return x

return x

share an encoder and a decoder. The i -th set of encoder ζ en and decoder ζ de also adopt the SENet structure, given by:

where r i is the compression ratio, typically set to ⌊ N i / 4 ⌋ , and

In fact, since the parameter count of ζ ϕ ( x ) itself is not substantial, to make LS more versatile, we can also individually set a scaling module into each long skip connection. This configuration typically does not lead to significant performance degradation and can even result in improved performance, while maintaining the same inference speed as the original LS.

B The proof of the Theorems and Lemma

B.1 Preliminaries

Diffusion Model (DM) . DDPM-alike DMs [1-6, 8, 9, 43] generates a sequence of noisy samples { x i } T i =1 by repeatedly adding Gaussian noise to a sample x 0 until attatining x T ∼ N ( 0 , I ) . This noise injection process, a.k.a. the forward process, can be formalized as a Markov chain q ( x 1: T | x 0 ) = ∏ T t =1 q ( x t | x t -1 ) , where q ( x t | x t -1 ) = N ( x t | √ α t x t -1 , β t I ) , α t and β t depend on t and satisfy α t + β t = 1 . By using the properties of Gaussian distribution, the forward process can be written as

where ¯ α t = ∏ t i =1 α i . Next, one can sample x t = √ ¯ α t x 0 + √ 1 -¯ α t ϵ t , where ϵ t ∼ N ( 0 , I ) . Then DMadopts a neural network ˆ ϵ θ ( · , t ) to invert the forward process and predict the noise ϵ t added at each time step. This process is to recover data x 0 from a Gaussian noise by minimizing the loss

UNet-alike Network . Since the network ˆ ϵ θ ( · , 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 in (16) that uses LSCs to connect distant parts of the network for long-range information transmission and aggregation

where x ∈ R m denotes the input, a i and b i ( i ≥ 1) are the trainable parameter of the i -th block, κ i > 0 ( i ≥ 1) are the scaling coefficients and are set to 1 in standard UNet. f N 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 block a i and b i as the stacked networks [46, 47] which is implemented as

with a ReLU activation function ϕ and learnable matrices W a i j , W b i j ∈ R m × m ( j ≥ 1) . Moreover, let f N also have the same architecture, i.e. f N ( x ) = W f i l ϕ ( W f i l -1 ...ϕ ( W f i 1 x )) .

Proposition B.1 (Amoroso distribution) . The Amoroso distribution is a four parameter, continuous, univariate, unimodal probability density, with semi-infinite range [83]. And its probability density function is

for x, a, θ, α, β ∈ R , α > 0 and range x ≥ a if θ > 0 , x ≤ a if θ < 0 . The mean and variance of Amoroso distribution are

and

Proposition B.2 (Stirling's formula) . For big enough x and x ∈ R + , we have an approximation of Gamma function:

Proposition B.3 (Scaled Chi distribution) . Let X = ( x 1 , x 2 , ...x k ) and x i , i = 1 , ..., k are k independent, normally distributed random variables with mean 0 and standard deviation σ . The statistic ℓ 2 ( X ) = √ ∑ k i =1 x 2 i follows Scaled Chi distribution [83]. Moreover, ℓ 2 ( X ) also follows Amoroso ( x | 0 , √ 2 σ, k 2 , 2) . By Eq. (19) and Eq. (20), the mean and variance of Scaled Chi distribution are

and

Lemma B.4. For big enough x and x ∈ R + , we have

And

Proof.

on the other hand, we have

Lemma B.5. [84, 85] Let x = ( x 1 , x 2 , ..., x N ) , where x i , 1 ≤ i ≤ N , are independently and identically distributed random variables with x i ∼ N ( µ, σ 2 ) . In this case, the mathematical expectation E ( ∑ N i =1 x 2 i ) can be calculated as σ 2 N + µ 2 N .

Lemma B.6. [86] Let F ij ∈ R N i × k × k be the j -th filter of the i -th convolutional layer. If all filters are initialized as Kaiming's initialization, i.e., N (0 , 2 N i × k × k ) , in i -th layer, F ij ( j = 1 , 2 , ..., N i +1 ) are i.i.d and approximately follow such a distribution during training:

LemmaB.7. [47, 87] For the feedforward neural network W l ϕ ( W l -1 ...ϕ ( W 1 x )) with initialization N (0 , 2 m ) , and W i , 1 ≤ i ≤ l, ∈ R m × m . ∀ ρ ∈ (0 , 1) , the following inequality holds with at least 1 -O ( l ) exp[ -Ω( mρ 2 /l )] :

where h a is ϕ ( W a h a -1 ) for a = 1 , 2 , ..., l -1 with h 0 = x .

Lemma B.8. For a matrix W ∈ R m × m and its elements are independently follow the distribution N (0 , c 2 ) , we have the follwing estimation while m is large enough

where s ∈ R m .

Proof. Let v i , 1 ≤ i ≤ m be the i -th column of matrix W , and we have v i ∼ N ( 0 , c 2 I m ) . We first prove the following two facts while m →∞ :

⟨ v i , v j ⟩ → 0 , 1 ≤ i < j ≤ m ;

For the fact (1), since Chebyshev inequality, for 1 ≤ i ≤ m and a given M , we have

from Proposition B.3 and Lemma. B.4, we can rewrite Eq. (29) when m →∞ :

For a small enough ϵ > 0 , let M = 1 /ϵ . Note that √ M 2 c = c/ √ 2 ϵ is a constant. When m →∞ , √ 2 c · Γ(( m +1) / 2) Γ( m/ 2) ≫ √ M 2 c . Hence, for any i ∈ [1 , m ] and any small enough ϵ , we have

So the fact (1) holds. Moreover, we consider the ⟨ v i , v j ⟩ , 1 ≤ i < j ≤ m .

Let v i = ( v i 1 , v i 2 , ..., v im ) and v j = ( v j 1 , v j 2 , ..., v jm ) . So ⟨ v i , v j ⟩ = ∑ m p =1 v ip v jp . Note that, v i and v j are independent, hence

Var ( v ip v jp ) = Var ( v ip ) Var ( v jp ) + ( E ( v ip )) 2 Var ( v jp ) + ( E ( v jp )) 2 Var ( v ip ) = c 4 , (33) since central limit theorem, we have

According to Eq. (22), Lemma B.4 and Eq. (34), when m →∞ , we have

Therefore, the fact (2) holds and we have ⟨ v , v ⟩ → 0 , 1 ≤ i < j ≤ m

Next, we consider ∥ W s ∥ 2 for any s ∈ R

i j . m .

̸

From the facts (1) and (2), when m is large enough, we have v T i v j → 0 , i = j , and ∥ v 1 ∥ 2 2 ≈ ∥ v 2 ∥ 2 2 ≈ · · · ≈ ∥ v m ∥ 2 2 := c v . Therefore, we have

Let k = 1 in Lemma B.6, and the convolution kernel degenerates into a learnable matrix. In this case, we can observe that if W ∈ R m × m is initialized using Kaiming initialization, i.e., N (0 , 2 /m ) , during the training process, the elements of W will approximately follow a Gaussian distribution N (0 , O ( m -1 )) . In this case, we can easily rewrite Lemma B.7 as the following Corollary B.9:

Corollary B.9. For the feedforward neural network W l ϕ ( W l -1 ...ϕ ( W 1 x )) with kaiming's initialization during training. ∀ ρ ∈ (0 , 1) , the following inequality holds with at least 1 -O ( l ) exp[ -Ω( mρ 2 /l )] :

where h a is ϕ ( W a h a -1 ) for a = 1 , 2 , ..., l -1 with h 0 = x .

Moreover, considering Lemma B.6 and Lemma B.8, we can conclude that the matrix W l in Corollary B.9 will not change the magnitude of its corresponding input ϕ ( W l -1 ...ϕ ( W 1 x )) . In other words, we have the following relationship:

Therefore, for a i and b i defined in Eq. (17), we have following corollary:

Corollary B.10. For a i and b i defined in Eq. (17), 1 ≤ i ≤ N , and their corresponding input x a and x b . ∀ ρ ∈ (0 , 1) , the following each inequality holds with at least 1 -O ( l ) exp[ -Ω( mρ 2 /l )] :

Lemma B.11. [47] Let D be a diagonal matrix representing the ReLU activation layer. Assume the training loss is 1 n ∑ n s =1 ℓ s ( W ) , where ℓ s ( W ) denotes the training loss of the s -th sample among the n training samples. For the W q p defined in Eq. (17) and f 0 ( x t ) defined in Eq. (16), where p ∈ { 1 , 2 , ..., l } , q ∈ { a i , b i } , the following inequalities hold with at least 1O ( nN ) exp[ -Ω( m )] :

and

where v ∈ R m .

B.2 The proof of Theorem 3.1.

Theorem 3.1 . Assume that all learnable matrices of UNet in Eq. (16) are independently initialized as Kaiming's initialization, i.e., N (0 , 2 m ) . Then for any ρ ∈ (0 , 1] , by minimizing the training loss Eq. (15) of DMs, with probability at least 1 -O ( N ) exp[ -Ω( mρ 2 )] , we have

where the hidden feature h i is the output of f i defined in Eq. (16), x t = √ ¯ α t x 0 + √ 1 -¯ α t ϵ t is the input of UNet; c 1 and c 2 are two constants; κ i is the scaling coefficient of the i -th LSC in UNet.

Proof. Let s represent the input of f i in Eq. (16), i.e., s = a i ◦ a i -1 ◦ · · · a 1 ◦ x t , and let h i = f i ( s ) . By expanding Eq. (16), we can observe that h i = f i ( s ) is mainly composed of the following N -i terms in the long skip connection:

When N is sufficiently large, h i 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:

where π j b := b i +1 ◦ b i +2 · · · ◦ b j , π j a := a j ◦ a j -1 · · · ◦ a i +1 .

In practice, the depth l of block a i or b i is a universal constant, therefore, from Corollary B.9, ∀ ρ ∈ (0 , 1) and j ∈ { i +1 , i +2 , ..., N } , with probability at least 1 -O ( N ) exp[ -Ω( mρ 2 )] , we have

where s.t. h i = f i ( s ) ≈ ∑ N j = i +1 κ j W b i +1 l ◦ h κ j . From Lemma B.6, the elements of W b i +1 l will approximately follow a Gaussian distribution N (0 , O ( m -1 )) , therefore, E ( W b i +1 l ◦ h κ j ) = 0 . Moreover, from Lemma B.8,

Therefore [47], h i ≈ ∑ N j = i +1 κ j W b i +1 l ◦ h κ j can be approximately regarded as a Gaussian variable with zero mean and covariance matrix ∑ N j = i +1 κ 2 j O ( m -1 ) ∥ h κ j ∥ 2 2 · I m × m . According to Lemma B.5, we have

From Eq. (46) and Eq. (49), we have

where we use c 1 to denote O (1) in Eq. (49) and we use the constant c 2 to compensate for the estimation loss caused by all the approximation calculations in the derivation process.

B.3 The proof of Lemma 3.2.

Lemma3.2 . For x t = √ ¯ α t x 0 + √ 1 -¯ α t ϵ t defined in Eq. (14) as a input of UNet, ϵ t ∼ N (0 , I ) , if x 0 follow the uniform distribution U [ -1 , 1] , then we have

Proof. For x t = √ ¯ α t x 0 + √ 1 -¯ α t ϵ t , we have

where x 0 ( j ) represents the j -th element of x 0 . The equality in Eq. (52) follows from the fact that, from Eq. (14) and Lemma B.5, we know that x t follows N ( x t ; √ ¯ α t x 0 , (1 -¯ α t ) I ) . Therefore, ∥ x t ∥ 2 2 follows a chi-square distribution, whose mathematical expectation is E t, x 0 (¯ α t ∑ m j =1 x 2 0 ( j ) + (1 -¯ α t ) m ) .

Moreover, since x 0 follow the uniform distribution U [ -1 , 1] , we have

where P x ∼ U [ -1 , 1] ( x ) is density function of the elements in x 0 . The last equality in Eq. (53) follows the fact that, according to the setup of the diffusion model, it is generally observed that ¯ α t monotonically decreases to 0 as t increases, with ¯ α t ≤ 1 . Consequently, we can easily deduce that E t ¯ α t is of the order O (1) . For instance, in the setting of DDPM, we have E t ¯ α t = E t ( ∏ t i =1 √ 1 -0 . 02 i 1000 ) ≈ 0 . 27 = O (1) .

Theorem 3.3 . Assume that all learnable matrices of UNet in Eq. (16) are independently initialized as Kaiming's initialization, i.e., N (0 , 2 m ) . Then for any ρ ∈ (0 , 1] , with probability at least 1 -O ( nN ) exp[ -Ω( m )] , for a sample x t in training set, we have

where x t = ¯ α t x 0 + 1 -¯ α t ϵ t denotes the noisy sample of the s -th sample, ϵ t ∼ N ( 0 , I ) , N is the number of LSCs, c 3 is a small constant. n is the size of the training set.

Proof. From Eq. (16) and Eq. (44), we can approximate the output of UNet f 0 ( x t ) by weight sum of N forward path g j , 1 ≤ j ≤ N , i.e.,

where g j denote b 1 ◦ b 2 ◦ · · · b j ◦ a j ◦ · · · a 1 ◦ x t , and 1 ≤ j ≤ N .

Moreover, we use g j,q p represent the feature map after W q p in forward path g j , i.e., g j,q p = W q p h j,q p -1 and h j,q p = ϕ ( g j,q p ) . Next, we first estimate ∥∇ W a i p ℓ s ( W ) ∥ 2 2 ,

From Lemma B.11 and Corollary B.9, with probability at least 1O ( nN ) exp[ -Ω( m )] , we have

where the constant c 3 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 = b i , i.e., ∥∇ W b i p ℓ s ( W ) ∥ 2 2 ≲ O ( ℓ s ( W ) · ∥ x t ∥ 2 2 · ∑ N j = i κ 2 j + c 3 ) .

B.4 The proof of Theorem 3.4.

Theorem 3.4 . For UNet in Eq. (16), assume M 0 = max {∥ b i ◦ a i ∥ 2 , 1 ≤ i ≤ N } and f N is L 0 -Lipschitz continuous. c 0 is a constant related to M 0 and L 0 . Suppose x ϵ δ t is an perturbated input of the vanilla input x t with a small perturbation ϵ δ = ∥ x ϵ δ t -x t ∥ 2 . Then we have

where x t = √ ¯ α t x 0 + 1 -¯ α t ϵ t , ϵ t ∼ N ( 0 , I ) , N is the number of the long skip connections.

Proof. From Eq. (16), we have

and

Using Taylor expansion, we have

Therefore,

Next, we estimate the ∥∇ z f 1 ( z ) | z = a 1 ◦ x t ∥ 2 , From Eq. (16), we have

Therefore, we have

From Eq. (62) and Eq. (64), let ∆ = ∥ UNet ( x ϵ δ t ) -UNet ( x t ) ∥ 2 , we have

Next, we demonstrate through empirical experiments that M 0 in Theorem 3.4 is generally greater than 1. As shown in Fig. 11, we directly computed ∥ b i ◦ a i ∥ 2 for the Eq. (16) on the CIFAR10 and CelebA datasets. We observe that ∥ b i ◦ a i ∥ 2 exhibits a monotonically increasing trend with training iterations and is always greater than 1 at any iteration. Since M 0 = max ∥ b i ◦ a i ∥ 2 , 1 ≤ i ≤ N , we conclude that M 0 ≥ 1 .

Norm =  1.0

Norm =  1.0

Norm

0

20

40

60

0

10

20

30

Training iteration / 5k

Norm

0

20

40

60

0

10

20

30

Training iteration / 5k

CIFAR10

CelebA

C Other experimental observations

Fig. 12 is a supplement to Fig. 1 in the main text. We investigate the case of UNet with a 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.

0

30000

20000

10000

Iterations

0

30000

20000

10000

Iterations

0

30000

20000

10000

Iterations

0

30000

20000

10000

Iterations

Norm of feature

1.0

0

3.0

2.0

×10000

×10000

×10000

×10000

2.0

1.0

0

Norm of feature

1.0

0

3.0

2.0

Norm of feature

2.0

1.0

0

Norm of feature

LS (ours)

CS (ours)

Org

LS (ours)

CS (ours)

Org

LS (ours)

CS (ours)

Org

LS (ours)

CS (ours)

Org

In the main text, we specify that the domain of κ in CS should be (0,1). This ensures stable forward and backward propagation of the model and enhances the model's robustness to input noise, as discussed in Section 3. Otherwise, it would exacerbate the instability of the model training.

In Fig. 13, we illustrate the impact of having κ > 1 on the performance of the diffusion model. We observe that generally, having κ > 1 leads to a certain degree of performance degradation, especially in unstable training settings, such as when the batch size is 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 κ align with the analysis presented in Section 3.

100

400

300

500

200

Training iteration (k)

100

400

300

500

200

Training iteration (k)

100

250

200

300

150

Training iteration (k)

100

250

200

300

150

Training iteration (k)

FID-10k↓

6

10

14

18

10

20

30

8

12

16

20

0

200

400

600

CS (ours)

=1.2

=1.5

CS (ours)

=1.2

=1.5

CS (ours)

=1.2

=1.5

CS (ours)

=1.2

=1.5

CIFAR10-bz=128

CIFAR10-bz=64

ImageNet64-bz=1024

ImageNet64-bz=256

Peer reviews:
comments: SUMMARY:
They state that diffusion models using Unet suffer from unstable training
and 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.

STRENGTHS:
I believe this paper’s contribution is simple and meaningful, and written
clearly. 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.

WEAKNESSES:
I did not see meaningful weaknesses in the paper.

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nU-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.\n\nSTRENGTHS:\n* 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).\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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\n\nLIMITATIONS:\nThe authors adequately addressed the limitations and introduced further directions.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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}}$?\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\nPlease kindly refer to the weakness section for more details.\n\nLIMITATIONS:\nLimitations are discussed in Sect. 6 of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\n-\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nMinor:\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\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}