PaperlayerView exact revision on arXiv

Source-linked AI summary

From the Loss Landscape to Diverse Feature Learning in Neural Networks

David Aram Yunis

arXiv:2608.28948v1cs.LGcs.AI

TL;DR

The paper asks why mode connectivity occurs and how neural-network optimization produces its solutions. It analyzes spectral dynamics across architectures and tasks, finding that low-rank dynamics connect mode connectivity and other phenomena while motivating a method for addressing spurious correlations.

  • Problem

    The causes of mode connectivity and the broader optimization processes underlying neural-network solutions remain poorly understood, despite their scientific and societal importance.

    Source: Introduction · p. 19, Background · p. 24

  • Method

    The paper empirically studies singular-value and singular-vector dynamics across architectures and tasks, then uses projected gradient descent to address spurious correlations.

    Source: §4.1 · p. 33, §5.1 · p. 57

  • Results

    The study finds that low-rank dynamics link diverse deep-learning phenomena, including mode connectivity, weight decay, pruning, and feature learning.

    Source: §4.1 · p. 33, §4.10 · p. 53

  • Takeaways & Limitations

    Spectral structure provides a unified lens on optimization and motivates manipulating low-rank bias to address simplicity bias and spurious correlations.

    Source: §4.10 · p. 55, §5.1 · p. 57

    The paper cannot detect a convexity change around the split point that gives rise to mode connectivity, possibly because its metrics sample the landscape sparsely.

    Source: §3.5 · p. 30

Abstract

from arXiv · show

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

Introduction

The introduction argues that neural networks’ expanding societal role makes their optimization and generalization behavior urgent to understand, since current remedies such as scaling face practical and scientific limits. The dissertation develops an account of loss-landscape structure and feature learning that aims to address distribution shifts, low-data regimes, and simplicity bias.

Introduction

  • Neural networks increasingly govern consequential domains, but their poorly understood optimization and generalization create risks of uninterpretable errors and social repercussions.They can memorize training data while still lacking a complete explanation for generalization, making understanding important for deployment and algorithm design.Source: Introduction · pp. 19–20
  • Scaling data and computation can improve performance, but it requires extreme resources, struggles under distribution shifts, and does not resolve the underlying scientific gap.The introduction presents simplicity bias as an example where scaling-based remedies encounter practical walls and leave understanding incomplete.Source: Introduction · pp. 19–20
  • Existing theory and empirical studies remain fragmented by restrictive settings, difficult-to-scale methodologies, or abstract representations, leaving contemporary neural-network phenomena poorly understood.The dissertation therefore seeks a broad understanding applicable to practical systems rather than only special theoretical regimes or small-scale analyses.Source: Introduction · p. 20

1.1 Contributions

  • It uses spectral dynamics to improve training, overcoming simplicity bias with ensembles and targeting low-data regimes before concluding with broader implications.The contribution sequence begins with background on Linear Mode Connectivity and ends by discussing the results’ implications.Source: §1.1 · pp. 20–21
  • The dissertation extends linear mode connectivity by identifying high-dimensional, diverse convex modes in neural-network loss landscapes and explains them through stable low-rank spectral dynamics.It further argues that spectral dynamics unify phenomena involving generalization and regularization.Source: §1.1 · p. 21

Background

Neural-network optimization is difficult because many minima exist but only some generalize well, while nonconvex landscapes can trap optimization. Linear mode connectivity reveals unexpectedly structured, broadly occurring low-loss paths whose mechanism and practical value motivate the dissertation.

Background

  • Many neural-network minima exist, but only some generalize well, raising questions about how optimization selects solutions and how better ones might be found.Neural networks can approximate arbitrary functions, yet their optimization is theoretically difficult (Choromanska et al., 2015; Kawaguchi, 2016; Keskar et al., 2016).Source: Background · p. 22
  • Unlike convex landscapes, neural-network nonconvexity permits disconnected optima and suboptimal points where optimization can become stuck.Convex landscapes allow paths to optima through local steps, whereas neural networks generally exhibit nonconvex landscapes with multiple disconnected craters (Vandenberghe, 2004; Sontag and Sussmann, 1989; Li et al., 2018; Frankle et al., 2020).Source: Background · p. 22, Figure 2.1 · p. 23
  • Mode connectivity shows that minima from the same initialization can be joined by simple nonlinear low-loss paths despite apparent nonconvexity.Piecewise-continuous paths with few segments and even piecewise-linear paths are possible under sufficient capacity (Draxler et al., 2018; Garipov et al., 2018; Kuditipudi et al., 2019). Parameter permutations, rescaling, and hyperparameters can produce distinct minima (Dinh et al., 2017; Keskar et al., 2016).Source: Background · p. 23
  • Linear mode connectivity connects same-initialization optima by a low-loss linear path, implying optimization remains in a structured convex subspace rather than exploring the entire landscape (Nagarajan and Kolter, 2019).The phenomenon was first observed at small scale, then confirmed for broader networks using a protocol with shared optimization before splitting trajectories (Frankle et al., 2020).Source: Background · pp. 23–24
  • LMC is nontrivial rather than merely local: the distances among initialization, split point, and optima are of comparable magnitude.LMC also appears with pretrained initialization and in Transformer language models, extending beyond the original image-classification setting (Juneja et al., 2022).Source: Background · p. 24
  • LMC supports model averaging for improved performance, out-of-distribution generalization, and combining tasks, while the dissertation investigates its low-rank-dynamics explanation and feature-learning applications.These applications include finding an optimal point between solutions, improving out-of-distribution generalization, and combining tasks into one model (Li et al., 2022; Ilharco et al., 2022).Source: Background · p. 24

