Title: Critical Initialization of Wide and Deep Neural Networks using Partial Jacobians: General Theory and Applications
Venue: NeurIPS 2023
Accepted: yes
Keywords: Criticality

Gaussian Process

Jacobian

LayerNorm

Residual connections

ResNet
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

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 ≤ l . We derive recurrence relations for the norms of partial Jacobians and utilize these relations to analyze criticality of deep fully connected neural networks with LayerNorm and/or residual connections. We derive and implement a simple and cheap numerical test that allows one to select optimal initialization for a broad class of deep neural networks; containing fully connected, convolutional and normalization layers. Using these tools we show quantitatively that proper stacking of the LayerNorm (applied to preactivations) and residual connections leads to an architecture that is critical for any initialization. Finally, we apply our methods to analyze ResNet and MLP-Mixer architectures; demonstrating the everywhere-critical regime. 4

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 process [37, 30] with a kernel - sometimes referred to as neural network Gaussian process (NNGP) kernel [30] - determined by the network architecture and hyperparameters ( e.g depth, precise choices of layers and the activation functions, as well as the distribution of weights and biases). A similar line of reasoning was earlier developed for recurrent neural networks [36]. Furthermore, for special choices of parameterization and MSE loss function, the training dynamics under gradient descent can be solved exactly in terms of the neural tangent kernel (NTK) [27, 31]. A large body of work was devoted to the calculation of the NNGP kernel and NTK for different architectures, calculation of the finite width corrections to these quantities, and empirical investigation of the training dynamics of wide networks [39, 54, 25, 14, 3, 33, 1, 18, 19, 45, 58, 47, 6, 5, 32, 63, 61, 60, 59, 35, 17, 2, 49, 34].

One important result that arose from these works is that the network architecture determines the most appropriate initialization of the weights and biases [44, 46, 30]. To state this result, we consider networks with/without LayerNorm [7] and residual connections [21]; the preactivations for which can be defined as follows

where ˜ h l j = LayerNorm( h l j ) and the parameter µ controls the strength of residual connections. For the input layer: h 1 i ( x ) = ∑ N 0 j =1 w 1 ij x j + b 1 i . In the ( l +1) -th layer, weights w l +1 ij ∈ R N l +1 × N l and biases b l +1 i ∈ R N l +1 × 1 are taken from normal distributions N (0 , σ 2 w /N l ) and N (0 , σ 2 b ) , respectively. Hyperparameters σ w and σ b need to be tuned. ϕ ( · ) is the activation function and x ∈ R N 0 × 1 is the input. For results discussed in this work, x can be sampled from either a realistic ( i.e. highly correlated) dataset or a high entropy distribution.

For a network of depth L , the network function is given by f ( x ) = h L ( x ) . Different network architectures and activation functions, ϕ , lead to different 'optimal' choices of ( σ w , σ b ) . The optimal choice can be understood, using the language of statistical mechanics, as a critical point (or manifold) in the σ b -σ w plane. The notion of criticality becomes sharp as the network depth, L , becomes large. Criticality ensures that NNGP kernel and the gradients' norm remain O ( L 0 ) as the network gets deeper [45]. Very deep networks will not train unless initialized critically, since the gradients explode or vanish exponentially. Note that high trainability does not imply that the trained model has great performance (test accuracy) after training.

1.1 Results

Here we focus on two main results of this work: (i) empirical method to check criticality of a neural network and (ii) an architecture based on layer normalization and residual connections that is critical for any initialization. First we introduce the notion of a partial Jacobian.

Definition 1.1. Let h l i ( x ) be preactivations of a neural network f ( x ) . The partial Jacobian J l 0 ,l ij is defined as derivative of preactivations at layer l with respect to preactivations at layer l 0 ≤ l

The partial Jacobian is a random matrix with vanishing mean at initialization. Fixing l 0 and varying l allows us to study the behavior of Jacobian with depth. On the other hand, varying both l 0 and l is necessary to study general networks with non-repeating building blocks. Next, we introduce a deterministic measure of the magnitude of J l 0 ,l ij -its squared Frobenius norm, averaged over parameter-initializations.

Definition 1.2. Let J l 0 ,l ij be a partial Jacobian of a neural network f ( x ) . Averaged partial Jacobian norm (APJN) 5 is defined as

where E θ indicates averaging over parameter-initializations.

APJN is a dominant factor in the magnitude of the gradients. In what follows, we show that criticality, studied previously in the literature, occurs when APJN either remains finite, or varies algebraically as l becomes large. To prove this we derive the recurrence relation for J l 0 ,l ( x ) in the limit N l →∞ and analyze it at large depth. Algebraic behavior of APJN with depth is characterized by an architecturedependent critical exponent, ζ , so that J l 0 ,l ( x ) ≈ l -ζ . Such behavior is familiar from statistical mechanics when a system is tuned to a critical point [10]. Away from criticality, there are two phases:

ordered and chaotic. In the ordered phase APJN vanishes exponentially with depth, whereas in the chaotic phase APJN grows exponentially

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) . Let f ( x ) be a deep MLP network with Lipschitz continuous activation ϕ ( · ) . Assume that LayerNorm is applied to preactivations and there are residual connections with strength µ acting according to (1) . In the sequential limit ( N 1 →∞ , . . . , N l →∞ , . . . , N L -1 → ∞ ) , the correlation length with a large L can be written as

where the non-negative constants A and B are given by

and ϕ ′ ( · ) is the derivative of ϕ ( · ) . We omitted the average over the neuron index k in the infinite width limit, and the result does not depend on k .

Remark 1.4 . As µ increases, the dependence on the initialization ( σ b , σ w ) 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 no 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 et al. [7]. LayerNorm also alters the criticality of J l 0 ,l ( x ) . It was noted in Xu et al. [57] that LayerNorm (applied to preactivations) regularizes the backward pass. We formalize this observation by showing that LayerNorm (applied to preactivations) dramatically enhances correlation length (which is not the case for LayerNorm applied to activations). This can be seen from Theorem 1.3, setting µ = 0 . When residual connections of strength 1 are combined with erf (or any other erf -like activation function, e.g. tanh), the neural network enters a subcritical phase with enhanced correlation length (see 4.3). A version of this result was discussed in Yang and Schoenholz [62]. When residual connections are introduced on top of LayerNorm, the correlation length ξ 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 initialization 6 [20], Fix-up [65] and ReZero [8] are special cases of our method. They conform to the notion of criticality using APJN, as defined in our work.

The partial Jacobian has been used to study generalization bounds in Arora et al. [4]. The Jacobian norm ( i.e. || J 0 ,l ij || 2 ) of trained feed-forward neural networks was studied in Novak et al. [38], where it was correlated with generalization. Partial Jacobians with l 0 = l -1 were studied in the context of RNNs [11, 9], referred to as state-to-state Jacobians. Furthermore, the relation between Partial Jacobians and Lyapunov exponents, as well as their impact on trainability, have been explored for RNNs [16, 15].

As the aspect ratio ( L/N ) of the network approaches 1, the finite width corrections to the Jacobian become more prominent. On the other hand, even with a small aspect ratio, the effect of the spectral density of Jacobian becomes important as the depth L becomes very large. Pennington et al. [43] studied the spectrum of the input-output Jacobian for MLPs. Xiao et al. [54] extended the analysis to CNNs, showing that very deep vanilla CNNs can be trained by achieving 'dynamical isometry'.

