Title: On the impact of activation  and normalization in obtaining  isometric embeddings at initialization
Venue: NeurIPS 2023
Accepted: yes
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

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.

1 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 et al., 2016] have significantly influenced network training and have thus become the building blocks of neural networks. The practical success of these components has inspired extensive theoretical studies on the intricate role of weight initialization [Saxe et al., 2013, Daneshmand et al., 2021], normalization [Yang et al., 2019, Kohler et al., 2019, Daneshmand et 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 at 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 et al., 2021, Joudaki et al., 2023]. These issue of degeneracy or rank deficiency has been observed in multilayer perceptrons (MLPs) [Saxe et al., 2013, Daneshmand et al., 2020], convolutional networks [Bjorck et 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

indicates that normalization layers can act effectively to circumvent such Gram degeneracy, thereby improving training [Yang et al., 2019, Daneshmand et 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. [2018]. 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 at 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

A 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 a fixed point of a 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-points are not unique, and they can be degenerate or non-degenerate Yang et al. [2019]. In this paper, we establish a convergence rate to a non-degenerate fixed-point for a 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 theoretical 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.

A broad spectrum of activation functions such as ReLU [Fukushima, 1969], GeLU [Hendrycks and Gimpel, 2016], SeLU [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 a theoretical framework to characterize the influence of a broad range of activations on intermediate representations in DNNs.

3 Preliminaries

Notation. Let ⟨ x, y ⟩ be the inner product of vectors x and y, and ∥ x ∥ 2 = ⟨ x, x ⟩ the squared Euclidean norm of x . For a matrix X, we write X i · and X · i for the i -th row and column of X, respectively. We use W ∼ N ( µ, σ 2 ) m × n to indicate that W is an m × n Gaussian matrix with i.i.d. elements from N ( µ, σ 2 ) . We denote by 0 n the zero vector of size n. Given vector x ∈ R n , x denotes the arithmetic mean of 1 n ∑ n i =1 x i . Lastly, I n is the identity matrix of size n.

Normalization layers. Let LN : R d → R d and BN : R d × n → R d × n , denote batch normalization and layer normalization respectively. Table 1 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 as we include centering in the network architecture defined below.

MLP setup. The subject of our analysis is an MLP with constant width d across the layers and L layers, which takes input x ∈ R d and maps it to output x L ∈ R d , with hidden representations as

While the original ordering of layer normalization and activation is different [Ba et al., 2016], Xiong et al. [2020] show that the above ordering is more effective for large language models.

Gram matrices and isometry. Given n data points { x i } i ≤ n ∈ R d , the Gram matrix G ℓ of the feature vectors x ℓ 1 , . . . , x ℓ n ∈ R d at layer ℓ of the network is defined as

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 × n positive semi-definite matrix. We define the isometry I ( G ) of G as the ratio of its normalized determinant to its normalized trace:

I ( G ℓ ) is a scale-invariant quantity measuring the parallelepiped volume spanned by the feature vectors x ℓ 1 , . . . , x ℓ n . For example, consider two points on a plane x 1 , x 2 ∈ R 2 with lengths a = | x 1 | , b = | x 2 | and angle θ = ∠ ( x 1 , x 2 ) . The ratio is given by ab sin( θ ) / ( a 2 + b 2 ) , which is maximized when a = b and θ = π/ 2 . This relationship between volume and isometry is visually clear n = 2 and n = 3 feature vectors in Figure 1.

Remarkably, I ( M ) has the following properties (see Lemma A.1 for formal statements and proofs):

∥

x

1

∥∥

x

2

∥

θ

1

2

(

∥

x

1

∥

2

+

∥

x

2

∥

2

)

Vol

(

{

x

i

}

i

n

)

2

/n

1

n

∑

n

i

∥

x

i

∥

2

=

X

⊤

X

)

1

/n

1

n

tr(

X

⊤

X

)

x ⃗

1

x ⃗

2

θ

x ⃗

1

x ⃗

2

θ

x ⃗

1

x ⃗

2

θ

x ⃗

1

x ⃗

2

θ

x ⃗

1

x ⃗

2

x ⃗

3

x ⃗

1

x ⃗

2

x ⃗

3

x ⃗

1

x ⃗

2

x ⃗

3

x ⃗

1

x ⃗

2

x ⃗

3

(i) Scaling-invariant: For all constants c > 0 , we have I ( G ) = I ( cG ) .

(ii) Range: I ∈ [0 , 1] where the boundaries 0 and 1 are achieved for to degenerate and identity matrices respectively.

We also define the isometry gap as negative logarithm of isometry -log I ( M ) . Based on these properties of isometry, isometry gap lies between 0 and ∞ , with 0 and ∞ indicating the perfect isometry (identity matrix) and degenerate matrices respectively. Isometry allows us to establish the inherent bias of normalization layers in the following section.

4 Isometry bias of normalization

This section is devoted to discussing the remarkable property of isometry in the context of normalization. We present a theorem that formalizes this property, followed by its geometric interpretation and implications.

Theorem 1. Given n samples { x i } i ≤ n ⊂ R d \ { 0 d } , their projection onto the unit sphere ˜ x i := x i / ∥ x i ∥ , and their respective Gram matrices G and ˜ G , the isometry obeys

where a i := ∥ x i ∥ , and ¯ a := 1 n ∑ n i a i .

Geometric Interpretation. Isometry can be considered a measurement of the 'volume' of the parallelepiped formed by sample vectors, made scale and dimension-independent. The normalization process effectively equalizes the edge lengths of this parallelepiped, enhancing the overall "volume" or isometry, provided there is a variance in the sample norms. Thus, projection onto the unit sphere x i → x i / ∥ x i ∥ makes the edge lengths of the parallelepiped equal while leaving the angles between its edges intact. From this geometric perspective, Theorem 1 implies that among parallelepiped 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 a cube, where max volume is realized when all edge lengths are equal. This fact that maximum volume is achieved when edge lengths are equal can be deduced from the arithmetic vs 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 us to decouple the role of angles and edge lengths in volume formulation. This fact is evident for n = 2 in Figure 1. Since normalization does not modify the angles between edges, the remainder of the proof falls back onto the case where edges form a cube.

Proof of Theorem 1 . Define D := diag( a 1 , . . . , a n ) . Observe that G = D ˜ GD, implying det( G ) = det( ˜ G ) det( D ) 2 . Because ˜ x i 's have norm 1 , diagonals of Gram after normalization are constant ˜ G ii = 1 , implying 1 n tr( ˜ G ) = 1 . We have

Theorem 1 further shows a subtle property of normalization: as long as there is some variation in the sample norms, i.e., ∥ x i ∥ '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 a function of variation of norms. Intuitively, terms a and 1 n ∑ n i ( a i -¯ a ) 2 can be interpreted as the average and variance of sample norms a 1 , . . . , a n . Thus, a higher variation in the norms a i 's leads to a larger increase in isometry after normalization.

4.1 Implications for layer (and batch) normalization

Theorem 1 reveals insights into the biases introduced by layer and batch normalization in neural networks, particularly highlighting the improvement in isometry not just limited to initialization but also persistent through the training process.

Corollary 2. Consider n vectors before and after layer-normalization { x i } i ≤ n ⊂ R d \ { 0 d } and { ˜ x i } i ≤ n , ˜ x i := LN ( x i ) . Define their respective Gram matrices G := [ ⟨ x i , x j ⟩ ] i,j ≤ n , and ˜ G := [ ⟨ ˜ x i , ˜ x j ⟩ ] i,j ≤ n . We have:

What makes the above result distinct from related studies [Daneshmand et al., 2021, 2020, Yang et al., 2019] is that the increase in isometry is not limited to 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 et 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 a link between these two different normalization techniques.

Corollary 3. Given n samples in a mini-batch before X ∈ R d × n , and after normalization ˜ X = BN ( X ) and define covariance matrices C := XX ⊤ and ˜ C := XX ⊤ . We have:

Gram matrices of 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

fc-0

act-0

norm-1

fc-1

act-1

norm-2

fc-2

act-2

norm-3

fc-3

act-3

norm-4

fc-4

norm-5

fc-5

act-5

norm-6

fc-6

norm-7

fc-7

act-7

norm-8

fc-8

norm-9

fc-9

act-9

norm-10

fc-10

act-10

Layer

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Isometry

(

G

)

initialization

fc-0

act-0

norm-1

fc-1

act-1

norm-2

fc-2

act-2

norm-3

fc-3

act-3

norm-4

fc-4

norm-5

fc-5

act-5

norm-6

fc-6

norm-7

fc-7

act-7

norm-8

fc-8

norm-9

fc-9

act-9

norm-10

fc-10

act-10

Layer

epoch 9

activation

normalization

et al., 2020] and that it orthogonalizes the representations [Daneshmand et al., 2021], and that it imposes isometry [Yang et al., 2019]. It is straightforward to verify that orthogonal matrices have the maximum isometry. Thus, the increase in 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 3 is 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 or improve isometry. To highlight the fact that the claims of Corollary 2 holds at 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 2.

5 Isometry bias of non-linear activation functions

So far, our focus was individual normalization layers. In this section, we extend our 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 modifications in the scale or shape of activations can lead to significant changes in the representations. A powerful tool to analyze activations is to express activations in the Hermite polynomial basis. Inspired by previous successful applications of Hermite polynomials in 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 He k ( x ) , is defined as

All square-integrable function σ with respect to the Gaussian kernel, which obeys ∫ ∞ -∞ σ ( x ) 2 e -x 2 / 2 dx < ∞ , can be expressed as a linear combination of Hermite polynomials

as σ ( x ) = ∑ k c k He k ( x ) with (see section A for more details):

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 ℓ . We use the mean-field dynamic of Gram matrices subject of previous studies [Yang et al., 2019, Schoenholz et al., 2017, Poole et al., 2016]. The mean-field dynamics of Gram matrices is given by

This equation gives the expected Gram matrix for layer ℓ +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 on the notion of isometry strength of the activation function, defined next.

Definition 3 (Isometry strength) . Given activation σ with Hermite expansion { c k } k ≥ 0 , define its isometry strength β σ as:

We can readily check from the definition that isometry strength β σ has the following basic properties: (i) it ranges between 1 and 2 , and (ii) it is 1 if and only if the activation is a linear function. Table 2 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 to normalization layers.

Theorem 4. Let σ be an activation function with a Hermite expansion and a non-linearity strength β σ , (see equation (11) ). Given non-degenerate input Gram matrix G 0 ∗ , then for sufficiently large layer ℓ ≳ β -1 σ ( -n log I ( G 0 ∗ ) + log(4 n )) , we have

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 ℓ ∗ 's will be degenerate.

Theorem 4 reveals the importance of non-linear Hermite coefficients ( c k , k ≥ 2 ) in activation function to ensure β σ > 1 and obtain isometry in depth. This connection between β σ and isometry is the rationale for referring to β σ 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.

Figure 3 compares the established bound on the isometry gap with those observed in practice, i.e. G ℓ , for three activations. We observe β σ 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, in the next section, we will discuss other insights from the above analysis.

0

1

2

3

4

5

6

7

8

9

10

layer

0.0

0.2

0.4

0.6

0.8

1.0

-log

(

G

)

Sigmoid,

=1.012

0

1

2

3

4

5

6

7

8

9

10

layer

=1.070

0

1

2

3

4

5

6

7

8

9

10

layer

ReLU,

=1.264

Trace

empirical (

log

(

G

))

theory (

)

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

Experimental Setup. Our experiments utilize MLPs with layer normalization and various activation functions for the task of image classification on the CIFAR10 dataset [Krizhevsky et al.]. For training, we use stochastic gradient descent (SGD) with a 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, a ℓ and x ℓ respectively denote the post-activation and post-normalization vector for layer ℓ.

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 in obtaining isometry. In the mean-field regime, pre-activations follow standard Gaussian distributions, and thus the average post-activation a ℓ = 1 d ∑ d i a ℓ i will converge to their expectation a ℓ = E z ∼N (0 , 1) [ σ ( X )] = c 0 . This insight suggests an alternative way of obtaining isometry by explicitly removing the offset term from activation, i.e., replacing activation σ ( x ) by σ ( x ) -c 0 . Strikingly, our experiments presented in Figure 4 indicate that such replacement can also impose isometry. This result provides novel insights into the role of centering in layer normalization.

0

10

20

30

10

1

10 0

log

(

G

)

relu

0

10

20

30

tanh

0

10

20

30

sigmoid

0

10

20

30

exp

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 ∑ d i =1 ( a ℓ i -a ℓ ) 2 in layer normalization converges to variance

var z ∼ N (0 , 1) ( σ ( z )) = ∑ ∞ k =1 c 2 k =: ¯ σ (1) . Thus, as the width increases, the layer normalization operator LN ( a ℓ -a ℓ ) will converge to ( a ℓ -c 0 ) / √ ¯ σ (1) . Figure 6 demonstrates that the constant scaling 1 / √ ¯ σ (1) , achieves comparable isometry to layer normalization for hyperbolic tangent and sigmoid and ReLU, while it is not effective for exp function. This observation calls for future research on the link between normalization and activation in deep neural networks.

0

10

20

30

10

1

10 0

10 1

log

(

G

)

relu

0

10

20

30

tanh

0

10

20

30

sigmoid

0

10

20

30

exp

0

10

20

30

10

1

10 0

10 1

log

(

G

)

relu

0

10

20

30

tanh

0

10

20

30

sigmoid

0

10

20

30

exp

6.3 Isometry strength correlates with SGD convergence rate in shallow MLPs.

Besides the direct consequences of our theory, we observe a striking correlation between the convergence of SGD and isometry strength in a specific range of neural network hyper-parameters. Figure 7 shows the convergence of SGD is faster for activations with a significantly larger isometry strength β σ (see Definition 3) for shallow MLPs, e.g., with 10 layers or less. We can speculate that this correlation reflects the input-output sensitivity of the networks with higher non-linearity. Surprisingly, this correlation does not extend to deeper networks. This discrepancy between shallow and deep networks regarding SGD convergence may be due to the issue of gradient explosion studied by Meterez et al. [2023]. This finding suggests multiple avenues for future research.

0

2

4

6

8

0.5

1.0

1.5

2.0

2.5

loss

Sigmoid,

=1.02

=1.07

ReLU,

=1.26

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. It is worth investigating whether we can impose isometry without layer normalization. Our empirical observations suggest that certain activations, such as ReLU, require layer normalization to attain isometry. In contrast, other activations, which can be considered as 'self-normalizing' (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, we believe Hermit expansion provides a theoretical grounding to analyze the isometry of SeLU.

0

5

10

15

20

25

30

10 0

log

(

G

)

relu

tanh

selu

0

5

10

15

20

25

30

10

1

10 0

log

(

G

)

x

+1 =

LN

(

(

W x

))

x

+1 =

(

LN

(

W x

))

(a)

(b)

Impact of the ordering of normalization and activation layers on isometry. Theorem 4 highlights that the ordering of activation and normalization layers has a critical impact on the isometry. Figure 8 demonstrates that a different ordering can lead to a non-isotropic Gram matrix. Remarkably, the structure analyzed in this paper is used in transformers [Vaswani et al., 2017].

Normalization'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 / √ width ) error per layer when the network width is 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 to the isometry bias induced by normalization, as proven in Theorem 1. 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.

References

Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415 , 2016.

0

10

20

30

10

1

10 0

log

(

G

)

d=100

0

10

20

30

d=500

0

10

20

30

d=1000

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 to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120 , 2013.

Hadi Daneshmand, Amir Joudaki, and Francis Bach. Batch normalization orthogonalizes representations in deep random networks. Advances in 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 arXiv:1902.08129 , 2019.

Jonas Kohler, Hadi Daneshmand, Aurelien Lucchi, Thomas Hofmann, Ming Zhou, and Klaus Neymeyr. Exponential convergence rates for batch normalization: The power of length-direction decoupling in non-convex optimization. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 806-815. PMLR, 2019.

Hadi Daneshmand, Jason D 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. On 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 on Artificial Intelligence and Statistics , pages 1924-1932, 2018.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems , 30, 2017.

Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya 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 , 35:27198-27211, 2022a.

Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard E. Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. In International Conference on Learning Representations , 2018.

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 depthand-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 normalization. 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 arXiv:2206.03126 , 2022b.

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

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

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

Bobby He, James Martens, Guodong Zhang, Aleksandar Botev, Andrew Brock, Samuel L Smith, and Yee Whye Teh. Deep transformers without shortcuts: 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 arXiv:1607.06450 , 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 normalization layers: Sharpness reduction. Advances in Neural Information Processing Systems , 35: 34689-34708, 2022.

Ekdeep S Lubana, Robert Dick, and Hidenori Tanaka. Beyond batchnorm: towards a unified understanding of normalization in deep learning. Advances in Neural Information Processing Systems , 34:4778-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.

Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 (canadian institute for advanced research). URL http://www.cs.toronto.edu/~kriz/cifar.html .

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.02012 , 2023.

Mufan Li, Mihai Nica, and Dan Roy. The future is log-gaussian: Resnets and their infinite-depth-andwidth limit at initialization. Advances in Neural Information Processing Systems , 34:7852-7864, 2021.

Semyon Aranovich Gershgorin. Uber die abgrenzung der eigenwerte einer matrix. News of the Russian Academy of Sciences. Mathematical series , (6):749-754, 1931.

F 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, high-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, R'emi Louf, Morgan Funtowicz, Jamie Brew, and Guillaume DulacArnold. Transformers: State-of-the-art natural language processing. Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pages 38-45, October 2020. URL https://www.aclweb.org/anthology/2020.emnlp-demos.6 .

Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics , pages 249-256. JMLR Workshop and Conference Proceedings, 2010.

Appendix outline

The appendix is partitioned into four main components, each serving its purpose as described:

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.1 elaborates on the basic properties of isometry.

Section A.2 provides an elaborate review of the mean-field Gram dynamics.

Section A.3 presents the Lyapunov function γ , and establishes that this function provides both upper and lower bounds for the isometry, thereby implying that geometric contraction of γ ( G ℓ ) indicates a geometric contraction of isometry gap -log I ( G ℓ ) .

Section A.4 proves that γ ( G ℓ ) exhibits an exponential contraction in depth with rate β σ .

Section B 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 isometrypreserving properties:

Lemma A.1. For PSD matrix M, the isometry defined in (3) obeys the following properties: 1) scale-invariance I ( cM ) = I ( M ) , 2) only takes value in the unit range I ( M ) ∈ [0 , 1] 3) it takes its maximum value if and only if M is identity I ( M ) = 1 ⇐⇒ M = I n , and 3) takes minimum value if and only if M is degenerate I ( M ) = 0 .

