SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Assume an extra Gaussian noise  $\epsilon_{\delta} \sim \mathcal{N}(0, \sigma_{\delta}^{2}\mathbf{I})$  is injected into  $\mathbf{x}_t$  which yields a new input  $\mathbf{x}_t' \sim \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t}\mathbf{x}_0, [(1 - \bar{\alpha}_t) + \sigma_{\delta}^2]\mathbf{I})$ . Accordingly, if the variance  $\sigma_{\delta}^{2}$  of extra noise is large, it can dominate the noise  $\epsilon_t$  in  $\mathbf{x}_t'$ , and thus hinders predicting the desired noise  $\epsilon_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  $\epsilon_{\delta}$  with different  $\sigma_{\delta}^{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 \circ a_i \|_2, 1 \leq i \leq N \}$  and  $f_N$  is  $L_0$ -Lipschitz continuous.  $c_0$  is a constant related to  $M_0$  and  $L_0$ . Suppose  $\mathbf{x}_t^{\epsilon_\delta}$  is an perturbed input of the vanilla input  $\mathbf{x}_t$  with a small perturbation  $\epsilon_\delta = \| \mathbf{x}_t^{\epsilon_\delta} - \mathbf{x}_t \|_2$ . Then we have
$$
\left\| \mathbf {U N e t} \left(\mathbf {x} _ {t} ^ {\epsilon_ {\delta}}\right) - \mathbf {U N e t} \left(\mathbf {x} _ {t}\right) \right\| _ {2} \leq \epsilon_ {\delta} \left[ \sum_ {i = 1} ^ {N} \kappa_ {i} M _ {0} ^ {i} + c _ {0} \right], \tag {7}
$$
where  $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon_t$ ,  $\epsilon_t \sim \mathcal{N}(\mathbf{0},\mathbf{I})$ ,  $N$  is the number of the long skip connections.
Theorem 3.4 shows that for a perturbation magnitude  $\epsilon_{\delta}$ , the robustness error bound of UNet in DMs is  $\mathcal{O}(\epsilon_{\delta}(\sum_{i=1}^{N}\kappa_{i}M_{0}^{i}))$ . For standard UNet ( $\kappa_{i} = 1\forall i$ ), this bound becomes a very large bound  $\mathcal{O}(NM_0^N)$ . This implies that standard UNet is sensitive to extra noise in the input, especially when LSC number  $N$  is large ( $M_0 \geq 1$  in practice, see Appendix). Hence, it is necessary to control the coefficients  $\kappa_{i}$  of LSCs which accords with the observations in Fig. 3 (c). Therefore, setting appropriate scaling coefficients  $\{\kappa_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 / \sqrt{2}$ -scaling technique in [27-31] which scales all LSC coefficients  $\kappa_{i}$  from one in UNet to a constant, and is called  $1 / \sqrt{2}$  ( $1 / \sqrt{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\sim 3.4$  we already show that adjusting the scaling coefficients  $\{\kappa_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  $\{\kappa_i\}$  of LSCs in UNet:
$$
\mathrm {C S}: \quad \kappa_ {i} = \kappa^ {i - 1}, \quad (i = 1, 2,.., N, \quad \forall \kappa \in (0, 1 ]). \tag {8}
$$
When  $i = 1$ , the scaling coefficient of the first long skip connection is 1, and thus ensures that at least one  $\kappa_{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  $\kappa_{i} = \kappa$ , e.g.  $\kappa = 1 / \sqrt{2}$  in the  $1 / \sqrt{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  $\mathcal{O}\big(M_0(\kappa M_0)^{N - 1}\big)$  which is  $\mathcal{O}(\kappa^{N - 2})$  times smaller than the bound  $\mathcal{O}(\kappa M_0^N)$  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  $\mathcal{O}\big(\kappa^{2i}m\big)$ . Compared with the oscillation range  $\mathcal{O}\big(Nm\big)$  of standard UNet, CS also greatly alleviates the oscillation since  $\kappa^{2i}\ll 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  $\kappa$  in Eq. (8). To this end, we use the widely used DDPM framework on the benchmarks (CIFAR10, CelebA and ImageNet) to estimate  $\kappa$ . Firstly, Theorem
3.1 reveals that the norm of UNet output is of order  $\mathcal{O}(\| \mathbf{x}_t\| _2^2\sum_{j = 1}^N\kappa_j^2)$ . Next, from the loss function of DMs, i.e., Eq. (2), the target output is a noise  $\epsilon_t\sim \mathcal{N}(0,\mathbf{I})$ . Therefore, we expect that  $\mathcal{O}(\| \mathbf{x}_t\| _2^2\sum_{j = 1}^N\kappa_j^2)\approx \| \epsilon_t\| _2^2$ , and we have  $\sum_{j = 1}^{N}\kappa_{j}^{2}\approx \mathcal{O}(\| \epsilon_{t}\|_{2}^{2} / \| \mathbf{x}_{t}\|_{2}^{2})$ . Moreover, Fig. 3 (b) shows the distribution of  $\| \epsilon_t\| _2^2 /\| \mathbf{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  $\| \epsilon_t\| _2^2 /\| \mathbf{x}_t\| _2^2$ . They may cause the output order  $\mathcal{O}(\| \mathbf{x}_t\| _2^2\sum_{j = 1}^N\kappa_j^2)$  to deviate more from the  $\| \epsilon_t\| _2^2$ , leading to unstable loss and impacting training. Therefore, it is advisable to appropriately increase  $\sum_{j = 1}^{N}\kappa_{j}^{2}$  to address the influence of long-tail samples. [61-65]. So the value of  $\sum_{j = 1}^{N}\kappa_{j}^{2}$  should be between the mean values of  $\| \epsilon_t\| _2^2 /\| \mathbf{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  $\kappa \in [0.5,0.95]$ . In practice, the hyperparameter  $\kappa$  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
Figure 4: The diagram of LS.
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  $\{\kappa_i\}$ . Specifically, for the  $i$ -th long skip connection, its input feature is  $x_{i} \in \mathbb{R}^{B \times N \times D}$ , where  $B$  denotes the batch size. For convolution-based UNet,  $N$  and  $D$  respectively denote the channel number and spatial dimension  $(H \times 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} \in \mathbb{R}^{B \times N \times D}$  into the tiny calibration network  $\zeta_{\phi}$  parameterized by  $\phi$  for prediction:
$$
\mathbf {L S}: \quad \kappa_ {i} = \sigma \left(\zeta_ {\phi} [ \operatorname {G A P} (x _ {i}) ]\right) \in \mathbb {R} ^ {B \times N \times 1}, 1 \leq i \leq N, \tag {9}
$$
where GAP denotes the global average pooling that averages  $x_{i}$  along the last dimension to obtain a feature map of size  $B\times N\times 1$ ;  $\sigma$  is a sigmoid activation function. After learning  $\{\kappa_i\}$ , LS uses  $\kappa_{i}$  to scale the input  $x_{i}$  via element-wise multiplication. In practice, network  $\zeta_{\phi}$  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  $\zeta_{\phi}$  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.
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 / \sqrt{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
Figure 5: Training curve comparison under different training settings in which all experiments here adopt 12-layered UViT and a batch size 128 on CIFAR10 unless specified in the figures.
Table 1: Synthesis performance comparison under different diffusion model settings.
Model (CIFAR10)	Type	#Param	FID↓
EDM [68]	Uncondition	56M	1.97
IDDPM [3]	Uncondition	53M	2.90
DDPM++ cont. [27]	Uncondition	62M	2.55
GenViT [70]	Uncondition	11M	20.20
UNet [4]	Uncondition	36M	3.17
UViT-S/2 [36]	Uncondition	44M	3.11
UNet+1/√2-CS	Uncondition	36M	3.14
UViT-S/2-1/√2-CS	Uncondition	44M	3.11
UNet+CS (ours)	Uncondition	36M	3.05
UNet+LS (ours)	Uncondition	36M	3.07
UViT-S/2+CS (ours)	Uncondition	44M	2.98
UViT-S/2+LS (ours)	Uncondition	44M	3.01
Model (MS-COCO)	Type	#Param	FID↓
VQ-Diffusion [41]	Text-to-Image	370M	19.75
Friro [72]	Text-to-Image	766M	8.97
UNet [36]	Text-to-Image	260M	7.32
UNet+1/√2-CS	Text-to-Image	260M	7.19
UViT-S/2+1/√2-CS	Text-to-Image	252M	5.88
UViT-S/2 [36]	Text-to-Image	252M	5.95
UViT-S/2 (Deep) [36]	Text-to-Image	265M	5.48
UNet+CS (ours)	Text-to-Image	260M	7.04
UNet+LS (ours)	Text-to-Image	260M	6.89
UViT-S/2+CS (ours)	Text-to-Image	252M	5.77
UViT-S/2+LS (ours)	Text-to-Image	252M	5.60
Model (ImageNet 64)	Type	#Param	FID↓
Glow [69]	-	-	3.81
IDDPM (small) [3]	Class-condition	100M	6.92
IDDPM (large) [3]	Class-condition	270M	2.92
ADM [71]	Class-condition	296M	2.07
EDM [68]	Class-condition	296M	1.36
UViT-M/4 [36]	Class-condition	131M	5.85
UViT-L/4 [36]	Class-condition	287M	4.26
UViT-M/4-1/√2-CS	Class-condition	131M	5.80
UViT-L/4-1/√2-CS	Class-condition	287M	4.20
UNet [3]	Class-condition	144M	3.77
UNet-1/√2-CS	Class-condition	144M	3.66
UNet+CS (ours)	Class-condition	144M	3.48
UNet+LS (ours)	Class-condition	144M	3.39
UViT-M/4+CS (ours)	Class-condition	131M	5.68
UViT-M/4+LS (ours)	Class-condition	131M	5.75
UViT-L/4+CS (ours)	Class-condition	287M	3.83
UViT-L/4+LS (ours)	Class-condition	287M	4.08
Model (CelebA)	Type	#Param	FID↓
DDIM [73]	Uncondition	79M	3.26
Soft Truncation [74]	Uncondition	62M	1.90
UViT-S/2 [36]	Uncondition	44M	2.87
UViT-S/2-1/√2-CS	Uncondition	44M	3.15
UViT-S/2+CS (ours)	Uncondition	44M	2.78
UViT-S/2+LS (ours)	Uncondition	44M	2.73
most cases, they consistently accelerate the learning speed by at least  $1.5 \times$  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 / \sqrt{2}$ -CS. All these results are consistent with the analysis in Section 4.1.
Figure 6: Synthesis performance of proposed methods under small batch size and deep architecture.
Robustness to Batch Size and Network Depth. UViT is often sensitive to batch size and network depth [36], since small batch sizes yield more noisy stochastic gradients that further aggravate gradient instability. Additionally, Section 3 shows that increasing network depth impairs the stability of both forward and backward propagation and its robustness to noisy inputs. So in Fig. 6, we evaluate our methods on standard UViT with different batch sizes and network depths. The results show that our CS and LS methods can greatly help UViT on CIFAR10 by showing big improvements on standard UViT and  $1/\sqrt{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
Table 2: Robustness of LS to network design
Ratio r	FID↓	Activation	FID↓	Modules	FID↓
4	3.06	ELU [75]	3.08	Learnable κi	3.46
8	3.10	ReLU [76] (ours)	3.01	IE [77]	3.09
16 (ours)	3.01	Mish [78]	3.08	SE [79] (ours)	3.01
improve LS performance compared with simply treating  $\kappa_{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.
Table 3: Other scaling
Mehtod (CIFAR10)	FID↓
stable τ = 0.1 [46]	3.77
stable τ = 0.3 [46]	3.75
ReZero [48]	3.46
Relationship between LS and CS. Taking UViT on CIFAR10, MS-COCO, ImageNet and CelebA as baselines, we randomly sample 30 Gaussian noises from  $\mathcal{N}(0,\mathbf{I})$  and measure the scaling coefficients, i.e.,  $\kappa_{i}\leq 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  $\kappa_{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.
Figure 7: The  $\kappa_{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,  $\kappa_{i} = \kappa^{N - i + 1}$  ( $\kappa < 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_{\mathrm{r}} = \sum_{i = 1}^{N}\kappa_{i}M_{0}^{i} = \kappa^{N}M_{0} + \kappa^{N - 1}M_{0}^{2} + \ldots +\kappa M_{0}^{N}$ , and could be very large, since  $M_0^N >1$  is large when  $N$  is large and scaling it by a factor  $\kappa$  could not sufficiently control its magnitude. In contrast, our default setting  $\kappa_{i} = \kappa^{i - 1}$  of CS can well control the main term in stability bound:  $S = \sum_{i = 1}^{N}\kappa_{i}M_{0}^{i} = M_{0} + \kappa^{1}M_{0}^{2} + \ldots +\kappa^{N - 1}M_{0}^{N}$ , where the larger terms  $M_0^{i + 1}$  are weighted by smaller coefficients  $\kappa^i$ . In this way,  $S$  is much smaller than  $S_{\mathrm{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 \leq i \leq N)$ , the learnable  $\kappa_{i}$  should be smaller to better control the
Table 4: The FID-10k  $\downarrow$  result of different direction of scaling.
Training step	5k	10k	15k	20k	25k	30k	35k	40k	45k
\( \kappa_i = \kappa^{N-i+1} \)	67.26	33.93	22.78	16.91	15.01	14.01	12.96	12.34	12.26
\( \kappa_i = \kappa^{i-1} \) (ours)	85.19	23.74	15.36	11.38	10.02	8.61	7.92	7.27	6.65
magnitude of  $M_0^i$  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  $\kappa$  as shown in Section 4.1. This means that one still needs to manually select  $\kappa$  from the range. In the future, we may be able to design better methods to assist DMs in selecting better values of  $\kappa$ , or to have better estimates of the range of  $\kappa$ . Another effective solution is to use our LS method which can automatically learn qualified scaling coefficient  $\kappa_{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 \times$  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.
[1] Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11461-11471, 2022.
[2] Tiankai Hang, Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, and Baining Guo. Efficient diffusion training via min-snr weighting strategy. arXiv preprint arXiv:2303.09556, 2023.
[3] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162-8171. PMLR, 2021.
[4] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
[5] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17981-17993, 2021.
[6] Florinel-Alin Croitoru, Vlad Hondru, Radu Tudor Ionescu, and Mubarak Shah. Diffusion models in vision: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023.
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  $\kappa$  greater than 1 in CS, and the experimental results are consistent with our theoretical analysis.
ImageNet  $64\times 64$
CelebA  $64\times 64$
CIFAR10  $32 \times 32$
MS-COCO
Figure 8: The examples of the generated images by UViT with our proposed LS.