On Convexity and Linear Mode Connectivity in Neural Networks

The section shows that early-split SGD trajectories form high-dimensional convex hulls of low loss, establishing a broad form of linear mode connectivity across neural-network settings. However, simple convexity measures do not identify its onset, and interpolation connects functionally diverse endpoints rather than similar predictors.

3.1 Introduction

  • Early-split SGD trajectories define high-dimensional convex hulls of low loss, showing that linear mode connectivity reflects convexity around an entire mode rather than merely between two endpoints.This behavior was confirmed across a SIREN regression task and several CIFAR-10 classifiers.Source: Yunis et al. (2022) · p. 25, §3.1 · p. 26, §3.2 · p. 27

3.3 Measures of Convexity

  • Hessian spectra and between-update loss curves remain locally near-convex throughout training, showing no distinguishing transition at the point where linear mode connectivity begins.The results instead suggest that convergence to different convex hulls occurs gradually rather than at a sharply defined training point.Source: §3.1 · p. 26, §3.3 · p. 28, §3.5 · p. 29

3.2 Convex Mode Connectivity

  • Convex combinations become lower-loss than endpoint averages after an early training phase, while random perturbations are almost always worse, supporting a convex region in the loss surface.The section calls the post-threshold period the LMC regime and measures it with convex mode connectivity.Source: §3.2 · p. 26, Figure 3.2 · p. 28

3.4 Diversity of Endpoints

  • Only a fraction of correct predictions are shared within a mode, and fewer are shared across hulls, even though interpolation improves accuracy inside the convex hull.This indicates that loss improvement is not driven by a common subset of predictions, while poor agreement across hulls weakens permutations as a sufficient explanation.Source: §3.1 · p. 26, §3.4 · p. 29, Figure 3.5 · p. 31

3.5 Discussion

  • The combined evidence is dichotomous: large low-loss convex hulls appear after early training, yet simple proxies detect no increase in convexity before or after their formation.The section leaves the source of this convexity unresolved despite its occurrence across settings.Source: §3.5 · pp. 29–30
  • The section hypothesizes that low-rank dynamics stabilize top singular vectors after trajectory splitting, producing endpoint models that can be averaged despite functional dissimilarity.This hypothesis is motivated by deep linear-model dynamics in which larger singular values grow faster under certain technical assumptions.Source: §3.5 · p. 31

Approaching Deep Learning through the Spectral Dynamics of Weights

This section examines spectral dynamics in neural-network training, connecting singular-value and vector evolution to rank, alignment, pruning, memorization, generalization, grokking, and mode connectivity across architectures and tasks. It develops low-rank dynamics as a framework for understanding primary feature learning and simplicity bias, while noting limits in their theoretical interpretation and proposing spectral interventions to address spurious correlations.

4.5.3 Alignment of Singular Vectors Between Layers

  • The observed rank-minimization dynamics disagree with theoretical assumptions that consecutive layers maintain aligned singular-vector bases under balanced initialization.In larger nonlinear networks, initial alignment is absent, any top-rank signal is weak and transient, and no global alignment trend is apparent.Source: §4.1 · p. 34, §4.2.1 · p. 35, §4.5.3 · p. 41

4.1 Introduction

  • Weight decay promotes rank minimization across architectures and tasks beyond its explicit role as a norm penalty, and in some cases also promotes inter-layer singular-vector alignment.These effects suggest that spectral dynamics may explain why weight decay benefits practical generalization more fully than norm-based accounts alone.Source: §4.1 · pp. 34–35
  • Random-label training produces high-rank, unaligned parameters, whereas true-label training produces lower-rank parameters with greater consecutive-layer alignment.The spectral distinction between memorizing and generalizing networks provides a route toward understanding their different solutions.Source: §4.1 · pp. 34–35
  • Global magnitude pruning preserves top singular vectors and behaves like low-rank pruning, while linear interpolation between models correlates strongly with shared top singular vectors.These observations place lottery tickets and linear mode connectivity under a common spectral-dynamics perspective.Source: §4.1 · pp. 34–35

