SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Darshil Doshi
ddoshi@umd.edu
Tianyu He*†
tianyuh@umd.edu
Andrey Gromov
gromovand@meta.com
Deep neural networks are notorious for defying theoretical treatment. However, when the number of parameters in each layer tends to infinity, the network function is a Gaussian process (GP) and quantitatively predictive description is possible. Gaussian approximation allows one to formulate criteria for selecting hyperparameters, such as variances of weights and biases, as well as the learning rate. These criteria rely on the notion of criticality defined for deep neural networks. In this work we describe a new practical way to diagnose criticality. We introduce partial Jacobians of a network, defined as derivatives of preactivations in layer  $l$  with respect to preactivations in layer  $l_0 \leq l$ . We derive recurrence relations for the norms of partial Jacobians and utilize these relations to analyze criticality of deep fully connected neural networks with LayerNorm and/or residual connections. We derive and implement a simple and cheap numerical test that allows one to select optimal initialization for a broad class of deep neural networks; containing fully connected, convolutional and normalization layers. Using these tools we show quantitatively that proper stacking of the LayerNorm (applied to preactivations) and residual connections leads to an architecture that is critical for any initialization. Finally, we apply our methods to analyze ResNet and MLP-Mixer architectures; demonstrating the everywhere-critical regime. $^4$
When the number of parameters in each layer becomes large, the functional space description of deep neural networks simplifies dramatically. The network function,  $f(x)$ , in this limit, is a Gaussian process [37, 30] with a kernel – sometimes referred to as neural network Gaussian process (NNGP) kernel [30] – determined by the network architecture and hyperparameters (e.g. depth, precise choices of layers and the activation functions, as well as the distribution of weights and biases). A similar line of reasoning was earlier developed for recurrent neural networks [36]. Furthermore, for special choices of parameterization and MSE loss function, the training dynamics under gradient descent can be solved exactly in terms of the neural tangent kernel (NTK) [27, 31]. A large body of work was devoted to the calculation of the NNGP kernel and NTK for different architectures, calculation of the finite width corrections to these quantities, and empirical investigation of the training dynamics of wide networks [39, 54, 25, 14, 3, 33, 1, 18, 19, 45, 58, 47, 6, 5, 32, 63, 61, 60, 59, 35, 17, 2, 49, 34].
*Equal Contributions
†Condensed Matter Theory Center & Department of Physics, University of Maryland, College Park, MD 20740
†Meta AI, Menlo Park, CA 94025
$^{4}$ Code to reproduce our results is available at https://github.com/ablghtianyi/PartialJac.git
One important result that arose from these works is that the network architecture determines the most appropriate initialization of the weights and biases [44, 46, 30]. To state this result, we consider networks with/without LayerNorm [7] and residual connections [21]; the preactivations for which can be defined as follows
$$
h _ {i} ^ {l + 1} (x) = \sum_ {j = 1} ^ {N _ {l}} w _ {i j} ^ {l + 1} \phi \left(\tilde {h} _ {j} ^ {l} (x)\right) + b _ {i} ^ {l + 1} + \mu h _ {i} ^ {l} (x), \tag {1}
$$
where  $\tilde{h}_j^l = \mathrm{LayerNorm}(h_j^l)$  and the parameter  $\mu$  controls the strength of residual connections. For the input layer:  $h_i^1 (x) = \sum_{j = 1}^{N_0}w_{ij}^1 x_j + b_i^1$ . In the  $(l + 1)$ -th layer, weights  $w_{ij}^{l + 1}\in \mathbb{R}^{N_{l + 1}\times N_l}$  and biases  $b_{i}^{l + 1}\in \mathbb{R}^{N_{l + 1}\times 1}$  are taken from normal distributions  $\mathcal{N}(0,\sigma_w^2 /N^l)$  and  $\mathcal{N}(0,\sigma_b^2)$ , respectively. Hyperparameters  $\sigma_w$  and  $\sigma_{b}$  need to be tuned.  $\phi (\cdot)$  is the activation function and  $x\in \mathbb{R}^{N_0\times 1}$  is the input. For results discussed in this work,  $x$  can be sampled from either a realistic (i.e. highly correlated) dataset or a high entropy distribution.
For a network of depth  $L$ , the network function is given by  $f(x) = h^{L}(x)$ . Different network architectures and activation functions,  $\phi$ , lead to different "optimal" choices of  $(\sigma_w, \sigma_b)$ . The optimal choice can be understood, using the language of statistical mechanics, as a critical point (or manifold) in the  $\sigma_b - \sigma_w$  plane. The notion of criticality becomes sharp as the network depth,  $L$ , becomes large. Criticality ensures that NNGP kernel and the gradients' norm remain  $O(L^0)$  as the network gets deeper [45]. Very deep networks will not train unless initialized critically, since the gradients explode or vanish exponentially. Note that high trainability does not imply that the trained model has great performance (test accuracy) after training.
1.1 Results
Here we focus on two main results of this work: (i) empirical method to check criticality of a neural network and (ii) an architecture based on layer normalization and residual connections that is critical for any initialization. First we introduce the notion of a partial Jacobian.
Definition 1.1. Let  $h_i^l (x)$  be preactivations of a neural network  $f(x)$ . The partial Jacobian  $J_{ij}^{l_0,l}$  is defined as derivative of preactivations at layer  $l$  with respect to preactivations at layer  $l_0\leq l$
$$
J _ {i j} ^ {l _ {0}, l} (x) = \frac {\partial h _ {j} ^ {l} (x)}{\partial h _ {i} ^ {l _ {0}} (x)}. \tag {2}
$$
The partial Jacobian is a random matrix with vanishing mean at initialization. Fixing  $l_{0}$  and varying  $l$  allows us to study the behavior of Jacobian with depth. On the other hand, varying both  $l_{0}$  and  $l$  is necessary to study general networks with non-repeating building blocks. Next, we introduce a deterministic measure of the magnitude of  $J_{ij}^{l_{0},l}$  — its squared Frobenius norm, averaged over parameter-initializations.
Definition 1.2. Let  $J_{ij}^{l_0,l}$  be a partial Jacobian of a neural network  $f(x)$ . Averaged partial Jacobian norm (APJN) $^5$  is defined as
$$
\mathcal {J} ^ {l _ {0}, l} (x) \equiv \mathbb {E} _ {\theta} \left[ \frac {1}{N _ {l}} \sum_ {j = 1} ^ {N _ {l}} \sum_ {i = 1} ^ {N _ {l _ {0}}} \left(\frac {\partial h _ {j} ^ {l} (x)}{\partial h _ {i} ^ {l _ {0}} (x)}\right) ^ {2} \right], \tag {3}
$$
where  $\mathbb{E}_{\theta}$  indicates averaging over parameter-initializations.
APJN is a dominant factor in the magnitude of the gradients. In what follows, we show that criticality, studied previously in the literature, occurs when APJN either remains finite, or varies algebraically as  $l$  becomes large. To prove this we derive the recurrence relation for  $\mathcal{I}^{l_0,l}(x)$  in the limit  $N_{l}\to \infty$  and analyze it at large depth. Algebraic behavior of APJN with depth is characterized by an architecture-dependent critical exponent,  $\zeta$ , so that  $\mathcal{I}^{l_0,l}(x)\approx l^{-\zeta}$ . Such behavior is familiar from statistical mechanics when a system is tuned to a critical point [10]. Away from criticality, there are two phases:
In practice, we compute APJN using a very cheap estimator, viz. Hutchinson's trace estimator[24].
ordered and chaotic. In the ordered phase APJN vanishes exponentially with depth, whereas in the chaotic phase APJN grows exponentially
$$
\mathcal {J} ^ {l _ {0}, l} \approx c _ {l _ {0}} e ^ {\pm \frac {l}{\xi}}. \tag {4}
$$
Here  $\xi$  is the correlation length. It characterizes how fast gradients explode or vanish. (See Section 2.1 for further discussion.)
Theorem 1.3 (Main result). Let  $f(x)$  be a deep MLP network with Lipschitz continuous activation  $\phi(\cdot)$ . Assume that LayerNorm is applied to preactivations and there are residual connections with strength  $\mu$  acting according to (1). In the sequential limit  $(N_1 \to \infty, \dots, N_l \to \infty, \dots, N_{L-1} \to \infty)$ , the correlation length with a large  $L$  can be written as
$$
\xi = \frac {1}{\left| \log \left[ (1 - \mu^ {2}) \frac {A}{B} + \mu^ {2} \right] \right|}, \tag {5}
$$
where the non-negative constants  $A$  and  $B$  are given by
$$
A = \sigma_ {w} ^ {2} \mathbb {E} _ {\theta} \left[ \phi^ {\prime} (\tilde {h} _ {k} ^ {L - 2}) ^ {2} \right], \quad B = \sigma_ {w} ^ {2} \mathbb {E} _ {\theta} \left[ \phi (\tilde {h} _ {k} ^ {L - 2}) ^ {2} \right] + \sigma_ {b} ^ {2},
$$
and  $\phi'(\cdot)$  is the derivative of  $\phi(\cdot)$ . We omitted the average over the neuron index  $k$  in the infinite width limit, and the result does not depend on  $k$ .
Remark 1.4. As  $\mu$  increases, the dependence on the initialization  $(\sigma_b,\sigma_w)$  gets weaker. When  $\mu = 1$ , the correlation length diverges; and the network is critical for any initialization, with  $\zeta = O(1)$ .
In practice, Theorem 1.3 and Remark 1.4 imply that different choices of initialization bear no effect on trainability of the network provided that LayerNorm and residual connections are arranged as stated.
Some of our results were either surmised or obtained in a different form in the literature. We find that LayerNorm ensures that NNGP kernel remains finite at any depth as suggested in the original work of Ba et al. [7]. LayerNorm also alters the criticality of  $\mathcal{I}^{l_0,l}(x)$ . It was noted in Xu et al. [57] that LayerNorm (applied to preactivations) regularizes the backward pass. We formalize this observation by showing that LayerNorm (applied to preactivations) dramatically enhances correlation length (which is not the case for LayerNorm applied to activations). This can be seen from Theorem 1.3, setting  $\mu = 0$ . When residual connections of strength 1 are combined with erf (or any other erf-like activation function, e.g. tanh), the neural network enters a subcritical phase with enhanced correlation length (see 4.3). A version of this result was discussed in Yang and Schoenholz [62]. When residual connections are introduced on top of LayerNorm, the correlation length  $\xi$  is further increased. If residual connections have strength  $\mu = 1$  the network enters a critical phase for any initialization. Importance of correct ordering of LayerNorm, residual connections and attention layers was discussed in Xiong et al. [56]. Benefit of combining BatchNorm and residual connections was mentioned in [12]. Several architectures with the same order of GroupNorm and residual connections were investigated in Yu et al. [64]. Existing initialization schemes such as He initialization[20], Fix-up [65] and ReZero [8] are special cases of our method. They conform to the notion of criticality using APJN, as defined in our work.
The partial Jacobian has been used to study generalization bounds in Arora et al. [4]. The Jacobian norm (i.e.  $||J_{ij}^{0,l}||^2$ ) of trained feed-forward neural networks was studied in Novak et al. [38], where it was correlated with generalization. Partial Jacobians with  $l_0 = l - 1$  were studied in the context of RNNs [11, 9], referred to as state-to-state Jacobians. Furthermore, the relation between Partial Jacobians and Lyapunov exponents, as well as their impact on trainability, have been explored for RNNs [16, 15].
As the aspect ratio  $(L / N)$  of the network approaches 1, the finite width corrections to the Jacobian become more prominent. On the other hand, even with a small aspect ratio, the effect of the spectral density of Jacobian becomes important as the depth  $L$  becomes very large. Pennington et al. [43] studied the spectrum of the input-output Jacobian for MLPs. Xiao et al. [54] extended the analysis to CNNs, showing that very deep vanilla CNNs can be trained by achieving "dynamical isometry".
He initialization is critical only for ReLU networks.
2 Recurrence Relations
Here we derive the infinite width recurrence relations for the APJN and the NNGP kernel. We use Lemma 2.2 to derive NNGP kernel recurrence relation, and leverage that to get the recurrence relation for APJN. Fixed point analyses of these relations help us define critical line and point. We start with the vanilla MLP with no LayerNorm and  $\mu = 0$ . Results with LayerNorm and/or residual connections, as well as modern architectures are presented in the following sections. (We refer the reader to Appendices D to F for proofs and detailed application of this recipe. Appendices I to K contain the calculations and results for various activation functions.)
Definition 2.1. We define averaged covariance of preactivations as follows
$$
\mathcal {K} ^ {l} \left(x, x ^ {\prime}\right) = \mathbb {E} _ {\theta} \left[ \frac {1}{N _ {l}} \sum_ {i = 1} ^ {N _ {l}} h _ {i} ^ {l} (x) h _ {i} ^ {l} \left(x ^ {\prime}\right) \right]. \tag {6}
$$
Lemma 2.2. When  $N_{l} \to \infty$  for  $l = 1, \ldots, L - 1$  sequentially, the expectation value over parameter initializations for a general function of preactivations:  $\mathcal{O}(h^{l}(x))$ , can be expressed as the averaging over the Gaussian process  $h^{l}(x)$  with covariance  $\mathcal{K}^{l}(x, x')$ .
$$
\mathbb {E} _ {\theta} \left[ \mathcal {O} \left(h _ {i} ^ {l} (x)\right) \right] = \frac {1}{\sqrt {2 \pi \mathcal {K} ^ {l} (x , x)}} \int d h ^ {l} \mathcal {O} \left(h _ {i} ^ {l} (x)\right) e ^ {- \frac {\left(h _ {i} ^ {l} (x)\right) ^ {2}}{2 \mathcal {K} ^ {l} (x , x)}}. \tag {7}
$$
This result has been established in Lee et al. [30]. Note that the density in (7) only depends on the diagonal part of the covariance matrix,  $\mathcal{K}^l(x,x)$ . We will refer to  $\mathcal{K}^l(x,x)$  as NNGP kernel.
Remark 2.3. In the sequential infinite width limit the means appearing in (9)-(11) are self-averaging and, therefore, deterministic. They converge in distribution to their averages over parameterizations.
$$
\frac {1}{N _ {l}} \sum_ {i = 1} ^ {N _ {l}} \phi \left(h _ {i} ^ {l}\right) ^ {2} \xrightarrow {N _ {l ^ {\prime} \leq l - 1} \rightarrow \infty} \mathbb {E} _ {\theta} \left[ \frac {1}{N _ {l}} \sum_ {i = 1} ^ {N _ {l}} \phi \left(h _ {i} ^ {l}\right) ^ {2} \right]. \tag {8}
$$
When performing analytic calculations we use the infinite width convention; whereas in our finite-width experiments we explicitly average over initializations of  $\theta^l$ .
Theorem 2.4. With 2.2, in the infinite width limit, the NNGP kernel  $\mathcal{K}^{l + 1}(x,x)$  is deterministic, and can be determined recursively via
$$
\mathcal {K} ^ {l + 1} (x, x) = \sigma_ {w} ^ {2} \mathbb {E} _ {\theta} \left[ \frac {1}{N _ {l}} \sum_ {i = 1} ^ {N _ {l}} \phi \left(h _ {i} ^ {l} (x)\right) ^ {2} \right] + \sigma_ {b} ^ {2}. \tag {9}
$$
Theorem 2.5. Let  $f(x)$  be an MLP network with a Lipschitz continuous activation function  $\phi (x)$ . In the infinite width limit, APJN  $\mathcal{I}^{l_0,l + 1}(x)$  is deterministic and satisfies a recurrence relation
$$
\mathcal {J} ^ {l _ {0}, l + 1} (x) = \chi_ {\mathcal {J}} ^ {l} \mathcal {J} ^ {l _ {0}, l} (x), \tag {10}
$$
where the factor  $\chi_{\mathcal{J}}^l$  is given by
$$
\chi_ {\mathcal {J}} ^ {l} = \sigma_ {w} ^ {2} \mathbb {E} _ {\theta} \left[ \frac {\sigma_ {w} ^ {2}}{N _ {l}} \sum_ {i = 1} ^ {N _ {l}} \left(\phi^ {\prime} \left(h _ {i} ^ {l} (x)\right)\right) ^ {2} \right]. \tag {11}
$$
2.4 is due to Neal [37] and Lee et al. [30]. 2.5 is new and is valid only in the limit of infinite width. The proof is in D. We will drop the explicit dependence on  $x$  to improve readability.
The expectation values that appear in (9)-(11) are evaluated using (7). When the integrals can be taken analytically, they lead to explicit equations for the critical lines and/or the critical points. Details of these calculations as well as the derivation of (9)-(11) can be found in the D. A subtlety emerges in (10) when  $l_0 = 0$ , where a correction of the order  $O(N_0^{-1})$  arises for non-scale invariant activation functions. This subtlety is discussed in the D.
When the depth of the network becomes large, the  $l$ -dependence of the expectation values that appear in (6), (11) saturate to a (possibly infinite) constant value; which means that  $\mathcal{K}^l$ ,  $\mathcal{J}^{l_0,l}$  and  $\chi_{\mathcal{J}}^l$  have
reached a fixed point. We denote the corresponding quantities as  $\mathcal{K}^{\star}$ ,  $\mathcal{J}^{l_0,\star}$ ,  $\chi_{\mathcal{J}}^{\star}$ . The existence of a fixed point is not obvious and should be checked on a case by case basis. Fixed point analysis for  $\mathcal{K}^l$  was done in Poole et al. [44] for bounded activation functions and in Roberts et al. [45] for the general case. The stability is formulated in terms of
$$
\chi_ {\mathcal {K}} ^ {\star} = \left. \frac {\partial \mathcal {K} ^ {l + 1}}{\partial \mathcal {K} ^ {l}} \right| _ {\mathcal {K} ^ {l} = \mathcal {K} ^ {\star}}. \tag {12}
$$
The norm of preactivations remains finite (or behaves algebraically) when  $\chi_{\mathcal{K}}^{\star} = 1$
Eq. (10) nicely expresses  $\mathcal{J}^{l_0,l + 1}$  as a linear function of  $\mathcal{J}^{l_0,l}$ . The behaviour of  $\mathcal{J}^{l_0,l + 1}$  at large  $l$  is determined by  $\chi_l^l$ . When  $\chi_l^l > 1$  partial Jacobians diverge exponentially, while for  $\chi_l^l < 1$  partial Jacobians vanish exponentially. Neural networks are trainable only up to a certain depth when initialized  $O(1)$  away from criticality, which is determined by the equation
$$
\chi_ {\mathcal {I}} ^ {\star} = 1. \tag {13}
$$
Eq. (13) is an implicit equation on  $\sigma_{b},\sigma_{w}$  and generally outputs a critical line in  $\sigma_{b} - \sigma_{w}$  plane. The parameter  $\chi_{\mathcal{J}}^{\star}$  has to be calculated on a case-by-case basis using either (11) or the method presented in the next section. Everywhere on the critical line,  $\mathcal{J}^{l_0,l}$  saturates to a constant or behaves algebraically.
When the condition  $\chi_{\mathcal{K}}^{\star} = 1$  is added, we are left with a critical point<sup>7</sup>. This analysis of criticality at infinite width agrees with Roberts et al. [45], where  $\chi_{\perp}$  is to be identified with  $\chi_{\mathcal{I}}^{\star}$ ; and Schoenholz et al. [46], Martens et al. [34], where their analysis based on the equivalent  $\chi_{1}$  or  $C'(1)$  only works for bounded activation functions. In particular, condition (13) together with  $\chi_{\mathcal{K}}^{\star} = 1$  ensures that NTK is  $O(1)$  at initialization.
2.1 Empirical Diagnostic of Criticality
APJN  $\mathcal{I}^{l_0,l}$  provides a clear practical way to diagnose whether the network is critical or not. Proper choice of  $l_{0}$  and  $l$  allows us to minimize the non-universal effects and cleanly extract  $\chi_{\mathcal{I}}^{\star}$ .
Recurrence relation (10), supplemented with the initial condition  $\mathcal{I}^{l_0,l_0 + 1} = \chi_{\mathcal{I}}^{l_0}$ , can be formally solved as
$$
\mathcal {J} ^ {l _ {0}, l} = \prod_ {\ell = l _ {0}} ^ {l - 1} \chi_ {\mathcal {I}} ^ {\ell}. \tag {14}
$$
We would like to obtain an estimate of  $\chi_{\mathcal{I}}^{\star}$  as accurately as possible. To that end, imagine that for some  $l' > l_0$  the fixed point has been essentially reached and  $\chi_{\mathcal{I}}^{l'} \approx \chi_{\mathcal{I}}^{\star}$ . Then the APJN
$$
\mathcal {J} ^ {l _ {0}, l} = \left(\chi_ {\mathcal {J}} ^ {\star}\right) ^ {l - l ^ {\prime} - 1} \cdot \prod_ {\ell = l _ {0}} ^ {l ^ {\prime}} \chi_ {\mathcal {J}} ^ {\ell} \tag {15}
$$
depends on the details of how the critical point is approached; which are encoded in the last factor.
Proposition 2.6. If the network is homogeneous, i.e., consists of (possibly complex) blocks of layers, periodically repeated  $L$  times; then the penultimate APJN provides an accurate estimate of  $\chi_{\mathcal{I}}^{\star}$ :
$$
\left. \mathcal {J} ^ {L - 2, L - 1} \right| _ {L \rightarrow \infty} = \chi_ {\mathcal {I}} ^ {\star}. \tag {16}
$$
This is a direct consequence of combining (9) and (11) as  $L$  goes to infinity. See 1 for numerical justification.
2.6 is the central result of this section and will be heavily used in the remainder of this work.
Note that for deep networks, away from criticality, APJN takes the form
$$
\mathcal {J} ^ {l _ {0}, l} \approx c _ {l _ {0}} e ^ {\pm \frac {l}{\xi}}, \quad \xi = | \log \chi_ {\mathcal {J}} ^ {\star} | ^ {- 1}, \tag {17}
$$
where  $c_{l_0}$  is a non-universal constant that depends on  $l_0$ . If the sign in (17) is positive  $(\chi_{\mathcal{I}}^{\star} > 1)$  the network is in the chaotic phase, while when the sign is negative  $(\chi_{\mathcal{I}}^{\star} < 1)$  the network is in the
${}^{7}$  Scale-invariant activation functions are more forgiving: away from the critical point,  ${\mathcal{K}}^{l}$  scales algebraically with  $l$  .
Figure 1: log-log plots of partial Jacobian  $\mathcal{J}^{0,l}$  vs.  $l$ . From left to right: (1)erf, (2)ReLU, (3)erf and GELU with LayerNorm applied to preactivations and residual connections of strength 1. The fluctuations get larger towards the output because the aspect ratio (i.e.  $l / N_{l}$ ) approaches  $1 / 4$ .
ordered phase.  $\xi$  has the meaning of correlation length: on the depth scale of approximately  $k\xi$  the gradients remain appreciable, and hence the network with the depth of  $\approx k\xi$  will train.
We used (16) to map out the  $\sigma_{b} - \sigma_{w}$  phase diagrams of various MLP architectures. The location of the critical line agrees remarkably well with our infinite width calculations. Results are presented in Fig. 2.
At criticality,  $\chi_{\mathcal{I}}^{\star} = 1$  and the correlation length diverges; indicating that gradients can propagate arbitrarily far. A more careful analysis of non-linear corrections shows that APJN can exhibit algebraic behavior with depth and can still vanish in the infinite depth limit, but much slower than the ordered phase.
2.2 Scaling at a Critical Point
At criticality  $\chi_{\mathcal{I}}^{l}$  saturates to a fixed value  $\chi_{\mathcal{I}}^{\star} = 1$ . If we are interested in  $\mathcal{I}^{l_0,l}$  with  $l - l_0 = O(L)$  then it is essential to know how exactly  $\chi_{\mathcal{I}}^{l}$  approaches 1.
Theorem 2.7. Assume that deep neural network  $f(x)$  is initialized critically. Then  $l \to \infty$  asymptotics of APJN is given by
$$
\mathcal {J} ^ {l _ {0}, l} (x) = O \left(l ^ {- \zeta}\right), \tag {18}
$$
where  $\zeta$  is the critical exponent Roberts et al. [45], see Appendix G for further details.
Critical exponents can be determined analytically in the limit of infinite width. Note that  $\chi_{\mathcal{J}}^l$ , given by (11), depends on  $\kappa^l$  by virtue of (7). Consequently, Eqs. (9)-(11) are coupled through non-linear (in  $\kappa^l$  and  $\mathcal{I}^{l_0,l}$ ) terms. These non-linear corrections are absent for any scale-invariant activation function, but appear for other activation functions.
We checked the scaling empirically by plotting  $\mathcal{J}^{0,l}$  vs.  $l$  in a log-log plot and fitting the slope. These results are presented in Fig.1, and are in excellent agreement with infinite width calculation.
3 Layer Normalization
The fact that critical initialization is concentrated on a single point  $(\sigma_w^\star, \sigma_b^\star)$  may appear unsettling because great care must be taken to initialize the network critically. The situation can be substantially improved by utilizing the normalization techniques known as LayerNorm [7] and GroupNorm [52]. Our results apply to GroupNorm verbatim in the case when the number of groups is much smaller than the width. LayerNorm can act either on preactivations or on activations (discussed in the Appendix D). Depending on this choice, criticality will occur on different critical lines in  $\sigma_b - \sigma_w$  plane. When LayerNorm is applied to preactivations the correlation length is enhanced, allowing for training much deeper networks even far away from criticality.
The LayerNorm applied to preactivations takes the following form