Proof of Lemma A.1. The scale-invariance is trivially true as scaling M by any constant will scale det( M ) 1 /n 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 ) = ( ∏ i λ i ) 1 /n / ( 1 n ∑ i λ i ) , where λ i 's are eigenvalues of M. By arithmetic vs geometric mean inequality over the eigenvalues we have ( ∏ i λ i ) 1 /n ≤ 1 n ∑ i λ i ) , which proves that I ( M ) ∈ [0 , 1] . Furthermore, the inequality is tight iff the values are all equal λ 1 = · · · = λ n , which holds only for identity M = I n . 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):

Assuming that inputs are encoded as columns of X , we can restate the MLP dynamics as follows

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

where ⊗ denotes Hadamard product, and subscript d emphasises the dependence of Gram on width d. Note that conditioned on the previous layer, rows of H ℓ and A ℓ are i.i.d. , because of independence of rows of W ℓ . Thus, by law of large numbers, in the infinitely wide network regime, µ i and s i will converge to the expected mean and variance respectively lim d →∞ µ i = E A ℓ 1 i , and lim d →∞ s i = √ Var ( A ℓ 1 i ) , for all i = 1 , . . . , n. By construction of ϕ, in the infinitely wide regime, we can rewrite Gram dynamics as lim d →∞ G ℓ d = 1 d ∑ d k =1 ϕ ( A ℓ i · ) ⊤ ϕ ( A ℓ i · ) . We can invoke the fact that rows of A ℓ are i.i.d. to conclude that G d is the sample Gram matrix that converges to its expectation

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 a potential