4.2.2 Low-Rank Properties

  • The section extends prior low-rank and weight-decay results beyond idealized settings by examining larger practical networks and a broad suite of experiments.Its central proposal is that spectral dynamics provide a unified view of previously isolated phenomena in deep learning.Source: §4.1 · p. 35, §4.2.2 · p. 36, §4.5 · p. 39

4.3 Common Quantities

  • The analysis measures singular-value evolution, normalized effective rank, and neighboring-layer singular-vector alignment to compare spectral dynamics across layers, architectures, and training time.The study focuses on weight matrices, excluding one-dimensional bias and normalization parameters, and uses literature-based hyperparameters across representative classification, generation, speech, and language-modeling tasks.Source: §4.3 · pp. 36–37, §4.5.1 · p. 39

4.4 Grokking and Rank Minimization

  • Grokking’s sudden validation-loss drop coincides with low-rank behavior emerging across weight matrices in both modular arithmetic Transformers and MNIST MLPs.Without sufficient weight decay, low-rank solutions do not develop in the reported experiments, linking further training after memorization to rank regularization.Source: §4.1 · pp. 33–34, Figure 4.2 · p. 37, §4.4 · p. 38

4.5 Spectral Dynamics Across Tasks

  • Across ConvNets, UNets, LSTMs, and Transformers, effective rank generally decreases during training as larger singular values grow faster, progressively simplifying network parameters.Top singular directions also become stable earlier than the end of training, and retaining the top half of singular values closely approximates full-model performance.Source: §4.1 · p. 33, §4.5.2 · pp. 40–41

4.6 The Effect of Weight Decay

  • Moderate weight decay promotes low-rank solutions and can improve generalization, whereas excessive decay causes norm collapse and performance failure when rank becomes too low.Lower rank generally corresponds to better generalization and a smaller training-validation gap, but the minimum usable rank depends on the task.Source: §4.6 · pp. 43–44, Figure 4.6 · p. 43

4.7 Additional Connections

  • True-label training produces lower-rank inner layers and persistent middle-layer alignment, while random-label training yields higher-rank solutions and alignment that eventually disappears.This pattern appears qualitatively across MLP, VGG, and LSTM experiments, though alignment is weaker in VGG and LSTM.Source: §4.7 · p. 44, Figure 4.7 · p. 44, Figure 4.9 · p. 46

Linear Mode Connectivity Correlates with Top Singular Vector Agreement

4.9.1 Top Singular Vectors Become Stable Earlier

  • Top singular vectors stabilize earlier than bottom vectors, providing a spectral basis for why later training preserves task-relevant directions.The result is shown by tracking singular-vector agreement over training and supports the subsequent pruning and connectivity analyses.Source: §4.9.1 · p. 48, Figure 4.11 · p. 48

4.9.2 Lottery Tickets Preserve Final Top Singular Vectors

  • Magnitude-pruning masks preserve top singular-vector dynamics and approximate low-rank structure, whereas random masks preserve them poorly and can prevent further learning after early rewinding.At 5% VGG sparsity, magnitude pruning has significantly lower training loss than equally sparse random pruning; analogous LSTM results are reported at 25% pruning.Source: §4.9.2 · p. 49, Figure 4.12 · p. 50, Figure 4.13 · p. 51

Perturbing Breaks Linear Mode Connectivity and Singular Vector Agreement Simultaneously

  • Increasing perturbations to emerging top components simultaneously increase interpolation barriers, eliminate linear mode connectivity, and destroy singular-vector agreement.The intervention supports a causal connection between top-direction agreement and LMC.Source: Perturbing Breaks Linear Mode… · p. 53, Figure 4.15 · p. 54

4.10 Discussion

  • Spectral dynamics connect low-rank bias, weight-decay regularization, label-dependent solutions, lottery-ticket pruning, and linear mode connectivity through singular-value and singular-vector behavior.The observations remain consistent across diverse settings without balanced initialization, linearity, or small weight-scale assumptions, although comprehensive theory is still missing.Source: §4.10 · p. 53
  • The framework’s main limitation is that linear alignment between neighboring layers is difficult to interpret after nonlinearities, motivating post-activation analysis in larger systems.Existing theoretical models also fail to capture the observed training complexities, including alignment behavior beyond balanced or near-zero initialization.Source: §4.5.3 · p. 42, §4.10 · p. 53
  • Studying low-rank models may support better optimizers, deployment-risk diagnosis, safety, and interpretability by reducing the model complexity under examination.The discussion frames the empirical exploration as supporting scientific understanding, theoretical assumptions, and engineering improvements.Source: §4.10 · p. 54
  • Low-rank dynamics are tightly related to convex-mode development and connect to prior lottery-ticket work (Frankle et al., 2020).The discussion presents these dynamics as central to neural-network optimization.Source: §4.10 · p. 54
  • Low-rank dynamics can be understood as the primary feature learning occurring during optimization.This perspective aligns with low-rank pruning for compression and performance and findings that different low-rank cores learn different functions (Lubana et al., 2023).Source: §4.10 · p. 55
  • Defining simplicity bias through low-rank bias suggests manipulating spectral dynamics to overcome spurious correlations.The discussion identifies this as the key idea developed in the following chapter.Source: §4.10 · p. 55

