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 dissertation addresses limited understanding of neural-network optimization, including why mode connectivity occurs and how it relates to broader risks and behavior. It studies spectral and low-rank dynamics across architectures and tasks, linking them to mode connectivity, generalization, pruning, and feature learning, while proposing a method to overcome simplicity bias.

  • Problem

    Neural-network optimization and mode connectivity remain poorly understood, while scaling with more data faces practical limits under changing distributions and leaves important societal questions unresolved.

  • Method

    The dissertation empirically studies spectral dynamics across ConvNets, UNets, LSTMs, and Transformers, then uses low-rank structure to develop projected gradient descent for addressing spurious correlations.

  • Results

    The work links low-rank dynamics to weight decay, rank and alignment differences, magnitude pruning, mode connectivity, and the sharing of top singular vectors between checkpoints.

  • Takeaways & Limitations

    The findings support using spectral dynamics and low-rank bias to understand diverse learned functions and to overcome simplicity bias in neural networks.

  • Takeaways & Limitations

    Convexity proxies detect no increase around the split point, leaving unresolved whether this reflects sparse metric sampling or greater underlying convexity.

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 dissertation addresses limited understanding of how neural networks optimize and generalize, especially under distribution shifts and low-data conditions. It develops a progression from analyzing loss-landscape structure to using spectral dynamics for feature learning and simplicity-bias mitigation.

  • Neural networks increasingly influence socially consequential decisions, making their underlying mechanisms important to understand.
  • Existing approaches often rely on collecting more data, but this can require extreme amounts, fail under changing distributions, and leave the underlying science unexplained.
  • The dissertation seeks a broad understanding of practical neural networks that supports more principled treatment of distribution shifts and low-data regimes.
  • Chapter 3: Chapter 3 studies high-dimensional, diverse convex modes in neural-network loss landscapes, extending linear mode connectivity.
  • Chapter 4: Chapter 4 explains these modes through stable low-rank dynamics of top singular vectors and connects them to generalization and regularization.
  • Chapter 5: Chapter 5 uses spectral dynamics and ensembles to overcome simplicity bias and improve training in low-data regimes.

Background

Neural-network optimization is difficult because useful networks have nonconvex loss landscapes with many possible optima. Linear mode connectivity reveals unexpectedly low-loss linear paths between some solutions, raising questions about why this apparent convexity occurs and how broadly it applies.

  • Neural networks can represent many functions, but finding solutions that generalize well remains a difficult optimization problem.
  • Nonconvex loss landscapes may contain separate optima and local minima that do not provide straightforward paths to a global solution.
  • Linear mode connectivity shows that optima from a shared initialization can be joined by a low-loss linear path, suggesting optimization remains within a convex subspace.
  • The phenomenon became more general under a protocol that shares early optimization before splitting trajectories, rather than remaining limited to a small-scale experiment.
  • LMC is not explained merely by infinitesimal local convexity because distances among the split point, optima, and initialization are comparable.
  • The dissertation investigates why LMC occurs, linking it to low-rank dynamics and using that connection to develop improved feature learning.

On Convexity and Linear Mode Connectivity in Neural Networks

The dissertation finds that SGD trajectories form high-dimensional convex low-loss regions, supporting convexity as an explanation for linear mode connectivity (LMC). However, simple convexity proxies do not identify an LMC transition, and interpolation can connect functionally diverse endpoints.

  • Convex mode connectivity: SGD trajectories sharing an early training segment eventually form a high-dimensional convex hull of low loss, extending linear mode connectivity beyond pairs of endpoints.The study calls this procedure convex mode connectivity (CMC) and samples convex combinations of multiple branch endpoints.
  • Measures of convexity: Hessian spectra remain slightly nonconvex and update-to-update loss is locally convex throughout training, so these measures do not reveal when LMC begins.The results therefore support local convexity without identifying a distinct transition into the LMC regime.
  • Functional diversity: Linear interpolation is only loosely related to functional similarity: endpoints and hulls share relatively few correct predictions, even when training accuracy is high.Improved accuracy inside the hull is not explained by a large common set of predictions; instead, diverse functions appear amenable to interpolation.
  • Convex mode connectivity: After an initial training phase, convex combinations of endpoints have lower loss than the endpoints, while random perturbations around endpoints are usually worse.This contrast provides evidence that LMC reflects a convex region in the loss surface rather than merely favorable endpoint geometry.
  • Discussion: The chapter concludes that convexity may arise from diverse, easily interpolated functions rather than a large shared subset of training predictions, but sparse measurements limit the interpretation.The authors note that failure to detect a convexity change could reflect sparse sampling or optimization that is more convex than expected.

