1 Introduction↩︎

Neural networks induce highly nonconvex loss landscapes that are complicated to study. Despite this complexity, the Hessian matrix of the loss, encoding only its second-order information, has played a central role in designing efficient optimization algorithms [1], [2], deriving measures of generalization [3], [4] – which, in turn, inspired the design of well-generalizing optimizers [5], [6] – as well as practical algorithms for pruning [7], [8], quantization [9], continual learning [10], [11], etc.

Our motivation resembles that of [12], [13], in that we aim to precisely characterize the full spectrum of the Hessian. However, while these works prioritize minimal assumptions at the cost of model complexity, we leverage standard assumptions in deep learning theory in exchange for networks with arbitrary widths and depths, and datasets with arbitrary number of samples, features and labels. Our methodology falls closest to [14], which studies the spectrum of the Hessian and training dynamics at the edge-of-stability (EoS) [15] for matrix factorization problems. In contrast, we consider the problem of supervised learning, with the overarching goal of understanding the influence of data-geometry on the sharpness of the learnt solution.

2 Setup↩︎

Consider the inputs \(\mathbf{x}_1, \ldots, \mathbf{x}_n \in \mathbb{R}^{1\times d_0}\) collected in \(\mathbf{X}\in \mathbb{R}^{n\times d_0}\), and outputs \(\mathbf{y}_1, \ldots, \mathbf{y}_n \in \mathbb{R}^{1\times d_L}\) collected in \(\mathbf{Y}\in \mathbb{R}^{n\times d_L}\). Our model is an \(L\)-layer linear neural network, defined as

with \(\mathbf{W}_{\ell} \in \mathbb{R}^{d_{\ell} \times d_{\ell-1}}\) is the learnable weight matrix in the \(\ell^{\textsuperscript{th}}\) layer, \(\forall \ell \in \left[L\right] \mathrel{\vcenter{:}}= \{1, 2, \ldots, L\}\).¹ In what follows, we denote the model parameters by \(\mathbf{W}_{1:L} \mathrel{\vcenter{:}}= \left(\mathbf{W}_1, \mathbf{W}_2, \ldots, \mathbf{W}_L\right)\). The loss function is the Mean-Squared Error (MSE):

where \(\left\|\cdot\right\|_F\) denotes the Frobenius norm of the argument.

2.1 Approximating the Hessian-spectrum↩︎

We use the generalized Gauss-Newton (GGN) approximation of the Hessian, which gets more and more precise as loss reduces with training:

with the total number of parameters given by \(p = \sum_{\ell=1}^L d_{\ell} d_{\ell-1}\), and the Jacobian w.r.t.predictions given by \(\mathbf{J}\mathrel{\vcenter{:}}= \partial \mathbf{f}\left(\mathbf{X}\right) / \partial \mathbf{W}_{1:L} = \begin{bmatrix} \mathbf{J}_1 & \mathbf{J}_2 & \ldots & \mathbf{J}_{L} \end{bmatrix} \in \mathbb{R}^{d_Ln\times p}\), where

where \(\otimes\) is the Kronecker outer-product. The GGN matrix shares its non-zero eigenvalues with the Neural Tangent Kernel (NTK) matrix [17], which is simpler to study:

Going forward, we will make use of the Singular Value Decomposition (SVD) of the feature matrix, \(\mathbf{X}= \mathbf{U}_{\mathbf{X}}\mathbf{\Sigma}\mathbf{V}_{\mathbf{X}}^\intercal\), where \(\mathbf{U}_{\mathbf{X}}\in\mathbb{R}^{n\times r}\) is a semi-orthogonal matrix, \(\mathbf{\Sigma}\in\mathbb{R}^{r\times r}\) is a diagonal matrix with \(\mathbf{\Sigma}_{11} \geq \mathbf{\Sigma}_{22} \geq \ldots \geq \mathbf{\Sigma}_{rr} > 0\), and \(\mathbf{V}_{\mathbf{X}}\in\mathbb{R}^{r\times r}\) is an orthogonal matrix; \(r\) is the rank of \(\mathbf{X}\).

This is a rather strong assumption, and not realistic in practice, but it helps draw insights from the bounds we derive on the eigenvalues of the NTK.

2.2 Shallow Networks↩︎

We visualize the case of isotropic features in 1, where we note that the approximation becomes increasingly precise as training progresses. 1 implies that, of the \(d_0d_1+d_1d_2\) eigenvalues of the Hessian, at most \(\min\left\{d_0,d_1\right\} \times \min\left\{d_1,d_2\right\}\) may be non-zero; we comment on this at the end of 3.

