[Paper Review] Gradient Descent Happens in a Tiny Subspace

Guy Gur-Ari, Daniel A. Roberts+1 more|arXiv (Cornell University)|Dec 12, 2018

Stochastic Gradient Optimization Techniques14 references110 citations

TL;DR

The gradient during training concentrates in the top Hessian subspace (dimension equal to the number of classes) and this subspace remains nearly fixed, suggesting gradient descent effectively operates in a tiny evolving subspace.

ABSTRACT

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.

Motivation & Objective

• Investigate how gradients and Hessian spectra behave during SGD in large-scale, overparameterized models.
• Characterize the subspaces of the Hessian (top vs. bulk) and how the gradient projects onto them.
• Examine whether the gradient dynamics are confined to a small, slowly changing subspace and its implications for optimization.
• Provide empirical evidence across architectures and datasets and offer a toy model to illustrate the mechanism.

Proposed method

• Decompose the Hessian into a top subspace spanned by its largest k eigenvectors (k = number of classes) and a bulk subspace.
• Measure the gradient’s projection onto the top subspace and quantify the fraction f_top = ||g_top||^2 / ||g||^2 over training.
• Compute Hessian-gradient overlap to assess alignment between g and Hg as a proxy for subspace concentration.
• Evaluate subspace overlap over training steps to determine preservation of the top subspace.
• Present a solvable toy model (softmax regression on Gaussian mixtures) to illustrate the observed dynamics.
• Use Lanczos method and Hessian-vector products to estimate top eigenvectors without forming the full Hessian.

Experimental results

Research questions

• RQ1 Does the gradient concentrate in the top Hessian subspace during training?
• RQ2 Is the top Hessian subspace approximately preserved over long training runs across architectures and datasets?
• RQ3 How does the gradient’s projection onto the top subspace relate to optimization progress?
• RQ4 Can a simple toy model replicate the observed dynamics and provide analytic insight?

Key findings

• Across architectures and datasets, the gradient quickly concentrates in the top Hessian subspace (dimension k, the number of classes).
• The top subspace remains approximately preserved over long training periods, while the bulk subspace changes more rapidly.
• The overlap between the gradient and its Hessian-projected component is high (near 1) in empirical measurements.
• A simple toy model reproduces the phenomenon: the gradient lies in the top subspace, and the Hessian has two large eigenvalues with the rest near zero.
• Including small variance or biases perturbs the spectrum but leaves the main top-subspace concentration intact.
• The observed dynamics suggest gradient descent effectively operates in a convex, low-dimensional subspace even in high-dimensional, nonconvex loss landscapes.