Approaching Deep Learning through the Spectral Dynamics of Weights

The dissertation uses spectral dynamics—the evolution of singular values and vectors during training—to connect rank minimization with generalization-related phenomena across diverse neural networks. It further relates these dynamics to weight decay, memorization, lottery tickets, linear mode connectivity, and diverse feature learning.

  • Spectral dynamics: Across practical networks and tasks, singular values evolve unequally, producing effective rank minimization and faster stabilization of top singular vectors.This pattern appears across image classification, image generation, speech recognition, and language modeling, without the consecutive-layer alignment assumed by some theories.
  • Grokking and regularization: Grokking is intimately linked to rank minimization, while sufficient weight decay is required for both grokking and low-rank weights.The sudden validation-loss drop coincides with the onset of low-rank behavior, and weight decay promotes rank minimization across architectures and tasks.
  • Grokking and regularization: Weight decay enhances rank minimization beyond its explicit role as a norm regularizer, further tying low-rank behavior to generalization.In some cases, it also appears to promote singular-vector alignment between consecutive weight matrices despite intervening nonlinearities.
  • Memorization and generalization: Training with random labels produces high-rank, unaligned solutions, whereas true-label training produces lower-rank solutions with aligned consecutive layers.This spectral distinction offers a perspective on the difference between memorizing and generalizing networks.
  • Additional connections: Magnitude pruning for lottery tickets acts similarly to low-rank pruning, while linear mode connectivity coincides with shared top singular vectors between checkpoints.These findings place lottery-ticket structure and mode connectivity within the same spectral-dynamics framework.

Overcoming Spurious Correlations through Diverse Feature Learning

The paper addresses simplicity bias by using spectral dynamics to drive neural networks toward more diverse features, reducing reliance on spurious correlations. Experiments span controlled image tasks, larger image classifiers, and language modeling, revealing both gains and limits of diversity objectives.

  • Projected Gradient Descent: Projected gradient descent off top singular-vector directions forces networks to learn progressively more diverse features without task-specific assumptions.In MNIST-CIFAR, later training rounds increasingly rely on CIFAR after the dominant MNIST signal is suppressed.
  • Controlled Image Experiments: MNIST-CIFAR models initially rely entirely on MNIST, but later PGD rounds produce over 70% error when CIFAR images are swapped.Swapping MNIST digits moves error from nearly 100% toward less than random chance, while CIFAR-swap error rises as CIFAR features are learned.
  • Spectral Dynamics: Top singular vectors evolve from MNIST digit detectors toward finer curve detectors and increasing CIFAR sensitivity across training rounds.The model is eventually forced to rely on CIFAR when the MNIST signal is destroyed.
  • Method Comparison: Vector-level gradient orthogonalization does not significantly learn CIFAR, indicating that the matrix structure of neural-network parameters matters.The paper attributes the failure to trivial orthogonality in high-dimensional spaces and contrasts it with matrix-aware methods.
  • Diversity Objectives: CKA-based diversity improves out-of-distribution linear probing beyond random ensembles in image classification, while remaining agnostic to the task.CKA directly minimizes representational similarity and avoids the permutation sensitivity of PGD.
  • Language Modeling: Language modeling exposes a boundary: random ensembles outperform a larger model out of distribution, while CKA does not improve over random ensembles.The paper suggests language modeling requires more capacity, making the desired semantic dissimilarity difficult to achieve without a pathological solution.