Consider the case where \(\mathbf{\Sigma}= \sqrt{n} \mathbf{I}_{d_0}\). The spectral norm (largest eigenvalue for symmetric matrices) of the NTK is given by \(\lambda_{\max}\left(\frac{1}{n}\mathbf{J}\mathbf{J}^\intercal\right) = \left\|\mathbf{W}_1\right\|_2^2 + \left\|\mathbf{W}_2\right\|_2^2\). This contradicts the conclusion drawn in [13], suggesting that the maximum eigenvalue is given by \(\max\left\{\left\|\mathbf{W}_1\right\|_2^2, \left\|\mathbf{W}_2\right\|_2^2\right\}\); we believe that the discrepancy is caused by the block-diagonal structure of the Hessian assumed in their analysis. This distinction is visualized in 2, where the sum of squared spectral-norms gets increasingly good at approximating the sharpness of the Hessian, as training progresses.

As for the sum of eigenvalues, \(\text{Tr}\left(\frac{1}{n}\mathbf{J}\mathbf{J}^\intercal\right) = d_2 \left\|\mathbf{W}_1\right\|_F^2 + d_0 \left\|\mathbf{W}_2\right\|_F^2\). This relation provides insight into the progressive flattening and subsequent sharpening observed in early- and intermediate-training [19], [20], respectively, before training reaches EoS [15], [21], [22] – sharpness decreases since the parameter norm decreases in early-training, as the model output evolves to align with the labels, and both start increasing after a point [23]. In 3, we demonstrate this relation for a shallow network as well as for a deep network, showing a high amount of agreement in the dynamics of \(\text{Tr}\left(\frac{1}{n}\mathbf{J}\mathbf{J}^\intercal\right)\) and the squared parameter-norm.

3 Strongly Balanced Deep Networks↩︎

In this section, we derive analogous results to 1 but for networks of arbitrary depth, while relying on an additional assumption on the balance between layers.

In essence, this implies that the left singular vectors of a layer are aligned with the right singular vectors of the next layer. Moreover, all layers share their non-zero singular values; we denote them by \(\sigma_i\left(\mathbf{W}\right)\). While arbitrary points in the parameter space need not correspond to balanced layers, initially balanced layers remain balanced under gradient flow training, provided their is no nonlinearity between them [24], [25]. Even if initially unbalanced, [23] showed that the layers become increasingly balanced under weight-decay training, while [14] showed that gradient descent training at EoS [15] balances the singular values of the weight matrices.Therefore, the results hereon can be seen as characterizing the local geometry of the minimizers.

Hence, of the \(p \mathrel{\vcenter{:}}= \sum_{\ell=1}^L d_{\ell}d_{\ell-1}\) eigenvalues of the Hessian, at most \(\left(\min\left\{d_0,d_1,\ldots,d_L\right\}\right)^2\) are non-zero. This corresponds to the empirical observations suggesting that the bulk of eigenvalues of the Hessian are near-zero [26]–[28]. We visualize the case of isotropic features in 4, where we note that the approximation becomes increasingly precise as training progresses.

We distinguish two cases: 1. when the features are highly anisotropic (typical case for real-world datasets), and 2. when the features are isotropic (more idealistic). We observe that in the former case, the lower bound is tight, while in the latter case the upper bound is tight; see 5. This suggests that the true sharpness floats between the two bounds as the degree of isotropy in the features changes.

4 Dependence on Data↩︎

We visualize this correspondence in the left panel in 6, where we note that the theory predicts the non-zero eigenvalues exactly. Dependence on dataset size. Assuming the entries of \(\mathbf{X}\) and \(\mathbf{Y}\) are constant in the datasets size, i.e.\(\mathbf{X}_{ij}, \mathbf{Y}_{ij} = \mathbf{\Theta}_n\left(1\right)\), their spectral-norms scale as \(\mathbf{\Theta}\left(\sqrt{n}\right)\), and hence, the sharpness of balanced solutions is independent of dataset size (see 8), in disagreement with [15], which suggests that sharpness increases with \(n\). Their observation may be better understood through the geometry of the loss landscape and the behavior of optimization algorithms, which are beyond the scope of this work.

Dependence on depth. Sharpness increases with depth when \(\left\|\mathbf{X}\right\|_2 \leq \left\|\mathbf{Y}\right\|_2\), or \(L \geq 2\ln \frac{\left\|\mathbf{X}\right\|_2}{\left\|\mathbf{Y}\right\|_2}\) – we visualize this in the middle panel in 6.

