SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 2:  $\chi_{\mathcal{I}}^{\star}$  empirical phase diagrams for an MLP with  $L = 50$ ,  $N = 500$ . The solid lines indicate the critical lines obtained through infinite width limit calculations, and the stars indicate the critical points. The dotted lines in the rightmost column correspond to the critical lines for  $\mu < 1$  case. For networks with LayerNorm and  $\mu = 1$ ,  $\chi_{\mathcal{I}}^{\star} = 1$  holds on the entire  $\sigma_b - \sigma_w$  plane. We also note that for erf activation, the case  $\mu = 1$  without LayerNorm is subcritical and has a large correlation length.
Definition 3.1 (Normalized preactivations).
$$
\tilde {h} _ {i} ^ {l} = \frac {h _ {i} ^ {l} - \mathbb {E} [ h ^ {l} ]}{\sqrt {\mathbb {E} [ (h ^ {l}) ^ {2} ] - \mathbb {E} [ h ^ {l} ] ^ {2}}} \xrightarrow {N _ {l} \rightarrow \infty} \frac {1}{\sqrt {\mathcal {K} ^ {l}}} h _ {i} ^ {l}, \tag {19}
$$
where we have introduced  $\mathbb{E}[h^l] = (\sum_{i=1}^{N_l} h_i^l) / N_l$ . In the limit of infinite width  $\mathbb{E}[h^l] = 0$  and  $\mathbb{E}[(h^l)^2] = \mathcal{K}^l$ , defined according to (6).
Normalized preactivations,  $\tilde{h}_j^l$ , are distributed according to  $\mathcal{N}(0,1)$  for all  $l,\sigma_w,\sigma_b$ . The norms are, therefore, always finite and the condition  $\chi_{\mathcal{K}}^{\star} = 1$  is trivially satisfied. This results in a critical line rather than a critical point.
The recurrence relations (9)-(11) for the NNGP and partial Jacobians are only slightly modified
$$
\mathcal {K} ^ {l + 1} = \sigma_ {w} ^ {2} \mathbb {E} _ {\theta} \left[ \frac {1}{N _ {l}} \sum_ {i = 1} ^ {N _ {l}} \phi \left(\tilde {h} _ {i} ^ {l}\right) ^ {2} \right] + \sigma_ {b} ^ {2} \quad \chi_ {\mathcal {I}} ^ {l} = \frac {\sigma_ {w} ^ {2}}{\mathcal {K} ^ {l}} \mathbb {E} _ {\theta} \left[ \frac {1}{N _ {l}} \sum_ {i = 1} ^ {N _ {l}} \phi^ {\prime} \left(\tilde {h} _ {i} ^ {l}\right) ^ {2} \right]. \tag {20}
$$
Assuming that the value of  $\chi_{\mathcal{I}}^{l}$  at the fixed point is  $\chi_{\mathcal{I}}^{\star}$ , the network is critical when (13) holds.
$\chi_{\mathcal{J}}^{l}$  (20) is depth independent and changes slowly with  $\sigma_w$  and  $\sigma_b$ . Thus,  $\chi_{\mathcal{J}}^{\star}$  remains close to 1 for a wider range of hyperparameters. Consequently, the correlation length is large even away from criticality, leading to a much higher trainability of deep networks.
4 Residual (Skip) Connections
Adding residual connections between the network layers is a widely used technique to facilitate the training of deep networks. Originally introduced [21] in the context of convolutional neural networks [29] (CNNs) for image recognition, residual connections have since been used in a variety of networks architectures and tasks [50, 13].
Consider (1) with non-zero  $\mu$  and without LayerNorm layers. Then the recurrence relations (9)-(11) for the NNGP kernel and  $\chi_{\mathcal{T}}^l$  are modified as follows
$$
\mathcal {K} ^ {l + 1} = \sigma_ {w} ^ {2} \mathbb {E} _ {\theta} \left[ \frac {1}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \phi \left(h _ {j} ^ {l}\right) ^ {2} \right] + \sigma_ {b} ^ {2} + \mu^ {2} \mathcal {K} ^ {l}, \quad \chi_ {\mathcal {I}} ^ {l} = \sigma_ {w} ^ {2} \mathbb {E} _ {\theta} \left[ \frac {1}{N _ {l}} \sum_ {k = 1} ^ {N _ {l}} \phi^ {\prime} \left(h _ {k} ^ {l}\right) ^ {2} \right] + \mu^ {2}. \tag {21}
$$
Figure 3: Trainability (Training Accuracy) of deep MLP ( $N_{l} = 500$ ,  $L = 50$ ) on FashionMNIST. The combination of LayerNorm and  $\mu = 1$  makes the network everywhere-trainable. The subcritical case of erf activation without LayerNorm, and  $\mu = 1$  also has enhanced trainability. The dashed white lines denote the (analytical) critical lines.
Remark 4.1. When  $\mu < 1$ , the fixed point value of NNGP kernel is scaled by  $(1 - \mu^2)^{-1}$ . For  $\mu = 1$ , the critical point is formally at  $(0,0)$ .
Remark 4.2. For  $\mu = 1$ , (21) implies that  $\chi_{\mathcal{I}}^l \geq 1$ , where the equality holds on the  $\sigma_w = 0$  axis. Consequently, APJN exponentially diverges as a function of depth  $l$  for all  $\sigma_w > 0$ . In this case,  $\sigma_w$  needs to be taken sufficiently close to 0 to ensure trainability at large depths.
When  $\mu < 1$ , residual connections amplify the chaotic phase and decrease the correlation length away from criticality for unbounded activation functions.
Solving the recurrence relations (21) for erf activation, we find an effect observed in Yang and Schoenholz [62] for tanh activation. They noted that tanh-like MLP networks with skip connections "hover over the edge of chaos". We quantify their observation as follows.
Theorem 4.3. Let  $f(x)$  be a deep MLP network with erf activation function and residual connections of strength  $\mu = 1$ . Then in the limit  $N_{l} \to \infty$
- The NNGP kernel  $K^l$  linearly diverges with depth  $l$ .
-  $\chi_{\mathcal{T}}^{l}$  approaches 1 from above (Fig. 2):  $\chi_{\mathcal{T}}^{l} \approx 1 + \tilde{c} / \sqrt{l}$ , where  $\tilde{c} = 2\sigma_w^2 / (\pi \sqrt{\sigma_w^2 + \sigma_b^2})$  is a non-universal constant. Consequently, APJN diverges as a stretched exponential:  $\mathcal{T}^{l_0,l} = O(e^{\sqrt{\frac{l}{\lambda}}})$ , where  $\lambda = 1 / (4\tilde{c}^2)$  is the new length scale.
We will refer to this case as subcritical. Although  $\chi_{\mathcal{I}}^{\star}$  reaches 1, the APJN still diverges with depth faster than any power law. The growth is controlled by the new scale  $\lambda$ . To control the gradient we would like to make  $\lambda$  large, which can be accomplished by decreasing  $\sigma_w$ . In this case, the trainability is enhanced (see Fig. 3). Similar results hold for tanh activation function [62], however in that case there is no explicit expression for  $\tilde{c}$ .
5 Residual Connections + LayerNorm
In practice, it is common to use a combination of residual connections and LayerNorm. Using (1), the recurrence relations (9)-(11) for NNGP and partial Jacobians are modified as follows
$$
\mathcal {K} ^ {l + 1} = \sigma_ {w} ^ {2} \mathbb {E} _ {\theta} \left[ \frac {1}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \phi \left(\tilde {h} _ {j} ^ {l}\right) ^ {2} \right] + \sigma_ {b} ^ {2} + \mu^ {2} \mathcal {K} ^ {l}, \quad \chi_ {\mathcal {I}} ^ {l} = \frac {\sigma_ {w} ^ {2}}{\mathcal {K} ^ {l}} \mathbb {E} _ {\theta} \left[ \frac {1}{N _ {l}} \sum_ {k = 1} ^ {N _ {l}} \phi^ {\prime} \left(\tilde {h} _ {k} ^ {l}\right) ^ {2} \right] + \mu^ {2}. \tag {22}
$$
Remark 5.1. For  $\mu < 1$ , (22) implies that the fixed point value of NNGP kernel is scaled by  $1 - \mu^2$ . Moreover, residual connections do not shift the phase boundary. The interference between residual connections and LayerNorm brings  $\chi_{\mathcal{I}}^{l}$  closer to 1 on the entire  $\sigma_b - \sigma_w$  plane (as can be seen from Fig. 2). Therefore the correlation length  $\xi$  is improved in both the phases, allowing for training of deeper networks. At criticality, Jacobians linearly diverge with depth.
As was mentioned before, the combination of LayerNorm and residual connections dramatically enhances correlation length, leading to a more stable architecture. This observation is formalized by 1.3. The proof leverages the solutions of (22) close to the fixed point, and is fleshed out in Appendix F.
Remark 5.2. When  $\mu = 1$ , the correlation length diverges for any initialization.
5.2 provides an alternative perspective on architecture design. On the one hand, one can use (16) to initialize a given network at criticality. Alternatively, one can use a combination of residual connections and LayerNorm to ensure that the network will train well, irrespective of the initialization.
Remark 5.3. When  $\mu = 1$ , the condition  $\chi_{\mathcal{I}}^{\star} = 1$  holds on the entire  $\sigma_b - \sigma_w$  and for any activation function  $\phi$  (see Fig. 2). NNGP kernel diverges linearly, while APJN diverges algebraically with the critical exponent of  $\zeta = O(1)$ . The exact value of the critical exponent depends on the activation function and the ratio  $\sigma_b / \sigma_w$ . The trainability is dramatically enhanced, as shown in Fig. 3.
Remark 5.4. Networks with BatchNorm [26], used in conjunction with residual connection of strength  $\mu = 1$ , also enjoy this everywhere criticality and enhanced trainability [63, 23] (See Appendix B).
6 Modern Architectures
ResNet110 ResNet is one of the most widely used architectures in computer vision[21]. Figure 4 shows the results for ResNetV2 [22]; with BatchNorm replaced with LayerNorm. (See Appendix B for BatchNorm results and discussions.)
At  $\mu = 1$ ,  $\sigma_w^2 - \sigma_b^2$  phase diagram shows everywhere-criticality, as expected from our theory. For  $\sigma_b^2 = 0$ ,  $\sigma_w^2 - \mu$  phase diagram gets closer to criticality as we increase  $\mu$ , which results in better training at higher  $\mu$ .
Additionally, networks with larger  $\mu$  enjoy better trainability due to the suppression of finite width corrections to our theory (decreased effective  $L / N$ ).
Figure 4: ResNet110(LayerNorm): Left to right: (1)  $\chi_{\mathcal{J}}^{\star}$  phase diagram w.r.t.  $(\sigma_w^2, \sigma_b^2)$  at  $\mu = 1$ . The network is everywhere-critical in this case. (2)  $\chi_{\mathcal{J}}^{\star}$  phase diagram w.r.t.  $(\sigma_w^2, \mu)$ . (3) Training accuracy w.r.t.  $\sigma_w^2$ , for different values of  $\mu$ . Trainability improves with higher  $\mu$ , as the network gets closer to everywhere-critical phase. For  $\mu = 0$  and small  $\sigma_w^2$  the network is in the ordered phase, and the output of every block is almost zero - this explains the poor trainability. (4) Training accuracy w.r.t.  $\mu$  at  $(\sigma_w^2, \sigma_b^2) = (2, 0)$ . We see a monotonic improvement in trainability.
MLP-Mixer MLP-Mixer architecture is a recent example of MLP approach to computer vision [48]. It can be analyzed using the tools presented above. Fig. 5 shows the phase and training diagrams. Further details can be found in the SM. Note that the original architecture uses  $\mu = 1.0$ .
Test accuracies for figures 4, 5 follow similar trends as train accuracies. (See Appendices B, H).
Figure 5: MLP-Mixer. Left to right: (1)  $(2)\mu = 0.5, 1.0$  phase diagrams  $(\chi_{\mathcal{I}}^{\star})$ . The network becomes everywhere-critical when  $\mu = 1$ . The solid black line indicates the empirical phase boundary. Stars denote points we selected to train on CIFAR-10. (3)  $(4)\mu = 0.5, 1.0$  MLP-Mixer training curves. All, but one, networks are  $L = 100$  blocks deep. We see that as we increase  $\mu$  from 0.5 to 1, the trainability of all networks increases, and they are less sensitive to initialization.
We have introduced partial Jacobians and their averaged norms as tools to analyze the propagation of gradients through deep neural networks at initialization. Using APJN evaluated close to the output,  $\mathcal{J}^{L - 2,L - 1}\approx \chi_{\mathcal{J}}^{\star}$ , we have introduced a very cheap and simple empirical test for criticality. We have also shown that criticality formulated in terms of partial Jacobians is equivalent to criticality studied previously in literature [44, 45, 34]. Additionally, APJN can be utilized to quantify criticality in inhomogeneous (i.e. no periodic stacking of blocks) networks [23].
We have investigated homogeneous architectures that include fully-connected layers, normalization layers and residual connections. In the limit of infinite width, we showed that (i) in the presence of LayerNorm, the critical point generally becomes a critical line, making the initialization problem much easier, (ii) LayerNorm applied to preactivations enhances correlation length leading to improved trainability, (iii) combination of  $\mu = 1$  residual connections and erf activation function enhances correlation length driving the network to a subcritical phase with APJN growing according to a stretched exponential law, (iv) combination of residual connections and LayerNorm drastically increases correlation length leading to improved trainability, (v) when  $\mu = 1$  and LayerNorm is applied to preactivations the network is critical on the entire  $\sigma_{b} - \sigma_{w}$  plane.
We have considered examples of modern architectures: ResNet and MLP-Mixer. We showed that at  $\mu = 1$ , they are critical everywhere, due to the interaction between LayerNorm and residual connections. We have studied ResNet at different  $\mu$ 's, showing that higher  $\mu$  enjoys better trainability, due to improved initialization and suppressed effective  $L / N$  corrections.
We have empirically demonstrated that deep (100 blocks) MLP-Mixer with  $\mu = 1$  trains well for various initializations. In comparison, for  $\mu = 0.5$ , it only trains well close to the critical line.
Our work shows that an architecture can be designed to have a large correlation length leading to a guaranteed trainability with SGD for any initialization scheme.
8 Limitations and Future Works
While our method is empirically applicable to Transformer architectures (see Appendix C for phase diagrams), there remains a notable absence of robust theoretical foundations for initializing Transformers. We aim to extend our theory to the attention mechanism and determine the optimal way to initialize Transformer architectures.
Another challenge arises from network topology. For graph neural networks, the generalization of our methods remains ambiguous. We aspire to elucidate this in future research.
Lastly, we want to scale our experiments to larger image datasets as well as language tasks, with the hope that our methods could help people train large models more efficiently.
We thank T. Can, G. Gur-Ari, B. Hanin, D. Roberts and S. Yaida for comments on the manuscript. A.G.'s work at the University of Maryland was supported in part by NSF CAREER Award DMR-2045181, Sloan Foundation and the Laboratory for Physical Sciences through the Condensed Matter Theory Center.
[53] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017.
[54] Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington. Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks. In International Conference on Machine Learning, pages 5393-5402. PMLR, 2018.
[55] Lechao Xiao, Jeffrey Pennington, and Sam Schoenholz. Disentangling trainability and generalization in deep learning, 2020.
[56] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533. PMLR, 2020.
[57] Jingjing Xu, Xu Sun, Zhiyuan Zhang, Guangxiang Zhao, and Junyang Lin. Understanding and improving layer normalization. arXiv preprint arXiv:1911.07013, 2019.
[58] Sho Yaida. Non-gaussian processes and neural networks at finite widths. In Mathematical and Scientific Machine Learning, pages 165-192. PMLR, 2020.
[59] Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.04760, 2019.
[60] Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. 2019.
[61] Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning, pages 11727-11737. PMLR, 2021.
[62] Greg Yang and Samuel S Schoenholz. Mean field residual networks: On the edge of chaos. arXiv preprint arXiv:1712.08969, 2017.
[63] Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S Schoenholz. A mean field theory of batch normalization. In International Conference on Learning Representations, 2018.
[64] Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, and Shuicheng Yan. Metaformer is actually what you need for vision. arXiv preprint arXiv:2111.11418, 2021.
[65] Hongyi Zhang, Yann N. Dauphin, and Tengyu Ma. Residual learning without normalization via better initialization. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=H1gsz30cKX.
We implemented our methods using PyTorch [42] hooks and an efficient Jacobian approximate algorithm [24].
Figure 1: We generated MNIST-like inputs, where all elements are sampled from the Gaussian distribution  $\mathcal{N}(0,1)$ .  $\mathcal{I}^{0,l}$  data was averaged over 100 different parameter-initializations. Networks were initialized width  $N_{l} = 1000$ . For erf plot we initialized at critical point  $(\sigma_w,\sigma_b) = (\sqrt{\frac{\pi}{4}},0)$ , used depth  $L = 250$  and the fitting was done with data points collected at depth  $l > 100$ ; for ReLU plot we initialized at critical point  $(\sigma_w,\sigma_b) = (\sqrt{2},0)$ , used depth  $L = 100$  and the fitting was done with all data points; for the  $\mu = 1$  Pre-LN plot, we initialized both networks at  $(\sigma_w,\sigma_b) = (\sqrt{2},0)$ , used depth  $L = 250$  and the fitting was done with  $l > 100$  data points.
Figure 2: All the phase diagrams were plotted using  $\chi_{\mathcal{J}}^{L-1}$  generated from networks with  $L = 50$  and  $N_l = 500$ . We used hooks to obtain the gradients that go into calculating  $\chi_{\mathcal{J}}^{L-1}$ .  $\chi_{\mathcal{J}}^{L-1}$  data was averaged over 100 different parameter-initializations. Inputs were generated from a normal Gaussian distribution and have dimension  $28 \times 28$ . Generating the data for the figure took approximately 2 days on Google Colab Pro (single Tesla P100 GPU).
Figure 3: In all cases, networks are trained for 10 epochs using stochastic gradient descent with CrossEntropy loss. We used the Fashion MNIST dataset [53]. All networks had depth  $L = 50$  and width  $N_{l} = 500$ . The learning rates were logarithmically sampled within  $(10^{-5}, 1)$ . Generating the data for the figure took approximately 12 days on Google Colab Pro (single Tesla P100 GPU).
Figure 4: (1) We made the  $\sigma_w^2 - \sigma_b^2$  phase diagram for ResNet110(LayerNorm) by averaging over 100 different parameter-initializations. The  $\sigma_w^2 - \mu$  phase diagram was made by averaging over 200 parameters initialization. (3)(4) We used SGD with momentum  $= 0.9$  and batch size 128. For selecting the learning rate we ran a grid-search over 0.001, 0.005, 0.01, 0.02, 0.5 for 10 epochs; with weight decay  $\lambda = 10^{-4}$ . All models were trained for 50 epochs and averaged over 3 random seeds. It takes 6 GPU days in total on a single NVIDIA RTX 3090 GPU.
Figure 5: (1)(2)We made the phase diagram for MLP-Mixer with 30 blocks and averaged over 100 different parameter-initializations. (3)(4)We used network with  $L = 100$ , patch size  $4 \times 4$ , hidden size  $C = 128$ , two MLP dimensions  $N_{tm} = N_{cm} = 256$ . The  $L = 32$  point has doubled widths. All networks have 10 million parameters. Notice that for all Mixer Layers we used NTK initialization. We trained all cases on CIFAR-10 dataset using vanilla SGD paired with CSE. Batch size  $\mathrm{bs} = 256$ , weight decay  $\lambda = 10^{-4}$  was selected from  $\{10^{-5}, 10^{-4}\}$ , mixup rate  $\alpha = 0.8$  was selected from  $\{0.4, 0.8\}$ . We also used RandArgument and horizontal flip with default settings in PyTorch. For all cases we searched learning rates within  $\{0.005, 0.01, 0.05, 0.1, 0.2, 0.5\}$ . We also tried a linear warm-up schedule for first 3000 iterations, but we did not see any improvement in performance. Generating the data for the figure took approximately 4 days on Google Colab Pro (single Tesla P100 GPU).
B Additional Discussion on ResNet, ResNet with BatchNorm
For Convolution Layers, the NNGP kernel is a 4-index tensor:  $\mathcal{K}_{\mu \nu ;ij}^{l}(x,x^{\prime})$ , where the Greek letters  $(\mu ,\nu)$  index the channels, whereas the Latin letters  $(i,j)$  index the pixels. The infinite width limit in this case is achieved by taking the number of channels to infinity (sequentially). In this limit, most of our equations for MLP can be easily rewritten using the convolutional NNGP kernel. However, in this case, the kernel is only diagonal in channel dimension:  $\mathcal{K}_{\mu \nu ;ij}^{l}(x,x^{\prime}) = \mathcal{K}_{ij}^{l}(x,x^{\prime})\delta_{\mu \nu}$ . This additional structure in the kernel makes it difficult to get a closed-form solution for  $\mathcal{J}^{l,l + 1}$  in general.
ResNet110 (LayerNorm) In Figure 4(2), the networks is critical close to  $\mu = 1$ , as expected from our analysis. One would naively expect the  $\mu < 1$  cases also to be critical, since for MLP with ReLU and Pre-LN,  $\sigma_{b} = 0$  is critical regardless of  $\sigma_{w}$  and  $\mu$ . However, in Figure 4(2) the region away from  $\mu = 1$  is in ordered phase. This is likely a result of the kernel  $\mathcal{K}_{\mu \nu ;ij}^{l}(x,x^{\prime})$  not being diagonal in spatial dimensions. We emphasize that the  $\mu = 1$  case stays unaffected by this, since the existence of criticality does not depend on the details of the NNGP kernel in this case. This can be readily seen from (71). We present the Numerical and training results in Figure 4.
ResNet110 (BatchNorm) The operation of BatchNorm on a preactivation (pre-BN) in an MLP can be described as follows:
$$
\tilde {h} _ {i} (x) = \frac {h _ {i} (x) - \mu_ {i , B}}{\sigma_ {i , B}},
$$
$$
\mu_ {i, B} = \frac {1}{| B |} \sum_ {x ^ {\prime} \in B} h \left(x ^ {\prime}\right) \quad a n d \quad \sigma_ {i, B} = \sqrt {\frac {1}{| B |} \sum_ {x ^ {\prime} \in B} \left(h \left(x ^ {\prime}\right) - \mu_ {i , B}\right) ^ {2}}, \tag {23}
$$
where  $B$  is the batch that  $x$  belongs to and  $|B|$  is batch size.
The works Yang et al. [63], He et al. [23] show that for large batch size, the effect of BatchNorm for NNGP kernel and Jacobian Norm is deterministic. We summarize the results for the pre-BN MLP setup:
$$
\mathcal {K} ^ {l + 1} (x, x ^ {\prime}) = \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi \left(\tilde {h} _ {j} ^ {l} (x)\right) \phi \left(\tilde {h} _ {j} ^ {l} \left(x ^ {\prime}\right)\right) \right] + \mu^ {2} \mathcal {K} ^ {l} (x, x ^ {\prime}) \tag {24}
$$
$$
\mathcal {J} ^ {l, l + 1} = \frac {\sigma_ {w} ^ {2}}{K ^ {l} (x , x) - K ^ {l} \left(x , x ^ {\prime}\right)} \mathbb {E} _ {\theta} \left[ \phi \left(\tilde {h} _ {j} ^ {l} (x)\right) \phi \left(\tilde {h} _ {j} ^ {l} (x)\right) \right] + \mu^ {2}, \tag {25}
$$
where  $x'$  in the APJN term can be any  $x' \neq x$ , since for a large batch size all choices are equivalent.
From the above result, we can see that most results we have for LayerNorm can be translated to BatchNorm with an easy replacement  $K^l(x,x) \to (\mathcal{K}^l(x,x) - \mathcal{K}^l(x,x'))$ . As a simple example, consider a pre-BN ResNet architecture, but with all the Convolutional layers replaced with Linear (Fully Connected) layers. For such a network, we have the following result for  $\mu < 1$ :
$$
\mathcal {J} ^ {l, l + 1} = \frac {\pi^ {2} (1 - \mu^ {2})}{(\pi - 1) ^ {2}} + \mu^ {2}. \tag {26}
$$
For  $\mu = 1$ , we have
$$
\mathcal {J} ^ {l, l + 1} = 1 + O \left(l ^ {- 1}\right) \tag {27}
$$
The ResNet results can then be obtained by replacing Fully Connected layers with Convolution layers, in a similar way as discussed in ResNet(LayerNorm) section. We show the numerical results for ResNet with BatchNorm in Figure 6.
Figure 6: ResNet110(BatchNorm): Left to right: (1)(2)  $\mathcal{J}^{L - 1,L}$  phase diagrams, with  $(\sigma_w^2,\sigma_b^2)$  and  $(\sigma_w^2,\mu)$ . (3)(4) Training curves w.r.t  $\sigma_w^2$  and  $\mu$ .
C Discussions on Transformers
As preliminary empirical results, we applied our method to Vision Transformers (ViTs). Figure 7 supports our argument in the main text that the Pre-LN architecture is generally insensitive to initialization when  $\mu = 1$ , where Post-LN is only stable with small  $\sigma_w^2$ ; Figure 8 further demonstrates that the attention mechanism is highly sensitive to initialization and thus vulnerable to the gradient vanishing/diverging problem.