cs.LGAug 28, 2026
SaveFrom the Loss Landscape to Diverse Feature Learning in Neural Networks
Organizations: TOYOTA TECHNOLOGICAL INSTITUTE AT CHICAGO Chicago, Illinois
Abstract
Over the course of the last decade, neural networks have grown from an academic curiosity to moving the markets of nations. Despite this explosion in both research and deployment, relatively little is understood about how they achieve the solutions they do. This is both scientifically relevant, and pressing for society. When neural networks make decisions across self-driving, construction, law, hiring and health, there have been and will continue to be unintended consequences. However, attempting to generalize the failures of the largest and most important production systems makes for a very difficult task. Yet signs of these failures exist at all scales of neural networks, so we should be able to study a much more tractable setting. All neural networks must undergo an optimization process, called training, to be useful. To a great degree, understanding neural networks is understanding their optimization: through what process and exposure to which data did they arrive at their results. Yet our knowledge on this topic as a field is quite imprecise. In particular, a curious phenomenon called mode connectivity, the ability to connect neural networks in the loss surface, defies explanation entirely. This dissertation elucidates, explains and exploits this special structure in the loss landscape...
Figures & tables
Figure 1: The hat structure of the loss landscape of neural networks. We demonstrate this structure, provide a preliminary explanation for it, and leverage the explanation to improve practical neural networks.
Figure 2.1 : Schematic of different varieties of loss landscapes. Color denotes height, where lower is better. (a) : Convex loss. It is possible to find a path to the global minimum through local steps, when a ball is dropped in the landscape it will roll to the bottom. (b) : Nonconvex loss. Unlike the convex case, a ball may get stuck in a particular local minimum without finding the global minimum. Neural networks are generally nonconvex, making optimization more difficult than the convex case.
Figure 3.1 : The algorithmic procedure and graphic for measuring convex mode connectivity.
Figure 3.2 : Top row: The empirical probability that a sample inside the convex hull (Convex) or a random perturbation around the endpoint (Random) have lower loss than their closest endpoints vs. the split epoch ( in Algorithm 1 ) for the many different trajectories. After an initial phase of training, the probability of lower loss inside the hull quickly converges to one, the loss inside the hull is always better than at the endpoints. Bottom row: Parameter distances between pairs of points in the CMC procedure. The distance between pairs of endpoints is of the same order of magnitude as other distances for all tasks. In other words, the apparent scale of convexity is similar to that of the total trajectory. Error bars for probabilities are computed by empirical standard deviations and , hence are not meaningful when values stretch greater than 1 or less than 0.
Figure 3.3 : Singular values for the matrix of stacked endpoints of the mode. We examine endpoints split after the point when the hull center performs better than the endpoints (see Figure 3.2 ). For points in space, the shape they define can have at most dimensions. We see in all cases that the magnitudes do not decay to zero until component . In the case that there are outputs and (ResNet), we see a roughly uniform mode. In cases with (SIREN, LeNet, and LeNet-m), it does not appear that there is any special fragmenting at . In particular, the LeNet plot is over 300 endpoints, and the tail does not look as if it will decay to zero any time soon. The computational expense of these experiments limits how many endpoints we train.
Figure 3.4 : Top row: The ratio of smallest to largest eigenvalue of the Hessian throughout training. This ratio is almost always negative, but quite small, making any local non-convexity appear unimportant. Bottom row: Probability of a convex line of loss between a pair of updates. Notice that this probability is close to 1 for the entire optimization, typically before the LMC regime in Figure 3.2 .
Figure 3.5 : Top row: Average accuracies at the endpoints of training (Endpoint), inside the convex mode defined by the endpoints (Convex), and at random directions around endpoints (Random). Notice that after a late enough split, samples in the hull are better than the endpoints. Bottom row: Average training accuracy over runs (Accuracy) compared to common points predicted correctly within a mode (Intra mode) and common points predicted correctly between hulls (Inter mode). We see that for models with an imperfect training accuracy, intra mode agreement is substantially lower than accuracy, and inter mode agreement is even lower. This suggests that endpoints of a mode are not similar functions.
Figure 4.1 : 4.1(a) Schematic for the spectral dynamics of a weight matrix. As training proceeds, top singular vectors become stable and top singular values grow disproportionately large. 4.1(b) Singular value evolution for a single intermediate matrix parameter in a Transformer, where each line is a single singular value. Darker colors correspond to larger singular values while lighter ones correspond to smaller values. We see a disproportionate trend where large singular values grow larger faster. 4.1(c) Previous works use SV basis alignment between layers to prove similar theoretical results, however actual alignment between consecutive layers is weak, which we describe in Section 4.5 . We explore these spectral dynamics of weights and connect them to generalization, regularization, and seemingly unrelated phenomena like linear mode connectivity.
Figure 4.2 : Grokking and Spectral Dynamics. Top row: Grokking for Transformers on modular addition ( Nanda et al., 2023 ) for a single seed due to different convergence times. See Figure B.7 for averaging over multiple seeds. Bottom row: Grokking for a 12-layer MLP on MNIST ( Fan et al., 2024 ) . (a) Training and validation error. Shaded regions correspond to standard deviation, where large deviations are due to different seeds converging at different times. (b) A visualization of singular value evolution for the first attention parameter and the second MLP layer, where each line represents a single singular value and the color represents the rank. (c) Effective rank of all layers (matrix parameters) over time (Eqn. 4.1 ). (d) A visualization of the alignment (Eqn. 4.3 ) between the embedding and the first attention parameter, and the first and second MLP layers, where the -axis corresponds to index of the diagonal. We see that grokking co-occurs with a transition to low-rank weights. In addition, there is an alignment that begins early in training that evolves up the diagonal. In the image classification case, we see a similar rank transition, though alignment does not evolve up the diagonal (see Figure B.8 ).
Figure 4.3 : Top row: Singular value evolution for a single matrix in the middle of each model. Each line represents a singular value, darker colors are larger values, while lighter colors are smaller values. Notice the unequal evolution where top singular values grow at a disproportionate rate. Bottom row: Normalized effective rank (Eqn. 4.1 ) evolution visualized in color for all matrix parameters across architectures and time. As we move down the -axis, the depth of the parameters in the model increases, while the -axis tracks training time. The axis label "Layer" is shorthand for "matrix parameter". For example, in the Transformer we visualize each of the parameters for each block. Notice decreasing effective rank with training time across nearly all parameters, though the magnitude differs across layers (indicated by absolute color). The block-like patterns for VGG are likely due to different channel dimension sizes. The banding in the UNet, LSTM, and Transformer is due to the differences between convolutional and linear layers, residual block connections, and attention and fully connected layers, respectively. The sharp transitions through training in VGG are due to 10 learning rate decays.
Figure 4.4 : Top row: Training loss. Bottom row: Validation loss. Red is the full model. Blue is post-training pruning the bottom half of the SVD for every matrix in the model that is not the final layer. Green is post-training pruning the top half of the SVD. Notice that for all models, keeping the top half of the SVD is close to the full model performance, supporting the idea that the top directions provide a good approximation to the function. We specifically avoid pruning the final layer as for some tasks it is so low-rank that pruning further affects performance in an anomalous way ( Frankle et al., 2020 ) .
Figure 4.5 : Neighboring layer alignment of singular vectors. Top row: The diagonal of the alignment matrix (Eqn. 4.3 ) vs. training time for a single pair of matrices in the middle of each model. We see a small amount of alignment in the top ranks between layers shortly after training begins, but this becomes more diffuse over time. Bottom row: Alignment metric (Eqn. 4.4 ) for pairs of matrices for depth vs. training time. It is hard to make out a global trend across models, though the LSTM shows a weak signal around Epoch 1 when the initial alignment occurs, and the Transformer case has a banding pattern with depth due to alignment between the query and key matrices that have no nonlinearity in between.
Figure 4.6 : Singular value evolution for a single matrix and normalized effective rank (Eqn. 4.1 ) across matrices over time, where the rows use differing amounts of weight decay. From top to bottom, for VGG we use coefficients , while for other networks we use coefficients . Higher weight decay coefficients promote more aggressive rank minimization. VGG uses SGD with momentum, while the rest use AdamW ( Loshchilov and Hutter, 2017 ) , which may explain the earlier norm collapse. We provide an additional version with shared color bars in Figure B.10 .
Figure 4.7 : Top row: results with true labels. Bottom row: results with random labels. We see that the middle layers have a lower effective rank (Eqn. 4.1 ) when using true labels and that alignment (Eqn. 4.3 ) in the middle layers persists throughout training, unlike in the random label case. We emphasize this alignment occurs despite the nonlinearities.
Figure 4.8 : Top row: Barrier size (Eqn. B.1 ) vs. split step. Middle row: singular vector agreement for a single matrix parameter between branch endpoints that share a common trunk. Bottom row: summary statistic for singular vector agreement across layers vs. split step. We see that as models exhibit LMC, they also share a small number of top singular vectors. We only show agreement for 100 dimensions as the rest are near zero.
Figure 4.9 : Dynamics with random labels for VGG. Top row: results with true labels. Bottom row: results with random labels. We see that the middle layers have a lower effective rank (Eqn. 4.1 ) when using true labels and that alignment (Eqn. 4.3 ) in the middle layers persists throughout training. The results are less stark in the VGG case, but similar to the MLP.
Figure 4.10 : Dynamics with random labels for LSTM. Top row: results with true labels. Bottom row: results with random labels. We see that the middle layers have a lower effective rank (Eqn. 4.1 ) when using true labels and that alignment (Eqn. 4.3 ) in the middle layers persists throughout training. Though the LSTM does not fit the random labels perfectly, the results are qualitatively similar to the other cases.
Figure 4.11 : Top row: Singular vector agreement for a single matrix in the middle of each model (diagonal of Eqn. 4.5 ). Notice top singular vectors become stable in direction earlier. Bottom row: Summary score for each matrix across architectures. As we move down the -axis, the depth of the parameters in the model increases, while the -axis tracks training time. The sharp transition midway through training in the VGG case is likely due to a 10x learning rate decay.
Figure 4.12 : Pruning results for VGG. Top row: Magnitude pruning. Bottom row: random pruning. (a): Training loss. We see that at 5% sparsity magnitude pruning is significantly better than random pruning of the same layerwise sparsity. (b): Singular vector alignment pre- and post-pruning at the end of training for a single layer (the 3rd convolution). We see that magnitude pruning approximates many more top singular vectors than random pruning at the same sparsity. (c): Singular vector alignment score pre- and post-pruning across all layers. Agreement is higher across all layers for magnitude pruning, though curiously deeper layers do not agree, likely as later layers are wider so weights are lower magnitude and more will be pruned by the unstructured process. (d): Singular vector alignment between the pruned and unpruned models along the training trajectory. We see that the magnitude pruning still has similar dynamics in its top singular vectors, while random pruning does not. (e): Singular vector alignment score between pruned and unpruned models across layers and time. Again evolution is similar for early layers with magnitude pruning, and completely different for random pruning.
Figure 4.13 : Pruning results for LSTM. Top row: Magnitude pruning. Bottom row: random pruning. See Figure 4.12 for details. Results are quite similar for the LSTM at 25% pruning as the VGG in Figure 4.12 .
Figure 4.14 : Top row: Barrier size vs. split step. Middle row: singular vector agreement for a single matrix parameter between branch endpoints that do not share a common trunk, but do share split time and branch data order. Bottom row: summary statistic for singular vector agreement across layers. We see that when branches do not share a common trunk, there is neither LMC nor singular vector agreement, even though the optimization is otherwise the same.
Figure 4.15 : Top row: Barrier size vs. perturbation magnitude. Middle row: singular vector agreement for a single matrix parameter between branch endpoints vs. perturbation magnitude. Bottom row: summary statistic for singular vector agreement across layers with perturbation magnitude. We see that whereas without perturbation models would exhibit LMC after training, with increasing perturbations the LMC property disappears simultaneously with the agreement in top singular vectors.
Figure 5.1 : Examples from the MNIST-CIFAR dataset in Shah et al. (2020) .
Figure 5.2 : Evaluation protocol for experiments. After training on MNIST-CIFAR, we swap the classes of either the MNIST half, or the CIFAR half of the image and evaluate on that setting. When models rely entirely on MNIST, error on the swapped task degrades to 100%, and similarly for CIFAR. On the right we plot the errors over successive rounds of PGD. We see that models start to learn CIFAR after the early rounds of PGD, coming to rely quite heavily on CIFAR in later rounds. Spikes correspond to restarting a new round of training.
Figure 5.3 : Left: Training and validation error for successive rounds with our proposed PGD method on top singular vectors. Right: Swapped-class errors throughout rounds. We see that as the round number increases, validation error decreases due to the complexity of the task increasing. In addition, the model increasingly relies more on CIFAR, as the swapped CIFAR error goes from near 0 to 60% and above. This indicates the model is learning to extract more complex signals which are less predictive but still useful.
Figure 5.4 : First layer singular vectors over rounds of PGD training. Each set of columns corresponds to the MNIST and CIFAR halves of the reshaped singular vector. Descending down the rows corresponds to increasing rank index when sorted by singular value. We see initially that the first layer singular vectors form digit classifiers exclusively. In middle rounds the network focuses on higher and higher frequency signals in MNIST. In late rounds, all but the high frequency signals are present in MNIST, hence the model starts to leverage the CIFAR half of the image.
Figure 5.5 : Left: Training and validation error for successive rounds with PGD where we flatten the gradient and singular vectors before projection. Right: Swapped-class errors throughout rounds. We see that no new features are learned when ignoring the matrix structure of the parameters, so the spectral design is critical.
Figure 5.6 : Left: Training and validation error for successive rounds with Muon. Right: Swapped-class errors throughout rounds. Even though Muon is incentivized to give relatively more gradient to less prominent features (through higher-rank upweighting), it is insufficient for picking up CIFAR.
| Model | Accuracy |
|---|---|
| Early | |
| Late |
Table 5.1 : Probing accuracy for CIFAR on middle layers for an ensemble of 5 models drawn from early or late rounds of PGD. We see that early rounds are substantially worse at extracting CIFAR information from the middle of the model, reinforcing the argument that PGD forces the model to learn more.
| Model | Accuracy |
|---|---|
| 4x wide | |
| CAT | |
| PGD |
Table 5.2 : OOD linear probing accuracy between different methods. We observe the same effect as Zhang and Bottou (2023) , that random ensembles (CAT) are much stronger than a large model that is parameter advantaged, but PGD ensembles do not appear to improve over random.
Figure 5.7 : Linear CKA ( Kornblith et al., 2019 ) between layers for random ensembles (CAT) and PGD ensembles. The values are quite similar, suggesting that PGD is managing to find similar solutions as simple random reruns, possibly permutations of the original weights. Hence, PGD does not capture every notion of diversity.
| Model | Accuracy |
|---|---|
| CAT | |
| CKA |
Table 5.3 : Top-1 OOD accuracy for random ensembles (CAT) v.s. CKA ensembles. We see that CKA improves by a substantial margin over CAT in the linear probing setup.
| Model | Cross Entropy |
|---|---|
| CAT | |
| CKA | |
| CKA max-in-depth | |
| CKA token pool | |
| CKA token cloud | |
| CKA random proj. |
Table 5.4 : OOD cross entropy for random ensembles (CAT) v.s. CKA ensembles under many different variations. Unlike the image classification case, we are unable to find a version of CKA that outperforms random ensembles.
Figure 5.8 : CKA for representation vectors and gradient vectors for the task loss alone at different layers for image classification. Left: We see that representation and gradient CKA are reasonably high for CAT. Right: The same, but for the CKA-trained models. We see that using CKA greatly affects the representation CKA as it optimizes for this, and it also affects the task gradient CKA.
Figure 5.9 : CKA between the gradient vectors for the task loss alone at different layers for language modeling. Left: We see CKA between representations and between gradients, which are both high for the CAT models. Right: The same, but for the CKA-trained models. We see that while representation CKA is low, gradient CKA is unaffected by our CKA objective, unlike image classification.
| Model | Cross Entropy |
|---|---|
| Random | |
| NCL JSD | |
| NCL MSE | |
| NCL correct only |
Table 5.5 : OOD cross entropy for random ensembles v.s. NCL ensembles We see no significant difference to random ensembles. NCL JSD corresponds to using Jensen-Shannon divergence for difference. MSE corresponds to MSE for the probability distribution, and Top-1 corresponds to only considering the difference between probabilities of the correct class.
| Model | Cross Entropy |
|---|---|
| Random | |
| Focal | |
| Focal top-p |
Table 5.6 : OOD cross entropy for random ensembles v.s. Focal scaling ensembles We see no significant difference to random ensembles. Top-p corresponds to scaling the top-p of the tokens by sequence by a constant additive factor.
| Model | Cross Entropy |
|---|---|
| Random | |
| Entropy | |
| Entropy top-p |
Table 5.7 : OOD cross entropy for random ensembles v.s. entropy scaling ensembles We see no significant difference to random ensembles. Top-p corresponds to scaling the top-p of the tokens by sequence by a constant additive factor.
| Model | Cross Entropy |
|---|---|
| Even | |
| Entropy | |
| KL disagreement |
Table 5.8 : OOD cross entropy for different ensemble weighting methods on the Entropy top-p runs. We use these as the base as we expect there may be nonuniformity compared to the Random ensembles.
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
| Abbrev. | Task | Dataset | Model | Optimizer | LR | Batch size |
|---|---|---|---|---|---|---|
| SIREN | Regress. | Bach Waveform | 5-layer SIREN | Adam | 1e-5 | 8192 |
| LeNet | Classif. | CIFAR-10 | LeNet-5 | SGD | 1e-2 | 128 |
| LeNet-m | Classif. | CIFAR-10 | LeNet-5 | SGD+mom. (0.9) | 1e-2 | 128 |
| ResNet | Classif. | CIFAR-10 | ResNet-20 | SGD+mom. (0.9) | 0.1 | 128 |
Table A.1 : Summary of tasks studied in this chapter.
Figure A.1 : Training loss curves for all tasks. When cross-referencing with Figure 3.2 , we see that the first split point that leads to a convex hull is quite early, long before the loss has converged, and often before it has even halved.
Figure B.1 : Spectral dynamics of Pythia suite. From top to bottom we examine the 160m, 410m, 1.4b and 2.8b parameter models. Notably, much less noise appears in the alignment plot with increasing scale. Presumably this could be due to the fact that larger dimensional vectors have higher probability to be orthogonal, which may play a role in making optimization easier. We see stronger alignment score (Eqn. 4.4 ) in all layers in the larger model, perhaps because of that cleaner signal.
Figure B.2 : Diagonal of alignment for a single pair over time (Eqn. 4.3 ) and alignment metric across pairs of matrices over time (Eqn. 4.4 ) where the y-axis represents depth. From top to bottom, for VGG we use coefficients , while for other networks we use coefficients . We see that the maximum alignment magnitude is higher with large weight decay, and in particular, the Transformer has the strongest alignment even when nonlinearities separate the MLP layers.
Figure B.3 : Training loss over time, where the rows use differing amounts of weight decay. From top to bottom, for VGG we use coefficients , while for other networks we use coefficients . We see that it is still possible to achieve low training loss under high weight decay, and as we increase the amount of weight decay, the gap between pruned and unpruned parameters closes, lending support to the idea that the parameters become lower rank.
Figure B.4 : Grokking and Spectral Dynamics in Modular addition. Top row: 30% data and no weight decay. 2nd row: 30% data and weight decay 1.0 (grokking), using hyperparameters from Nanda et al. (2023) . 3rd row: 70% data with no weight decay (slingshot), using hyperparameters from Thilak et al. (2022) . Bottom row: 90% data and no weight decay. (a): Training and validation error. (b): Singular value evolution is visualized for the first attention parameter, where each line represents a single singular value and the color represents the rank. (c): Effective rank of all layers (Eqn. 4.1 ). (d): Alignment (Eqn. 4.3 ) between the embedding and the first attention parameter is also visualized, where the y-axis corresponds to index of the diagonal. One can see that grokking co-occurs with low-rank weights. In addition, there is an alignment that begins early in training that evolves up the diagonal. Without weight decay and with less data, neither grokking nor the other phenomena occur during the entire training budget, but using more data, even without weight decay, leads to low-rank solutions from the beginning of training. The slingshot case follows a similar trend, though the validation loss is gradually fit. Across cases with good generalization, parameters are lower rank, and alignment is also more prevalent in the top ranks.
Figure B.5 : Varying initialization scale affects the dynamics of an MLP trained on CIFAR10. From top to bottom we use default PyTorch initialization multiplied by the constants . Notably, alignment occurs slightly more strongly with smaller initialization, as predicted by Woodworth et al. (2020) for deep linear models.
Figure B.6 : The effect of varying learning rate on an MLP trained on CIFAR10. From top to bottom we use learning rates . Notably, alignment increases with larger learning rate, agreeing with the deep linear setting in Ghosh et al. (2025) .
Figure B.7 : Plots for averages of 3 seeds of grokking modular addition with Transformers (Figure 4.2 ). We see a stark transition to low-rank when validation error decreases.
Figure B.8 : Alignment matrix (Eqn. 4.3 ) over time. We see that the majority of signal is concentrated in the diagonal, an example of why we consider the diagonal in the rest of the work.
Figure B.9 : Different architectures all trained on CIFAR10. From top to bottom, MLP, LeNet-5 and VGG-16. We see roughly consistent spectral dynamics across architectures within the same task.
Figure B.10 : Duplicate of Figure 4.6 with shared color scaling between settings.