2 Recurrence Relations

Here we derive the infinite width recurrence relations for the APJN and the NNGP kernel. We use Lemma 2.2 to derive NNGP kernel recurrence relation, and leverage that to get the recurrence relation for APJN. Fixed point analyses of these relations help us define critical line and point. We start with the vanilla MLP with no LayerNorm and µ = 0 . Results with LayerNorm and/or residual connections, as well as modern architectures are presented in the following sections. (We refer the reader to Appendices D to F for proofs and detailed application of this recipe. Appendices I to K contain the calculations and results for various activation functions.)

Definition 2.1. We define averaged covariance of preactivations as follows

Lemma 2.2. When N l →∞ for l = 1 , . . . , L -1 sequentially, the expectation value over parameter initializations for a general function of preactivations: O ( h l ( x )) , can be expressed as the averaging over the Gaussian process h l ( x ) with covariance K l ( x, x ′ ) .

This result has been established in Lee et al. [30]. Note that the density in (7) only depends on the diagonal part of the covariance matrix, K l ( x, x ) . We will refer to K l ( x, x ) as NNGP kernel.

Remark 2.3 . In the sequential infinite width limit the means appearing in (9)-(11) are self-averaging and, therefore, deterministic. They converge in distribution to their averages over parameterizations.

When performing analytic calculations we use the infinite width convention; whereas in our finitewidth experiments we explicitly average over initializations of θ l .

Theorem 2.4. With 2.2, in the infinite width limit, the NNGP kernel K l +1 ( x, x ) is deterministic, and can be determined recursively via

Theorem 2.5. Let f ( x ) be an MLP network with a Lipschitz continuous activation function ϕ ( x ) . In the infinite width limit, APJN J l 0 ,l +1 ( x ) is deterministic and satisfies a recurrence relation

where the factor χ l J is given by

2.4 is due to Neal [37] and Lee et al. [30]. 2.5 is new and is valid only in the limit of infinite width. The proof is in D. We will drop the explicit dependence on x to improve readability.

The expectation values that appear in (9)-(11) are evaluated using (7). When the integrals can be taken analytically, they lead to explicit equations for the critical lines and/or the critical points. Details of these calculations as well as the derivation of (9)-(11) can be found in the D. A subtlety emerges in (10) when l 0 = 0 , where a correction of the order O ( N -1 0 ) arises for non-scale invariant activation functions. This subtlety is discussed in the D.

When the depth of the network becomes large, the l -dependence of the expectation values that appear in (6), (11) saturate to a (possibly infinite) constant value; which means that K l , J l 0 ,l and χ l J have

reached a fixed point. We denote the corresponding quantities as K ⋆ , J l 0 ,⋆ , χ ⋆ J . The existence of a fixed point is not obvious and should be checked on a case by case basis. Fixed point analysis for K l was done in Poole et al. [44] for bounded activation functions and in Roberts et al. [45] for the general case. The stability is formulated in terms of

The norm of preactivations remains finite (or behaves algebraically) when χ ⋆ K = 1 .

Eq. (10) nicely expresses J l 0 ,l +1 as a linear function of J l 0 ,l . The behaviour of J l 0 ,l +1 at large l is determined by χ l J . When χ l J > 1 partial Jacobians diverge exponentially, while for χ l J < 1 partial Jacobians vanish exponentially. Neural networks are trainable only up to a certain depth when initialized O (1) away from criticality, which is determined by the equation

Eq. (13) is an implicit equation on σ b , σ w and generally outputs a critical line in σ b -σ w plane. The parameter χ ⋆ J has to be calculated on a case-by-case basis using either (11) or the method presented in the next section. Everywhere on the critical line, J l 0 ,l saturates to a constant or behaves algebraically.

When the condition χ ⋆ K = 1 is added, we are left with a critical point 7 . This analysis of criticality at infinite width agrees with Roberts et al. [45], where χ ⊥ is to be identified with χ ⋆ J ; and Schoenholz et al. [46], Martens et al. [34], where their analysis based on the equivalent χ 1 or C ′ (1) only works for bounded activation functions. In particular, condition (13) together with χ ⋆ K = 1 ensures that NTK is O (1) at initialization.

2.1 Empirical Diagnostic of Criticality

APJN J l 0 ,l provides a clear practical way to diagnose whether the network is critical or not. Proper choice of l 0 and l allows us to minimize the non-universal effects and cleanly extract χ ⋆ J .

Recurrence relation (10), supplemented with the initial condition J l 0 ,l 0 +1 = χ l 0 J , can be formally solved as

We would like to obtain an estimate of χ ⋆ J as accurately as possible. To that end, imagine that for some l ′ > l 0 the fixed point has been essentially reached and χ l ′ J ≈ χ ⋆ J . Then the APJN

depends on the details of how the critical point is approached; which are encoded in the last factor.

Proposition 2.6. If the network is homogeneous, i.e. , consists of (possibly complex) blocks of layers, periodically repeated L times; then the penultimate APJN provides an accurate estimate of χ ⋆ J :

This is a direct consequence of combining (9) and (11) as L goes to infinity. See 1 for numerical justification.

2.6 is the central result of this section and will be heavily used in the remainder of this work.

Note that for deep networks, away from criticality, APJN takes the form

where c l 0 is a non-universal constant that depends on l 0 . If the sign in (17) is positive ( χ ⋆ J > 1 ) the network is in the chaotic phase , while when the sign is negative ( χ ⋆ J < 1 ) the network is in the

10 0

10 1

10 2

10

2

10

1

10 0

0,

erf

Experiment

Theory,

=1.000

Fitting,

=0.968

10 0

10 1

10 2

10 0

2×10 0

3×10 0

ReLU

Experiment

Theory,

=0.000

Fitting,

=0.004

10 0

10 1

10 2

10 0

10 2

LN-GELU, LN-erf;

=1.0

GELU Experiment

GELU Theory,  =

1.072

GELU Fitting,  =

1.068

erf Experiment

erf Theory,  =

1.226

erf Fitting,  =

1.217

ordered phase . ξ has the meaning of correlation length: on the depth scale of approximately kξ the gradients remain appreciable, and hence the network with the depth of ≈ kξ will train.

We used (16) to map out the σ b -σ w phase diagrams of various MLP architectures. The location of the critical line agrees remarkably well with our infinite width calculations. Results are presented in Fig. 2.

At criticality, χ ⋆ J = 1 and the correlation length diverges; indicating that gradients can propagate arbitrarily far. A more careful analysis of non-linear corrections shows that APJN can exhibit algebraic behavior with depth and can still vanish in the infinite depth limit, but much slower than the ordered phase.

2.2 Scaling at a Critical Point

At criticality χ l J saturates to a fixed value χ ⋆ J = 1 . If we are interested in J l 0 ,l with l -l 0 = O ( L ) then it is essential to know how exactly χ l J approaches 1 .

Theorem 2.7. Assume that deep neural network f ( x ) is initialized critically. Then l →∞ asymptotics of APJN is given by

where ζ is the critical exponent Roberts et al. [45], see Appendix G for further details.