A Problem with Optimizing CKA

The dissertation examines diverse feature learning as a way to learn more from less data, using spectral dynamics to address spurious correlations and ensemble diversity. It reports gains for larger image-classification networks and out-of-distribution language-model ensembles, while finding several diversity objectives insufficient for language models.

  • Motivation: Diverse feature learning targets the greedy acquisition of simple features and aims to overcome spurious correlations when training with less data.The dissertation frames this as a route toward extracting more information from limited data.
  • Method: Spectral dynamics yield a projected-gradient method that overcomes simplicity bias on MNIST-CIFAR without task-specific assumptions.The method uses singular vectors to define simplicity and is presented as task-agnostic in the toy setting.
  • Image classification: Larger image-classification networks beat simple random ensembles, which are described as stronger than large models.The result follows from using model-wise similarity directly as a diversity loss after the initial method became ineffective on harder tasks.
  • Language modeling: Ensembles of small language models perform better out-of-distribution than larger models, but earlier diversity losses do not help language models.The dissertation studies this discrepancy through extensive ablations.
  • Open problem: The dissertation identifies structured data subsets as an unresolved requirement because difficult-example selection can reproduce random-ensemble solutions or induce spurious correlations.Finding such subsets without problem-specific heuristics remains difficult.

A.1 Task Details

This section records the task summary and introduces SIREN as a non-piecewise-linear model used to test mode connectivity beyond ReLU networks.

  • Task details: SIREN is a fully connected sine-activation network tuned to reconstruct an audio waveform from timestep inputs.It is included because prior LMC experiments used ReLU activations.
  • Task details: SIREN broadens the activation-function setting for mode-connectivity experiments beyond piecewise-linear ReLU models.The stated motivation is to observe a system that is not piecewise linear.
  • Task details: The chapter includes a summary table of the tasks studied.

A.2 Metric Calculation Details

The metric details define measurements for Hessian curvature, update-level non-convexity, mode connectivity, and singular-value structure. These measurements compare loss behavior across training and characterize the geometry of connected solutions.

  • Metric calculations: Hessian eigenvalues are calculated with power iteration using Hessian-vector products.
  • Metric calculations: Update-level non-convexity is measured by sampling two stochastic updates and estimating how often their loss gap is negative.The resulting probability is evaluated at different points throughout training.
  • Mode connectivity: The first split producing a convex hull occurs early, before training loss has converged and often before it has halved.
  • Spectral structure: Rank-minimization analyses rely on balanced initialization, which imposes a relationship between consecutive matrices and their singular-vector bases.The decomposition yields matching singular vectors up to signs and permutations under stated conditions.

B.2 Experimental Details

The experiments span image, diffusion, speech, and language models, using multiple architectures, optimizers, datasets, and alignment procedures. Singular-value and alignment measurements are adapted to each model’s parameterization.

  • Experimental setup: All experiments average plots over three random seeds, a small number chosen because the study requires many runs and typically has tight error bars.
  • Alignment: Alignment is computed for consecutive layers that directly feed one another, while residual-connection effects are ignored to reduce comparisons.
  • Vision models: The image experiments use VGG-16 on CIFAR-10 and a residual UNet on MNIST generation, with SVDs adapted to linear, convolutional, and transposed-convolutional parameters.The VGG-16 training uses SGD, while the UNet uses AdamW.
  • Speech model: The speech experiment trains a bidirectional LSTM on LibriSpeech with AdamW and cosine annealing over 50 epochs.Its matrix-only parameters make SVD computation direct, with alignment connections defined across input and hidden parameters.
  • Language model: The language experiment trains a 160-million-parameter Transformer on Wikitext-103 and computes SVDs and alignments for attention and MLP matrices.The analysis treats query, key, and value projections as one matrix before internal head splitting.

