SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhongzhan Huang $^{1,2}$  Pan Zhou $^{2*}$  Shuicheng Yan $^{2}$  Liang Lin $^{1*}$
$^{1}$ Sun Yat-Sen University,  $^{2}$ Sea AI Lab
huangzhzh23@mail2.sysu.edu.cn;zhoupan@sea.com;
shuicheng.yan@gmail.com; linliang@ieee.org
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 \times$  training acceleration on different diffusion models with UNet or UViT backbones. Click here for Code.
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
Figure 1: The feature oscillation in UNet.
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/\sqrt{2}$ -constant scaling  $(1/\sqrt{2}$ -CS), that scales the coefficients of LSC (i.e.  $\{\kappa_i\}$  in Fig. 2) from one used in UNet to  $1/\sqrt{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/\sqrt{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
Figure 2: The diagram of UNet.
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  $\mathcal{O}\big(m\sum_{i = 1}^{N}\kappa_{j}^{2}\big)$ , where  $\{\kappa_i\}$  denotes the scaling coefficients of  $N$  LSCs and the input dimension  $m$  of UNet is often of hundreds. Since standard UNet sets  $\kappa_{i} = 1 (\forall i)$ , the above oscillation range becomes  $\mathcal{O}(mN)$  and is large, which partly explains the oscillation and instability of hidden features as observed in Fig. 1. The  $1 / \sqrt{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  $\mathcal{O}\big(m\sum_{i = 1}^{N}\kappa_{i}^{2}\big)$ . In standard UNet, this bound becomes a large bound  $\mathcal{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 / \sqrt{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  $\mathcal{O}\big(\sum_{i = 1}^{N}\kappa_{i}M_{0}^{i}\big)$  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 / \sqrt{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  $\{\kappa_i\}$  as a serir of exponentially-decaying constants, i.e.,  $\kappa_{i} = \kappa^{i - 1}$  with a constant  $\kappa \in (0,1]$ . As a result, CS greatly stabilizes the UNettraining by largely reducing the robustness error bound from  $\mathcal{O}(M_0^N)$  to  $\mathcal{O}\big(M_0(\kappa M_0)^{N - 1}\big)$ , which is also  $\mathcal{O}(\kappa^{N - 2})$  times smaller than the bound  $\mathcal{O}(\kappa M_0^N)$  of the universal scaling method, i.e.,  $\kappa_{i} = \kappa$ , like  $1 / \sqrt{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 \times$  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  $\{\mathbf{x}_i\}_{i=1}^T$  by repeatedly adding Gaussian noise to a sample  $\mathbf{x}_0$  until attaining  $\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ . This noise injection process, a.k.a. the forward process, can be formalized as a Markov chain  $q(\mathbf{x}_{1:T}|\mathbf{x}_0) = \prod_{t=1}^{T} q(\mathbf{x}_t|\mathbf{x}_{t-1})$ , where  $q(\mathbf{x}_t|\mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t|\sqrt{\alpha_t}\mathbf{x}_{t-1}, \beta_t\mathbf{I})$ ,  $\alpha_t$  and  $\beta_t$  depend on  $t$  and satisfy  $\alpha_t + \beta_t = 1$ . By using the properties of Gaussian distribution, the forward process can be written as
where  $\bar{\alpha}_t = \prod_{i=1}^t \alpha_i$ . Next, one can sample  $\mathbf{x}_t = \sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon_t$ , where  $\epsilon_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ . Then DM adopts a neural network  $\hat{\epsilon}_{\theta}(\cdot, t)$  to invert the forward process and predict the noise  $\epsilon_t$  added at each time step. This process is to recover data  $\mathbf{x}_0$  from a Gaussian noise by minimizing the loss
$$
\ell_ {\text {s i m p l e}} ^ {t} (\theta) = \mathbb {E} \| \epsilon_ {t} - \hat {\epsilon} _ {\theta} \left(\sqrt {\bar {\alpha} _ {t}} \mathbf {x} _ {0} + \sqrt {1 - \bar {\alpha} _ {t}} \epsilon_ {t}, t\right) \| _ {2} ^ {2}. \tag {2}
$$
UNet-alike Network. Since the network  $\hat{\epsilon}_{\theta}(\cdot, 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
$$
\mathbf {U N e t} (x) = f _ {0} (x), \quad f _ {i} (x) = b _ {i + 1} \circ \left[ \kappa_ {i + 1} \cdot a _ {i + 1} \circ x + f _ {i + 1} \left(a _ {i + 1} \circ x\right) \right], \quad 0 \leq i \leq N - 1 \tag {3}
$$
where  $x \in \mathbb{R}^m$  denotes the input,  $a_i$  and  $b_i$  ( $i \geq 1$ ) are the trainable parameter of the  $i$ -th block,  $\kappa_i > 0$  ( $i \geq 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  $\circ$ , it can be designed to implement different networks.
W.l.o.g, in this paper, we consider  $\circ$  as matrix multiplication, and set the  $i$ -th block as a stacked network [46, 47] which is implemented as  $a_{i} \circ x = \mathbf{W}_{l}^{a_{i}} \phi(\mathbf{W}_{l-1}^{a_{i}} \dots \phi(\mathbf{W}_{1}^{a_{i}} x))$  with a ReLU activation function  $\phi$  and learnable matrices  $\mathbf{W}_{j}^{a_{i}} \in \mathbb{R}^{m \times m}$  ( $j \geq 1$ ). Moreover, let  $f_{N}$  also have the same architecture, i.e.  $f_{N}(x) = \mathbf{W}_{l}^{f_{i}} \phi(\mathbf{W}_{l-1}^{f_{i}} \dots \phi(\mathbf{W}_{1}^{f_{i}} x))$ . Following [36], the above UNet has absorbed the fusion operation of the feature  $a_{i} \circ x$  and  $f_{i}$  ( $a_{i} \circ 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.  $\mathbf{x}_{i + 1} = \mathbf{x}_i + \kappa_i f_i(\mathbf{x}_i)$  where  $f_{i}(\mathbf{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  $(\kappa_{i} = 1,\forall i)$  often suffers from the unstable training issue in diffusion models (DMs), while scaling the long skip connection (LSC) coefficients  $\kappa_{i}(\forall 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 coefficients 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
Figure 3: Distribution of  $\mathbf{x}_0$ ,  $\| \mathbf{x}_t\| _2^2$  and  $\| \epsilon_t\| _2^2 /\| \mathbf{x}_t\| _2^2$ , and effect of extra noise to performance.
the coefficients of LSC in UNet helps stabilize 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.,  $\mathcal{N}(0,\frac{2}{m})$ . Then for any  $\rho \in (0,1]$ , by minimizing the training loss Eq. (2) of DMs, with probability at least  $1 - \mathcal{O}(N)\exp [-\Omega (m\rho^2)]$ , we have
$$
(1 - \rho) ^ {2} \left[ c _ {1} \| \mathbf {x} _ {t} \| _ {2} ^ {2} \cdot \sum_ {j = i} ^ {N} \kappa_ {j} ^ {2} + c _ {2} \right] \lesssim \| h _ {i} \| _ {2} ^ {2} \lesssim (1 + \rho) ^ {2} \left[ c _ {1} \| \mathbf {x} _ {t} \| _ {2} ^ {2} \cdot \sum_ {j = i} ^ {N} \kappa_ {j} ^ {2} + c _ {2} \right], \tag {4}
$$
where the hidden feature  $h_i$  is the output of  $f_i$  defined in Eq. (3),  $\mathbf{x}_t = \sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon_t$  is the input of UNet;  $c_1$  and  $c_2$  are two constants;  $\kappa_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  $\mathcal{O}\left(\|\mathbf{x}_t\|_2^2 \sum_{j=i}^{N} \kappa_j^2\right)$  which is decided by the scaling coefficients  $\{\kappa_i\}$  and the UNet input  $\mathbf{x}_t$ . So when  $\|\mathbf{x}_t\|_2^2$  or  $\sum_{i=1}^{N} \kappa_i^2$  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  $\| \mathbf{x}_t\| _2^2$  is around feature dimension  $m$  in DDPM.
Lemma 3.2. For  $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon_t$  defined in Eq. (1) as a input of UNet,  $\epsilon_t\sim \mathcal{N}(0,\mathbf{I})$ , if  $\mathbf{x}_0$  follows the uniform distribution  $U[-1,1]$ , then we have
$$
\mathbb {E} \| \mathbf {x} _ {t} \| _ {2} ^ {2} = (1 - 2 \mathbb {E} _ {t} \bar {\alpha} _ {t} / 3) m = \mathcal {O} (m). \tag {5}
$$
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  $\kappa_{i} = 1$ , the oscillation range of hidden feature  $\| h_i\| _2^2$  becomes a very large bound  $\mathcal{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  $\{\kappa_{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 / \sqrt{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  $\mathbf{W}$  to denote all the learnable matrices of the  $i$ -th block (see Eq. (3)), i.e.,  $\mathbf{W} = [\mathbf{W}_l^{a_1}, \mathbf{W}_{l-1}^{a_1}, \dots, \mathbf{W}_1^{a_1}, \mathbf{W}_l^{a_2}, \dots, \mathbf{W}_2^{b_1}, \mathbf{W}_1^{b_1}]$ . Assume the training loss is  $\frac{1}{n} \sum_{s=1}^{n} \ell_s(\mathbf{W})$  where  $\ell_s(\mathbf{W})$  denotes the training loss of the  $s$ -th sample among the  $n$  training samples. Next, we analyze the gradient  $\nabla \ell(\mathbf{W})$  of each training sample, and summarize the main results in Theorem 3.3, in which we omit the subscript  $s$  of  $\nabla \ell_s(\mathbf{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.,  $\mathcal{N}(0,\frac{2}{m})$ . Then for any  $\rho \in (0,1]$ , with probability at least  $1 - \mathcal{O}(nN)\exp [-\Omega (m)]$ , for a sample  $\mathbf{x}_t$  in training set, we have
$$
\left\| \nabla_ {\mathbf {W} _ {p} ^ {q}} \ell_ {s} (\mathbf {W}) \right\| _ {2} ^ {2} \lesssim \mathcal {O} \left(\ell (\mathbf {W}) \cdot \| \mathbf {x} _ {t} \| _ {2} ^ {2} \cdot \sum_ {j = i} ^ {N} \kappa_ {j} ^ {2} + c _ {3}\right), \quad (p \in \{1, 2, \dots , l \}, q \in \{a _ {i}, b _ {i} \}), \tag {6}
$$
where  $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon_t$  denotes the noisy sample of the  $s$ -th sample,  $\epsilon_t \sim \mathcal{N}(\mathbf{0},\mathbf{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  $\mathbf{W}_p^q$  in UNet can be bounded by  $\mathcal{O}(\ell(\mathbf{W}) \cdot \| \mathbf{x}_t\|_2^2 \sum_{j=i}^{N} \kappa_j^2)$  (up to the constant term). Since Lemma 3.2 shows  $\|\mathbf{x}_t\|_2^2$  is at the order of  $\mathcal{O}(m)$  in expectation, the bound of gradient norm becomes  $\mathcal{O}(m \ell_s(\mathbf{W}) \sum_{j=i}^{N} \kappa_j^2)$ . Consider the feature dimension  $m$  is often hundreds and the loss  $\ell(\mathbf{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 ( $\kappa_i = 1 \forall i$ ) would become  $\mathcal{O}(m N \ell_s(\mathbf{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  $\{\kappa_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/\sqrt{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  $\epsilon_t \sim \mathcal{N}(0, \mathbf{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  $\epsilon_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  $\epsilon_{t}$  from the UNet input  $\mathbf{x}_t\sim \mathcal{N}(\mathbf{x}_t;\sqrt{\bar{\alpha}_t}\mathbf{x}_0,(1 - \bar{\alpha}_t)\mathbf{I})$