Critical exponents can be determined analytically in the limit of infinite width. Note that χ l J , given by (11), depends on K l by virtue of (7). Consequently, Eqs. (9)-(11) are coupled through non-linear (in K l and J l 0 ,l ) terms. These non-linear corrections are absent for any scale-invariant activation function, but appear for other activation functions.

We checked the scaling empirically by plotting J 0 ,l vs. l in a log -log plot and fitting the slope. These results are presented in Fig.1, and are in excellent agreement with infinite width calculation.

3 Layer Normalization

The fact that critical initialization is concentrated on a single point ( σ ⋆ w , σ ⋆ b ) may appear unsettling because great care must be taken to initialize the network critically. The situation can be substantially improved by utilizing the normalization techniques known as LayerNorm [7] and GroupNorm [52]. Our results apply to GroupNorm verbatim in the case when the number of groups is much smaller than the width. LayerNorm can act either on preactivations or on activations (discussed in the Appendix D). Depending on this choice, criticality will occur on different critical lines in σ b -σ w plane. When LayerNorm is applied to preactivations the correlation length is enhanced, allowing for training much deeper networks even far away from criticality.

The LayerNorm applied to preactivations takes the following form

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

2

b

ReLU,

=0.0

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

ReLU,

=0.9

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

ReLU,

=1.0

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

LN-ReLU,

=0.0

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

LN-ReLU,

=0.9

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

LN-ReLU,

=1.0

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

2

b

erf,

=0.0

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

erf,

=0.9

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

erf,

=1.0

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

LN-erf,

=0.0

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

LN-erf,

=0.9

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

LN-erf,

=1.0

1

2

3

4

2

w

0.0

0.2

0.4

0.6

0.8

1.0

2

b

GELU,

=0.0

1

2

3

4

2

w

0.0

0.2

0.4

0.6

0.8

1.0

GELU,

=0.9

1

2

3

4

2

w

0.0

0.2

0.4

0.6

0.8

1.0

GELU,

=1.0

1

2

3

4

2

w

0.0

0.2

0.4

0.6

0.8

1.0

LN-GELU,

=0.0

1

2

3

4

2

w

0.0

0.2

0.4

0.6

0.8

1.0

LN-GELU,

=0.9

1

2

3

4

2

w

0.0

0.2

0.4

0.6

0.8

1.0

LN-GELU,

=1.0

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Definition 3.1 (Normalized preactivations) .

where we have introduced E [ h l ] = ( ∑ N l i =1 h l i ) /N l . In the limit of infinite width E [ h l ] = 0 and E [( h l ) 2 ] = K l , defined according to (6).

Normalized preactivations, ˜ h l i , are distributed according to N (0 , 1) for all l, σ w , σ b . The norms are, therefore, always finite and the condition χ ⋆ K = 1 is trivially satisfied. This results in a critical line rather than a critical point.

The recurrence relations (9)-(11) for the NNGP and partial Jacobians are only slightly modified

Assuming that the value of χ l J at the fixed point is χ ⋆ J , the network is critical when (13) holds.

χ l J (20) is depth independent and changes slowly with σ w and σ b . Thus, χ ⋆ J remains close to 1 for a wider range of hyperparameters. Consequently, the correlation length is large even away from criticality, leading to a much higher trainability of deep networks.

4 Residual (Skip) Connections

Adding residual connections between the network layers is a widely used technique to facilitate the training of deep networks. Originally introduced [21] in the context of convolutional neural networks [29] (CNNs) for image recognition, residual connections have since been used in a variety of networks architectures and tasks [50, 13].

Consider (1) with non-zero µ and without LayerNorm layers. Then the recurrence relations (9)-(11) for the NNGP kernel and χ l J are modified as follows

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

2

b

ReLU,

=0.0

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

ReLU,

=0.9

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

ReLU,

=1.0

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

LN-ReLU,

=0.0

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

2

b

LN-ReLU,

=0.9

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

LN-ReLU,

=1.0

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

2

b

erf,

=0.0

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

erf,

=0.9

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

erf,

=1.0

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

LN-erf,

=0.0

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

LN-erf,

=0.9

1

2

3

4

5

0.0

0.5

1.0

1.5

2.0

LN-erf,

=1.0

1

2

3

4

5

2

w

0.0

0.5

1.0

1.5

2.0

2

b

GELU,

=0.0

1

2

3

4

5

2

w

0.0

0.5

1.0

1.5

2.0

GELU,

=0.9

1

2

3

4

5

2

w

0.0

0.5

1.0

1.5

2.0

GELU,

=1.0

1

2

3

4

5

2

w

0.0

0.5

1.0

1.5

2.0

LN-GELU,

=0.0

1

2

3

4

5

2

w

0.0

0.5

1.0

1.5

2.0

LN-GELU,

=0.9

1

2

3

4

5

2

w

0.0

0.5

1.0

1.5

2.0

LN-GELU,

=1.0

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

Remark 4.1 . When µ < 1 , the fixed point value of NNGP kernel is scaled by (1 -µ 2 ) -1 . For µ = 1 , the critical point is formally at (0 , 0) .

Remark 4.2 . For µ = 1 , (21) implies that χ l J ≥ 1 , where the equality holds on the σ w = 0 axis. Consequently, APJN exponentially diverges as a function of depth l for all σ w > 0 . In this case, σ w 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 (21) for erf activation, we find an effect observed in Yang and Schoenholz [62] for tanh activation. They noted that tanh -like MLP networks with skip connections "hover over the edge of chaos". We quantify their observation as follows.

Theorem 4.3. Let f ( x ) be a deep MLP network with erf activation function and residual connections of strength µ = 1 . Then in the limit N l →∞

The NNGP kernel K l linearly diverges with depth l .

χ l J approaches 1 from above (Fig. 2) : χ l J ≈ 1 + ˜ c/ √ l , where ˜ c = 2 σ 2 w / ( π √ σ 2 w + σ 2 b ) is a non-universal constant. Consequently, APJN diverges as a stretched exponential : J l 0 ,l = O ( e √ l λ ) , where λ = 1 / (4˜ c 2 ) is the new length scale.

We will refer to this case as subcritical . Although χ ⋆ J reaches 1 , the APJN still diverges with depth faster than any power law. The growth is controlled by the new scale λ . To control the gradient we would like to make λ large, which can be accomplished by decreasing σ w . In this case, the trainability is enhanced (see Fig. 3). Similar results hold for tanh activation function [62], however in that case there is no explicit expression for ˜ c .

5 Residual Connections + LayerNorm

In practice, it is common to use a combination of residual connections and LayerNorm. Using (1), the recurrence relations (9)-(11) for NNGP and partial Jacobians are modified as follows

Remark 5.1 . For µ < 1 , (22) implies that the fixed point value of NNGP kernel is scaled by 1 -µ 2 . Moreover, residual connections do not shift the phase boundary. The interference between residual connections and LayerNorm brings χ l J closer to 1 on the entire σ b -σ w plane (as can be seen from Fig. 2). Therefore the correlation length ξ is 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 a more stable architecture. This observation is formalized by 1.3. The proof leverages the solutions of (22) close to the fixed point, and is fleshed out in Appendix F.

Remark 5.2 . When µ = 1 , the correlation length diverges for any initialization.