Overcoming Spurious Correlations through Diverse Feature Learning

This section develops generic approaches for overcoming simplicity bias and learning diverse features, evaluating projected gradient descent, random ensembles, and CKA-based diversity methods across controlled spurious-correlation and language-modeling settings. It examines their effects on out-of-distribution transfer while testing multiple CKA variants and low-dimensional projections.

5.1 Introduction

  • The section identifies simplicity bias as a generic obstacle: optimization favors minimal feature sets, while existing remedies require misclassification reweighting, test-distribution information, dataset assumptions, or task-specific augmentations.The proposed methods aim to avoid these requirements and extend diverse feature learning beyond classification toward language modeling.Source: Overcoming Spurious Correlations… · p. 56, §5.1 · p. 57

Overcoming Spurious Correlations through Diverse Feature Learning

  • The proposed approach iteratively projects gradients away from top singular-vector directions, forcing successive models to learn more complex and diverse features from the same data.The method reinitializes and retrains models while accumulating orthogonal directions to avoid across rounds.Source: §5.1 · p. 57, §5.2.1 · p. 59

5.2 Simplicity Bias and Spectral Dynamics

  • The MNIST-CIFAR testbed demonstrates simplicity bias because both MNIST and CIFAR are perfectly label-correlated, yet standard training uses only MNIST despite scale, regularization, adversarial training, and ensembling.The dataset concatenates MNIST and CIFAR images, with class labels determined by matched digit-object pairs.Source: §5.2 · p. 58

5.2.2 Experiments

  • On MNIST-CIFAR, projected gradient descent preserves near-zero training error while shifting reliance from MNIST to CIFAR, with swapped-CIFAR errors eventually exceeding 70%.Early models rely almost entirely on MNIST; later rounds reduce swapped-MNIST error below random chance and increase sensitivity to CIFAR.Source: §5.2.2 · pp. 59–60
  • Singular-vector visualizations show PGD progressing from MNIST digit detectors to higher-frequency MNIST features and then CIFAR-sensitive features.These results indicate that blocking simple directions can force a simple MLP to overcome its simplicity bias.Source: §5.2.2 · p. 60, Figure 5.4 · p. 62

5.3.1 A Problem with Projected Gradient Descent

  • In larger systems, naive PGD fails because orthogonality can preserve semantic content, motivating CKA minimization as a permutation-invariant objective for functional diversity.PGD and random ensembles have similar OOD performance and CKA similarity, suggesting PGD may rediscover permutations of the same features.Source: §5.3.1 · pp. 63–64

5.3.2 Inter-Example Similarity as a Measure of Functional Diversity

  • CKA-based training samples a previous ensemble member, penalizes representation similarity at selected middle layers, and adds the resulting term to the task loss.The method samples prior models because comparing against all of them is prohibitive in disk space or memory.Source: §5.3.2 · pp. 64–65

Image Classification

Language Modeling

  • The experiment compares a Transformer three times wider with an ensemble of six standard-width models at similar compute, evaluating validation loss in-distribution and Simple English Wikipedia loss out-of-distribution.The related dataset has high lexical overlap but substantially different topic distribution, targeting general language structure without linear probing.Source: Language Modeling · p. 66
  • Random ensembles achieve better out-of-distribution performance than the larger model, despite slightly worse in-distribution performance, while no CKA variant outperforms them.The comparison reports best results after hyperparameter tuning.Source: Language Modeling · p. 66, Table 5.4 · p. 66
  • CKA over maximum pairwise similarity across layers remains insufficient to improve over random ensembles.This variant was introduced to address the possibility that CKA was applied to the wrong transformer modules or encouraged repeated features across residual layers.Source: Language Modeling · pp. 66–67
  • CKA applied to pooled sequence representations or token-cloud Gram matrices also fails to beat random ensembles.These variants address possible high-dimensional objectives and token-level information permutation as trivial solutions.Source: Language Modeling · p. 67
  • Random matrix projections that optimize CKA in changing low-dimensional spaces likewise remain insufficient, leaving CKA’s language-modeling failure unexplained.Changing the projection each time is intended to prevent the model from exploiting a fixed coordinate subset.Source: Language Modeling · p. 67

A Problem with Optimizing CKA … Convex Mode Connectivity

The section argues that representation-level diversity objectives can be pathological because they may change internal structure without changing model functions, motivating output-space diversity. Across language-model experiments, proposed diversity, boosting, entropy, and weighting methods generally fail to improve over random ensembles, while the thesis identifies task difficulty and structured data subsets as central unresolved challenges.