Here we will introduce a Lyapunov function that enables us to precisely quantify the isometry of activations in deep networks:

Definition 4. Given a positive semidefinite matrix G ∈ R n × n , we define γ : R n × n → R ≥ 0 as:

̸

Remarkably, γ 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 ii = 1 , i = 1 , . . . , n. It holds:

Thus, we may apply Theorem A.2 iteratively to prove that in the mean-field, the Lyapunov function γ ( G ℓ ) decays at an exponential rate β σ . A straightforward induction over layers leads to a decay rate in γ , which is presented in the next corollary.

Corollary A.3. If activation σ has Hermite expansion with non-linearity strength β σ , the mean-field Gram matrices G ℓ ∗ , obey Lyapunov of these Gram matrices decays at an exponential rate β σ :

where G 0 ∗ denotes the input Gram matrix.

In Figure A.1, we experimentally validated the above equation for MLPs with a finite width and activations { tanh ,relu, sigmoid } where we observe that the mean-field analysis well predicates the decay in γ .

0

10

20

30

10 1

10 3

(

G

)

relu

0

10

20

30

tanh

0

10

20

30

sigmoid

theory

empirical

Interestingly, we can connect the Lyapunov γ to the isometry, by proving an upper and lower based on the determinant G based on γ ( G ) , when G is PSD and has unit diagonals:

Lemma A.4. For PSD matrix G with unit diagonals holds:

̸

where the lower bound holds if ( n -1) | G ij | ≤ 1 .

̸

Now, we are ready to prove the main theorem.

Theorem A.5 (Restated Theorem 4) . Let σ be an activation function with a Hermite expansion and a nonlinearity strength β σ , (see equation (11) ). Given non-degenerate input Gram matrix G 0 ∗ , then for sufficiently large layer ℓ ≳ β -1 σ ( -n log I ( G 0 ∗ ) + log(4 n )) , we have

Proof of Theorem 4. First, note that we have the following transformation

̸

̸

̸

By upper bound of Lemma A.4 we have max i = j | G ij | ≤ √ 1 -det( G ) ≤ 1 -det( G ) / 2 , implying γ ( G 0 ) ≤ 2 det( G 0 ) -1 , which allows us to

which concludes the proof.

Proof of Lemma A.4. Lower bound. The lower bound is a result of Gershgorin circle theorem [Gershgorin, 1931], which implies that every eigenvalue must be within the disc [1 -( n -1) max i = j | G ij | , 1+( n -1) max i = j | G ij | ] . Thus, the determinant is lower-bounded by (1 -( n -1) max i = j | G ij | ) n .

̸

Upper bound. Since G is PSD, we can write G = X ⊤ X, which implies that columns of X are unit norm ∥ x i ∥ = 1 , and G encodes the angles between them cos ∠ ( x i , x j ) = ⟨ x i , x j ⟩ = G ij . Furthermore, we have det( G ) = vol ( X ) 2 , where the volume refers to the parallelepiped spanned by the columns of X. With this formulation, we write volume recursively as volume spanned by columns 1 up to n -1 , times the projection distance of x n from their span