5.2 provides an alternative perspective on architecture design. On the one hand, one can use (16) to initialize a given network at criticality. Alternatively, 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 χ ⋆ J = 1 holds on the entire σ b -σ w 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 /σ w . The trainability is dramatically enhanced, as shown in Fig. 3.

Remark 5.4 . Networks with BatchNorm [26], used in conjunction with residual connection of strength µ = 1 , also enjoy this everywhere criticality and enhanced trainability [63, 23] (See Appendix B).

6 Modern Architectures

ResNet110 ResNet is one of the most widely used architectures in computer vision[21]. Figure 4 shows the results for ResNetV2 [22]; with BatchNorm replaced with LayerNorm. (See Appendix B for BatchNorm results and discussions.)

At µ = 1 , σ 2 w -σ 2 b phase diagram shows everywhere-criticality, as expected from our theory. For σ 2 b = 0 , σ 2 w -µ phase diagram gets closer to criticality as we increase µ , which results in better training at higher µ .

Additionally, networks with larger µ enjoy better trainability due to the suppression of finite width corrections to our theory (decreased effective L/N ).

MLP-Mixer MLP-Mixer architecture is a recent example of MLP approach to computer vision [48]. It can be analyzed using the tools presented above. Fig. 5 shows the phase and training diagrams. Further details can be found in the SM. Note that the original architecture uses µ = 1 . 0 .

Test accuracies for figures 4, 5 follow similar trends as train accuracies. (See Appendices B, H).

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, J L -2 ,L -1 ≈ χ ⋆ J , 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 to improved trainability, (iii) combination of µ = 1 residual connections and erf activation function enhances correlation length driving the network to a subcritical phase with APJN growing according to a stretched exponential law, (iv) combination of residual connections and LayerNorm drastically increases correlation length leading to improved trainability, (v) when µ = 1 and LayerNorm is applied to preactivations the network is critical on the entire σ b -σ w plane.

We have considered examples of modern architectures: ResNet and MLP-Mixer. We showed that at µ = 1 , they are critical everywhere, due to the interaction between LayerNorm and residual connections. We have studied ResNet at different µ '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-Mixer with µ = 1 trains well for various initializations. In comparison, for µ = 0 . 5 , it only trains well close to the critical line.

Our work shows that an architecture can be designed to have a large correlation length leading to a guaranteed trainability with SGD for any initialization scheme.

8 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 to initialize Transformer architectures.

Another challenge arises from network topology. For graph neural networks, the generalization of our 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.

Acknowledgments and Disclosure of Funding

We thank T. Can, G. Gur-Ari, B. Hanin, D. Roberts and S. Yaida for comments on the manuscript. A.G.'s work at the University of Maryland was supported in part by NSF CAREER Award DMR2045181, Sloan Foundation and the Laboratory for Physical Sciences through the Condensed Matter Theory Center.

References

Kyle Aitken and Guy Gur-Ari. On the asymptotics of wide networks with polynomial activations. arXiv preprint arXiv:2006.06687 , 2020.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via overparameterization. In International Conference on Machine Learning , pages 242-252. PMLR, 2019.

Anders Andreassen and Ethan Dyer. Asymptotics of wide convolutional neural networks. arXiv preprint arXiv:2008.08675 , 2020.

Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via a compression approach. In International Conference on Machine Learning , pages 254-263. PMLR, 2018.

Sanjeev 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.

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. In International Conference on Learning Representations , 2019.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

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, Proceedings of the Thirty-Seventh Conference on Uncertainty in Artificial Intelligence , volume 161 of Proceedings of Machine Learning Research , pages 1352-1361. PMLR, 27-30 Jul 2021. URL https://proceedings.mlr.press/v161/bachlechner21a.html .

Tankut Can, Kamesh Krishnamurthy, and David J Schwab. Gating creates slow modes and controls phase-space complexity in grus and lstms. In Mathematical and Scientific Machine Learning , pages 476-511. PMLR, 2020.

John Cardy. Scaling and renormalization in statistical physics , volume 5. Cambridge university press, 1996.

Minmin Chen, Jeffrey Pennington, and Samuel Schoenholz. Dynamical isometry and a mean field theory of rnns: Gating enables signal propagation in recurrent neural networks. In International Conference on Machine Learning , pages 873-882. PMLR, 2018.

Soham De and Sam Smith. Batch normalization biases residual blocks towards the identity function 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 Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper_files/paper/2020/file/ e6b738eca0e6792ba8a9cbcba6c1881d-Paper.pdf .

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 https://openreview.net/forum? id=YicbFdNTTy .

Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations , 2019.

Rainer Engelken. Gradient flossing: Improving gradient descent through dynamic control of jacobians. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. URL https://openreview. net/forum?id=jEQRoJzDx8 .

Rainer Engelken, Fred Wolf, and L. F. Abbott. Lyapunov spectra of chaotic recurrent neural networks. Phys. Rev. Res. , 5:043044, Oct 2023. doi: 10.1103/PhysRevResearch.5.043044. URL https://link. aps.org/doi/10.1103/PhysRevResearch.5.043044 .

Adrià Garriga-Alonso, Carl Edward Rasmussen, and Laurence Aitchison. Deep convolutional networks as shallow gaussian processes. arXiv preprint arXiv:1808.05587 , 2018.

Mario Geiger, Arthur Jacot, Stefano Spigler, Franck Gabriel, Levent Sagun, Stéphane d'Ascoli, 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.

Boris Hanin. Random neural networks in the infinite width limit as gaussian processes. arXiv preprint arXiv:2107.01562 , 2021.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision , pages 1026-1034, 2015.

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.1109/CVPR.2016.90.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In Computer Vision - ECCV 2016 , pages 630-645, 2016.

Tianyu He, Darshil Doshi, and Andrey Gromov. Autoinit: Automatic initialization via jacobian tuning, 2022.

Judy Hoffman, Daniel A. Roberts, and Sho Yaida. Robust learning with jacobian regularization. arXiv preprint arXiv:1908.02729 , 2019.

Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nngp and ntk for deep attention networks. In International Conference on Machine Learning , pages 4376-4386. PMLR, 2020.

Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning , pages 448-456. PMLR, 2015.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint arXiv:1806.07572 , 2018.

Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.

Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE , 86(11):2278-2324, 1998.

Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. In International Conference on Learning Representations , 2018.

Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, 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.

Jaehoon Lee, Samuel S Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. 2020.

Aitor Lewkowycz and Guy Gur-Ari. On the training dynamics of deep networks with l _ 2 regularization. arXiv preprint arXiv:2006.08643 , 2020.

James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszcz, Valentin Dalibard, Jascha SohlDickstein, and Samuel S Schoenholz. Rapid training of deep neural networks without skip connections or normalization layers using deep kernel shaping. arXiv preprint arXiv:2110.01765 , 2021.

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.

Lutz Molgedey, J Schuchhardt, and Heinz G Schuster. Suppressing chaos in neural networks by noise. Physical review letters , 69(26):3717, 1992.

Radford M Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks , pages 29-53. Springer, 1996.

Roman Novak, Yasaman Bahri, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Sensitivity and generalization in neural networks: an empirical study. In International Conference on Learning Representations , 2018.