A Problem with Optimizing CKA

  • High representation CKA can coexist with similar outputs in language models, showing that CKA may be fooled without producing meaningful functional diversity.Language-model CKA-trained models reduce representation CKA while gradient CKA remains high; the core limitation is that CKA measures implementation structure rather than actual functional change.Source: A Problem with Optimizing CKA · pp. 67–68, Figure 5.9 · p. 68

5.4 Diverse Feature Learning in Output Space

  • Output-space diversity defines distinct ensemble members by requiring different predictions, motivating negative correlation learning and subsequent sequential diversity methods.Negative correlation learning combines ensemble task loss with a diversity term, maximizing disagreement between the ensemble and individual members using JSD or MSE variants.Source: §5.4 · p. 68, §5.4.1 · pp. 68–69

5.4.1 Negative Correlation Learning

  • NCL ensembles show no significant improvement over random ensembles, including a top-1 variant intended to avoid pathological high-dimensional probability-distribution solutions.The experiments train models sequentially because joint optimization can collapse, but straightforward NCL variants remain insufficient.Source: Table 5.5 · p. 69, §5.4.1 · p. 69

5.4.2 Boosting Incorrect Predictions

  • Focal-style boosting of previous models’ errors preserves the original task loss but fails to outperform random ensembles, even with top-p variants and hyperparameter search.Later models upweight errors from sampled previous models using normalized weights, while retaining some weight on the original loss to avoid shrinking-sample overfitting.Source: §5.4.2 · p. 70, Table 5.6 · p. 70

5.4.3 Boosting Low-Entropy Predictions

  • Entropy-based boosting likewise cannot improve upon random ensembles, likely because forcing difficult or low-confidence predictions can sacrifice performance elsewhere.Both continuous and threshold entropy-scaling variants show no significant difference from random ensembles.Source: Table 5.7 · p. 71, §5.4.3 · p. 71

5.4.4 Ensemble Weighting

  • Heuristic per-token ensemble weighting using entropy or KL relevance also fails to beat uniform random ensembling, because disagreement may reflect either correct diversity or incorrect predictions.The authors suggest structured subsets or holdout-based global weighting as possible ways forward, while noting that difficult examples may induce solutions similar to training on the full dataset.Source: §5.4.4 · pp. 71–72, Table 5.8 · p. 71

Conclusion

  • The thesis frames diverse feature learning as an alternative to scaling: spectral dynamics explain simplicity bias and enable task-agnostic gains in toy settings, but finding structured subsets that avoid spurious correlations remains unresolved.Related work spans downstream-information methods, pretrained features, architectural modifications, objective-based approaches, and ensembling; the thesis instead targets learning dynamics from scratch and seeks to do more with less.Source: §5.6 · p. 74, Conclusion · p. 75

5.6 Discussion

  • Larger image networks beat random ensembles, but language-model diversity losses do not help, revealing a task-dependent limit to output-space ensemble diversification.The discussion reports success over random ensembles for larger image classification networks, but failure of prior diversity losses for language models after extensive ablations.Source: §5.6 · p. 74

A.1 Task Details

The chapter studies mode connectivity across a nonlinear SIREN audio-reconstruction task and LeNet-5 and ResNet-20 image-classification tasks, using constant learning rates and convergence-based stopping criteria.

Task Details

  • SIREN uses sine activations to reconstruct an audio waveform from its timestep, providing a non-piecewise-linear setting beyond prior ReLU experiments.The task has no generalization component, so training stops when training loss converges.Source: §A.1 · p. 92
  • LeNet-5 and ResNet-20 are trained on CIFAR-10 as simple systems without mode connectivity at initialization.These models extend the study beyond SIREN and use tuned hyperparameters with training stopped when validation error converges.Source: §A.1 · p. 92
  • All experiments use constant learning rates, with task-specific stopping based on converged training loss or validation error.Hyperparameter settings are summarized in Table A.1, and full training loss curves appear in Figure A.1.Source: §A.1 · p. 92, Table A.1 · p. 92

A.2 Metric Calculation Details … B.1 Explanation of Balancedness

The section defines metrics for Hessian curvature and update-induced non-convexity, then explains balanced initialization through matching singular vectors in consecutive matrices. These constructions connect measurable training geometry with assumptions used in deep matrix-factorization analyses.

A.2 Metric Calculation Details

  • Hessian eigenvalues are calculated with power iteration using Hessian-vector products, following Li et al. (2018).Source: §A.2 · p. 92
  • The method measures local convexity between two stochastic parameter updates, ΔW1 and ΔW2, from parameters W.Source: §A.2 · p. 92
  • Training curves show that the first split producing a convex hull often occurs early, before convergence and sometimes before the loss has halved.Source: Figure A.1 · p. 93
  • The probability that the update-induced gap is negative provides a non-convexity metric evaluated at multiple training points.Source: §A.2 · p. 93

