Critical Initialization of Wide and Deep Neural Networks using Partial Jacobians: General Theory and Applications
Authors: Darshil Doshi
Tianyu He
Andrey Gromov
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Criticality
Gaussian Process
Jacobian
LayerNorm
Residual connections
ResNet
TL;DR: (i) We introduce a new diagnostic for critical initialization in a wide class of deep neural networks. (ii) We show that a combination of Normalization layers and residual connections leads to everywhere-critical architectures.

Abstract: 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.

Critical Initialization of Wide and Deep Neural Networks using Partial Jacobians: General Theory and Applications
Darshil Doshi ddoshi @umd edu
Tianyu He* tianyuh@umd.cm edu
Andrey Gromov #
gromovandemeta..com
Abstract
Deep neural networks are notorious for defying theoretical treatment. However, when the number of parameters each layer tends to infinity, the network function is a Gaussian process (GP) and quantitatively predictive description is possible. Gaussian approximation allows one formulate criteria for selecting hyperparam- eters, 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 new practical way diagnose criticality. We introduce partial Jacobians network, defined as derivatives of preactivations in layer with respect preactivations in layer lg I L. We derive recurrence relations for the norms of partial Jacobians and utilize these relations to analyze criticality 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 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 preactivations) and residual connections leads to an architecture that is critical for any initialization. Finally, apply methods to analyze ResNet and MLP-Mixer architectures; demonstrating the everywhere-critice regime. f
1
Introduction
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 kernel
30 with kernel sometimes referred to neural network Gaussian process (NNGP) 30 determined by the network architecture and hyperparameters (e.g depth, precise choices of layers and the activation functions, well as the distribution of weights and biases). A similar line 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 large body of work was devoted the calculation of the NNGP kernel and NTK for different architectures, calculation the finite width corrections to these quantities, and empirical investigation of the training dynamics of
process 377
wide networks 39 54 25
33
58
60 59
Equual Contributions
Condensed Matter Theory Center & Department of Physics, University of Maryland, College Park, MD
20740
Mett AI, Menlo Park, CA 94025
1 "Code reproduce our results available httppp///gghhhhhoooo////aaaaaaas
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