Roman Novak, Lechao Xiao, Yasaman Bahri, Jaehoon Lee, Greg Yang, Jiri 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.

Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A Alemi, Jascha Sohl-Dickstein, and Samuel S Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In International Conference on Learning Representations , 2019.

Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In International Conference on Learning Representations , 2020.

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 DeVito, 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é-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32 , pages 8024-8035. Curran Associates, Inc., 2019.

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-First International Conference on Artificial Intelligence and Statistics , volume 84 of Proceedings of Machine Learning Research , pages 1924-1932. PMLR, 09-11 Apr 2018. URL https://proceedings.mlr. press/v84/pennington18a.html .

Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. Advances in neural information processing systems , 29:3360-3368, 2016.

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.

Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. 2016.

Vaishaal Shankar, Alex Fang, Wenshuo Guo, Sara Fridovich-Keil, Jonathan Ragan-Kelley, Ludwig Schmidt, and Benjamin Recht. Neural kernels without tangents. In International Conference on Machine Learning , pages 8614-8623. PMLR, 2020.

Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy. Mlp-mixer: An all-mlp architecture for vision. arXiv preprint arXiv:2105.01601 , 2021.

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.

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. V on Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper_files/ paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf .

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.

Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV) , pages 3-19, 2018.

Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017.

Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington. Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks. In International Conference on Machine Learning , pages 5393-5402. PMLR, 2018.

Lechao Xiao, Jeffrey Pennington, and Sam Schoenholz. Disentangling trainability and generalization in deep learning, 2020.

Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning , pages 10524-10533. PMLR, 2020.

Jingjing Xu, Xu Sun, Zhiyuan Zhang, Guangxiang Zhao, and Junyang Lin. Understanding and improving layer normalization. arXiv preprint arXiv:1911.07013 , 2019.

Sho Yaida. Non-gaussian processes and neural networks at finite widths. In Mathematical and Scientific Machine Learning , pages 165-192. PMLR, 2020.

Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.04760 , 2019.

Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. 2019.

Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning , pages 11727-11737. PMLR, 2021.

Greg Yang and Samuel S Schoenholz. Mean field residual networks: On the edge of chaos. arXiv preprint arXiv:1712.08969 , 2017.

Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S Schoenholz. A mean field theory of batch normalization. In International Conference on Learning Representations , 2018.

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:2111.11418 , 2021.

Hongyi Zhang, Yann N. Dauphin, and Tengyu Ma. Residual learning without normalization via better initialization. In International Conference on Learning Representations , 2019. URL https: //openreview.net/forum?id=H1gsz30cKX .

A Experimental Details

We implemented our methods using PyTorch [42] hooks and an efficient Jacobian approximate algorithm [24].

Figure 1: We generated MNIST-like inputs, where all elements are sampled from the Gaussian distribution N (0 , 1) . J 0 ,l data was averaged over 100 different parameter-initializations. Networks were initialized width N l = 1000 . For erf plot we initialized at critical point ( σ w , σ b ) = ( √ π 4 , 0) , used depth L = 250 and the fitting was done with data points collected at depth l > 100 ; for ReLU plot we initialized at critical point ( σ w , σ b ) = ( √ 2 , 0) , used depth L = 100 and the fitting was done with all data points; for the µ = 1 Pre-LN plot, we initialized both networks at ( σ w , σ b ) = ( √ 2 , 0) , used depth L = 250 and the fitting was done with l > 100 data points.

Figure 2: All the phase diagrams were plotted using χ L -1 J generated from networks with L = 50 and N l = 500 . We used hooks to obtain the gradients that go into calculating χ L -1 J . χ L -1 J data was averaged over 100 different parameter-initializations. Inputs were generated from a normal Gaussian distribution and have dimension 28 × 28 . Generating the data for the figure took approximately 2 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 L = 50 and width N l = 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 4: (1)We made the σ 2 w -σ 2 b phase diagram for ResNet110(LayerNorm) by averaging over 100 different parameter-initializations. The σ 2 w -µ phase diagram was made by averaging over 200 parameters initialization. (3)(4) 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 -4 . All models were trained for 50 epochs and averaged over 3 random seeds. It takes 6 GPU days in total on a single NVIDIA RTX 3090 GPU.

Figure 5: (1)(2)We made the phase diagram for MLP-Mixer with 30 blocks and averaged over 100 different parameter-initializations. (3)(4)We used network with L = 100 , patch size 4 × 4 , hidden size C = 128 , two MLP dimensions N tm = N cm = 256 . The L = 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 dataset using vanilla SGD paired with CSE. Batch size bs = 256 , weight decay λ = 10 -4 was selected from { 10 -5 , 10 -4 } , mixup rate α = 0 . 8 was selected from { 0 . 4 , 0 . 8 } . We also used RandAgument and horizontal flip with default settings in PyTorch. For all cases we searched learning rates within { 0 . 005 , 0 . 01 , 0 . 05 , 0 . 1 , 0 . 2 , 0 . 5 } . We also tried a 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 4 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: K l µν ; ij ( x, x ′ ) , where the Greek letters ( µ, ν ) 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: K l µν ; ij ( x, x ′ ) = K l ij ( x, x ′ ) δ µν . This additional structure in the kernel makes it difficult to get a closed-form solution for J l,l +1 in general.

ResNet110 (LayerNorm) In Figure 4(2), the networks is critical close to µ = 1 , as expected from our analysis. One would naively expect the µ < 1 cases also to be critical, since for MLP with ReLU and Pre-LN, σ b = 0 is critical regardless of σ w and µ . However, in Figure 4(2) the region away from µ = 1 is in ordered phase. This is likely a result of the kernel K l µν ; ij ( x, x ′ ) not being diagonal in spatial dimensions. We emphasize that the µ = 1 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.

ResNet110 (BatchNorm) The operation of BatchNorm on a preactivation (pre-BN) in an MLP can be described as follows:

where B is the batch that x belongs to and | B | is batch size.

The works Yang et al. [63], 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:

̸

where x ′ in the APJN term can be any x ′ = x , since for a 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 K l ( x, x ) → ( K l ( x, x ) -K l ( x, x ′ ) ) . As a 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 :

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 6.

C Discussions on Transformers

As preliminary empirical results, we applied our method to Vision Transformers (ViTs). Figure 7 supports our 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 σ 2 w ; Figure 8 further demonstrates that the attention mechanism is highly sensitive to initialization and thus vulnerable to the gradient vanishing/diverging problem.

For µ = 1 , we have

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

2

w

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

(Pre-LN, with MLP)

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

2

w

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

(Post-LN, with MLP)

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

2

w

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

(Pre-LN, No MLP)

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

2

w

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

(Post-LN, No MLP)

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

1.8

2.0

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

2

w

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

log

(No LN, with MLP)

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

2

w

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

log

(No LN, No MLP)

< -1

0

> 1

D Technical details for Jacobians and LayerNorm

We will drop the dependence of h l i ( x ) on x throughout the Appendices. It should not cause any confusion since we are always considering a single input.

D.1 NNGP Kernel

First, we derive the recurrence relation for the NNGP kernel Eq.(9). As mentioned in the main text, weights and biases are initialized (independently) from standard normal distribution N (0 , σ 2 w / fan _ in) . We then have

by definition.

We would like to prove theorem 2.4, as a consequence of lemma 2.2. The proof of lemma 2.2 can be found in [45].

Proof of theorem 2.4. One can prove this by definition with lemma 2.2.

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 (APJN) ( l > l 0 )