̸

where span refers to the space of all linear combinations of these vectors. Note that projection distance of x i onto x j can be written as sin ∠ ( x i , x j ) = √ 1 -⟨ x i , x j ⟩ 2 = √ 1 -G 2 ij . Since the projection distance onto the linear span cannot be greater than projection distance onto a single vector, which is bounded by √ 1 -max i = j G 2 ij . This concludes the upper bound that det( G ) ≤ 1 -max i = j G 2 ij .

A.4 Dual activation and proof of Thm. A.2

According to Daniely et al. [2016], we leverage the notion dual activation associated with the activation function σ, which is defined in following.

Definition 5. Given activation σ : R → R that is square integrable with respect to the Gaussian kernel, define its dual activation ˆ σ : R → R , and its mean reduced dual ¯ σ : R → R as:

The following lemma connects the dual activation and its mean-reduced version with the Hermite expansion of σ :

Lemma A.6. Given two standard Gaussian variables X,Y ∼ N (0 , 1) with covariance E XY = ρ, and activation σ with normalized Hermite coefficients c kk , we have

̸

we have

̸

̸

Finally, we have the tools to prove the first main theorem.

Proof of Theorem A.2. Let a = σ ( h ) for h ∼ N (0 , G ) . Note that by definition of the dual-activation we have E aa ⊤ = [ˆ σ ( G ij )] i,j ≤ n . Since we assumed G has unit diagonals, we have h i ∼ N (0 , 1) , which implies that E a i = E h i ∼ N (0 , 1) σ ( h i ) = c 0 , implying that E ( a i -E a i )( a j -E a j ) ⊤ = ˆ σ ( G ij ) -c 2 0 = ¯ σ ( G ij ) . Furthermore, the variance can be driven as Var( a i ) = E a 2 i -( E a i ) 2 = ∑ ∞ k =0 c 2 k = ∑ ∞ k =1 c 2 k = ¯ σ (1) for all i = 1 , . . . , n. Thus, we have E ϕ ( a i ) ϕ ( a j ) = ¯ σ ( G ij ) / ¯ σ (1) . In the matrix form we have

