On the impact of activation  and normalization in obtaining  isometric embeddings at initialization
Authors: Amir Joudaki
Hadi Daneshmand
Francis Bach
Venue: NeurIPS 2023 poster (2023)
Keywords: dynamical isometry
Lyapunov analysis
random neural networks

Abstract: In this paper, we explore the structure of the penultimate Gram matrix in deep neural networks, which contains the pairwise inner products of outputs corresponding to a batch of inputs. In several architectures it has been observed that this Gram matrix becomes degenerate with depth at initialization, which dramatically slows training. Normalization layers, such as batch or layer normalization, play a pivotal role in preventing the rank collapse issue. Despite promising advances, the existing theoretical results do not extend to layer normalization, which is widely used in transformers, and can not quantitatively characterize the role of non-linear activations.  To bridge this gap, we prove that layer normalization, in conjunction with activation layers, biases the Gram matrix of a multilayer perceptron towards the identity matrix at an exponential rate with depth at initialization. We quantify this rate using the Hermite expansion of the activation function.

On the impact of activation and normalization in obtaining isometric embeddings at initialization
Amir Joudaki ETH Zurich juir.joudkiieee @inf ethz.
Hadi Daneshmand MIT//LDS--FDSI chadi@mit... edu
Francis Bach INRIA-ENS-PAS EELS-S Paris francis bach@inria.ff
Abstract
In this paper, explore the structure of the penultimate Gram matrix in deep neural networks, which contains the pairwise inner products of outputs corresponding batch of inputs. In several architectures it has been observed that this Gram matrix becomes degenerate with depth at initialization, which dramatically slows training. Normalization layers, such as batch or layer normalization, play pivotal role in preventing the rank collapse issue. Despite promising advances, the existing theoretical results do extend layer normalization, which widely used in transformers, and can not quantitatively characterize the role of non-linear activations. To bridge this gap, we prove that layer normalization, in conjunction with activation layers, biases the Gram matrix of multilayer perceptron towards the identity matrix at an exponential rate with depth at initialization. We quantify this rate using the Hermite expansion of the activation function.
Introduction
Optimization of deep neural networks is a challenging non-convex problem. Various components and optimization techniques have been developed over the last decades to make optimization feasible. Components such as activation functions [Hendrycks and Gimpel, 2016], normalization layers [Ioffe and Szegedy, 2015], and residual connections [He al., 2016] have significantly influenced net- work training and have thus become the building blocks of neural networks. The practical success of these components has inspired extensive theoretical studies the intricate role of weight ini- tialization [Saxe et al., 2013, Daneshmand et al., 2021], normalization [Yang et al., 2019, Kohler al., 2019, Daneshmand al., 2021, 2023, Joudaki et al., 2023] and activation layers [Pennington et al., 2018, Joudaki et al., 2023], on neural network training. For example, the training of large language models hinges on carefully utilizing residual connections, normalization layers, and tailored activations [Vaswani et al., 2017, Radford et al.. 2018]. Noci et al. [2022a] highlight that the absence or improper utilization of these components can substantially slow training.
To delve deeper into the influence of normalization and activation layers on training, one line of research has studied neural networks initialization (Pennington et al., 2018, de G. Matthews et al., 2018, Jacot et al., 2018, Yang et al., 2019, Li et al., 2022]. Several studies have focused on the Gram matrix, which captures the inner products of intermediate representations for a batch of inputs, revealing that Gram matrices become degenerate as the network depth increases [Saxe et al., 2013, Daneshmand al., 2021, Joudaki al., 2023]. These issue of degeneracy rank deficiency has been observed multilayer perceptrons (MLPs) [Saxe et al., 2013, Daneshmand et al., 2020], convolutional networks [Bjorck al., 2018], and transformers [Dong et al., 2021], posing challenges to the training process [Noci et al., 2022b, Pennington et al., 2018, Xiao et al., 2018]. Research
Code available at: http://gttaaa..oooooouuudleeesse
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