where we used the chain rule and took the expectation value over w l +1 . Next, we take the chain rule again:

where we integrate (by parts) over w l to get the second line. We take N l -1 →∞ and then N l →∞ to get the third line. We rearrange terms and use the Eq.(11) 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 l i is independent of h l -1 i . Thus we proved

The critical line is defined by requiring χ ⋆ J = 1 , where critical points are reached by further requiring χ ⋆ K = 1 .

As we mentioned in the main text, l 0 = 0 is subtle since the input dimension is fixed N 0 , which can not be assumed to be infinity. Even though for a dataset like MNIST, usually N 0 is not significantly smaller than width N l . We show how to take finite O ( N -1 0 ) correction into account by using one example.

Lemma D.1. Consider a one hidden layer network with a finite input dimension N 0 . In the infinite width limit, the APJN is still deterministic and the first step of the recurrence relation is modified to:

where J 0 , 1 = σ 2 w .

Proof.

where to get the result we used integrate by parts, then explicitly integrated over w 1 ij . We have introduced a coefficient of finite width corrections, χ l ∆ , defined as follows.

Definition D.2 (Coefficient of Finite Width Corrections) .

Remark D.3 . Notice that the correction to J 0 , 2 is order O ( N -1 0 ) . If one calculates the recurrence relation for deeper layers, the correction to J 0 ,l will be O ( ∑ l l ′ =0 N -1 l ′ ) , which means the contribution from hidden layers can be ignored in infinite width limit.

The J 0 , 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 (APJN with l 0 = 0 ) . The APJN (with l 0 = 0 ) of a given network can be written as

Note that APJN with l 0 > 0 does not receive the O ( N -1 0 ) correction.

D.3 APJN and gradients

As mentioned in the main text, APJN is an 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) L 2 norm of the gradient of any flattened

parameter matrix θ l , at initialization. In the infinite width limit, one gets

where ∥·∥ 2 denotes the L 2 norm and ∥·∥ F denotes the Frobenius norm.

D.4 LayerNorm on Pre-activations

Definition D.5 (Layer Normalization) .

where γ l i and β l i are learnable parameters.

Remark D.6 . With only LayerNorm, the (1) is simplified to

Remark D.7 . In the limit of infinite width, using the law of large numbers, the average over neurons E [ · · · ] can be replaced by the average of parameter-initializations E θ [ · · · ] . Additionally, in this limit, the preactivations are i.i.d. Gaussian distributed : h l ∼ N (0 , K l ) .

The normalized preactivation then simplifies to the form of Eq.(19).

Remark D.8 . At initialization, the parameters γ l i and β l i take the values 1 and 0 , respectively. This leads to the form in equation (19). In infinite width limit, it has the following form

Lemma D.9. With LayerNorm on preactivations, the gaussian average is modified to

Proof. By definition ˜ h l i is sampled from a standard normal distribution N (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 LayerNorm on preactivations is

Proof.

Theorem D.11. In the infinite width limit the recurrence relation for partial Jacobian with LayerNorm on preactivations is

Proof.

D.5 LayerNorm on activations

The general definition of LayerNorm on activations is given as follows.

Definition D.12 (LayerNorm on Activations) .

Remark D.13 . The recurrence relation for preactivations (Eq.(1)) gets modified to

Remark D.14 . At initialization, the parameters γ l i and β l i take the values 1 and 0 , respectively. This leads to the form

where the first line follows from the fact that at initialization, the parameters γ l i and β l i take the values 1 and 0 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 h l -1 ∼ N (0 , K l -1 = σ 2 w + σ 2 b ) . Again this can be summarized as

Next, we 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 LayerNorm on activations is

Proof.

Theorem D.17. In the infinite width limit the recurrence relation for partial Jacobian with LayerNorm on activations is

where χ l J ≡ σ 2 w E θ [ ϕ ′ ( h l ) 2 ) ] E θ [ ϕ ( h l ) 2 ] -E θ [ ϕ ( h l )] 2 .

Proof.

E Residual Connections

Definition E.1. We define residual connections by the modified the recurrence relation for preactivations (Eq.(1))

where the parameter µ controls the strength of the residual connection.

Remark E.2 . Note that this definition requires N l +1 = N l . We ensure this by only adding residual connections to the hidden layers, which are of the same width. More generally, one can introduce a tensor parameter µ ij .

Remark E.3 . In general, the parameter µ could be layer-dependent ( µ l ). But we suppress this dependence here since we are discussing self-similar networks.

Theorem E.4. In the infinite width limit, the recurrence relation for the NNGP kernel with residual connections is changed by an additional term controlled by µ

Proof.

where we used the fact N l +1 = N l to get the last line.

Theorem E.5. In the infinite width limit, the recurrence relation for partial Jacobians with residual connections has a simple multiplicative form

where the recurrence coefficient is shifted to χ l J = σ 2 w E θ [ ϕ ′ ( h l k ) ϕ ′ ( h l k ) ] + µ 2 .

Proof.

F Residual Connections with LayerNorm on Preactivations (Pre-LN)

We recall the recurrence relation (1):

Theorem F.1. In the infinite width limit, the recurrence relation for the NNGP kernel is then modified to

Proof.

Remark F.2 . For µ < 1 , the recursion relation has a fixed point

where the average here is exactly the same as cases for LayerNorm applied to preactivations without residue connections. l ⋆ labels some very large depth l .

Remark F.3 . For µ = 1 case, the solution of (62) is

which is linearly growing since the expectation does not depend on depth. K 0 is the NNGP kernel after the input layer.

Theorem F.4. In the infinite width limit, the recurrence relation for Jacobians changes by a constant shift in the recursion coefficient.

where for this case

Proof.

Remark F.5 . Assuming l ≥ l ⋆ , we combine (64) and (67) to derive Equation (5) in the Theorem 1.3.

where A = ∑ N l k =1 E θ [ ϕ ′ ( ˜ h l k ) ϕ ′ ( ˜ h l k ) ] and B = σ 2 w N l ∑ N l j =1 E θ [ ϕ ( ˜ h l j ) ϕ ( ˜ h l j ) ] + σ 2 b . Recall that ξ = | log χ ⋆ J | -1 . This directly leads us to Equation (5).

Note that in Equation (5), we discard the average over neurons in A and B , since we are in the infinite width limit.

Remark F.6 . As we mentioned above µ = 1 needs extra care. Plugging µ = 1 into the result (65) and (67) we find out that

Thus, we get everywhere criticality: critical behavior independent of the choice of initialization ( σ b , σ w ) . It also follows that the correlation length ξ diverges in this case. (71) leads to power law behaviour in Jacobians (with exponent ζ ) at large depth. Note that the exponent ζ is not universal.

Remarks F.5 and F.6 together serve as the complete proof of Theorem 1.3

G Critical Exponents

To prove theorem 2.7, we first need to find the critical exponent of the NNGP kernel [45].

Lemma G.1. In the infinite width limit, consider a critically initialized network with an activation function ϕ . The scaling behavior of the fluctuation δ K l ≡ K l - K ⋆ in non-exponential. If the recurrence relation can be expand to leading order δK l as δ K l +1 ≈ δ K l -c n ( δ K l ) n for n ≥ 2 . The solution of δ K l is

where ζ K = 1 n -1 .

Remark G.2 . The constant c n and the order of the first non-zero term n is determined by the choice of activation function.