One important result that arose from these works is that the network architecture determines the most appropriate initialization of the weights and biases 44 networks with/without LayerNorm 7 and residual connections
30
To state this result, we consider
[1] the preactivations for which
can be defined as follows
(+ (z)
Σ u (1))(() + + ph((z),
where hj LayerNorm(tt)) and the parameter μ controls the strength of residual connections. For the input layer: (x) = - w/ya 67. the (L 1)-th layer, weights w &R and biases 1/+11 E M are taken from normal distributions, o( and and of.). respectively. Hyperparameters ow and ob need to be tuned. Φ(-) is the activation function and E R Nox1 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 network of depth L, the network function is given by f(x) = (x). Different network architectures and activation functions, Φ, lead to different "optimal" choices of (aw. Tb). The optimal choice can be understood, using the language of statistical mechanics, in the σg-Oa 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 (I.) 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.
critical point (or manifold)
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 (x) be preactivations of a neural network (=x) The partial Jacobian J is defined as derivative of preactivations at layer i with respect to preactivations layer to 1
ahl (z)
Jl (z)
(2)
ah (x)
The partial Jacobian is random matrix with vanishing mean at initialization. Fixing lo and varying I allows us to study the behavior of Jacobian with depth. On the other hand, varying both lo and I is necessary study general networks with non-repeating building blocks. Next, introduce a deterministic measure of the magnitude of fla. its squared Frobenius norm, averaged over parameter-mitializatoonss
Definition 1.2. Let d be a partial Jacobian of a neural network f(xx). Averaged partial Jacobian norm (APJN is defined as
N₁
Ni
(h,(aa) a (z)
27
flbl(z)= = Ep
(3)
N
where Eg indicates averaging over parcamter-imittalizatiins
APJN is a dominant factor in the magnitude of the gradients. In what follows, we show that criticality, studied previously the literature, occurs when APJN either remains finite, or varies algebraically as I becomes large. To prove this we derive the recurrence relation for f'o:(2)) in the limit N + oc and analyze it large depth. Algebraic behavior of APJN with depth characterized by an architecture- dependent critical exponent, Ç. so that (x). x I-S, Such behavior is familiar from statistical mechanics when a system is tundd to critical point 10 Away from criticality, there are two phases:
In practice, compute APJN using very cheap estimator, viz. Hutchinson's 's trace estimator 24

---

ordered and chaotic. In the ordered phase APIN vanishes exponentially with depth, whereas in the chaotic phase APJN grows exponentially
glad c
(4)
Here { is the correlation length. It characterizes how fast gradients explode or vanish. (See Section 2.1 for further discussion.)
Theorem 1.3 (Main result). Lei f(x) be deep MLP network with Lipschitz continuous activation
(()). Assume that LayerNorm is applied to preactivations and there are residual connections with strength μ acting according to Q In the sequential limit (N1 ......, Ni ..,..., N.,,,NL111
o0), the correlation length with large can be written as
8= log [11
(5)
12) 12]l
where the non negatiitevve constants and are given by
A A= (ht
B B= of [ L 23p]
+o0
and d.) is the derivative of o(). We omitted the average over the neuron index l in the infinite width limit, and the result does not depend k.
Remark 1.4. As increases, the dependence on the initialization (on, Tw) gets weaker. When μ 1, the correlation length diverges; and the network is critical for any initialization, with Ç = O(1).
In practice, Theorem 1.3 and Remark 1.4 imply that different choices of initialization bear effect on trainability of the network provided that LayerNorm and residual connections are arranged as stated.
1.2 Related Work
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 al. 77 LayerNorm also alters the criticality of flor was noted in Xu 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 μ = 0. When residual connections strength are combined with erf (or any other erf- like activation function, e.g. tanh), the neural network enters suberitical 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 top of LayerNorm, the correlation length & is further increased. If residual connections have strength = 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 initializatior
[20), Fix-up [65] and ReZero [ 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 study generalization bounds in Arora et al. 44. The Jacobian norm (Le. [|79/1113 of trained feed-forward neural networks was studied in Novak et al. (38]. where it correlated with generalization. Partial Jacobians with Iy 1 were studied in the context of RNNs [I]]] 99. referred to as state-to-stad Jacobians. Furthermore, the relation between Partial Jacobians and Lyapunov exponents, as well as their impact on trainability, have been explored for
RNNs [6 15
As the aspect ratio (I/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".
'ee initialization is critical only for ReLU networks.
3

---

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 APIN. Fixed point analyses of these relations help us define critical line and point. We start with the vanilla MLP with no LayerNorm and connections, as well as modern architectures are presented in the following sections. (We refer the reader to Appendices to F for proof' and detailed application of this recipe. Appendices to K contain the calculations and results for various activation functions.)
O. Results with LayerNorm and/or residual
Definition 2.1. We define averaged covariance of preactivations follows
N:
(x, 2') Eg
Σ M(()))(())) N
1
(6)
Lemma 2.2. When Ni- - - L- sequentially, the expectation value over parameter initializations for a general function of preactivations: O(h' (z)), can be expressed as the averaging over the Gaussian process (x) with covariance K K(
Bo [O(h(()))]
(I(()) (2.)).)
V 22TK3 (z,
This result has been established in Lee et al. [30 Note that the density in 7 only depends on the diagonal part the covariance matrix, KC' (', z). We will refer to IC' (z, (0, x) as NNGP kernel.
Remark 2.3. In the sequential infinite width limit the means appearing and, therefore, deterministic. They converge in distribution to their averages parameterizations.
(-- [1 are self-averaging
N;
o(h)) N001-1-000
Es
o(h))
N1
N
When performing analytic calculations we use the infinite width convention; whereas in our finite- width experiments we explicitly average over initializations of 0'.
Theorem 2.4. With in the infinite width limit, the NNGP kernel KC (zz, x)) is deterministic, and can be determined recursively via
N Σ ((((())))
KC- (x, o
+ +
N
Theorem 2.5. Let (1: be an MLP network with Lipschitz continuous activation function ((x). In the infinite width limit, APJN flo.++(x) is deterministic and satisfies a recurrence relation
Flod+ (x) = Flud
(10)
where the factor X'y is given by
X = Es
(o (d(()())))
(11)
N
2.4 is due to Neal 37 and Lee et al.
2.5 is new and is valid only in the limit of infinite width. We will drop the explicit dependence on x to improve readability.
30
The proof is in D
The expectation values that appear in 9
are evaluated using 7 When the integrals can be taken analytically, they lead explicit equations for the critical lines and/or the critical points. Details of these calculations as well as the derivation of (9)-(I) can be found in the A subtlety emerges in (10) when In = O, where a correction of the order O arises for non-scale invariant activation functions. This subtlety is discussed in the D
When the depth the network becomes large, the I-dependence of the expectation values that appear in 6 [II saturate (possibly infinite) constant value; which means that IC', Tloo and X²y have

---

reached a fixed point. We denote the corresponding quantities as K "".'"oo*.." .X'. The existence of fixed point obvious and should be checked on case by case basis. Fixed point analysis for K.' was done in Poole et al. [4 for bounded activation functions and in Roberts et al. 45 for the general case. The stability is formulated in terms of
ak!+1 XIC
(12)
ak! K'-K-
The norm of preactivations remains finite (or behaves algebraically) when X²: = 1.
Eq. (T0) nicely expresses Jlo.+1 as a linear function of Joo''. The behaviour Jlod+1 at large I is determined by X'7- When x'7 > partial Jacobians diverge exponentially, while for x < 1 partial Jacobians vanish exponentially. Neural networks are trainable only up to a certain depth when initialized (1) away from criticality, which is determined by the equation
=
(13)
Eq. 13 is an implicit equation on ob. the and generally outputs critical line in db-O plane. The parameter X has to be calculated on case-by-case basis using either (TT) or the method presented in the next section. Everywhere on the critical line, Jlo. saturates to a constant or behaves algebraically.
When the condition x² = is added, we are left with a critical poin This analysis of criticality at infinite width agrees with Roberts et al. [45), where X to be identified with X²; and Schoenholz al. [46]. Martens et al. (34), where their analysis based on the equivalent X1 or (1( only works
for bounded activation functions. In particular, condition (33 together with X = 1 ensures that
NTK is (1) initialization.
2.1 Empirical Diagnostic of Criticality
APJN Too provides a clear practical way to diagnose whether the network is critical or not. Proper choice of in and allows minimize the non-universal effects and cleanly extract X²--
Recurrence relation (10). supplemented with the initial condition
Jlo.o+1
X can be formally
solved as
Jood
II X E=lo
(14)
We would like obtain an estimate of X3 as accurately as possible. To that end, imagine that for some l' > lo the fixed point has been essentially reached and X"` R Xyy. Then the APJN
II
X3
(15)
E=la
depends the details of how the critical point is approached; which 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 Xy:
J1-221-11
Xy
(16)
L
This is direct consequence of combining and 11 as L goes to infinity. Seel 7 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
glodt 9
+
= = log
(17)
where Ct2 is a non-universal constant that depends on 4. If the sign in (7 is positive (x² > 1) the network is in the chaotic phase, while when the sign is negative (X²r 1) the network is in the
activation functions are more forgiving: away from the critical point, kc scales algebraically
with I.
5

---

ert
ReLU
LN-GELU, LN-erf: μ
Experiment
Experiment
CEW/TheRY En -1012
Theory, <= 1.000
Theory, <= 0.000
Fitting, 968
Fitting, 0.004
10a
on Repedteet
2m--221
x 10
S 10
10"
10-2
10a
10a
10
10a
10'
10
Figure 1: log-log plots of partial Jacobian To. vs. From left right: (I)rf, (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. I/N) approaches 1/4.
ordered phase. & has the meaning of correlation length: on the depth scale of approximately kg the gradients remain appreciable, and hence the network with the depth of = ky will train.
We used (66) map out the a-- 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.
At criticality, x = 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 Critical Point
At criticality Xy saturates to fixed value X 1. we are interested in Flod with 1o = O(L) then it essential know how exactly X'y approaches
Theorem 2.7. Assume that deep neural network f(x) is initialized critically. Then I oo asymp- totics of APJN is given by
(18)
where ( the critical exponent Roberts al. 45
Appendix for further details.
Critical exponents can determined analytically in the limit of infinite width. Note that x'y given by (IT), depends on K by virtue of (7). Consequently, Eqs. 9 II are coupled through non-linear (in KC and J'o-)) 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 Jo. vs. in log-log plot and fitting the slope. These results are presented in [ and are excellent agreement with infinite width calculation.
Layer Normalization
The fact that critical initialization is concentrated single point (oth oī)) 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 [ and GroupNorm [52]. Our results apply to GroupNorm verbatim in the case when the number of groups is much smaller than the width. LayerNorm Depending on this choice, criticality will occur on different critical lines in σ plane. When LayerNorm is applied to preuctivations the correlation length is enhanced, allowing for training much deeper networks even far away from criticality.
act either preactivations
activations (discussed in the Appendix(D)
The LayerNorm applied preactivations takes the following form

---

3.0
Reull,
LM
LN-AEU
0.4
LK-erf
0.0
0.
16.
44-
60
44- C6-
0.6
36-
60
0.3
Figure 2: X- empirical phase diagrams for MLP with L= 50, 500. The solid lines indicate the critical lines obtained through infinite width limit calculations, and the stars indicate the critical points. The dotted lines in the rightmost column correspond to the critical lines for μ < 1 case. For networks with LayerNorm and 1, X holds on the entire dg-Ow plane. We also note that for erf activation, the case μ = without LayerNorm subcritical and has large correlation length.
Definition 3.1 (Normalized preactivations).
i [E[( ETh
h1 Biny
N N10000
1
h
(19)
where we have introduced E[h'] = I /N/ In the limit of infinite width E[h'] and [[(h))]] = K', defined according to 6
Normalized preactivations, in are distributed according to A'0 1) for all ww Tg. The norms are, therefore, always finite and the condition x = 1 is trivially satisfied. This results in a critical line rather than a critical point.
The recurrence relations
9
for the NNGP and partial Jacobians are only slightly modified
N;
N;
k+ = ogEn
2(h))
+ +o}
X
Eg
8' holds.
(20)
NJ 1=1
N
1=1
Assuming that the value of X/y at the fixed point is X2y, the network is critical when
13)
X's (20) is depth independent and changes slowly with aw and σg- Thus, X²y remains close to 1 for wider range of hyperparameters. Consequently, the correlation length large even away from criticality, leading to much higher trainability of deep networks.
Residual (Skip) Connections
Adding residual connections between the network layers is widely used technique to facilitate the training of deep networks. Originally introduced [2T] in the context of convolutional neural
networks (29 (CNNs) for image recognition, residual connections have since been used in variety
of networks architectures and tasks [50 13].
Consider ① with non-zero μ and without LayerNorm layers. Then the recurrence relations for the NNGP kernel and x-y are modified as follows
N;
N Σ
k1+1
22 Eo
oh2))
+ to + 22, X/y = = E.
('L)))2
+ +2²
(21)
N 1=1
N
k=1
7

---

Kel 0.0
S4lU
UM
0.0
0.6
UK-rnt
LK-
68
5.
60
0.7
Figure 3: Trainability (Training Accuracy) deep MLP (N₃ 500, L 50) on FNSTO The combination of LayerNorm and makes the network of erf activation without LayerNorm, and = also has enhanced trainability. The dashed white lines denote the (analytical) critical lines.
The suberitical case
Remark 4.1. When < 1, the fixed point value of NNGP kernel is scaled by (1 - 112)-11 For the critical point formally (0,0)
Remark 4.2. For = 1, 21 implies that x'y 2 1, where the equality holds on the ow = 0 axis. Consequently, APJN exponentially diverges as function of depth / for all o O. In this case, a. needs to be taken sufficiently close to 0 to ensure trainability at large depths.
When μ < 1, residual connections amplify the chaotic phase and decrease the correlation length away from criticality for unbounded activation functions.
Solving the recurrence relations (2T) for erf activation, we find an effect observed in Yang and Schoenholz (62 for tanh activation. They noted that tanh-like MIP networks with skip connections "hover over the edge of chaos". We quantify their observation as follows.
Theorem 4.3. Let (x) be a deep MLP network with erf iiniitatio and residual connections of strength μ = 1. Then in the limit N₁ + 00
The NNGP kernel Kt linearly diverges with depth 1.
X's approaches 1 frrm above (Fig. 2 X'y = + ē/Vl, where ē = 20 (T oo +o is a non-universal constant. Consequently, APJN diverges as stretched exponential flo.l = O(e O(eV))))))))))) where À = 1/(42)) ) is the new length scale.
We will refer to this case as subcritical. Although X3 reaches 1, the APJN still diverges with depth faster than any power law. The growth is controlled by the new scale A. To control the gradient would like to make À large, which can accomplished by decreasing ow. In this case, the trainability enhanced (see Fig. 3. Similar results hold for tanh activation function [2] however in that case there is no explicit expression for ē.
Residual Connections + LayerNorm
In practice, it common to use a combination of residual connections and LayerNorm. Using (.). for NNGP and partial Jacobians are modified as follows
the recurrence relations 99 [1]
N; N
N
Xl+1
-o
Σ
+o + 22 X3
=
2 Kt
Eo
1
o' (RL))
+ +2². (22)
N
1=1
k=
8

---

Remark 5.1. For μ < 1, (22) implies that the fixed point value of NNGP kernel is scaled by 1 - μ². Moreover, residual connections do not shift the phase boundary. The interference between residual connections and LayerNorm brings X'y closer to on the entire ds-Ow plane (as can be seen from Fig. 2). Therefore the correlation length improved in both the phases, allowing for training of deeper networks. At criticality, Jacobians linearly diverge with depth.
As was mentioned before, the combination of LayerNorm and residual connections dramatically enhances correlation length, leading to more stable architecture, This observation is formalized by The proof leverages the solutions of (2) close the fixed point, and is fleshed out in Appendix
1.3
Remark 5.2. When μ = 1, the correlation length diverges for any initialization.
5.2
provides an alternative perspective architecture design. On the one hand, (16) to initialize a given network at criticality. Alternativcly, one can use a combination of residual connections and LayerNorm to ensure that the network will train well, irrespective of the initialization.
Remark 5.3. When μ = 1, the condition X holds on the entire σ, - aw and for any activation function á (see Fig. 2}. NNGP kernel diverges linearly, while APJN diverges algebraically with the critical exponent of Ç O(1). The exact value of the critical exponent depends on the activation function and the ratio σb, Tw- The trainability is dramatically enhanced, as shown in Fig.
Remark 5.4. Networks with BatchNorm [26], used in conjunction with residual connection of strength (See Appendix B).
1 1, also enjoy this everywhere criticality and enhanced trainability 63 23
Modern Architectures
ResNet110 ResNet is one the most widely used architectures in computer vision 21] Figure shows the results for ResNetV2 [22]; with BatchNorm replaced with LayerNorm. (See Appendix for BatchNorm results and discussions.)
1 o2 - of phase diagram shows everywhere cerrwhee--ritticllttyy as expected from our theory. For of = 0, o la phase diagram gets closer to criticality as we increase which results in better training higher μ..
Additionally, networks with larger j enjoy better trainability due to the suppression of finite width corrections to our theory (decreased effective L/N).
-2.
23-20.0]-00
145
0524
335
0134
022
27
Figure 4: ResNet100LayerNomm)) Left to right: (1) X² phase diagram w.r.t. (o 0 at 1 1. The network is everywhere-critical in this case. (2) X"y phase diagram w.r.t. (o ud) (3) Training accuracy w.r.t. o² for different values of j. Trainability improves with with higher μt, the network gets closer to everywhere-critice phase. For μ = 0 and small o² the network in the ordered phase, and the output every block is almost zero this explains the poor trainability. (4) Training accuracy w.r.t. at (o o.) = (2, 0). We see monotonic improvement in trainability.
MLP- Mixerr MLP-Mixer architecture is a recent example of MLP approach to computer vision
can be analyzed using the tools presented above. Fig. 5hows the phase and training diagrams. Further details can be found in the SM. Note that the original architecture uses l 1.0.
48
Test accuracies for figures
4 follow similar trends train accuracies. (See Appendices B

---

2-20
25-
7. L3
575
5-
50
10 54
24
3
125 200
Figure 5: MLP-Mixer. Left to right: (1)(2) jr 0.5, 1.0 phase diagrams (x3). The network becomes everywhere-critee critical when jr = 1. The solid black line indicates the empirical phase boundary. Stars denote points we selected train on CIFAR-10. (3)(4) μr = 0.5, 1.0 MLP-Mixer training curves. All, but one, networks are L = 100 blocks deep. We see that as we increase from 0.5 to 1, the trainability of all networks increases, and they are less sensitive to initialization.
7 Conclusions
We have introduced partial Jacobians and their averaged norms as tools to analyze the propagation of gradients through deep neural networks at initialization. Using APJN evaluated close to the output, J1--21.-1 x X we have introduced a very cheap and simple empirical test for criticality. We have also shown that criticality formulated in terms of partial Jacobians is equivalent to criticality studied previously in literature [44 45 34 Additionally, APJN can be utilized to quantify criticality in
inhomogeneous (i.e. no periodic stacking of blocks) networks
23
We have investigated homogeneous architectures that include fully-connected layers, normalization layers and residual connections. In the limit of infinite width, we showed that (i) in the presence of LayerNorm, the critical point generally becomes a critical line, making the initialization problem much easier, (ii) LayerNorm applied to preactivations enhances correlation length leading improved trainability, (iii) combination of 1 residual connections and erf activation function enhances correlation length driving the network
suberitical phase with APJN growing according stretched exponential law, (iv) combination of residual connections and LayerNorm drastically increases correlation length leading to improved trainability, (v) when μ = and LayerNorm applied to preactivations the network is critical on the entire T w plane.
We have considered examples of modern architectures: ResNet and MLP-Mixer. We showed that at 14 1. they are critical everywhere, due to the interaction between LayerNorm and residual connections. We have studied ResNet at different p's, showing that higher μ enjoys better trainability, due to improved initialization and suppressed effective L/N corrections.
We have empirically demonstrated that deep (100 blocks) MLP- with μ 1 trains well for various initializations. In comparison, for = 0.5, only trains well close to the critical line.
Our work shows that an architecture can be designed to have large correlation length leading to guaranteed trainability with SGD for any initialization scheme.
Limitations and Future Works
While our method is empirically applicable to Transformer architectures (see Appendix C for phase diagrams), there remains a notable absence of robust theoretical foundations for initializing Transformers. We aim to extend our theory to the attention mechanism and determine the optimal way initialize Transformer architectures.
Another challenge arises from network topology. For graph neural networks, the generalization of methods remains ambiguous. We aspire to elucidate this in future research.
Lastly, we want to scale our experiments to larger image datasets as well as language tasks, with the hope that our methods could help people train large models more efficiently.
10

---

Acknowledgments and Disclosure of Funding
We thank 'T. Can, G. Gur-Ari, B. Hanin, D. Roberts and S. Yaida for comments the manuscript. A.G.'s work at the University of Maryland was supported in part by NSF CAREER Award DMR- 2045181, Sloan Foundation and the Laboratory for Physical Sciences through the Condensed Matter Theory Center.
References
[1] Kylc Aitken and Guy Gur-Ari. On the asymptotics of wide networks with polynomial activations. arXiv preprint arXiv:2006. 06687, 2020.
[2] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pages 242-252. PMIR, 2019.
[3] Anders Andreassen and Ethan Dyer. Asymptoties of wide convolutional neural networks. arXiv preprint arXiv:2008. 08675, 2020.
[4] Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets a compression approach. In International Conference on Machine Learning, pages 234-663. PMLR, 2018
[5] Sanjecv Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, pages 8141-8150, 2019.
[6] Sanjeev Arora, Simon S Du, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang, and Dingli Yu. Harnessing the power of infinitely wide deep nets on small-data tasks. International Conference on Learning Representations, 2019.
[7] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 660.064400 2016.
[8] Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley. Rezero is all you need: fast convergence at large depth. In Cassio de Campos and Marloes H. Maathuis, editors, Proceedingy of the Thinr-Seventh Conference on Uncertainty in Artificial Intelligncce volume 161 of Proceedings of Machine Learning Research, pages 1352-1361. PMIR, 27-30 Jul 2021. URL.
.ma.iinggsmm.....mmm.mmm.......m mlr, .ww..ppes///ee http///achleehhner2/ html
[9] Tankut Can, Kamesh Krishnamurthy, and David J Schwab. Gating creates slow modes and controls phase-space complexity grus and Istms. In Mathematical and Scientific Machine Learning, pages 476-511. PMLR, 2020.
[10] John Cardy. Scaling and renormalization in statistical physics, volume 5. Cambridge university press. 1996.
[11] Minmin Chen, Jeffrey Pennington, and Samuel Schoenholz. Dynamical isometry and a mean field theory of runs: Gating enables signal propagation in recurrent neural networks. In International Conference on Machine Learning, pages 873-882, PMLR, 2018
[12] Soham De and Sam Smith. Batch normalization biases residual blocks towards the identity func- tion in deep networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems. volume 33, pages 19964-19975. Curran Asso- ciates, Inc.. 2020. URL htw.88accabe7722baaascccccaooaee Papeee .ddf
neurips. meurips.cc/paper falerffilllepppae httppe/fillepepae////// /fiee/
[13] Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn. Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL http:pss//ppeeeevve..nttttttt id-YicbFaNTTy
[14] Ethan Dyer and Guy Gur-Ari, Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations, 2019.
[15] Rainer Engelken. Gradient flossing: Improving gradient descent through dynamic control of jacobians. In Thirty-sevetth Conference on Neural Information Processing Systems, 2023. URL htthttt:::////// httpss://oeenreview net/forumd iarrorummmidjdjeeeeee JQQoo

---

[16] Rainer Engelken, Fred Wolf. and L. F. Abbott. Lyapunov spectra of chaotic recurrent neural networks.
Phys. Rev. Rev. 5:043044, Oct 2023. doi: 10.1 1003/hhy.eeeesssacch.......3333
URL https://link
aps. hhw..e/e.....././hheeeeeeeeer.. 5. 043044
[17] Adrià Garriga-Alos -looso, Carl Edward Rasmussen, and Laurence Aitchison. Deep convolutional networks as shallow gaussian processes, arXiv preprint arXiv: 1808.05587, 2018.
[18] Mario Geiger, Arthur Jacot, Stefano Spigler. Franck Gabriel, Levent Sagun, Stéphane d' Ascooli Giulio Biroli, Clément Hongler, and Matthieu Wyart. Scaling description of generalization with number of parameters in deep learning. Journal of Statistical Mechanics: Theory and Experiment, 2020(2)023401, 2020.
[19] Boris Hanin. Random neural networks in the infinite width limit gaussian processes. arXiv preprint 1oXiv:2007.01562 2021.
[20] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. Proceedings of the IEEE international conference on computer vision, pages 1026 1034, 2015.
[21] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770-778, 2016. doi: 10.1 LO9//CPR.2016....
[22] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In Computer Vision ECCV 2016, pages 630-645, 2016.
[23] Tianyu He, Darshil Doshi, and Andrey Gromov. Autoinit: Automatic initialization via jacobian tuning, 2022.
[24] Judy Hoffman, Daniel A. Roherts, and Sho Yaida. Robust learning with jacobian regularization. arXiv preprint arXiv: 02729, 2019.
[25] Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nugp and ntk for deep attention networks. In International Conference on Machine Learning, pages 4376-4386. PMLR, 2020.
[26] Sergey Toffe and Christian Szegedy. Batch normalization: Accelcrating deep network training by reducing internal covariate shift. In International conference machine learning, pages 448-456 PMIR, 2015.
[27] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint /1t.::88000777777 2018.
[28] Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.
[29] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, (6(12)) 2:222222324, 1998.
[30] Jachoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dicksttin. Deep neural networks as gaussian processes. In International Conference on Learning Representations, 2018.
[31] Jachoon Lee, Lechao Xiao, Samuel Schocnholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32:8572-8583, 2019.
[32] Jaehoon Lee, Samuel S Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: empirical study. 2020.
[33] Aitor Lewkowycz and Guy Gur-Ari. On the training dynamics of deep networks with / regularization. arXiv preprint arXiv:2006.0 08643, 2020.
[34] James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszez, Valentin Dalibard, Jascha Sohl- Dickstein, and Samuel Schoenholz. Rapid training deep neural networks without skip connections or normalization layers using deep kernel shaping. arXiv preprint
2021.
[35] Alexander G de G Matthews, Jiri Hron, Mark Rowland, Richard E Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. In International Conference on Learning Representations, 2018.
12

---

[36] Lutz. Molgedey. J Schuchhardt, and Heinz G Schuster. Suppressing chaos in neural networks by noise. Physical review letters, 69((6):3717, 1992.
[37] Radford M Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks, pages 29-53. Springer, 1996.
[38] Roman Novak, Yasaman Bahri, Daniel A Abolañia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Sensi- tivity and generalization neural networks: an empirical study. In International Conference on Learning Representations, 2018.
[39] Roman Novak, Lechao Xiao, Yasaman Bahri, Jaehoon Lee. Greg Yang. Hron, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-dickstein. Bayesian deep convolutional networks with many channels are gaussian processes. In International Conference on Learning Representations, 2018.
[40] Roman Novak, Lechao Xiao, Jiri Hron, Jachoon Lee, Alexander A Alemi, Jascha Sohl-Dickstein, and Samuel S Schoenholz. Neural tangents: Fast and casy infinite neural networks in python. In International Conference on Learning Representations, 2019.
[41] Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Neural tangents: Fast and casy infinite neural networks in python. In International Conference on Learning Representations, 2020.
[42] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf. Edward Yang. Zachary Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bus Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 8024 8035, Curran Associates, Inc., 2019.
[43] Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. The emergence of spectral universality in deep networks. In Amos Storkey and Fernando Perez-Cruz, editors, Proceedings of the Twenty-Fiie First International Confenence on Artificial Intelligence and Statistics, volume 84 of Proceedings of Machine Learning Research, pages 1924-1932. PMLR, 09-11 Apr 2018. URL https://proceeddii press/v84 1/////44///e//n htt.///nnnnnntoneeee. html
mlr
[44] Ben Poolc,
Subhaneil Lahiri, Maithra Raghu, Jascha bohh-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. Advances neural information processing systems, 29:3360-3368, 2016.
[45] Daniel A. Roberts, Sho Yaida, and Boris Hanin. The Principles of Deep Learning Theory: An Effective Theory Approach to Understanding Neural Networks Cambridge University Press, 2022.
[46] Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein Deep information propagation. 2016.
[47] Vaishaal Shankar, Alex Fang, Wenshuo Guo, Sara Fridovich-Keil, Jonathan Ragan-Kelley, Ludwig Schmidt, and Benjamin Recht, Neural kernels without tangents. International Conference on Machine Learning. pages 8614-8623. PMLR, 2020.
[48] Ilya Tolstikhin, Neil Houlshy, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dasovitskiy. Mlp-mixer: An all-mlp architecture for vision. arXiv preprint urXiv:2105.0 006,0 2021.
[49] Russell Tsuchida, Tim Pearce, Chris van der Heide, Fred Roosta, and Marcus Gallagher. Avoiding kernel fixed points: Computing with elu and gelu infinite networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 9967-9977, 2021.
[50] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, & ukasz Kaiser, and Illia Polosukhin, Attention is all you need. In I. Guyon, U. Von Luxburg. S. Bengio, H. Wallach, R. Fergus, Vishwanathan, and R. Garnett, editors, Advances Neural Information Processing Systems. volume 30. Curran Associates, Inc., 2017. URL http::tttp://occesee (aFeai/e/afeee22255eeeeeeeeeaaa
heerrip..c/pppae happp/ppee___eeeee_eeeeepppepeee files/ pddf
[51] Christopher Williams. Computing with infinite networks. In M. C. Mozer, M. Jordan, and T. Petsche, editors, Advances in Neural Information Processing Systems, volume 9, MIT Press, 1997.
[52] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference computer vision (ECCV), pages 39, 2018
13

---

[53] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist. novel image dataset for benchmarking machine learning algorithms, 2017.
[54] Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington. Dy- namical isometry and field theory of cans: How to train 10,000-layer vanilla convolutional neural networks. International Conference on Machine Learning, pages 5393-5402. PMLR, 2018.
[55] Lechao Xiao, Jeffrey Pennington, and Sam Schoenholz. Disentangling trainability and generalization in deep learning. 2020.
[56] Ruibin Xiong, Yunchang Yang, Di He. Kai Zheng. Shuxin Zheng. Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu, layer normalization in the transformer architecture. International Conference on Machine Learning, pages 10$24-10533, PMLR, 2020.
[57] Jingjing Xu, Sun, Zhiyuan Zhang. Guangxiang Zhao, and Junyang Lin. Understanding improving layer normalization. arXiv preprint urXiv:9911.07013, 2019.
[58] Sho Yaida. Non-gaussian processes and neural networks finite widths, Mathematical and Scientific Machine Learning, pages 165-192. PMLR, 2020.
[59] Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior. gradient independence, and neural tangent kernel derivation. arXiv preprint 19Xiv:1902.04766 2019.
[60] Greg Yang. Wide feedforward or recurrent neural networks any architecture are gaussian processes. 2019.
[61] Greg Yang and Edward Hu. Tensor programs iv: Feature learning in anfinite-width neural networks. In International Conference on Machine Learning, pages 11727-11737. PMLR, 2021.
[62] Greg Yang and Samuel Schoenholz. Mean field residual networks: On the edge of chaos. arXiv preprint 17Xiv:11/.008699, 2017.
[63] Greg Yang. Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, Samuel S Schoenholz. A mean field theory of batch normalization. In International Conference on Learning Representations, 2018.
[64] Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, and Shuicheng Yan. Metaformer is actually what you need for vision. arXiv preprint arXiv:21 1418, 2021.
[65] Hongyi Zhang, Yann Dauphin, and Tengyu Ma. Residual learning without normalization better initialization. In International Conference on Learning Representations, 2019. URL https //openreview.net.u nett/forrumiicHIIyse23tt
14

---

A Experimental Details
We implemented our methods using PyTorch algorithm (24].
42 hooks and an efficient Jacobian approximate
Figure I We generated MNIST-like inputs, where all elements are sampled from the Gaussian distribution (0, 1). Joo data was averaged over 100 different aurameter-iiitiaizzate
Networks
were initialized width N₁ 1000. For erf plot we initialized at critical point (a σb) (V O), used depth L = 250 and the fitting was done with data points collected at depth > 100; for ReLU plot we initialized at critical point (aw:; 0b) = 2 0), used depth L = 100 and the fitting was done with all data points: for the μ Pre-LN plot, we initialized both networks at (a 0b) = (v2, 0), used depth L 250 and the fitting was done with / 100 data points.
Figure 2 All the phase diagrams were plotted using XJ 1-1 generated from networks with L 50 and N₁- 500. We used hooks to obtain the gradients that go into calculating X X5-1 data was averaged over 100 different
inaamettr-iitialization Inputs were generated from normal Gaussian distribution and have dimension 28 28. Generating the data for the figure took approximately days on Google Colab Pro (single Tesla P100 GPU).
Figure 3 In all cases, networks are trained for 10 epochs using stochastic gradient descent with CrossEntropy loss. We used the Fashion MNIST dataset [53] All networks had depth 7, = 50 and width Ni- 500. The learning rates were logarithmically sampled within (10- 5, 1). Generating the data for the figure took approximately 12 days on Google Colab Pro (single Tesla P100 GPU).
Figure (1)We made the o² of phase diagram for ResNetl 1a 10(LayeNNorm) by averaging over 100 different aurameter-initaalizatio
The o 1 phase diagram was made by averaging over 200 parameters initialization. (3) We used SGD with momentum= 0.9 and batch size 128. For selecting the learning rate we ran a grid-search over 0.001, 0.005, 0.01, 0.02, 0.5 for 10 epochs; with weight decay À = 10-1 All models were trained for 50 epochs and averaged over 3 random sceds. It takes 6 GPU days in total
single NVIDIA RTX 3090 GPU.
Figure 3 (e) made the phase diagram for MLP-Mixer with 30 blocks and averaged over 100 different iisaamtter-nitializatioi
(3) used network with L 100, patch size 4 x 4, hidden size = 128, two MLP dimensions NE. Nen = 256. The I, = 32 point has doubled widths. All networks have 10 million parameters. Notice that for all Mixer Layers we used NTK initialization. We trained all cases on CIFAR-10 10 dataset using vanilla SGD paired with CSE. Batch size = 256, weight decay x= = 10- was selected from {10-5 10 11-1). mixup rate = 0.8 was selected from (0.4, 0.8). We also used RandAgument and horizontal flip with default settings PyTorch. For all cases we searched learning rates within (0.005, 0.01, 0.05, 0.1, 0.2, 0.5). We also tried linear warm-up schedule for first 3000 iterations, but we did not see any improvement in performance. Generating the data for the figure took approximately days on Google Colab Pro (single Tesla P100 GPU).
B Additional Discussion on ResNet, ResNet with BatchNorm
For Convolution Layers, the NNGP kernel is a 4-index tensor: KL Ka (x,,), where the Greek letters (n,, v) index the channels, whereas the Latin letters (i,j) index the pixels. The infinite width limit in this case is achieved by taking the number of channels to infinity (sequentially). In this limit, most of our equations for MLP can be easily rewritten using the convolutional NNGP kernel. However, in this case, the kernel is only diagonal in channel dimension: Kewwa (zr, x') = Ki, (2, )0) This additional structure in the kernel makes it difficult to get a closed-form solution for J1,+1 in general.
ResNet1 10 (LayerNorm) In Figure 2), the networks is critical close to μr 1, as expected from our analysis. One would naively expect the μ < cases also to be critical, since for MLP with ReLU and Pre-LN, = 0 is critical regardless of o, and p. However, in Figure 2 the region away from J = in ordered phase. This is likely a result spatial dimensions. We emphasize that the μ = I case stays unaffected by this, since the existence of criticality does not depend on the details of the NNGP kernel in this case. This can be readily seen from 71. We present the Numerical and training results in Figure
4
the kernel Ki (x. x') not being diagonal in
15

---

ResNet110 (BatchNorm) be described as follows:
The operation of BatchNorm on preactivation (pre-BN) in an MLP can
hx(x)= hx(x) PK.B
(23)
P4,
Σ h(a')
and Ji,
Σ r'EB
(h(x²) -1,,n)
B
B
dEEB
where B is the batch that x belongs to and [B] is batch size.
The works Yang al. 663 He et al. 23 show that for large batch size, the effect of BatchNorm for NNGP kernel and Jacobian Norm is deterministic. We summarize the results for the pre-BN MLP setup:
o K² (x,x')
N Eg
((( + p2
(24)
Ni
glut1 (x
KJ xx) Ea
1²
(25)
where x' in the APJN term can be any x' + since for large batch size all choices are equivalent.
From the above result, we can see that most results we have for LayerNorm can be translated to BatchNorm with an easy replacement (x æ) + (x' x)) - - (z, z²)). As simple example, consider a pre-BN ResNet architecture, but with all the Convolutional layers replaced with Linear (Fully Connected) layers. For such a network, we have the following result for < 1:
guut1 >² 12
+
(26)
(x 1)2
For = 1, we have
JuH4 1
(27)
The ResNet results can then be obtained by replacing Fully Connected layers with Convolution layers, in a similar way as discussed in ResNet(LayerNorm) section. We show the numerical results for ResNet with BatchNorm in Figure
SwI0
Figure 6: ResNetl (ONNee110(BtccNormm)) Left to right: (1)(2) JL-1.Z phase diagrams, with (a of) and
(o u). (3)(4) Training curves w.r.t o² and μ.
C Discussions on Transformers
As preliminary empirical results, we applied our method to Vision Transformers (ViTs). Figure supports argument in the main text that the Pre-LN architecture is generally insensitive to initialization when μ = 1, where Post-LN is only stable with small o2. Figure further demonstrates that the attention mechanism is highly sensitive to initialization and thus vulnerable to the gradient vanishing/diverging problem.
16

---

Preck, with MIP
(Pockt NJ with MUP
60
C7
97-
65
36
11-
0.
65
35
62
32
Figure 7: XJ Phase diagram for Vision Transformer (ViT) (L = 24, Theada 8, demba = 256). We present cases with pre/post-LN, with/without MLP layers. Top-left: Usual ViT becomes everywhere-critical when 1; ordered otherwise. Bottom-left: Usual ViT without MLP layers. Top-right: VIT with post-LN (instead of pre-LN). Always ordered for finite dw. Bottom-right -right VIT with post-LN, without MLP layers.
67
E7
66
66
65
65
25
Figure log X.7 Phase diagram for stacked Attention layers with/without MLP layers (L 24, Nheads 8, demba 256). We use log X and cut off the colorbar at X 10 because the X3 values this case span from 50 orders of magnitude.
D Technical details for Jacobians and LayerNorm
We will drop the dependence of (z) on throughout the Appendices. It should not cause any confusion since we are always considering a single input.
D.I NNGP Kernel
First, we derive the recurrence relation for the NNGP kernel Pq. As mentioned in the main text, weights and biases are initialized (independently) from standard normal distribution (0,, 2 /fan_jnnn We then have
2 dam 85n and Eg 1/ = 05
Eng (uot whal =
(28)
N
17

---

by definition.
We would like prove theorem 2.4 found in 45
consequence of lemma 2.2
The proof of lemma 2.2 be
Proof of theorem 2.4
One can prove this by definition with lemma 2.2
N1+1
1
KC+1
Eglh
Ni+1
N141
N₂
No
1
Eg
Σ
w ( 6,43
Σ
with
o(hh)) (ch))))hh))))))))b))) +
N141
1=1
N141
1
E.
# 1+1 f ((h))d((h))) + b
N1+1
N1
1
Ee
Σ
(h (1))) + +o?
N1+1
N
i=1
2
Σ
Eg [+(H)))(h))]] + oo
(29)
N
==1
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 (APIN)
(1 > 1o)
Ni+1 Ni
ah!t1
0h!+1
Joo.
Ea
N141
on
Ni+1 No
Ni
an!+1 anl and o
Ni
an!t1 an ah
Es
Σ
N411
Dhk
Nit No
N
ahk Th! ahly ah"
Eg
(wit 1+1 0' (((hb)
(w 1+1 o' (k(w))
N4+1
j=l k,m=1
Ni+1 Nig
Nr
ohk alin Th to an
Eg
1+1 w ø' (h)))' (Hh
N411
Witk
1=1 j=1 k,m=1
No Nip
ahk anl anly anly
Eg
8' ('hh)o'(h)))
N1+1
N
1=1
Nin
]
ahk ahk
Eg
o'(h))o(bb)
(30)
ah Th
18

---

where we used the chain rule and took the expectation value over w Next, we take the chain rule again:
Ni
Nig
N₁
dl on
glo.41
Ea
o (1
w w w (l (0((()))))))
Ni
ahi
m,n==
Nig No-1
8 p((1rm1 )6' (h 1 0h 1
Egg
((
Σ
Σ
O(1/Ni-1)
NiN-1
Th
0h
Nig M
o (hm )o' ( ( 0h !
m
Eg (o' (d(h)))'(1))) Ee
Σ
N
ah
f=1
Ni
tee Eg Ni-1
Nio No-1
@ dhm!!
Eg [d'(h))
Σ Σ
Φ p(hhm
( (0'(m()))
Ni
ah ohly
mmm
X Thod
(31)
where we integrate (by parts) over w to get the second line. We take Ni-1 and then N; + oo to get the third line. We rearrange terms and use the (q. 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 is independent of hh-1 Thus we proved
glalt1= = x'o
(32)
The critical line is defined by requiring X = 1, where critical points are reached by further requiring X² 1.
As we mentioned in the main text, lg = 0 subtle since the input dimension fixed No. which can not be assumed to be infinity. Even though for dataset like MNIST, usually No is not significantly smaller than width N₁. We show how to take finite O(N correction into account by using one example.
Lemma D.1. Consider one hidden layer network with a finite input dimension No. the infinite width limit, the APJN is still deterministic and the first step of the recurrence relation is modified to:
20
No
Jo.2
X
X 2 No
JO.
(33)
No
where J0,1 2
19

---

Proof.
N₂ No
ahe ah? Jh 0h
Jo.2
N2
1 B
No
thg Sh!
Σ
W im o' (1))) (H)))
22
ahy ang
f=1
k,m=1
No
No
Σ
Esl Bol w why wiy (12) ((h))
22
k,m=1
No
Eg w wky o' (hk) (1k))
Ni
20
No
X
XA
No
k
20
No
X
X
Jo.1
(34)
No
No
where get the result we used integrate by parts, then explicitly integrated over w/y- We have introduced a coefficient of finite width corrections, X'a defined as follows.
Definition D.2 (Coefficient of Finite Width Corrections).
X
Eold (1)) (M.) + ""' ("!()))((
(35)
N
Remark D.3. Notice that the correction to J is order O(N If one calculates the recurrence relation for deeper layers, the correction to Jot will O OE N7 which means the contribution from hidden layers can be ignored in infinite width limit.
The J0,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 (APIN with = -). The APJN (with lo of a given network can be written as
20
Jod 22
XJ
No
XA
Σ
No
II X
(36)
1=2
Note that APJN with Io > does not receive the O correction.
D.3 APJN and gradients
As mentioned in the main text, APJN is 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) Le norm of the gradient of any flattened
20

---

parameter matrix 0', at initialization. In the infinite width limit, one gets
ac ohf
oh 0h 00%
2
Th ah
th
all
0c ac
oh 0h4
O
Σ
Th oh
0h 1-1
ahi
Thim
36%
OC 7c
Oh
oh
Ohlt
Oh!
Σ
8ymm
ahy ane
9h 1
on
Dhn
agi
ac Oc
1
014
Oh
Σ
OE Jl++1
on on
Oh 0h
0h 4-1
ag
all
F
OC Oc
ah 0h
0h
E
ورو
th oh!
ah 1h1-1 5
1-11
00
OC Oc
Oh
8i J 1.
gL++1
th! ah
00
0C
2
Oh'
-L-1L
Jl.1+1
11L
a0
OC
Oh
J..
(37)
ONL
a0e
F
where ||-1|2 denotes the L2 norm and ||-||æ-
denotes the Frobenius norm.
D.4 LayerNorm on Pre-activations
Definition D.5 (Layer Normalization).
E(a'] E[ E 7/2 + 8
(38)
where of and ( are learnable parameters.
Remark D.6. With only LayerNorm, the is simplified to
1+1
Σ +
(39)
Remark D.7. the limit of infinite width, using the law of large numbers, the average over neurons E [ can be replaced by the average of arrameter-iittializaioon Eg [...]. Additionally, in this limit, the preactivations are i.i.d. Gaussian distributed: h (O,
E
Th' = E3 hh = =0
(40)
E
2 (h)) = E. (n')
(hd)
(41)
The normalized preactivation then simplifies to the form of Eq. 19)).
Remark D.8. At initialization, the parameters 7f and /8| take the values and 0, respectively. This leads to the form in equation (19). In infinite width limit, has the following form
h Eg[h']
(42)
Ea ( EN
Lemma D.9. With LayerNorm on preactivations, the gaussian average is modified
Eg [o(it)] = 2 / (h!) --
(43)
21

---

Proof. By definition K is sampled from a standard normal distribution (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
K1+1
Es [0()))(((5)]
+
(44)
Proof.
N111
1C+1
1
Σ
Ep
h
1+
N141
i=1
11
N₂
Nr.
1
Σ
w ((h) ( + 0,41
Σ
w +062)+++++)
Ed
N1+1
i=1
[o(R))d(it)]] +
Ee
(45)
Theorem D.11. In the infinite width limit the recurrence relation for partial Jacobian with LayerNorm on preactivations is
gla.t1 = xzJl
(46)
where X²3
a N.K
No Eim R.g
(1((1)y).
Proof.
Jlo.l11
1
Eg
Ni+1 No
ahlt ah!t! 1+1
N1+1
oh ah
1
NiL1 Ni
an!+1
a anl
an!+1 oi ah!
EE
N1+1
ait. anl on
oihn The oh
i=1
1
Ed
Not1
Nip
N
1+1 o (((W)) Wikk
1
(wsttl o' (((hmm)
Ohk an! an an
N4+1
Vxcl
VxCl
1=1
1=1
k,m=1
a NYC
Ni
Noo
ohk onk Thly only
Σ
Es
d'(h)jo'(h
kml
22
Ni
Σ
Eg
6 (
flad
NiCC
k=1
glo.
(47)
D.5 LayerNorm on activations
The general definition of LayerNorm on activations is given as follows. Definition D.12 (LayerNorm on Activations).
o(h)) - E((((')
(h!)
VE(o(H))] - Elp(h!)12 oo + B/.
(48)
22

---

Remark D.13. The recurrence relation for preactivations (Eq. gets modified to
Σ
w 1+1 ((h,) b
(49)
Remark D.14. At initialization, the parameters 7f and & take the values 1 and 0, respectively. This leads to the form
((h) B[o(h']
o(h))
v (1 - B[o(h')} ((h)) - Eo [o(h']]
(50)
Ea (o(h)))) Ee (o(h1)]
where the first line follows from the fact that initialization, the parameters 7 and Bf take the values 1 and O 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 / N(0, o + of) Again this can be summarized as
Es [O(h)]]
dhl O(h))
(51)
V + o
Next, calculate the modifications to the recurrence relations for the NNGP kernel and Jacobians. Theorem D.16. In the infinite width limit the recurrence relation for the NNGP kernel with Layer- Norm on activations
K++1 oz +0
(52)
Proof.
Nos1
K
Eg hh
N141
Nr
Eg
(11)) ( 1,41
Σ
wott ((())( 20()) +
N1+1
N:
(b(())))
o
N
1=1
((h)) Es [o(h']]
Σ
+ +0b
VEO Eo ( 19] Eo [o(N)))
N
Es
(o(h)) - Eo [)(1))))2 EA (o(h))²] Ep [ + o3
N
a. + +oa
(53)
Theorem D.17. In the infinite width limit the recurrence relation for partial Jacobian with LayerNorm on activations is
glo.t1 glo.
(54)
Es (
where X = 20
23

---

Proof.
M441 No
Thit
oh !+1
glo.t1
Ey
Nit1
Th
ohly
M41 Nio
Ni
ohlt1
onle ahl o
Ni
Thlt1 om Ohin only
Σ mm=1
Eg
N141
j=l
M411 Nio
ahl ah'm only anly
Eg
wikh (1()))
(with (1()))
Ni+1
1=1 k.m=1
N2
Mig
Ni. ahk onk
Es
d'(h))
Ni
an an
N2
((h)) 27
Ee
Flod
N1
Eo (0(h))2] E. ( Tlod
Es
[h')3]
(55)
E Residual Connections
Definition E.1. We define residual connections by the modified the recurrence relation for preactiva-
tions (Eq. 11)
N₁ o(hh)
h
> j=l
w i+1
b1+1 + phl,
(56)
where the parameter j controls the strength of the residual connection.
Remark E.2. Note that this definition requires N411 = N₁. We ensure this by only adding residual connections to the hidden layers, which are of the same width. More generally, introduce tensor parameter Pij.
Remark E.3. In general, the parameter μ could be layer-dependent (4'). 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
2² N
kd+1
Es [o(h))o(h)) + of + w
(57)
24

---

Proof.
Ni+1
Rug [h+ 1
N1+1
Ne+1
N
++ (())) + 1 pul Σ wht (1/1)) + + mur))
Σ
N1+1
/=1
No+1
N1
N2
Eg
Wij
w f o(h))o(())) + 05+ 0 + ps hhhh
N111
j=1
Ni+1
N₁
N+1
Eg
Th
Σ
(())))()))) + o
12²
1
Σ im1
Eo [n/n!]]
N141
Ni
Ni+1
Eg
[o(h))
of + 12K1
(58)
N
where we used the fact N141 Ni 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
gla.t1 = X'y
(59)
where the recurrence coefficient shifted to X'y Bo (0(())0 + 222.
Proof.
Nip1 Nig
o an!+!
Jlo.411
Eg
N1+1
1
No+1 Nip
Ni
anl+1 ah
anlt on ahl anly
Eg
Σ
N1+1
ahd
No+1 Nip
N
@ahk oh anly to an!
}
Eg
(wit d' d((1)) + udox) (with o' (bmmm) + udim)
N1+1
1=1 j=1 k,m=1
No+i Nig
Ohk. ahi Th ah!y
En
(w
1+1 1+1 witt wim ob (B)m)))))))))) + 20
N++1
k,m=1
d
N;
Eg
Nin
Thk onl T th
+
N
Ea
12
Σ
No
onl and ahly only
N
N
km
Am
j=1
Nr Nip
ahk ank
Eo (1() (Mb)) + (2?) Eo
Ni
ah"y
ohl
Eo (o'(N') (1/))d' (M4) + + glo.l
flod+1= = flod
(60)
25