indicates that normalization layers can act effectively to circumvent such Gram degeneracy, thereby improving training [Yang et al., 2019, Daneshmand al., 2020, 2021, Bjorck et al., 2018].
Analyses based of neural networks in the mean-field, i.e., the infinite width regime, have revealed profound insights about initialization by characterizing the local solutions to the Gram dynamics Yang et al. [2019], Pennington et al. [201 8]. However, these results do not guarantee global convergence towards the mean-field solutions or depend on technical assumptions that are challenging to verify numerically [Daneshmand et al., 2021, 2020, Joudaki et al., 2023]. Furthermore, all these theories primarily pertain to the network initialization, where parameters are typically random and do not necessarily hold during or after the training. In this work, our objective is to bridge these existing gaps.
Contributions. Building upon existing literature that elucidates the spectral properties of the Gram matrix, we introduce the concept of isometry, which quantifies the similarity of the Gram matrix to the identity. Our initial theoretical finding demonstrates that isometry does not decrease under conditions of (batch and layer) normalization. This finding illuminates the bias of normalization layers towards isometry at various stages, namely initialization, during, and post-training.
We subsequently extend our analysis to explore the impact of non-linear activations on the isometry of intermediate representations in MLPs. Within the mean-field regime, we establish that non-linear activations incline the intermediate representations towards isometry at an exponential rate in depth. Our principal contribution is quantifying this rate by utilizing the Hermit polynomial expansion of activations. Intriguingly, our empirical experiments unveil a correlation between this rate and the convergence of stochastic gradient descent in MLPs equipped with layer normalization and standard activations used in practice.
2 Related works
line of research investigates the interplay between signal propagation of the network and training. The existing literature postulates that in order to ensure fast training [Schoenholz et al., 2017, Poole et al., 2016], the network output must be sensitive to input changes, quantified by the spectrum of input-input Jacobean. This hypothesis is employed by Xiao et al. [2018] to train a 10.000-layer CNN using proper weight initialization without stabilizing components such as skip connection or normalization layers. He et al. [2023] demonstrate the critical role of the Jacobean spectra in large language models. In this paper, we analyze the spectrum of Gram matrices that connect to the spectral properties of input-output Jacobean.
Mean-field theory has been extensively used to characterize Gram matrix dynamics in the limit of infinite width. In this setting, the Gram matrix is fixed point of recurrence equation that depends on the network architecture [Schoenholz et al., 2017, Yang et al., 2019, Pennington et al., 2018]. This fixed-point analysis can provide insights into the structure and spectral properties of Gram matrices in deep neural networks, thereby shedding light on the degeneracy of Gram matrices in networks [Schoenholz et al., 2017, Yang et al., 2019]. However, often fixed-pointe are not unique, and they can be degenerate or non-degenerate Yang et al. [2019]. In this paper, we establish a convergence rate to non-degenerate fixed- fixed-point for family of MLPs.
Batch normalization [Ioffe and Szegedy, 2015] and layer normalization [Ba et al., 2016] layers are widely used in deep neural networks (DNNs) to improve training. Batch normalization ensures that each feature within a layer across a mini-batch has zero mean and unit variance. In contrast, layer normalization centers and divides the output of each layer by its standard deviation. There have been numerous thcoretical studies on the effects of batch normalization due to its popularity Yang et al. [2019]. Daneshmand et al. [2021], Joudaki et al. [2023]. While layer normalization has been the subject of increasing interest due to its application in transformers Xiong et al. [2020], there are relatively fewer studies on its theoretical underpinnings. While we primarily focus on layer normalization, we define and characterize a property that is shared between batch and layer normalization.
broad spectrum of activation functions such as ReLU [Fukushima, 1969], GeLU [Hendrycks and Gimpel, 2016). SeL.U [Klambauer et al., 2017], and Hyperbolic Tangent, and Sigmoid, are used in DNNs. These functions have various computational and statistical consequences in deep learning. Despite this diversity, only the design of SeLU activation is theoretically motivated [Klambauer et al.,

---

2017], while a broader theoretical understanding of activations remains elusive. To address this issue, we develop theoretical framework to characterize the influence of broad range of activations on intermediate representations in DNNs.
3 Preliminaries
Notation. Let (Ir, y) be the inner product of vectors and y, and ||0|||2 Euclidean norm of x. For a matrix X, we write X,. and X. for the i-th row and column of X, respectively. We use W o o2)mxx indicate that W an m X n Gaussian matrix with i.dd. elements from Nu,, )²). We denote by O, the zero vector of size n. Given vector E R" I denotes the arithmetic mean of Ei= T4. Lastly, In the identity matrix of size n.
(I) the squared
Normalization layers. Let LN IR Rd and BN R2 Rdx= denote batch normalization and layer normalization respectively. Table I summarizes the definition of normalization layers. In our notations, we separate centering from normalization in layer (batch) normalization. Similarly, we split batch normalization into centering and normalization steps in our definitions. This notation allows us to decouple the effect of normalization from the centering. However, we will not depart from the standard MLP architectures we include centering in the network architecture defined below.
Width Depth
dEN N LEM N X Eddx a: R R LN(x)=
Batch size
n E N x E Rd W .....
Input
Input batch
Gaussian weights
W N(0, 1)dxd
Activation
Centering
x T
Layer Norm
Batch Norm
BN(X);:
X VE X
V
Table 1: Building blocks consider in this work.
MLP setup. The subject of our analysis is an MLP with constant width d across the layers and layers, which takes input E IRd and maps it output x Rd, with hidden representations
1 = 1 LN(he hh), he = o(W e= 0...... 2 Va LN(
Vi
input.
While the original ordering of layer normalization and activation is different [Ba et al., 2016], Xiong al. [2020] show that the above ordering is more effective for large language models.
Gram matrices and isometry. Given n data points (s:)ign E Rd, the Gram matrix G of the feature vectors ,f, E Rdd at layer & of the network is defined as
Gº = (2 135
f = 0,1...,,LL
We define the notion of isometry to measure how much G is close to a scaling factor of the identity matrix.
Definition 1. Let G be an n x n positive semi-definite matrix. We define the isometry I(G) of G as the ratio of its normalized determinant to its normalized trace:
LG)
I(G) = Lir(G)
(3)
I(G') is a scale-invariant quantity measuring the parallelepiped volume spanned by the feature vectors rf..... ... of For example, consider two points plane I1, R² with lengths = | b [xa and angle 8 = Z(x,, 22). The ratio is given by ab sin(O) (a)/)((aa (a2 + b²). which is maximized when a = and 0 2/2. This relationship between volume and isometry is visually clear 2 and feature vectors in Figure 1.
Remarkably, I(M) has the following properties (see Lemma A.l for formal statements and proofs):
3

---

42
32
0
31
222
22
2
mg sin
(ol (s..) 9/o
det(X x) tr( X)
++
geometric interpretation of isometry: higher volume, the second row, corresponds to higher value for isometry.
Figure 1:
Scaling-invariant: For all constants > 0, we have I(G) = I(GG).
Range: I [0, 1] where the boundaries and are achieved for degenerate and identity
matrices respectively.
We also define the isometry gap negative logarithm of isometry - log [(M). Based on these properties of isometry, isometry gap lies between 0 and co, with 0 and do indicating the perfect isometry (identity matrix) and degenerate matrices respectively. Isometry allows establish the inherent bias of normalization layers in the following section.
4 Isometry bias of normalization
This section devoted to discussing the remarkable property of isometry in the context of normaliza- tion. We present theorem that formalizes this property, followed by its geometric interpretation and implications.
Theorem 1. Given n samples (wilisn Rdd {0a), their projection onto the unit sphere 3; I;/[|z,]l, and their respective Gram matrices and G, the isometry obeys
I(G)
1 E.
(a, a)²
I(G)
1+
ā²
where a; : || and T O₁
Geometric Interpretation. Isometry can be considered measurement of the "volume" of the parallelepiped formed by sample vectors, made scale and pnneension-indppendeee The normalization process effectively equalizes the edge lengths of this parallelepiped, enhancing the overall "volume" or isometry, provided there a variance the sample norms. Thus, projection onto unit sphere - ><;||<>;|| makes the edge lengths of the parallelepiped equal while leaving the angles between its edges intact. From this geometric perspective, Thcorem implies that among parallclepiped with similar angles between their edges and fixed total squared edge lengths, the one with equal edge lengths has the highest volume (and thereby isometry).
The proof of Theorem 1 is intuitive for the special case of vectors forming cube, where max volume realized when all edge lengths are equal. This fact that maximum volume achieved when edge lengths are equal can be deduced from the arithmetic geometric mean inequality. Strikingly, the proof for the general case is nearly as simple as this special case. The high-level intuition behind the proof is that the determinant allows to decouple the role of angles and edge lengths in volume formulation. This fact evident for n in Figure Since normalization does not modify the angles between edges, the remainder of the proof falls back onto the case where edges form cube.

---

Proof of Theorem 1. Define D == diag(e1, ..,) Observe that G DED, implying det(G) =
GC) det(D)?. Because T; s have norm 1, diagonals of Gram after normalization constant Ga = implying +,2 [r(()) = We have
I(G) tr((G) det (G)
(5)
I(G)
Ar)()) det(G)
E. af
det (T)
(6)
det(G))3/"(II a/)
1 Σ, )
1 1 ag
(II: a.) 2//n ( E, a:)²
det(D) II
0
(7)
21 1 Σ a3
0
2
a #
(8)
Σ a;)²
i
E. (a: a)²
ā ag.
(9)
d²
Theorem 1 further shows a subtle property of normalization: as long as there is some variation in the sample norms, i.e. ||z,||'s are not all equal, the post-normalization Gram has strictly higher isometry than the pre-normalization Gram matrix. It further quantifies the improvement in isometry as function of variation of norms. Intuitively, terms and ΣV (a, a)² can be interpreted the average and variance of sample norms d1...., On. Thus, a higher variation in the norms a, 's leads to larger increase isometry after normalization.
4.1 Implications for layer (and batch) normalization
Theorem I reveals insights into the biases introduced by layer and batch normalization in neural networks, particularly highlighting the improvement in isometry not just limited initialization but also persistent through the training process.
Corollary Consider n vectors before and after layer-normalization (xihenn C IR \ {Oa) and (T, Jig: 34 == LN(x,). Define their respective Gram matrices G = K MyMissn and G = [6x,, 3y Misss We have:
1 E. a)2 (1+
1 a
I I(ā) 2 I(G)
where a; := |1
n Σ
a.
a
What makes the above result distinct from related studies [Daneshmand et al., 2021, 2020, Yang al., 2019] is that the increase isometry not limited random initialization. Thus, layer normalization increases the isometry even during and after training. This calls for future research on the role of this inherent bias in enhanced optimization and generalization performance with batch normalization [Ioffe and Szegedy, 2015, Yang et al., 2019, Lyu al., 2022, Kohler et al., 2019].
Despite the seemingly vast differences between layer normalization and batch normalization [Lubana et al., 2021], the following corollary shows link between these two different normalization tech- niques.
Corollary 3. Given n. samples in mini-batch before X E Rdxx, and after normalization X BN(X) and define covariance matrices C == xxx and C == XXT, We have:
I I(O) I(C) (1 (1+ 1. - 2)2
d
where == |
a²
i=1
Gram matrices networks with batch normalization have been the subject of many previous studies at network initialization: it has been postulated that BN prevents rank collapse issue [Daneshmand
5

---

Initi izwtion
0.59
0.25
0.29-
0:35-
nommalzatioo
0.10
0.03
8.00
nin
nom
nom
Figure 2: Validation of Corollary Isometry (y-axis) vs different layer of MLP: Normalization layers (shaded blue) across all layers and configurations maintain or increase isometry both before (left) and after (right) training. validating Corollary 2. Hyper purameters: activation: tauh, depth: 10, width: 1000, batch-size: 512, training SGD on training set of CIFAR10. with Ir 0.01.. Layer names are encoded as type-index, where type can be fe: fully connected, norm: LayerNorm, and act: activation.
al., 2020] and that it orthogonalizes the representations [Daneshmand al., 2021], and that imposes isometry [Yang et al., 2019]. It is straightforward to verify that orthogonal matrices have the maximum isometry. Thus, the increase isometry links to the orthogonalization of hidden representation characterized by Daneshmand et al. [2021]. While all previous results heavily rely on Gaussian random weights to establish this inherent bias, Corollary not limited to random weights.
4.2 Empirical validation of Corollary 2 in an MLP setup
We can validate Corollary 2 by tracking the isometry of various layers of an MLP with layer normalization. Figure 2 shows the isometry of intermediate representations in an MLP with layer normalization and hyperbolic tangent on CIFAR10 dataset. Shades in the figure mark layers illustrate that the isometry of the Gram matrix is non-decreasing after each layer normalization layer. We can see in Figure 2 that both before (left) and after training (right), the normalization layers maintain improve isometry. To highlight the fact that the claims of Corollary 2 holds all times and not only for initialization, Figure 2 tracks isometry of various layers both at initialization (left) and after training (right). This can be verified by the fact that isometry in the normalization layers (shaded blue) is either stable or increased, which validates Corollary
5 Isometry bias of non-linear activation functions
So far, our focus individual normalization layers. In this section, extend analysis to all layers when weights are Gaussian. Inspired by the isometry bias of normalization, we analyze how other components of neural networks influence the isometry of the Gram matrices, denoted by I(G') with a specific focus on non-linear activations.
5.1 Hermite expansion of activation functions
Analyzing Gram matrix dynamics for non-linear activation is challenging since even small modifi- cations in the scale or shape of activations can lead to significant changes in the representations. powerful tool to analyze activations is to express activations in the Hermite polynomial basis. Inspired by previous successful applications of Hermite polynomials neural network analysis [Daniely et al., 2016, Yang, 2019], we explore their impact on the isometry of activation functions.
Definition 2. Hermite polynomial of degree k. denoted by Hex(æ), is defined as
dk
Hex(x) = (
dx -
All square-integrable function a with respect to the Gaussian kernel, which obeys m e==/2 can be expressed linear combination of Hermite polynomials

---

as [x)) Ex (r) with (see section A for more details):
== (z) (x)))
The subsequent section will discuss how to leverage the Hermite expression of the activation to analyze the dynamics of Gram matrix isometry.
5.2 Non-linear activations bias Gram dynamics towards isometry
In this section, we analyze how I(G²) changes with L. We use the mean-field dynamic of Gram matrices subject of previous studies [Yang al., 2019, Schoenholz al., 2017, Poole et al., 2016]. The mean-fiele field dynamics of Gram matrices is given by
G = =EINN(, (((())))
where [o(a)]; == (a: - Ea;)/ / Var a,
(10)
This equation gives the expected Gram matrix for layer l + 1, based on the Gram matrix G from the previous layer, and Φ is mean-field regime counterpart for layer normalization operator (see section A for more details). The sequence G approximates the dynamics of G and this correspondence becomes exact for infinitely wide MLPs. In the rest of this section, we analyze the above dynamical system. Our theory relies the notion of isometry strength of the activation function, defined next.
Definition 3 (Isometry strength). Given activation with Hermite expansion {ck) define its isometry strength 3g as:
Bg :- 2-
(1I)
10
1=1
We can readily check from the definition that isometry strength B, has the following basic properties (i) it ranges between and 2, and (ii) it is if and only if the activation is linear function. Table presents the isometry strength of certain activations in closed form. With this definition, we can finally analyze Gram matrix mean-field dynamics. Interestingly, the negative log of isometry log I G..)) can serve as a Lyapunov function for the above dynamics. The following theorem proves non-linear activations also impose isometry similar normalization layers.
Theorem 4. Let be activation function with Hermite expansion and non-linearity strength Pos (see equation (11)). Given non non-deeeeeeate input Gram matrix Go,, then for sufficiently large layer i 2 (-n log [[G)) + log( (ng(4n)), we have
log I(G'S) S eexp(=elogBBo n log I(G )++ + 1og(4n)).
(12)
Note that the condition on input being non-degenerate is essential to reach isometry through depth. For example, if the input batch contains a duplicated sample, their corresponding representations across all layers will remain duplicated, implying that all G will be degenerate.
Theorem 4 reveals the importance of non-linear Hermite coefficients (Ck, k 2 2) in activation function to ensure 33 > and obtain isometry in depth. This connection between 8rr and isometry is the rationale for referring to &or as the isometry strength. This constant can be computed in closed form for various activations, as shown in Table 2, and for all other activations, it can be computed numerically by sampling.
He (x)
Hez(x)
Sine sin(x)
Exponential exp(x
Step
ReLU
2²
2)
1 0
max (r, 0)
3
1
2
2
2
Table 2: Isometry strength 3o (see Definition 3) for various activation functions.
Figure 3 compares the established bound on the isometry gap with those observed in practice, i.e. G for three activations. We observe 3r predicts the decay rate in isometry of Gram matrices G.. While so far, we have only discussed the direct results of our theory for isometry of Gram matrices, the next section, will discuss other insights from the above analysis.
7

---

Sigmold, 9,-1.012 1.012
Tanh, By
RoLU, As
10-
0.8
0.6
Trace empinical [-log z(Gp)
(P5')
0.0
layer
layer
layer
Figure 3: Validation of Theorem 4. Solid blue traces show the isometry of MLP with n 100, d = 5000, and various activations a. Solid lines show the average of #0 independent runs. The dashed traces are thcoretical upper hounds given in Theorem 4, with constant C
6 Implications of our theory
In this section, we elucidate the implications of our theory, beginning with insights into layer normalization through the Hermit expansion of activation functions, followed by examination of its impact on training.
Layer normalization primarily involves two steps: (i) centering and (ii) normalizing the norms. Through the Hermit expansion of activation functions, we unravel the underlying intricacies of these components and propose alternatives based on insights from the Hermit expansion.
Our experiments utilize MLPs with layer normalization and various activation functions for the task of image classification on the CIFAR 10 dataset [Krizhevsky et al.]. For training. use stochastic gradient descent (SGD) with fixed step size of 0.01. Unless stated otherwise, the MLP has constant width of d 1000 is maintained across hidden layers, and a batch size of n = 10. Throughout this section, at and xf respectively denote the post-activation and nost-normalization vector for layer f.
Experimental Setup.
6.1 Centering and Hermit expansion
One of the key insights of our mean-field theory is that the centering step in layer normalization is crucial obtaining isometry. In the mean-field regime, pre- ae-avivviaioos follow standard Gaussian distributions, and thus the average post-activation af = 1/2 Σ of will converge to their expectation co. This insight suggests an alternative way of obtaining isometry by explicitly removing the offset term from activation, i.e., replacing activation (xx) by (xx) co- Strikingly, our experiments presented in Figure indicate that such replacement can also impose isometry. This result provides novel insights into the role of centering in layer normalization.
at Eo. (X))
tanh
sigmoid
exp
10a
10-1
10
20
30
10
20
30
10
20
30
0
10
20
30
Figure 4: Layer vs. mean-field centering in obtaining isometry. Layer centering (orange): x e+1 = LN(aa' - Mean-field centering (blue): ff+ LN(a" co).
6.2 Normalization and Hermit expansion
Theorem 4 further reveals the importance of normalization of norms in addition to centering to achieve isometry. Figure 5 underlines the importance of the normalization for different activations, where we observe the isometry gap may increase without normalization. Similar to our mean-field analysis of centering, the factor 1 Σd [ a²)² in layer normalization converges to variance
8

---

v(.. = Ex =: 5(1. Thus, the width increases, the layer normalization operator LN - a²) will converge to (af - co)/ (T)). Figure 6 demonstrates that the constant scaling 1/vō(1) achieves comparable isometry to layer normalization for hyperbolic tangent and sigmoid and ReLU, while is not effective for exp function. This observation calls for future research on the link between normalization and activation in deep neural networks.
relu
tanh
sigmoid
exp
10a
logz(G")
100
10
10
20
30
10
20
30
10
20
30
0
10
20
30
1
Figure 5: Comparing no normalization (bluc) =+1 = (a (0) with layer normalization (orange)
++ LN(a* co) in obtaining isometry.
relu
tanh
sigmoid
exp
log (G9
10
10
10-1
0
10
20
30
10
20
30
0
10
20
30
0
10
20
30
Figure 6: Comparing layer normalization (blue) f+1
LN(aL - co) vs. mean-field normalization (orange)
-ro in obtaining isometry.
(1
6.3 Isometry strength correlates with SGD convergence rate in shallow MLPs.
Besides the direct consequences of our theory, we observe striking correlation between the conver- gence of SGD and isometry strength in specific range of neural network hyper-parameters. Figure 7 shows the convergence of SGD is faster for activations with a significantly larger isometry strength 3, (see Definition 3) for shallow MLPs, e.g., with 10 layers or less. We can speculate that this correlation reflects the onput-output sensitivity of the networks with higher non-linearity Surprisingly, this correlation does extend deeper networks. This discrepancy between shallow and deep networks regarding SGD convergence may be due to the issue of gradient explosion studied by Meterez.c et al. [2023]. This finding suggests multiple avenues for future research.
Sigmoid, A, 1.02
Tanh; 1.07
105s
1.5
1.-
e5och
Figure 7: B, correlates with SGD training speed. Training loss (y-axis) vs epoch (x-axis) on a training dataset of CIFAR10. Each curve shows the MLP with corresponding activation, with color codes denoting isometry strength S.. Hyper parameters: depth: 10, width: 1000, batch size: 512, Average of #5 independent runs.

---

7 Discussion
In this study, we explored the influence of layer normalization and nonlinear activation functions on the isometry of MLP representations. Our findings open up several avenues for future research.
Self normalized activations. is worth investigating whether we can impose isometry without layer normalization. Our empirical observations suggest that certain activations, such ReLU, require layer normalization to attain isometry. In contrast, other activations, which can be considered "seff-normalizingg (e.g., SeLU [Klambauer et al., 2017] and hyperbolic tangent), can achieve isometry with only offset and scale adjustments (see Figure 8). We experimentally show how we can replace centering and normalization by leveraging Hermit expansion of activation. Thus, believe Hermit expansion provides a theoretical grounding to analyze the isometry of SeLU.
10a
10at
5
10
15
20
25
30
relu
tanh
selu
x LM(otW5s
x oLLMW')))))
(a)
(b)
Figure 8: (a) Exploring potential of achieving isometry with self-normalized activations. Batch size n = 10. width d = 1000. (b) Impact of the order of activation and normalization layers. Batch size = 10, width
d= = 1000.
Impact of the ordering of normalization and activation layers on isometry. Theorem 4 highlights that the ordering activation and normalization layers has a critical impact on the isometry. Figure demonstrates that different ordering can lead to non-isotropic Gram matrix. Remarkably, the structure analyzed in this paper used transformers [Vaswani et al., 2017].
Normalization's 's role in stabilizing mean-field accuracy through depth. Numerous theoretical studies conjecture that mean-field predictions may not be reliable for considerably deep neural networks [Li et al., 2021, Joudaki et al., 2023]. Mean-field analysis incurs a O(1/v widdhh) error per layer when the network width finite. This error may accumulate with depth, making mean-field predictions increasingly inaccurate with an increase in depth. However, Figure 9 illustrates that layer normalization controls this error accumulation through depth. This might be attributable the isometry bias induced by normalization, as proven in Theorem I. Similarly, batch normalization also prevents error propagation with depth by imposing the same isometry [Joudaki et al., 2023]. This observation calls for future research on the essential role normalization plays in ensuring the accuracy of mean-field predictions.
Acknowledgements
Amir Joudaki funded through Swiss National Science Foundation Project Grant #200550 Andre Kahles, and partially funded by ETH Core funding award to Gunnar Ratsch. Hadi Daneshmand acknowledges support from the NSF TRIPODS program (DMS-2022448).
References
Dan Hendrycks and Kevin Gimpel. arXiv:1606.08415, 2016.
Gaussian error linear units (gelus). arXiv preprint
10

---

d=100
d=500
d=1000
-log -log I(G
10"
10-1
0
10
20
30
0
10
20
30
0
10
20
30
Figure 9: Role normalization stabilizing mean-field accuracy for ReLU, batch size n = 10. Layer normalization (orange): 2f11 = INN c))
Mean-field normalization (blue): x f+1
(1))
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, pages 448 456. pmlr, 2015.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: 2013.
Hadi Daneshmand, Amir Joudaki, and Francis Bach. Batch normalization orthogonalizes repre- sentations in deep random networks. Advances Neural Information Processing Systems, 34: 4896-4906, 2021.
Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S Schoenholz. A mean field theory of batch normalization. arXiv preprint
2019.
Jonas Kohler, Hadi Daneshmand, Aurelien Lucchi, Thomas Hofmann, Ming Zhou, and Klaus Neymeyr. Exponential convergence rates for batch normalization: The power of length-directee decoupling in non-convex optimization. In The 22nd International Conference Artificial Intelligence and Statistics, pages 806-815. PMLR, 2019.
Hadi Daneshmand, Jason Lee, and Chi Jin. Efficient displacement convex optimization with particle gradient descent. International Conference on Machine Learning, 2023.
Amir Joudaki, Hadi Daneshmand, and Francis Bach. bridging the gap between mean field and finite width in deep random neural networks with batch normalization. International Conference on Machine Learning. 2023.
Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. The emergence of spectral universality in deep networks. In International Conference Artificial Intelligence and Statistics, pages 1924 1932, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
Alec Radford, Karthik Narasimhan, Tim Salimans, llya Sutskever, et al. Improving language understanding by generative pre-training. 2018
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 15:27198-27211 2022a.
Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard E. Turner, and Zoubin Ghahramani. Gaussian process behaviour wide deep neural networks, International Conference on Learning Representations, 2018

---

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems, 31, 2018.
Mufan Bill Li, Mihai Nica, and Daniel M. Roy. The neural covariance sde: Shaped infinite depth- and-width networks at initialization. Advances in Neural Information Processing Systems, 2022.
Hadi Daneshmand, Jonas Kohler, Francis Bach, Thomas Hofmann, and Aurelien Lucchi. Batch normalization provably avoids ranks collapse for randomly initialised deep networks. Advances in Neural Information Processing Systems, 33:18387-18398, 2020.
Nils Bjorck, Carla P. Gomes, Bart Selman, and Kilian Q Weinberger. Understanding batch normal- ization. Advances in Neural Information Processing Systems, 31, 2018.
Yihe Dong, Jean- Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pages 2793 2803, 2021.
Lorenzo Noci. Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto. Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. arXiv preprint ar (iXxx:2206003/22 2022b.
Lechao Xiao, Yasaman Bahri, Jascha Sohl- bilk-iicstten, Samuel Schoenholz, and Jeffrey Pennington. Dynamical isometry and a mean field theory of CNNs: How to train 10,000-layer vanilla convo- lutional neural networks. In International Conference Machine Learning, pages 5393-5402, 2018.
Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference Learning Representations, 2017.
Ben Poole. Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity deep neural networks through transient chaos. Advances in Neural Information Processing Systems, 29, 2016.
Bobby He, James Martens, Guodong Zhang, Aleksandar Botev, Andrew Brock, Samuel L. Smith, and Yee Whye Teh. Deep transformers without shorteuts: Modifying self-attention for faithful signal propagation. arXiv preprint arXiv:2302. 10322, 2023.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. arXiv preprint wrXiv:1607.06400 2016.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533, 2020.
Kunihiko Fukushima. Visual feature extraction by a multilayered network of analog threshold elements. IEEE Transactions on Systems Science and Cybernetics, 1969.
Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. Self-normalizing neural networks. Advances in Neural Information Processing Systems, 30, 2017.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normal- ization layers: Sharpness reduction. Advances in Neural Information Processing Systems, 35: 34689-34700 34708, 2022.
Ekdeep S Lubana, Robert Dick, and Hidenori Tanaka. Beyond batchnorm: towards unified understanding of normalization in deep learning. Advances in Neural Information Processing Systems, 34:4778-477 4791, 2021.
Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity. Advances in Neural Information Processing Systems, 29, 2016.
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv: 1902.04760, 2019.
12

---

Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 10 (canadian institute for advanced research). URL http: //wp:.///www../ cs. toronto. edu/ cifar .tml.
Alexandru Meterez, Amir Joudaki, Francesco Orabona, Alexander Immer, Gunnar Rätsch, and Hadi Daneshmand. Towards training without depth limits: Batch normalization without gradient explosion. arXiv preprint arXiv:2310.00012, 2023.
Mufan Li, Mihai Nica, and Dan Roy. The future is log-gaussian: Resnets and their infinit--depth-and width limit at initialization. Advances in Neural Information Processing Systems, 34:7852-7888 7864, 2021.
Semyon Aranovich Gershgorin. Uber die abgrenzung der eigenwerte einer matrix. News of the Russian Academy Sciences. Mathematical series, (6):749-754, 1931.
Gustav Mehler. Ueber die entwicklung einer function von beliebig vielen variablen nach laplaceschen functionen höherer ordnung. Journal für die Reine und Angewandte Mathematik (in German), 1866.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, nigh-performance deep learning library. Advances in neural information processing systems, 32. 2019.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, e'm Louf, Morgan Funtowic7, Jamie Brew, and Guillaume Dulac- Arnol.. Transformers: State-of-the-
natural language processing. Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38-45, October 2020. URL https: //wp::////www.// aclweb. org/anthologgy2220 emnlp-demos. demos 6.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference artificial intelligence and statistics, pages 149-256. JMLR Workshop and Conference Proceedings, 2010.
13

---

Appendix outline
The appendix is partitioned into four main components, each serving its purpose as described:
1. Section A details the all the proofs, with most of it dedicated to proof of Theorem 4 alongside numerical confirmation of significant steps
Section A. elahorates on the hasic properties of isometry.
Section A.2 provides an elaborate review of the mean- -ield Gram dynamics.
Section A.3 presents the Lyapunov function y, and establishes that this function provides both upper and lower bounds for the isometry, thereby implying that geometric contraction of 7G') indicates a geometrie contraction of isometry gap log I(G').
Section A.4 proves that ((')) exhibits an exponential contraction in depth with rate Pr.
2. Section outlines our rebuttal responses to the reviews that we chose to leave out of the main text.
Section B.1 gives additional details concerning the experiments reported in the main text and appendix.
Section B.2 explores the effect of gain on isometry, and links the rate to the associated isometry strength.
Section B.3 explores the effect of varying widths of hidden layers on the isometry.
Section B.4 explores the notion of isometry for representations in language models.
A Proofs
A.1 Basic properties of isometry
Basic properties of isometry It is straightforward to check isometry obeys the following basic isometry- preserving properties:
Lemma A.1. For PSD matrix M, the isometry defined (3) obeys the following properties: 1) iecleeiivvriance I(cM) = I(M), 2) only takes value in the unit range I(M) [0, 1] 3) takes its maximum value if and
only if M is identity I(M) =1 + M = Ia, and 3) takes minimum value if and only if M is degenerate I(M)= 0.
Proof of Lemma A.I. The scale-invariance trivially true as scaling M by any constant will scale det(M)) 1/ and tr (M) by the same amount. The proof of other properties is a straightforward consequence of writing the isometry in terms of the eigenvalues I(M) = QI, A./ /C; E, Ar), where A,'s are eigenvalues M. By arithmetic vs geometric mean inequality over the eigenvalues we have (II, A,) < 1 Σ, Ar), which proves that Z(M) E j0, 1]. Furthermore, the inequality is tight iff the values are all equal X = .... An which holds only for identity M = In. Finally, isometry is zero iff at least one eigenvalue is zero, which is the case for degenerate matrix M.
A.2 Mean-field Gram Dynamics
Recall the mean-field Gram dynamics stated in equation (10):
Ge++ (o, ((a(()))(((()))
where [d(a)), (a. Ea,)/ Sa..)/Varoo o,,
(13)
Assuming that inputs are encoded columns of X, we can restate the MLP dynamics follows
x 1 LN(X X),
inputs
(14)
Vd
Hº = W'x' x w vr'x"w" ~ NOO, 1)add
preactivation
(15)
= O(TT),
activations
(16)
x+ LN(A* - -A)
centering & normalization,
(17)
Vi
where centering and layer normalization are applied column-wise, as defined in the main text. Observe that Gram matrix of representations can be written as
Gf+1 x++1 e+1
(18)
A
M1
A
M4
92
1
(A1.). ==
1
Σ (AC, mr)2, (19)
d
==
d k=1
ā
k=
14

---

where @ denotes Hadamard product, and subscript emphasises the dependence of Gram width d. Note that conditioned on the previous layer. of H' and A are ii.d. because of independence of of W Thus, by law of large numbers, the infinitely wide network regime, f4, and 8; will converge to expected mean and / (Va((Ai..)) for all L...... By construction
variance respectively lima- 14i EAfi, and limg- rewrite Gram dynamics lim Co à E ((AA) (TA)(.A.))).
of Φ, in the infinitely wide regime,
We can invoke the fact that rows of A are i..dd to conclude that Ga is the sample Gram matrix that converges to its expectation
lim Go Eo(Af))o(A))
(l =: G
(20)
where, denotes the mean-field regime d This concludes the connection between the mean-field Gram dynamics and infinitely wide Gram dynamics.
A.3
Introducing potential
Here we will introduce Lyapunov function that enables us to precisely quantify the isometry of activations in deep networks:
Definition 4. Civen positive semidefinite matrix G _=x" define 7 R R220
|G:yl
(G)
:y G.=Gu/VGGGA
(21)
G
Remarkably, 7 obey exhibits an geometric contraction under one MLP layer update, which is stated in the following theorem:
Theorem A.2. Let G be PSD matrix with unit diagonals G.. 1, L...., holds:
(En..No..)) ((((())))(((()))) & rGG)\/Bo
where [o(a]] = (a: Eai)/ //v Var ag.
(22)
Thus, we may apply Theorem A.2 iteratively prove that the mean- field, the Lyapunov function y (G') decays at an exponential rate 3g. A straightforward induction over layers leads to a decay rate in 7, which is presented in the next corollary.
Corollary A.3. If activation has Hermite expansion with non-linearity strength So, the mean-field Gram matrices G., obey Lyapunov these Gram matrices decays at exponential rate S
<(.)) 7C.))8
(23)
where G. denotes the imput Gram matrix.
In Figure A.I, we experimentally validated the above equation for MI.Ps with finite width and activations (tanh relu, sigmoid} where we observe that the mean-field analysis well predicates the decay in T.
relu
tanh
sigmold
10a
v(G
theory
10
empirical
0
10
20
30
0
10
20
30
0
10
20
30
1
Figure A.1: (G vs depth e, for MLP with 10, d 10000, various activations a. Solid lines shows average of 10 independent runs. The dashed traces the theoretical upper bounds given in Corollary A.3.
Interestingly, we can connect the Lyapunov y to the isometry, by proving an upper and lower based on the determinant G based on 7(G), when G is PSD and has unit diagonals:
Lemma A.4. For PSD matrix G with unit diagonals holds:
1 1- (m 1) max [Gts]] S (G) S max G
(24)
(8)
where the lower bound holds if (n L)]Gul | 51.
15

---

Now, we are ready to prove the main theorem.
Theorem A.S (Restated Theorem 4). Let be activation function with a Hermite expansion and a non- linearity strength Bu. (see equation (11)). Given non-degenerate input Gram matrix Go,, then for sufficiently large layer 2 30 (-nlog (G.)) + log (4n)). we have
- llog I I exp(-f log 3. - log.) + log(4n)).
(25)
Proof of Theorem 4. First, note that we have the following transformation
7()
max [Gt
(26)
x
(+)
we have
det(G 2 (1 (n - 1) max 1)
Lemma A.4
(27)
det(G()) = (1 (n - 10--))((()))))) /(1 ((1++)
Invoke (26)
(28)
2 1 (1 (1-(n11))G5))
(29)
log GG. < log(1 (n ()(()))
(30)
S (
for for) <000)>1/// <1/(2n 2)
(31)
S 22ny(Go)8,*
(32)
By upper bound of Lemma A.4 we have maxişj |Gay S v1 - det((G) S det(G)/2, implying (Go) -
2 deetGoa) which allows us to
log I(G)) S 4n del 30²
(33)
I eeppl_llos Bo n log I(Go) + log(4n))
(34)
which concludes the proof.
Proof of Lemma A.4. Lower bound. The lower bound is result of Gershgorin circle thearem [Gershgorin, 1931], which implies that every eigenvalue must be within the disc [1- (n- 1)) maxe7) [G:yl, (+(n=1)) max./; [Gay]] |]- Thus, the determinant
dower-bounded by (1 (n 1) max.fs [Gty )"
Upper bound. Since G is PSD, we can write = x X, which implies that columns X are unit norm ||=;]| = 1, and G encodes the angles hetween them cos 2(x;, det(G) vol(X)? where the volume refers to the parallelepiped spanned by the columns of X. With this formulation, write volume recursively volume spanned by columns up L, times the projection distance of s., from their span
(z,, = Gūj- Furthermore, we have
vol .,) = dist( 2.)))))).....)))llllll....
1n-1),
where span refers to the space of all linear combinations of these vectors. Note that projection distance of æ, onto I, be written sin xx,,,,,) VI 2,)² / G Since the projection distance onto the linear span cannot be greater than projection distance onto single vector, which is bounded by 11-mmxx:s Ciyy This concludes the upper bound that det(G) < 1 max:zj Ghj-
A.4 Dual activation and proof of Thm. A.2
According Daniely al. [2016], we leverage the notion dual activation associated with the activation function a, which defined in following.
Definition S. Given activation a : R - R that is square integrable with respect to the Gaussian kernel, define its dual activation ē R- -> RR and its mean reduced dual ā R - IR as:
o (p) == (x))
(p) = E [[o (z)) - (y)))
(35)
The following lemma connects the dual activation and its mean mean-reduced version with the Hermite expansion of 0:
Lemma A.6. Given two standard Gaussian variables X,Y ~ N(0, 1) with covariance EXY = P, and activation o with normalized Hermite coefficients Ckk, we have
80
6 ((p)=
Σ E
x(p)
(36)
k=
1=1
16

---

Finally, we have the tools to prove the first main theorem.
Proof of Theorem A.2. Let = (hh) for A N (0, G). Note that by definition of the dual-activation have Ena" = [G Since we assumed G has unit diagonals, we have h, ~ [ 1), which implies that (G = oG.,). Furthermore, the variance can be driven as Var(a.) Eaq (Ba,)* = -k o Ex= c (1) for all = I...., n. In the matrix form we have
Ea.- -E,, (h,) cu, implying that E(a, - Ea,)aa Ea,)"
Thus, we have (o(a)))(a,) (G,,)
Ein [(lo(t)))(((())]]] - (IOGG
(37)
The remainder proof relies the following contractive property of Gram matrix potential:
Lemma A.7. Consider activation a; with normalized Hermite coefficients {cs (cx) For all p (0. 1), the mean-reduced dual activation ō obeys
[o(p))/o(1) (0)1/a(1)) S 3
1p
(38)
1- l
which the right hand-side is strictly larger if some nonlinear coefficient is nonzero # for some k 2 2.
Thus we can apply Lemma A.7 on each element A to conclude that
|G(G)))//(()) (1)
|G.yl
- I Gag
B7
Lemma A.7 for all }j
(39)
|G:yl
5 5 max *fj 1 Gtyl
(40)
-/- rG),
(41)
(42)
since the incquality holds for any value of i + J. we can take the maximum over i + jj to write:
5 (Ga))/ S B
(G))/((L) = max-
(43)
( 1 - (G,))\/)(()
which concludes the proof.
L
Proof of Lemma A.7. Note the ratio is invariant scaling of Σ of. Hence, we assume Eee
km 4
without loss of generality. With this simplification, we have o (1) 1.
For the positive range p E 0, 1]
we have
P (() ) 11= --1-= ( 113
0
ctpp
(44)
p) Σm
cp(1 - ph)
k= 2hp-11
(45)
000
1-pt
1
T-p
(46)
1=
21
chpp-1
(47)
1 Cept
+ > x
2 ck
Dk-1 2 1-1
max #c
(48)
cptt
+ > ch
===1 c
(49)
1= 2
Ets c2 cf
(50)
B,'
Thus for p E [0, 1] we have
(p)
< 1- p
P 3o
(51)
1- (p)
By Jensen inequality for convex function
+- |1|| we have
[5()) = Σ xp *1
Σ
Aeel = =((p))
(52)
k=1
17