B.1 Explanation of Balancedness

  • Prior work proposes rank minimization as an implicit-regularization description for deep matrix factorization, but relies on balanced initialization (Arora et al., 2019; Milanesi et al., 2021).Arora et al. (2018) (Appendix A) provides a detailed argument for this setting.Source: §B.1 · p. 94
  • Balanced initialization yields matching singular vectors for consecutive matrices, with equality up to signs or block-diagonal permutations when repeated diagonal values occur.If all diagonal elements are distinct and the matrices are square, the singular-vector factors match up to signs.Source: §B.1 · p. 94

B.2 Experimental Details … B.2.4 Language Modeling with Transformers

The experiments span classification, generation, speech recognition, and language modeling, using three-seed averages and architecture-specific SVD and alignment procedures. Each model uses task-specific datasets, optimization settings, and parameter representations for measuring singular values and inter-layer alignment.

B.2 Experimental Details

  • All experiments average plots over 3 random seeds, while alignment compares only consecutive directly connected layers and ignores residual connections.The authors report tight error bars but limited resources prevented substantially more runs.Source: §B.2 · p. 94

B.2.1 Image Classification with VGG

  • VGG-16 is trained on CIFAR-10 for 164 epochs with SGD, batch size 128, initial learning rate 0.1, momentum 0.9, scheduled decay, and weight decay 0.0001.Data augmentation is omitted; the learning rate drops by 10 at epochs 82, 120, and 160.Source: §B.2.1 · pp. 94–95
  • For VGG-16, linear weights use direct SVDs, convolutional weights use flattened (c_out, c_in · h · w) tensors, and alignment uses the next layer’s center spatial position.The individual singular-value plot uses convolutional layer 12, while the alignment plot connects layers 11 and 12.Source: §B.2.1 · p. 95

B.2.2 Image Generation with UNets

  • A 4-layer residual UNet is trained for MNIST generation with AdamW, batch size 128, learning rate 0.0003, and 100 epochs.The model uses swish activations and linear, convolutional, and transposed-convolutional layers; transposed convolutions require an extra dimension transpose for SVDs.Source: §B.2.2 · p. 95

B.2.3 Speech Recognition with LSTMs

  • A bidirectional LSTM is trained for LibriSpeech speech recognition with AdamW, batch size 32, learning rate 0.0003, weight decay 0.1, and 50-epoch cosine annealing.SVDs use matrix parameters; alignment emphasizes the gate submatrix and connects input and hidden parameters across directions and depth.Source: §B.2.3 · pp. 95–96

B.2.4 Language Modeling with Transformers

  • The Transformer language model uses the 160-million-parameter Pythia configuration on Wikitext-103, with 12 layers, dimension 768, 12 heads, and context length 2048.It is trained with AdamW, batch size 256, learning rate 0.0006, weight decay 0.1, gradient clipping, warmup, and cosine decay to 10%.Source: §B.2.4 · p. 96
  • For Transformer SVDs, query-key-value parameters are treated as one (3d_model, d_model) matrix, alongside output and MLP linear matrices.The individual singular-value plot uses W_1 of the eighth-layer MLP.Source: §B.2.4 · p. 96
  • Transformer alignment averages per-head comparisons for W_Q-W_K and W_V-W_O, and also examines MLP and cross-block connections.The individual layer alignment plot compares W_1 and W_2 in the eighth-layer MLP.Source: §B.2.4 · p. 96

B.2.5 Spectral Dynamics with Scale (Pythia)

Applying the Section 4.5 perspective to Pythia models reveals broadly similar spectral rank dynamics across scales, with scale-dependent singular-value evolution and later alignment in larger models.

Spectral Dynamics with Scale (Pythia)

  • The analysis uses Pythia training trajectories spanning 70m to 12b parameters, but SVD and alignment memory limits restrict the largest analyzed model to 2.8b parameters.The models were not trained by the authors because they lacked the resources to do so.Source: §B.2.5 · p. 96
  • Across Pythia scales, rank dynamics are similar, but the 2.8b model has proportionally fewer singular values becoming large than the 410m model.The comparison uses the 7th-layer MLP, which is present at all scales.Source: §B.2.5 · p. 96, Figure B.1 · p. 97
  • Alignment is absent beyond the top rank and occurs much later in the largest model, while larger models show stronger alignment scores across layers.The alignment plot also contains less noise as model scale increases.Source: §B.2.5 · p. 96, Figure B.1 · p. 97

B.2.6 Weight Decay Experiments

Weight decay experiments show that even high weight decay does not cause models to entirely break down, while pruning more closely approximates full-model performance as weight decay increases.

B.2.6 Weight Decay Experiments

  • Across weight-decay settings from 0 to 10.0, models do not entirely break down, and pruned models approximate full models better at higher weight decay.Figure B.3 compares trained and pruned model performance across tasks.Source: §B.2.6 · p. 97

B.2.7 Grokking Experiments

The grokking experiments examine alignment across network types and implement Transformer, slingshot, and deep-MLP settings based on prior work with specified architectural and training modifications. Larger weight decay increases maximum alignment magnitude, with the Transformer showing the strongest alignment despite nonlinear MLP separation.

