SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In practice, the depth  $l$  of block  $a_{i}$  or  $b_{i}$  is a universal constant, therefore, from Corollary B.9,  $\forall \rho \in (0,1)$  and  $j\in \{i + 1,i + 2,\dots,N\}$ , with probability at least  $1 - \mathcal{O}(N)\exp [-\Omega (m\rho^2)]$ , we have
$$
\left\| h _ {\kappa_ {j}} \right\| _ {2} \in \mathcal {O} \left(\left(1 \pm \rho\right) \| \mathbf {x} _ {t} \| _ {2}\right), \tag {46}
$$
$$
h _ {\kappa_ {j}} = \phi \left(\mathbf {W} _ {l - 1} ^ {b _ {i + 1}} \dots \phi \left(\mathbf {W} _ {1} ^ {b _ {i + 1}} r\right)\right), \quad r = b _ {i + 2} \dots \circ b _ {j} \circ \pi_ {a} ^ {j} \circ s, \tag {47}
$$
s.t.  $h_i = f_i(s) \approx \sum_{j=i+1}^{N} \kappa_j \mathbf{W}_l^{b_{i+1}} \circ h_{\kappa_j}$ . From Lemma B.6, the elements of  $\mathbf{W}_l^{b_{i+1}}$  will approximately follow a Gaussian distribution  $\mathcal{N}(0, \mathcal{O}(m^{-1}))$ , therefore,  $\mathbb{E}(\mathbf{W}_l^{b_{i+1}} \circ h_{\kappa_j}) = 0$ . Moreover, from Lemma B.8,
$$
\mathbf {V a r} \left(\mathbf {W} _ {l} ^ {b _ {i + 1}} \circ h _ {\kappa_ {j}}\right) \approx h _ {\kappa_ {j}} ^ {\mathrm {T}} \mathcal {O} \left(m ^ {- 1}\right) \cdot \mathbf {I} _ {m \times m} h _ {\kappa_ {j}} = \mathcal {O} \left(m ^ {- 1}\right) \| h _ {\kappa_ {j}} \| _ {2} ^ {2} \cdot \mathbf {I} _ {m \times m}. \tag {48}
$$
Therefore [47],  $h_i \approx \sum_{j=i+1}^{N} \kappa_j \mathbf{W}_l^{b_{i+1}} \circ h_{\kappa_j}$  can be approximately regarded as a Gaussian variable with zero mean and covariance matrix  $\sum_{j=i+1}^{N} \kappa_j^2 \mathcal{O}(m^{-1}) \| h_{\kappa_j} \|_2^2 \cdot \mathbf{I}_{m \times m}$ . According to Lemma B.5, we have
$$
\left\| h _ {i} \right\| _ {2} ^ {2} \approx \mathbb {E} \left(\sum_ {j = i + 1} ^ {N} \kappa_ {j} \mathbf {W} _ {l} ^ {b _ {i + 1}} \circ h _ {\kappa_ {j}}\right) = m \cdot \sum_ {j = i + 1} ^ {N} \kappa_ {j} ^ {2} \mathcal {O} (m ^ {- 1}) \| h _ {\kappa_ {j}} \| _ {2} ^ {2} = \sum_ {j = i + 1} ^ {N} \kappa_ {j} ^ {2} \mathcal {O} (1) \| h _ {\kappa_ {j}} \| _ {2} ^ {2}. \tag {49}
$$
From Eq. (46) and Eq. (49), we have
$$
(1 - \rho) ^ {2} \left[ c _ {1} \| \mathbf {x} _ {t} \| _ {2} ^ {2} \cdot \sum_ {j = i + 1} ^ {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 + 1} ^ {N} \kappa_ {j} ^ {2} + c _ {2} \right], \tag {50}
$$
where we use  $c_{1}$  to denote  $\mathcal{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.
Lemma 3.2. For  $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon_t$  defined in Eq. (14) as a input of UNet,  $\epsilon_t\sim \mathcal{N}(0,\mathbf{I})$  if  $\mathbf{x}_0$  follow 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 {51}
$$
Proof. For  $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon_t$ , we have
$$
\mathbb {E} \| \mathbf {x} _ {t} \| _ {2} ^ {2} = \mathbb {E} _ {t, \mathbf {x} _ {0}} \left(\bar {\alpha} _ {t} \sum_ {j = 1} ^ {m} \mathbf {x} _ {0} ^ {2} (j) + (1 - \bar {\alpha} _ {t}) m\right), \tag {52}
$$
where  $\mathbf{x}_0(j)$  represents the  $j$ -th element of  $\mathbf{x}_0$ . The equality in Eq. (52) follows from the fact that, from Eq. (14) and Lemma B.5, we know that  $\mathbf{x}_t$  follows  $\mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t) \mathbf{I})$ . Therefore,  $\| \mathbf{x}_t \|_2^2$  follows a chi-square distribution, whose mathematical expectation is  $\mathbb{E}_{t,\mathbf{x}_0} (\bar{\alpha}_t \sum_{j=1}^{m} \mathbf{x}_0^2(j) + (1 - \bar{\alpha}_t)m)$ .
Moreover, since  $\mathbf{x}_0$  follow the uniform distribution  $U[-1,1]$ , we have
$$
\begin{array}{l} \mathbb {E} \| \mathbf {x} _ {t} \| _ {2} ^ {2} = \mathbb {E} _ {t} \bar {\alpha} _ {t} \cdot \sum_ {j = 1} ^ {m} \int_ {- \infty} ^ {\infty} x ^ {2} P _ {x \sim U [ - 1, 1 ]} (x) d x + \mathbb {E} _ {t} (1 - \bar {\alpha} _ {t}) m \tag {53} \\ = \mathbb {E} _ {t} \bar {\alpha} _ {t} \cdot m / 3 + \mathbb {E} _ {t} (1 - \bar {\alpha} _ {t}) m = (1 - 2 \mathbb {E} _ {t} \bar {\alpha} _ {t} / 3) m = \mathcal {O} (m). \\ \end{array}
$$
where  $P_{x\sim U[-1,1]}(x)$  is density function of the elements in  $\mathbf{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  $\bar{\alpha}_{t}$  monotonically decreases to 0 as  $t$  increases, with  $\bar{\alpha}_{t} \leq 1$ . Consequently, we can easily deduce that  $\mathbb{E}_t\bar{\alpha}_t$  is of the order  $\mathcal{O}(1)$ . For instance, in the setting of DDPM, we have  $\mathbb{E}_t\bar{\alpha}_t = \mathbb{E}_t\left(\prod_{i=1}^t \sqrt{1 - \frac{0.02i}{1000}}\right) \approx 0.27 = \mathcal{O}(1)$ .
Theorem 3.3. Assume that all learnable matrices of UNet in Eq. (16) 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_ {s} (\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 {54}
$$
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.  $n$  is the size of the training set.
Proof. From Eq. (16) and Eq. (44), we can approximate the output of UNet  $f_{0}(\mathbf{x}_{t})$  by weight sum of  $N$  forward path  $g^{j}, 1 \leq j \leq N$ , i.e.,
$$
f _ {0} \left(\mathbf {x} _ {t}\right) \approx \sum_ {j = 1} ^ {N} \kappa_ {i} \cdot g ^ {j}. \tag {55}
$$
where  $g^{j}$  denote  $b_{1} \circ b_{2} \circ \dots \circ b_{j} \circ a_{j} \circ \dots \circ a_{1} \circ \mathbf{x}_{t}$ , and  $1 \leq j \leq N$ .
Moreover, we use  $g_{p}^{j,q}$  represent the feature map after  $\mathbf{W}_p^q$  in forward path  $g^{j}$ , i.e.,  $g_{p}^{j,q} = \mathbf{W}_{p}^{q}h_{p - 1}^{j,q}$  and  $h_p^{j,q} = \phi (g_p^{j,q})$ . Next, we first estimate  $\| \nabla_{\mathbf{W}_p^{a_i}}\ell_s(\mathbf{W})\| _2^2$
$$
\begin{array}{l} \| \nabla_ {\mathbf {W} _ {p} ^ {a _ {i}}} \ell_ {s} (\mathbf {W}) \| _ {2} ^ {2} = \| \sum_ {j = i} ^ {N} \partial \ell_ {s} (\mathbf {W}) / \partial g ^ {j} \cdot \partial g ^ {j} / \partial g _ {p} ^ {j, a _ {i}} \cdot \partial g _ {p} ^ {j, a _ {i}} / \partial \mathbf {W} _ {p} ^ {a _ {i}} \| _ {2} ^ {2} \\ \leq \sum_ {j = i} ^ {N} \| \partial \ell_ {s} (\mathbf {W}) / \partial g ^ {j} \cdot \partial g ^ {j} / \partial g _ {p} ^ {j, a _ {i}} \cdot \partial g _ {p} ^ {j, a _ {i}} / \partial \mathbf {W} _ {p} ^ {a _ {i}} \| _ {2} ^ {2} \\ \leq \sum_ {j = i} ^ {N} \| \mathbf {D} \left(\mathbf {W} _ {p + 1} ^ {a _ {i}}\right) ^ {\mathrm {T}} \dots \mathbf {D} \left(\mathbf {W} _ {l} ^ {a _ {i}}\right) ^ {\mathrm {T}} \dots \mathbf {D} \left(\mathbf {W} _ {l} ^ {b _ {1}}\right) ^ {\mathrm {T}} \cdot \partial \ell_ {s} (\mathbf {W}) / \partial g ^ {j} \| _ {2} ^ {2} \cdot \| h _ {p - 1} ^ {j, a _ {i}} \| _ {2} ^ {2} \tag {56} \\ \end{array}
$$
From Lemma B.11 and Corollary B.9, with probability at least  $1 - \mathcal{O}(nN) \exp[-\Omega(m)]$ , we have
$$
\begin{array}{l} \| \nabla_ {\mathbf {W} _ {p} ^ {a _ {i}}} \ell_ {s} (\mathbf {W}) \| _ {2} ^ {2} \leq \sum_ {j = i} ^ {N} \mathcal {O} \left(\| \partial \ell_ {s} (\mathbf {W}) / \partial g ^ {j} \| _ {2} ^ {2}\right) \cdot \mathcal {O} (1) \cdot \| \mathbf {x} _ {t} \| _ {2} ^ {2} \\ = \sum_ {j = i} ^ {N} \left(\mathcal {O} \left(\kappa_ {j} ^ {2} \ell_ {s} (\mathbf {W})\right) \cdot \mathcal {O} (1) \cdot \| \mathbf {x} _ {t} \| _ {2} ^ {2} \right. \tag {57} \\ = \mathcal {O} \left(\ell_ {s} (\mathbf {W}) \cdot \| \mathbf {x} _ {t} \| _ {2} ^ {2} \cdot \sum_ {j = i} ^ {N} \kappa_ {j} ^ {2} + c _ {3}\right) \\ \end{array}
$$
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.,  $\| \nabla_{\mathbf{W}_p^{b_i}}\ell_s(\mathbf{W})\| _2^2\lesssim \mathcal{O}\left(\ell_s(\mathbf{W})\cdot \| \mathbf{x}_t\| _2^2\cdot \sum_{j = i}^{N}\kappa_j^2 +c_3\right)$ .
B.4 The proof of Theorem 3.4.
Theorem 3.4. For UNet in Eq. (16), 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 {58}
$$
Proof. From Eq. (16), we have
$$
\mathbf {U N e t} (\mathbf {x} _ {t}) = f _ {0} (\mathbf {x} _ {t}) = b _ {1} \circ \left[ \kappa_ {1} \cdot a _ {1} \circ \mathbf {x} _ {t} + f _ {1} \left(a _ {1} \circ \mathbf {x} _ {t}\right) \right], \tag {59}
$$
$$
\mathbf {U N e t} \left(\mathbf {x} _ {t} ^ {\epsilon_ {\delta}}\right) = f _ {0} \left(\mathbf {x} _ {t} ^ {\epsilon_ {\delta}}\right) = b _ {1} \circ \left[ \kappa_ {1} \cdot a _ {1} \circ \mathbf {x} _ {t} ^ {\epsilon_ {\delta}} + f _ {1} \left(a _ {1} \circ \mathbf {x} _ {t} ^ {\epsilon_ {\delta}}\right) \right]. \tag {60}
$$
Using Taylor expansion, we have
$$
\begin{array}{r l} \mathbf {U N e t} (\mathbf {x} _ {t} ^ {\epsilon_ {\delta}}) - \mathbf {U N e t} (\mathbf {x} _ {t}) & = b _ {1} \circ \left[ \kappa_ {1} \cdot a _ {1} \circ (\mathbf {x} _ {t} ^ {\epsilon_ {\delta}} - \mathbf {x} _ {t}) + f _ {1} (a _ {1} \circ \mathbf {x} _ {t} ^ {\epsilon_ {\delta}}) - f _ {1} (a _ {1} \circ \mathbf {x} _ {t}) \right] \\ & = b _ {1} \circ a _ {1} \circ \left[ \kappa_ {1} \cdot (\mathbf {x} _ {t} ^ {\epsilon_ {\delta}} - \mathbf {x} _ {t}) + (\mathbf {x} _ {t} ^ {\epsilon_ {\delta}} - \mathbf {x} _ {t}) ^ {T} \nabla_ {z} f _ {1} (z) | _ {z = a _ {1} \circ \mathbf {x} _ {t}} \right]. \end{array} \tag {61}
$$
$$
\begin{array}{r l} \| \mathbf {U N e t} \left(\mathbf {x} _ {t} ^ {\epsilon_ {\delta}}\right) - \mathbf {U N e t} \left(\mathbf {x} _ {t}\right) \| _ {2} & = \| b _ {1} \circ a _ {1} \circ \left[ \kappa_ {1} \cdot \left(\mathbf {x} _ {t} ^ {\epsilon_ {\delta}} - \mathbf {x} _ {t}\right) + \left(\mathbf {x} _ {t} ^ {\epsilon_ {\delta}} - \mathbf {x} _ {t}\right) ^ {T} \nabla_ {z} f _ {1} (z) \mid_ {z = a _ {1} \circ \mathbf {x} _ {t}} \right] \| _ {2} \\ & \leq \epsilon_ {\delta} \cdot \| b _ {1} \circ a _ {1} \| _ {2} \cdot \left(\kappa_ {1} + \| \nabla_ {z} f _ {1} (z) | _ {z = a _ {1} \circ \mathbf {x} _ {t}} \| _ {2}\right) \end{array} \tag {62}
$$
Next, we estimate the  $\| \nabla_{z}f_{1}(z)\mid_{z = a_{1}\circ \mathbf{x}_{t}}\|_{2}$  , From Eq. (16), we have
$$
\begin{array}{l} \nabla_ {z} f _ {1} (z) = \nabla_ {z} b _ {2} \circ [ \kappa_ {2} \cdot a _ {2} \circ z + f _ {2} (a _ {2} \circ z) ] \\ = \kappa_ {2} \cdot b _ {2} \circ a _ {2} \circ \mathbf {I} + b _ {2} \circ a _ {2} \circ \nabla_ {u} f _ {2} (u) | _ {u = a _ {2} \circ z}, \tag {65} \\ \end{array}
$$
$$
\begin{array}{l} \left\| \nabla_ {z} f _ {1} (z) \right| _ {z = a _ {1} \circ \mathbf {x} _ {t}} \| _ {2} = \| (\kappa_ {2} \cdot b _ {2} \circ a _ {2} \circ \mathbf {I} + b _ {2} \circ a _ {2} \circ \nabla_ {u} f _ {2} (u) | _ {u = a _ {2} \circ z}) | _ {z = a _ {1} \circ \mathbf {x} _ {t}} \| _ {2} \\ = \| \kappa_ {2} \cdot b _ {2} \circ a _ {2} \circ \mathbf {I} + b _ {2} \circ a _ {2} \circ \nabla_ {u} f _ {2} (u) | _ {u = a _ {2} \circ a _ {1} \circ \mathbf {x} _ {t}} \| _ {2} \tag {64} \\ \leq \| b _ {2} \circ a _ {2} \| _ {2} \cdot \left(\kappa_ {2} + \| \nabla_ {u} f _ {2} (u) | _ {u = a _ {2} \circ a _ {1} \circ \mathbf {x} _ {t}} \| _ {2}\right). \\ \end{array}
$$
From Eq. (62) and Eq. (64), let  $\Delta = \| \mathbf{UNet}(\mathbf{x}_t^{\epsilon_6}) - \mathbf{UNet}(\mathbf{x}_t)\|_2$ , we have
$$
\begin{array}{l} \Delta \leq \epsilon_ {\delta} \cdot \| b _ {1} \circ a _ {1} \| _ {2} \cdot \left(\kappa_ {1} + \| \nabla_ {z} f _ {1} (z) | _ {z = a _ {1} \circ \mathbf {x} _ {\ell}} \| _ {2}\right) \\ \leq \epsilon_ {\delta} \cdot \| b _ {1} \circ a _ {1} \| _ {2} \cdot \left[ \kappa_ {1} + \| b _ {2} \circ a _ {2} \| _ {2} \cdot \left(\kappa_ {2} + \| \nabla_ {u} f _ {2} (u) | _ {u = a _ {2} \circ a _ {1} \circ \mathbf {x} _ {t}} \| _ {2}\right) \right] \\ \leq \epsilon_ {\delta} \cdot \left(\kappa_ {1} \| b _ {1} \circ a _ {1} \| _ {2} + \kappa_ {2} \| b _ {1} \circ a _ {1} \| _ {2} \cdot \| b _ {2} \circ a _ {2} \| _ {2} + \| b _ {1} \circ a _ {1} \| _ {2} \cdot \| b _ {2} \circ a _ {2} \| _ {2} \cdot \| \nabla_ {u} f _ {2} (u) | _ {u = a _ {2} \circ a _ {1} \circ \mathbf {x} _ {\ell}} \| _ {2}\right) \\ \leq \epsilon_ {\delta} \left[ \sum_ {i = 1} ^ {N} \left(\kappa_ {i} \prod_ {j = 1} ^ {i} \| b _ {j} \circ a _ {j} \| _ {2}\right) + \prod_ {j = 1} ^ {N} \| b _ {j} \circ a _ {j} \| _ {2} \cdot L _ {0} \right] \\ \leq \epsilon_ {\delta} \cdot \left(\kappa_ {1} M _ {0} + \kappa_ {2} M _ {0} ^ {2} \dots + \kappa_ {N} M _ {0} ^ {N} + c _ {0}\right), \tag {65} \\ \end{array}
$$
where  $c_{0} = M_{0}^{N}L_{0}$ , and hence we have  $\| \mathbf{UNet}(\mathbf{x}_t^{\epsilon_\delta}) - \mathbf{UNet}(\mathbf{x}_t)\|_2 \leq \epsilon_\delta \left[\sum_{i=1}^{N} \kappa_i M_0^i + c_0\right]$  for UNet in Eq. (16) and  $\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})$ , with small perturbation  $\epsilon_\delta$ .
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 \circ a_i \|_2$  for the Eq. (16) on the CIFAR10 and CelebA datasets. We observe that  $\| b_i \circ 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 \circ a_i \|_2, 1 \leq i \leq N$ , we conclude that  $M_0 \geq 1$ .
Figure 11: The norm of  $\| b_{i}\circ a_{i}\|_{2}$  on different dataset.
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.
Figure 12: More examples for feature oscillation under different random seeds (UNet's depth = 12).
In the main text, we specify that the domain of  $\kappa$  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  $\kappa > 1$  on the performance of the diffusion model. We observe that generally, having  $\kappa > 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  $\kappa > 1$ . These experimental results regarding  $\kappa$  align with the analysis presented in Section 3.
Figure 13: The situation that the  $\kappa_{i} = \kappa > 1$  under different settings.