The remainder proof relies on the following contractive property of Gram matrix potential:

Lemma A.7. Consider activation σ, with normalized Hermite coefficients { c k } k ≥ 0 . For all ρ ∈ (0 , 1) , the mean-reduced dual activation ¯ σ obeys

̸

which the right hand-side is strictly larger if some nonlinear coefficient is nonzero c k = 0 for some k ≥ 2 .

̸

Thus we can apply Lemma A.7 on each element i = j to conclude that

̸

̸

(42)

since the inequality holds for any value of i = j, we can take the maximum over i = jj to write:

̸

̸

̸

which concludes the proof.

Proof of Lemma A.7. Note the ratio is invariant to scaling of ∑ ∞ k =1 c 2 k . Hence, we assume ∑ ∞ k =1 c 2 k = 1 without loss of generality. With this simplification, we have ¯ σ (1) = 1 . For the positive range ρ ∈ [0 , 1] we have

By Jensen inequality for convex function x ↦→| x | we have

Thus for ρ ∈ [0 , 1] we have

Because x ↦→ x/ (1 -x ) is monotonically increasing for x ∈ [0 , 1] , we have

Where we invoked the inequality that was proven for ρ ∈ [0 , 1] , because | ρ | ∈ [0 , 1] .

The following lemma, which is a consequence of Mehler's formula, is at the hart of proof of Lemma A.6: Lemma A.8 (Consequence of Mehler's kernel) . If X,Y ∼ N (0 , 1) with covariance EXY = ρ we have

where δ nk is the Dirac delta.

