Source-linked AI summary

Gradient Descent Happens in a Tiny Subspace

Guy Gur-Ari, Daniel A. Roberts, Ethan Dyer

arXiv:1812.04754v1cs.LGcs.AIstat.ML

TL;DR

Deep learning uses SGD successfully despite highly nonlinear and nonconvex loss landscapes, but the dynamics underlying that success remain unclear. The paper analyzes gradients and Hessians across deep-learning settings and a solvable classification model, finding that gradients concentrate in a small, approximately preserved top Hessian subspace, with implications for learning and optimization.

  • Problem

    SGD trains nearly every large-scale machine-learning model, yet its success remains somewhat mysterious in highly nonlinear and nonconvex deep-learning loss landscapes.

  • Method

    The paper measures gradient overlap with the top Hessian subspace across architectures and datasets, then analyzes a solvable softmax-regression Gaussian-mixture model.

  • Results

    Across the studied scenarios, the gradient concentrates in the top Hessian subspace, whose dimension equals the number of classes, while the subspace is approximately preserved over long training periods.

  • Takeaways & Limitations

    The concentration of learning in a small subspace with positive Hessian eigenvalues may partially explain how deep networks train well despite nonconvex loss functions and motivates second-order optimization ideas.

  • Takeaways & Limitations

    The toy model is illustrative rather than definitive, and practical methods based on these observations are left for future work.

Abstract

from arXiv · show

We show that in a variety of large-scale deep learning scenarios the gradient dynamically converges to a very small subspace after a short period of training. The subspace is spanned by a few top eigenvectors of the Hessian (equal to the number of classes in the dataset), and is mostly preserved over long periods of training. A simple argument then suggests that gradient descent may happen mostly in this subspace. We give an example of this effect in a solvable model of classification, and we comment on possible implications for optimization and learning.

1 INTRODUCTION

The paper investigates gradient and Hessian dynamics during SGD, finding that gradients quickly concentrate in the Hessian’s small top subspace while that subspace remains approximately preserved over long training periods.

  • Motivation: SGD is ubiquitous in large-scale machine learning, but its success remains poorly understood in highly nonlinear, nonconvex loss landscapes.The paper studies gradient and Hessian dynamics during SGD to shed light on this question.
  • Hessian structure: Overparameterized models typically exhibit many flat directions, reflected by a Hessian spectrum with a large bulk of small eigenvalues and a smaller top component of larger positive eigenvalues.The top subspace has dimension k, equal to the number of classes in the classification objective.
  • Main observations: The gradient quickly moves into the Hessian’s top subspace, although its direction within that subspace appears random relative to the eigenvector basis.The top subspace is spanned by the leading Hessian eigenvectors.
  • Main observations: The top Hessian eigenvectors evolve nontrivially but generally avoid mixing with bulk eigenvectors, preserving the top subspace approximately over hundreds of training steps or more.These observations hold across fully connected, convolutional, and ResNet-18 architectures and across datasets.
  • Implication: If the gradient is mostly contained in the top subspace, projecting it there should produce a similar leading-order loss decrease under gradient descent.For a step −ηg, the leading-order loss change is δL = −η ∥g∥2, assuming the linear approximation is valid.
  • Toy model: A solvable softmax-regression model on a Gaussian mixture reproduces the observed concentration, with the gradient in the top Hessian subspace and zero bulk eigenvalues in the zero-variance limit.The authors emphasize that this toy model is illustrative rather than a definitive explanation; adding small variance is argued not to change the qualitative conclusions.

2 THE GRADIENT AND THE TOP HESSIAN SUBSPACE