B.2.5 Spectral Dynamics with Scale (Pythia)

The Pythia analysis extends spectral-dynamics measurements across model scales, finding broadly similar rank dynamics but cleaner and stronger alignment in larger models.

  • Experimental scope: The study uses existing Pythia training trajectories from 70m to 12b parameters but limits SVD and alignment computations to the 2.8b model.The resource constraint arises from memory requirements when computing SVDs and alignment.
  • Spectral dynamics across scales: Training produces similar rank dynamics across Pythia scales, with longer training contracting singular values.The comparison uses the 7th-layer MLP, which is present at every examined scale.
  • Spectral dynamics across scales: The 2.8b model has proportionally fewer singular values evolving to large magnitudes than the 410m model.
  • Alignment with scale: Larger Pythia models show less noise and stronger alignment scores across layers.The passage suggests higher-dimensional vectors may be more likely to be orthogonal, potentially contributing to easier optimization, but presents this as a presumption.

B.2.6 Weight Decay Experiments

Weight-decay experiments show that models can retain low training loss at high regularization, while pruned models increasingly approximate their full counterparts.

  • Performance under weight decay: High weight decay does not entirely break down model training, and pruned-model performance becomes closer to full-model performance as weight decay increases.The experiments evaluate weight-decay settings from 0 through 10.0.

B.2.7 Grokking Experiments

The grokking experiments examine Transformer alignment under varying weight decay alongside deep-MLP and modular-arithmetic settings, with the Transformer showing especially strong alignment.

  • Transformer setup: Transformer experiments use settings based on prior work but replace learned positional encodings with sinusoidal encodings.
  • Alignment results: Across weight-decay settings, larger weight decay produces higher maximum alignment magnitude, with the Transformer showing the strongest alignment.The Transformer remains strongly aligned even when nonlinearities separate its MLP layers.
  • Slingshot setup: The slingshot experiment uses modular addition modulo 97 and a 1-layer architecture rather than the referenced 2-layer architecture.Its x-axis is linear rather than logarithmic for presentation clarity.
  • Deep-MLP setup: The deep-MLP experiment uses a 12-layer, width-400 ReLU network trained with MSE loss on MNIST.It uses 2000 examples, batch size 100, weight decay 0.01, and initialization scale 8.

B.2.8 Random Label Experiments

Random-label experiments compare training configurations across VGG, MLP, and LSTM settings, including the effect of weight decay and model capacity on pruned performance and memorization.

  • CIFAR10 random labels: A 4-layer CIFAR10 MLP compares random and true labels using SGD with momentum 0.9, learning rate 0.001, and 300 training epochs.A constant learning rate is used instead of the schedule in the referenced setting to avoid conflating the results.
  • VGG setup: VGG experiments omit weight decay and learning-rate scheduling, using a constant learning rate of 0.01.
  • LSTM capacity: The LSTM training budget is extended to 200 epochs, but the network lacks sufficient capacity to memorize the data completely.
  • Pruning and weight decay: Increasing weight decay narrows the training-loss gap between pruned and unpruned parameters.

B.2.9 Magnitude Pruning Experiments

These experiments evaluate magnitude pruning and lottery-ticket-style retraining across VGG and LSTM models, while measuring loss barriers and singular-vector agreement between connected solutions.

  • Magnitude pruning: Global magnitude pruning retains the top 5% of weights after flattening linear and convolutional weights, excluding the last linear layer.The retained weights are reshaped into layerwise masks, producing different sparsity levels across layers.
  • Retraining: Pruned networks are rewound to epoch 4, after the point of linear mode connectivity, and retrained with masked weights and gradients.Results are averaged over 3 random seeds.
  • Architecture-specific setting: For the LSTM, the mask reaches only 25% sparsity because higher sparsity causes larger performance degradation.
  • Connectivity analysis: The connectivity analysis uses 5 evenly spaced checkpoints from the first five epochs, three trunks, and three branches per trunk.This yields 9 averaged branches for the reported plots.
  • Connectivity analysis: Barrier values are computed between branch endpoints that share a common trunk, with zero barrier defining linear mode connectivity.
  • Spectral analysis: Singular-vector agreement is measured between corresponding parameters and across 11 endpoint interpolations with matched split epochs and branch seeds.