Proof. We can expand the recurrence relation for the NNGP kernel (9) to the second order of δ K l = K l -K ⋆ on both sides.

Use power law ansatz δ K l = Al -ζ K then

Multiply l ζ K on both side then use Taylor expansion ( l l +1 ) ζ K ≈ 1 -ζ K l

For arbitrary l , the only non-trivial solution of the equation above is

̸

Proof of theorem 2.7. We will assume c 2 = 0 . Then use lemma G.1, we can expand χ l J in terms of δ K l . To leading order l -1

Consider a sufficiently large l . In this case O ( l -1 ) approximation is valid. We write recurrence relations of Jacobians as

When c n = 0 for all n ≥ 2 , from lemma G.1 we have δ K l = 0 . Thus the Jacobian saturates to some constant.

We checked the scaling empirically by plotting J 0 ,l vs. l in a log -log plot and fitting the slope. These results are presented in Fig.1.

H MLP-Mixer

In this section we would like to analyze an architecture called MLP-Mixer [48], which is based on multi-layer perceptrons (MLPs). A MLP-Mixer (i) chops images into patches, then applies affine transformations per patch, (ii) applies several Mixer Layers, (iii) applies pre-head LayerNorm, Global Average Pooling, and an output affine transformation. We will explain the architecture by showing forward pass equations.

Suppose one has a single input with dimension ( C in , H in , W in ) . We label it as x µi , where the Greek letter labels channels and the Latin letter labels flattened pixels.

First of all the (i) is realized by a special convolutional layer, where kernel size f is equal to the stride s . Then, the first convolution layer can be written as

where f is the size of filter and s is the stride. In our example f = s . Notice in PyTorch both bias and weights are sampled from a uniform distribution U ( -√ k, √ k ) , where k = ( C in f 2 ) -1 .

Notice that the output of Conv2d: h 0 µi ∈ R C × N p , where C stands for channels and N p = H in W in /f 2 stands for patches, both of them will be mixed later by Mixer layers.

Next, we stack l Mixer Layers. A Mixer Layer contains LayerNorms and two MLPs, where the first one mixed patches i, j (token mixing) with a hidden dimension N tm , the second one mixed channels µ, ν (channel-mixing) with a hidden dimension N cm . Notice that for Mixer Layers we use the standard parameterization.

First LayerNorm. It acts on channels µ .

where we defined a channel mean E C [ h 6 l ρi ] ≡ 1 C ∑ C ρ =1 h 6 l ρi and channel variance Var C ≡ E C [ ( h 6 l ρi ) 2 ] -( E C [ h 6 l ρi ] ) 2 .

First MlpBlock. It mixes patches i, j , preactivations from different channels that share the same weight and bias.

-6 l +1 : Linear Affine Layer.

where N tm stands for the hidden dimension of "token mixing".

-6 l +3 : Residual Connections.

-6 l +2 : Affine Layer.

Second LayerNorm. It again acts on channels µ .

Second MlpBlock. It mixes channels µ, ν , preactivations from different patches share the same weight and bias.

-6 l +6 . Residual Connections.

Suppose the network has L Mixer layers. After those layers the network has a pre-head LayerNorm layer, a global average pooling layer, and an output layer. The pre-head LayerNorm normalizes over channels µ can be described as the following

We plotted the phase diagram using the following quantity from repeating Mixer Layers:

I Results for Scale Invariant Activation Functions

Definition I.1 (Scale invariant activation functions) .

where Θ( x ) is the Heaviside step function. ReLU is the special case with a + = 1 and a -= 0 .

I.1 NNGP Kernel

First evaluate the average using lemma 2.2

Thus we obtain the recurrence relation for the NNGP kernel with scale invariant activation function.

Global Average Pool over patches i .

Output Layer

Finite fixed point of the recurrence relation above exists only if

For σ 2 w = 2 a 2 + + a 2 -case, finite fixed point exists only if σ 2 b = 0 .

I.2 Jacobian(s)

The calculation is quite straightforward, by definition

where we used the property xδ ( x ) = 0 for Dirac's delta function to get the first line.

Thus the critical line is defined by

As a result where the critical point is located at

I.3 Critical Exponents

Since the recurrence relations for the NNGP kernel and Jacobians are linear. Then from lemma G.1 and theorem 2.7

I.4 LayerNorm on Pre-activations

Use lemma D.9 and combine all known results for scale-invariant functions

is always true. The equality only holds at σ b = 0 line.

In this case,

For ReLU with a + = 1 and a -= 0 , the network is at critical line when

I.5 LayerNorm on Activations

First we substitute K l -1 = σ 2 w + σ 2 b into known results

There is a new expectation value we need to show explicitly

Thus

The critical line is defined by χ ⋆ J = 1 , which can be solved as

For ReLU with a + = 1 and a -= 0

I.6 Residual Connections

The recurrence relation for the NNGP kernel can be evaluated to be

The condition for the existence of fixed point

leads us to

For σ 2 w = 2(1 -µ 2 ) a 2 + + a 2 -, finite fixed point exists only if σ 2 b = 0 . (Diverges linearly otherwise)

The recurrence coefficient for Jacobian is evaluated to be

The critical line is defined as

The critical point is located at (√ 2(1 -µ 2 ) a 2 + + a 2 -, 0 ) .

For ReLU, the critical point is at ( √ 2(1 -µ 2 ) , 0 ) .

I.7 Residual Connections with LayerNorm on Preactivations (Pre-LN)

Again use lemma D.9 and combine all known results for scale-invariant functions

Similar to the case without residue connections

is always true. The equality only holds at σ b = 0 line for µ < 1 .

Notice there is a very special case µ = 1 , where the whole σ b -σ w plane is critical.

J Results for erf Activation Function

Definition J.1 (erf activation function) .

J.1 NNGP Kernel

To evaluate lemma 2.2 exactly, we introduce two dummy variables λ 1 and λ 2 [51].

We use the special case where λ 1 = λ 2 = 1 .

Thus the recurrence relation for the NNGP kernel with erf activation function is

As in the scale-invariant case, finite fixed point only exists when

Numerical results show the condition is satisfied everywhere in σ b -σ w plane, where χ ⋆ K = 1 is only possible when K ⋆ = 0 .

J.2 Jacobians

Follow the definition

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 ( σ b , σ w ) = (0 , √ π 4 ) , with K ⋆ = 0 . Now suppose l is large enough such that the deviation of K l from fixed point value K ⋆ is small. Define δ K l ≡ K l -K ⋆ . Eq.(121) can be rewritten as

From lemma G.1

Next we analyze critical exponent of Jacobians by expanding (123) around K ⋆ = 0 critical point ( σ b , σ w ) = (0 , √ π 4 ) .

To leading order l -1 we have

To find phase boundary χ ⋆ J = 1 , we need to combine Eq.(121) and Eq.(123) and evaluate them at K ⋆ .

One can solve the equations above and find the critical line

Critical point is reached by further requiring χ ⋆ K = 1 . Since χ ⋆ K ≤ χ ⋆ J , the only possible case is K ⋆ = 0 , which is located at

Thus the recurrence relation for partial Jacobian, at large l , takes form

At large l with a non-universal constant c l 0 .

The critical exponent is which is the same as ζ K .