B.2.7 Grokking Experiments

  • The Transformer follows Nanda et al. (2023), using sinusoidal rather than learned positional encodings.Source: §B.2.7 · p. 97
  • Larger weight decay increases maximum alignment magnitude, and the Transformer shows the strongest alignment even when nonlinearities separate its MLP layers.Alignment is measured both for a single pair over time and across matrix pairs over network depth.Source: Figure B.2 · p. 98
  • The slingshot experiment follows Thilak et al. (2022) with a 1-layer architecture from Nanda et al. (2023) and addition modulo 97.The original grokking plot log-scales the x-axis, whereas this experiment does not.Source: §B.2.7 · p. 98
  • The deep-MLP experiment follows Fan et al. (2024) using a 12-layer, width-400 ReLU MLP trained with MSE on MNIST.It uses 2000 examples, batch size 100, weight decay 0.01, and initialization scale 8 (Liu et al., 2023).Source: §B.2.7 · p. 98

B.2.8 Random Label Experiments

The experiments compare training with random versus true labels under constant-learning-rate settings across MLP, VGG, and LSTM networks. They also examine how weight decay affects training loss and the gap between pruned and unpruned parameters.

MLP setup

  • A 4-layer MLP is trained on CIFAR10 with random or true labels for 300 epochs using momentum 0.9 and a constant learning rate of 0.001.The constant learning rate is intended to avoid conflating results with learning-rate scheduling used by Zhang et al. (2021).Source: §B.2.8 · p. 98

VGG setup

  • VGG experiments omit weight decay and learning-rate scheduling, using a constant learning rate of 0.01.Source: §B.2.8 · p. 98

LSTM setup

  • The LSTM random-label experiment extends training to 200 epochs, but the network lacks sufficient capacity to memorize the data completely.Source: §B.2.8 · p. 98

Weight-decay analysis

  • Low training loss remains achievable under high weight decay, while increasing weight decay narrows the gap between pruned and unpruned parameters.This supports the idea that the parameters become lower rank.Source: Figure B.3 · p. 99

B.2.9 Magnitude Pruning Experiments · B.2.10 LMC Experiments

The experiments define global magnitude-pruning masks and retrain rewound networks, while LMC experiments quantify barriers and singular-vector agreement across branched checkpoints. Procedures include architecture-specific sparsity limits, normalization recomputation, and controlled interpolation comparisons.

B.2.9 Magnitude Pruning Experiments

  • Global magnitude pruning keeps the top 5% of weights after flattening all linear and convolutional weights, excluding the final linear layer.The retained weights are reshaped into layerwise masks, producing different sparsity levels across layers.Source: §B.2.9 · pp. 99–100
  • Random masks use the per-layer sparsities produced by the global magnitude mask.Source: §B.2.9 · p. 100
  • Pruned networks are retrained from epoch 4, with masked weights and gradients fixed at zero, and results averaged over three random seeds.Epoch 4 is the point after LMC identified by Frankle et al. (2020).Source: §B.2.9 · p. 100
  • For the LSTM, the same pruning procedure reaches only 25% sparsity because higher sparsity causes larger performance degradations.Source: §B.2.9 · p. 100

B.2.10 LMC Experiments

  • LMC experiments use 10 initializations from checkpoints saved throughout the first five epochs, training three trunks and three branches per trunk.Plots are averaged over the resulting nine branches.Source: §B.2.10 · p. 100
  • Checkpoint barriers are computed between branches from a shared checkpoint, with zero barrier defining linear mode connectivity.The barrier computation follows Neyshabur et al. (2020).Source: §B.2.10 · p. 100
  • After interpolation, batch-normalization parameters for VGG-16 and group-normalization parameters for UNet are recomputed, while singular-vector agreement is measured between branch endpoints.Source: §B.2.10 · p. 100
  • LMC and singular-vector agreement are measured at 11 evenly spaced interpolations between branch endpoints sharing split epoch and branch seed but differing in trunk initialization.Source: §B.2.10 · p. 100

B.2.11 Perturbed LMC Experiments

The experiments perturb weights after dynamics stability, when LMC is expected at the end of training, using normalized random Gaussian noise. The output layer is excluded, and the Transformer input layer is also excluded for computational reasons.

B.2.11 Perturbed LMC Experiments

  • Weights are perturbed after dynamics stability, with epoch 4 sufficiently late in all cases, to test LMC near the end of training.The perturbations use randomly sampled normal noise 𝜖∼𝒩(0, 𝐼) scaled so that ∥𝜖∥=η∥W∥.Source: §B.2.11 · p. 100
  • The perturbation magnitude varies over η∈{0.0, 0.1, 0.25, 0.5, 1.0, 2.5}.Source: §B.2.11 · p. 100
  • The output layer is not perturbed because it substantially affects optimization, while the Transformer input layer is excluded because it is too computationally expensive.Source: §B.2.11 · p. 100