The gradient quickly concentrates in the Hessian’s top-k subspace, while that subspace remains approximately preserved during training across several architectures and datasets. A normalized gradient–Hessian overlap provides an efficient proxy for this concentration, supporting the view that optimization occurs largely in a small, slowly evolving subspace.

  • The gradient lies predominantly in the Hessian’s top subspace, and this concentration develops quickly during training.
  • The top subspace is spanned by the top k Hessian eigenvectors, where k equals the number of classes.
  • 2.2 HESSIAN-GRADIENT OVERLAP: The Hessian-gradient overlap is used as a computationally efficient proxy for the fraction of the gradient in the top subspace.
  • 2.2 HESSIAN-GRADIENT OVERLAP: Under a spectrum with k positive eigenvalues and the remaining eigenvalues near zero, the estimated overlap is 3/4 ≈ 0.87, whereas an unrestricted random vector has much smaller overlap.
  • 2.3 EVOLUTION OF THE TOP SUBSPACE: After about 100 steps, the top-k subspace overlap remains significant over much longer intervals, unlike the next-k subspace whose overlap decays quickly.
  • 2.3 EVOLUTION OF THE TOP SUBSPACE: The approximately time-independent top-subspace projector, together with gradient concentration, supports the idea that gradient descent occurs in a tiny subspace.

3 A TOY MODEL

The toy classification model reproduces the paper’s observed concentration of gradients in a small top Hessian subspace and clarifies how noise, biases, and multiple classes modify that structure.

  • Model setup: The toy model uses softmax regression on a two-class Gaussian mixture, analyzed first in the zero-variance limit with orthogonal class means.Small sample noise is later introduced to test whether the qualitative behavior persists.
  • Exact structure: Under these assumptions, the loss has 2d − 2 flat directions, the Hessian has rank 2, and the gradient always lies in its top subspace.The two nontrivial Hessian eigenvectors span the top subspace, while the kernel is the bulk subspace.
  • Exact structure: The solvable model has two positive Hessian eigenvalues, with the top subspace always preserved and the gradient evolving within it.The solution family is parameterized by arbitrary components orthogonal to the two class means and vectors within their span.
  • Perturbations: Small sample noise produces small nonzero bulk eigenvalues while the gradient still evolves into the top subspace.The noise breaks some translation symmetry and reduces the number of flat directions; the resulting optimization-trajectory change is expected to be suppressed by σ2.
  • Generalizations: Adding biases breaks the top-subspace eigenvalue degeneracy, and the gradient becomes aligned with the eigenvector having the smaller positive eigenvalue.For k > 2 classes, the top Hessian subspace contains k positive eigenvalues; with biases, one is smaller and the gradient aligns with it.
  • Dynamics: The linearized gradient update explains alignment with the smallest nonzero eigenvalue, but this explanation may fail at late times when higher-order terms matter.The argument applies within the top subspace; it does not drive the gradient into exactly flat bulk directions because the gradient vanishes there.

4 DISCUSSION

The discussion interprets the empirical pattern as learning concentrated in a slowly varying, convex top Hessian subspace and identifies implications for second-order optimization and future study.

  • Main interpretation: Across architectures, training methods, and tasks, the Hessian splits into two slowly varying subspaces, while the gradient lies in the k-dimensional top subspace defined by the number of classes.The top subspace is associated with the k largest Hessian eigenvalues.
  • Main interpretation: Because the gradient essentially occupies a convex subspace with positive Hessian eigenvalues, this concentration may partially explain effective training despite nonconvex loss landscapes.The paper presents this as a possible explanation rather than a definitive guarantee.
  • Open questions: Future work must clarify the top subspace’s eigenvalue hierarchy, internal mixing, and whether its eigenvectors have transparent meanings useful for feature extraction.The authors also note that explicitly projecting updates onto the top subspace could alter the spectrum or unfreeze the subspace.
  • Optimization implications: A low-rank Newton approximation is suggested because the gradient predominantly lives in a tiny, nearly fixed top subspace.Inverting the Hessian there is well-defined and computationally simple.
  • Optimization implications: The top-subspace Hessian has strictly positive eigenvalues, so the proposed Newton approximation is expected to descend rather than climb.The discussion notes that similar considerations may apply to other second-order methods.

A NUMERICAL METHODS

The empirical analysis avoids explicitly representing the Hessian by using Lanczos iterations based on Hessian-vector products.

  • Numerical methods: Lanczos efficiently computes the Hessian’s top eigenvectors by repeatedly evaluating Hessian-vector products Hv.These products can be computed using common automatic-differentiation packages such as TensorFlow or PyTorch.

B HESSIAN SPECTRUM