---

Residual Connections with LayerNorm on Preactivations (Pre-LN)
We recall the recurrence relation
1+1 h
Σ
++ w
o(h)) + 6,t1 + wh!
(61)
Theorem F.1. In the infinite width limit, the recurrence relation for the NNGP kernel is then modified
to
k²+1
Eo (0()))(()))
oo + 12k
(62)
Ni
Proof.
Ni+1
1
K041
Eg hh
(+ 11
N1+1
1
Nus1
Σ
Ee
[
(ld)) (li)
F +
6,tt + und
Σ
No
with ((h)) a(h)))((h))))))))))) 6,t + wid)
N1+1
02
N;
En [(((5))(()]]
+ oz + 1!K!.
(63)
N1
j=1
Remark F.2. For μ < 1, the recursion relation has fixed point
IC
o²
Nia Eo
(hl )o(k)))
}]
of 1 122
(64)
Nr. 122)
where the average here is exactly the same as cases for LayerNorm applied to preactivations without residue connections. 1* labels some very large depth
Remark F.3. For 1 case, the solution of (62) is
a N Eg
KC = KCO
(0
+ of
(65)
N
which is linearly growing since the expectation does not depend depth. KC 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.
glo.t1 = glot
(66)
where for this case
Ni Es
d'( ((t)) 122
(67)
NIC
26

---

Proof.
Nii1 Nio
t ah!
Flodt1
Eg
Ni+1
Thly
i=1 j=1
Nit1 Nig
Nr
anlt oit an
Ni
anlt aiin al'm Shi ahin only
Σ m==
Ee
N1+1
oht ahl o
i=
N441 Nig
Nr
with o' (iit)
w// of d'(kim)) Vxc
only an!n (
]
Eg
+ pdas)
+ uddi)
Ni+1
VKy
anly only
1=1 1=1
k,m==
t MiK
Ni
Ni
ont onl
}
Eg
o'(k))o + 122
Σ
ohly oh
k=
N₂
Eo (0 (hA) + n² glod
NiKC
k=1
(68)
Remark F.5. Assuming/ 1 2 1", combine 64
and
67
to derive Equation 55 in the Theorem 1.3
Eg
o' (1)) ((
X3 = (1 - 11²²)
+ 212
N₁
Ep
@b ((h))
+ +o
= (1 - n
2²
(69)
B
where A = Ek= Eg o (h2)) (hhh) and B
Σ j= Es (o()))(())) + oz.. Recall that
Ni
& = |log>> This directly leads us to Equation
5
{ = llog (xt)
(70)
|log ((1 μ²) 1/4 n2²)
Note that in Equation ( we discard the average over neurons in A and B, since we are in the infinite width limit.
Remark F.6. As mentioned above μ needs extra care. Plugging μ into the result 65
and 67 we find out that
2 Σ k=
Eg
D' (1 )o (74.)
X ||=1
1
NiKo
Ei
0
Σn 1=
Ed
((h)) + Nio3)
110 (
(71)
Thus, we get everywhere criticality: critical behavior independent of the choice initialization ow)) It also follows that the correlation length & diverges this case. (7I) leads to power law behaviour in Jacobians (with exponent () at large depth. Note that the exponent is not universal.
Remarks
E5 and F.6 together serve as the complete proof of Theorem 1.3
G Critical Exponents
To prove theorem 2.7
first need to find the critical exponent of the NNGP kernel 45
27

---

Lemma G.1. In the infinite width limit, consider critically initialized network with an activation
function G. The scaling behavior of the fluctuation åk K'' K in non-exponential. If the recurrence relation can be expand to leading order ok as 5k =o 5IC - cn(k)) for n 2 2. The solution of 8K! is
1
8K'
1 1-5e
(72)
On (n
where (x: n-I
Remark G.2. The constant On 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 5K' = K KC* on both sides.
8K z0 ōk - (,11)))"
(73)
Use power law ansatz ö1C = 5k then
(1+ = 1 Sx - CnAl nix
(74)
Multiply [&x on both side then use Taylor expansion
( Sc ss S2
SK
Cn (n-1))
(75)
1
For arbitrary 1, the only non-trivial solution of the equation above is
A 1
1 n
(n 1) and Gx
(76)
Proof of theorem 2.7 We will assume 0o + Then lemma G.I 8K! To leading order 1-1
we can expand X in terms of
xy R0 d, djok!
=1 -1- d1
1-
(77)
02
Consider a sufficiently large I. In this case O(1-1) approximation is valid. We write recurrence relations of Jacobians
1-1
Jlod
II
-
d1
1-1
gloid
C2
P==lo
Rd Cgg
(78)
When 6, = 0 for all n 2 2, from lemma constant.
G.1 we have oK = 0. Thus the Jacobian saturates to some
We checked the scaling empirically by plotting Jo. vs. in log-log plot and fitting the slope. These results are presented in Fig
H MLP-Mixer
In this section we would like to analyze architecture called MLP-Mixer 48), which is based on multi-layer perceptrons (MLPs). 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.
28

---

Suppose one has a single input with dimension (Cin. Han Win). We label it as Tui. 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 Then, the first convolution layer can be written
ΣΣ w j=
pers V..+ + bp
(79)
where ff is the size of filter and is the stride. our example f = s. Notice in PyTorch both bias and weights are sampled from a uniform distribution 21( -V, v k). where k: = (2in
Ee [W pure W po:yy 1
3C Supdvodis.
(80)
Eo Eo 3Cin Snoiss-
(81)
Notice that the output of Conv2d: bo E RCxN,, where stands for channels and N, Hin Wine/ f2 stands for patches, both of them will be mixed later by Mixer layers.
Next, we stack Mixer Layers. Mixer Layer contains LayerNorms and two MLPs, where the first one mixed patches i.j (token mixing) with hidden dimension Ntm. the second one mixed channels jr, v (channel-mixing) with a hidden dimension Nem- Notice that for Mixer Layers we use the standard parameterization.
First LayerNorm. It acts channels J..
h1 Eo
0 pi
V Vare 62) pi.
(82)
where we defined channel mean Ee [ha Eo hot and channel variance Vare
Ee (15)) 21 (h)))
First MlpBlock. mixes patches j,j, preactivations from different channels that share the same weight and bias.
67 + 1: Linear Affine Layer.
N
1,6+1 pÿ
+ 6gtt1 jk
66+1
(83)
67 + 2: Affine Layer.
Now
61+2
61+2
(1
61+1)
+ 001+2
+
(84)
j=1
where Nem stands for the hidden dimension of "token mixing". 6/ + 3: Residual Connections.
h 6L+3
1,6422 + jch
(85)
Second LayerNorm. It again acts on channels μ..
1 6643 u1
Ec: (1 61+3
R 66443 pi
(86)
Varg 61+33
Second MlpBlock, mixes channels μ, v, preactivations from different patches share the same weight and bias.
29

---

6l + 4: Linear Affine Layer.
h hutt1
W
up 11
61+ 66+4 6L+3
+ 16144
(87)
67 + 5. Affine Layer.
Nem
6645 je
w 6l45 (hh + bp 101+5 pu
(88)
6/ + 6. Residual Connections.
1 61+5 js
G1+6 jsi
+ jah
GL++3 js
(89)
Suppose the network has Mixer layers. After those layers the network has a pre-head LayerNorm layer, a global average pooling layer, and output layer. The pre-head LayerNorm normalizes over channels can be described as the following
69 HI
(90)
Varc (1
Global Average Pool over patches i.
(91)
Np
Output Layer
fy
Σ Upuhyf + by
(92)
We plotted the phase diagram using the following quantity from repeating Mixer Layers:
N,
C N
Jh ah oh
ah oh-
X's
lim L-+00
Es
(93)
NOC
==
I Results for Scale Invariant Activation Functions
Definition I.1 (Scale invariant activation functions).
d(x) a +(x) a_ x 0(-z),
(94)
where ((r) is the Heaviside step function. ReLU is the special case with a4 and a 0.
L.1 NNGP Kernel
First evaluate the average using lemma 2.2
Es (((h))(()))
dhl (a+, + (1))
2RT
0 + a² Kc
(95)
Thus we obtain the recurrence relation for the NNGP kernel with scale invariant activation function.
0 (a +
a
k++1
K + o
(96)
30

---

Finite fixed point of the recurrence relation above exists only if
d (0_ + a
X
< 1.
(97)
2
As result
ou S 2
(98)
a
+ a
For 2 a= case, finite fixed point exists only if of = 0.
L.2 Jacobian(s)
The calculation is quite straightforward, by definition
X'y =o l [d
/
-upa 2x
2 (a + a²)
(99)
2
where we used the property ar)(x)) for Dirac's delta function to get the first line.
Thus the critical line is defined by
2 a
ww
o
(100)
For ReLU with a4 = and a_ = 0, the network is at critical line when
Tw = 22,
(101)
where the critical point is located at
(ow, σb) (V2.0).
(102)
I.3 Critical Exponents
Since the recurrence relations for the NNGP kernel and Jacobians are linear. Then from lemma G. and theorem 2.7
ÇK = and = 0.
(103)
L.4 LayerNorm on Pre-activations
Use lemma D.9 and combine all known results for scale-invariant functions
I.9maD..9mm
X
Eg
(o(()))'((()))) (hbb)
NiK
R6-1
2 (0 +a2)
(104)
2 (a + +a²) + 20
In this case,
X <
(105)
always true. The equality only holds at 06 line.
31

---

1,5 LayerNorm on Activations
First we substitute K'-1 o + of into known results
af
+ +
Eg (d(( (k))
(106)
at + a²
Eg [o(M))(()]]
(o + og.
(107)
2
There is new expectation value we need to show explicitly
Egg [cb(n]]]
di(p(h)) 2
(2 + t
cnt)
dhi(a+ - a-)iEe
( + +0%
2 + og
(a+ -a-)
(108)
2π
Thus
Taa + a²)
X3
2 + og
Taa a²) (a+ )3
(109)
The critical line is defined by X = 1, which can be solved as
(a+ a_)
a6 V ( + +a²) (a+ a ow
(110)
For ReLU with a4 = 1 and a = 0
σ. ow
00.6830
(111)
1.6 Residual Connections
The recurrence relation for the NNGP kernel can be evaluated to be
0 (a
a²)
K
K +
oh + 1 K!!
(112)
The condition for the existence of fixed point
0 (af
+ a2)
X
11² <
(113)
2
leads us to
0 2(1 122)
(114)
a + a²
For σ3 221 2(1-#²) ai +o?
finite fixed point exists only of 0. (Diverges linearly otherwise)
The recurrence coefficient for Jacobian is evaluated to be
2 (a² + a2
X
(115)
2
The critical line is defined as
2(1 n22)
T 1
a + a²
(116)
The critical point located
0). For ReLU, the critical point is at 112)1 12%,.0)
(v 2( 20111²² 1²)
22 a=+a= +
32

---

1.7 Residual Connections with LayerNorm on Preactivations (Pre-LN)
Again use lemma D.9 and combine all known results for scale-invariant functions
Xy lim 10000
Eg
of d'(k))
u²
Ni
R2-1
o (a + d²) - 122)
+ 11² o ( + aL3) + 20
20 (1 )
(117)
a + a²=) + 20
Similar to the case without residue connections
X 1
(118)
always true. The equality only holds at ob = line for μ < 1.
Notice there is a very special case = 1, where the whole σb aw plane is critical.
J Results for erf Activation Function
Definition J.1 (erf activation function).
(xx) V 1" - dt.
(119)
J.1 NNGP Kernel
To evaluate lemma
2.2 exactly, we introduce two dummy variables A and Xa 51
I dd2
de
Eg [d(Ar
= ddi
d'idA2 Eo [+(N11))((2ab))
/ dd1 I ddz2 / dh!
3 (h!)² e (AP+ (1E)
/ dd1 I
4/C
ddi (1 + 21 (A7 + X3))
2 aresin
2K A1
(120)
1 2K + 22)
We the special where A
A2 1.
Thus the recurrence relation for the NNGP kernel with erf activation function is
202 aresin
2X
k++1
oj
(121)
I
2K
As in the scale-invariant case, finite fixed point only exists when
4o2
1
X
π (1 21 VI+ $1.
(122)
Numerical results show the condition is satisfied everywhere in σb - ow plane, where X² is only possible when K 0.
33

---

J.2 Jacobians
Follow the definition
Xy o Bo
/ dii -2(h1)
(05) 2K
V
40
(123)
π
v1 1
To find phase boundary Xy K
1, we need to combine Eq.
121 and Eq
123 and evaluate them at
20 in
2K
K
aresin
+ 0
(124)
1
2K
4
1
XJ
1.
(125)
V1 4/C
One can solve the equations above and find the critical line
160 n²
20
160 π² 16 20 + π2
aresin
(126)
4x²
Critical point is reached by further requiring X² 1. Since X² S Xy. the only possible case is K O, which is located at
(owi Tb) (VV 7.
(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 (as, o.) (0, VF, with K Now suppose large enough such
that the deviation of XC! from fixed point value X is small. Define òIC' = K' 1 Eq. [2T can be rewritten as
28K
8K = aresin 1 28K' 2
(128)
roK 2(0x!
From lemma G.I
1 A= and (x 1.
(129)
Next analyze critical exponent of Jacobians by expanding 123 around K 0 critical point (0%. aw) = (0, /
To leading order 1-1 we have
X x1 28K
(130)
~1 7
Thus the recurrence relation for partial Jacobian, at large l, takes form
glo.411 = (1 (1-1) zlo..
(131)
34

---

At large/ I
gloot ct 1-1
(132)
with non-universal constant Qu-
The critical exponent
1=1,
(133)
which is the same as ÇX-
J.4
LayerNorm on Pre-activations
Use lemma D.9 we have
Ni NIK
Xy
Σ Eg
('(h)))' (ile)
18-1=1
2 m
(134)
v5
20 aresin (2) + ToO
The critical line is then defined by
σ
aresin
L5
0.324o
(135)
J.S LayerNorm on Activations
Due to the symmetry of erf activation function Egg [o(ht)] results.
0, we only need to modify our known
4
1
Es (((()))((()))
(136)
/1 + 4(0 + of
2( + o
Es [0(1))(()))
aresin
(137)
1 2 + +
Thus
20 m
X = V1+ + 4(0 + +
(138)
2 +o 1+2(0%++0))
aresin
where the phase boundary is defined by the transcendental equation x 1.
J.6 Residual Connections
The recurrence relation for the NNGP kernel can be evaluated to be
20 m
2K
aresin
+ + + p²x
(139)
I 2K
Finite fixed point only exists when
2
X
+ + <
(140)
π (1
2K V 4K
Notice that X² I X²y still holds, where the equality holds only when K = 0.
The recurrence coefficient for Jacobian is evaluated to be
4
X
* + 112
(141)
π
VI
35

---

The critical line is defined as
160 T2 22
20
160 x² 22)2 160 e + ² - 11223
aresin
(142)
4m 1222
n
Critical point is reached by further requiring X KC* = O, which is located at
Since X S Xyy, the only possible case is
T( 112) 0
(o σb)
(143)
Note that for 1, needs to put extra effort into analyzing the scaling behavior. First we notice that K'' monotonically increases with depth - the recurrence relation for the NNGP kernel at large / (or large XC) is
K+ o o + +
(144)
which regulates the first term in 141
For μ = large depth
4
(145)
Co + 4 + of
Here Co is constant that depends the input.
We approximate the asymptotic form of log Floo follows
log Joo. log
IIx X3
I'=lo
4
log
1+
π Co + 4 + oz
I'=lo
40
dl' log
1
l0
Co 1 + +og)l'
2čvl + Ollog),
(146)
2n
where č
/ +og
We conclude that at large depth, the APJN for 1, erf networks can be written as
glod ~ 0 (2ov(+(()))))
(147)
This result checks out empirically, as shown in Figure 9
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.
oZ
No
X lim
Σ
Eg
(h (1t)
10000
NiK
Re-1=1
46 (1
112)
+ 2²²
(148)
V5 220 aresin
3
The critical line is then defined by
2 aresin
L5
3
(149)
0.324o
'We used NTK parameterization for this experiment. However, we emphasize that it does not affect the final result.
36

---

1020
Theory, mnooy.t-1am 2-1424
Fiting,
Experiment, 2.000
Theery, E
1034
Fitting 0-81
Experiment, ag-
1011
Theery, 0.5.64
Fiting,
10a
10a
10
7.5
15.0
/7
v for Figure 9: log(T
1, 0f = 0, erf.
K
Results for GELU Activation Function
Definition K.1 (GELU activation function).
x
x
Φ(x)
+
erf
2
V2
x
2
-t
1
+
6
dt
(150)
2
0
K.1 NNGP Kernel
Use lemma
2.2 for GELU
dn! (11)²
h1
Eg [o()))(h))]
+ erf
21c1
V2
v22
[ an!
(n1)2
111
eet²
ert2
e
4
22
K!
h
(A
dh's (nt)) eer
4
V32nKi
2
K!
KJ
(hL)²
dl erf2
2R!
V32TK!
(IC)
h
h!
h
h/
dhl
erf
erf'
erf
erf"
V32mm
/2
V2
K
K!
K
2K
arcsin
+
(151)
2m
KC
(1
11+K3)11
21
where from the third line to the fourth line we used integrate by parts twice, and to get the last line used results from erf activations.
Thus the recurrence relation for the NNGP kernel is
K'
K
K!
(K))2
kt+1
aresin
+
02 + +
(152)
2x
1+
1
T(1 UI1+KVVVII + 2
As result
K
4(K*)) + 11
X
+
aresin
(153)
4
2
1
(1
37

---

K.2 Jacobians
Follow the definition
Xy = Eo [ (h)))'
2
1
h
e
hg
dh4
+ erf
+
2x
2
2
2
/2)
22T
dh4
+
1
erf
h
2
M
1/2
11
e (066))
(h.)²
04
2m
4
4
V2T
2T
0
K
+ (3+55
aresin
(154)
2T
1
+1+ (1 21
where we dropped odd function terms get the third line, and to get the last line 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 K
and X3
0 +
KC*
KC
1
K
aresin
1 + +
(155)
2T
1+ K+
T (1+ 2K
ow
2
KC*
K (3++)
aresin
1+K=-
=1.
(156)
4
2x
(1+ X++ 2
Cancel the aresin term, One and σr then can be written as a function K
(3+
2
K"
ow 2 1
arcsin
(157)
x (1++))
1 1+*
K
ab 21T + 2k*) Ow
(158)
One can then scan K to draw the critical line.
In order to locate the critical point, we further require X²: 1. To locate the critical point, we solve X X 0 instead. We have
2 -( -KC* 0,
(159)
2m + KK 21C-)
which has non-negative solutions out of three
3 + V17
KC and K
(160)
2
One can then solve on and au by plugging corresponding K* values.
(ow, σb) (), for 1 =0,
(161)
(aw, 0b) ~ 1..408004100 for K 3 V17
(162)
2
K.3
Critical Exponents
GELU behaves in different way compared to erf. First discuss the K 0 critical point, which is located at (ar, ow) = (0, 2). We expand (52), and keep next to leading order òK! = K. K
skt+1 8K' 6 (8x)))2
(163)
π
38

---

From lemma G.1
A π and Ck = 1,
(164)
which is not possible since 81C 2 0 for this case. This result means scaling analysis is not working here.
Next, consider the other fixed point with K 3+ (ob; Tw) (0.416, 1.408). Expand the
NNGP kernel recurrence relation again.
8K'- ~ N 0.00014(00x))
(165)
Following the same analysis, find
OK' R -7142.9-
(166)
Looks like scaling analysis works for this case, since K > O. The solution shows that the critical point half-stable(45). If K' KC* the fixed point is repealing, while when AC' K* the fixed point is attractive. However, the extremely large coefficient in the scaling behavior of öAC embarrasses the analysis. Since for any network with a reasonable depth, the deviation 8K is not small.
Now we can expand x'y at some large depth, up to leading order 1-1
66.668
Xy ~ 1
(167)
1
Then
8 2 Cio 06.668
(168)
where % a positive non-universal constant.
Critical exponent
< = 66.668
(169)
Which in practice is not traceable.
K.4
LayerNorm on Pre-activations
Use lemma
D.9
we have
N
X
G N₁K
E. (((
k=1
R1-1=1
a2 (6π + 4v3)
(170)
0 (6T 3v3) + 18T0
The critical line then at
= (6V3T) 3π
00.175
(171)
K.S LayerNorm on Activations
First, need to evaluate new expectation value
dh 1/ 2
/2 )
od
Eo [((ht]]
1 + erf
2
v 2T (o + of) 22 + of
(172)
V 1m + oy + +
where we used integration by parts to get the result.
39

---

The other integrals are modified to Eo (
1
o + of
+o o
aresin
1 d + t
(1 o1 +of)11 2(0 + of)
(173)
Ea [(()))(())]]
ob
2
+
ob
aresin
o + of 1 + to6
1 o
2x
T + d + 0
+ 2 to
(174)
One can then combine those results find x,y
T (1 + a + oo,)
+ 2 aresin
2( 10
1+0
π + +z oo + to,) -2 t0 + V 11+2
+
2( + oz) + t + + aresin ao
[+0g toj
(175)
The critical line defined by x = 1, one can numerically solve it by scanning over σg and ow
K.6 Residual Connections
The recurrence relation for the NNGP kernel is
K
K'
K
(KC)2 T1 +(1+ + 2K'
K
4
+
2x
aresin
1+
2 + 0 + 12,1
(176)
Fixed point exists if
2
K
(10*+ 11004)) 55C (1+ 210-)
X
aresin
+
+ 11
(177)
2T
1+
The recurrence coefficient for Jacobian is
02
a
K
(3*3)
aresin
1+K*
+
+
(178)
4
2x
(1+13)(1+2K+))
Phase boundary is shifted
23
2
K
)
= 2
2
1+
+ arcsin
(179)
T(1 + +(1+K=) 2k+
1 1+1*
K
2 + 2/C+) ww.
(180)
One can again scan over 1 to draw the critical line.
In order to locate the critical point, we further require X 1. To locate the critical point, we solve X5 X = O instead. We have
a
-XC* 0,
(181)
2x + (1 210
which has two non-negative solutions out of three
KC and K 3 V17
(182)
2
One can then solve σb and aw by plugging corresponding K* values.
( ob) = (2v1 (2,0), for K x==0 =0,
(183)
(am, σb) x (1.408 12200.16 42), for K* 3+ VI7
(184)
40

---

K.7
Residual Connections with LayerNorm on Preactivations (Pre-LN)
Use lemma D.9 and the results we had without residue connections for GELU.
Ni
XJ
lim 1000
0 w
Eg
08 ((
+ u²
NiK
k=1
Re-1=1
o (6 31 2
+ 112 d (6T + 3v + 18T0
(v30 -18T0 (1 112)
(185)
8 (6T + +3 + 18T0
The critical line is then at
96 (6v3T)
(186)
00.175
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.
(8)
experiment Theory
Theary
Fitting C 0.007
Fitting, (= 0.009
5x10
3
4x
2x
10
3x.10-
10
10r
Figure 10: log
log plot of partial Jacobian Jool vs. for (A) LN-erf and (B) erf-LN.
In figure I1 we tested 6k samples from CIFAR-10 dataset [8 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 of with LayerNorm applied to preactivations, where the prefactor 28. The prefactor we had is the same as vanilla cases [55] The difference is from the fact that they used logio and used log.-
we explore the broad range o of the performance of MLP network with erf activation In figure 12
function and LayerNorm preativations. The network has depth L. = 50 and width N, = 500; and trained using SGD on Fashion MNIST. The learning rates are chosen based on logarithmic scan with short training time.
41

---

LN-
LN-er
IN-GFLU
201
222
23C
413
415
20500
20/00
211
23
20
6.5
2.0
2.5
3.5
6.0
95
1.9
2.6
2.5
3.5
6.5
05
1.0
2.0
2.5
35
4.5
Figure 11: Test accuracy for LayerNorm applied to preactivations. of = 0.5 for all cases, Correlation lengths are calculated using analytical results of xzy.
LN-Erf, of = 0.1
80
Training Accuracy (%)
40
10a
101
10P
10a
10"
10a
Figure 12: Training performance MLP networks with erf activation function; and LayerNorm applied to preactivations. It continues to train for several orders of magnitude of o (with learning- rate tuning).
42


Reviews:
[{"review_id": "gaKhtDjhEK", "reviewer": "Reviewer_RvQ9", "summary": "The paper presents a theoretical framework for understanding the trainability of deep neural networks with LayerNorm and residual connections. The authors derive analytical expressions for the neural network Gaussian process (NNGP) kernel and the partial Jacobian norm (PJN) for a wide range of activation functions. They show that the combination of LayerNorm and residual connections leads to an everywhere-critical regime, where the network can be trained effectively irrespective of the initialization. The authors also provide insights into the role of the hyperparameters and the activation function in determining the trainability of the network. The paper's contributions include a theoretical understanding of the trainability of deep neural networks with LayerNorm and residual connections and insights into the design of effective architectures.", "questions": "See the weaknesses part.", "limitations": "See the weaknesses part.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The main strengths of the paper are:\n1. Introduces partial Jacobians and their averaged norms as tools to analyze the propagation of gradients through deep neural networks at initialization.\n2. Presents a very cheap and simple empirical test for criticality using APJN evaluated close to the output.\n3. Shows that criticality formulated in terms of partial Jacobians is equivalent to criticality studied previously in the literature.\n4. Investigates homogeneous architectures that include fully-connected layers, normalization layers, and residual connections, and shows that the combination of LayerNorm and residual connections can drastically increase correlation length leading to improved trainability.\n5. Considers examples of modern architectures, ResNet and MLP-Mixer, and shows that they are critical everywhere at µ = 1 (i.e., with residual connections) due to the interaction between LayerNorm and residual connections.\n6. Empirically demonstrates that deep MLP-Mixer with µ = 1 trains well for various initializations.", "weaknesses": "Some potential limitations of the paper could include:\n1. The analysis is limited to homogeneous architectures with fully-connected layers, normalization layers, and residual connections. The results may not generalize to other types of architectures or layers (e.g., attention layers).\n2. The paper focuses on the infinite width limit, which may not be directly applicable to finite-width networks commonly used in practice.\n3. The empirical test for criticality based on the averaged partial Jacobian norm may not be sufficient to fully capture the behavior of deep neural networks (in terms of accuracy for instance).\n4. The paper does not provide a comprehensive comparison with other methods for improving the trainability of deep neural networks, such as weight initialization schemes or adaptive optimization algorithms.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback.\n\n## Infinite width limit\n\nFormally, our theoretical results become crisp when $L/N \\sim o(1)$ (width $N$, depth $L$). In practice, our conclusions apply whenever this ratio is small. For instance, all the phase diagrams in figures 2,4,5 are obtained on real, finite width, deep networks (MLP: N=500, L=50; ResNet and MLP-Mixer: N=512, L=32). One can readily see the agreement between finite-width phase diagrams and infinite-width predictions (dashed lines).\n\nMoreover, in practice, very deep architectures typically use residual connections. Residual connections make networks effectively shallow, leading to an effectively small ratio $L/N$.\n\nWe direct the reviewer to \"Global Response\" for more details.\n\n## Inhomogeneous architectures\n\nThe utility of Averaged Partial Jacobian Norm (APJN) can be readily generalized to inhomogeneous architectures. These architectures can contain differently initialized fully-connected layers, convolutional layers, normalization layers, residual connections, and even attention layers. In these case, due to inhomogeneity, the fixed point analysis and scaling discussion no longer applies. However, it is possible to impose criticality with a stricter requirement: $\\mathcal J^{l,l+1}=1$ for all layers. This requirement fixes the hyperparameters layer-by-layer, which can be implemented using our code. The additional coputational cost required in this case is no more than other methods for critical initializations. Though more computational power is required in this case, we view this as a price one must pay no matter which methods one uses.\n\n## Attention\n\nWe present empirical phase diagrams for the Vision Transformer (ViT) with various settings, in the accompanying on-page pdf. Our empirical methods correctly measure the gradient scale, including the fact that pre-LN transformers are better behaved than post-LN transformers. \n\nOn the theoretical side, the analysis is more involved. Attention layers can be viewed as fully-connected layers with data-dependent weights. It is known that the commonly used scalar-dot product attention layer converges to a Gaussian process in the limit where the number of heads and the embedding dimension go to infinity [1]. For most real-word models $L / n_{heads}$ is $O(1)$, which makes infinite width analysis unfavorable even with residual connections. Furthermore, softmax function in the attention prevents getting closed-form results. Despite the lack of proper theoretical understanding, our tools can remarkably analyse criticality for such networks (See attached one-page PDF).\n\n## Accuracy\n\nThe objective of our methods is not to predict the best performance based on initialization. In general, it is hard to isolate the effect of initialization from factors like dataset, optimization, regularization etc. Our goal is to provide a universal framework for initializing a network that facilitates good training performance. Such a framework is invaluable, for instance, in Neural Architecture Search (NAS); where it is important to distinguish the effect of architecture from initialization. Our experiments show that APJN manages to achieve this in a diverse set of cases.\n\n## Comparison\n\nWe will add more related works and explain the difference and connections between known methods and ours. Here we just want to summarize the key point briefly:\n\n- Popular initialization schemes, for example, He initialization, is a special case of our method. These schemes apply to specific cases (e.g. He init is critical for ReLU networks). \n\n- References to prior methods that utilize the infinite width limit to fix initializations (e.g. $\\sigma_w^2$) are included in lines 146-150. \n\n- Methods focusing on scaling, including Fix-up[2], ReZero[3], and other works, those methods can also be summarized as initialize with $\\mathcal J^{l, l+1} = 1$.\n\n- Good initialization is essential even with adaptive methods and should be used in conjunction. For Adam, at early training time, the gradient update is bounded by the learning rate $\\eta$ [4][5]. At layer $l$, the difference between the feature-updates of two adjacent input vectors $x$ and $x+\\epsilon$ is upper bounded by $O(\\eta \\cdot \\epsilon \\cdot \\mathcal J^{0, l})$ (up to some architecture/loss dependent constant). In this case, $\\mathcal J^{0, l}$ regulates the relative updates for different features, preventing them from being too large or too small, facilitating better training.\n\n[1] J. Hron, et al., Infinite attention: NNGP and NTK for deep attention networks, 2020.\n\n[2] H. Zhang, Y. N. Dauphin, T. Ma, Fixup Initialization: Residual Learning Without Normalization, 2020.\n\n[3] T. Bachlechner, et al., ReZero is All You Need: Fast Convergence at Large Depth, 2020.\n\n[4] D. P. Kingma, J. Ba, Adam: A Method for Stochastic Optimization, 2017.\n\n[5] J, Ma, D. Yarats, On the Adequacy of Untuned Warmup for Adaptive Optimization, 2021.\n"}, {"review_id": "P4WwaGFPi3", "reviewer": "Reviewer_EkqX", "summary": "The paper studies the effect of the expected value of the Jacobian norm of a particular layer with respect to a previous layer as the depth of a neural network (NNs) increases. The study is done under the assumption of infinite width NNs, and with the goal of assessing sensitivity to initialisation hyperparameters (the standard deviation of the normal initialiser of the weight and biases of the network) depending on architecture and as depth increases. The importance of understanding this setting comes from the connection with exploding and vanishing gradients that the network is likely to exhibit in training, unless the initialisation (in expectation) is in the critical region. The conditions for criticality across architectures are studied in this work.\n\nThe settings studied include feed-forward networks, residual networks with and without layer-norm applied to the pre-activation, as well as group norm and briefly Batch Norm in conjunction with residual networks.\n\nReasons for rating: While the contribution of the paper is worthwhile, the clarity could be greatly improved. At the moment, the paper resembles more a series of facts rather than a clear logical deduction. I worry that this will limit the value it can bring to the ML community.\n", "questions": "Suggestions for improvement:\n * Please see above in terms of weaknesses.\n  * For Thm 2.4 (and similarly in the Appendix proofs), why not write the recurrence relationship as K^{(l+1)}(x, x) = \\sigma_w^2 K^{(l)}(x, x) +  \\sigma_b^2 ?\n  * There are many sentences in the paper that require additional citations, such as line 213.\n * Figure 1b shows a lot of noise as the number of layers increases, can the authors comment on that? \n* The authors should provide additional details on how the figures are obtained, at the moment these are scant. \n", "limitations": "I believe that the biggest limitation of the paper as is now is its presentation, which can be drastically improved. This applies both to the writing and the details around the figures.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Strengths:\n * The paper provides a theoretical assessment of why, in the infinite width limit, certain neural networks can achieve better trainability. In particular, the authors study the effect of residual connections and layer norm, both which have become a staple of deep learning. This type of study can help explain why certain architectures are less sensitive to initialisation compared to others.\n * Code reproducing the figures is provided in the supplementary material as notebooks (and a clear readme.txt detailing the results).\n * Proofs in the attached SM are generally readable and can be followed.\n", "weaknesses": "Weaknesses:\n* The clarity of the paper can be drastically improved. \n  * In the paper, the same recipe of proof is applied repeatedly in multiple settings. I think the paper would benefit from clarifying that recipe in the main text, and working through the logical steps one by one in one example. At the moment, there is no clear explanation of the logical flow of what is occurring. My understanding of that recipe is:  In the setting studied (for that specific architecture), in the infinite width limit find a recurrence relationship between the NNGP kernels at layer l+1 and l when evaluated at the same datapoint. Use that recurrence relationship, together with the recurrence relationship for APJN to find the conditions of criticality for initialisations. I urge the authors to clarify this recipe (in more detail) in the main manuscript. as well as work through an example.\n * There are many sentences for which it is unclear where they come from, either from a proof in the Appendix or from a previous work or the authors consider it trivial. Such an example is Eq 4, which potentially follows from Thm 1.3, but the order of the two is unclear. Similarly, the proof of Thm 1.3 in the Appendix is not well delimited or clarified (given that it is the main result). \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback and detailed suggestions on the presentation.\n\n## Clarity \n\nWe believe that the logical flow of the current manuscript is similar to what the reviewer has proposed. We state the outline of our paper here -- we will add a version of this in the Introduction of the final version:\n\nIn section 2 we introduce the recurrence relation for NNGP kernel and Jacobians with fully connected networks. Based on these quantities we define criticality. In Section 3 and Section 4, we add LayerNorm and residual connections, describing their effects on criticality. Then in Section 5, we combine LayerNorm and residual connections to show the emergence of everywhere-criticaly. Finally, we use our methods on real-world architectures like ResNet and MLP-Mixer in Section 6.\n\n## Proof of Eq 4\n\nWe use Section 1.1 to summarize our theoretical results. Eq. 4 can be viewed as the definition of correlation length $\\xi$. Theorem 1.3 then utilizes this definition and states the effect of LayerNorm and residual connections on correlation length. \nAn additional explanation for Eq. 4 is provided in Section 2.1. In the final version, we will add a reference to Section 2.1 as a motivation for Eq 4.\n\n## Proof of Thm 1.3\n\nTheorem 1.3 directly follows from Eq. 77 (Appendix F), upon taking infinite width and large depth limits. (and taking the logarithm, as mentioned in the text below Eq 77). We will flesh out these steps more explicitly in the final version.\n\n## Form of recursion equations:\n\nWe offer multiple responses to 2 possible interpretations of the reviewer's question:\n\n(i) If the reviewer is suggesting the specific form: The form that the reviewer mentioned, $K^{(l+1)}(x, x) = \\sigma_w^2 K^{(l)}(x, x) + \\sigma_b^2$, only holds for linear networks. For other cases $K^{(l+1)}(x, x)$ is a non-linear map of $K^{(l)}(x, x)$, which depends on the activation functions. \n\n(ii) If the reviewer means to ask why we keep the summation and the factors of $1 / N_l$ in the equations: It is possible to introduce some extra symbols to simplify the equations notationally. However, we believe that such additional notation will come at the price of clarity and lucidity. The current way of writing formulae makes the underlying calculations apparent. This should make the equations accessible to a broader audience.\n\n## Additional references:\n\nWe thank the reviewer for pointing out places where more references are potentially warranted. We will scan the text and add the relevant references.\n\n## Noise in Figure 1(b)\n\nThe apparent fluctuations in Figure 1(b) come from two factors:\n\n(i) The Y-scale of Figure 1(b) is more zoomed in compared to other plots, making the fluctuations appear strong. We chose this scale because the underlying curve of $\\mathcal J^{0, l}$ is a constant -- without increase or decrease in the trend, fluctuation offers the only natural choice for the scale of the y-axis.\n\n(ii) The existence of the fluctuations is a result of the depth-to-width ratio ($L/N$). $L/N = 250/1000 = 0.25$ in this case. This results in strong fluctuations even after averaging.\n\n## Details of figures\n\nThe experimental details of all the figures, along with computational resources for reproducing them, are fleshed out in Appendix A (due to space limitations). In the final version, we will expand the figure-captions by adding further experimental details."}, {"review_id": "3jjnnv27mn", "reviewer": "Reviewer_8MUF", "summary": "The paper addresses the theoretical treatment of deep neural networks and introduces a novel practical approach to identify criticality within these networks. The authors work in the setting where the number of parameters per layer approaches infinity, enabling the formulation of quantitatively predictive descriptions ( establish criteria for hyperparameter selection). These criteria are based on the notion of criticality.¨\n\nTo identify criticality, the paper introduces partial Jacobians, which represent derivatives of preactivations in layer l with respect to preactivations in layer l0 (where l0 ≤ l). Recurrence relations for the norms of these partial Jacobians are derived and utilized to analyze criticality in deep fully connected neural networks featuring LayerNorm and/or residual connections.\n\nThe authors devise a straightforward and cost-effective numerical test to determine the optimal initialization for various types of deep neural networks, including fully connected, convolutional, and normalization layers. They present quantitative evidence demonstrating that arranging LayerNorm (applied to preactivations) and residual connections appropriately leads to an architecture that exhibits criticality regardless of the initialization.\n\nFinally, the paper applies these methods to investigate the ResNet and MLP-Mixer architectures, revealing the presence of an everywhere-critical regime within these modern models.", "questions": "Dear authors, I have a few questions regarding your work that would greatly aid my comprehension:\n\n* Can you provide a rationale for selecting the architectures used in your study? Are there any other potential architectures that could have been considered within the framework?\n\n* In addition to LayerNorm and residual connection what other technics used in the field could have been explored in your framework ?\n\nThank you in advance.", "limitations": "To enhance the paper, it would be valuable to provide a clear outline of the limitations of the work. Some potential limitations to consider include: \n\nThe paper focuses on testing a specific set of realistic architectures, which means there is room for exploring other architectures beyond those examined in the study. Although the obtained results are encouraging, extending the analysis to different architectures would provide a more comprehensive understanding of the findings.\n\nAdditionally, the paper focuses on investigating the effects of LayerNorm and residual connections. While these factors are certainly important, exploring the impact of other components or techniques apart from LayerNorm and recurrent layers could broaden the scope of the research. Examining different elements in the network architecture could reveal additional insights into their contributions and interactions.\n\nAnother important limitation of the paper is that it specifically focuses on the infinite width limit of deep neural networks. While analyzing neural networks in the infinite width limit can provide valuable theoretical insights, it may not fully capture the behavior and characteristics of networks with finite widths. To address this limitation, it would be valuable to extend the analysis to include investigations and experiments on networks with various finite widths. This would provide a more realistic and comprehensive understanding of the behavior, performance, and criticality of deep neural networks in practical settings.\n\nIncluding a note on the computational complexity required to conduct these experiments would be beneficial. Acknowledging the computational demands of the experiments would provide a better understanding of the resources needed for replication or further research.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This research represents a significant advancement in the field. A notable contribution of the paper is the introduction of partial Jacobians and their averaged norms as powerful tools for analyzing gradient propagation in deep neural networks at initialization. The paper particularly investigates the implications of LayerNorm and residual connections, shedding light on their impact on trainability.\n It is encouraging that the theoretical prediction derived from this Framework match previously empirically observed or theoretically proven findings. Additionally, the paper strengthens its contributions by testing the findings on more realistic datasets, thereby enhancing the robustness and applicability of the research.\n\nThe paper is highly comprehensive and lucid, greatly facilitating understanding. It effectively summarizes the contributions, allowing readers to grasp the main points without constant reference to the appendix. Furthermore, the authors have shared the relevant code with their submission, ensuring transparency and enabling the replication of results. The inclusion of detailed step-by-step derivations significantly aids comprehension and is highly appreciated by readers.\n", "weaknesses": "These suggestions are intended to improve the overall clarity and accessibility of the research.\n\nTo enhance the organization of the paper, it would be beneficial to include dedicated sections for limitations, future work and contributions. Additionally, the current title of the section labeled \"results\" may be misleading and should be reconsidered. Furthermore, FashonMNIST is not mentioned in the main text.\n\nThe limitations outlined below highlight some of the weaknesses of the paper, despite the fact that they represent the current best efforts in the field.\n\nAlthough it may seem like a minor detail, I believe it holds significant importance. For the sake of accessibility, completeness, and readability, I would suggest including relevant lemmas or theorems used from other sources in the appendix to provide a more comprehensive understanding of the research.  Furthermore, It would be helpful if you could extend slightly the captions in the figures for the camera-ready version for the same reason.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their encouraging feedback and suggestions for improvement.\n\n## Presentation and Clarity\n\nWe welcome the suggestions to improve the presentation of our paper:\n- We will include a discussion on limitations and future work in Section 7 (Conclusion).\n\n- We will change the Section title to \"Main Result\", as it states \"everywhere-criticality\", a key contribution of our work.\n\n- We mention FashionMNIST in the caption for Figure 3. \n\n- We will extend the figure captions (especially Figure 3) in the final version of the paper. (See Global Response for further discussion.)\n\n## Answers to Questions:\n\n1. The applicability of *numerical* Averaged Patial Jacobian Norm (APJN) is quite general. The choice of modern architectures like ResNet and MLP-Mixer was made to demonstrate its utility in SOTA models. It is possible to extend this method to models with Attention layers (eg. Transformers). We find that APJN correctly predicts the everywhere-criticality of pre-LN Transformer. For further details, we refer the reviewer to the General Response and the accompanying one-page pdf.\n\n2. As we mention in the paper, the analysis also extends to other normalization techniques such as BatchNorm and GroupNorm. Dropout can also be included. In general, one can divide the network into \"blocks\" of layers and study APJN wrt outputs of these blocks. Since these blocks can contain arbitrary layers/techniques within them, any/all models with this \"feedforward\" structure fall within the domain of applicability of APJN.\n\n## Infinite Width Limit\n\nWe would like to emphasize that although our theoretical analyses are performed in infinite width limit, all training (including MLP, ResNet and MLP-Mixer) are performed on real, finite-width networks. For example, all MLP experiments were performed with N=500, L=50.\n\nFormally, the notion of criticality remains crisp for networks with $L/N \\sim o(1)$. $L/N$ corrections to the infinite width analysis have been calculated in other works [1][2]. Finite $L/N$ creates fluctuations in partial Jacobian-norm around its mean value. Nevertheless, for practical purposes, the notion of criticality remains largely unchanged. This can be readily seen by the agreement of our finite-width experiments (phase diagrams) and infinite-width predictions (dashed lines) in Figure 2.\n\nWe direct the reviewer to Global Response for further remarks.\n\n## Extension to other architectures and techniques\n\n- As mentioned above, empirical application of APJN is quite general. On the theoretical side, a comprehensive extension of our methods to Transformers would be very useful; especially due to the ubiquity of Transformers across various tasks/modalities. \n\n- We study normalization techniques like LayerNorm as well as residual connections because of their prevalence in modern architectures. We encourage further studies of our methods on domain-specific techniques. \n\n## Computational Resources\n\nAppendix A outlines the computational resources utilized in performing each of our experiments. In general, the computation of empirical APJN takes $<1\\\\%$ of the resources required for training. We are happy to provide further details.\n\n[1] S. Yaida, Non-Gaussian processes and neural networks at finite widths, 2020.\n\n[2] B. Hanin, Random Fully Connected Neural Networks as Perturbatively Solvable Hierarchies, 2022."}, {"review_id": "ljNTlTbwcV", "reviewer": "Reviewer_NnYC", "summary": "This paper studies criticality of deep neural networks at initialization. The authors propose a new practical way to diagnose criticality by introducing the partial Jacobian of the network and analyzing the averaged partial Jacobian norm (APJN) and its recurrence relation at large depth. The authors then apply their method to analyze criticality in fully connected networks with LayerNorm and/or residual connections, providing theoretical analysis for infinitely wide networks and a numerical test to select optimal initialization and identifying conditions on network architecture that allow for criticality for any initialization. ", "questions": "1. How is the APJN defined in this paper related to the order parameters used in previous works [Poole et al, Schoenholz et al.] for criticality (specifically the $c^*$ and $\\chi_1$ defined in Eq.4-5 in [Schoenholz et al.]). It would be nice if the authors can add a paragraph on how the APJN is related to these order parameters, and how the empirical test they introduced would be superior to evaluating scalings of the order parameters directly (providing a simpler numerical test? Unbounded activation functions?). \n2. Line 146: I am a bit confused what are the respective roles of $\\chi_K$ and $\\chi_J$ in the trainability of the network and how they are related to each other and the conditions on $\\chi_1$ and $c^*$  in [Poole et al.], can you clarify in this paragraph? \n3. Section 2: There are two limits here, the infinite width and the infinite depth. It seems that it is important for the scope of this paper to take the infinite width limit (successively) first, is that correct?\n4. Figure 3: This figure perhaps need a better caption. Are the white dashed lines given by the infinite width analysis? Why does there seem to be a larger discrepancy in networks with LayerNorm? \n", "limitations": "The authors mentioned in various parts of the paper but did not specifically summarize the limitations, I would recommend adding a paragraph in the conclusion section. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper presents theoretical analysis of the APJN in the infinite width limit for various network architectures, and thorough validation by numerics. \t\n2. The paper extensively explored application of the theoretical analysis and numerical test on modern architectures (ResNet and MLPmixer), providing practical insights for initializing neural networks for improved trainability.  \n", "weaknesses": "1. As most other works studying critical initialization, the work focuses on networks at Initialization with Gaussian random weights, and there is no learning in the network. While the authors briefly mentioned NTK (line 150), I would like to see a discussion on how their approach may be extended for analysis beyond networks at initialization and perhaps shed light on the learning dynamics of the network. \n2. In section 1.2, the authors discussed several related works and how part of their results were previously obtained in a different form in these works. I would recommend that the authors further stress the novelty of their approach/analysis in this section. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback. \n\n## Training Dynamics and Relation to NTK\n\nThe training regime that we target enjoys feature learning and large learning rates, and hence, is far from the NTK regime. While it is possible to connect APJN with NTK and analyze linear dynamics, that is not the purpose of our work. (We can readily add the quantitative relation between APJN and NTK in the Appendix.) Moreover, dynamics beyond the NTK regime for non-linear networks with $L>2$ is a famously unsolved problem.\n\n## Novelty\n\nThe reviewer has correctly identified the key advantages of our work. We elaborate them here:\n\n- All previous diagnostics of criticality require either closed-form solutions or an integral formula followed by numerical estimation [1,2]. This limits their usage to general/unseen architectures. Empirical APJN circumvents these limitations, and is useful as a numerical test for *general* real-world architectures. \\\nIn addition to the architectures presented in the text, one can analyze inhomogeneous networks by dividing them into \"blocks\" of layers and considering block-to-block Jacobians. \n\n- $\\chi_{\\mathcal J}$ can be numerically calculated with a single layer-to-layer gradient, which is readily obtained using Autodiff. Consequently, our method is simple and computationally cheap.\n\n- The characterization of everywhere-critical regime using normalization layers and residual connections. This finding, to our knowledge, is new.\n\nIn Section 1.2 of the final version, we will emphasize our contributions and the novelty of our approach.\n\n## The role of $\\chi_{\\mathcal J}, \\chi_{\\mathcal K}$ and relation to $c^*, \\chi_1$:\n\n$\\chi_{\\mathcal J}^l$ is defined as the layer-to-layer APJN $\\mathcal J^{l,l+1}$. For large enough $l$ it converges to the fixed point $\\chi_{\\mathcal J}^*$. The constraint $\\chi_{\\mathcal J} = 1$ regulates the backward pass (gradient propagation) in the network. For Fully connected networks (FCNs), $\\chi_\\mathcal J^*$ can be viewed as a generalization of $\\chi_1$ defined in the works [1,2]. The key advantages of our method over theirs are:\n\n- Unlike $\\chi_1$, $\\chi_{\\mathcal J}^*$ can be applied to both bounded and unbounded activation functions. (Note that $c^*$ in the aforementioned works fails to capture the critical behavior for unbounded activations such as ReLU -- $c^*=1$ in both phases.) We point this out in line 148.\n\n- As mentioned in the \"Novelty\" section above, empirical $\\chi_{\\mathcal J}$ is computationally cheap and has general applicability compared to the aforementioned works.\n\nWhile $\\chi_{\\mathcal J}=1$ regulates the backward pass of the network (i.e. gradient updates), $\\chi_{\\mathcal K}=1$ regulates the forward pass of the networks (i.e. preactivation norms). In practice, we find that the condition on $\\chi_{\\mathcal J}$ is more important for trainability. \n\nWe will state these points more clearly in the final version.\n\n## Order of limits in Section 2\n\nYes, for the scope of our work, it is important to take the width to infinity first. More precisely, we take the successive layer widths to infinity (sequential limit in Theorem 1.3). In practice, this translates to $L/N \\sim o(1)$.\n\n## Figure 3 Caption and Clarifications\n\nWe had to truncate the image captions due to space constraints. In the final version, we will add the following caption for Figure 3:\n\n\"Figure 3: Trainability (Training accuracy) of deep MLP on FashionMNIST $(N_l=500, L=50)$. The dotted white line denotes the (analytical) critical lines. Note that the combination of LayerNorm and residual connections significantly improve trainability, with the $\\mu=1$ case being everywhere-trainable.\"\n\nYes, the white dashed lines are critical lines from theory.\n\nWe are unsure of what the reviewer means by \"larger discrepancy in networks with LayerNorm\". We answer this question with two possible interpretations. We urge the reviewer to clarify their question so we can better answer it.\n\n- Networks with LayerNorm seem trainable even far from the critical line (dotted white line): LayerNorm, especially in conjunction with residual connections drastically improves the correlation length. For real, finite-depth networks, correlation length sets the scale for trainable depth. This results in finite-depth networks training well in a finite region around criticality.\n\n- Critical lines have seemingly different slopes in Figures 2 and 3: This is simply the result of the span of the Y axes of the plots. They are identical otherwise.\n\n[1] B. Poole, et al., Exponential expressivity in deep neural networks through transient chaos, 2016.\n\n[2] S. S. Schoenholz, Deep Information Propagation, 2016."}]

Author rebuttal:
# Global Response to All Reviewers

## Infinite width limit

For analytical results, we first take the infinite width limit and then a large depth limit. Stated differently, it means that the depth-to-width ratio $L/N \sim o(1)$. In practice, this assumption holds as long as $L/N$ is small. Moreover, residual connections make networks effectively shallow, significantly decreasing the effective $L/N$ ratio, and allowing our conclusions to be applied at even larger depths.

All of our phase diagrams and training results are obtained from real networks with reasonable width and depth (MLP: N=500, L=50; ResNet and MLP-Mixer: N=512, L=32). This should convince the reviewers that although our theory is formulated in the infinite width limit, our results concerning criticality describe real-world models; albeit with an error (fluctuation) of up to $O(L/N)$.

## Transformer Architectures

We have added phase diagrams for Vision Transformers (ViTs) in the global response PDF file. As predicted, $\mu=1$ with pre-LN is everywhere critical. Removing LayerNorm or using smaller $\mu$ leads to non-critical initializations. We will add these results to the Appendix of the final version. 

On the theoretical side, the analysis is more involved. It is known that the commonly used scalar-dot product attention layer converges to a Gaussian process in the limit where the number of heads and the embedding dimensions go to infinity [1]. However, real-world models are far from this limit; wherein $L/n_{heads}$ is $O(1)$ (even with residual connections). This makes infinite-width analysis unfavorable. Moreover, Softmax function in attention prevents closed-form results for most calculations. Despite these limitations, people have found a particular width scaling of initialization and learning rate for pre-LN transformers [2].

Despite the lack of comprehensive theoretical understanding, our tools can extract meaningful information about the magnitude of gradients for attention layers (See attached one-page pdf). It identifies criticality and correctly predicts that pre-LN transformers are better-behaved than post-LN transformers [3].

## Captions and Clarity

Due to space limitations, we had oversimplified some of the Figure-captions. We will write more detailed captions in the final version. Details of experiments and computational resources are fleshed out in Appendix A.

[1] J. Hron, et al., Infinite attention: NNGP and NTK for deep attention networks, 2020.

[2] E. Dinan, et al., Effective Theory of Transformers at Initialization, 2023.

[3] R. Xiong, et al., On Layer Normalization in the Transformer Architecture, 2020.

Meta-review: The paper introduces a novel practical approach to identify criticality within neural networks. Using the infinite width limit, the authors derive analytical expressions for the neural network Gaussian process kernel and the partial Jacobian norm for a wide range of settings. The authors show that the combination of LayerNorm and residual connections leads to an everywhere-critical regime, where the network can be trained effectively irrespective of the initialization. The paper's contributions include a theoretical understanding of the trainability of deep neural networks with LayerNorm and residual connections via averaged partial Jacobian Norm and insights into the design of effective architectures. The authors apply these methods to investigate the ResNet and MLP-Mixer architectures, revealing the presence of an everywhere-critical regime within these modern settings.


[Strength highlighted by reviewers]

- Introduces partial Jacobians and their averaged norms as tools to analyze the propagation of gradients through deep neural networks at initialization
- The paper presents theoretical analysis of the APJN in the infinite width limit for various network architectures, and thorough validation by numerics.
- Presents a computationally cheap and simple empirical test for criticality using APJN evaluated close to the output.
- Shows that criticality formulated in terms of partial Jacobians is equivalent to criticality studied previously in the literature.
- Investigates homogeneous architectures that include fully-connected layers, normalization layers, and residual connections, and shows that the combination of LayerNorm and residual connections can drastically increase correlation length leading to improved trainability.
- The paper explored application of the theoretical analysis and numerical test on modern architectures (ResNet and MLPmixer), providing practical insights  that they are critical everywhere at µ = 1 (i.e., with residual connections) due to the interaction between LayerNorm and residual connections.
- The paper provides a theoretical assessment of why, in the infinite width limit, certain neural networks can achieve better trainability. In particular, the authors study the effect of residual connections and layer norm, both which have become a staple of deep learning. This type of study can help explain why certain architectures are less sensitive to initialisation compared to others.
- Code reproducing the figures is provided in the supplementary material as notebooks

[Weakness highlighted by reviewers] Main issues raised by reviewers are on clarity and suggested many forms / directions suggesting organizing presentations based on clear logical deduction.

[Overall] After the rebuttal, all reviewers supported accepting the paper (score of 7, 7, 6, 6) saying "the contribution of the paper is worthwhile", " a significant advancement in the field", "powerful tools for analyzing gradient propagation in deep neural networks at initialization", "paper is highly comprehensive and lucid, greatly facilitating understanding". While there are still issues on clarity remaining after rebuttal, the AC believes this can be effectively addressed if authors incorporate many promised changes suggested by the reviewers. I strongly encourage the authors to take these changes seriously to improve clarity and presentation to help the readers.
Decision: Accept (spotlight)