Dependence on features. For a centered feature matrix \(\mathbf{X}\), i.e.the empirical mean of each feature-dimension is zero, \(\{\sigma_{i}^{2}\left(\mathbf{X}\right)/n\}_{i=1}^{d_0}\) represents the amount of variance explained by the principle components. Under 1, this is given by \(\mathbf{\Sigma}_{11}^2\). Hence, larger the magnitude of the features, the more spread out they are, and larger is the sharpness of the solution.

Dependence on labels. For one-hot encoded labels, \(\{\sigma_{i}^{2}\left(\mathbf{Y}\right)/n\}_{i=1}^{d_L}\) is the empirical distribution of labels, i.e.proportion of labels from each class in the training set; in what follows, we will use \(\alpha_i \mathrel{\vcenter{:}}= \sigma_{i}^{2}\left(\mathbf{Y}\right)/n\) to denote these proportions. This implies that the model learns sharper solutions when there is one class with disproportionately large number of labels; we illustrate this in the right panel in 6. Such datasets are, intuitively, simpler to learn, e.g.as an extreme case, a dataset with all inputs belonging to one class should be easier to learn than a dataset with uniform label distribution. This intuition is in disagreement with [15], which suggests that the model achieves lower peak-sharpness on simpler datasets.

Ablations. We repeat the experiment in 6, shedding each of 1 2 one-at-a-time. In the left panel in 7, we present the results shedding the assumption on whitened inputs. In this case, our theory breaks down immediately – the sharpness of the learnt solution turns out to have inconsistent variations with maximum label proportion in the dataset, suggesting that peak-sharpness is strongly influenced by the input-geometry. In the middle panel, we present the results shedding the assumption on balanced weights, and use the standard Kaiming uniform initialization [29] implemented in the PyTorch framework. Here, we note the theoretical predictions made in previous section stay intact, as the final sharpness increasing with label imbalance. Finally, we inspect if incorporating nonlinearity breaks our theory. In the right panel, we present results for a Tanh-MLP, noting that while the results are slightly different from the theoretical predictions, they are qualitatively similar, allowing us to extend our conclusions on the effects of label distribution to such nonlinear networks.

5 Omitted Proofs↩︎

6 Other Results↩︎

6.1 Summary of Eigenvalues↩︎

Several works have used Hessian-based measures like the maximum eigenvalue (spectral norm for symmetric matrices) or the sum of eigenvalues (trace) to measure generalization. Intuitively, at a local minimum, the spectral norm measures the worst-case change in loss caused by a small perturbation in any direction, while the trace measures the expected change in loss caused by a random isotropic perturbation with small variance.

6.2 Other Figures↩︎

7 Experimental Details↩︎