J.4 LayerNorm on Pre-activations

Use lemma D.9, we have

The critical line is then defined by

J.5 LayerNorm on Activations

Due to the symmetry of erf activation function E θ [ ϕ ( h l i ) ] = 0 , we only need to modify our known results.

Thus where the phase boundary is defined by the transcendental equation χ l J = 1 .

J.6 Residual Connections

The recurrence relation for the NNGP kernel can be evaluated to be

Finite fixed point only exists when

Notice that χ ⋆ K ≤ χ ⋆ J still holds, where the equality holds only when K ⋆ = 0 .

The recurrence coefficient for Jacobian is evaluated to be

The critical line is defined as

Critical point is reached by further requiring χ ⋆ K = 1 . Since χ ⋆ K ≤ χ ⋆ J , the only possible case is K ⋆ = 0 , which is located at

Note that for µ = 1 , one needs to put extra effort into analyzing the scaling behavior. First we notice that K l monotonically increases with depth l - the recurrence relation for the NNGP kernel at large l (or large K l ) is

which regulates the first term in (141).

For µ = 1 at large depth

Here C 0 is a constant that depends on the input.

We can approximate the asymptotic form of log J l 0 ,l as follows

where ˜ c = 2 σ 2 w π √ σ 2 w + σ 2 b .

We conclude that at large depth, the APJN for µ = 1 , erf networks can be written as

This result checks out empirically, as shown in Figure 9. 8

J.7 Residual Connections with LayerNorm on Preactivations (Pre-LN)

Use lemma D.9 and the results we had without residue connections for erf with LayerNorm on preactivations.

The critical line is then defined by

≈

0

.

324

σ

w

.

K.1 NNGP Kernel

Use lemma 2.2 for GELU

where from the third line to the fourth line we used integrate by parts twice, and to get the last line we used results from erf activations.

Thus the recurrence relation for the NNGP kernel is

As a result

2.5

5.0

7.5

10.0

12.5

15.0

17.5

20.0

10

1

10 2

10 5

10 8

10 11

10 14

10 17

10 20

Experiment,

2

w

=5.00

Theory,

c

=1.424

Fitting,

c

=1.414

Experiment,

2

w

=2.00

Theory,

c

=0.900

Fitting,

c

=0.851

Experiment,

2

w

=0.79

Theory,

c

=0.564

Fitting,

c

=0.548

K Results for GELU Activation Function

Definition K.1 (GELU activation function) .

K.2 Jacobians

Follow the definition

χ l J = σ 2 w E θ [ ϕ ′ ( h l i ) ϕ ′ ( h l i ) ]

where we dropped odd function terms to get the third line, and to get the last line we used known result for erf in the second term, integration by parts in the third term.

Here to get the critical line is harder. One can use the recurrence relation for the NNGP kernel at fixed point K ⋆ and χ ⋆ J = 1

Cancel the arcsin term, σ w and σ b then can be written as a function of K ⋆

One can then scan K ⋆ to draw the critical line.

In order to locate the critical point, we further require χ ⋆ K = 1 . To locate the critical point, we solve χ ⋆ J -χ ⋆ K = 0 instead. We have

which has two non-negative solutions out of three

One can then solve σ b and σ w by plugging corresponding K ⋆ values.

K.3 Critical Exponents

GELU behaves in a different way compared to erf. First we discuss the K ⋆ = 0 critical point, which is located at ( σ b , σ w ) = (0 , 2) . We expand Eq.(152), and keep next to leading order δ K l = K l -K ⋆

From lemma G.1

which is not possible since δ K l ≥ 0 for this case. This result means scaling analysis is not working here.

Next, we consider the other fixed point with K ⋆ = 3+ √ 17 2 at ( σ b , σ w ) = (0 . 416 , 1 . 408) . Expand the NNGP kernel recurrence relation again.

Following the same analysis, we find

Then

Looks like scaling analysis works for this case, since K ⋆ > 0 . The solution shows that the critical point is half-stable[45]. If K l < K ⋆ , the fixed point is repealing, while when K l > K ⋆ , the fixed point is attractive. However, the extremely large coefficient in the scaling behavior of δ K l embarrasses the analysis. Since for any network with a reasonable depth, the deviation δ K l is not small.

Now we can expand χ l J at some large depth, up to leading order l -1 .

where c l 0 is a positive non-universal constant.

Critical exponent

Which in practice is not traceable.

K.4 LayerNorm on Pre-activations

Use lemma D.9, we have

The critical line is then at

K.5 LayerNorm on Activations

First, we need to evaluate a new expectation value

where we used integration by parts to get the result.

The other integrals are modified to

One can then combine those results to find χ l J

The critical line defined by χ l J = 1 , one can numerically solve it by scanning over σ b and σ w .

K.6 Residual Connections

The recurrence relation for the NNGP kernel is

Fixed point exists if

The recurrence coefficient for Jacobian is

Phase boundary is shifted

One can again scan over K ⋆ to draw the critical line.

In order to locate the critical point, we further require χ ⋆ K = 1 . To locate the critical point, we solve χ ⋆ J -χ ⋆ K = 0 instead. We have

which has two non-negative solutions out of three

One can then solve σ b and σ w by plugging corresponding K ⋆ values.

K.7 Residual Connections with LayerNorm on Preactivations (Pre-LN)

Use lemma D.9 and the results we had without residue connections for GELU.

The critical line is then at

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.

10 0

10 1

10 2

2×10 1

3×10 1

4×10 1

(A)

LN-erf, experiment

Theory

Fitting,  =

0.007

10 0

10 1

10 2

3×10 1

4×10 1

5×10 1

6×10 1

(B)

erf-LN, experiment

Theory

Fitting,  =

0.009

In figure 11, we tested 6 k samples from CIFAR-10 dataset[28] with kernel regression based on neural tangents library [40] [31] [41]. Test accuracy from kernel regression reflects the trainability (training accuracy) with SGD in the ordered phase. We found that the trainable depth is predicted by the correlation length cξ with LayerNorm applied to preactivations, where the prefactor c = 28 . The prefactor we had is the same as vanilla cases in [55]. The difference is from the fact that they used log 10 and we used log e .

In figure 12, we explore the broad range in σ 2 w of the performance of MLP network with erf activation function and LayerNorm on preativations. The network has depth L = 50 and width N l = 500 ; and is trained using SGD on Fashion MNIST. The learning rates are chosen based on a logarithmic scan with a short training time.

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

2

w

20

40

60

80

100

Depth

LN-ReLU

28

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

2

w

20

40

60

80

100

Depth

LN-erf

28

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

2

w

20

40

60

80

100

Depth

LN-GELU

28

8.0

13.0

18.1

23.1

28.2

33.2

38.2

43.3

48.3

10 0

10 1

10 2

10 3

10 4

10 5

2

w

20

40

60

80

Training Accuracy (%)

LN-Erf,

2

b

=0.1

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThis 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\nWEAKNESSES:\nThese 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\nQUESTIONS:\nDear 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.\n\nLIMITATIONS:\nTo 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nStrengths:\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\nWEAKNESSES:\nWeaknesses:\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\nQUESTIONS:\nSuggestions 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\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nSome 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.\n\nQUESTIONS:\nSee the weaknesses part.\n\nLIMITATIONS:\nSee the weaknesses part.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}