Critical Initialization of Wide and Deep Neural Networks through Partial Jacobians: General Theory and Applications
Authors: Darshil Doshi
Tianyu He
Andrey Gromov
Venue: Submitted to ICLR 2023 (2023)
Keywords: Criticality
Gaussian Process
Jacobian
LayerNorm
Residual connections

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 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. 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 the MLP-Mixer architecture and show that it is everywhere critical.

Under review as a conference paper at ICLR 2023
CRITICAL INITIALIZATION OF WIDE AND DEEP NEU- RAL NETWORKS THROUGH PARTIAL JACOBIANS: GENERAL THEORY AND APPLICATIONS
Anonymous authors Paper under double-bling review
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 formulate criteria for selecting hyperparame- ters, such as variances of weights and biases, as well as the learning rate. These criteria rely on the notion of criticality defined for deep neural networks. In this work we describe new practical way to diagnose criticality. We introduce par- tial Jacobians of a network, defined as derivatives of preactivations in layer with respect preactivations in layer lo S I. 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; including fully connected, convolutional and attention layers. Using these tools we show quan- titatively that proper stacking of the LayerNorm (applied preactivations) and residual connections leads to an architecture that is critical for any initialization. Finally, we apply our methods to analyze the MLP-Mixer architecture and show that it is everywhere critical.
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(), in this limit, is Gaussian process (Neal 1996. Lee et al. 2018) with kernel sometimes referred to as neural network Gaussian process (NNGP) kernel (Lee et al. 2018) determined by the network architecture and hyperparameters (e.g depth, precise choices of layers and the activation functions, well as the distribution of weights and biases). Similar line of reasoning was earlier developed for recurrent neural networks (Molgedey et al. 1992).. Furthermore, for special choices of parameterization and MSE loss function, the training dynamics under gradient descent can solved exactly terms of the neural tangent kernel (NTK) (Jacot et al. 2018 Lee al. (019).. 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 these quantities, and empirical investigation of the training dynamics of wide networks (Novak et al., al., 20186) Xiao et al., 2018 Hron et al. 2020 Dyer & Gur-Ari 2019)
Andreassen & Dyer [020 Lewkowycz
Gur-Ari 2020 Aitken
Gur-Ari [2020] Geiger et al.
2020] Hanin [202 Roberts et al. 2022 Yaida 2020 Shankar et al., 2020 Arora et al. 2019b a.
Lee et al.) 2020 Yang et 2018 Yang & Hu [021 Yang 2019h 9b a Matthews et al. [2018
Garriga-
Alonso et al. 2018 Allen -huu et al. [2019] suchida et al. 2021
Martens et al. 20211)
One important result that arose from these works is that the network architecture determines the most
appropriate initialization of the weights and biases (Poole et al. 2016; Schoenholz et al.) 2016 Lec jtt al., 2018). To state this result, we consider networks with/without LayerNorm (Ba et al.) 2016) and residual connections (He et al. 2016); the preactivations for which can be defined as follows
Ni
N ++ (x) = Σ w ((H((z)) + + + ph((2),

---

Under review as a conference paper at ICLR 2023
where htj LayerNorm(H,) and the parameter μ controls the strength of residual connections. For the input layer: (x)) Eو= wije, + b7. In the (++ 1)-th layer, weights w/j E R RRN4+1 X and biases 6,+1 E I are taken from normal distributions o o²/N") and (0,,, of), respec- tively. Hyperparameters on and 0 need be tuned. <(-) the activation function and IR Nox1 is the input. For results discussed in this work, can be sampled from either realistic (i.e. highly correlated) dataset or a high entropy distribution.
For network of depth L, the network function is given by f(x) = (x)). Different network ar- chitectures and activation functions, é, lead to different "optimal" choices of (aw, Tb). The optimal choice can be understood. using the language of statistical mechanics, as a critical point (or mani- fold) in the a6-0w plane. The notion of criticality becomes sharp as the network depth, L, becomes large. Criticality ensures that both NNGP and the norm of gradients remain (L)) as the network gets deeper (Roberts et al., 2022). Very deep networks will not train unless initialized critically, since the gradients explode vanish exponentially. Moreover, 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 partial Jacobian.
Definition 1.1. Let (z) be preactivations of neural network f(x). The partial Jacobian J/ is defined as derivative of preactivations at layer / with respect to preactivations at layer to
(11(2))
(x)
oh (x)
The partial Jacobian is a random matrix with vanishing mean at initialization. We introduce deter- ministic measure of the magnitude of ! its squared Frobenius norm, averaged over parameter- initializations.
Definition 1.2. Let J be a partial Jacobian of a neural network f(x). Averaged partial Jacobian norm (APJN) is defined as
N₂
Nio
Old,(x)
27
flo.((z) = Eg
(3)
N
ah (x))
j=1
where Eg indicates averaging over parameter-inittilizaaiiee initializations.
In what follows, we show that criticality. studied previously in literature, occurs when APJN either remains finite, or varies algebraically as I becomes large. To prove this we derive the recurrence relation for (z) in the limit N; oc and analyze at large depth. Algebraic behaviour of APJN with depth is characterized by an archieecturedeppedderr critical exponent, Ç. so that Jio. (2) R I-S_ Such behaviour is familiar from statistical mechanics when system tuned [ardy][196] Away from criticality, there two phases; ordered and chaotic. In the ordered phase APJN vanishes exponentially with depth, whereas in the chaotic phase APJN grows exponentially
critical point
glost no ch,ett
(4)
Here& & is the correlation length. It characterizes how fast gradients explode or vanish.
Theorem 1.3 (Main result). Let fxx be a deep MLP network with Lipschitz continuous activation
o(.). Assume that the LayerNorm is applied to preactivations and there are residual connections
with strength μ acting according to (.. In the limit N₁ + so the correlation length is bounded from below for of < co
E log (1 - 12) A/ + 12
(5)
2

---

Under review as a conference paper at ICLR 2023
where the non-negative constants A and B given by
A F.g
Σ (hL))
B= Eag
(hL))
(6)
N₁
where )()) is the derivative of Φ(-). When μ = the correlation length diverges and the network is critical for any initialization, with (
In practice Theorem 1.3 means that different choices of initialization bear no effect on trainability of the network provided that LayerNorm and residual connections arranged 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. (011)) LayerNorm also alters the criticality of Jo (x). It was noted in Xu et al. (2019) that LayerNorm (applied to preactivations) regularizes the backward pass. We formalize this observation by showing that LayerNorm (applied preactivations) dramatically enhances correla- tion length (which is not the case for LayerNorm applied to activations). This can be seen from Theorem I.3 setting μ 0. When residual connections of strength
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 Theorem 4.3) A version of this result was discussed in Yang & Schoenholz (011)) When residual connections are introduced on top of LayerNorm, the correlation length & is further increased. If residual connections have strength 1 the network enters critical phase for any initialization. Importance of correct ordering of LayerNorm, residual connections and attention layers was discussed Xiong et al. (2020). Several architectures with the same order of GroupNorm and residual connections were investigated in Yu et al. (2021).
The partial Jacobian has been used to study generalization bounds in Arora et al. (018). The Jacobian norm (i.e. || (12) of trained feed-forward neural networks was studied in Novak et al. (2018a), where was correlated with generalization. Partial Jacobians with 4o = (- were studied in the context of RNNs (Chen al. 2018 Can et al. 2020), where they were referred to state-to- state Jacobians.
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 small aspect ratio, the effect of the spectral density of Jacobian becomes important the depth L becomes very large. |Pennington et al. (2018) study the spectrum of the input-output Jacobian for MLPs. Xiao et al. (2018) extend the analysis to CNNs, showing that very deep vanilla CNNs can be trained by achieving "dynamical isometry".
RECURRENCE RELATIONS
Here derive the infinite width recurrence relations for the APJN and the norm of preactivations. We use (I) in its simplest form, which has no LayerNorm and = O.
Definition 2.1. define averaged covariance of preactivations as follows
N: MoN
(x, = Egg
1 N
Σ
Lemma 2.2. When N₁ - for 1,.., 1, the expectation value over parameter initializa- tions for general function of preactivations: O'h' (x)), ) can be expressed the averaging over the Gaussian process 1 (x(x) with covariance (x,xx))
1
/ an' dh'o
(o) (x..).
Eg [o(h((=))
2 (2,x)
This result has been established in Lee et al. (011). Note that the density in 8 only depends on the diagonal part of the covariance matrix, K K'Tr,, F.). We will refer to KC' (2,,a)) as NNGP kernel.
3

---

Under review as a conference paper at ICLR 2023
Remark 2.3. In the infinite width limit the means appearing in (10) (2) self-averaging and, therefore, deterministic. They converge in distribution to their averages over parameterizations.
N₂ N
N
1
N1-00
Eg
p(h)) N
= Eo [o(n]]]]
(9)
When performing analytic calculations we use the infinite width convention; whereas in our finite- width experiments we explicitly average over initializations of 6'_ Theorem 2.4. With Lemma 2.2 in the infinite width limit, the NNGP kernel Kt+ (Ir, is determin- istic, and can determined recursively via
N
K xx, o
1 Nr
Σ o(h(())??
+ +
(10)
Theorem 2.5. Let (() be an MLP network with Lipschitz continuous activation function ((x). In the infinite width limit, APJN Jla.41((x) (xx) is deterministic and satisfies recurrence relation
flout ('od))))
(11)
where the factor X3 given by
X
Eg
10
( (d(((())))
(12)
Theorem 2.4 due to Lee et al. (2018). Theorem 2.5 is new and is valid only in the limit of infinite We will drop the explicit dependence on a: to improve readability.
width. The proof is in Appendix B
12) are evaluated using 8 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 (1O)((I2) can be found in the Appendix. A subtlety emerges in (TT) when to invariant activation functions. This subtlety is discussed in the Appendix B
The expectation values that appear in 10
where a correction of the order O(N, arises for non-scale
When the depth the network becomes large, the I-dependence of the expectation values that appear in (7). 122 saturate to (possibly infinite) constant value; which means that IC', Jlo.l and X'y have reached a fixed point. We denote the corresponding quantities KC* T .x. The existence of fixed point is not obvious and should be checked on a case by case basis. Fixed point analvsis for K was done in Poole et al. (2016) for bounded activation functions and in Roberts et al. (0022) for the general case. The stability is formulated in terms of
ak!+1
XiC
(13)
aKC K'=K-
The norm of preactivations remains finite (or behaves algebraically) when x²: 1.
Eq. (1]) nicely expresses Jla.+1 as a linear function of Jla. The behaviour of Jood+1 at large 1 is determined by Xyy When x²y > partial Jacobians diverge exponentially, while for xyy < 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
(14) Eq. (14) is an implicit equation on Fb, Tw and generally outputs a critical line in σg-O plane. The parameter X"y has to be calculated on a case-by-case basis using either (12) or the method presented in the next section. Everywhere on the critical line, Flo. saturates a constant or behaves algebraically.
When the condition xi= = 1 is added, are left with a critical poin[] This analysis of criticality at infinite width agrees with Roberts et al. (2022), where X1 is to be identified with X;; and Schoen- holz et al. (2016); Martens al. (202T) where their analysis based on the equivalent X1 C'(I) only works for bounded activation functions. In particular, condition 14 together with Xx: = ensures that NTK is O(1) initialization.
'Scalc-invariant activation functions more forgiving: away from the critical point K. scales algehraically
with I.
4

---

Under review as a conference paper at ICLR 2023
2.1
EMPIRICAL DIAGNOSTIC OF CRITICALITY
APJN Tloi! provides a clear practical way to diagnose whether the network is critical or not. Proper choice of ln and allows us minimize the non-universal effects and cleanly extract x²-
Recurrence relation supplemented with the initial condition J'o solved as
Xy. can be formally
1-1
Jlod = IIxx X2y
(15)
We would like to obtain an estimate of X accurately as possible. To that end, imagine that for some l' > to the fixed point has been essentially reached and X'y =s Xy.. Then the APJN
glod = (x3)
II
X
(16)
t=
depends on the details of how the critical point is approached; which are encoded in the last factor. Proposition 2.6. If the network f (x) is homogeneous, i.e. consists of a (possibly complex) block of layers, periodically repeated L times. Then the penultimate APJN provides an accurate estimate of X²::
J1-
X²g
(17)
10000
This is direct consequence of combining Appendix Cor numerical justification.
10
and
12
as L goes infinity. See Figure 7 in
Proposition 2.6 is the central result of this section and will heavily used in the remainder of this work.
Note that for deep networks, away from criticality, APJN takes form
glool a ctett
= log [ogox>>- x
(18)
where Ct, is non-universal constant that depends on to. If the sign in (18) is positive (x > 1) the network is in the chaotic phase, while when the sign is negative (x's < <) the network is in the ordered phase. & has the meaning of correlation length: on the depth scale of approximately &g the gradients remain appreciable, and hence the network with the depth of z kg will train.
We used 17) map out the dy-Ow phase diagrams of various MLP architectures. The partial Jacobians are calculated numerically with N₁ = 500, L = 50 and averaged over initializations. The details further elaborated in Appendix A The location of the critical line agrees remarkably well with our infinite width calculations. Results are presented in Fig. [ One fortunate outcome of both theory and experiment is that when LayerNorm applied preactivations, ReLU networks can still be initialized using He initialization (He et al. 2015) which, our convention, (v2, 0).
At criticality, X and the correlation length diverges; indicating that gradients can propagate arbitrarily far. A more careful analysis of non- non-linear corrections shows that APJN can exhibit alge- braic behaviour 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 Xy saturates to fixed value X3 1. If we interested in Flod with / 4 = ((L) then it essential know how exactly X,y approaches
Theorem 2.7. Assume that deep neural network fxx) is initialized critically, Then oc asymp- totics of APJN is given by
= =x)(())=0(115 O(-
(19)
where Ç is the critical exponent Roberts et al. (2022). Critical exponents can be determined analyti- cally in the limit of infinite width. (For detailed discussion, Appendix C
5

---

Under review as a conference paper at ICLR 2023
3 LAYER NORMALIZATION
The fact that critical initialization is concentrated on a single point (g) 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 (Ba et al. 2016) and GroupNorm (Wu & He 2018). 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 B). Depending this choice, criticality will occur on different critical lines in σ6-0 plane. When LayerNorm is applied to preactivations the correlation length is enhanced, allowing train much deeper networks even far away from criticality.
The LayerNorm applied to preactivations takes the following form Definition 3.1 (Normalized preactivations).
Rt
(20)
v [ E
In the limit of infinite width B[44] 0 and where we have introduced B[W] N. Em m4
B[(h')] = Kt, defined according to
7
Normalized preactivations, h/. are distributed according to W(0, 1) for all I, ow. Sb. The norms are, therefore, always finite and the condition X = 1 is trivially satisfied. This results in a critical line rather than a critical point.
The recurrence relations
TO
12 for the NNGP and partial Jacobians are only slightly modified
N
Ni
k² = Eo
1
o(h!)²
+ +oz
X
Eg
o' d'(h))²
(21)
N
NI
Assuming that the value of X'y at the fixed point is xý. the network is critical when
14 holds.
X 21 changes very slowly with and is also bounded from below, elaborated in the next section. Thus, X remains close to for a very wide range of hyperparameters. Consequently, the correlation length large even away from criticality. This leads much higher trainability of deep networks with LayerNorm on preactivations even away from criticality.
RESIDUAL (SKIP) CONNECTIONS
Adding residual connections between the network layers is widely used technique facilitate the training of deep networks. Originally introduced (He et al.] (2016) in the context of convolutional neural networks (LeCun al.] (998) (CNNs) for image recognition, residual connections have since been used in a variety of networks architectures and tasks.
Consider 1 with non-zero y and without LayerNorm layers. Then the recurrence relations I0 for the NNGP kernel and x"y are modified as follows
12
Ni Σ
N Σ
k!+1 = Eo
o(h))
+ +of + u2' XS = o
d'(hh)
+ +
(22)
N
N
Remark 4.1. When μ < 1, the fixed point value of NNGP kernel is scaled by (1 - -12) For = 1, the critical point formally (0,0).
Remark 4.2. For Consequently, APJN exponentially diverges as a function of depth / for all aw > 0. In this case, a. needs to be taken sufficiently close ( to ensure trainability large depths.
22 implies that X 2 1, where the equality holds the ow 0 axis.
When μ residual connections amplify the chaotic phase and decrease the correlation length away from criticality for unbounded activation functions.
Solving the recurrence relations (22 for erf activation, find an effect observed in Yang & Schoen- holz (2017) for tanh activation. They noted that tanh-like MLP networks with skip connections "hover over the edge of chaos". We quantify their observation follows.

---

Under review as a conference paper at ICLR 2023
4.-
0.6
36-
6.
0.3
2.0
44-
12
82
12
Figure 1: Xyy phase diagrams for ReLU (first row), erf (second row) and GELU (third row): with residual connections of variable strengths (first three columns) and with LayerNorm (last three columns) are shown. The solid lines indicate the critical lines obtained through infinite width limit calculations; while the stars indicate the critical points. The dotted lines in the rightmost column correspond to the critical lines for p. < 1 case. For networks with LayerNorm and μ X holds on the entire dy-Ow plane, for all activation functions that we considered. We also note that for erf activation, the case without LayerNorm subcritical and has a large correlation length.
{0.0, 0.9, 1.0). Both cases: without LayerNorm
Theorem 4.3. Let f(z) be a deep MLP network with erf activation function and residual connec- tions of strength μ = 1. Then in the limit N1 - 3
The NNGP kernel K' linearly diverges with depth I.
Xyy approaches from above (as can be from Fig. / (o + ) is a non-universal constant.
1 Xyy 1 2/V/, where
APJN diverges stretched exponential: Flost O(oVF), where À 1/(46?) is the new length scale.
We will refer this case subcritical. Although X reaches 1, the APJN still diverges with depth faster than any power law. The growth is controlled by the new scale A. To control the gradient we would like to make large, which can be accomplished by decreasing dw- In this case the trainability is enhanced (see Fig. 2 Similar results hold for tanh activation function Yang & Schoenholz 2017), however in that case there no explicit expression for č.
RESIDUAL CONNECTIONS + LAYERNORM
In practice, common use combination of residual connections and LayerNorm.
Using (I. the recurrence relations (T0) 12) for the NNGP and partial Jacobians are modified as follows
KC+1 E
(kt)2 N Σ
o0 + 1K X²
a2 KT
E.
1 N
Σ
+ 1=2
(23)
Remark 5.1. For J < 1. 23 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 X closer to 1 the entire as-ow plane (as be seen from

---

Under review as a conference paper at ICLR 2023
83.0
6.-
45
65
0.5
60
12-
14
14 19
10
Figure 2: Trainability (Training Accuracy) of deep MLP networks (N₂ = 500, = 50) featuring ReLU (first row), erf (second row) and GELU (third row); with residual connections of variable strengths μ = {0.0, 0.9, 1.0]. Both cases: without LayerNorm (first three columns) and with Lay- erNorm (last three columns) shown. Cases with LayerNorm and μ (last column) train at all values of oe and of we considered, in agreement with theory.
Fig. [. Therefore the correlation length & is improved in both the phases, allowing for training of deeper networks. At criticality, Jacobians linearly diverge with depth.
As mentioned before, the combination of LayerNorm and residual connections dramatically enhances correlation length, leading to more stable architecture. This observation is formalized by Theorem I.3 The proof leverages the properties of solutions of (23 close to the fixed point, and is fleshed out in Appendix F
Remark 5.2. When μ 1, the correlation length diverges for any initialization.
Remark 5.2 provides an alternative perspective on architecture design. On the one hand, given a add extra
neural network architecture layers, such as a combination residual connections and I LayerNorm, to ensure that the network is always critical and will train well no matter which initialization scheme used.
use (17 to initialize it critically. Alternatively,
Remark 5.3. When μ 1, the condition X holds the entire σ) - ow and for any activation function & (see Fig. 1. NNGP kernel diverges linearly, while APJN diverges algebraically with the critical exponent of (1)) The exact value of the critical exponent depends on the activation function and the ratio <h.. The trainability is dramatically enhanced as can be seen from Fig. 2
Remark 5.4. Networks with BatchNorm (Toffe Szegedy 2015), used in conjunction with residual connection of strength μ = 1, also enjoy this everywhere criticality and enhanced trainability Yaag
al. 2018 He et 2022).
MLP-MIXER
MLP-Mixes architecture is recent example of MLP approach to computer vision (Tolstikhin al.) 2021). Its main ingredients are: patches, MLP layers, LayerNorm and residual connections. As such it can be analyzed using the tools and results presented above. The detailed summary of the MLP-Mixer is presented in the Appendix[F while here we will state the results.
When μ 1 the MLP-Mixer is everywhere critical due the interaction between LayerNorm and residual connections. When < we can identify a critical line by numerically evaluating X'y using (17). The phase diagrams are presented in Fig.
To illustrate the importance of critical initialization we trained MLP-Mixer at μ = 1 for various initializations, including highly unconventional of 10, 10. While the final performance
8

---

Under review as a conference paper at ICLR 2023
varies by few percent, the network trains well at large depth of L 100 mixer blocks. We also trained MLP-Mixer at μ = 0.5. In this case, the model trains well at L = 100 when initialized critically: while far away from the critical line (o 10, o0 10) it starts training after 30 epochs, cases. The learning curves are presented in Fig. 3 We emphasize that were interested in trainability and, consequently, did not tune the hyperparameters to achieve the best generalization.
and the gap between it and critical initialization is larger than
55
30-
1-10
55
1.5
415
1-61-d-31
40-120
80
280
Figure 3: from left to right: (1)(2) μ 0.5 and μr = 1.0 phase diagrams plotted using X" from repeating Mixer Layers of the GELU MLP-Mixer. Black line indicates the empirical phase bound- ary. Stars indicate points we selected to train on CIFAR-10. (o" = 10, of = 10) point is outside the phase diagrams. (3)(4) μ 0.5 and indicate training and validation accuracies, respectively. All the networks are = 100 blocks deep, except for one, which is L 32 blocks deep; all networks have 10 million parameters.
MLP-Mixer training curves. Solid and dashed lines
7 CONCLUSIONS
We have introduced partial Jacobians and their averaged norms as a tool analyze the propagation of gradients through deep neural networks at initialization. Using APJN evaluated close to the output, J -1.-2.1-11 Xyy we have introduced 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 (Poole et al. 20166] Roberts al., 2022; Martens et al. 202T). APJÑ will play an important role in quantifying the criticality of inhomogeneous (i.e. no periodic stacking of blocks) networks.
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 critical line, making the initialization prob- lem much easier, (ii) LayerNorm applied to preactivations enhances correlation length leading to improved trainability, (iii) combination of μ = residual connections and erf activation function enhances correlation length driving the network to a subcritical phase with APJN growing according stretched exponential law, (iv) combination of residual connections and LayerNorm drastically increases correlation length leading to improved trainability, (v) when μ = and LayerNorm is applied to preactivations the network is critical on the entire Ty-O plane.
We have considered the example of modern high performance architecture - the MLP-Mixer. We showed that is critical everywhere and not sensitive to initialization at μ due to the interaction between LayerNorm and residual connections. We have also studied MLP-Mixer at I 0.5 and showed that it is critical along a line (as expected for an architecture with LayerNorm). We demonstrated empirically that deep (100 blocks) MI.P-Mxer trains for variety of initializations at = but only trains well close to the critical line for μ = 0.5.
Our work shows that architecture can be designed to have large correlation length leading to guaranteed trainability with SGD for any initialization scheme.
8 ETHICS STATEMENT
We have read the Code of Ethics; have adhered to while writing this paper; and will adhere to during the paper submission, review, and discussion process.

---

Under review as a conference paper at ICLR 2023
9 REPRODUCIBILITY STATEMENT
We provide the details of all the experiments in Appendix [A] including hyperparameter choices, GPU specifications and GPU hours. We also provide the source code in the Supplementary Material.
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 over-
parameterization. In International Conference on Machine Learning, pp. 242-252. PMLR, 2019.
Anders Andreassen and Ethan Dyer. Asymptotics of wide convolutional neural networks. arXiv preprint
000i::2000.00665, 2020.
Sanjeev Arora, Rong Gc, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via compression approach. International Conference on Machine Learning, pp. 254 263. PMLR, 2018.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact com- putation with an infinitely wide neural net. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, pp. 8141-8150, 2019a.
Sanjeev Arora, Simon S Du. Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang, and Dingli Yu. Harness- ing the power of infinitely wide deep nets on small-data tasks. In International Conference on Learning Representations, 2019h.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.
Layer normalization.
arXiv preprint
arXiv: 107.00650 2016.
Tankut Can, Kamesh Krishnamurthy, and David Schwab. Gating creates slow modes and controls phase-space complexity in grus and Istms. In Mathematical and Scientific Machine Learning, pp. 476-511, PMLR, 2020.
John Cardy. Scaling and renormalization in statistical physics, volume Cambridge university press, 1996.
Minmin Chen, Jeffrey Pennington, and Samuel Schoenbolz. Dynamical isometry and field theory of rmns: Gating enables signal propagation recurrent neural networks. In International Conference on Machine Learning, pp. 873-882, PMLR, 2018.
Ethan Dyer and Guy Gur-Ari. Asymptotics wide networks from feynman diagrams. International Con- ference on Learning Representations, 2019. 2019..
Adrià Garriga-Alone Alooso, Carl Edward Rasmussen, and Laurence Aitchison. Deep convolutional networks as shallow gaussian processes. arXiv preprint
055877 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()):023401, 2020.
Boris Hanin. Random neural networks in the infinite width limit gaussian processes. arXiv preprint arXiv:2107.01566 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human- level performance on imagenet classification. Proceedings of the IEFE international conference on com- puter vision, pp. 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), pp. 770-778, 2016. doi: 10. 109/......00...9..
Tianyu He, Darshil Doshi, and Andrey Gromov. Autoinit: Automatic initialization via jacobian tuning, 2022.
Judy Hoffmman, Daniel A. Roberts, and Sho Yaida. Robust learning with jacobian regularization. arXiv preprint arXiv: 1908:000.0000 02729, 2019.
Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nugp and ntk for deep attention networks. In International Conference Machine Learning, pp. 4376-4386. PMLR, 2020.
10

---

and Under review as a conference paper at ICLR 2023
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. International conference on machine learning, pp. 448 456. PMLR, 2015.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization neural networks. arXiv preprint 180.::1800.0777,, 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 applicd to document recognition. Proceedings of the IEEE, 86(11) 36(11):227822324, 1998.
Jachoon 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 Represen- tations, 2018
Jachoon Lec, Lechao Xiao, Samuel Schocnholz, Yasaman Bahri, Roman Novak, Jascha 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.
Jachoon 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 1.2 regularization, arXiv preprint ar0iv:2006.000 08643, 2020.
James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszez, Valentin Dalibard, Jascha Sohl- Dickstein, and Samuel S Schoenholz. Rapid training of deep neural networks without skip connections or normalization layers using deep kereel shaping. arXiv preprint arXiv:21 10.01765, 2021.
Alexander G de G Matthews, Jiri Hron, Mark Rowland, Richard E Turner, and Zoubin Ghahramani. Gaussian process behaviour wide deep neural networks. In International Conference on Learning Representations, 2018.
Lutz Molgedey. 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, pp. 29-53. Springer, 1996.
Roman Novak, Yasaman Bahri, Daniel A Abolañia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Sensitiv- ity and generalization neural networks: an empirical study. In International Conference on Learning Representations. 2018a.
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 Learning Representations, 2018b.
Roman Novak, Lechao Xiao, Jiri Hron, Jachoon Lee, Alexander Alemi, Jascha Sohl-Dickstein, and Samuel Schoenholz. Neural tangents: Fast and easy infinite neural networks 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é-Buus Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Informa- tion Processing Systems 32, pp. 8024 8035, Curran Associates, Inc., 2019.
Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. The emergence spectral universality in deep networks. In Amos Storkey and Fernando Perez- (eds.), Proceedings of the Twenty-Fiss International Conference on Artificial Intelligence and Statistics. volume 84 of Proceedings of Machine Learning Re- search, pp. 1924-1932. PMI.R, 09-11 1 Apr 2018. URL penni ngten 1 8a. htm1
ngs ml presss/887

---

Under review as a conference paper at ICLR 2023
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expres- sivity in deep neural networks through transient chaos. Advances in neural information processing systems. 29:3360- -33300-3668,.20 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-Dicke bohl-Dickstein. Deep information propaga- tion. 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. pp. 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. MIp-mixer: An all-mlp architecture for vision. arXiv preprint arXiv:2 105.0160) 2021.
Russell Tsuchida, Tim Pearce, Chris van der Heide, Fred Roosta, and Marcus Gallagher. Avoiding keree fixed points: Computing with elu and gelu infinite networks. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 9967-9977, 2021.
Christopher Williams. Computing with infinite networks. In M. C. Mozer, M. Jordan, and T. Petsche (eds.),
Advances 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), pp. 3-19, 2018.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist novel image dataset for benchmarking machine learning algorithms, 2017.
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein Samuel Schoenholz. and Jeffrey Pennington. Dynamical isometry and mean field theory of cnns: How train 10,000-layer vanilla convolutional neural networks. In International Conference on Machine Learning, pp. 5393-5402. PMLR, 2018.
Lechao Xiao, Jeffrey Pennington, and Sam Schoenholz. Disentangling trainability and generalization in deep learning, 2020.
Ruibin Xiong, Yunchang Yang. Hc, 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, pp. 10524- 10533. PMLR, 2020.
Jingjing Xu, Xu Sun, Zhiyuan Zhang, Guangxiang Zhao, and Junyang Lin. Understanding and improving layer normalization. arXiv preprint (9iv:/91110000)) 2019.
Sho Yaida. Non-gaussial processes and neural networks at finite widths. In Mathematical and Scientific Machine Learning. pp. 165 192. PMLR, 2020.
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kemel derivation. arXiv preprint arXiv:1902.047330 2019a.
Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. 2019b.
Greg Yang and Edward Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning, pp. 117271 11737, PMLR, 2021.
Greg Yang and Samuel Schoenholz. Mean field residual networks: On the edge chaos. arXiv preprint arXiv:1712.08969 2017.
Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel Schoenholz. A mean field theory of batch normalization. In International Conference on Learning Representations, 2018.
Weihao Yu, Luo, Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, and Shuicheng Yan. Metaformer is actually what you need for vision. arXiv preprint arXiv:2111.11416. 2021.
12

---

Under review as a conference paper at ICLR 2023
A EXPERIMENTAL DETAILS
We implemented our methods using Py PyTorch (Paszke et al., 2019) hooks and an efficient Jacobian approximate algorithm (Hoffman et al., 2019).
Figure All the phase diagrams were plotted using X generated from networks with L 50 and Nj = 500. We used hooks to obtain the gradients that go into calculating x X data was averaged 100 different parameter-imitializatoon Inputs were generated from normal Gaussian distribution and have dimension 28 x 28. Generating the data for the figure took approximately days on Google Colab Pro (single Tesla P100 GPU).
Figure 2 In all cases, networks are trained for 10 epochs using stochastic gradient descent with CrossEntropy loss. We used the Fashion MNIST dataset Xiao et al. (2017). All networks had depth
L = 50 and width N; = 500. The learning rates were logarithmically sampled within (10-B, 1). Generating the data for the figure took approximately days on Google Colab Pro (single Tesla
P100 GPU).
(e) made the phase diagram for MLP-Mixer with blocks and averaged over 100 different garammter-initializaiin
Figure 3
(c)(d)We used network with T, = 100, patch size 4 x 4, hidden size = 128, two MLP dimensions Nem Nem = 256. The = 32 point has doubled widths. All networks have 10 million parameters. Notice that for all Mixer Layers
used NTK initialization. weight decay À = 10 was selected from {10-5 10 1104) mixup rate
We trained all cases on CIFAR-10 10 dataset using vamilla SGD paired with CSE. Batch size bs = 256, 0.8 selected from (0.4,0.8), We also used RandAgument and horizontal flip with default settings in PyTorch. For all cases searched learning rates within (0.005, 0.01, 0.05, 0.1, 0.2, 0.5). We also tried linear warm-up schedule for first 3000 iterations, but we did not see any improvement in performances. Generating the data for the figure took approximately days on Google Colab Pro (single Tesla P100 GPU).
B TECHNICAL DETAILS FOR JACOBIANS AND LAYERNORM
We will drop the dependence (z( on x throughout the Appendices. should not cause any confusion since are always considering a single input.
B. NNGP KERNEL
First, we derive the recurrence relation for the NNGP kernel E0. As mentioned in main text, weights and biases are initialized (independently) from standard normal distribution N0, og. (.),nn)). We then have
o dmm 8in and Ep [66666 = of
Eo[w,, [uo, wmnll
(24)
N1
by definition.
We would like to prove Theorem 2.4 as a consequence of Lemma 2.2 The proof Lemma 2.2 can
be found in Roberts et al. (2022).
13

---

Under review as a conference paper at ICLR 2023
Proof of Theorem
2.4
One can prove this by definition with Lemma 2.2
11
K++1
1
+ 1
Ni+1
1
N111
N2
Nr
(at2o(())) + 0,422 v
Ey
Σ
w # (b() + 0,43
Σ
1+3
N141
N141
N2.
1
Σ
E
wkj 1+1 1+1 ((h))d(h)) + b.t1 w
N1+1
i=1
No+1
1
Es
(h1))((1)) + tob
N1+1
N
1=1
Eg [+(H))o(tb)] o
(25)
B.2 JACOBIANS
Next, we prove Theorem 2.5
Proof of Theorem
2.5
We start from the definition of the partial Jacobian (1 > ())
No+1 Nig
ah!t1
oh!
Jlo.411
Es
N+11
ah
ah"
i=
j=1
No+1 No
No
ah onl Thl ah
Ni a ah m a ohly
Eg
N1+1
Ni++ Nig
N2
Ohl on ohly on
Eg
obbb)) (u 0' (((mm))
N1+1
j=1 k,m=1
NN Nip
N
ah ah! ah an
Ea
Wik wim of ((m))
N1+1
j=1 k,m=1
Noi
Nip
N:
di Eo Ni
ahk ank on an
N41
Ni Eg
ahl anl Nia
o N
o' ((H)))o(H))
(26)
ohly th
In the infinite width limit, the sum neurons in layer self-averages due to the law of large numbers. Also the distribution of h1 is independent of h/. This allows us to represent the expec- tation value of a product as product of expectation values. (This holds for to o We will show momentarily that the lo case acquires corrections due to finite input width No). Thus we have
Ni Nip Shk anl
flo.+11 oz 1o' (1 (H)))
N k=1j=1
ahly oh?y
= o (o' (1h))'' (htb Jlo.
glol+1 = =
(27)
14

---

Under review as a conference paper at ICLR 2023
where in the first line used Lemma 2.2 The first integral is taken over the distribution of h4g and the second integral is taken over the distribution of h4- Xy is defined by Eq. 12
L
The critical line is defined by requiring X = where critical points are reached by further requir- ing X²e 1.
As we mentioned in main text, lo = is subtle since the input dimension fixed No. which can not be assumed to be infinity. Even though for dataset like MNIST, usually No is not significantly smaller than width Nj. We show how to take finite O correction into account by using one example.
Lemma B.1. Consider one hidden layer network with finite imput dimension No- In the infinite width limit, the Jacobian is still deterministic and the first step of the recurrence relation modified to:
2
No
Jo.2 = X3
+
No
A
2 k
No
Jo.1.
(28)
where Jo,1 o
Proof.
N₂ No
ah? ah?
Jo.2
Eg
N2
Th oh!
j=1
1 Eg
N₂ No
th Sh The on
W w im
o'
2₂
k,m=1
N2
No
N1
Σ
Eso [w wim why wn o (z
N₂
i=1
j=
k,m=1
No
Ni
2
Eg
E wky of (1))
N
20
No
1
XJ
+
XA
Σ
No
No
20
No
XJ
No
No
Ju.z
(29)
k
where get the result we used integrate by parts, then explicitly integrated over why.
We defined new quantity X'a
Definition B.2 (Coefficient of Finite Width Corrections).
X
N Σ (M))""") + 6" (K))
(30)
Remark B.3. Notice that the correction to Jo. is order O 1). If one calculate the recurrence lation for deeper layers, the correction to Jo. will O N7 1, which means the contribution from hidden layers can be ignored in infinite width limit.
The Jo.2 example justified factorization of the integral when we go from the last line of Eq 26 to Eq.(27]. Finally, the full Jacobian in infinite width limit can be written
15

---

Under review as a conference paper at ICLR 2023
Theorem B.4 (Partial Jacobian), The partial Jacobian of a given network can be written as
2
No
Jod = o
X5
+
I
(31)
No
No
where any partial Jacobian with to 0 does not receive an O(No correction.
B.3 LAYERNORM ON PRE-ACTIVATIONS
Definition B.5 (Layer Normalization).
11 E[4"] R1 B[ E
7/ +
(32)
where of and 3 are learnable parameters.
Remark B.6. With only LayerNorm, the is simplified to
Σ m (())) + D
(33)
Remark B.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 ggrameter-initiaizzttoo Eg F..... Additionally, in this limit, the preactivations are i.i.d. Gaussian distributed h² (oO,
E
[h'] = Eg [h'] = =0,
(34)
E
(h')
27 = E. (h') ( KC'
(35)
The normalized preactivation then simplifies to the form of Eq. 20
Remark B.8. At initialization, the parameters 7, and 8 take the values and 0, respectively. This leads to the form in equation (20). infinite width limit has following form
h Eg[h
(36)
[()) Bolly
Lemma B.9, With LayerNorm on preactivations, the gaussian average is modified to
- f dhd O(h)) (2g)
Es [o(kt)]
(37)
Proof. By definition Tr is sampled from a standard normal distribution // (0, 1), then use Lemma 2.2 get the final form.
Theorem B.10. In the infinite width limit the recurrence relation for the NNGP kernel with Layer- Norm on preactivations is
Ni
KP+1
0
Ea
[0()))(())
+
(38)
N
Proof.
N111
1
k
Ey
Ih 111+1 / 41
N1+1
i=1
N111
Ni
Nr
1
Ey
Σ
w ++ ((h)) + 1/21
Σ
w (()))))))))))))))())) +(41)
N1+1
N
a
Eg
[ + +o
(39)
N
16

---

Under review as a conference paper at ICLR 2023
Theorem B.11. In the infinite width limit the recurvence relation for partial Jacobian with Layer- Norm on preactivations is
glalt1 Xy
(40)
where X,
a² NiK
Em Es (0(h1)2)
Proof.
Ni+1 Nop
oh ah
Jlo.t11
Ea
Ni+1
oh an
i=1
N Nip
No
ah dît oh
a oi an N
Eg
Ni+1
oik anl an
aiin oh only
i=1 j=1
1
Eo
No+1
Nig
Σ
N
1+1 Witk
0' ((()))
VKc )
1
(wite &' (((hm))
VKc
@hle ohnn anly any
N1+1
j=1 k,m=1
oe N1A
Ni
Nio
oh on Th Tnly
Eg
o'(H))o'(h))
N₁
d
Σ
Es
(h(h)))())))
jood
NiC
k=1
(41)
B.4 LAYERNORM ON ACTIVATIONS
The general definition of LayerNorm on activations is given as follows. Definition B.12 (LayerNorm Activations).
((h)) - E(o('))
(h!)) - [1)) 1i + B
(42)
Remark B.13. The recurrence relation for preactivations (Eq. Th) gets modified to
Σ w o(h)) + +
(43)
Remark B.14. At initialization, the parameters of and 8 take the values and respectively. This leads to the form
((h)) E(((h))
((h))
B E(O(N']
(hh)) E. [d(h))]
(44)
LS (6(h))a] - Ee [o(h'))
where the first line follows from the fact that at initialization, the parameters of and 8 take the values 1 and 0 respectively. In the second line, we have invoked the infinite width limit.
Remark 1.. Evaluating Gaussian average in this case is similar to cases in previous section. The only difference being that the averages taking over the distribution h !-1 (' = 0 + af). Again this can be summarized as
(nty
(o + tog) / dht O(h)) 3eक
Ep [O(h)]]
(45)
17

---

Under review as a conference paper at ICLR 2023
Next, we calculate the modifications to the recurrence relations for the NNGP kernel and Jacobians. Theorem B.16. In the infinite width limit the recurrence relation for the NNGP kernel with Layer- Norm activations is
k
oe + +og
(46)
Proof.
kC+1
Eg [ h
N1+1
i=1
N1-1
N₂
Nr
wh ((hh)) + o5 >
Eg
Σ
((1)) + +b,t1
Σ
N1+1
1=1
1=1
Eg
(((1))
06
N
27
N2
o(h)) Eo [o(h']
Es
+ oo
N₂
VED Eo [6((1))] Eg [6(N))
j=l
N₂
Es
(((h)) Eo [6(N))])²]
0
Es (o((h))] - Eo [0()))] + of
N
j=1
22 + of.
(47)
Theorem B.17. In the infinite width limit the recurrence relation for partial Jacobian with Layer- Norm on activations is
glo.t1 = xlglod
(48)
Es {
where X3 = 0
Proof.
1
Ni41 Nio
ahlt1 ahl
Jlod41
Eg
Ni+1
ahl
ahly
1
N141 Nio
ah!+1 onl ahj on )
Σ
anet1 ow
Eg
NN+1
ahim only
i=1 j=l
mm=1
1
Eg
Ni+i Noo
Nr
with (((h))
(with o' ('(hm))
Ohl ahin on anly
}
NN+1
i=1 J=1
k.m=1
Ni
Mia
Nip Thk ohl
Eg
d'(h)))'()))
N
oh an
k=
N
2
Ea
(((h))
Flo.
N
Es
[0'(
0²
Eo [0(N))2] E, 2 Tloi
(49)
n
18

---

Under review as a conference paper at ICLR 2023
C CRITICAL EXPONENTS
To prove Theorem 2.7 we first need to find the critical exponent of the NNGP kernel Roberts et al. (2022).
Lemma C.1. In the infinite width limit, consider critically initialized network with activation function Φ The scaling behavior of the fluctuation 8K' K'' KC* in non-exponential. If the recurrence relation can be expand leading order oK as 8KC+ 5K - c. (8 for n. 2. The solution of dIC
1 1-5
8K'
(50)
On (n
where ÇK n-I
Remark C.2. The constant On and the order of first non-zero term n determined by the choice of activation function.
Proof. We can expand the recurrence relation for the NNGP kernel 10 to second order of ōK =
K KC on both side.
SK se ok (
(51)
Use power law ansatz 8K' A1 then
(1 11+ =1 -- An -nkK
(52)
Multiply /&x on both side then use Taylor expansion
( Sc ss S
SK
Cn Al (n-
(53)
1
For arbitrary 1, the only non-trivial solution of the equation above is
1 n
(n - and (K=
(54)
Proof of Theorem 2.7 We will assume 02 * O. Then use Lemma C.I of 8K! To leading order 1-1
we can expand x²y in terms
Xy m1 d1 OK
=1 - d1 1-|
(55)
C2
Consider a sufficiently large In this case 0(1-1) approximation valid. We write recurrence relations of Jacobians
1-1
d1 x-1
flo.
II
1
Flo.d
P'=lo
23 Cig
[ C
(56)
When c, = 0 for all n 2 2, from Lemma C.1 we have SIC 0. Thus the Jacobian saturates to some constant.
We checked the scaling empirically by plotting Jod vs. in log-log plot and fitting the slope. These results are presented Fig The agreement with infinite width calculation (following sec- tions) is excellent.
We note that for this particular experiment, used NTK parameterization for MLP. However, empha- size that this does affect the results.
19

---

Under review as a conference paper at ICLR 2023
ReLJ
UN UN u=10
Copetmert fitrrg [-1360
bostrent ralleg 0004
Gau Ciparment:
1x
GOL Fiitng -2.080
experieet Treery -1.225
2x307
31-
31
10
10
10 10
10
Figure 4: log-log plot of the partial Jacobian Jo. vs. / for erf. ReLU and erf together with GELU (with LayerNorm applied to preactivations and residual connections of strength 1) activation func- tions. The critical exponents predicted from the infinite width analysis are in agreement with the data. The fluctuations get larger towards the output because the aspect ratio (i.e. L/N) approaches 1/4.
D RESIDUAL CONNECTIONS
Definition D.I. We define residual connections by the modified the recurrence relation for preacti- vations (Eq. (T)
R 1+1
Σ + o(h)) b,+1 + ph!,
(57)
where the parameter controls the strength of the residual connection.
Remark D.2. Note that this definition requires No+1 = N₁. We ensure this by only adding residual connections to the hidden layers, which are of the same width. More generally, can introduce tensor parameter Pij.
Remark D.3. In general, the parameter μ could be layer-dependent (4')). But we suppress this dependence here since we are discussing self-similar networks.
Theorem D.4. In the infinite width limit, the recurrence relation for the NNGP kernel with residual connections is changed by an additional term controlled by
N. Es [o(h))((h))] + m2k
k++1
(58)
N
Proof.
No Σ
kt+1
Egg
[1 [11
Ni+1
No
Ni
Nr
Σ Egg
Σ
1+1 o(h)) + + + wn! 1
Σ
wit ((1h)) + + + unn!
N1+1
1=1
Nes1
N₂
Eg
w ((h))((11))) + + b/+16/+1 + mn
N111
j=1
Nea1
-² Ni
1 Ni+1
E.
(hh) (1))) + of
112
2 Ea
N141
N
Not1
j=1
Ni
Eg [0(1))
+ o + 121
(59)
N
20

---

Under review as a conference paper at ICLR 2023
where we used the fact N4+1 Ni get the last line.
Theorem D.S. In the infinite width limit, the recurrence relation for partial Jacobians with residual connections has simple multiplicative form
glo.t1 X Jloit
(60)
where the recurrence coefficient shifted to Xyy Eo [d + 122.
Proof.
No+1 Nio
ah ah!
Jlo.t11
Eg
Ni+1
T a
Ni41 Nig
Ni ahl+1 an m Thim Thly wwit 2l (1() + prox) (witte of (lim) + pdon)
11 onl Th on
Eg
N1+1
j=1
No+1 Nio
N
Only anly ]
Ohit oh
Eo
N1+1
j=1 k,m=1
Nii1 Nog Σ
a oh on ah
Eg
(wit w o' (l p² dmm)))
N1+1
Nig Tnl ah
]
Nia
}
1
Th ahi Dh anly
Eg
('(h))o'hh))
Eg
4²
N
oh o
N
k=1
Ni.
ahk ank
Eo (d((h))) o' (1k)] + 2² Es
Ni
2 a
Eo (1/))')) (1+) n²) glo.
glolt1 = fio
(61)
L
E RESIDUAL CONNECTIONS WITH LAYERNORM ON PREACTIVATIONS (PRE-LN)
We recall the recurrence relation
N: d(h)) + b//+1 + phÇ
h 1+1 = Σ
14 wet
(62)
Theorem E.1. In the infinite width limit, the recurrence relation for the NNGP kernel is then modi- fied to
kt+1
Σ Eg (0(5))(()) + t p2K!
(63)
N
21

---

Under review as a conference paper at ICLR 2023
Proof.
Ni+1
k+t1
1
Eg [h
Ni+1
Ni+1
N₂
1
Σ Eg
moy (({)) 6,th + and
[
with
o(h))
o,tt + run))
N1+1
j=1
a²
N;
En
(
+ o + n²k!
(64)
N
j=1
Remark E.2. For μ 1, the recursion relation has fixed point
N E.
of n
KC
(o(H)) )o(()))
(65)
N₂-
where the average here is exactly the same as cases for LayerNorm applied to preactivations without residue connections. 1* labels some very large depth
Remark E.3. For μ case, the solution of 63 is
N
KC' = KC +
Σ
Ea
+ of
(66)
N
P=11
j=1
which is linearly growing since the expectation does not depend on depth. KC is the NNGP kernel after the input layer.
Theorem E.4. In the infinite width limit, the recurrence relation Jacobians changes by a constant shift in the recursion coefficient.
gla.t1 x'gl
(67)
where for this case
N Eo
(((h)))0(()) +12
(68)
NiK
k=1
Proof.
No41 Noo
thl+1 ant41
glalt1=
Es
N1+1
ah
ahly
1=1
1
N4t Nig
N: an!+1 ah onk
on! añi, ahin
Eg
NN41
ahk.
Thl on
ahin Thi only
j=1
Nit1 No
1+1 oi (h1)) Wik
+ pdar) )
(with of ('(hm)) VK!
+ mddr) )
anh on!
]
Eg
N111
VK
only only
k-m=1
02 NIK
N
(T (Th.) A²2)
Nig ahk ahi
}
Es
o'
T an
o2 N
Nr
Ea
(0'(6)) n²²
n²)
glod
(69)
n
22

---

Under review as a conference paper at ICLR 2023
Remark E.5. One directly results from cases without residue connections. We will momen- tarily see that the phase boundary does not change with residual connections when μt < 1. However, the correlation length decays way slower when the network is initialized far from criticality. Remark E.6. As we mentioned above μ 1 needs extra care. Plug in the result (66 and p 1 find out that
2 Ek= Bg
o' (1 )o' (71))
X3 1==1
+1
NIKO
+
d2 Σラ= Eg
+ +Nioff ((h)))((1))
11)
~100
(70)
which leads to power law behaved Jacobians at large depth. Where the exponent is not universal.
Recall that { = log x-1"", then Theorem 1.3 is summary of 68 and 77 in I + oo limit.
F MLP-MIXER
In this section we would like to analyze an architecture called MLP-Mixer Tolstikhin et al. (2021) which is based multi-layer perceptrons (MIPs). 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, an output affine transformation. We will explain the architecture by showing forward pass equations.
Suppose has a single input with dimension (C.., Hi Wn.). We label Fpis where the Greek letter labels channels and the Latin letter labels flattened pixels.
First of all the (i) realized by a special convolutional layer, where kernel size is equal to the stride s. Then first convolution layer can be written as
f ΣΣ w
by
(71)
where ff is the size of filter and is the stride. In our example f Notice in PyTorch both bias and weights sampled from uniform distribution 21(- Vk), where k 22)
Ee W-O pare w:O po:yl 3C Supgvodes
1
(72)
Eo (b uu bu
3Cin Sno
(73)
Notice that the output of Conv2d: No E RCNN, where C stands for channels and N₂ Hnn Win/f f² stands for patches, both of them will be mixed later by Mixer layers.
Next stack | Mixer Layers. A Mixer Layer contains LayerNorms and two MLPs, where the first one mixed patches i,j (token mixing) with hidden dimension Nim. the second one mixed channels μr, (channel-mixing) with hidden dimension Nem- Notice that for Mixer Layers use the standard parameterization.
First LayerNorm. acts on channels μa.
7 6J Ht
(74)
V Varc (4
h where we defined channel mean EEc [10 = De
and channel variance Varc =
ps
Ec (Ma)) '] (ht))
First MlpBlock. mixes patches i,j. preactivations from different channels share the same weight and bias.
23

---

Under review as a conference paper at ICLR 2023
6l + 1: Linear Affine Layer.
h pp
Wikh jk
61+1
11 + by
(75)
6l + 2: Affine Layer.
Non
h 4-
W
61+2
(1h((tt)))) jg ) + bl/+2
(76)
where Ntm stands for hidden dimension "token mixing". 6/ + 3: Residual Connections.
h GL+3 4ss
1,6122 + jhh 61 ai
(77)
Second LayerNorm. It again acts on channels μ..
h 16.4 4u
Ec [h 61+331 ps
h
6144
(78)
Varc
61 pi
61+31
Second MlpBlock. mixes channels μ, v, preactivations from different patches share the same weight and bias.
6l + 4: Linear Affine Layer.
66+4 h
Σ p=1
W up
+ 1l141
(79)
6l + 5. Affine Layer.
Nem
66+5
w 61+5 yu
(244 + 66l+5 (11
(80)
- 67 + 6. Residual Connections.
h 46 pi
+ jhh
61+3 us
(81)
Suppose the network has L Mixer layers. After those layers the network has pre-head LayerNorm layer, a global average pooling layer and a output layer. The pre-head LayerNorm normalizes over channels J be described as the following
AGL
IC [hat]]
(82)
Varc [1
Global Average Pool over patches i.
N₂ Np
(83)
Output Layer
C
fu
Σ Eyuht + by
(84)
We plotted phase diagram using the following quantity from repeating Mixer Layers:
N,
C N
1hsl pi
ahsl ui 6L--
X3 = lim 1-000
Ea
(85)
N₂C
06--66 pk
==1
24

---

Under review as a conference paper at ICLR 2023
G RESULTS FOR SCALE INVARIANT ACTIVATION FUNCTIONS
Definition G.1 (Scale invariant activation functions).
<(z) = a x (0(x) + a x (-=)),
(86)
where (2)) is the Heaviside step function. ReLU is the special case with at = and a = 0.
G.1 NNGP KERNEL
First evaluate the average using Lemma 2.2
dh (a) + +c=) (h))² e- o.
Eg [d()))(())
k
a a
(87)
2
Thus we obtain the recurrence relation for the NNGP kernel with scale invariant activation function.
(a} + a²
k++1
+
(88)
2
Finite fixed point of the recurrence relation above exists only if
の (m + +
X
1.
(89)
2
As result
2
(90)
a + a
For 2 + case, finite fixed point exists only if of 0. ta²
G.2 JACOBIAN(S)
The calculation is quite straight forward, by definition
X Lo l 1 (H)))
of
/
obb -
V 2 (a², + a²)
(91)
2
where used the property 20(x) for Dirac's delta function to get the first line.
Thus the critical line is defined by
0
a
+ o
(92)
For ReLU with a. = 1 and a_ = O, the network is at critical line when
on /2.
(93)
where the critical point is located at
(ow, σb) (V2.0).
(94)
G.3 CRITICAL EXPONENTS
C.1 Since the recurrence relations for the NNGP kernel and Jacobians are linear. Then from Lemmal
and Theorem 2.7
Gx and 0.
(95)
25

---

Under review as a conference paper at ICLR 2023
G.4 LAYERNORM ON PRE-ACTIVATIONS
Use Lemma B.9 and combine all known results for scale invariant functions
T
Ni
X3
Eg
(o(()))
NIK
k=1
16-1=1
a (a + a==))
(96)
o w (a} + aá_) + 20
For this case,
X3 - 1
(97)
always true. The equality only holds at o, line.
G.5 LAYERNORM ON ACTIVATIONS
First we substitute K3-1 = o of into known results
a +
Eg [k())))
(98)
a3 + a2
Eg [o(h))(hh)]]
(o" + +0))
(99)
2
There is new expectation value we need to show explicitly
Ee [sb(h]]]
d 20/103+00)-
2T + tog)
dhi(a+ - a-)lie
21o t.
22 + +o%)
02 + oh
(a+ -a-)
(100)
2T
Thus
π (a + a²)
X3
(101)
2 + o
(a
x2) (a+ j3
The critical line is defined by X = 1, which can be solved as
(a+ )=)²
σ6 V T(a + + (a+ a )
(102)
For ReLU with a. and a 00
0
00.6830
(103)
G.6 RESIDUAL CONNECTIONS
The recurrence relation for the NNGP kernel can be evaluated to be
2 (a?
a2
+ a²)
x
K +
of xx
(104)
2
The condition for the existence of fixed point
a (a}
+ +
X
+ 2 <
(105)
2
26

---

Under review as a conference paper at ICLR 2023
leads us to
a 2(1 μ²) 4 a²
(106)
For " 201=#)) a +a²
finite fixed point exists only if of 0. (Diverges linearly otherwise)
The recurrence coefficient for Jacobian is evaluated to be
ote (a_ a
(107)
2
The critical line is defined as
2(1 122)
Ow
a1 a
2
(108)
The critical point located at (v 1
at +o= 00.
For ReLU, the critical point is at (V2(1 -2))
G.7 RESIDUAL CONNECTIONS WITH LAYERNORM ON PREACTIVATIONS (PRE-LN)
Again use Lemma
B.9 and combine all known results for scale invariant functions
N₂
lim
w NiK
E0
(TL)) (1A))
oz
+ 12
R1-1=1
2 (a a² (1 112)
+ 11a of w (a + az) + 20
20 112)
(109)
(a + +a=) 201
Similar to the without residue connections
XJ <
(110)
always true. The equality only holds 0b line for p 1.
Notice there is a very special case = 1, where the whole o, ow plane is critical.
H RESULTS FOR ERF ACTIVATION FUNCTION
Definition H.1 (erf activation function).
1" e dt.
x(x) V 3
(111)
H.1 NNGP KERNEL
To evaluate Lemma 2.2 exactly, we introduce two dummy variables A and X Williams (197).
ff adi / adz dindd2
de
Eo [d(Ai
Eo [+(A1)))()2M)))
/ dd1 / diz I an V (nd)² e +
/ ddi /
4K
ddx2 x 7 2X (A2 + x3))
2
2K. X A2
arcsin
(112)
1 1 + X2)
27

---

Under review as a conference paper at ICLR 2023
We use the special case where X1 X2 1.
Thus the recurrence relation for the NNGP kernel with erf activation function is
20 aresin π
2K'
1C
+ +oğ
(113)
1
2K'
As in scale invariant case, finite fixed point only exists when
4 a (1+ 2K- 4K 1. 1.
1
X
(114)
Numerical results show the condition is satisfied everywhere in 0 ow plane, where X² only possible when K = O.
is
H.2 JACOBIANS
Follow the definition
X²y = Eo (((1))0(1)))
4
/ ddi e 221/)
oap> akT
v 3m!
40
(115)
vI 4KC
To find phase boundary X3 1, we need to combine Eq. 113 and Eq. (15) and evaluate them at K
20 t0
2K
K
aresin
+ +o.
(116)
π
1+
2K*
4
1
X
=1
(117)
V1 + 4K
One can solve equations above and find the critical line
1661 4
n²
20
16 π² 16 n + x²
σ
aresin
(118)
4T
π
Critical point is reached by further requiring X2 1. Since Xie S X²y the only possible case is K O, which is located at
(a ob) =
(
π 0
(119)
H.3 CRITICAL EXPONENTS
We show how to extract critical exponents of the NNGP kernel and Jacobians erf activation function.
Critical point for erf is at (ab, (w) (0, V with KC 0. Now suppose is large enough such
that the deviation of K!! from fixed point value XC small. Define SK! = K! K Eq. can be rewritten as
8K -2 arcsin 28K!
1
(120)
rld 2(8xx
From Lemma C
1 and ÇK 1.
(121)
2
28

---

Under review as a conference paper at ICLR 2023
Next analyze critical exponent of Jacobians by expanding (113) around K 0 critical point (0b, Uw) = (0, 7
To leading order 1 we have
X m1 28k
(122)
=1
Thus the recurrence relation for partial Jacobian, at large 1, takes form
glolt1 (1 (1- 1) gho..
(123)
At large
Jlail = co 10.0 1-11
(124)
with non-universal constant Cig. The critical exponent is
<=1 1,
(125)
which is the same Çi--
H.4 LAYERNORM ON PRE- ACTIVAATASS
Use Lemma B.9
we have
Ni
Σ Eg
NYK!
k=1
16-1=1
4o m
(126)
v3
20 aresin (2) + TOス
The critical line is then defined by
2
aresin
V L
3
0.324o
(127)
H.5 LAYERNORM ON ACTIVATIONS
Due to the symmetry of erf activation function LEg [o(h)] results.
0. we only need to modify our known
4
Es (M))
(128)
/I + 4(02 +
2(02 + of)
E. (o(h))(())) aresin
(129)
1 2 + +o))
Thus
20 2
X
(130)
/I + 4(02, + +
+(02+
aresin
1+2
where the phase boundary is defined by the transcendental equation X
29

---

Under review as a conference paper at ICLR 2023
H.6 RESIDUAL CONNECTIONS
The recurrence relation for the NNGP kernel can be evaluated to be
20 2K + + 1²
2K
K
aresin
(131)
n
Finite fixed point only exists when
40
X
+ S 1.
(132)
π
2K*+ + 4K*
Notice that x S X still holds, where the equality holds only when K 0.
The recurrence coefficient for Jacobian evaluated to be
4
+ 4K 42
(133)
The critical line is defined as
160 x2 22)2
20
160 x2 22222 160 w + 2² - 12
aresin
(134)
4m
122)
Critical point is reached by further requiring X = 1. Since Xi: S X3. the only possible case is K 0, which located at
T T - 122)
(aw. Th)
(135)
4
Note that for = 1, one needs to put extra efforts into analyzing the scaling behavior. First we notice that KC' monotonically increases with depth 1 the recurrence relation for the NNGP kernel at large (or large IC') is
k²+1 ~ oz + of
(136)
which regulates the first term in (13
For μ = large depth
40
Xyy
(137)
T
Co + 4 o
Here Co is constant that depends the input.
We can approximate the asymptotic form of log Flod as follows
log Flo. = log
y
I==lo
40
log
1+
m Co
4(0 of
r==o
402 "
dl' log
1 // + 4(0 + o
la
22vl O(log),
(138)
where ē T / +
We conclude that at large depth, the APJN for 1, erf networks can be written as
gloot ~ O (22v(++()0)g))
(139)
This result checks empirically, shown in Figure
We used NTK parameterization for this experiment. However, emphasize that it does not affect the final result.
30

---

Under review as a conference paper at ICLR 2023
Experiment, 5500
Fiting
10
Experiment, of -200
Theery, 2
1074
Fitiigg 6-8851
Experieent oi = 0.79
1011
6-5.664
Fittngg 0.548
10a
10
10a
10-1
2.5
5.0
7.5
10.0
12.5
15.0
17.5
20.0
Figure log(T for = of =0, erf.
H.7 RESIDUAL CONNECTIONS WITH LAVERNORM PREACTIVATIONS (PRE-LN)
Use Lemma
B.9 and results we had without residue connections for erf with LayerNorm on preac-
tivations.
aZ
No
X
lim 1+000
Σ Eg
o (12)) (1k)
2
NiK
R1-1=1
40 (1
122)
+ 42²
(140)
V5 2202 aresin
no²
The critical line then defined by
2
9
V
LV5
aresin
3
T
(141)
2J 0.324o
RESULTS FOR GELU ACTIVATION FUNCTION
Definition I.1 (GELU activation function).
T
((x)
+
erf
2
V2
1
1
2
e -
dl
(142)
2
31

---

Under review as a conference paper at ICLR 2023
I.1 NNGP KERNEL
Use Lemma 2.2 for GELU
(
h
Eg [o(h))(hh)]]
dh
+ erf
212
2
V2)
[ an!
(n1) 4
11
(
eete
+ ert
e
2
122
K
11
(A))
dh: 2
4
V32TKi
V2
K!
K
14
(h()a
dh erf2
2k!
V32TK
V22
(KC))2
11 22
11
h
1
al!
err
erf'
erf
est"
20
V/2mT
V2
K
K!
KY
2K
aresin
(143)
2x
(1
(1 K
2K
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
K'
K
K'
(K))²
k++1
aresin
0 + +o
(144)
2x
1+ K2
# + (1+13 2X
As result
2
K
+ (1 5K
X
+
aresin
(145)
4
2x
1+1(*
(1+
I.2 JACOBIANS
Follow the definition
X33 =o Eo 1 (h())((h))
(N1)2
2
2
did
+ jerf"
1
h
+
e
h!
V k
2
V2
V2T
2
dhp
+ erf
1
hl
2
heef
/2
11
(k/)2
(h/)2
e
2
4
V2
V2T
2T
02
+ 02
Kt
+ +534
4
2T
arcsin
1
(++ 2(C)
(146)
where dropped odd function terms get the third line, and to get the last line we used known result for erf in the second term, integrate by parts in the third term.
Here to get the critical line is harder. One can the recurrence relation for the NNGP kernel at fixed point KC* and X3
2 K
K
22K
K
+
2TT
aresin
I+
=1+K**)
2k
K + of
(147)
0
02
K"*
X
arcsin
+
11
(148)
4
2T
1+
(1
K*) (1++)))
32

---

Under review as a conference paper at ICLR 2023
Cancel the aresin term, ow and dg then can be written function of K**
2K + (3++5%
K
Ow = 1+
+ aresin
(149)
T1 + (1+ 2KC
1 1+
9b V2T 2K*+) 0m
(150)
One can then scan K to draw the critical line.
In order to locate critical point, further require X² 1. To locate the critical point, solve X3 Xx = O instead. We have
- -2CC) 0,
(151)
2n 21 PC+ (1 [1[1+222) 2K--)
which has two non-negative solutions out of three
K and K 3 VI7
(152)
2
One can then solve σb and aw by plugging corresponding K values.
(ow. σb) = (2,0), for K =0,
(153)
(ow, ~ (1.408, 0.416), for K 3 VI7
(154)
1.3 CRITICAL EXPONENTS
GELU behaves in a different way compare to erf. First we discuss the K = 0 critical point, which
located at (o,, an:)) (0, 2). We expand Eq. [14). and keep leading order 8A KC K
8k SK.² 6 (oxc))2
(155)
From Lemma
A and Sx = 1,
(156)
which is not possible since SKC 2 0 for this case. This result means scaling analysis is not working here.
Next, we consider the other fixed point with KC
3+v17 2
at (ab, aw) = (0.416, 1.408). Expand the
NNGP kernel recurrence relation again.
8K'-1 R si
0.00014(6X))).
(157)
Following the same analysis, we find
5k -114297-
(158)
Looks like scaling analysis works for this case, since KC* >> O. The solution shows that the critical point is Rabbesta al. (2022). If K' < KC, the fixed point is repealing, while when AC' > K** the fixed point is attractive. However, the extremely large coefficient in the scaling behavior of 5K! embarrasses the analysis. Since for any network with reasonable depth, the deviation 5/ is not small.
Now expand x'y some large depth, up leading order 1-1
66.668
X ~
(159)
Then
8 ~ ct 1-6
66. 668
(160)
where G₁ is positive non-universal constant.
Critical exponent
< = 66.668
(161)
Which in practice is not traceable.
33

---

Under review as a conference paper at ICLR 2023
1.4 LAYERNORM ON PRE-ACTIVATIONS
Use Lemma B.9 we have
Ni
XJ =
F NiK!
Es
o'(h))o'(h)))
k=1
13 1=1
T (6 + 4v3)
(162)
0 (6 3v33) 18T0
The critical line is then at
σ6 (6v3T)
00.175
(163)
1.5 LAYERNORM ON ACTIVATIONS
First we need to evaluate a new expectation value
dh h1
0d52
Eg [((((]]]
1 +
erf
2
2T ( + of) 2 + of 21 + 02 +
2
V2
(164)
where we used integrate by parts get the result. The other integrals are modified to
22 + of 1 o + f
(o + + df)
Eo [ol (
= 4 + 2n
arcsin
(1+ os oft) 2(0 + of
(165)
of
o of
t + of 1+ - -ob
(0 o
Eg [o(h))o(h)]]
=
+
aresin
+
4
2x
T(1 oZ o
2 + + (166)
One can then combine those results to find X
0 (1+ 0. + oz
π + 2 aresin
2( +03
I+o
XJ T + +1) 03 + o 20 + ozf)² +
2( + of) + of aresin
o2
[1+2(g=+of)
1+0
(167)
The critical line defined by X/y 1, one can numerically solve it by scanning over as and d.r
1.6 RESIDUAL CONNECTIONS
The recurrence relation for the NNGP kernel is
K
K
K!
(xC))2
arcsin
(+Ki)
+
2
ob + +
(168)
2m
T(1+K') 2K
Fixed point exists if
2
K*
4(K*)) + 11 + +5K*
Xx
4
2T
aresin
K**
(1+ (
+ <
(169)
34

---

Under review as a conference paper at ICLR 2023
The recurrence coefficient for Jacobian is
K 1+k
13 +5**
X3
aresin
+
+ +
(170)
4
2x
(1 (1+
Phase boundary is shifted
(3+
K
= 2 1
2
1+
+ aresin
(171)
(1 + (1++))+ (1+ 2K*
1+ KC+
K
(172)
2 + 2
One can again sean over KC draw the critical line.
In order to locate critical point, we further require X² = 1. To locate the critical point, we solve X Xk O instead. We have
2 - - 0,
(173)
2x (1 + (1 + 210-)
which has two non-negative solutions out of three
K and K 3 3 V17
(174)
2
One can then solve ob and aw by plugging corresponding K values.
(aw, σb) = (2V1 -2,0), for K = 0,
(175)
(am: 2 (1.408 1.2.0.466/1- - n²2, for K 3 VI7
(176)
1.7 RESIDUAL CONNECTIONS WITH LAYERNORM ON PREACTIVATIONS (PRE-LN)
Use Lemma B.9 and results had without residue connections for GELU.
N;
lim 1000
Eg
o' ((
+ +u²
NiK
R1-1=1
2 (6T + 4v (1 2
+ 11 2 (6T + 3/3) + 18T0
(V30 18TO (1 112)
(177)
22 (6T + 33 + 18TO
The critical line is then at
σ6 (6v 3r
(178)
just like without residue connections.
J 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 the training and convergence in this case, compared to standard initialization.
In figure 6 compare the performance of deep MLP networks with and without LayerNorm. We note that the case with LayerNorm applied to preactivations continues to train at very large value of o2. In all cases, networks are trained using stochastic gradient descent with MSE. We used the Fashion MNIST MitaseeXiao et al. (2017). All networks had depth L 50 and width N, = 500. The learning rates logarithmically sampled
35

---

Under review as a conference paper at ICLR 2023
within (10-8 1,008,00 10%) for ReLU, (10 10) for LN-ReLU and ReLU-LN;
within (10-9, 1) for erf. LN-erf and erf-LN;
within (10-8, 10) for GELU, (10-3, 10) for LN-GELU and GELU-LN, where Amax the largest eigenvalue of NTK for each Ow.
191
Figure 6: Performance of deep MLP networks at and away from criticality, with and without Layer- Norm. The blue plateau, corresponding LayerNorm applied to preactivations, continues train at very large values of o² without the need to tune the learning rate.
In figure 7 we showed empirically that the critical exponent of partial Jacobians are vanished for erf with LayerNorm.
40
6x10
UNerr experiment
LL experiment
Theary
Theory
Fittin,C C- 0.007
Fitting, (- 660009
4x:10
10
102
163
3x.10 13F
10
10F
Figure 7: log log plot of partial Jacobian Tod vs. I for (A) LN-erf and (B) erf-LN.
In figure we tested 6k samples from CIFAR- 10 datase Krizhevsky et al. (209 with kernel re- gression based on neural tangents library/Novak Novak al.) (2019) Lee et al. (2019 Novak et al. (2020). Test accuracy from kernel regression reflects the trainability (training accuracy) with SGD in or- dered phase. We found that the trainable depth be predicted by the correlation length of with LayerNorm applied to preactivations, where the prefactor 28. The prefactor we had is the same as vanilla cases in Xiao al. (2020). The difference is from the fact that they used log and we used logg.
In figure we explore the broad range in 2 of the performance of MLP network with erf activation function and LayerNorm preativations. The network has depth 50 and width Nj 500; and trained using SGD on Fashion MNIST. The learning rates are chosen based on logarithmic scan with short training time.
36

---

Under review as a conference paper at ICLR 2023
LN
UNeer
IN-GELU
413
07
03-
80-
312
0440
212
42
40
211
130
55
6.0
85
1.0
1.5
2.0
2.5
1.5
2.0
2.5
3.5
90
Figure 8: Test accuracy for LayerNorm applied to preactivations. of = 0.5 for all cases. Correlation lengths calculated using analytical results of X'y.
LN-Erf, of 0.1
Traning Rcuusyy 160
Figure 9: Training performance of MLP networks with erf activation function; and LayerNorm applied to preactivations. It continues to train for several orders of magnitude of o (with learning- rate tuning).
37


Reviews:
[{"review_id": "b8IVJJFqVpp", "reviewer": "Reviewer_sMx1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper investigates criticality conditions in deep networks, using the NNGP (neural network Gaussian process) framework in the infinite width limit, focusing on architectures with LayerNorm and residual connections.\nIt introduces an estimator for trainability or training stability, the Averaged partial Jacobian norm (APJN), related to the scale (number of layers) over which gradients explode or vanish, which can help selecting the scale of parameter initialization.\nIt then shows theoretically and empirically that some architectures are always trainable, as they are critical regardless of the initialization scale, confirming previous experimental findings.", "summary_of_the_review": "The paper provides a theoretical analysis of criticality in different variants of deep neural architectures, proposes a practical surrogate for trainability, and provides experiments supporting its usefulness.\nThese are useful, well-supported results, that confirm and justify some previous findings, the presentation is clear enough to understand and reproduce.\nI think this is an overall good paper, and should be accepted.\n\n**Update after discussions**\n\nAfter live discussions with AC and other reviewers, I've come to add additional weight to the following negative points:\n* clarity still leaves to be desired, the reader has to work too much to extract the main points\n* the insight has limited value given the existing best practices: the proposed indicator is not really tested outside of current usual settings.\n\nAccording to that discussion, I'm updating my rating. I still think the paper should be accepted, and would have rated it a 7 (rather than 6) if the option was available.", "strength_and_weaknesses": "Strengths\n-------------\n1. A formal analysis, building on existing work on criticality, and expanding it to popular architecture choices\n2. Theoretical confirmation of practical observations regarding critical behavior of networks with various non-linearities, skip-connections, and LayerNorm\n3. Practical confirmation on finite-depth and width MLP Mixer networks\n4. The presentation is overall clear and easy to follow, previous work and novel contributions are usually well introduced\n5. Source code is included\n\nWeaknesses\n------------------\n1. Some confusion in notations (see below)\n2. Some clarity concerns: the definition of APJN always refer to an $(l_0, l)$ pair of layers, but it's not clear what the role of $l_0$ is. It looks like a free parameter one could chose, but what are the consequences of picking one? I get that setting $l_0 = 0$ would have challenges as $N_0$ cannot go to infinity, but could one simply set $l_0 = 1$ without loss of generality?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "MFfs97ZN-T", "reviewer": "Reviewer_7FZf", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the weight initialization scales/magnitudes for sufficiently wide and sufficiently deep neural networks. Noticing that an inappropriate initialization magnitude may multiplicatively accumulate via forward and backward propagation through the large number of layers leading to non-trainability of the network due to exponential explosion or vanishing, the authors focus on the critical initialization magnitudes in order to keep the trainability. Particularly, the paper shows that, with Layer Normalization, the criticality of initialization holds for a large range of magnitudes (or even all initialization).\n", "summary_of_the_review": "Initialization scale has long been an essential topic in training deep neural networks. While many aspects/results of the analysis have been known by the machine learning community, I have never seen the conclusion that, with a proper combination of LayerNorm and skip connection, any initialization scale does not result in gradient explosion and vanishing. ", "strength_and_weaknesses": "The main result of the paper is novel, while the other results and discussions are somewhat incremental. \n\n> As the gradient explosion or vanishing had been one of the major issues in training deep neural networks, many researchers have realized the importance of the magnitude of initialization, and proposed methods to solve or mitigate this issue. Many of these attempts were indeed about finding the critical initialization, such as Kaiming initialization, NTK initialization etc.. For example, in the NTK initialization [1], there is a scaling factor $1/\\sqrt{m_l}$ on each layer so that the output of each neuron is at the order of a constant (neither too large nor too small). In addition, in [2], there is an extra factor $c_\\sigma$ to normalize each layer in order to make the initialization critical. Given these prior results, the result and discussion in Section 2 about MLP seem repetitive and not new. \n\n> Theorem 2.7 is more like a definition of $\\zeta$ instead of a theorem. It is not a surprise at all that APJN is exponential to the number of layers, as the network is MLP. \n\n> The major contribution tells an interesting result. Although it has been wide known that both LayerNorm and skip connect mitigate the gradient explosion and vanishing issues, it is still surprising to see that the combination of the two can totally address this issue, in the sense that any initialization does not lead to gradient explosion or vanishing. \n\nOther comments:\n\n> Below Eq.(2), I cannot quite agree that the partial Jacobian is a random matrix. Once the network is trained away from its initialization, I don’t see how the partial Jacobian is still random.\n\n> In Theorem 2.4, Eq. (10), the NNGP kernel is not in a recursive format. \n\nReferences:\n[1] Neural Tangent Kernel: Convergence and Generalization in Neural Networks.\n[2] Gradient Descent Finds Global Minima of Deep Neural Networks\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Ki0QJduZ7U", "reviewer": "Reviewer_6A97", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper discusses a notion for assessing the trainability (meaning, do the gradients avoid vanishing or exploding during training) of deep neural networks of various architectures -- the Frobenius norm, squared, of the averaged partial Jacobian (APJN) between the pre-activations at layers $\\ell_0 \\leq \\ell$. The paper applies the APJN as an empirical diagnostic of the criticality (is the network trainable, in the above sense) and also to demonstrate that for deep MLP networks with residual connections and with Lipschitz activations with LayerNorm applied to pre-activations, the initialization scheme does not impact the criticality of the network. This result formalizes a view on deep network architecture design -- either fix a network for which the criticality (and thus, loosely, trainability) is sensitive to initialization, and find the right initialization, or define a network using methods like LayerNorm and residual connections so that the network is trainable regardless of initialization. Additionally, an upper bound on the rate of divergence of the gradients in non-critical settings depending on residual connection strength $\\mu$ and the activation functions is given. Several experiments are performed to validate the theoretical contribution as well.", "summary_of_the_review": "Overall, the paper introduced a seemingly novel diagnostic for evaluating the criticality/trainability of deep neural networks and provided some thorough characterizations of the trainability in a variety of situations, in the infinite-width limit. The paper backed up its results with experiments, and used the results to highlight and justify the ways methods like LayerNorm and residual connections can be viewed as a method of making the initialization problem easier to solve. The settings studied are limited, and it is not directly clear how these issues (trainability, etc) impact generalization performance, but it seems to be a useful approximation to understand better and to more thoroughly study. The ideas in the paper may also be helpful to further develop both the practice and theory of deep networks.\n\n==== POST-REBUTTAL ====\n\nAfter reading the other reviews and rebuttals, I remain convinced the paper should be accepted.\n\n==== POST-DISCUSSION =====\n\n After the discussion, overall, I think the paper proposes an interesting measure to investigate when deciding how to initialize deep networks, and provides some justification via a combination of existing theory and limited experiments that it should be interesting. My feeling was still that the paper was relatively clear, though perhaps there is some over-claiming in the theorem statements and attribution should be a bit clearer/cleaned up. I think it is reasonable to accept (I would vote something closer to 7 if that were an option), but I do agree that the paper would be significantly improved by a) cleaning up the theorems / equations (probably don't need to write down each case the way they did), b) tweaking the language to avoid some over-claiming, and c) additional experiments validating the measure in practice. However, I view these changes along the axis from being a decent paper to a great paper and still lean towards acceptance.", "strength_and_weaknesses": "* Strengths: The paper provides a nice characterization of their proposed measure, the APJN, and suggests that it is a useful diagnostic to assess the trainability of deep neural networks empirically. The paper also more formally makes some connections present in the existing literature, and presents a nice takeaway that more precisely characterizes the importance of LayerNorm and residual connections from the perspective of initialization. In particular I appreciated the development of the correlation length concept, it was clearly explained and provided nice intuition. Finally, several experiments on smaller datasets are conducted to validate the theory and the gap between the infinite width setting and realistic settings.\n\n* Weaknesses: The primary weaknesses in my view are as follows: \n     * I believe that the notion of trainability/criticality in this paper has not been directly shown to imply generalization in the literature, and it is more of a heuristic notion. However, many papers study this notion and so it is not a big issue. It would however be nice to clarify with more detail the relation between criticality/trainability and actual model performance in a more formal sense (or at least highlight that this is unknown). It would also be useful to provide a more formal definition of criticality from the previous works that studied it, and to make the connection between the criticality conditions that APJN satisfies and the existing notions of criticality clearer. It would be good to clarify the meaning (and justification) of \"Very deep networks will not train unless initialized critically.\" at the end of the first section.\n     * I don't see a justification of Proposition $2.6$ anywhere -- is this just a (reasonable) assumption for the empirical results? It would be nice to highlight / include a pointer to Figures 3 and 4 which essentially are justifying this assumption, if I understood correctly.\n     * At least one other paper (https://arxiv.org/abs/1802.05296) studies the notion of partial Jacobians and provides conditions that depend on the norms of these partial Jacobians (see Definition 5, the interlayer cushion, in the mentioned paper) in the setting of giving generalization bounds. It probably makes sense to mention this work and any others, as while I do not think I have seen the APJN specifically in previous works, there are probably quite a few more than just the three papers that were cited in the related work that study it. This issue is not a big deal as long as the APJN has not been introduced before or analyzed, but it does impact a little bit the claims of introducing this notion in my opinion.\n    * I don't understand why one would ever care about $\\ell_0$ other than $\\ell_0 = 1$ or $\\ell_0 = L - 1$ (the latter being a computationally tractable, and apparently accurate, approximation) -- is there any reason to care about other $\\ell_0$ as a diagnostic? If not, then perhaps this could be clarified when the paper mentions that the growth of the APJN only depends on $\\ell_0$ via some constant.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "5H77bERcdd", "reviewer": "Reviewer_PgVN", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work introduces the concept of Averaged Partial Jacobian Norm (APJN).\nThe APJN can be used to derive criticality measures &mdash;i.e. how trainable a network is&mdash; for arbitrary parts of a neural network.\nFrom the APJN, an empirical estimator can be derived to empirically compute how \"trainable\" a given neural network is.\nFurthermore, the criticality of layer normalisation, skip connections and the MLP-Mixer architecture is analysed.", "summary_of_the_review": "I would argue that the most interesting contribution of this paper would be a practical estimator for criticality.\nHowever, the proposed estimator is not properly tested to be able to assess when/how this could be used.\nMost of the theoretical analyses turn out to be known results from a body of work on mean field theory.\nTherefore, I do not see what added value the APJN has since it basically builds on the same principles.\nFinally, I found the paper very hard to read and I believe that a rewrite would be necessary to really benefit the broader ML community.\n\n---\n\n### update\n\nAfter the discussion with the authors, I see that there were some misconceptions from my side that lead to an overly negative review.\nOn top of the empirical measure for criticality, the effects due to the combination of layernorm and skip connections could be considered a novel (theoretical) contribution.\nHowever, the bias parameter of a linear layer in front of a normalisation layer is typically initialised to all zeros in practice and thus $\\sigma_b^2 = 0$. As a result, everywhere criticality should also be achieved with layernorm only.\nOverall, I would argue that the paper in its current form still does not do a good job of clearly communicating the insights (and their limitations).\nTherefore, I believe that the ideas in this paper would benefit from a more structural rewrite than the rebutttal period would allow.\n\nTherefore, I decided the correctness score from 2 to 3 and my recommendation from 3 to 5.", "strength_and_weaknesses": " - [+] An empirical estimate of criticality should be useful to easily tune certain hyper-parameters of neural networks (e.g. initialisation).\n - [-] The quality of the proposed empirical estimate is not properly tested.\n - [-] There are a lot of forward and backward references, making the paper hard to read.\n - [-] Assuming that the new architecture refers to applying Layer Normalisation before the activation function, there is nothing *new* about it (Ba et al., 2016).\n - [-] Most of the results from the criticality analysis are already known.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Reject