Data Whitening. Given the SVD of a feature matrix \(\mathbb{R}^{n\times d_0} \ni \mathbf{X}' = \mathbf{U}_{\mathbf{X}}\mathbf{\Sigma}\mathbf{V}_{\mathbf{X}}^\intercal\), we redefine the features to be proportional to the polar factor, i.e.\(\mathbf{X}= \sqrt{n}\sigma_{\mathbf{x}}\mathbf{U}_{\mathbf{X}}\mathbf{V}_{\mathbf{X}}^\intercal\). Sampling Balanced Weights. Assume we are given a sequence of \(k=\min\left\{d_0,d_1,\ldots,d_L\right\}\) singular values; we arrange them to create a diagonal matrix, \(\mathbf{S}\in \mathbb{R}^{k\times k}\). To construct a sequence of balanced weight matrices, we start by sampling the right singular vectors of \(\mathbf{W}_1\) by orthogonalizing a random \(d_0\times k\) matrix with \(\mathcal{N}\left(0,1\right)\) entries; denote it by \(\mathbf{V}_1 \in \mathbb{R}^{d_0\times k}\). Then, for each \(\ell \in \left[L\right]\), do

1. Sample the left singular vectors of \(\mathbf{W}_{\ell}\) by orthogonalizing a random \(d_{\ell}\times k\) matrix with \(\mathcal{N}\left(0,1\right)\) entries; denote it by \(\mathbf{U}_{\ell} \in \mathbb{R}^{d_1\times k}\).

2. Define the weight matrix as \(\mathbf{W}_{\ell} \mathrel{\vcenter{:}}= \mathbf{U}_{\ell}\mathbf{S}\mathbf{V}_{\ell}^\intercal\in \mathbb{R}^{d_{\ell}\times d_{\ell-1}}\).

3. Set the right singular vectors of the next layer as the left singular vectors of the previous layers, \(\mathbf{V}_{\ell+1} \mathrel{\vcenter{:}}= \mathbf{U}_{\ell}\).

Subsampling. With real-world datasets MNIST [30], FashionMNIST [31] and CIFAR10 [32], we use subsets of \(n=5,000\) samples. If the maximum label proportion \(\alpha_{\max}\) is set, we take \(\lfloor(1-\alpha_{\max})n\rfloor\) samples from each of the classes \(1\!-\!9\), where \(\lfloor\cdot\rfloor\) denotes the floor function, and the rest of the samples from class \(0\).

Training. All models are trained using full-batch gradient descent for \(10,000\) epochs, with an initial learning rate \(\eta_{\max}\) which is successively halved anytime \(<0.1\%\) improvement in loss is observed over \(100\) epochs.

1: Top-10 eigenvalues of the Hessian, computed numerically using LOBPCG and analytically using 1. Datasets are orthogonalized with \(\sigma_{\mathbf{x}}=1\), and our model is a 2-layer Linear-MLP with width 256, trained with \(\eta_{\max}=1e-3\).

2: Top eigenvalue of the Hessian, along with the predictions in 1 and [13]. Datasets are orthogonalized with \(\sigma_{\mathbf{x}}=1\), and our model is a 2-layer Linear-MLP with width 256, trained with \(\eta_{\max}=1e-3\).

3: Sum of eigenvalues of the Hessian is computed numerically using Hutchinson’s method [18]. Datasets are orthogonalized with \(\sigma_{\mathbf{x}}=1\), and our models are a 2-layer Linear-MLP with width 256 and a 4-layer Linear-MLP with width 64, trained with \(\eta_{\max}=5e-4\).

4: Top-10 eigenvalues of the Hessian, computed numerically using LOBPCG and analytically using 2. Datasets are orthogonalized with \(\sigma_{\mathbf{x}}=1\), and our model is a 4-layer Linear-MLP with width 64, trained with \(\eta_{\max}=1e-2\).

5: The OLS solution is computed as \(\mathbf{W}^\intercal\mathrel{\vcenter{:}}= \mathbf{W}_1^\intercal\mathbf{W}_2^\intercal\ldots\mathbf{W}_L^\intercal= \mathbf{X}^{\dagger}\mathbf{Y}\), and a network with balanced layers is constructed to simulate this solution by setting \(\mathbf{V}_1\) and \(\mathbf{U}_L\) as the right and left singular vectors of \(\mathbf{W}\). Our model is a 2-layer Linear-MLP with width 128.

6:

• Left: We sample \(\mathbf{X}' \in \mathbb{R}^{20\times 100}\) with standard normal entries, and then orthogonalize it with \(\sigma_{\mathbf{x}} = 3\). We sample \(d=20\) singular values from \(\mathcal{U}\left[0,1\right]\), and use it sample a sequence of \(L=5\) balanced weight matrices, with \(\mathbf{W}_{\ell} \in \mathbb{R}^{d\times d}\). Finally, set \(\mathbf{Y}= \mathbf{W}_5\mathbf{W}_4\ldots\mathbf{W}_1\mathbf{X}\) so that the GGN approximation is exact.

• Right: Spectral norm of the Hessian is computed numerically using LOPPCG [33], [34] and analytically using 2 3. Datasets are orthogonalized with \(\sigma_{\mathbf{x}}=1\), and our model is a 4-layer Linear-MLP with width 64, initialized with balanced weights, and trained with \(\eta_{\max}=1e-2\).

7: Spectral norm of the Hessian is computed numerically using LOBPCG and analytically using 2 3.

• Left: Datasets are not orthogonalized, and our model is a 4-layer Linear-MLP with width 64, initialized with balanced weights.

• Middle: Datasets are orthogonalized with \(\sigma_{\mathbf{x}}=1\), and our model is a 4-layer Linear-MLP with width 64, initialized with Kaiming uniform initialization [29] (standard in PyTorch).

• Right: Datasets are orthogonalized with \(\sigma_{\mathbf{x}}=1\), and our model is a 4-layer Tanh-MLP with width 64, initialized with balanced weights.

All models are trained with \(\eta_{\max}=1e-2\).

8: Spectral norm of the Hessian, computed numerically using LOPPCG and analytically using 2 3. Datasets are orthogonalized with \(\sigma_{\mathbf{x}}=1\), and our model is a 4-layer Linear-MLP with width 64, trained with \(\eta_{\max}=1e-2\).

References↩︎