SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Theorem D.16. In the infinite width limit the recurrence relation for the NNGP kernel with Layer-Norm on activations is
$$
\mathcal {K} ^ {l + 1} = \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}. \tag {52}
$$
$$
\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} \widetilde {\phi \left(h _ {j} ^ {l}\right)} + b _ {i} ^ {l + 1}\right) \left(\sum_ {k = 1} ^ {N _ {l}} w _ {i k} ^ {l + 1} \widetilde {\phi \left(h _ {k} ^ {l}\right)} + b _ {i} ^ {l + 1}\right) \right] \\ = \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \widetilde {\phi (h _ {j} ^ {l})} ^ {2} \right] + \sigma_ {b} ^ {2} \\ = \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \left(\frac {\phi (h _ {j} ^ {l}) - \mathbb {E} _ {\theta} [ \phi (h ^ {l}) ]}{\sqrt {\mathbb {E} _ {\theta} [ \phi (h ^ {l}) ^ {2} ] - \mathbb {E} _ {\theta} [ \phi (h ^ {l}) ] ^ {2}}}\right) ^ {2} \right] + \sigma_ {b} ^ {2} \\ = \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \frac {\mathbb {E} _ {\theta} \left[ \left(\phi (h _ {j} ^ {l}) - \mathbb {E} _ {\theta} \left[ \phi (h ^ {l}) \right]\right) ^ {2} \right]}{\mathbb {E} _ {\theta} \left[ \phi (h ^ {l}) ^ {2} \right] - \mathbb {E} _ {\theta} \left[ \phi (h ^ {l}) \right] ^ {2}} + \sigma_ {b} ^ {2} \\ = \sigma_ {w} ^ {2} + \sigma_ {b} ^ {2}. \tag {53} \\ \end{array}
$$
Theorem D.17. In the infinite width limit the recurrence relation for partial Jacobian with LayerNorm on activations is
$$
\mathcal {J} ^ {l _ {0}, l + 1} = \chi_ {J} ^ {l} \mathcal {J} ^ {l _ {0}, l}, \tag {54}
$$
where  $\chi_J^l\equiv \sigma_w^2\frac{\mathbb{E}_\theta[\phi'(h^l)^2)]}{\mathbb{E}_\theta[\phi(h^l)^2] - \mathbb{E}_\theta[\phi(h^l)]^2}.$
$$
\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 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} \widetilde {\phi^ {\prime} (h _ {k} ^ {l})}\right) \left(w _ {i m} ^ {l + 1} \widetilde {\phi^ {\prime} (h _ {m} ^ {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}} \sum_ {k = 1} ^ {N _ {l}} \sum_ {j = 1} ^ {N _ {\ell_ {0}}} \mathbb {E} _ {\theta} \left[ \widetilde {\phi^ {\prime} \left(h _ {k} ^ {l}\right)} \widetilde {\phi^ {\prime} \left(h _ {k} ^ {l}\right)} \left(\sum_ {j = 1} ^ {N _ {\ell_ {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}} \sum_ {k = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \widetilde {\phi^ {\prime} \left(h _ {k} ^ {l}\right)} ^ {2} \right] \mathcal {J} ^ {l _ {0}, l} \\ = \sigma_ {w} ^ {2} \frac {\mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h ^ {l}\right) ^ {2} \right]}{\mathbb {E} _ {\theta} \left[ \phi \left(h ^ {l}\right) ^ {2} \right] - \mathbb {E} _ {\theta} \left[ \phi \left(h ^ {l}\right) \right] ^ {2}} \mathcal {J} ^ {l _ {0}, l} \\ = \chi_ {J} ^ {l} \mathcal {J} ^ {l _ {0}, l}, \tag {55} \\ \end{array}
$$
E Residual Connections
Definition E.1. We define residual connections by the modified the recurrence relation for preactivations (Eq.(1))
$$
h _ {i} ^ {l + 1} = \sum_ {j = 1} ^ {N _ {l}} w _ {i j} ^ {l + 1} \phi \left(h _ {j} ^ {l}\right) + b _ {i} ^ {l + 1} + \mu h _ {i} ^ {l}, \tag {56}
$$
where the parameter  $\mu$  controls the strength of the residual connection.
Remark E.2. Note that this definition requires  $N_{l + 1} = N_l$ . We ensure this by only adding residual connections to the hidden layers, which are of the same width. More generally, one can introduce a tensor parameter  $\mu_{ij}$ .
Remark E.3. In general, the parameter  $\mu$  could be layer-dependent  $(\mu^l)$ . But we suppress this dependence here since we are discussing self-similar networks.
Theorem E.4. In the infinite width limit, the recurrence relation for the NNGP kernel with residual connections is changed by an additional term controlled by  $\mu$
$$
\mathcal {K} ^ {l + 1} = \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} + \mu^ {2} \mathcal {K} ^ {l}. \tag {57}
$$
$$
\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 \left(h _ {j} ^ {l}\right) + b _ {i} ^ {l + 1} + \mu h _ {i} ^ {l}\right) \left(\sum_ {k = 1} ^ {N _ {l}} w _ {i k} ^ {l + 1} \phi \left(h _ {k} ^ {l}\right) + b _ {i} ^ {l + 1} + \mu h _ {i} ^ {l}\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} + \mu^ {2} h _ {i} ^ {l} h _ {i} ^ {l} \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] + \mu^ {2} \frac {1}{N _ {l + 1}} \sum_ {i = 1} ^ {N _ {l + 1}} \mathbb {E} _ {\theta} \left[ h _ {i} ^ {l} h _ {i} ^ {l} \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} + \mu^ {2} \mathcal {K} ^ {l}, \tag {58} \\ \end{array}
$$
where we used the fact  $N_{l + 1} = N_l$  to get the last line.
Theorem E.5. In the infinite width limit, the recurrence relation for partial Jacobians with residual connections has a simple multiplicative form
$$
\mathcal {J} ^ {l _ {0}, l + 1} = \chi_ {\mathcal {I}} ^ {l} \mathcal {J} ^ {l _ {0}, l}, \tag {59}
$$
where the recurrence coefficient is shifted to  $\chi_{\mathcal{J}}^{l} = \sigma_{w}^{2}\mathbb{E}_{\theta}\left[\phi^{\prime}(h_{k}^{l})\phi^{\prime}(h_{k}^{l})\right] + \mu^{2}$ .
$$
\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) + \mu \delta_ {i k}\right) \left(w _ {i m} ^ {l + 1} \phi^ {\prime} \left(h _ {m} ^ {l}\right) + \mu \delta_ {i m}\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}} \left(w _ {i k} ^ {l + 1} w _ {i m} ^ {l + 1} \phi^ {\prime} \left(h _ {k} ^ {l}\right) \phi^ {\prime} \left(h _ {m} ^ {l}\right) + \mu^ {2} \delta_ {i k} \delta_ {i m}\right) \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {m} ^ {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 _ {\ell_ {0}}} \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}}\right) \right] + \frac {1}{N _ {l}} \sum_ {k = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \mu^ {2} \left(\sum_ {j = 1} ^ {N _ {\ell_ {0}}} \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}}\right) \right] \\ = \left(\sigma_ {w} ^ {2} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {k} ^ {l}\right) \phi^ {\prime} \left(h _ {k} ^ {l}\right) \right] + \mu^ {2}\right) \mathbb {E} _ {\theta} \left[ \frac {1}{N _ {l}} \sum_ {k = 1} ^ {N _ {l}} \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] \\ = \left(\sigma_ {w} ^ {2} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} \left(h _ {k} ^ {l}\right) \phi^ {\prime} \left(h _ {k} ^ {l}\right) \right] + \mu^ {2}\right) \mathcal {J} ^ {l _ {0}, l} \\ \end{array}
$$
$$
\mathcal {J} ^ {l _ {0}, l + 1} = \chi_ {\mathcal {I}} ^ {l} \mathcal {J} ^ {l _ {0}, l}. \tag {60}
$$
F Residual Connections with LayerNorm on Preactivations (Pre-LN)
We recall the recurrence relation (1):
$$
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} + \mu h _ {i} ^ {l}. \tag {61}
$$
Theorem F.1. In the infinite width limit, the recurrence relation for the NNGP kernel is then modified to
$$
\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} + \mu^ {2} \mathcal {K} ^ {l}. \tag {62}
$$
$$
\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} + \mu h _ {i} ^ {l}\right) \left(\sum_ {k = 1} ^ {N _ {l}} w _ {i k} ^ {l + 1} \phi (\tilde {h} _ {k} ^ {l}) + b _ {i} ^ {l + 1} + \mu h _ {i} ^ {l}\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} + \mu^ {2} \mathcal {K} ^ {l}. \tag {63} \\ \end{array}
$$
Remark F.2. For  $\mu < 1$ , the recursion relation has a fixed point
$$
\mathcal {K} ^ {\star} = \frac {\sigma_ {w} ^ {2}}{N _ {l ^ {\star}} (1 - \mu^ {2})} \sum_ {j = 1} ^ {N _ {l ^ {\star}}} \mathbb {E} _ {\theta} \left[ \phi \left(\tilde {h} _ {j} ^ {l ^ {\star}}\right) \phi \left(\tilde {h} _ {j} ^ {l ^ {\star}}\right) \right] + \frac {\sigma_ {b} ^ {2}}{1 - \mu^ {2}}. \tag {64}
$$
where the average here is exactly the same as cases for LayerNorm applied to preactivations without residue connections.  $l^{\star}$  labels some very large depth  $l$ .
Remark F.3. For  $\mu = 1$  case, the solution of (62) is
$$
\mathcal {K} ^ {l} = \mathcal {K} ^ {0} + \sum_ {l ^ {\prime} = 1} ^ {l} \left(\frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi \left(\tilde {h} _ {j} ^ {l ^ {\prime}}\right) \phi \left(\tilde {h} _ {j} ^ {l ^ {\prime}}\right) \right] + \sigma_ {b} ^ {2}\right). \tag {65}
$$
which is linearly growing since the expectation does not depend on depth.  $\mathcal{K}^0$  is the NNGP kernel after the input layer.
Theorem F.4. In the infinite width limit, the recurrence relation for Jacobians changes by a constant shift in the recursion coefficient.
$$
\mathcal {J} ^ {l _ {0}, l + 1} = \chi_ {\mathcal {J}} ^ {l} \mathcal {J} ^ {l _ {0}, l}, \tag {66}
$$
where for this case
$$
\chi_ {\mathcal {I}} ^ {l} = \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] + \mu^ {2}. \tag {67}
$$
$$
\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(\frac {w _ {i k} ^ {l + 1} \phi^ {\prime} (\tilde {h} _ {k} ^ {l})}{\sqrt {\mathcal {K} ^ {l}}} + \mu \delta_ {i k}\right) \left(\frac {w _ {i m} ^ {l + 1} \phi^ {\prime} (\tilde {h} _ {m} ^ {l})}{\sqrt {\mathcal {K} ^ {l}}} + \mu \delta_ {i k}\right) \left(\frac {\partial h _ {k} ^ {l}}{\partial h _ {j} ^ {l _ {0}}} \frac {\partial h _ {m} ^ {l}}{\partial h _ {j} ^ {l _ {0}}}\right) \right] \\ = \mathbb {E} _ {\theta} \left[ \left(\frac {\sigma_ {w} ^ {2}}{N _ {l} \mathcal {K} ^ {l}} \sum_ {k = 1} ^ {N _ {l}} \phi^ {\prime} (\tilde {h} _ {k} ^ {l}) \phi^ {\prime} (\tilde {h} _ {k} ^ {l}) + \mu^ {2}\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] \\ = \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] + \mu^ {2}\right) \mathcal {J} ^ {l _ {0}, l} \\ = \chi_ {\mathcal {I}} ^ {l} \mathcal {J} ^ {l _ {0}, l}, \tag {68} \\ \end{array}
$$
Remark F.5. Assuming  $l \geq l^{\star}$ , we combine (64) and (67) to derive Equation (5) in the Theorem 1.3.
$$
\begin{array}{l} \chi_ {\mathcal {I}} ^ {*} = (1 - \mu^ {2}) \frac {\frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {k = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} (\tilde {h} _ {k} ^ {l}) \phi^ {\prime} (\tilde {h} _ {k} ^ {l}) \right]}{\frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi (\tilde {h} _ {j} ^ {l}) \phi (\tilde {h} _ {j} ^ {l}) \right] + \sigma_ {b} ^ {2}} + \mu^ {2} \\ = \left(1 - \mu^ {2}\right) \frac {A}{B} + \mu^ {2}, \tag {69} \\ \end{array}
$$
where  $A = \sum_{k=1}^{N_l} \mathbb{E}_{\theta}\left[\phi'(\tilde{h}_k^l)\phi'(\tilde{h}_k^l)\right]$  and  $B = \frac{\sigma_w^2}{N_l}\sum_{j=1}^{N_l}\mathbb{E}_{\theta}\left[\phi(\tilde{h}_j^l)\phi(\tilde{h}_j^l)\right] + \sigma_b^2$ . Recall that  $\xi = |\log \chi_{\mathcal{J}}^\star|^{-1}$ . This directly leads us to Equation (5).
$$
\xi = \frac {1}{\left| \log \left(\chi_ {\mathcal {I}} ^ {\star}\right) \right|} = \frac {1}{\left| \log \left((1 - \mu^ {2}) \frac {A}{B} + \mu^ {2}\right) \right|} \tag {70}
$$
Note that in Equation (5), we discard the average over neurons in  $A$  and  $B$ , since we are in the infinite width limit.
Remark F.6. As we mentioned above  $\mu = 1$  needs extra care. Plugging  $\mu = 1$  into the result (65) and (67) we find out that
$$
\begin{array}{l} \left. \chi_ {\mathcal {I}} ^ {l} \right| _ {\mu = 1} = \frac {\sigma_ {w} ^ {2} \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]}{N _ {l} \mathcal {K} ^ {0} + \sum_ {l ^ {\prime} = 1} ^ {l} \left(\sigma_ {w} ^ {2} \sum_ {j = 1} ^ {N _ {l}} \mathbb {E} _ {\theta} \left[ \phi \left(\tilde {h} _ {j} ^ {l}\right) \phi \left(\tilde {h} _ {j} ^ {l}\right) \right] + N _ {l} \sigma_ {b} ^ {2}\right)} + 1 \\ \sim 1 + O \left(\frac {1}{l}\right). \tag {71} \\ \end{array}
$$
Thus, we get everywhere criticality: critical behavior independent of the choice of initialization  $(\sigma_{b},\sigma_{w})$ . It also follows that the correlation length  $\xi$  diverges in this case. (71) leads to power law behaviour in Jacobians (with exponent  $\zeta$ ) at large depth. Note that the exponent  $\zeta$  is not universal.
Remarks F.5 and F.6 together serve as the complete proof of Theorem 1.3.
G Critical Exponents
To prove theorem 2.7, we first need to find the critical exponent of the NNGP kernel [45].
Lemma G.1. In the infinite width limit, consider a critically initialized network with an activation function  $\phi$ . The scaling behavior of the fluctuation  $\delta \mathcal{K}^l \equiv \mathcal{K}^l - \mathcal{K}^\star$  in non-exponential. If the recurrence relation can be expanded to leading order  $\delta K^l$  as  $\delta \mathcal{K}^{l + 1} \approx \delta \mathcal{K}^l - c_n (\delta \mathcal{K}^l)^n$  for  $n \geq 2$ . The solution of  $\delta \mathcal{K}^l$  is
$$
\delta \mathcal {K} ^ {l} = \frac {1}{c _ {n} (n - 1)} l ^ {- \zeta_ {\mathcal {K}}}, \tag {72}
$$
where  $\zeta_{\mathcal{K}} = \frac{1}{n - 1}$
Remark G.2. The constant  $c_{n}$  and the order of the first non-zero term  $n$  is determined by the choice of activation function.
Proof. We can expand the recurrence relation for the NNGP kernel (9) to the second order of  $\delta \mathcal{K}^l = \mathcal{K}^l -\mathcal{K}^\star$  on both sides.
$$
\delta \mathcal {K} ^ {l + 1} \approx \delta \mathcal {K} ^ {l} - c _ {n} (\delta \mathcal {K} ^ {l}) ^ {n}. \tag {73}
$$
Use power law ansatz  $\delta \mathcal{K}^l = A l^{-\zeta_{\kappa}}$  then
$$
(l + 1) ^ {- \zeta \kappa} = l ^ {- \zeta \kappa} - c _ {n} A l ^ {- n \zeta \kappa}. \tag {74}
$$
Multiply  $l^{\zeta \kappa}$  on both side then use Taylor expansion  $\left(\frac{l}{l + 1}\right)^{\zeta \kappa} \approx 1 - \frac{\zeta \kappa}{l}$
$$
\frac {\zeta \kappa}{l} = c _ {n} A l ^ {- (n - 1) \zeta \kappa}. \tag {75}
$$
For arbitrary  $l$ , the only non-trivial solution of the equation above is
$$
A = \frac {1}{c _ {n} (n - 1)} \text {a n d} \zeta_ {\mathcal {K}} = \frac {1}{n - 1}. \tag {76}
$$
Proof of theorem 2.7. We will assume  $c_{2} \neq 0$ . Then use lemma G.1, we can expand  $\chi_{\mathcal{I}}^{l}$  in terms of  $\delta \mathcal{K}^{l}$ . To leading order  $l^{-1}$
$$
\begin{array}{l} \chi_ {\mathcal {I}} ^ {l} \approx 1 - d _ {1} \delta \mathcal {K} ^ {l} \\ = 1 - \frac {d _ {1}}{c _ {2}} l ^ {- 1}. \tag {77} \\ \end{array}
$$
Consider a sufficiently large  $l$ . In this case  $O(l^{-1})$  approximation is valid. We write recurrence relations of Jacobians as
$$
\begin{array}{l} \mathcal {J} ^ {l _ {0}, l} = \prod_ {l ^ {\prime} = l _ {0}} ^ {l - 1} \left(1 - \frac {d _ {1}}{c _ {2}} l ^ {\prime - 1}\right) \mathcal {J} ^ {l _ {0}, l _ {0}} \\ \approx c _ {l _ {0}} \cdot l ^ {- \zeta}. \tag {78} \\ \end{array}
$$
When  $c_{n} = 0$  for all  $n \geq 2$ , from lemma G.1 we have  $\delta \mathcal{K}^{l} = 0$ . Thus the Jacobian saturates to some constant.
We checked the scaling empirically by plotting  $\mathcal{I}^{0,l}$  vs.  $l$  in a log-log plot and fitting the slope. These results are presented in Fig.1.
H MLP-Mixer
In this section we would like to analyze an architecture called MLP-Mixer [48], which is based on multi-layer perceptrons (MLPs). A MLP-Mixer (i) chops images into patches, then applies affine transformations per patch, (ii) applies several Mixer Layers, (iii) applies pre-head LayerNorm, Global Average Pooling, and an output affine transformation. We will explain the architecture by showing forward pass equations.
Suppose one has a single input with dimension  $(C_{in},H_{in},W_{in})$ . We label it as  $x_{\mu i}$ , where the Greek letter labels channels and the Latin letter labels flattened pixels.
First of all the (i) is realized by a special convolutional layer, where kernel size  $f$  is equal to the stride  $s$ . Then, the first convolution layer can be written as
$$
h _ {\mu i} ^ {0} = \sum_ {j = 1} ^ {f ^ {2}} \sum_ {\nu = 1} ^ {C _ {i n}} W _ {\mu \nu ; j} ^ {0} x _ {\nu , j + (i - 1) s ^ {2}} + b _ {\mu i} ^ {0}, \tag {79}
$$
where  $f$  is the size of filter and  $s$  is the stride. In our example  $f = s$ . Notice in PyTorch both bias and weights are sampled from a uniform distribution  $\mathcal{U}(-\sqrt{k},\sqrt{k})$ , where  $k = (C_{in}f^{2})^{-1}$ .
$$
\mathbb {E} _ {\theta} \left[ W _ {\mu \nu ; i} ^ {0} W _ {\rho \sigma ; j} ^ {0} \right] = \frac {1}{3 C _ {i n} f ^ {2}} \delta_ {\mu \rho} \delta_ {\nu \sigma} \delta_ {i j}, \tag {80}
$$
$$
\mathbb {E} _ {\theta} \left[ b _ {\mu i} ^ {0} b _ {\nu j} ^ {0} \right] = \frac {1}{3 C _ {i n} f ^ {2}} \delta_ {\mu \nu} \delta_ {i j}. \tag {81}
$$
Notice that the output of Conv2d:  $h_{\mu i}^{0} \in \mathbb{R}^{C \times N_{p}}$ , where  $C$  stands for channels and  $N_{p} = H_{in}W_{in} / f^{2}$  stands for patches, both of them will be mixed later by Mixer layers.
Next, we stack  $l$  Mixer Layers. A Mixer Layer contains LayerNorms and two MLPs, where the first one mixed patches  $i,j$  (token mixing) with a hidden dimension  $N_{tm}$ , the second one mixed channels  $\mu, \nu$  (channel-mixing) with a hidden dimension  $N_{cm}$ . Notice that for Mixer Layers we use the standard parameterization.
- First LayerNorm. It acts on channels  $\mu$ .
$$
\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 {82}
$$
where we defined a channel mean  $\mathbb{E}_C[h_{\rho i}^{6l}] \equiv \frac{1}{C}\sum_{\rho = 1}^C h_{\rho i}^{6l}$  and channel variance  $\mathrm{Var}_C \equiv \mathbb{E}_C\left[\left(h_{\rho i}^{6l}\right)^2\right] - \left(\mathbb{E}_C[h_{\rho i}^{6l}]\right)^2$ .
- First MlpBlock. It mixes patches  $i, j$ , preactivations from different channels that share the same weight and bias.
-  $6l + 1$ : Linear Affine Layer.
$$
h _ {\mu j} ^ {6 l + 1} = \sum_ {k = 1} ^ {N _ {p}} w _ {j k} ^ {6 l + 1} \tilde {h} _ {\mu k} ^ {6 l} + b _ {j} ^ {6 l + 1}. \tag {83}
$$
-  $6l + 2$ : Affine Layer.
$$
h _ {\mu i} ^ {6 l + 2} = \sum_ {j = 1} ^ {N _ {t m}} w _ {i j} ^ {6 l + 2} \phi \left(h _ {\mu j} ^ {6 l + 1}\right) + b _ {i} ^ {6 l + 2}, \tag {84}
$$
where  $N_{tm}$  stands for the hidden dimension of "token mixing".
-  $6l + 3$ : Residual Connections.
$$
h _ {\mu i} ^ {6 l + 3} = h _ {\mu i} ^ {6 l + 2} + \mu h _ {\mu i} ^ {6 l}. \tag {85}
$$
- Second LayerNorm. It again acts on channels  $\mu$ .
$$
\tilde {h} _ {\mu i} ^ {6 l + 3} = \frac {h _ {\mu i} ^ {6 l + 3} - \mathbb {E} _ {C} \left[ h _ {\rho i} ^ {6 l + 3} \right]}{\sqrt {\operatorname {V a r} _ {C} \left[ h _ {\rho i} ^ {6 l + 3} \right]}}. \tag {86}
$$