B.3 Limitations … Diverse Feature Learning

The appendix reports limitations from computational constraints, while additional experiments examine initialization, learning rate, grokking, alignment, and architecture-dependent spectral dynamics. It also documents compute, software, and seed practices for the experiments.

B.3 Limitations

  • The study uses at most 3 random seeds, omits biases and normalization, does not align layers across residual connections, and does not fully test models larger than the 12-layer Transformer.Error bars are generally tight except for the grokking generalization epoch, and results on the Pythia suite are encouraging.Source: §B.3 · pp. 100–101

B.4.2 Effect of learning rate

  • Learning-rate effects are mixed: 0.01 gives the strongest balancedness, while the figure caption reports alignment increasing with larger learning rate.The tested learning rates are 0.001, 0.01, and 0.1.Source: §B.4.2 · p. 101, Figure B.6 · p. 105

B.4.5 Image classification

  • MLP, LeNet-5, and VGG-16 show qualitative agreement in spectral dynamics and a trend toward rank minimization on CIFAR10.Figure B.9 presents the architectures side by side for comparison.Source: §B.4.5 · p. 101, Figure B.9 · p. 106

B.5 Compute Resources

  • Storing all experiment checkpoints requires on the order of 10 terabytes.Source: §B.5 · p. 102

B.6 Code Sources / Diverse Feature Learning

B.4.1 Effect of initialization scale

  • Smaller initialization produces slightly stronger alignment in an MLP trained on CIFAR10, matching the predicted deep-linear trend.The experiment varies default PyTorch initialization by multipliers from 0.001 to 1.0.Source: §B.4.1 · p. 101, Figure B.5 · p. 104

B.4.3 Individual seed plots for grokking

  • In grokking modular addition with Transformers, the transition to low error coincides with a transition to low rank.Figure B.7 averages 3 seeds, while Figure B.4 shows that good generalization is associated with lower-rank parameters and more prevalent top-rank alignment.Source: §B.4.3 · p. 101, Figure B.4 · p. 103, Figure B.7 · p. 105

B.4.4 Full alignment matrix evolution

  • The full alignment matrix concentrates most of its signal on the diagonal, supporting the study’s focus on diagonal alignment.Off-diagonal elements showed little signal anecdotally in the CIFAR10 MLP experiment.Source: §B.4.4 · p. 101, Figure B.8 · p. 105

C.1 PGD Experiments with Image Classification

The experiments train 4-layer MLPs on MNIST-CIFAR binary classification using PGD applied to all layers except the last. Training uses Adam for up to 50 PGD rounds, removing three additional ranks per round.

C.1 PGD Experiments with Image Classification

  • The experiments use 4-layer MLPs with hidden dimension 256 for MNIST-CIFAR binary classification.Models are trained with Adam at learning rate 0.005, batch size 128, and no regularization.Source: §C.1 · p. 107
  • PGD is applied to every layer except the last, which is excluded because of its low rank.Biases are also left unperturbed because ablating them produced no difference in the results.Source: §C.1 · p. 107
  • Training runs for 50 PGD rounds or until training stops when no ranks remain, projecting out three additional ranks per round.Source: §C.1 · p. 107

C.2 CKA Experiments with Image Classification … C.4 Ensemble Weighting

The experiments evaluate CKA, diversity methods, and ensemble weighting across image classification and language modeling, using systematic sweeps over layers, coefficients, thresholds, and temperatures. Language-model ensembles are compared with a similarly parameterized but wider model.

C.2 CKA Experiments with Image Classification

  • Image-classification experiments use ResNet-18 with Adam across ImageNet and CIFAR100, restricting ImageNet to 10 classes and selecting CKA after projection layers with γ=10.PGD removes 5% of singular-value probability mass per round; experiments train 10 rounds of 50 epochs, with a 4×-width larger model in one setting.Source: §C.2 · p. 107

C.3 Ensembling in Language Modeling

  • Language-model experiments use an 8-layer Transformer trained on BabiStories for six rounds, and compare ensembles with a 3×-wider model of similar parameter count but greater compute.Unless otherwise noted, ensemble loss is measured after averaging model probability distributions.Source: §C.3 · p. 108

C.3.1 CKA

  • Language-model CKA sweeps layer locations and coefficients, with the best setting applying CKA only to MLP layers at a coefficient selected from 0.1–30.0.This choice parallels applying CKA after final convolutional layers in residual blocks.Source: §C.3.1 · p. 108

C.3.4 Entropy weighting

  • Entropy weighting uses the same tuning process as focal weighting.Source: §C.3.4 · p. 108

C.4 Ensemble Weighting

  • Ensemble-weighting evaluations sweep temperatures from 0.01–100 and interpolate with the uniform distribution using α values 0.2, 0.5, and 1.0.Source: §C.4 · p. 108

Related papers

Loading 2608.28948v1…