Proof of Lemma A.6. Let σ ( x ) = ∑ k =0 c k He k ( x ) , denote the Hermite expansion of σ. Thus, we have

where in the last line we applied result of Lemma A.8. For the mean-reduced dual kernel ¯ σ, observe that E X ∼ N (0 , 1) σ ( X ) = c 0 . Thus, the reduction of mean will cancel the k = 0 term, which concludes the proof that ¯ σ ( ρ ) = ∑ k =1 c 2 k ρ k .

Proof of Lemma A.8. The property can be deduced from Mehler's formula Mehler [1866]. The formula states that

where the m ! factor difference is due to the definition of Hermite polynomials with an additional 1 / √ m ! compared to the one used in Mehler's kernel. Observe that the left hand side is equal to p ( x, y ) /p ( x ) p ( y ) , where p ( x, y ) is the joint PDF of ( X,Y ) , and p ( x ) , p ( y ) are PDF of X and Y respectively. Therefore, we can take the expectation using the expansion

where in the last line we used the orthogonality property E X ∼ N (0 , 1) H k ( x ) H n ( X ) = δ nk .

B Additional experiments

B.1 Details about empirical validations

Hardware. Experiments that did not require training a model were run on a AMD Ryzen 9 3900X 12-Core CPU, which takes about 5 minutes overall. All experiments that required training a model were trained a single NVIDIA GeForce RTX 3090 GPU, which takes about 1 minutes per each training task.

Figures. The solid lines in 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.

Codes and reproducibility. We implemented our experiments in Python using the PyTorch framework Paszke et al. [2019]. All the figures are reproducible with the code attached in the supplementary.

Training procedure For all training-related experiments, the isometry or isometry gap are 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 of CIFAR10 (epoch 0 means network is at initialization).

Pre-trained large language models The pre-trained language models and 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. A 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 in optimizing neural networks. These initialization strategies apply a gain value to the weights, which is a scaling factor, to ensure a 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-linearity strength, denoted as β σ , as a function of gain α. For certain instances, such as ReLU, sine, and exponential activations, we are capable of deriving β σ ( α ) in a closed form. Table B.1 presents a few of these cases.

For a more extensive selection of activation functions, we have numerically computed the non-linearity strength as a function of gain α , as visualized in Figure B.1. Leveraging Theorem 4, these computed values provide an estimation for the isometry strength for various activations. This correlation proves to be remarkably predictive, as shown in Figure B.2. Remarkably, in the case of ReLU activation, its unique characteristics lead to both our closed form β σ (refer Table B.1) and rate towards isometry (refer Figure B.2) remaining consistent across different values of α .

0

2

4

6

8

10

1.0

1.2

1.4

1.6

1.8

2.0

0

sigmoid

tanh

exp

selu

alpha

Comparison to Xavier gain for initialization Inspired by the results so far, we can compare mean-field centering and normalization to Xavier gain for 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 can explain this starker contrast by the fact that both activations have a significant offset term c 0 , 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 a choice to streamline our proof and notation. Since our primary result is derived from the mean-field regime, the

0

10

20

30

10

2

10

1

10 0

log

(

G

)

0

10

20

30

relu

0

10

20

30

tanh

alpha = 0.5

alpha = 1.0

alpha = 2.0

0

10

20

30

10 0

log

(

G

)

selu

0

10

20

30

leaky_relu

0

10

20

30

relu

0

10

20

30

tanh

MF

xavier

only criterion for it to hold is for the width to be sufficiently large to approximate the mean-field regime. Our experiments in Figure B.4 substantiate this claim that the specific sizes of 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.

act-0

fc-0

act-1

norm-1

fc-1

act-2

norm-2

fc-2

act-3

norm-3

fc-3

norm-4

fc-4

act-5

norm-5

fc-5

Layer

0.2

0.4

0.6

0.8

1.0

1.2

Isometry Gap

dataset = MNIST

act-0

fc-0

act-1

norm-1

fc-1

act-2

norm-2

fc-2

act-3

norm-3

fc-3

norm-4

fc-4

act-5

norm-5

fc-5

Layer

dataset = CIFAR10

[1000, 1000, 1000, 1000, 1000]

[512, 1024, 2048, 4096, 8192]

[8192, 4096, 2048, 1024, 512]

[1024, 2048, 1024, 2048, 1024]

B.4 Isometry in pre-trained large language models

Since our theory for normalization is not limited to initialization, we can expand our search for isometry to 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 can view each token as a sample and define Gram over different tokens. Thus, isometry here quantifies the similarity between various tokens within one sample. As can be seen in the figure below, LayerNorm layers (shaded in red) in the last six layers of the pre-trained 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 to the pre-trained model.

