SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Amir Joudaki
amir.joudaki@inf.ethz.ch
Hadi Daneshmand
MIT/LIDS-FODSI
dhadi@mit.edu
INRIA-ENS-PSL Paris
francis.bach@inria.fr
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.
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
Code is available at: https://github.com/ajoudaki/deepnet-isometry
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 on 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.,
[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.
Notation. Let  $\langle x,y\rangle$  be the inner product of vectors  $x$  and  $y$ , and  $\|x\|^2 = \langle x,x\rangle$  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 \sim N(\mu, \sigma^2)^{m \times n}$  to indicate that  $W$  is an  $m \times n$  Gaussian matrix with i.i.d. elements from  $N(\mu, \sigma^2)$ . We denote by  $\mathbf{0}_n$  the zero vector of size  $n$ . Given vector  $x \in \mathbb{R}^n$ ,  $\overline{x}$  denotes the arithmetic mean of  $\frac{1}{n}\sum_{i=1}^{n}x_i$ . Lastly,  $I_n$  is the identity matrix of size  $n$ .
Normalization layers. Let  $\mathrm{LN}:\mathbb{R}^d\to \mathbb{R}^d$  and  $\mathrm{BN}:\mathbb{R}^{d\times n}\rightarrow \mathbb{R}^{d\times 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.
Width	d ∈ N	Batch size	n ∈ N
Depth	L ∈ N	Input	x ∈ Rd
Input batch	X ∈ Rd×n	Gaussian weights	W1, ..., WL ∼ N(0,1)d×d
Activation	σ: R → R	Centering	x - x̄
Layer Norm	LN(x) = x/√1/d ∑i d xi2	Batch Norm	BN(X)ij = Xij/√1/n ∑k d Xik2
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 \in \mathbb{R}^d$  and maps it to output  $x^L \in \mathbb{R}^d$ , with hidden representations as
$$
\left\{ \begin{array}{l l} x ^ {\ell + 1} = \frac {1}{\sqrt {d}} \mathrm {L N} \left(h _ {\ell} - \bar {h} _ {\ell}\right), & h _ {\ell} = \sigma \left(W ^ {\ell} x ^ {\ell}\right), \quad \ell = 0, \dots , L - 1 \\ x ^ {0} := \frac {1}{\sqrt {d}} \mathrm {L N} (x - \bar {x}), & \text {i n p u t .} \end{array} \right. \tag {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\leq n} \in \mathbb{R}^{d}$ , the Gram matrix  $G^{\ell}$  of the feature vectors  $x_{1}^{\ell},\ldots ,x_{n}^{\ell} \in \mathbb{R}^{d}$  at layer  $\ell$  of the network is defined as
$$
G ^ {\ell} := \left[ \langle x _ {i} ^ {\ell}, x _ {j} ^ {\ell} \rangle \right] _ {i, j \leq n}, \quad \ell = 0, 1, \dots , L. \tag {2}
$$
We define the notion of isometry to measure how much  $G^{\ell}$  is close to a scaling factor of the identity matrix.
Definition 1. Let  $G$  be an  $n \times n$  positive semi-definite matrix. We define the isometry  $\mathcal{I}(G)$  of  $G$  as the ratio of its normalized determinant to its normalized trace:
$$
\mathcal {I} (G) := \frac {\det  (G) ^ {1 / n}}{\frac {1}{n} \operatorname {t r} (G)}. \tag {3}
$$
$\mathcal{I}(G^{\ell})$  is a scale-invariant quantity measuring the parallelepiped volume spanned by the feature vectors  $x_{1}^{\ell}, \ldots, x_{n}^{\ell}$ . For example, consider two points on a plane  $x_{1}, x_{2} \in \mathbb{R}^{2}$  with lengths  $a = |x_{1}|, b = |x_{2}|$  and angle  $\theta = \angle (x_{1}, x_{2})$ . The ratio is given by  $ab \sin(\theta) / (a^{2} + b^{2})$ , which is maximized when  $a = b$  and  $\theta = \pi/2$ . This relationship between volume and isometry is visually clear  $n = 2$  and  $n = 3$  feature vectors in Figure 1.
Remarkably,  $\mathcal{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  $\mathcal{I}(G) = \mathcal{I}(cG)$ .
(ii) Range:  $\mathcal{I} \in [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 \mathcal{I}(M)$ . Based on these properties of isometry, isometry gap lies between 0 and  $\infty$ , with 0 and  $\infty$  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\leq n} \subset \mathbb{R}^{d} \setminus \{\mathbf{0}_{d}\}$ , their projection onto the unit sphere  $\widetilde{x}_i \coloneqq x_i / \| x_i \|$ , and their respective Gram matrices  $G$  and  $\widetilde{G}$ , the isometry obeys
$$
\mathcal {I} (\widetilde {G}) \geq \mathcal {I} (G) \left(1 + \frac {\frac {1}{n} \sum_ {i} ^ {n} \left(a _ {i} - \bar {a}\right) ^ {2}}{\bar {a} ^ {2}}\right), \tag {4}
$$
where  $a_{i}:=\|x_{i}\|$ , and  $\bar{a}:=\frac{1}{n}\sum_{i}^{n}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} \to 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 \coloneqq \mathrm{diag}(a_1, \ldots, a_n)$ . Observe that  $G = D\widetilde{G}D$ , implying  $\operatorname*{det}(G) = \operatorname*{det}(\widetilde{G})\operatorname*{det}(D)^2$ . Because  $\widetilde{x}_i$ 's have norm 1, diagonals of Gram after normalization are constant  $\widetilde{G}_{ii} = 1$ , implying  $\frac{1}{n}\mathrm{tr}(\widetilde{G}) = 1$ . We have
$$
\begin{array}{l} \frac {\mathcal {I} (\widetilde {G})}{\mathcal {I} (G)} = \frac {\frac {1}{n} \operatorname {t r} (G)}{\frac {1}{n} \operatorname {t r} (\widetilde {G})} \frac {\det  (\widetilde {G}) ^ {1 / n}}{\det  (G) ^ {1 / n}} (5) \\ = \frac {\frac {1}{n} \sum_ {i} ^ {n} a _ {i} ^ {2}}{1} \frac {\det (\widetilde {G}) ^ {1 / n}}{\det (\widetilde {G}) ^ {1 / n} \left(\prod_ {i} ^ {n} a _ {i} ^ {2}\right) ^ {1 / n}} (6) \\ = \frac {\left(\frac {1}{n} \sum_ {i} a _ {i}\right) ^ {2}}{\left(\prod_ {i} ^ {n} a _ {i}\right) ^ {2 / n}} \frac {\frac {1}{n} \sum_ {i} ^ {n} a _ {i} ^ {2}}{\left(\frac {1}{n} \sum_ {i} a _ {i}\right) ^ {2}} \quad \det  (D) = \prod_ {i} ^ {n} a _ {i} ^ {2} (7) \\ \geq 1 \cdot \frac {\frac {1}{n} \sum_ {i} ^ {n} a _ {i} ^ {2}}{\left(\frac {1}{n} \sum_ {i} ^ {n} a _ {i}\right) ^ {2}} \quad \frac {1}{n} \sum_ {i} ^ {n} a _ {i} \geq \left(\prod_ {i} ^ {n} a _ {i}\right) ^ {\frac {1}{n}} (8) \\ = 1 + \frac {\frac {1}{n} \sum_ {i} ^ {n} \left(a _ {i} - \bar {a}\right) ^ {2}}{\bar {a} ^ {2}}, \quad \bar {a} := \frac {1}{n} \sum_ {i} ^ {n} a _ {i}. (9) \\ \end{array}
$$
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  $\overline{a}$  and  $\frac{1}{n}\sum_{i}^{n}(a_{i} - \bar{a})^{2}$  can be interpreted as the average and variance of sample norms  $a_1,\ldots ,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 \leq n} \subset \mathbb{R}^d \setminus \{\mathbf{0}_d\}$  and  $\{\widetilde{x}_i\}_{i \leq n}, \widetilde{x}_i := \mathrm{LN}(x_i)$ . Define their respective Gram matrices  $G := [\langle x_i, x_j \rangle]_{i,j \leq n}$ , and  $\widetilde{G} := [\langle \widetilde{x}_i, \widetilde{x}_j \rangle]_{i,j \leq n}$ . We have:
$$
\mathcal {I} \left(\widetilde {G}\right) \geq \mathcal {I} (G) \left(1 + \frac {\frac {1}{n} \sum_ {i} ^ {n} \left(a _ {i} - \bar {a}\right) ^ {2}}{\bar {a} ^ {2}}\right), \quad \text {w h e r e} a _ {i} := \| x _ {i} \|, \bar {a} := \frac {1}{n} \sum_ {i} ^ {n} a _ {i}.
$$
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 \in \mathbb{R}^{d \times n}$ , and after normalization  $\widetilde{X} = \mathrm{BN}(X)$  and define covariance matrices  $C := XX^{\top}$  and  $\widetilde{C} := XX^{\top}$ . We have:
$$
\mathcal {I} \left(\widetilde {C}\right) \geq \mathcal {I} (C) \left(1 + \frac {\frac {1}{d} \sum_ {i} ^ {d} \left(a _ {i} - \bar {a}\right) ^ {2}}{\bar {a} ^ {2}}\right), \quad \text {w h e r e} a _ {i} := \| X _ {i}. \|, \bar {a} := \frac {1}{n} \sum_ {i = 1} ^ {d} a _ {i}.
$$
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
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.
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 hold 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  $\mathcal{I}(G^{\ell})$  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  $H_{k}(x)$ , is defined as
$$
H e _ {k} (x) := (k!) ^ {- \frac {1}{2}} (- 1) ^ {k} e ^ {\frac {x ^ {2}}{2}} \frac {d ^ {k}}{d x ^ {k}} e ^ {- \frac {x ^ {2}}{2}}.
$$
All square-integrable function  $\sigma$  with respect to the Gaussian kernel, which obeys  $\int_{-\infty}^{\infty}\sigma (x)^{2}e^{-x^{2} / 2}dx < \infty$ , can be expressed as a linear combination of Hermite polynomials
as  $\sigma (x) = \sum_{k}c_{k}He_{k}(x)$  with (see section A for more details):
$$
c _ {k} := \mathbb {E} _ {x \sim \mathcal {N} (0, 1)} [ \sigma (x) H e _ {k} (x) ].
$$
The subsequent section will discuss how to leverage the Hermite expression of the activation to analyze the dynamics of Gram matrix isometry.
5.2 Non-linear activations bias Gram dynamics towards isometry
In this section, we analyze how  $\mathcal{I}(G^{\ell})$  changes with  $\ell$ . 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
$$
G _ {*} ^ {\ell + 1} = \mathbb {E} _ {h \sim N \left(0, G _ {*} ^ {\ell}\right)} \left[ \phi (\sigma (h)) \phi (\sigma (h)) ^ {\top} \right], \quad \text {w h e r e} [ \phi (a) ] _ {i} := \left(a _ {i} - \mathbb {E} a _ {i}\right) / \sqrt {\operatorname {V a r} a _ {i}}. \tag {10}
$$
This equation gives the expected Gram matrix for layer  $\ell + 1$ , based on the Gram matrix  $G_{*}^{\ell}$  from the previous layer, and  $\phi$  is mean-field regime counterpart for layer normalization operator (see section A for more details). The sequence  $G_{*}^{\ell}$  approximates the dynamics of  $G^{\ell}$ , 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.