SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \chi_ {\mathcal {I}} ^ {j} = \sigma_ {w} ^ {2} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {i} ^ {l}\right) \phi^ {\prime} \left(h _ {i} ^ {l}\right) \right] \\ = \frac {\sigma_ {w} ^ {2}}{\sqrt {2 \pi \mathcal {K} ^ {l}}} \int d h _ {i} ^ {l} \left[ \frac {1}{2} + \frac {1}{2} \operatorname {e r f} \left(\frac {h _ {i} ^ {l}}{\sqrt {2}}\right) + \frac {e ^ {- \frac {\left(h _ {i} ^ {l}\right) ^ {2}}{2}} h _ {i} ^ {l}}{\sqrt {2 \pi}} \right] ^ {2} e ^ {- \frac {\left(h _ {i} ^ {l}\right) ^ {2}}{2 \mathcal {K} ^ {l}}} \\ = \frac {\sigma_ {w} ^ {2}}{\sqrt {2 \pi \mathcal {K} ^ {l}}} \int d h _ {i} ^ {l} \left[ \frac {1}{4} + \frac {1}{4} \operatorname {e r f} \left(\frac {h _ {i} ^ {l}}{\sqrt {2}}\right) ^ {2} + \frac {h _ {i} ^ {l} \operatorname {e r f} \left(\frac {h _ {i} ^ {l}}{\sqrt {2}}\right) e ^ {- \frac {(h _ {i} ^ {l}) ^ {2}}{2}}}{\sqrt {2 \pi}} + \frac {e ^ {- (h _ {i} ^ {l}) ^ {2}} (h _ {i} ^ {l}) ^ {2}}{2 \pi} \right] e ^ {- \frac {(h _ {i} ^ {l}) ^ {2}}{2 \mathcal {K} ^ {l}}} \\ = \frac {\sigma_ {w} ^ {2}}{4} + \frac {\sigma_ {w} ^ {2}}{2 \pi} \left[ \arcsin \left(\frac {\mathcal {K} ^ {l}}{1 + \mathcal {K} ^ {l}}\right) + \frac {\mathcal {K} ^ {l} \left(3 + 5 \mathcal {K} ^ {l}\right)}{\left(1 + \mathcal {K} ^ {l}\right) \left(1 + 2 \mathcal {K} ^ {l}\right) ^ {\frac {3}{2}}} \right], \tag {154} \\ \end{array}
$$
where we dropped odd function terms to get the third line, and to get the last line we used known result for erf in the second term, integration by parts in the third term.
Here to get the critical line is harder. One can use the recurrence relation for the NNGP kernel at fixed point  $\mathcal{K}^{\star}$  and  $\chi_{\mathcal{I}}^{\star} = 1$
$$
\mathcal {K} ^ {\star} = \frac {\sigma_ {w} ^ {2}}{4} \mathcal {K} ^ {\star} + \frac {\sigma_ {w} ^ {2}}{2 \pi} \left[ \arcsin \left(\frac {\mathcal {K} ^ {\star}}{1 + \mathcal {K} ^ {\star}}\right) + \frac {\sigma_ {w} ^ {2} \mathcal {K} ^ {\star}}{\pi (1 + \mathcal {K} ^ {\star}) \sqrt {1 + 2 \mathcal {K} ^ {\star}}} \right] \mathcal {K} ^ {\star} + \sigma_ {b} ^ {2}, \tag {155}
$$
$$
\chi_ {\mathcal {I}} ^ {\star} = \frac {\sigma_ {w} ^ {2}}{4} + \frac {\sigma_ {w} ^ {2}}{2 \pi} \left[ \arcsin \left(\frac {\mathcal {K} ^ {\star}}{1 + \mathcal {K} ^ {\star}}\right) + \frac {\mathcal {K} ^ {\star} \left(3 + 5 \mathcal {K} ^ {\star}\right)}{\left(1 + \mathcal {K} ^ {\star}\right) \left(1 + 2 \mathcal {K} ^ {\star}\right) ^ {\frac {3}{2}}} \right] = 1. \tag {156}
$$
Cancel the arcsin term,  $\sigma_w$  and  $\sigma_b$  then can be written as a function of  $\mathcal{K}^{\star}$
$$
\sigma_ {w} = 2 \left[ 1 + \frac {2 \mathcal {K} ^ {\star} (3 + 5 \mathcal {K} ^ {\star})}{\pi (1 + \mathcal {K} ^ {\star}) (1 + 2 \mathcal {K} ^ {\star}) ^ {\frac {3}{2}}} + \frac {2}{\pi} \arcsin \left(\frac {\mathcal {K} ^ {\star}}{1 + \mathcal {K} ^ {\star}}\right) \right] ^ {- \frac {1}{2}}, \tag {157}
$$
$$
\sigma_ {b} = \frac {\mathcal {K} ^ {\star}}{\sqrt {2 \pi} (1 + 2 \mathcal {K} ^ {\star}) ^ {\frac {3}{4}}} \sigma_ {w}. \tag {158}
$$
One can then scan  $\kappa^{\star}$  to draw the critical line.
In order to locate the critical point, we further require  $\chi_{\mathcal{K}}^{\star} = 1$ . To locate the critical point, we solve  $\chi_{\mathcal{J}}^{\star} - \chi_{\mathcal{K}}^{\star} = 0$  instead. We have
$$
\frac {\sigma_ {w} ^ {2} \left[ \left(\mathcal {K} ^ {\star}\right) ^ {3} - 3 \left(\mathcal {K} ^ {\star}\right) ^ {2} - 2 \mathcal {K} ^ {\star} \right]}{2 \pi \left(1 + \mathcal {K} ^ {\star}\right) ^ {2} \left(1 + 2 \mathcal {K} ^ {\star}\right) ^ {\frac {3}{2}}} = 0, \tag {159}
$$
which has two non-negative solutions out of three
$$
\mathcal {K} ^ {\star} = 0 \text {a n d} \mathcal {K} ^ {\star} = \frac {3 + \sqrt {1 7}}{2}. \tag {160}
$$
One can then solve  $\sigma_{b}$  and  $\sigma_w$  by plugging corresponding  $K^{\star}$  values.
$$
\left(\sigma_ {w}, \sigma_ {b}\right) = (2, 0), \text {f o r} \mathcal {K} ^ {\star} = 0, \tag {161}
$$
$$
\left(\sigma_ {w}, \sigma_ {b}\right) \approx (1. 4 0 8, 0. 4 1 6), \text {f o r} \mathcal {K} ^ {\star} = \frac {3 + \sqrt {1 7}}{2}. \tag {162}
$$
K.3 Critical Exponents
GELU behaves in a different way compared to erf. First we discuss the  $\mathcal{K}^{\star} = 0$  critical point, which is located at  $(\sigma_b,\sigma_w) = (0,2)$ . We expand Eq.(152), and keep next to leading order  $\delta \mathcal{K}^l = \mathcal{K}^l -\mathcal{K}^\star$
$$
\delta \mathcal {K} ^ {l + 1} \approx \delta \mathcal {K} ^ {l} + \frac {6}{\pi} \left(\delta \mathcal {K} ^ {l}\right) ^ {2}. \tag {163}
$$
$$
A = - \frac {\pi}{6} \text {a n d} \zeta_ {\mathcal {K}} = 1, \tag {164}
$$
which is not possible since  $\delta \mathcal{K}^l \geq 0$  for this case. This result means scaling analysis is not working here.
Next, we consider the other fixed point with  $\mathcal{K}^{\star} = \frac{3 + \sqrt{17}}{2}$  at  $(\sigma_b,\sigma_w) = (0.416,1.408)$ . Expand the NNGP kernel recurrence relation again.
$$
\delta \mathcal {K} ^ {l + 1} \approx \delta \mathcal {K} ^ {l} + 0. 0 0 0 1 4 \left(\delta \mathcal {K} ^ {l}\right) ^ {2}. \tag {165}
$$
Following the same analysis, we find
$$
\delta \mathcal {K} ^ {l} \approx - 7 1 4 2. 9 l ^ {- 1}. \tag {166}
$$
Looks like scaling analysis works for this case, since  $\mathcal{K}^{\star} > 0$ . The solution shows that the critical point is half-stable[45]. If  $\mathcal{K}^l < \mathcal{K}^\star$ , the fixed point is repealing, while when  $\mathcal{K}^l > \mathcal{K}^\star$ , the fixed point is attractive. However, the extremely large coefficient in the scaling behavior of  $\delta \mathcal{K}^l$  embarrasses the analysis. Since for any network with a reasonable depth, the deviation  $\delta \mathcal{K}^l$  is not small.
Now we can expand  $\chi_{\mathcal{I}}^l$  at some large depth, up to leading order  $l^{-1}$ .
$$
\chi_ {\mathcal {I}} ^ {l} \approx 1 - \frac {6 6 . 6 6 8}{l}. \tag {167}
$$
$$
\delta \mathcal {J} ^ {l _ {0}, l} \approx c _ {l _ {0}} l ^ {- 6 6. 6 6 8}, \tag {168}
$$
where  $c_{l_0}$  is a positive non-universal constant.
Critical exponent
$$
\zeta = 6 6. 6 6 8. \tag {169}
$$
Which in practice is not traceable.
K.4 LayerNorm on Pre-activations
$$
\begin{array}{l} \chi_ {\mathcal {I}} ^ {l} = \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] \Bigg | _ {\tilde {\mathcal {K}} ^ {l - 1} = 1} \\ = \frac {\sigma_ {w} ^ {2} (6 \pi + 4 \sqrt {3})}{\sigma_ {w} ^ {2} (6 \pi + 3 \sqrt {3}) + 1 8 \pi \sigma_ {b} ^ {2}}. \tag {170} \\ \end{array}
$$
The critical line is then at
$$
\begin{array}{l} \sigma_ {b} = \left(6 \sqrt {3} \pi\right) ^ {- \frac {1}{2}} \sigma_ {w} \\ \approx 0. 1 7 5 \sigma_ {w}. \tag {171} \\ \end{array}
$$
K.5 LayerNorm on Activations
First, we need to evaluate a new expectation value
$$
\begin{array}{l} \mathbb {E} _ {\theta} \left[ \phi (h _ {i} ^ {l}) \right] = \frac {1}{\sqrt {2 \pi (\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2})}} \int d h _ {i} ^ {l} \frac {h _ {i} ^ {l}}{2} \left[ 1 + \mathrm {e r f} \left(\frac {x}{\sqrt {2}}\right) \right] e ^ {- \frac {(h _ {i} ^ {l}) ^ {2}}{2 (\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2})}} \\ = \frac {\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}}{\sqrt {2 \pi \left(1 + \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right)}}, \tag {172} \\ \end{array}
$$
where we used integration by parts to get the result.
The other integrals are modified to
$$
\mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {i} ^ {l}\right) \phi^ {\prime} \left(h _ {i} ^ {l}\right) \right] = \frac {1}{4} + \frac {1}{2 \pi} \left[ \arcsin \left(\frac {\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}}{1 + \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}}\right) + \frac {\left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) \left[ 3 + 5 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) \right]}{\left(1 + \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) \left[ 1 + 2 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) \right] ^ {\frac {3}{2}}} \right], \tag {173}
$$
$$
\mathbb {E} _ {\theta} \left[ \phi (h _ {i} ^ {l}) \phi (h _ {i} ^ {l}) \right] = \frac {\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}}{4} + \frac {\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}}{2 \pi} \arcsin \left(\frac {\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}}{1 + \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}}\right) + \frac {\left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) ^ {2}}{\pi (1 + \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}) \sqrt {1 + 2 (\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2})}}.
$$
One can then combine those results to find  $\chi_{\mathcal{I}}^{l}$
$$
\chi_ {\mathcal {J}} ^ {l} = \frac {\sigma_ {w} ^ {2} \left(1 + \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) \left[ \pi + 2 \arcsin \left(\frac {\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}}{1 + \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}}\right) + \frac {2 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) \left(3 + 5 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right)\right)}{\left(1 + \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) \left(1 + 2 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right)\right) ^ {\frac {3}{2}}} \right]}{\pi \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) \left(1 + \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) - 2 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) ^ {2} + \frac {4 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) ^ {2}}{\sqrt {1 + 2 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right)}} + 2 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) \left(1 + \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) \arcsin \left(\frac {\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}}{1 + \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}}\right)}. (1 7 5)
$$
The critical line defined by  $\chi_{\mathcal{I}}^l = 1$ , one can numerically solve it by scanning over  $\sigma_b$  and  $\sigma_w$ .
K.6 Residual Connections
The recurrence relation for the NNGP kernel is
$$
\mathcal {K} ^ {l + 1} = \left[ \frac {\mathcal {K} ^ {l}}{4} + \frac {\mathcal {K} ^ {l}}{2 \pi} \arcsin \left(\frac {\mathcal {K} ^ {l}}{1 + \mathcal {K} ^ {l}}\right) + \frac {(\mathcal {K} ^ {l}) ^ {2}}{\pi (1 + \mathcal {K} ^ {l}) \sqrt {1 + 2 \mathcal {K} ^ {l}}} \right] \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2} + \mu^ {2} \mathcal {K} ^ {l}. \tag {176}
$$
Fixed point exists if
$$
\chi_ {\mathcal {K}} ^ {\star} = \frac {\sigma_ {w} ^ {2}}{4} + \frac {\sigma_ {w} ^ {2}}{2 \pi} \left[ \arcsin \left(\frac {\mathcal {K} ^ {\star}}{1 + \mathcal {K} ^ {\star}}\right) + \frac {4 (\mathcal {K} ^ {\star}) ^ {3} + 1 1 (\mathcal {K} ^ {\star}) ^ {2} + 5 \mathcal {K} ^ {\star}}{(1 + \mathcal {K} ^ {\star}) ^ {2} (1 + 2 \mathcal {K} ^ {\star}) ^ {\frac {3}{2}}} \right] + \mu^ {2} \leq 1. \tag {177}
$$
The recurrence coefficient for Jacobian is
$$
\chi_ {\mathcal {J}} ^ {\star} = \frac {\sigma_ {w} ^ {2}}{4} + \frac {\sigma_ {w} ^ {2}}{2 \pi} \left[ \arcsin \left(\frac {\mathcal {K} ^ {\star}}{1 + \mathcal {K} ^ {\star}}\right) + \frac {\mathcal {K} ^ {\star} \left(3 + 5 \mathcal {K} ^ {\star}\right)}{\left(1 + \mathcal {K} ^ {\star}\right) \left(1 + 2 \mathcal {K} ^ {\star}\right) ^ {\frac {3}{2}}} \right] + \mu^ {2}. \tag {178}
$$
Phase boundary is shifted
$$
\sigma_ {w} = 2 \sqrt {1 - \mu^ {2}} \left[ 1 + \frac {2 \mathcal {K} ^ {\star} (3 + 5 \mathcal {K} ^ {\star})}{\pi (1 + \mathcal {K} ^ {\star}) (1 + 2 \mathcal {K} ^ {\star}) ^ {\frac {3}{2}}} + \frac {2}{\pi} \arcsin \left(\frac {\mathcal {K} ^ {\star}}{1 + \mathcal {K} ^ {\star}}\right) \right] ^ {- \frac {1}{2}}, \tag {179}
$$
$$
\sigma_ {b} = \frac {\mathcal {K} ^ {\star}}{\sqrt {2 \pi} (1 + 2 \mathcal {K} ^ {\star}) ^ {\frac {3}{4}}} \sigma_ {w}. \tag {180}
$$
One can again scan over  $\kappa^{\star}$  to draw the critical line.
$$
\frac {\sigma_ {w} ^ {2} \left[ \left(\mathcal {K} ^ {\star}\right) ^ {3} - 3 \left(\mathcal {K} ^ {\star}\right) ^ {2} - 2 \mathcal {K} ^ {\star} \right]}{2 \pi \left(1 + \mathcal {K} ^ {\star}\right) ^ {2} \left(1 + 2 \mathcal {K} ^ {\star}\right) ^ {\frac {3}{2}}} = 0, \tag {181}
$$
$$
\mathcal {K} ^ {\star} = 0 \text {a n d} \mathcal {K} ^ {\star} = \frac {3 + \sqrt {1 7}}{2}. \tag {182}
$$
$$
\left(\sigma_ {w}, \sigma_ {b}\right) = \left(2 \sqrt {1 - \mu^ {2}}, 0\right), \text {f o r} \mathcal {K} ^ {\star} = 0, \tag {183}
$$
$$
\left(\sigma_ {w}, \sigma_ {b}\right) \approx \left(1. 4 0 8 \sqrt {1 - \mu^ {2}}, 0. 4 1 6 \sqrt {1 - \mu^ {2}}\right), \text {f o r} \mathcal {K} ^ {\star} = \frac {3 + \sqrt {1 7}}{2}. \tag {184}
$$
K.7 Residual Connections with LayerNorm on Preactivations (Pre-LN)
Use lemma D.9 and the results we had without residue connections for GELU.
$$
\begin{array}{l} \chi_ {\mathcal {I}} ^ {*} = \lim  _ {l \rightarrow \infty} \left(\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] \Bigg | _ {\bar {\mathcal {K}} ^ {l - 1} = 1} + \mu^ {2}\right) \\ = \frac {\sigma_ {w} ^ {2} (6 \pi + 4 \sqrt {3}) (1 - \mu^ {2})}{\sigma_ {w} ^ {2} (6 \pi + 3 \sqrt {3}) + 1 8 \pi \sigma_ {b} ^ {2}} + \mu^ {2} \\ = 1 - \frac {\left(\sqrt {3} \sigma_ {w} ^ {2} - 1 8 \pi \sigma_ {b} ^ {2}\right) \left(1 - \mu^ {2}\right)}{\sigma_ {w} ^ {2} \left(6 \pi + 3 \sqrt {3}\right) + 1 8 \pi \sigma_ {b} ^ {2}}. \tag {185} \\ \end{array}
$$
$$
\begin{array}{l} \sigma_ {b} = \left(6 \sqrt {3} \pi\right) ^ {- \frac {1}{2}} \sigma_ {w} \tag {186} \\ \approx 0. 1 7 5 \sigma_ {w}, \\ \end{array}
$$
just like without residue connections.
L Additional Experimental Results
In the following training results, we used NTK parameterization for the linear layers in the MLP. We emphasize that this choice has little effect on the training and convergence in this case, compared to standard initialization.
In figure 10, we showed empirically that the critical exponent of partial Jacobians are vanished for erf with LayerNorm.
Figure 10: log - log plot of partial Jacobian  $\mathcal{J}^{0,l}$  vs.  $l$  for (A) LN-erf and (B) erf-LN.
In figure 11, we tested  $6k$  samples from CIFAR-10 dataset[28] with kernel regression based on neural tangents library [40] [31] [41]. Test accuracy from kernel regression reflects the trainability (training accuracy) with SGD in the ordered phase. We found that the trainable depth is predicted by the correlation length  $c\xi$  with LayerNorm applied to preactivations, where the prefactor  $c = 28$ . The prefactor we had is the same as vanilla cases in [55]. The difference is from the fact that they used  $\log_{10}$  and we used  $\log_e$ .
In figure 12, we explore the broad range in  $\sigma_w^2$  of the performance of MLP network with erf activation function and LayerNorm on preativations. The network has depth  $L = 50$  and width  $N_l = 500$ ; and is trained using SGD on Fashion MNIST. The learning rates are chosen based on a logarithmic scan with a short training time.
Figure 11: Test accuracy for LayerNorm applied to preactivations.  $\sigma_b^2 = 0.5$  for all cases. Correlation lengths are calculated using analytical results of  $\chi_{\mathcal{I}}^l$ .
Figure 12: Training performance of MLP networks with erf activation function; and LayerNorm applied to preactivations. It continues to train for several orders of magnitude of  $\sigma_w^2$  (with learning-rate tuning).