B.2.11 Perturbed LMC Experiments

The perturbed LMC experiments test stability after dynamics have become stable by applying normalized random perturbations to weights at late training.

  • Perturbation protocol: Weights are perturbed after dynamics stability, using normal noise with norm η∥W∥ for η values from 0.0 to 2.5.The output layer is excluded, and the Transformer input layer is also excluded because of computational cost.

B.3 Limitations

The study reports computational, parameterization, architectural, and scale limitations, alongside supplementary experiments examining initialization, learning rate, rank dynamics, and alignment.

  • Study limitations: The experiments use no more than 3 random seeds, and the grokking generalization epoch has wider error bars than other results.
  • Study limitations: The analysis omits one-dimensional parameters and alignment across residual connections, leaving some parameter interactions unexamined.
  • Study limitations: Results are limited to models no larger than a 12-layer Transformer, although Pythia-suite results are described as encouraging.
  • Supplementary dynamics: Smaller initialization produces slightly stronger alignment, consistent with predictions from deep linear models.
  • Supplementary dynamics: Learning rate effects are mixed: learning rate 0.01 produces the strongest balancedness, while larger or smaller values have weaker effects.
  • Supplementary dynamics: During grokking on modular addition, the transition to low validation error coincides with a transition to low rank.
  • Cross-architecture comparisons: Across CIFAR10 architectures, spectral dynamics qualitatively agree in their trend toward rank minimization, while alignment signal is concentrated mainly on the diagonal.
  • Computational scope: The experiments require roughly 100 GPUs, days to rerun end-to-end, and about 10 terabytes for checkpoint storage.

C.1 PGD Experiments with Image Classification

The PGD experiments test rank-removal training on binary image classification using small MLPs, excluding the final layer and biases from projection.

  • Experimental setup: The experiments train 4-layer, width-256 MLPs on MNIST-CIFAR binary classification with Adam and no regularization.Training uses learning rate 0.005, batch size 128, and lasts up to 50 PGD rounds.
  • Projection procedure: Each PGD round projects away 3 additional ranks from every layer except the final layer.The final layer is excluded because it has low rank.
  • Projection procedure: Biases are not projected because their ablation produced no difference in the reported results.

C.2 CKA Experiments with Image Classification

The CKA experiments evaluate layer placement and coefficient settings across image-classification and language-model experiments, alongside several other diversity methods.

  • Image classification: Image-classification experiments use ResNet-18 with Adam, learning rate 3e-4, batch size 512, and weight decay 1e-4.Experiments span ImageNet to CIFAR100, with ImageNet restricted to 10 classes for speed.
  • Language modeling: Language-model experiments train an 8-layer Transformer with 512 hidden dimensions and 8 heads on BabiStories.Models use six rounds of ten epochs, batch size 128, sequence length 128, and a warmup stable decay schedule.
  • CKA experiments: CKA experiments sweep the loss across attention, MLP, early, late, middle, and all layers, with coefficients from 0.1 to 30.0.The best setting uses only MLP layers.
  • Diversity-method tuning: NCL experiments tune γ from 0.01 to 10 while comparing correct-class, top-k, and top-p variants.Top-k values are 1, 5, and 10; top-p thresholds range from 0.05 to 0.9.
  • Diversity-method tuning: Focal, entropy, and ensemble weighting methods receive separate hyperparameter sweeps over coefficients, temperatures, or interpolation weights.Focal weighting uses coefficients from 0.01 to 10, while ensemble temperatures range from 0.01 to 100 and α from 0.2 to 1.0.
Loading 2608.28948v1…