---

Because >> 2/(1 - x) is monotonically increasing for [0, 1], we have
[x(p)]
2 ( S |p lp g0
o((p))
(53)
[p(pp
Where invoked the inequality that proven for p [0, 1]. because E [0.
The following lemma, which a consequence of Mchler's formula, the hart of proof of Lemma A.6: Lemma A.8 (Consequence of Mehler's kerne)., X, ~ (0, 1) with covariance EXY = we have Ex,y He, (X) (Y) =p" p"8nk
where 8 is the Diruc delta.
Prouf of Lemma A.6. Let Σx=o 4 (x), denote Hermite expansion of Thus, we have
(p) = Exxro(X)o(M)
(54)
G.arBx,y Hen (X) Hex
(55)
Σ AP
(56)
where in the last line we applicd result of Lemma A.8. For the mean-reduced dual kernel ō, observe that Ex (x) cp. Thus, the reduction of mean will cancel the & = 0 term, which concludes the proof that {pp) = Ei- Cppp.
Proof of Lemma A.8. The property be deduced from Mchler's formula Mchler [1866). The formula states that
p² + 1²) 2ryp
p²
exp
2(1 - p²)
(57)
Σ Hem(a) Hem((y,
(58)
where the mu! factor difference is due to the definition of Hermite polynomials with an additional 1/Vm! compared the used in Mehler's kernel. Observe that the left hand side is equal to p(z, w)/p(=)p(s)) where p(a, y) the joint PDF (X, ), and p(x). are PDF and respectively. Therefore, we can take expectation using the expansion
Ex,y (H)a(X)Hee(YT))
He, Hex w)dady
(59)
20
/ Hon
Σ pm
(60)
m=0
Σ p "Xx=N( [Hem(X) (X) By-N(o.1) [He Hem
(61)
p 6nk
(62)
where in the last line we used the orthogonality property Fx (. 5nk-
B Additional experiments
B.1 Details about empirical validations
Hardware. Experiments that did not require training model were AMD Ryzen 3900X 12-Core CPU, which takes about minutes overall. All experiments that required training model were trained a single NVIDIA GeForce RTX 3090 GPU, which takes about minutes per each training task.
Figures. The solid lines all plots represent the average performance over multiple independent runs, and the shaded regions indicate the confidence intervals. Unless stated otherwise, each average is computed over #10 independent runs.
18

---

Codes and reproducibility. We implemented our experiments in Python using the Py PyTorch frame- work Paszke et al. [2019]. All the figures reproducible with the code attached in the supplementary.
Training procedure For all training-related experiments, the isometry or isometry gap computed per each batch by sampling a few of the batches randomly, and then averaged over. The epoch i corresponds to the network at after i steps of training on the training set CIFARIO 10 (epoch means network is at initialization).
Pre-trained large language models The pre-trained language models their default configuration was downloaded from Huggingface Wolf et al. [2020] library.
B.2 Quantifying the influence of gain on isometry through non-linearity strength
The concept of gain in neural networks is vital and closely connected with the weights initialization. neural network with properly initialized weights can learn faster, have a lesser chance of getting stuck at sub-optimal solutions, and provide better generalization. The impact of gain can be visualized through the lens of weight initialization strategies such as Xavier normalization [Glorot and Bengio. 2010], which has shown significant effectiveness optimizing neural networks. These initialization strategies apply gain value the weights. which scaling factor, ensure good signal flow through many layers during the forward and backward passes. The gain value essentially determines the variance of the weights in the initialization stage.
As an extension to our prior investigations, we delve into understanding the influence of gain on isometry, predominantly through our calculated metric, the non- non-linearity strength, denoted as 3a, as function of gain o. For certain instances, such ReLU, sine, and exponential activations, closed form. Table B.1 presents few of these cases.
capable of deriving (a) in
exp(ax))
sin(ar)
max(ax, 0)
Bg
2-20"² +a² 1-eat
2(-1+e
30
3xt
2-2m
Table B.1: Relationship between non-linearity strength and gain.
For more extensive selection of activation functions, have numerically computed the non-linearity strength function of gain a, visualized in Figure B.I. Leveraging Theorem 4, these computed values provide an estimation for the isometry strength for various activations. This correlation proves he remarkably predictive, as shown in Figure B.2. Remarkably, in the case of ReLU activation, its unique characteristics lead to both our closed form 3. (refer Table B. and rate towards isometry (refer Figure B.2) remaining consistent across different values of
2.0
3.8
sigmold
1.6
sin exp
1.4
122
1.0
alpha
Figure B.I: Analysis of gain through the lens of isometry strength.
Comparison Xavier gain for initialization Inspired by the results so far, we can compare mean-field centering and normalization to Xavier gain activations. Figure B.3 demonstrates that all mean-field based gains improve the isometry when compared with Xavier initialization. However, this is markedly stronger for ReLU and leaky ReLU. We explain this starker contrast by the fact that both activations have significant offset term co, which is not corrected by the Xavier initialization.
B.3 Varying width of hidden layers
While in our theoretical setup, we assume the network width is constant across the layers, this is only choice streamline our proof and notation. Since our primary result is derived from the mean-field regime, the
19

---

sin
relu
tanh
-log z(
10
alpha
10
alpha = 1.0
alpha = 2.0
10-2
0
10
20
30
0
10
20
30
0
10
20
30
Figure B.2: Influence of gain on the attainment of isometry.
selu
leaky_reih relu
relu
tanh
lag NG)
10"
MF
10
30
10
30
10
30
10
20
Figure B.3: Isometry vs depth for mean-field centering and normalization Xavier initialization.
only criterion for it hold for the width sufficiently large to approximate the mean-field regime. Our experiments in Figure B.4 substantiate this claim that the specific sizes hidden layers, as long as they are large, will not impact our main results on the isometry. We empirically validate this for four different configurations and show that the decay of the isometry gap remains largely consistent across these configurations.
dataset MNIST
hidden dimersions 1000]
1.2
[512. 1024, 2048. 4096. 8192] [8192 4096, 2048, 1024, 512]
1.0
Hooetry Gap
0.8
0.6
0.2
a550
KC
RC2
Ic
norma
res
ac5o
namm
[co
Jom
Pite
P
cuuna
Layer
Layer
Figure B.4: Theorem 4 holds for variable width MLP. Isometry gap of various layers for MLP
with ReLU activation and fixed width (blue), growing width (orange), narrowing width (green), and interchanging width (red) on MNIST (left) and CIFAR10 (right) training data. The weights of MLP are random. These plots show that our central theory holds regardless of the dataset or shape of MLP widths.
B.4 Isometry in pre-trained large language models
Since our theory for normalization other architectures. Figure B.5 shows the important role of normalization in the pre-trained GPT2 network. However, we need to adjust the notion of isometry with the architecture of layer norm in a transformer in mind. In fact, the mean and standard deviation are computed over features separately for each token. Thus, to adapt the notion of isometry, we view each token sample and define Gram over different tokens. Thus, isometry here quantifies the similarity between various tokens within one sample. As can seen the figure below, LayerNorm layers (shaded in red) in the last six layers of the pre- ve-ttrined GPT2 increase the isometry between tokens, which is consistent with our theory of layer normalization. It is crucial that our theory holds deterministically, which extends the pre- mre-trained model.
limited initialization, we can expand search isometry
20

---

One caveat in interpreting results is that, in practice. LayerNorm layers have learnable parameters that make them deviate from theory. It would be fruitful study the effects of learned parameters discer it from the role centering and normalization for a future study.
GPTZ, wie angi shi
Lysstem
the Coraaa
Figure B.5: Validation of corollary 2 in GPT2. The last six transformer layers of pre-trained GPT2, LayerNorm layers (shaded; red) decrease the isometry gap.
B.5 Tracking isometry at initialization and optimization for more activations
Here there are more numerical experiments related to to tracking isometry before and after training.
opoch
Signld, P,- 1.02
someryy 1007
0.4
SELU, 5.-1) 1.03
2:06
033
Retli R.
actiation
nomattorion
Figure B.6: Validation of Corollary and Theorem 4 for multiple activations.
The following plot shows that isometry gap remains relatively stable
arttvencon on
prenk
2:5
2.07
177
030
Figure B.7: Isometry gap change during training is relatively small.
21


Reviews:
[{"review_id": "JxUHt8jwiK", "reviewer": "Reviewer_tNrQ", "summary": "This paper studies the isometric properties of a randomly initialized neural network. The authors show layer normalization and proper activation functions can mitigate rank collapse. In addition, they also quantify the normalization bias for different type of normalization layers. They use the Hermite expansion of the activation function to highlight the importance of higher order Hermite coefficients in the bias towards isometry. The paper provides theoretical results and empirical evidence to support its findings and discusses the potential implications for future research on neural network architectures and training algorithms.", "questions": "I have some questions regarding the coherence of this paper, for example:\n1) what is the motivation of introducing normalization bias in Section 3.2, how does it related to the isometry bias?\n2) Is definition 1 the definition for ‘isometry’, or a measure of ‘isometry property’?\n3) Is the isometry gap (-logI(M)) ranges from 0 to -\\infinity? (It is \\infinity on line 100). \n\nI’m also curious how can the insights from this paper be used to design more efficient and effective deep learning systems? What is mainly the assumptions you make that present challenges of applying these insights in practice?", "limitations": "This paper is unlikely to have potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "(+) This paper provides analysis of the isometry properties using the penultimate Gram matrix in neural networks. It complements previous study by analyzing the role of layer normalization (while previous work mainly focus on batch norm). It also proposes a few useful measurement of isometry bias.\n\n(+) The authors provide theoretical results and empirical evidence (mainly figure 2 and figure 3) to show that activation and normalization techniques can bias the Gram matrix towards isometry at initialization, which can improve the training dynamics of deep neural networks.\n\n(+) The paper discusses the potential implications of the findings for future research on neural network architectures and training algorithms, which can inspire new directions for improving the performance and efficiency of deep learning systems.", "weaknesses": "(-) I feel this paper is not very coherent between different sections, though possibly because I’m not an expert in this field. I can understand each section but does not find the connection between different sections. \n\n(-) The empirical evidence presented in the paper is quite limited. Only very few activation functions are considered and it also assume the MLPs are of fixed width. The conclusion might be hard to generalize or be helpful to the practice.\n\n(-) The analysis of section 3 and section 4 seems to be two different systems. It does not explain how they are connected and lack a unified theory of them.\n\n(-) There are also a few confusing statements such as the authors show higher He coefficients have negative impact on isometry properties. However, in the abstract, they also state “ highlighting the importance of higher order (>2) Hermite coefficients in the bias towards isometry“, which implies higher order Hermite coefficients help isometry. \n\nOverall, I think this paper has limited applicability and does not inspire new ways on network design (either normalization or activation functions). ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time reviewing our work, providing constructive feedback, and pointing out potential points of confusion. \n\n*Detailed responses:*\n>  There are also a few confusing statements such as the authors show higher He coefficients have negative impact on isometry properties. However, in the abstract, they also state “ highlighting the importance of higher order (>2) Hermite coefficients in the bias towards isometry“, which implies higher order Hermite coefficients help isometry.\n\nTheoretical result aligns with the statement in the abstract. There is a subtle difference between isometry ($\\mathcal{I}$) and isometry gap ($-\\log\\mathcal{I}$): The negative logorithmic of isometry is the isometry gap that quantifies how far a matrix is from isometry. Thus Figure 5 shows that  $He_2, He_3,$ lead to a decay of isometry gap, which implies higher isometric properties. This is consistent with our main result and the summary stated in abstract. Extremium values of both isometry and isometry gap are elaborated in lines 97-101 of the main text: isometry lies between $0$ (degenerate matrix) and $1$ (perfect isometry, ie, orthogonal matrix), while isometry gap lies between $-\\infty$ (degnerate) and $0$ (orthogonal matrix, ie, perfect isometry). According to Theorem 4, the higher order Hermite cofficient imposes a faster decay in isometry gap, hence a higher faster convergence for isometry to 1. For example, there is no decay in $-\\log \\mathcal{I}$ for linear activations. Therefore, linear activation does not impose isometry which is experimentally substantiated in Figure 3. \n\n> The empirical evidence presented in the paper is quite limited. Only very few activation functions are considered \n\nTo address the reviewer's concerns about limitations of our theory, we have expanded empirical results to include activations `Identity, ReLU, PReLU, Tanh, SiLU(Swish), ELU, GELU, SELU.` Figure 2 in the PDf attached to the general response depicts the isometry across layers of an MLP for various activations, along with their $\\beta_0$ value. As can be seen, the value of $\\beta_0$ accurately predicts the isometry gap decay.\n\n\n> I feel this paper is not very coherent between different sections, though possibly because I’m not an expert in this field. I can understand each section but does not find the connection between different sections. The analysis of section 3 and section 4 seems to be two different systems. It does not explain how they are connected and lack a unified theory of them.\n\nBoth results are characterizing isometry across the layers of MLP. While Section 3 focus on a single normalization, Section 4 investigates the isometry for MLPs with layer normalization and non-linear activations. We will add an outline before section 3 to highlight the connection between different sections. In Figure 2 attached to the general response, we illustrate how these two results connects and when we can invoke results in these sections for deep neural networks. We will elaborate on the coherence . \n\n> and it also assume the MLPs are of fixed width. The conclusion might be hard to generalize or be helpful to the practice.\n\nThe assumption that MLP has fixed width is merely to ease the notation. In fact, one can readily extend the results to an MLP with variable widths across layersas long as these widths are sufficiently large, ie., $1/\\sqrt{width}$ is small). To substantiate this, we have added experiments with variable width (see PDF attached to general response, Figure 4).\n\n > Overall, I think this paper has limited applicability and does not inspire new ways on network design (either normalization or activation functions).  I’m also curious how can the insights from this paper be used to design more efficient and effective deep learning systems? What is mainly the assumptions you make that present challenges of applying these insights in practice?\n\nThe general response demonstrates the following applications: \n- **Predicting training speed with $\\beta_0$**: The non-linearity coefficient $\\beta_0$ is a strong indicator for training convergence speed (see general response and Figure 1 of attached PDF to the general response)\n- **Explaining isometry bias of activations** Theorem 4 and non-linearity strength $\\beta_0$ accurately predict the isometry of activations for a wide range of activations  (see Figure 2 of attached PDF to the general response)\n- **Explaining isometry bias of layer normalization in MLPs and transformers** Theorem 1 and corollary 2 prove the isometry-enhancing properties of  LayerNorm in MLPs (Figure 2, attached PDF) and transformers (Figure 3, attached PDF)\n\n> what is the motivation of introducing normalization bias in Section 3.2, how does it related to the isometry bias?\n\nThey both refer to the \"isometry bias of normalization.\" To avoid creating the impression that they are separate concepts, we will replace occurrences with \"isometry bias of normalization\" in the revised manuscript.  We introduced the normalization bias to show normalization layers impose isometry across the layers of neural networks. When normalization bias is significantly greater than zero, normalization layers constantly increases the isometry across the layers. Figure 2 and 3 illustrate the consequence of having high normalization bias. As we observe in Figure 3, having high normalization bias (in Figure 2) concludes $-\\log I$ decays across the layers of neural networks. \n\n> Is definition 1 the definition for ‘isometry’, or a measure of ‘isometry property’?\n\nIt is the definition for isometry. Notably, the terms of isometry is slightly different from isometric maps that preserves the distances. \n\n> Is the isometry gap $(-\\log I(M))$ ranges from 0 to $-\\infty$? (It is \\infinity on line 100).\n\n*Isometry gap ranges between $0$ and $\\infty$:* Since isometry is smaller than $1,$ its logarithm is non-positive $\\log I(M)\\le 0$. hence Isometry gap  is non-negative $-\\log I(M)\\ge 0$. "}, {"review_id": "dwvQM4Zor0", "reviewer": "Reviewer_nfuS", "summary": "This paper studies the isometry of Gram matrix under the effect of BN, LN and activation at initialization. For BN and LN, some results are obtained. For activation, most are empirical results.", "questions": "See weekness.", "limitations": "See weekness.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Study multiple factors that affect the isometry of Gram matrix.", "weaknesses": "1) The novelty of the theory for BN and LN is small.\n2) For activation, the conclusion is wekk. It is not strange to see different activation function have different effects. But so what and why?\n3) The paper does not justify why study isometry. Its connection to neural networks, such as training speed or generalization, should be demonstrated either empirically or theoretically.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time reviewing our work, providing constructive feedback, and pointing out potential points of confusion. \n\n> This paper studies the isometry of Gram matrix under the effect of BN, LN and activation at initialization. For BN and LN, some results are obtained. For activation, most are empirical results.\n\nTheorem 4 is a theoretical result on activations and one of our main contribution. This theorem is characterizing the isometric properties of activations. The empirical evidence presented are only meant to substantiate and complement this theoretical contribution. \n\n> It is not strange to see different activation function have different effects. But so what and why? The paper does not justify why study isometry. Its connection to neural networks, such as training speed or generalization, should be demonstrated either empirically or theoretically.\n\n\nWe directly address these points in supplementary experiments PDF attached to the general response (See figures 1-4).\n\n\n>  The novelty of the theory for BN and LN is small.\n\nThe isometry bias of normalization has been the subject of various theoretical studies, including [1,2,3,4, and 5]. While [1-4] establish a local isometry bias in a local neighborhood of certain inputs, we prove the global isometry bias for a wide range of inputs. Furthermore, [5] only proves the global stability for MLP with linear activations and BN. Our theoretical analysis derives a clear connection between non-linear activations and isometry. While pervious results only characterizes this bias for networks with random weights, Theorem 1 proves the isometry result for all networks with normalization layers. \n\n\n---\n\n**References** \n\n[1] Yang, Greg, et al. \"A mean field theory of batch normalization.\" ICLR (2019).\n\n[2] Li, Mufan, Mihai Nica, and Dan Roy. \"The neural covariance SDE: Shaped infinite depth-and-width networks at initialization.\" Advances in Neural Information Processing Systems (2022).\n\n[3] Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations, 2017.\n\n[4] Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. Advances in Neural Information Processing Systems, 2016.\n\n[5] Hadi Daneshmand, Amir Joudaki, and Francis Bach. Batch normalization orthogonalizes representations in deep random networks. Advances in Neural Information Processing Systems 2021. \n"}, {"review_id": "EVd4my3VIe", "reviewer": "Reviewer_D7v5", "summary": "The paper is a theoretical analysis in the mean-filed regime the of the second-to-last gram-matrix of MLP. It proves that the presence of layer normalization in conjunction with non-linear activation functions biases the input-output mapping at initialization towards an isometry.", "questions": "1. What kind of input datasets are used in Fig. 2-3-4 ? \n\n2. Why the unit-sphere in d dimension is denoted as $\\sqrt d$-sphere and not $d$-sphere? \n\n3. Minor: typo l 146: comma in place of full stop.", "limitations": "Not really. Perhaps the main limitation is that the author limit their analysis on MLP architectures. This makes the results of their paper not immediately applicable to transformers (due to the presence of the attention block) which are of much greater practical interest. Maybe this should be stated as a limitation of the current work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper presents a novel theoretical contribution that aims to enhance our understanding of the interaction between layer normalization and non-linear activation functions at initialization. The proof relies on general assumptions and provides a justification for the utilization of pre-layer normalized representations in order to prevent training instabilities associated with the collapse of the rank of the input-output gram matrix in a multi-layer perceptron (MLP) network.", "weaknesses": "The experimental tests are not well described, for instance the authors do not specify the datasets used neither in the \nmain text and nor in the captions (see also questions sections).\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback and their constructive comments. \n\n> ## Limitations:\n> - Not really. Perhaps the main limitation is that the author limit their analysis on MLP architectures. This makes the results of their paper not immediately applicable to transformers (due to the presence of the attention block) which are of much greater practical interest. Maybe this should be stated as a limitation of the current work.\n\nThis is a great suggestion by the reviewer. \nLet us remark that Theorem 1, and corrolaries 2 & 3 hold for any architecture and regardless of the weight distributions. Figure 3 in the `Rebuttal PDF` shows after each Layer Normalization layer in GPT2 (pretrained), the isometry gap decays. Given that the established property of isometry for normalization was deterministic, this is a consequence of these theoretical results.  \n\n> ## Weaknesses:\n> The experimental tests are not well described, for instance the authors do not specify the datasets used neither in the main text and nor in the captions (see also questions sections).\n\nWe thank the reviewer for pointing out these potential points of ambiguity regarding the experimental section. We will address this by adding detailed explanatioins regarding the experiments. \n\n> ## Questions:\n> - What kind of input datasets are used in Fig. 2-3-4 ?\n\nThe input used in these figures wes artificially created to have very low isometry (close to degenerate), so that it highlights the effects of the normalization and activation layers. More specifically, the rows of $d\\times n$ are drawn from $N(0, C)$ where $C$ has a highly skewed eigenvalue distribution. Same holds for figures 5,7, 8.  We thank the reviewer for pointing out these missing information in the main text. These details have been added to captions and main text. \n\n> - Why the unit-sphere in d dimension is denoted as $\\sqrt{d}$ sphere and not $d$ sphere?\n\n$\\sqrt{d}$ denotes the radius of the sphere, but not the dimensionality. The classical definition of LN projects each sample onto the sphere with radius $\\sqrt{d}$ as can be seen below: \n\\begin{align}\nLN(x) = \\frac{(x-\\bar x)}{\\sqrt{\\frac1d \\sum_i^d(x_i-\\bar x)^2}} \\implies \\\\|LN(x)\\\\| = \\frac{\\sqrt{\\sum_i^d(x_i-\\bar x)^2}}{\\sqrt{\\frac1d\\sum_i(x_i-\\bar x)^2}} = \\sqrt{d},\n\\end{align}\n\n> - Minor: typo l 146: comma in place of full stop.\n\nfixed. \n"}, {"review_id": "EVa9bnpYZH", "reviewer": "Reviewer_Qj2n", "summary": "The paper investigates normalization and non-linear activation functions from the perspective of isometry. ", "questions": "- One question while reading this paper is, why do we want to preserve the distance and angles between the input data points after mapping them to the feature space? Suppose in image recognition, the desired representation would not preserve the distance and angles in the input data points?\n- I think the previous question leads to another question which is also discussed in the paper: why would we want to use isometry to measure and explain layers? Despite many interesting points in the paper, I would suggest the authors add more motivation or probably a few past literature to better justify the motivation of this interesting work. ", "limitations": "Here are some minor comments:\n1. I would recommend the authors polish more the paper. For example, for the figures, I recommend the authors add more captions for readers to better follow and understand the message. \n2. I would also suggest improving the presentation a bit. For example, leaving two plots in the last page (page 9) may not be a good idea for the readers. ", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- This paper provides interesting analysis into normalization and non-linear activation from the perspective of isometry. \n- The paper has provided interesting discussion on future research. ", "weaknesses": "- Overall the paper looks a bit rushed or not polished. There are many typos for example, for example line 31 \"about the the\".\n- The paper may benefit more from slightly larger scale data. For example, using MNIST could be a good way to improve the paper without too much additional work. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive comments helping us to improve the writing. \n\n\n> The paper may benefit more from slightly larger scale data. For example, using MNIST could be a good way to improve the paper without too much additional work.\n\nTo address the reviewer's concern about scale of data and susbstantiate the validtiy of our theorems, we repeated experiments for MNIST (Figure 4), and CIFAR100 (Figure 1) and also Wikidataset for large language models (Figure 2) in the attached PDF to the general response. Let us remark that as theoretical analysis establishes, results holds indepdent from the dataset as long as some mild conditions are met, i.e., samples in the input batch are  linearly independent.  \n\n\n\n\n> One question while reading this paper is, why do we want to preserve the distance and angles between the input data points after mapping them to the feature space? Suppose in image recognition, the desired representation would not preserve the distance and angles in the input data points?\n\nPervious studies demonstrate isometry at initialization leads to a faster convergence of gradient descent (see general response). This is further illuserated in Figure 1 in document attached to the rebuttal respone. In this Figure, activation functions with a stronger isometry bias (i.e. a larger constant $\\beta_0$) lead to a faster convergence of stochastic gradient descent. Remarkably, this is only the matter of initialization. The representations will change after training. \n\n\n> I think the previous question leads to another question which is also discussed in the paper: why would we want to use isometry to measure and explain layers? Despite many interesting points in the paper, I would suggest the authors add more motivation or probably a few past literature to better justify the motivation of this interesting work.\n\nWe thank the reviewer for raising this lack of clarity on othe motivation and background. several works have attributed the success of normalization and residual/skip connections in training deep neural networks to a faitful signal propagation (also termed dynamic isometry) through deep between layers [1-7]. Here are a few past literature that are directly related to the notion of isometry.  \n\n- Dynamic isometry or signal propagation property postulates that in order to ensure a fast training[1,2], the network output must be sensitive to changes in the input. This hypothesis is employed by [4] to train a 10,000 layer CNN by a proper weight initialization, without skip connection or normalization layers.\n- In a similar attempt, [5] develops multiple methods to obtain isometry in transformer architecture that trains with a comparable speed to the model with skip and normalization layers. \n- [6] proposes shaping activations to impose isometry and empirically shows this allows training of very deep neural networks. \n- [7] designs activation functions that impose isometry, hence significantly enhance the training of networks with batch normalization. \n\nWe included an experimental result in the general response demonstrating applications of results, and will provide an extensive literature review, motivating our theoretical study. \n\n> # Limitations:\n> Here are some minor comments:\n> - I would recommend the authors polish more the paper. For example, for the figures, I recommend the authors add more captions for readers to better follow and understand the message.\n\nWe thank the reviewer for pointing out this lack of clarity in captions. We have added key details regarding the data and experimental procedure to each figure. \n\n\n> - I would also suggest improving the presentation a bit. For example, leaving two plots in the last page (page 9) may not be a good idea for the readers. Overall the paper looks a bit rushed or not polished. There are many typos for example, for example line 31 \"about the the\".\n\nWe will improve the writing. Thank you very much for the typo. \n\n---\n\n**References** \n[1] Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations, 2017.\n\n[2] Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. Advances in Neural Information Processing Systems, 2016.\n\n[3] Hadi Daneshmand, Amir Joudaki, and Francis Bach. Batch normalization orthogonalizes representations in deep random networks. Advances in Neural Information Processing Systems 2021. \n\n[4] Xiao, Lechao, et al. \"Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks.\" International Conference on Machine Learning 2018.\n\n[5] He, Bobby, et al. \"Deep transformers without shortcuts: Modifying self-attention for faithful signal propagation.\" arXiv preprint arXiv:2302.10322 (2023)\n\n[6] Zhang, Guodong, Aleksandar Botev, and James Martens. \"Deep learning without shortcuts: Shaping the kernel with tailored rectifiers.\" (2022).\n\n[7] Klambauer, Günter, et al. \"Self-normalizing neural networks.\" Advances in neural information processing systems 30 (2017).\n"}, {"review_id": "uyk3f70BWN", "reviewer": "Reviewer_B3t5", "summary": "A study of layernorm and Gramm matrix isometry is presented in both theory and emprical results. The results explain effectiveness of transformer (and similar) architechtures. Application of layernorm after activation seems to be a correct strategy when we are interested to isometry of output Gramm matrix. ", "questions": "Gap between theory (Theorem 4) and observed in Fig. 4 is quite big. As a future work, is there a clear path for obtaining tighter bound?", "limitations": "-", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "Good theoretical contributions in the field of explaining the effect of layernorm in deep neural networks via the concept of Gramm matrix isometry. Authors also added a large number of illustrative empirical results that help to drive home the main points of the paper. ", "weaknesses": "My main issue with the paper is not with the theory, but how it is presented in the main paper. The appendices appear to be fairly thorough mathematical exposition, but same rigorous exposition is absent in the main paper. Main paper should be readable without referencing to the appendices. \n\nHere under, I itemize places where I feel more rigour in exposition would be needed: \nFrom Theorem 1: \"in isometry as a function of fluctuations in norms\". What specifically is ment by \"fluctuations in norms\"?\n\nIsn't this statement obvious \"corollary 2 shows that the isometry of layer normalization is deterministic and does not rely on the random weights.\" ? In the setup of Corollary 2, no weight matrix is applied. So how this question even would be raised? Completely different issue is that if you would compare \nI(G_l) and I(G_{l+1}), there weight matrix and activation function are applied. Next sentence about empirical proof also is out of place in this context. \n\nIn Section 3.1. what is ment by the sentence \"which implies orthogonalization of these results.\"? Orthogonalization of what?\n\nIt would be good to expand the properties of Def 3 a bit more. Now we only have that it is \"non-linearity strength for activations\". Paper has half a page space for extra content. I would appreciate a few paragraphs of extra content in this context. \n\nAbout Eq. (7), you mention that G_*^l is a sequence that approximates the G^l. From these paragraphs it is hard to see where the sequence is coming from. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe would like to thank the reviewer for their positive review on our work\n\n\n> My main issue with the paper is not with the theory, but how it is presented in the main paper. ... some rigorous exposition is absent in the main paper. Main paper should be readable without referencing to the appendices.\n\nWe thank the reviewer for these constructive feedback and will address it in the revised manuscript. \n\n> ... From Theorem 1: \"in isometry as a function of fluctuations in norms\". What specifically is ment by \"fluctuations in norms\"?\n\nTo address this ambiguity we will replace \"fluctuation of norms\" by \"variations\"\n\n> - Isn't this statement obvious \"corollary 2 shows that the isometry of layer normalization is deterministic and does not rely on the random weights.\" ? In the setup of Corollary 2, no weight matrix is applied. So how this question even would be raised? Completely different issue is that if you would compare $I(G_l)$ and $I(G_{l+1})$, there weight matrix and activation function are applied. Next sentence about empirical proof also is out of place in this context.\n\nThis is a great and nuanced observation. We agree with the reviewer that the current phrasing does seem ambiguous. These two paragraphs (lines 124-130) stress that the non-decreasing isometry under `Layer Norm` is not a probablistic statement, but holds at all times. \n\nIn order to track and analyze evolution of isometry through depth, we are fundumentally faced with the following chain $\\dots \\to x^l \\to h^l \\to x^{l+1} \\to \\dots$ where $h^{l} = \\sigma(W^l \\tilde{x}^l)$ and $x^{l+1} = \\text{LayerNorm}(h^l).$ since $x^l$'s and $h^l$'s clearly depend on the preceding weights, one may attempt to anlyze their isometry by relying on the stochasticity of weights, and perhaps on their particular distribution (namely Gaussian weights.) Corrollary 2 and the following paragraphs assert that this isometry-preserving property does not rely on such stochastic nor mean-field approximations and holds at all times, implying that isometry of hidden representations $x_l$ cannot be smaller than isometry of hidden activations $h^l$. \n\nThe empirical evidence in the attached PDF to general response solidifies that, regardless other components, normalization layers preserve or enhance isometry/ Namely in MLP at initialization (Figure 2 left), and after training (Figure 2, right), and in pretrained transformers (Fig 3). \n\n> - In Section 3.1. what is ment by the sentence \"which implies orthogonalization of these results.\"? Orthogonalization of what?\n\nAssume that $\\mathcal{I}(h_\\ell) = 1$, then outputs becomes orthogonal to each other since the maximum determinant is obtained for orthogonal vectors. We will add more details about the notion of orthogonality and isometry. \n\n\n> - It would be good to expand the properties of Def 3 a bit more. Now we only have that it is \"non-linearity strength for activations\". Paper has half a page space for extra content. I would appreciate a few paragraphs of extra content in this context.\n\n\nWe acknowledge the need for further elaboration on non-linearity strength. We have a novel experiment confirming that the linearity strengths predicts the convergence of SGD in early epochs, discussed in the general response. We will include additional paragraphs to explain this concept more thoroughly in the revised manuscript, emphasizing its computation, and its influence on training and isometry. \n\n\n> - About Eq. (7), you mention that G_*^l is a sequence that approximates the G^l. From these paragraphs it is hard to see where the sequence is coming from.\n\nAppendix A.1 elaborates on the link between $G_{\\*}^l$ and $G^{l}$. The sequence $G_{\\*}^l$ refers to the mean-field approximation of the Gram matrix dynamics in the infinitely wide network regime, derived based on pervious mean field approaches in [1,2,3]. Here's the Breakdown of the ideas of why $G_{\\*}^l$ is an approximate of $G^l$.  $G^l$ is the actual Gram matrix of the representations at layer $l$. Given the stochastic weights, the dynamics of this sequence is *stochastic*.  The mean-field approximation $G_{\\*}^l$ is constructed to represent the average, or mean dynamics, such that conditionally it holds \n$\nE[G^{l+1} | G^l=G^l_\\*]=G_\\*^{l+1}.\n$\nIn contrast to the exact Gram sequence $G^l$ the mean-field Gram sequence $G_\\star^l$ is *deterministic.* As we increase the width of the neural network, we observe that dynamics of Gram matrices conditioned on the previous layer becomes highly concentrated around their mean. This implies that at infinite width, evolution of mean-field Gram matrices $G^l_*$ will exactly describe the evolution of Gram matrices $G_l$'s.  \n\n\n> ## Questions:\n> - Gap between theory (Theorem 4) and observed in Fig. 4 is quite big. As a future work, is there a clear path for obtaining tighter bound?\n\nThis is an excellent question. A likely path to a tighter bound is finding a tighter bound between $\\gamma(G)$ (Definition 4, appendix) and isometry gap. \n\n**References**\n\n[1] Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations, 2017.\n\n[2] Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. Advances in Neural Information Processing Systems, 2016.\n\n[3] Yang, Greg, et al. “A mean field theory of batch normalization.” ICLR (2019).\n"}]

Author rebuttal:
We appreciate constructive reviews that helped us to improve the paper. Let us recount our main contributions by showing excerpts from the reviews. 

This excerpt from `B3t5`'s captures the main part of contribution: 
> A study of layernorm and Gramm matrix isometry is presented in both theory and emprical results... Application of layernorm after activation seems to be a correct strategy when we are interested to isometry of output Gramm matrix.

This view is shared by reviewer `D7v5` highlighting applications of our result to training instabilities in deep neural networks 
> The proof relies on general assumptions and provides a justification for the utilization of pre-layer normalized representations in order to prevent training instabilities associated with the collapse of the rank of the input-output gram matrix in a multi-layer perceptron (MLP) network.

Furthermore, reviewer `Qj2n` views our analysis interesting:
> This paper provides interesting analysis into normalization and non-linear activation from the perspective of isometry.


Reviewers have also raised some concerns which we have  addressed with experiments and references:


## Background and motivation for isometry in the literature. 
Several studies have attributed the success of normalization and residual/skip connections in training deep neural networks to a faitful signal propagation (also termed dynamic isometry) through deep between layers [1-7]. Related literature demonstrate the important role of isometry in training stabilizing the training of deep neural networks: 

- Dynamic isometry or signal propagation property postulates that in order to ensure a fast training[1,2], the network output must be sensitive to changes in the input. This hypothesis is employed by [4] to train a 10,000 layer CNN by a proper weight initialization, without skip connection or normalization layers.

- [5] develops multiple methods to obtain isometry in transformer architecture that trains with a comparable speed to the model with skip and normalization layers.

- [6] proposes shaping activations to impose isometry and empirically shows this allows training of very deep neural networks.
- [7] designs activation functions that impose isometry, hence significantly enhance the training of networks with batch normalization.

## A  supplementary experiment
To elaborate on applications of our theoretical results, we supplement our findings with an experiment showing that: 
 
 **Non-linearity strength $\beta_0$ is a strong predictor for training performance** 
 
In Figure 1 of the attached PDF, we compare the convergence rates of stochastic gradient descent on a Multi-Layer Perceptron (MLP) with various activations for which $\beta_0$ (as defined in Eq. 6) is monotonically increasing. We observe activation functions with higher values of $\beta_0$ (as per Eq. 6) result in accelerated convergence for stochastic gradient descent. It is important to note that $\beta_0$ governs the isometry bias of the activation, as proven in Theorem 4. Consequently, our findings carry direct practical implications, shedding light on the popularity of certain activation functions such as ReLU in deep learning.

## Generality of our results for MLPs and transformers
Several of the questions raised are related to generality and broader applicability of our theoretical results. In response, we have provided the following explanatioins and clarifications: 
- **Isometry of LayerNorm after training**: Theorem 1 and corollary 2 are valid in any setting, regardless of the architecture and the distribution of weights This is demonstrated in Figure 2 of the attached PDF
- **Isometry of LayerNorm in transformers** Isometry-enhancing property of LayerNorm, as predicted by corollary 2, is architecture-independent. Figure 3 of the attached PDF shows isometry bias of normalization layers for a pretrained GPT2 architecture. 
- **Broad range of activations in MLPs** As demonstrated in Figure 2 of the attached PDF, the non-linearity strength  


We hope that with these additional clarifications and new empirical evidence, the reviewer's concerns are addressed.  


---

**References** 

[1] Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations, 2017.

[2] Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. Advances in Neural Information Processing Systems, 2016.

[3] Hadi Daneshmand, Amir Joudaki, and Francis Bach. Batch normalization orthogonalizes representations in deep random networks. Advances in Neural Information Processing Systems 2021. 

[4] Xiao, Lechao, et al. "Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks." International Conference on Machine Learning 2018.

[5] He, Bobby, et al. "Deep transformers without shortcuts: Modifying self-attention for faithful signal propagation." arXiv preprint arXiv:2302.10322 (2023)

[6] Zhang, Guodong, Aleksandar Botev, and James Martens. "Deep learning without shortcuts: Shaping the kernel with tailored rectifiers." (2022).

[7] Klambauer, Günter, et al. "Self-normalizing neural networks." Advances in neural information processing systems 30 (2017).

Meta-review: This paper considers the Gram matrix at the penultimate layer, and it proves that layer normalization biases such matrix towards an isometry. The rate of convergence is quantified via the Hermite expansion of the activation function, and it is exponential in the network depth. As it is typical of this line of work, the analysis is carried out at initialization. Empirical results are presented to validate the theory.

The authors have provided a comprehensive rebuttal addressing the concerns of the reviewers, who have now reached a consensus towards accepting this paper (although with different levels of support). After my own reading of the reviews, rebuttal and paper, I agree with this view and find the main results valuable and insightful. Hence, I am happy to recommend acceptance. I warmly encourage the authors to incorporate their rebuttal in the final version.
Decision: Accept (poster)