One caveat in interpreting our results is that, in practice, LayerNorm layers have learnable parameters that make them deviate from our theory. It would be fruitful to study the effects of learned parameters to discern it from the role of centering and normalization for a future study.

h.7.Add_1

h.7.ln_1

h.7.attn.c_attn

h.7.attn.c_proj

h.7.attn.resid_dropout

h.7.attn

h.7.Add_2

h.7.ln_2

h.7.mlp.c_fc

h.7.mlp.act

h.7.mlp.c_proj

h.7.mlp.dropout

h.7.mlp

h.7

h.8.ln_1

h.8.attn.c_attn

h.8.attn.c_proj

h.8.attn.resid_dropout

h.8.attn

h.8.Add_2

h.8.ln_2

h.8.mlp.c_fc

h.8.mlp.act

h.8.mlp.c_proj

h.8.mlp.dropout

h.8.mlp

h.8

h.9.Add_1

h.9.ln_1

h.9.attn.c_attn

h.9.attn.c_proj

h.9.attn.resid_dropout

h.9.attn

h.9.Add_2

h.9.ln_2

h.9.mlp.c_fc

h.9.mlp.act

h.9.mlp.c_proj

h.9.mlp.dropout

h.9.mlp

h.9

h.10.Add_1

h.10.ln_1

h.10.attn.c_attn

h.10.attn.c_proj

h.10.attn.resid_dropout

h.10.Add_2

h.10.ln_2

h.10.mlp.c_fc

h.10.mlp.act

h.10.mlp.c_proj

h.10.mlp.dropout

h.10.mlp

h.10

h.11.Add_1

h.11.ln_1

h.11.attn.c_attn

h.11.attn.c_proj

h.11.attn.resid_dropout

h.11.attn

h.11.Add_2

h.11.ln_2

h.11.mlp.c_fc

h.11.mlp.act

h.11.mlp.c_proj

h.11.mlp.dropout

h.11.mlp

h.11

0

2

4

6

Isometry Gap

GPT2, last 6 layers, wiki (simple english)

LayerNorm

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.

fc-0

act-0

norm-1

fc-1

act-1

norm-2

fc-2

act-2

norm-3

fc-3

act-3

norm-4

fc-4

norm-5

fc-5

act-5

norm-6

fc-6

norm-7

fc-7

act-7

norm-8

fc-8

norm-9

fc-9

act-9

norm-10

fc-10

act-10

Layer

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Isometry

(

G

)

initialization

fc-0

act-0

norm-1

fc-1

act-1

norm-2

fc-2

act-2

norm-3

fc-3

act-3

norm-4

fc-4

norm-5

fc-5

act-5

norm-6

fc-6

norm-7

fc-7

act-7

norm-8

fc-8

norm-9

fc-9

act-9

norm-10

fc-10

act-10

Layer

epoch 9

Sigmoid,

=1.02

SELU,

=1.03

ELU,

=1.06

=1.07

SiLU,

=1.2

ReLU,

=1.27

activation

normalization

The following plot shows that isometry gap remains relatively stable.

0

2

4

6

8

10

layer_idx

0.50

0.75

1.00

1.25

1.50

1.75

2.00

2.25

2.50

iso_gap

activation

Sigmoid,

=1.02

=1.07

ReLU,

=1.26

0

9

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nA 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.\n\nSTRENGTHS:\nGood 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.\n\nWEAKNESSES:\nMy 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.\n\nQUESTIONS:\nGap 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\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

comments: SUMMARY:
The paper investigates normalization and non-linear activation functions from the perspective of isometry.

STRENGTHS:
- This paper provides interesting analysis into normalization and non-linear activation from the perspective of isometry. 
- 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".
- 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.

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?
- 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:
1. 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. 
2. 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

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

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 
main text and nor in the captions (see also questions sections).

QUESTIONS:
1. What kind of input datasets are used in Fig. 2-3-4 ? 

2. Why the unit-sphere in d dimension is denoted as $\sqrt d$-sphere and not $d$-sphere? 

3. 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

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

STRENGTHS:
Study multiple factors that affect the isometry of Gram matrix.

WEAKNESSES:
1) The novelty of the theory for BN and LN is small.
2) For activation, the conclusion is wekk. It is not strange to see different activation function have different effects. But so what and why?
3) 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.

QUESTIONS:
See weekness.

LIMITATIONS:
See weekness.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n(+) 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.\n\nWEAKNESSES:\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. \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).\n\nQUESTIONS:\nI 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?\n\nLIMITATIONS:\nThis paper is unlikely to have potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}