SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- Second MlpBlock. It mixes channels  $\mu, \nu$ , preactivations from different patches share the same weight and bias.
-  $6l + 4$ : Linear Affine Layer.
$$
h _ {\nu i} ^ {6 l + 4} = \sum_ {\rho = 1} ^ {C} w _ {\nu \rho} ^ {6 l + 4} \tilde {h} _ {\rho i} ^ {6 l + 3} + b _ {\nu} ^ {6 l + 4}. \tag {87}
$$
-  $6l + 5$ . Affine Layer.
$$
h _ {\mu i} ^ {6 l + 5} = \sum_ {\nu = 1} ^ {N _ {c m}} w _ {\mu \nu} ^ {6 l + 5} \phi \left(h _ {\nu i} ^ {6 l + 4}\right) + b _ {\mu} ^ {6 l + 5}. \tag {88}
$$
-  $6l + 6$ . Residual Connections.
$$
h _ {\mu i} ^ {6 l + 6} = h _ {\mu i} ^ {6 l + 5} + \mu h _ {\mu i} ^ {6 l + 3}. \tag {89}
$$
Suppose the network has  $L$  Mixer layers. After those layers the network has a pre-head LayerNorm layer, a global average pooling layer, and an output layer. The pre-head LayerNorm normalizes over channels  $\mu$  can be described as the following
$$
\tilde {h} _ {\mu i} ^ {6 L} = \frac {h _ {\mu i} ^ {6 L} - \mathbb {E} _ {C} \left[ h _ {\rho i} ^ {6 L} \right]}{\sqrt {\operatorname {V a r} _ {C} \left[ h _ {\rho i} ^ {6 L} \right]}}. \tag {90}
$$
Global Average Pool over patches  $i$ .
$$
h _ {\mu} ^ {p} = \frac {1}{N _ {p}} \sum_ {i = 1} ^ {N _ {p}} \tilde {h} _ {\mu i} ^ {6 L}. \tag {91}
$$
Output Layer
$$
f _ {\mu} = \sum_ {\nu = 1} ^ {C} w _ {\mu \nu} h _ {\nu} ^ {p} + b _ {\mu}. \tag {92}
$$
We plotted the phase diagram using the following quantity from repeating Mixer Layers:
$$
\chi_ {\mathcal {I}} ^ {\star} = \lim  _ {L \rightarrow \infty} \left(\frac {1}{N _ {p} C} \sum_ {i = 1} ^ {N _ {p}} \sum_ {\mu = 1} ^ {C} \mathbb {E} _ {\theta} \left[ \sum_ {\rho = 1} ^ {C} \sum_ {k = 1} ^ {N _ {p}} \frac {\partial h _ {\mu i} ^ {6 L}}{\partial h _ {\rho k} ^ {6 L - 6}} \frac {\partial h _ {\mu i} ^ {6 L}}{\partial h _ {\rho k} ^ {6 L - 6}} \right]\right). \tag {93}
$$
I Results for Scale Invariant Activation Functions
Definition I.1 (Scale invariant activation functions).
$$
\phi (x) = a _ {+} x \Theta (x) + a _ {-} x \Theta (- x), \tag {94}
$$
where  $\Theta (x)$  is the Heaviside step function. ReLU is the special case with  $a_{+} = 1$  and  $a_{-} = 0$ .
I.1 NNGP Kernel
First evaluate the average using lemma 2.2
$$
\begin{array}{l} \mathbb {E} _ {\theta} \left[ \phi (h _ {i} ^ {l}) \phi (h _ {i} ^ {l}) \right] = \frac {1}{\sqrt {2 \pi \mathcal {K} ^ {l}}} \int d h _ {i} ^ {l} \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right) \left(h _ {i} ^ {l}\right) ^ {2} e ^ {- \frac {(h _ {i} ^ {l}) ^ {2}}{2 K ^ {l}}} \\ = \frac {a _ {+} ^ {2} + a _ {-} ^ {2}}{2} \mathcal {K} ^ {l}. \tag {95} \\ \end{array}
$$
Thus we obtain the recurrence relation for the NNGP kernel with scale invariant activation function.
$$
\mathcal {K} ^ {l + 1} = \frac {\sigma_ {w} ^ {2} \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right)}{2} \mathcal {K} ^ {l} + \sigma_ {b} ^ {2}. \tag {96}
$$
Finite fixed point of the recurrence relation above exists only if
$$
\chi_ {\mathcal {K}} ^ {\star} = \frac {\sigma_ {w} ^ {2} \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right)}{2} \leq 1. \tag {97}
$$
$$
\sigma_ {w} ^ {2} \leq \frac {2}{a _ {+} ^ {2} + a _ {-} ^ {2}}. \tag {98}
$$
For  $\sigma_w^2 = \frac{2}{a_+^2 + a_-^2}$  case, finite fixed point exists only if  $\sigma_b^2 = 0$ .
I.2 Jacobian(s)
The calculation is quite straightforward, by definition
$$
\begin{array}{l} \chi_ {\mathcal {I}} ^ {l} = \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[ a _ {+} \Theta (h _ {i} ^ {l}) - a _ {-} \Theta (h _ {i} ^ {l}) \right] ^ {2} e ^ {- \frac {(h _ {i} ^ {l}) ^ {2}}{2 \mathcal {K} ^ {l}}} \\ = \frac {\sigma_ {w} ^ {2} \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right)}{2}, \tag {99} \\ \end{array}
$$
where we used the property  $x\delta(x) = 0$  for Dirac's delta function to get the first line.
Thus the critical line is defined by
$$
\sigma_ {w} = \sqrt {\frac {2}{a _ {+} ^ {2} + a _ {-} ^ {2}}}. \tag {100}
$$
For ReLU with  $a_{+} = 1$  and  $a_{-} = 0$ , the network is at critical line when
$$
\sigma_ {w} = \sqrt {2}, \tag {101}
$$
where the critical point is located at
$$
\left(\sigma_ {w}, \sigma_ {b}\right) = (\sqrt {2}, 0). \tag {102}
$$
I.3 Critical Exponents
Since the recurrence relations for the NNGP kernel and Jacobians are linear. Then from lemma G.1 and theorem 2.7
$$
\zeta_ {\mathcal {K}} = 0 \text {a n d} \zeta = 0. \tag {103}
$$
I.4 LayerNorm on Pre-activations
Use lemma D.9 and combine all known results for scale-invariant functions
$$
\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} \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right)}{\sigma_ {w} ^ {2} \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right) + 2 \sigma_ {b} ^ {2}}. \tag {104} \\ \end{array}
$$
In this case,
$$
\chi_ {\mathcal {I}} ^ {l} \leq 1 \tag {105}
$$
is always true. The equality only holds at  $\sigma_{b} = 0$  line.
I.5 LayerNorm on Activations
First we substitute  $\mathcal{K}^{l - 1} = \sigma_w^2 +\sigma_b^2$  into known results
$$
\mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {i} ^ {l}\right) \phi^ {\prime} \left(h _ {i} ^ {l}\right) \right] = \frac {a _ {+} ^ {2} + a _ {-} ^ {2}}{2}, \tag {106}
$$
$$
\mathbb {E} _ {\theta} \left[ \phi \left(h _ {i} ^ {l}\right) \phi \left(h _ {i} ^ {l}\right) \right] = \frac {a _ {+} ^ {2} + a _ {-} ^ {2}}{2} \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right). \tag {107}
$$
There is a new expectation value we need to show explicitly
$$
\begin{array}{l} \mathbb {E} _ {\theta} \left[ \phi (h _ {i} ^ {l}) \right] = \frac {1}{\sqrt {2 \pi (\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2})}} \int_ {- \infty} ^ {\infty} d h _ {i} ^ {l} \phi (h _ {i} ^ {l}) e ^ {- \frac {1}{2} h _ {i} ^ {l} (\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}) ^ {- 1} h _ {i} ^ {l}} \\ = \frac {1}{\sqrt {2 \pi (\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2})}} \int_ {0} ^ {\infty} d h _ {i} ^ {l} (a _ {+} - a _ {-}) h _ {i} ^ {l} e ^ {- \frac {(h _ {i} ^ {l}) ^ {2}}{2 (\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2})}} \\ = \left(a _ {+} - a _ {-}\right) \sqrt {\frac {\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}}{2 \pi}}. \tag {108} \\ \end{array}
$$
$$
\chi_ {\mathcal {I}} ^ {l} = \frac {\sigma_ {w} ^ {2}}{\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}} \cdot \frac {\pi \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right)}{\pi \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right) - \left(a _ {+} - a _ {-}\right) ^ {2}}. \tag {109}
$$
The critical line is defined by  $\chi_{\mathcal{I}}^{\star} = 1$ , which can be solved as
$$
\sigma_ {b} = \sqrt {\frac {\left(a _ {+} - a _ {-}\right) ^ {2}}{\pi \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right) - \left(a _ {+} - a _ {-}\right) ^ {2}}} \sigma_ {w}. \tag {110}
$$
For ReLU with  $a_{+} = 1$  and  $a_{-} = 0$
$$
\begin{array}{l} \sigma_ {b} = \sqrt {\frac {1}{\pi - 1}} \sigma_ {w} \\ \approx 0. 6 8 3 \sigma_ {w}. \tag {111} \\ \end{array}
$$
I.6 Residual Connections
The recurrence relation for the NNGP kernel can be evaluated to be
$$
\mathcal {K} ^ {l + 1} = \frac {\sigma_ {w} ^ {2} \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right)}{2} \mathcal {K} ^ {l} + \sigma_ {b} ^ {2} + \mu^ {2} \mathcal {K} ^ {l}. \tag {112}
$$
The condition for the existence of fixed point
$$
\chi_ {\mathcal {K}} ^ {\star} = \frac {\sigma_ {w} ^ {2} \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right)}{2} + \mu^ {2} \leq 1 \tag {113}
$$
leads us to
$$
\sigma_ {w} ^ {2} \leq \frac {2 \left(1 - \mu^ {2}\right)}{a _ {+} ^ {2} + a _ {-} ^ {2}}. \tag {114}
$$
For  $\sigma_w^2 = \frac{2(1 - \mu^2)}{a_+^2 + a_-^2}$ , finite fixed point exists only if  $\sigma_b^2 = 0$ . (Diverges linearly otherwise)
The recurrence coefficient for Jacobian is evaluated to be
$$
\chi_ {\mathcal {I}} ^ {\star} = \frac {\sigma_ {w} ^ {2} \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right)}{2} + \mu^ {2}. \tag {115}
$$
The critical line is defined as
$$
\sigma_ {w} = \sqrt {\frac {2 \left(1 - \mu^ {2}\right)}{a _ {+} ^ {2} + a _ {-} ^ {2}}}. \tag {116}
$$
The critical point is located at  $\left(\sqrt{\frac{2(1 - \mu^2)}{a_+^2 + a_-^2}},0\right)$ .
For ReLU, the critical point is at  $\left(\sqrt{2(1 - \mu^2)},0\right)$
I.7 Residual Connections with LayerNorm on Preactivations (Pre-LN)
Again use lemma D.9 and combine all known results for scale-invariant functions
$$
\begin{array}{l} \chi_ {\mathcal {I}} ^ {\star} = \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 | _ {\tilde {\mathcal {K}} ^ {l - 1} = 1} + \mu^ {2}\right) \\ = \frac {\sigma_ {w} ^ {2} (a _ {+} ^ {2} + a _ {-} ^ {2}) (1 - \mu^ {2})}{\sigma_ {w} ^ {2} (a _ {+} ^ {2} + a _ {-} ^ {2}) + 2 \sigma_ {b} ^ {2}} + \mu^ {2} \\ = 1 - \frac {2 \sigma_ {b} ^ {2} \left(1 - \mu^ {2}\right)}{\sigma_ {w} ^ {2} \left(a _ {+} ^ {2} + a _ {-} ^ {2}\right) + 2 \sigma_ {b} ^ {2}} \tag {117} \\ \end{array}
$$
Similar to the case without residue connections
$$
\chi_ {\mathcal {I}} ^ {l} \leq 1 \tag {118}
$$
is always true. The equality only holds at  $\sigma_{b} = 0$  line for  $\mu < 1$ .
Notice there is a very special case  $\mu = 1$ , where the whole  $\sigma_{b} - \sigma_{w}$  plane is critical.
J Results for erf Activation Function
Definition J.1 (erf activation function).
$$
\phi (x) = \frac {2}{\sqrt {\pi}} \int_ {0} ^ {x} e ^ {- t ^ {2}} d t. \tag {119}
$$
J.1 NNGP Kernel
To evaluate lemma 2.2 exactly, we introduce two dummy variables  $\lambda_{1}$  and  $\lambda_{2}[51]$ .
$$
\begin{array}{l} \mathbb {E} _ {\theta} \left[ \phi \left(\lambda_ {1} h _ {i} ^ {l}\right) \phi \left(\lambda_ {2} h _ {i} ^ {l}\right) \right] = \int d \lambda_ {1} \int d \lambda_ {2} \frac {d ^ {2}}{d \lambda_ {1} d \lambda_ {2}} \mathbb {E} _ {\theta} \left[ \phi \left(\lambda_ {1} h _ {i} ^ {l}\right) \phi \left(\lambda_ {2} h _ {i} ^ {l}\right) \right] \\ = \int d \lambda_ {1} \int d \lambda_ {2} \int d h _ {i} ^ {l} \frac {4}{\sqrt {2 \pi^ {3} \mathcal {K} ^ {l}}} \left(h _ {i} ^ {l}\right) ^ {2} e ^ {- \left(\lambda_ {1} ^ {2} + \lambda_ {2} ^ {2} + \frac {1}{2 \mathcal {K} ^ {l}}\right) \left(h _ {i} ^ {l}\right) ^ {2}} \\ = \int d \lambda_ {1} \int d \lambda_ {2} \frac {4 \mathcal {K} ^ {l}}{\pi \left(1 + 2 \mathcal {K} ^ {l} \left(\lambda_ {1} ^ {2} + \lambda_ {2} ^ {2}\right)\right)} \\ = \frac {2}{\pi} \arcsin \left(\frac {2 \mathcal {K} ^ {l} \lambda_ {1} \lambda_ {2}}{1 + 2 \mathcal {K} ^ {l} \left(\lambda_ {1} ^ {2} + \lambda_ {2} ^ {2}\right)}\right). \tag {120} \\ \end{array}
$$
We use the special case where  $\lambda_1 = \lambda_2 = 1$ .
Thus the recurrence relation for the NNGP kernel with erf activation function is
$$
\mathcal {K} ^ {l + 1} = \frac {2 \sigma_ {w} ^ {2}}{\pi} \arcsin \left(\frac {2 \mathcal {K} ^ {l}}{1 + 2 \mathcal {K} ^ {l}}\right) + \sigma_ {b} ^ {2}. \tag {121}
$$
As in the scale-invariant case, finite fixed point only exists when
$$
\chi_ {\mathcal {K}} ^ {\star} = \frac {4 \sigma_ {w} ^ {2}}{\pi} \frac {1}{(1 + 2 \mathcal {K} ^ {\star}) \sqrt {1 + 4 \mathcal {K} ^ {\star}}} \leq 1. \tag {122}
$$
Numerical results show the condition is satisfied everywhere in  $\sigma_{b} - \sigma_{w}$  plane, where  $\chi_{\mathcal{K}}^{\star} = 1$  is only possible when  $\mathcal{K}^{\star} = 0$ .
J.2 Jacobians
Follow the definition
$$
\begin{array}{l} \chi_ {\mathcal {I}} ^ {j} = \sigma_ {w} ^ {2} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {i} ^ {j}\right) \phi^ {\prime} \left(h _ {i} ^ {j}\right) \right] \\ = \frac {4 \sigma_ {w} ^ {2}}{\sqrt {2 \pi^ {3} \mathcal {K} ^ {l}}} \int d h _ {i} ^ {l} e ^ {- 2 (h _ {i} ^ {l}) ^ {2}} e ^ {- \frac {(h _ {i} ^ {l}) ^ {2}}{2 \mathcal {K} ^ {l}}} \\ = \frac {4 \sigma_ {w} ^ {2}}{\pi} \frac {1}{\sqrt {1 + 4 \mathcal {K} ^ {l}}}. \tag {123} \\ \end{array}
$$
To find phase boundary  $\chi_{\mathcal{J}}^{\star} = 1$ , we need to combine Eq.(121) and Eq.(123) and evaluate them at  $\kappa^{\star}$ .
$$
\mathcal {K} ^ {\star} = \frac {2 \sigma_ {w} ^ {2}}{\pi} \arcsin \left(\frac {2 \mathcal {K} ^ {\star}}{1 + 2 \mathcal {K} ^ {\star}}\right) + \sigma_ {b} ^ {2}, \tag {124}
$$
$$
\chi_ {\mathcal {J}} ^ {\star} = \frac {4 \sigma_ {w} ^ {2}}{\pi} \frac {1}{\sqrt {1 + 4 \mathcal {K} ^ {\star}}} = 1. \tag {125}
$$
One can solve the equations above and find the critical line
$$
\sigma_ {b} = \sqrt {\frac {1 6 \sigma_ {w} ^ {4} - \pi^ {2}}{4 \pi^ {2}} - \frac {2 \sigma_ {w} ^ {2}}{\pi} \arcsin \left(\frac {1 6 \sigma_ {w} ^ {4} - \pi^ {2}}{1 6 \sigma_ {w} ^ {4} + \pi^ {2}}\right)}. \tag {126}
$$
Critical point is reached by further requiring  $\chi_{\mathcal{K}}^{\star} = 1$ . Since  $\chi_{\mathcal{K}}^{\star} \leq \chi_{\mathcal{J}}^{\star}$ , the only possible case is  $\mathcal{K}^{\star} = 0$ , which is located at
$$
\left(\sigma_ {w}, \sigma_ {b}\right) = \left(\sqrt {\frac {\pi}{4}}, 0\right). \tag {127}
$$
J.3 Critical Exponents
We show how to extract critical exponents of the NNGP kernel and Jacobians of erf activation function.
Critical point for erf is at  $(\sigma_b, \sigma_w) = (0, \sqrt{\frac{\pi}{4}})$ , with  $\mathcal{K}^{\star} = 0$ . Now suppose  $l$  is large enough such that the deviation of  $\mathcal{K}^l$  from fixed point value  $\mathcal{K}^{\star}$  is small. Define  $\delta \mathcal{K}^l \equiv \mathcal{K}^l - \mathcal{K}^{\star}$ . Eq.(121) can be rewritten as
$$
\delta \mathcal {K} ^ {l + 1} = \frac {1}{2} \arcsin \left(\frac {2 \delta \mathcal {K} ^ {l}}{1 + 2 \delta \mathcal {K} ^ {l}}\right) \tag {128}
$$
$$
\approx \delta \mathcal {K} ^ {l} - 2 (\delta \mathcal {K} ^ {l}) ^ {2}.
$$
From lemma G.1
$$
A = \frac {1}{2} \text {a n d} \zeta_ {\mathcal {K}} = 1. \tag {129}
$$
Next we analyze critical exponent of Jacobians by expanding (123) around  $\mathcal{K}^{\star} = 0$  critical point  $(\sigma_b,\sigma_w) = (0,\sqrt{\frac{\pi}{4}})$ .
To leading order  $l^{-1}$  we have
$$
\begin{array}{l} \chi_ {\mathcal {I}} ^ {l} \approx 1 - 2 \delta K ^ {l} \tag {130} \\ \approx 1 - \frac {1}{l}. \\ \end{array}
$$
Thus the recurrence relation for partial Jacobian, at large  $l$ , takes form
$$
\mathcal {J} ^ {l _ {0}, l + 1} = \left(1 - \frac {1}{l}\right) \mathcal {J} ^ {l _ {0}, l}. \tag {131}
$$
At large  $l$
$$
\mathcal {J} ^ {l _ {0}, l} = c _ {l _ {0}} l ^ {- 1}, \tag {132}
$$
with a non-universal constant  $c_{l_0}$ .
The critical exponent is
$$
\zeta = 1, \tag {133}
$$
which is the same as  $\zeta_{\mathcal{K}}$
J.4 LayerNorm on Pre-activations
Use lemma D.9, we have
$$
\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 {4 \sigma_ {w} ^ {2}}{\sqrt {5} \left[ 2 \sigma_ {w} ^ {2} \arcsin \left(\frac {2}{3}\right) + \pi \sigma_ {b} ^ {2} \right]}. \tag {134} \\ \end{array}
$$
The critical line is then defined by
$$
\begin{array}{l} \sigma_ {b} = \sqrt {\frac {2}{\pi} \left[ \frac {2}{\sqrt {5}} - \arcsin \left(\frac {2}{3}\right) \right]} \sigma_ {w} \\ \approx 0. 3 2 4 \sigma_ {w}. \tag {135} \\ \end{array}
$$
J.5 LayerNorm on Activations
Due to the symmetry of erf activation function  $\mathbb{E}_{\theta}\left[\phi (h_i^l)\right] = 0$ , we only need to modify our known results.
$$
\mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {i} ^ {l}\right) \phi^ {\prime} \left(h _ {i} ^ {l}\right) \right] = \frac {4}{\pi} \frac {1}{\sqrt {1 + 4 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right)}}, \tag {136}
$$
$$
\mathbb {E} _ {\theta} \left[ \phi \left(h _ {i} ^ {l}\right) \phi \left(h _ {i} ^ {l}\right) \right] = \frac {2}{\pi} \arcsin \left(\frac {2 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right)}{1 + 2 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right)}\right). \tag {137}
$$
$$
\chi_ {\mathcal {I}} ^ {l} = \frac {2 \sigma_ {w} ^ {2}}{\sqrt {1 + 4 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right)}} \cdot \frac {1}{\arcsin \left(\frac {2 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right)}{1 + 2 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right)}\right)}, \tag {138}
$$
where the phase boundary is defined by the transcendental equation  $\chi_{\mathcal{I}}^{l} = 1$
J.6 Residual Connections
$$
\mathcal {K} ^ {l + 1} = \frac {2 \sigma_ {w} ^ {2}}{\pi} \arcsin \left(\frac {2 \mathcal {K} ^ {l}}{1 + 2 \mathcal {K} ^ {l}}\right) + \sigma_ {b} ^ {2} + \mu^ {2} \mathcal {K} ^ {l}. \tag {139}
$$
Finite fixed point only exists when
$$
\chi_ {\mathcal {K}} ^ {\star} = \frac {4 \sigma_ {w} ^ {2}}{\pi} \frac {1}{(1 + 2 \mathcal {K} ^ {\star}) \sqrt {1 + 4 \mathcal {K} ^ {\star}}} + \mu^ {2} \leq 1. \tag {140}
$$
Notice that  $\chi_{\mathcal{K}}^{\star} \leq \chi_{\mathcal{J}}^{\star}$  still holds, where the equality holds only when  $\mathcal{K}^{\star} = 0$ .
$$
\chi_ {\mathcal {J}} ^ {\star} = \frac {4 \sigma_ {w} ^ {2}}{\pi} \frac {1}{\sqrt {1 + 4 \mathcal {K} ^ {\star}}} + \mu^ {2}. \tag {141}
$$
$$
\sigma_ {b} = \sqrt {\frac {1 6 \sigma_ {w} ^ {4} - \pi^ {2} (1 - \mu^ {2}) ^ {2}}{4 \pi^ {2} (1 - \mu^ {2})} - \frac {2 \sigma_ {w} ^ {2}}{\pi} \arcsin \left(\frac {1 6 \sigma_ {w} ^ {4} - \pi^ {2} (1 - \mu^ {2}) ^ {2}}{1 6 \sigma_ {w} ^ {4} + \pi^ {2} (1 - \mu^ {2}) ^ {2}}\right)}. \tag {142}
$$
$$
\left(\sigma_ {w}, \sigma_ {b}\right) = \left(\sqrt {\frac {\pi \left(1 - \mu^ {2}\right)}{4}}, 0\right). \tag {143}
$$
Note that for  $\mu = 1$ , one needs to put extra effort into analyzing the scaling behavior. First we notice that  $\mathcal{K}^l$  monotonically increases with depth  $l$  – the recurrence relation for the NNGP kernel at large  $l$  (or large  $\mathcal{K}^l$ ) is
$$
\mathcal {K} ^ {l + 1} \approx \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2} + \mathcal {K} ^ {l}, \tag {144}
$$
which regulates the first term in (141).
For  $\mu = 1$  at large depth
$$
\chi_ {\mathcal {I}} ^ {l} \sim 1 + \frac {4 \sigma_ {w} ^ {2}}{\pi \sqrt {C _ {0} + 4 \left(\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}\right) l}}. \tag {145}
$$
Here  $C_0$  is a constant that depends on the input.
We can approximate the asymptotic form of  $\log \mathcal{J}^{l_0,l'}$  as follows
$$
\begin{array}{l} \log \mathcal {J} ^ {l _ {0}, l} = \log \left(\prod_ {l ^ {\prime} = l _ {0}} ^ {l} \chi_ {\mathcal {J}} ^ {l ^ {\prime}}\right) \\ = \sum_ {l ^ {\prime} = l _ {0}} ^ {l} \log \left(1 + \frac {4 \sigma_ {w} ^ {2}}{\pi \sqrt {C _ {0} + 4 (\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}) l ^ {\prime}}}\right) \\ \approx \int_ {l _ {0}} ^ {l} d l ^ {\prime} \log \left(1 + \frac {4 \sigma_ {w} ^ {2}}{\pi \sqrt {C _ {0} + 4 (\sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}) l ^ {\prime}}}\right) \\ \sim 2 \tilde {c} \sqrt {l} + O (\log l), \tag {146} \\ \end{array}
$$
where  $\tilde{c} = \frac{2\sigma_w^2}{\pi\sqrt{\sigma_w^2 + \sigma_b^2}}$
We conclude that at large depth, the APJN for  $\mu = 1$ , erf networks can be written as
$$
\mathcal {J} ^ {l _ {0}, l} \sim O \left(e ^ {2 \tilde {c} \sqrt {l} + O (\log l)}\right). \tag {147}
$$
This result checks out empirically, as shown in Figure 9.8
J.7 Residual Connections with LayerNorm on Preactivations (Pre-LN)
Use lemma D.9 and the results we had without residue connections for erf with LayerNorm on preactivations.
$$
\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} \left(\tilde {h} _ {k} ^ {l}\right) \phi^ {\prime} \left(\tilde {h} _ {k} ^ {l}\right)\right] \Bigg | _ {\tilde {\mathcal {K}} ^ {l - 1} = 1} + \mu^ {2}\right) \\ = \frac {4 \sigma_ {w} ^ {2} \left(1 - \mu^ {2}\right)}{\sqrt {5} \left[ 2 \sigma_ {w} ^ {2} \arcsin \left(\frac {2}{3}\right) + \pi \sigma_ {b} ^ {2} \right]} + \mu^ {2}. \tag {148} \\ \end{array}
$$
$$
\sigma_ {b} = \sqrt {\frac {2}{\pi} \left[ \frac {2}{\sqrt {5}} - \arcsin \left(\frac {2}{3}\right) \right]} \sigma_ {w} \tag {149}
$$
$$
\approx 0. 3 2 4 \sigma_ {w}.
$$
<sup>8</sup>We used NTK parameterization for this experiment. However, we emphasize that it does not affect the final result.
Figure 9:  $\log (\mathcal{J}^{l_0,l}) - \sqrt{l}$  for  $\mu = 1$ ,  $\sigma_b^2 = 0$ , erf.
K Results for GELU Activation Function
Definition K.1 (GELU activation function).
$$
\begin{array}{l} \phi (x) = \frac {x}{2} \left[ 1 + \operatorname {e r f} \left(\frac {x}{\sqrt {2}}\right) \right] \\ = \frac {x}{2} \left[ 1 + \frac {2}{\sqrt {\pi}} \int_ {0} ^ {\frac {x}{\sqrt {2}}} e ^ {- t ^ {2}} d t \right]. \tag {150} \\ \end{array}
$$
K.1 NNGP Kernel
Use lemma 2.2 for GELU
$$
\begin{array}{l} \mathbb {E} _ {\theta} \left[ \phi \left(h _ {i} ^ {l}\right) \phi \left(h _ {i} ^ {l}\right) \right] = \frac {1}{\sqrt {2 \pi \mathcal {K} ^ {l}}} \int d h _ {i} ^ {l} \frac {\left(h _ {i} ^ {l}\right) ^ {2}}{4} \left[ 1 + \operatorname {e r f} \left(\frac {h _ {i} ^ {l}}{\sqrt {2}}\right) \right] ^ {2} e ^ {- \frac {\left(h _ {i} ^ {l}\right) ^ {2}}{2 \mathcal {K} ^ {l}}} \\ = \frac {1}{\sqrt {2 \pi \mathcal {K} ^ {l}}} \int d h _ {i} ^ {l} \frac {\left(h _ {i} ^ {l}\right) ^ {2}}{4} \left[ 1 + \operatorname {e r f} ^ {2} \left(\frac {h _ {i} ^ {l}}{\sqrt {2}}\right) \right] e ^ {- \frac {\left(h _ {i} ^ {l}\right) ^ {2}}{2 \mathcal {K} ^ {l}}} \\ = \frac {\mathcal {K} ^ {l}}{4} + \frac {1}{\sqrt {3 2 \pi \mathcal {K} ^ {l}}} \int d h _ {i} ^ {l} \left(h _ {i} ^ {l}\right) ^ {2} \operatorname {e r f} ^ {2} \left(\frac {h _ {i} ^ {l}}{\sqrt {2}}\right) e ^ {- \frac {\left(h _ {i} ^ {l}\right) ^ {2}}{2 \mathcal {K} ^ {l}}} \\ = \frac {\mathcal {K} ^ {l}}{4} + \frac {\mathcal {K} ^ {l}}{\sqrt {3 2 \pi \mathcal {K} ^ {l}}} \int d h _ {i} ^ {l} \operatorname {e r f} ^ {2} \left(\frac {h _ {i} ^ {l}}{\sqrt {2}}\right) e ^ {- \frac {\left(h _ {i} ^ {l}\right) ^ {2}}{2 \mathcal {K} ^ {l}}} \\ + \frac {(\mathcal {K} ^ {l}) ^ {2}}{\sqrt {3 2 \pi \mathcal {K} ^ {l}}} \int d h _ {i} ^ {l} \left[ \operatorname {e r f} ^ {\prime} \left(\frac {h _ {i} ^ {l}}{\sqrt {2}}\right) \operatorname {e r f} ^ {\prime} \left(\frac {h _ {i} ^ {l}}{\sqrt {2}}\right) + \operatorname {e r f} \left(\frac {h _ {i} ^ {l}}{\sqrt {2}}\right) \operatorname {e r f} ^ {\prime \prime} \left(\frac {h _ {i} ^ {l}}{\sqrt {2}}\right) \right] e ^ {- \frac {\left(h _ {i} ^ {l}\right) ^ {2}}{2 \mathcal {K} ^ {l}}} \\ = \frac {\mathcal {K} ^ {l}}{4} + \frac {\mathcal {K} ^ {l}}{2 \pi} \left[ \arcsin \left(\frac {\mathcal {K} ^ {l}}{1 + \mathcal {K} ^ {l}}\right) + \frac {2 \mathcal {K} ^ {l}}{(1 + \mathcal {K} ^ {l}) \sqrt {1 + 2 \mathcal {K} ^ {l}}} \right], \tag {151} \\ \end{array}
$$
where from the third line to the fourth line we used integrate by parts twice, and to get the last line we used results from erf activations.
Thus 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}. \tag {152}
$$
$$
\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]. \tag {153}
$$
K.2 Jacobians