SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 7:  $\chi_{\mathcal{I}}$  Phase diagram for Vision Transformer (ViT) ( $L = 24$ ,  $n_{\mathrm{heads}} = 8$ ,  $d_{\mathrm{embd}} = 256$ ). We present cases with pre/post-LN, with/without MLP layers. Top-left: Usual ViT - becomes everywhere-critical when  $\mu = 1$ ; ordered otherwise. Bottom-left: Usual ViT without MLP layers. Top-right: ViT with post-LN (instead of pre-LN). Always ordered for finite  $\sigma_w$ . Bottom-right: ViT with post-LN, without MLP layers.
Figure 8:  $\log \chi_{\mathcal{I}}$  Phase diagram for stacked Attention layers with/without MLP layers ( $L = 24$ ,  $n_{\mathrm{heads}} = 8$ ,  $d_{\mathrm{embd}} = 256$ ). We use  $\log \chi_{\mathcal{I}}$  and cut off the colorbar at  $\chi_{\mathcal{I}} = 10^1$  because the  $\chi_{\mathcal{I}}$  values in this case span from  $\sim 50$  orders of magnitude.
D Technical details for Jacobians and LayerNorm
We will drop the dependence of  $h_i^l (x)$  on  $x$  throughout the Appendices. It should not cause any confusion since we are always considering a single input.
D.1 NNGP Kernel
First, we derive the recurrence relation for the NNGP kernel Eq.(9). As mentioned in the main text, weights and biases are initialized (independently) from standard normal distribution  $\mathcal{N}(0,\sigma_w^2 /\mathrm{fan\_in})$ . We then have
$$
\mathbb {E} _ {\theta} \left[ w _ {i j} ^ {l} w _ {m n} ^ {l} \right] = \frac {\sigma_ {w} ^ {2}}{N _ {l - 1}} \delta_ {i m} \delta_ {j n} \text {a n d} \mathbb {E} _ {\theta} \left[ b _ {i} ^ {l} b _ {j} ^ {l} \right] = \sigma_ {b} ^ {2} \delta_ {i j} \tag {28}
$$
by definition.
We would like to prove theorem 2.4, as a consequence of lemma 2.2. The proof of lemma 2.2 can be found in [45].
Proof of theorem 2.4. One can prove this by definition with lemma 2.2.
$$
\begin{array}{l} \mathcal {K} ^ {l + 1} \equiv \frac {1}{N _ {l + 1}} \sum_ {i = 1} ^ {N _ {l + 1}} \mathbb {E} _ {\theta} \left[ h _ {i} ^ {l + 1} h _ {i} ^ {l + 1} \right] \\ = \frac {1}{N _ {l + 1}} \sum_ {i = 1} ^ {N _ {l + 1}} \mathbb {E} _ {\theta} \left[ \left(\sum_ {j = 1} ^ {N _ {l}} w _ {i j} ^ {l + 1} \phi \left(h _ {j} ^ {l}\right) + b _ {i} ^ {l + 1}\right) \left(\sum_ {k = 1} ^ {N _ {l}} w _ {i k} ^ {l + 1} \phi \left(h _ {k} ^ {l}\right) + b _ {i} ^ {l + 1}\right) \right] \\ = \frac {1}{N _ {l + 1}} \sum_ {i = 1} ^ {N _ {l + 1}} \mathbb {E} _ {\theta} \left[ \sum_ {j = 1} ^ {N _ {l}} \sum_ {k = 1} ^ {N _ {l}} w _ {i j} ^ {l + 1} w _ {i k} ^ {l + 1} \phi \left(h _ {j} ^ {l}\right) \phi \left(h _ {k} ^ {l}\right) + b _ {i} ^ {l + 1} b _ {i} ^ {l + 1} \right] \\ = \frac {1}{N _ {l + 1}} \sum_ {i = 1} ^ {N _ {l + 1}} \mathbb {E} _ {\theta} \left[ \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \phi \left(h _ {j} ^ {l}\right) \phi \left(h _ {j} ^ {l}\right) + \sigma_ {b} ^ {2} \right] \\ = \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi \left(h _ {j} ^ {l}\right) \phi \left(h _ {j} ^ {l}\right) \right] + \sigma_ {b} ^ {2}. \tag {29} \\ \end{array}
$$
D.2 Jacobians
Next, we prove theorem 2.5 in the main text.
Proof of theorem 2.5. We start from the definition of the averaged partial Jacobian norm (APJN)  $(l > l_0)$
$$
\begin{array}{l} \mathcal {J} ^ {l _ {0}, l + 1} \equiv \frac {1}{N _ {l + 1}} \mathbb {E} _ {\theta} \left[ \sum_ {i = 1} ^ {N _ {l + 1}} \sum_ {j = 1} ^ {N _ {l _ {0}}} \frac {\partial h _ {i} ^ {l + 1}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {i} ^ {l + 1}}{\partial h _ {j} ^ {l _ {0}}} \right] \\ = \frac {1}{N _ {l + 1}} \mathbb {E} _ {\theta} \left[ \sum_ {i = 1} ^ {N _ {l + 1}} \sum_ {j = 1} ^ {N _ {l _ {0}}} \left(\sum_ {k = 1} ^ {N _ {l}} \frac {\partial h _ {i} ^ {l + 1}}{\partial h _ {k} ^ {l}} \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}}\right) \left(\sum_ {m = 1} ^ {N _ {l}} \frac {\partial h _ {i} ^ {l + 1}}{\partial h _ {m} ^ {l}} \frac {\partial h _ {m} ^ {l}}{\partial h _ {j} ^ {l _ {0}}}\right) \right] \\ = \frac {1}{N _ {l + 1}} \mathbb {E} _ {\theta} \left[ \sum_ {i = 1} ^ {N _ {l + 1}} \sum_ {j = 1} ^ {N _ {l _ {0}}} \sum_ {k, m = 1} ^ {N _ {l}} \left(w _ {i k} ^ {l + 1} \phi^ {\prime} \left(h _ {k} ^ {l}\right)\right) \left(w _ {i m} ^ {l + 1} \phi^ {\prime} \left(h _ {m} ^ {l}\right)\right) \left(\frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {m} ^ {l}}{\partial h _ {j} ^ {l _ {0}}}\right) \right] \\ = \frac {1}{N _ {l + 1}} \mathbb {E} _ {\theta} \left[ \sum_ {i = 1} ^ {N _ {l + 1}} \sum_ {j = 1} ^ {N _ {l _ {0}}} \sum_ {k, m = 1} ^ {N _ {l}} w _ {i k} ^ {l + 1} w _ {i m} ^ {l + 1} \phi^ {\prime} \left(h _ {k} ^ {l}\right) \phi^ {\prime} \left(h _ {m} ^ {l}\right) \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {m} ^ {l}}{\partial h _ {j} ^ {l _ {0}}} \right] \\ = \frac {1}{N _ {l + 1}} \sum_ {i = 1} ^ {N _ {l + 1}} \sum_ {j = 1} ^ {N _ {l _ {0}}} \sum_ {k = 1} ^ {N _ {l}} \frac {\sigma_ {w} ^ {2}}{N _ {l}} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {k} ^ {l}\right) \phi^ {\prime} \left(h _ {k} ^ {l}\right) \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}} \right] \\ = \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {k = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {k} ^ {l}\right) \phi^ {\prime} \left(h _ {k} ^ {l}\right) \left(\sum_ {j = 1} ^ {N _ {l _ {0}}} \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}}\right) \right]. \tag {30} \\ \end{array}
$$
where we used the chain rule and took the expectation value over  $w^{l + 1}$ . Next, we take the chain rule again:
$$
\begin{array}{l} \mathcal {J} ^ {l _ {0}, l + 1} = \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {k = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {k} ^ {l}\right) \phi^ {\prime} \left(h _ {k} ^ {l}\right) \left(\sum_ {j = 1} ^ {N _ {l _ {0}}} \sum_ {m, n = 1} ^ {N _ {l - 1}} w _ {k m} ^ {l} w _ {k n} ^ {l} \phi^ {\prime} \left(h _ {m} ^ {l - 1}\right) \phi^ {\prime} \left(h _ {n} ^ {l - 1}\right) \frac {\partial h _ {m} ^ {l - 1}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {n} ^ {l - 1}}{\partial h _ {j} ^ {l _ {0}}}\right) \right] \\ = \frac {\sigma_ {w} ^ {4}}{N _ {l} N _ {l - 1}} \sum_ {k = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {k} ^ {l}\right) \phi^ {\prime} \left(h _ {k} ^ {l}\right) \left(\sum_ {j = 1} ^ {N _ {l _ {0}}} \sum_ {m = 1} ^ {N _ {l - 1}} \phi^ {\prime} \left(h _ {m} ^ {l - 1}\right) \phi^ {\prime} \left(h _ {m} ^ {l - 1}\right) \frac {\partial h _ {m} ^ {l - 1}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {m} ^ {l - 1}}{\partial h _ {j} ^ {l _ {0}}}\right) + O (1 / N _ {l - 1}) \right] \\ = \frac {\sigma_ {w} ^ {4}}{N _ {l} N _ {l - 1}} \sum_ {k = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {k} ^ {l}\right) \phi^ {\prime} \left(h _ {k} ^ {l}\right) \right] \mathbb {E} _ {\theta} \left[ \left(\sum_ {j = 1} ^ {N _ {l _ {0}}} \sum_ {m = 1} ^ {N _ {l - 1}} \phi^ {\prime} \left(h _ {m} ^ {l - 1}\right) \phi^ {\prime} \left(h _ {m} ^ {l - 1}\right) \frac {\partial h _ {m} ^ {l - 1}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {m} ^ {l - 1}}{\partial h _ {j} ^ {l _ {0}}}\right) \right] \\ = \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {k = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {k} ^ {l}\right) \phi^ {\prime} \left(h _ {k} ^ {l}\right) \right] \cdot \frac {\sigma_ {w} ^ {2}}{N _ {l - 1}} \mathbb {E} _ {\theta} \left[ \left(\sum_ {j = 1} ^ {N _ {l _ {0}}} \sum_ {m = 1} ^ {N _ {l - 1}} \phi^ {\prime} \left(h _ {m} ^ {l - 1}\right) \phi^ {\prime} \left(h _ {m} ^ {l - 1}\right) \frac {\partial h _ {m} ^ {l - 1}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {m} ^ {l - 1}}{\partial h _ {j} ^ {l _ {0}}}\right) \right] \\ = \chi_ {\mathcal {I}} ^ {l} \mathcal {J} ^ {l _ {0}, l}, \tag {31} \\ \end{array}
$$
where we integrate (by parts) over  $w^l$  to get the second line. We take  $N_{l-1} \to \infty$  and then  $N_l \to \infty$  to get the third line. We rearrange terms and use the Eq.(11) to get the fourth and fifth lines. Notice that to get the third line we used the fact in the infinite width limit, the distribution of  $h_i^l$  is independent of  $h_i^{l-1}$ . Thus we proved
$$
\mathcal {J} ^ {l _ {0}, l + 1} = \chi_ {\mathcal {I}} ^ {l} \mathcal {J} ^ {l _ {0}, l}. \tag {32}
$$
The critical line is defined by requiring  $\chi_{\mathcal{J}}^{\star} = 1$ , where critical points are reached by further requiring  $\chi_{\mathcal{K}}^{\star} = 1$ .
As we mentioned in the main text,  $l_{0} = 0$  is subtle since the input dimension is fixed  $N_{0}$ , which can not be assumed to be infinity. Even though for a dataset like MNIST, usually  $N_{0}$  is not significantly smaller than width  $N_{l}$ . We show how to take finite  $O(N_{0}^{-1})$  correction into account by using one example.
Lemma D.1. Consider a one hidden layer network with a finite input dimension  $N_0$ . In the infinite width limit, the APJN is still deterministic and the first step of the recurrence relation is modified to:
$$
\mathcal {J} ^ {0, 2} = \left(\chi_ {\mathcal {I}} ^ {1} + \frac {2 \sigma_ {w} ^ {2}}{N _ {0}} \chi_ {\Delta} ^ {1} \sum_ {k} ^ {N _ {0}} \frac {1}{N _ {0}} h _ {k} ^ {0} h _ {k} ^ {0}\right) \mathcal {J} ^ {0, 1}, \tag {33}
$$
where  $\mathcal{J}^{0,1} = \sigma_w^2$
$$
\begin{array}{l} \mathcal {J} ^ {0, 2} = \frac {1}{N _ {2}} \mathbb {E} _ {\theta} \left[ \sum_ {i = 1} ^ {N _ {2}} \sum_ {j = 1} ^ {N _ {0}} \frac {\partial h _ {i} ^ {2}}{\partial h _ {j} ^ {0}} \frac {\partial h _ {i} ^ {2}}{\partial h _ {j} ^ {0}} \right] \\ = \frac {1}{N _ {2}} \mathbb {E} _ {\theta} \left[ \sum_ {i = 1} ^ {N _ {2}} \sum_ {j = 1} ^ {N _ {0}} \sum_ {k, m = 1} ^ {N _ {1}} w _ {i k} ^ {2} w _ {i m} ^ {2} \phi^ {\prime} \left(h _ {k} ^ {1}\right) \phi^ {\prime} \left(h _ {m} ^ {1}\right) \frac {\partial h _ {k} ^ {1}}{\partial h _ {j} ^ {0}} \frac {\partial h _ {m} ^ {1}}{\partial h _ {j} ^ {0}} \right] \\ = \frac {1}{N _ {2}} \sum_ {i = 1} ^ {N _ {2}} \sum_ {j = 1} ^ {N _ {0}} \sum_ {k, m = 1} ^ {N _ {1}} \mathbb {E} _ {\theta} \left[ w _ {i k} ^ {2} w _ {i m} ^ {2} w _ {k j} ^ {1} w _ {m j} ^ {1} \phi^ {\prime} \left(h _ {k} ^ {1}\right) \phi^ {\prime} \left(h _ {m} ^ {1}\right) \right] \\ = \sum_ {j = 1} ^ {N _ {0}} \sum_ {k = 1} ^ {N _ {1}} \frac {\sigma_ {w} ^ {2}}{N _ {1}} \mathbb {E} _ {\theta} \left[ w _ {k j} ^ {1} w _ {k j} ^ {1} \phi^ {\prime} \left(h _ {k} ^ {1}\right) \phi^ {\prime} \left(h _ {k} ^ {1}\right) \right] \\ = \sigma_ {w} ^ {2} \left(\chi_ {\mathcal {I}} ^ {1} + \frac {2 \sigma_ {w} ^ {2}}{N _ {0}} \chi_ {\Delta} ^ {1} \sum_ {k} ^ {N _ {0}} \frac {1}{N _ {0}} h _ {k} ^ {0} h _ {k} ^ {0}\right) \\ = \left(\chi_ {\mathcal {I}} ^ {1} + \frac {2 \sigma_ {w} ^ {2}}{N _ {0}} \chi_ {\Delta} ^ {1} \sum_ {k} ^ {N _ {0}} \frac {1}{N _ {0}} h _ {k} ^ {0} h _ {k} ^ {0}\right) \mathcal {J} ^ {0, 1}, \tag {34} \\ \end{array}
$$
where to get the result we used integrate by parts, then explicitly integrated over  $w_{ij}^{1}$ . We have introduced a coefficient of finite width corrections,  $\chi_{\Delta}^{l}$ , defined as follows.
Definition D.2 (Coefficient of Finite Width Corrections).
$$
\chi_ {\Delta} ^ {l} = \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {i = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi^ {\prime \prime} \left(h _ {i} ^ {l}\right) \phi^ {\prime \prime} \left(h _ {i} ^ {l}\right) + \phi^ {\prime \prime \prime} \left(h _ {i} ^ {l}\right) \phi^ {\prime} \left(h _ {i} ^ {l}\right) \right]. \tag {35}
$$
Remark D.3. Notice that the correction to  $\mathcal{J}^{0,2}$  is order  $O(N_0^{-1})$ . If one calculates the recurrence relation for deeper layers, the correction to  $\mathcal{J}^{0,l}$  will be  $O(\sum_{l'=0}^{l} N_{l'}^{-1})$ , which means the contribution from hidden layers can be ignored in infinite width limit.
The  $\mathcal{J}^{0,2}$  example justifies the factorization of the integral when we go from the last line of Eq.(30) to Eq.(32).
Finally, the full Jacobian in infinite width limit can be written as
Lemma D.4 (APJN with  $l_0 = 0$ ). The APJN (with  $l_0 = 0$ ) of a given network can be written as
$$
\mathcal {J} ^ {0, l} = \sigma_ {w} ^ {2} \left(\chi_ {\mathcal {J}} ^ {1} + \frac {2 \sigma_ {w} ^ {2}}{N _ {0}} \chi_ {\Delta} ^ {1} \sum_ {k} ^ {N _ {0}} \frac {1}{N _ {0}} h _ {k} ^ {0} h _ {k} ^ {0}\right) \prod_ {l ^ {\prime} = 2} ^ {l - 1} \chi_ {\mathcal {J}} ^ {l ^ {\prime}}. \tag {36}
$$
Note that APJN with  $l_0 > 0$  does not receive the  $O(N_0^{-1})$  correction.
D.3 APJN and gradients
As mentioned in the main text, APJN is an important tool in studying the exploding and vanishing gradients problem. Its utility stems from the fact that it is a dominant factor in the norm of the gradients. This can be readily by looking at the (squared)  $L_{2}$  norm of the gradient of any flattened
parameter matrix  $\theta^l$ , at initialization. In the infinite width limit, one gets
$$
\begin{array}{l} \left\| \nabla_ {\theta^ {l}} \mathcal {L} \right\| _ {2} ^ {2} = \left(\sum_ {a l l} \frac {\partial \mathcal {L}}{\partial h _ {i} ^ {L}} \frac {\partial h _ {i} ^ {L}}{\partial h _ {j} ^ {L - 1}} \dots \frac {\partial h _ {k} ^ {l + 1}}{\partial h _ {m} ^ {l}} \frac {\partial h _ {m} ^ {l}}{\partial \theta_ {n} ^ {l}}\right) ^ {2} \\ = \sum_ {a l l} \left(\frac {\partial \mathcal {L}}{\partial h _ {i} ^ {L}} \frac {\partial \mathcal {L}}{\partial h _ {i ^ {\prime}} ^ {L}}\right) \left(\frac {\partial h _ {i} ^ {L}}{\partial h _ {j} ^ {L - 1}} \frac {\partial h _ {i ^ {\prime}} ^ {L}}{\partial h _ {j ^ {\prime}} ^ {L - 1}}\right) \dots \left(\frac {\partial h _ {k} ^ {l + 1}}{\partial h _ {m} ^ {l}} \frac {\partial h _ {k ^ {\prime}} ^ {l + 1}}{\partial h _ {m ^ {\prime}} ^ {l}}\right) \left(\frac {\partial h _ {m} ^ {l}}{\partial \theta_ {n} ^ {l}} \frac {\partial h _ {m ^ {\prime}} ^ {l}}{\partial \theta_ {n} ^ {l}}\right) \\ = \sum_ {a l l} \left(\frac {\partial \mathcal {L}}{\partial h _ {i} ^ {L}} \frac {\partial \mathcal {L}}{\partial h _ {i ^ {\prime}} ^ {L}}\right) \left(\frac {\partial h _ {i} ^ {L}}{\partial h _ {j} ^ {L - 1}} \frac {\partial h _ {i ^ {\prime}} ^ {L}}{\partial h _ {j ^ {\prime}} ^ {L - 1}}\right) \dots \left(\frac {\partial h _ {k} ^ {l + 1}}{\partial h _ {m} ^ {l}} \frac {\partial h _ {k ^ {\prime}} ^ {l + 1}}{\partial h _ {m ^ {\prime}} ^ {l}}\right) \delta_ {m m ^ {\prime}} \left\| \frac {\partial h ^ {l}}{\partial \theta^ {l}} \right\| _ {F} ^ {2} \\ = \sum_ {a l l} \left(\frac {\partial \mathcal {L}}{\partial h _ {i} ^ {L}} \frac {\partial \mathcal {L}}{\partial h _ {i ^ {\prime}} ^ {L}}\right) \left(\frac {\partial h _ {i} ^ {L}}{\partial h _ {j} ^ {L - 1}} \frac {\partial h _ {i ^ {\prime}} ^ {L}}{\partial h _ {j ^ {\prime}} ^ {L - 1}}\right) \dots \delta_ {k k ^ {\prime}} \mathcal {J} ^ {l, l + 1} \left\| \frac {\partial h ^ {l}}{\partial \theta^ {l}} \right\| _ {F} ^ {2} \\ = \sum_ {i, i ^ {\prime}, j, j ^ {\prime}} \left(\frac {\partial \mathcal {L}}{\partial h _ {i} ^ {L}} \frac {\partial \mathcal {L}}{\partial h _ {i ^ {\prime}} ^ {L}}\right) \left(\frac {\partial h _ {i} ^ {L}}{\partial h _ {j} ^ {L - 1}} \frac {\partial h _ {i ^ {\prime}} ^ {L}}{\partial h _ {j ^ {\prime}} ^ {L - 1}}\right) \delta_ {j j ^ {\prime}} \dots \mathcal {J} ^ {l, l + 1} \left\| \frac {\partial h ^ {l}}{\partial \theta^ {l}} \right\| _ {F} ^ {2} \\ = \sum_ {i, i ^ {\prime}} \left(\frac {\partial \mathcal {L}}{\partial h _ {i} ^ {L}} \frac {\partial \mathcal {L}}{\partial h _ {i ^ {\prime}} ^ {L}}\right) \delta_ {i i ^ {\prime}} \mathcal {J} ^ {L - 1, L} \dots \mathcal {J} ^ {l, l + 1} \left\| \frac {\partial h ^ {l}}{\partial \theta^ {l}} \right\| _ {F} ^ {2} \\ = \left\| \frac {\partial \mathcal {L}}{\partial h ^ {L}} \right\| _ {2} ^ {2} \mathcal {J} ^ {L - 1, L} \dots \mathcal {J} ^ {l, l + 1} \left\| \frac {\partial h ^ {l}}{\partial \theta^ {l}} \right\| _ {F} ^ {2} \\ = \left\| \frac {\partial \mathcal {L}}{\partial h ^ {L}} \right\| _ {2} ^ {2} \mathcal {J} ^ {l, L} \left\| \frac {\partial h ^ {l}}{\partial \theta^ {l}} \right\| _ {F} ^ {2}, \tag {37} \\ \end{array}
$$
where  $\| \cdot \| _2$  denotes the  $L_{2}$  norm and  $\| \cdot \| _F$  denotes the Frobenius norm.
D.4 LayerNorm on Pre-activations
Definition D.5 (Layer Normalization).
$$
\tilde {h} _ {i} ^ {l} = \frac {h _ {i} ^ {l} - \mathbb {E} \left[ h ^ {l} \right]}{\sqrt {\mathbb {E} \left[ \left(h ^ {l}\right) ^ {2} \right] - \mathbb {E} \left[ h ^ {l} \right] ^ {2}}} \gamma_ {i} ^ {l} + \beta_ {i} ^ {l}, \tag {38}
$$
where  $\gamma_i^l$  and  $\beta_i^l$  are learnable parameters.
Remark D.6. With only LayerNorm, the (1) is simplified to
$$
h _ {i} ^ {l + 1} = \sum_ {j = 1} ^ {N _ {l}} w _ {i j} ^ {l + 1} \phi \left(\tilde {h} _ {j} ^ {l}\right) + b _ {i} ^ {l + 1}. \tag {39}
$$
Remark D.7. In the limit of infinite width, using the law of large numbers, the average over neurons  $\mathbb{E}[\dots]$  can be replaced by the average of parameter-initializations  $\mathbb{E}_{\theta}[\dots]$ . Additionally, in this limit, the preactivations are i.i.d. Gaussian distributed:  $h^l \sim \mathcal{N}(0,\mathcal{K}^l)$ .
$$
\mathbb {E} \left[ h ^ {l} \right] = \mathbb {E} _ {\theta} \left[ h ^ {l} \right] = 0, \tag {40}
$$
$$
\mathbb {E} \left[ \left(h ^ {l}\right) ^ {2} \right] = \mathbb {E} _ {\theta} \left[ \left(h ^ {l}\right) ^ {2} \right] = \mathcal {K} ^ {l}. \tag {41}
$$
The normalized preactivation then simplifies to the form of Eq.(19).
Remark D.8. At initialization, the parameters  $\gamma_i^l$  and  $\beta_i^l$  take the values 1 and 0, respectively. This leads to the form in equation (19). In infinite width limit, it has the following form
$$
\tilde {h} _ {i} ^ {l} = \frac {h _ {i} ^ {l} - \mathbb {E} _ {\theta} [ h ^ {l} ]}{\sqrt {\mathbb {E} _ {\theta} \left[ (h ^ {l}) ^ {2} \right] - \mathbb {E} _ {\theta} [ h ^ {l} ] ^ {2}}}. \tag {42}
$$
Lemma D.9. With LayerNorm on preactivations, the gaussian average is modified to
$$
\mathbb {E} _ {\theta} \left[ O \left(\tilde {h} _ {i} ^ {l}\right) \right] = \frac {1}{\sqrt {2 \pi}} \int d \tilde {h} _ {i} ^ {l} O \left(\tilde {h} _ {i} ^ {l}\right) e ^ {- \frac {\left(\tilde {h} _ {i} ^ {l}\right) ^ {2}}{2}}. \tag {43}
$$
Proof. By definition  $\tilde{h}_i^l$  is sampled from a standard normal distribution  $\mathcal{N}(0,1)$ , then use lemma 2.2 to get the final form.
Theorem D.10. In the infinite width limit the recurrence relation for the NNGP kernel with Layer-Norm on preactivations is
$$
\mathcal {K} ^ {l + 1} = \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi \left(\tilde {h} _ {j} ^ {l}\right) \phi \left(\tilde {h} _ {j} ^ {l}\right) \right] + \sigma_ {b} ^ {2}. \tag {44}
$$
$$
\begin{array}{l} \mathcal {K} ^ {l + 1} = \frac {1}{N _ {l + 1}} \sum_ {i = 1} ^ {N _ {l + 1}} \mathbb {E} _ {\theta} \left[ h _ {i} ^ {l + 1} h _ {i} ^ {l + 1} \right] \\ = \frac {1}{N _ {l + 1}} \sum_ {i = 1} ^ {N _ {l + 1}} \mathbb {E} _ {\theta} \left[ \left(\sum_ {j = 1} ^ {N _ {l}} w _ {i j} ^ {l + 1} \phi (\tilde {h} _ {j} ^ {l}) + b _ {i} ^ {l + 1}\right) \left(\sum_ {k = 1} ^ {N _ {l}} w _ {i k} ^ {l + 1} \phi (\tilde {h} _ {k} ^ {l}) + b _ {i} ^ {l + 1}\right) \right] \\ = \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi \left(\tilde {h} _ {j} ^ {l}\right) \phi \left(\tilde {h} _ {j} ^ {l}\right) \right] + \sigma_ {b} ^ {2}. \tag {45} \\ \end{array}
$$
Theorem D.11. In the infinite width limit the recurrence relation for partial Jacobian with LayerNorm on preactivations is
$$
\mathcal {J} ^ {l _ {0}, l + 1} = \chi_ {\mathcal {J}} ^ {l} \mathcal {J} ^ {l _ {0}, l}, \tag {46}
$$
where  $\chi_J^l = \frac{\sigma_w^2}{N_l\mathcal{K}^l}\sum_{i = 1}^{N_l}\mathbb{E}_\theta \left[\phi '(\tilde{h}_i^l)^2\right].$
$$
\begin{array}{l} \mathcal {J} ^ {l _ {0}, l + 1} = \frac {1}{N _ {l + 1}} \mathbb {E} _ {\theta} \left[ \sum_ {i = 1} ^ {N _ {l + 1}} \sum_ {j = 1} ^ {N _ {l _ {0}}} \frac {\partial h _ {i} ^ {l + 1}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {i} ^ {l + 1}}{\partial h _ {j} ^ {l _ {0}}} \right] \\ = \frac {1}{N _ {l + 1}} \mathbb {E} _ {\theta} \left[ \sum_ {i = 1} ^ {N _ {l + 1}} \sum_ {j = 1} ^ {N _ {l _ {0}}} \left(\sum_ {k = 1} ^ {N _ {l}} \frac {\partial h _ {i} ^ {l + 1}}{\partial \tilde {h} _ {k} ^ {l}} \frac {\partial \tilde {h} _ {k} ^ {l}}{\partial h _ {k} ^ {l}} \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}}\right) \left(\sum_ {m = 1} ^ {N _ {l}} \frac {\partial h _ {i} ^ {l + 1}}{\partial \tilde {h} _ {m} ^ {l}} \frac {\partial \tilde {h} _ {m} ^ {l}}{\partial h _ {m} ^ {l}} \frac {\partial h _ {m} ^ {l}}{\partial h _ {j} ^ {l _ {0}}}\right) \right] \\ = \frac {1}{N _ {l + 1}} \mathbb {E} _ {\theta} \left[ \sum_ {i = 1} ^ {N _ {l + 1}} \sum_ {j = 1} ^ {N _ {l _ {0}}} \sum_ {k, m = 1} ^ {N _ {l}} \left(w _ {i k} ^ {l + 1} \phi^ {\prime} (\tilde {h} _ {k} ^ {l}) \frac {1}{\sqrt {\mathcal {K} ^ {l}}}\right) \left(w _ {i m} ^ {l + 1} \phi^ {\prime} (\tilde {h} _ {m} ^ {l}) \frac {1}{\sqrt {\mathcal {K} ^ {l}}}\right) \left(\frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {m} ^ {l}}{\partial h _ {j} ^ {l _ {0}}}\right) \right] \\ = \frac {\sigma_ {w} ^ {2}}{N _ {l} \mathcal {K} ^ {l}} \sum_ {k = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} (\tilde {h} _ {k} ^ {l}) \phi^ {\prime} (\tilde {h} _ {k} ^ {l}) \left(\sum_ {j = 1} ^ {N _ {l _ {0}}} \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}}\right) \right] \\ = \frac {\sigma_ {w} ^ {2}}{N _ {l} \mathcal {K} ^ {l}} \sum_ {k = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} (\tilde {h} _ {k} ^ {l}) \phi^ {\prime} (\tilde {h} _ {k} ^ {l}) \right] \mathcal {J} ^ {l _ {0}, l} \\ = \chi_ {\mathcal {J}} ^ {l} \mathcal {J} ^ {l _ {0}, l}, \tag {47} \\ \end{array}
$$
D.5 LayerNorm on activations
The general definition of LayerNorm on activations is given as follows.
Definition D.12 (LayerNorm on Activations).
$$
\widetilde {\phi \left(h _ {i} ^ {l}\right)} = \frac {\phi \left(h _ {i} ^ {l}\right) - \mathbb {E} \left[ \phi \left(h ^ {l}\right) \right]}{\sqrt {\mathbb {E} \left[ \phi \left(h ^ {l}\right) ^ {2} \right] - \mathbb {E} \left[ \phi \left(h ^ {l}\right) \right] ^ {2}}} \gamma_ {i} ^ {l} + \beta_ {i} ^ {l}. \tag {48}
$$
Remark D.13. The recurrence relation for preactivations (Eq.(1)) gets modified to
$$
h _ {i} ^ {l + 1} = \sum_ {j = 1} ^ {N _ {l}} w _ {i j} ^ {l + 1} \widetilde {\phi \left(h _ {j} ^ {l}\right)} + b _ {i} ^ {l + 1}. \tag {49}
$$
Remark D.14. At initialization, the parameters  $\gamma_i^l$  and  $\beta_i^l$  take the values 1 and 0, respectively. This leads to the form
$$
\begin{array}{l} \widetilde {\phi \left(h _ {i} ^ {l}\right)} = \frac {\phi \left(h _ {i} ^ {l}\right) - \mathbb {E} \left[ \phi \left(h ^ {l}\right) \right]}{\sqrt {\mathbb {E} \left[ \phi \left(h ^ {l}\right) ^ {2} \right] - \mathbb {E} \left[ \phi \left(h ^ {l}\right) \right] ^ {2}}} \\ = \frac {\phi \left(h _ {i} ^ {l}\right) - \mathbb {E} _ {\theta} \left[ \phi \left(h ^ {l}\right) \right]}{\sqrt {\mathbb {E} _ {\theta} \left[ \phi \left(h ^ {l}\right) ^ {2} \right] - \mathbb {E} _ {\theta} \left[ \phi \left(h ^ {l}\right) \right] ^ {2}}}, \tag {50} \\ \end{array}
$$
where the first line follows from the fact that at initialization, the parameters  $\gamma_i^l$  and  $\beta_i^l$  take the values 1 and 0 respectively. In the second line, we have invoked the infinite width limit.
Remark D.15. Evaluating the Gaussian average in this case is similar to the cases in the previous section. The only difference is that the averages are taking over the distribution  $h^{l - 1} \sim \mathcal{N}(0, \mathcal{K}^{l - 1} = \sigma_w^2 + \sigma_b^2)$ . Again this can be summarized as
$$
\mathbb {E} _ {\theta} \left[ O \left(h _ {i} ^ {l}\right) \right] = \frac {1}{\sqrt {2 \pi \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right)}} \int d h _ {i} ^ {l} O \left(h _ {i} ^ {l}\right) e ^ {- \frac {\left(h _ {i} ^ {l}\right) ^ {2}}{2 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right)}}. \tag {51}
$$
Next, we calculate the modifications to the recurrence relations for the NNGP kernel and Jacobians.