The Hessian spectrum separates into a small top component and a large bulk, while top eigenvectors evolve over training without concentrating on particular parameters. The maximal eigenvalue grows initially, remains large, and later decays in both examined architectures.

  • Spectrum structure: The Hessian spectrum separates into the largest k eigenvalues forming a top component and the remaining eigenvalues forming a bulk.In the MNIST example, 10 nontrivial eigenvalues form a visible tail; CIFAR100 and synthetic datasets show the same class-linked pattern.
  • Spectrum structure: A density undercurrent appears around the mean 100th eigenvalue in CIFAR100, while a Gaussian fits the top-subspace portion suggestively.The spectrum was estimated from the top 1000 eigenvalues of a fully connected (256, 256) model averaged over 200 realizations.
  • Eigenvector dynamics: The maximal Hessian eigenvector changes during training and has many nonzero components across layers rather than being dominated by one parameter group.Zoomed views show it is never sparse and is a time-evolving linear combination of parameters with different coefficients.
  • Eigenvalue dynamics: The maximal eigenvalue in both a fully connected MNIST model and ResNet-18 on CIFAR10 first grows, stays large, and then decays.The authors describe the plots as qualitative because they are too rough to determine the exact decay form, and reserve averaging many realizations for future work.

C k IS FOR CLASSES

Experiments connect the nearly preserved Hessian subspace to the number of classification classes. Gradient overlap increases after training and concentrates in the top eigenvectors without reducing the gradient to a single Hessian eigenvector.

  • Subspace dimension: The section tests whether the nearly preserved Hessian subspace has dimension k, where k is linked to the number of classes.It compares subspaces of dimensions 2 through 20 in fully connected MNIST and ResNet-18 CIFAR10 experiments.
  • Subspace dimension: Top subspaces of dimensions 15 and 20 are significantly less preserved, while preservation peaks near the number of classes.Adding eigenvectors does not always increase preservation; dimensions below 10 are mostly preserved among themselves in the 10-class examples.
  • Gradient overlap: After training, the gradient spreads across the top eigenvectors and has negligible weight on eigenvectors beyond the top 10.The gradient is therefore not simply aligned with one Hessian eigenvector.

D ADDITIONAL EXPERIMENTS

Additional experiments vary optimization and model settings, including learning rate, batch size, architecture, activation, labels, and optimizer. Across these scenarios, the gradient quickly enters the top subspace and remains there.

  • Robustness experiments: Across varied learning rates, batch sizes, architectures, activations, label assignments, and optimizers, the gradient quickly converges to the top subspace and stays there.The experiments include fully connected MNIST, convolutional CIFAR10, random MNIST labels, and Adam instead of SGD.

E ANALYTIC EXAMPLE: DETAILED CALCULATIONS

A solvable two-sample, two-class softmax model reproduces the observed low-dimensional structure. Its loss has many flat directions, the Hessian has at most rank two, and late-training gradients lie in the fixed nontrivial eigenspace.

  • Model and trajectory: The analytic example uses a reduced two-sample, two-class softmax-regression problem and solves its gradient-descent trajectory in the continuous-time limit.The parameter trajectory is expressed using positive constants and components orthogonal to both data vectors.
  • Connection to experiments: The empirical and analytic results together show that gradients dynamically concentrate in the top subspace without becoming a single Hessian eigenvector.The additional experiments use the top-subspace gradient fraction across multiple settings.
  • Hessian structure: The model has 2d − 2 flat directions, so its Hessian has rank at most 2 and the gradient remains inside the nontrivial eigenspace.This provides a solvable analogue of the low-dimensional phenomenon observed empirically.
  • Gradient dynamics: The gradient approaches a solution-parameter-independent direction as t →∞, with the explicit trajectory expression decaying as 2(µ2 − µ1)(ηt)^-1 + O(t^-2).The asymptotic expression is given for the two-class model's gradient components.
  • Hessian structure: At leading order, two nontrivial Hessian eigenvectors share eigenvalue (ηt)^-1, all remaining eigenvalues vanish, and the top subspace is fixed.The gradient is contained within this fixed top Hessian subspace in the late-training limit.
Loading 1812.04754v1…