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.
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.,
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.
Table 1: Building blocks we consider in this work.
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
input.
(1)
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:
.
(3)
) 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
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):
Figure 1: A geometric interpretation of isometry: higher volume, in the second row, corresponds to higher value for isometry.
(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.
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 ā := 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.
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 -ā) 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.
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
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:
where
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
2 act-2 norm-3 fc-3 act-3 norm-4 fc-4 act-4 norm-5 fc-5 act-5 norm-6 fc-6 act-6 norm-7 fc-7 act-7 norm-8 fc-8 act-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 norm-0 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 act-4 norm-5 fc-5 act-5 norm-6 fc-6 act-6 norm-7 fc-7 act-7 norm-8 fc-8 act-8 norm-9 fc-9 act-9 norm-10 fc-10 act-10 Layer epoch 9 activation normalization Figure 2: Validation of Corollary 2 Isometry (y-axis) vs different layer of an 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 parameters: activation: tanh, depth: 10, width: 1000, batch-size: 512, training SGD on training set of CIFAR10. with lr = 0.01.. Layer names are encoded as type-index, where type can be fc: fully connected, norm: LayerNorm, and act: activation. 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.
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.
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.
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.
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(4n)), 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.
Table 2: Isometry strength β σ (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 β σ 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. 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 ℓ.
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. Mean-field centering (blue): x ℓ+1 = LN(a ℓ -c0).
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 Figure 6: Comparing layer normalization (blue) x ℓ+1 = LN(a ℓ -c0) vs. mean-field normalization (orange)
in obtaining isometry.
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.
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 ) 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. Mean-field normalization (blue):
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.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 βσ. 2. 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.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 = In, and 3) takes minimum value if and only if M is degenerate
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 = In. Finally, isometry is zero iff at least one eigenvalue is zero, which is the case for degenerate matrix M.
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 si will converge to the expected mean and variance respectively lim d→∞ µi = EA ℓ 1i , and lim d→∞ si = Var(A ℓ 1i ), for all i = 1, . . . , n. By construction of ϕ, in the infinitely wide regime, we can rewrite Gram dynamics as
). 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.
Here we will introduce a Lyapunov function that enables us to precisely quantify the isometry of activations 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 Gii = 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.
where G 0 * denotes the input Gram matrix.
In 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)|Gij| ≤ 1.
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 Eaa ⊤ = [σ(Gij)] i,j≤n . Since we assumed G has unit diagonals, we have hi ∼ N (0, 1), which implies that Eai = E h i ∼N (0,1) σ(hi) = c0, implying that E(ai -Eai)(aj -Eaj) ⊤ = σ(Gij)-c 2 0 = σ(Gij). Furthermore, the variance can be driven as 1) for all i = 1, . . . , n. Thus, we have Eϕ(ai)ϕ(aj) = σ(Gij)/σ(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
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
Thus for ρ ∈ [0, 1] we have
By Jensen inequality for convex function x → |x| we have
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.
σ exp(αx) sin(αx) max(αx, 0)
Table B.1: Relationship between non-linearity strength and gain. 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 c0, which is not corrected by the Xavier initialization.
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 0 log (G ) selu 0 10 20 30 leaky_relu 0 10 20 30 relu 0 10 20 30 tanh MF xavier Figure B.3: Isometry vs depth for mean-field centering and normalization to Xavier initialization.
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 norm-0 fc-0 act-1 norm-1 fc-1 act-2 norm-2 fc-2 act-3 norm-3 fc-3 act-4 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 norm-0 fc-0 act-1 norm-1 fc-1 act-2 norm-2 fc-2 act-3 norm-3 fc-3 act-4 norm-4 fc-4 act-5 norm-5 fc-5
Layer dataset = CIFAR10 hidden dimensions [1000, 1000, 1000, 1000, 1000] [512, 1024, 2048, 4096, 8192] [8192, 4096, 2048, 1024, 512] [1024, 2048, 1024, 2048, 1024]
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. B.5 Tracking isometry at initialization and optimization for more activations there are more numerical experiments related to to tracking isometry before and after training. norm-0 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 act-4 norm-5 fc-5 act-5 norm-6 fc-6 act-6 norm-7 fc-7 act-7 norm-8 fc-8 act-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 norm-0 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 act-4 norm-5 fc-5 act-5 norm-6 fc-6 act-6 norm-7 fc-7 act-7 norm-8 fc-8 act-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 Tanh, = 1.07 SiLU, = 1.2 ReLU, = 1.27 activation normalization Figure B.6: Validation of Corollary 2 and Theorem 4 for multiple activations.
The following plot shows that isometry gap remains relatively stable.