Research map / Optimization / Implicit bias and generalization
Does SGD really happen in tiny subspaces?
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that projecting SGD updates onto the top-k Hessian dominant subspace stalls training while projecting onto the bulk subspace matches SGD, and attributes the gradient-dominant alignment to SGD noise in the gradient-flow regime and to self-stabilization in the edge-of-stability regime.
Problem
Gradients during neural network training approximately align with the low-rank dominant Hessian eigenspace; the question is whether networks can be trained within that dominant subspace, which would enable more efficient training.
Approach
Dominant-subspace projected updates (Dom-SGD, Dom-GD, Dom-SAM) fail to decrease loss despite gradients aligning with the dominant subspace, while bulk-subspace projected updates train as well as the originals; the alignment is spurious and caused by SGD stochastic noise in an ill-conditioned valley landscape (GF regime) or by self-stabilization (EoS regime).
Key results
- Dom-SGD fails to further decrease training loss across MLP/CNN/Transformer settings, and slowly increases loss in the long run, while Bulk-SGD matches SGD despite removing the majority of the update.
- Switching from SGD to GD makes the gradient-dominant alignment chi_k drop from about 1 to about 0 almost immediately, showing stochastic noise is the main cause in the gradient flow regime.
- In a 2D ill-conditioned quadratic toy model, GD gradients have alignment near 0 while SGD gradients stay near 1, reproducing the phenomenon.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| project each SGD update onto the dominant subspace (ablation) | SGD · MNIST-5k, CIFAR10-5k, and SST2-1k | training loss | Hurt |
| project each SGD update onto the bulk subspace (ablation) | SGD · MNIST-5k, CIFAR10-5k, and SST2-1k | training loss | No clear effect |
| project each GD update onto the dominant subspace (ablation) | GD in the EoS regime · MNIST-5k, CIFAR10-5k, and SST2-1k | training loss | Hurt |
| project each GD update onto the bulk subspace (ablation) | GD in the EoS regime · MNIST-5k, CIFAR10-5k, and SST2-1k | training loss | No clear effect |
| project each SAM update onto the dominant subspace (ablation) | SAM · MNIST-5k | training loss | Hurt |
| project each SAM update onto the bulk subspace (ablation) | SAM · MNIST-5k | training loss | No clear effect |
| add cross-entropy loss rather than MSE loss (ablation) | Dom-SGD · MNIST-5k, CIFAR10-5k, and SST2-1k | training loss | Hurt |
| use VGG11 or ResNet8 as the architecture (ablation) | Dom-SGD · CIFAR10-5k | training loss | Hurt |
| switch from SGD to full-batch GD | SGD in the GF regime · MNIST-5k | gradient alignment with dominant subspace | Hurt |
| switch from full-batch GD to SGD | GD in the GF regime · MNIST-5k | gradient alignment with dominant subspace | Helped |
| increase injected Gaussian noise scale | Noisy Gradient Descent · MNIST-5k | gradient alignment with dominant subspace | Helped |
| decrease SGD batch size | SGD · MNIST-5k | gradient alignment with dominant subspace | Helped |
Surprising findings
- Projecting SGD updates onto the dominant subspace stops training even though the dominant component constitutes most of the gradient.
- Alignment between gradient and dominant subspace disappears when SGD is replaced by full-batch GD in the gradient flow regime, despite nearly identical trajectories.
Limitations
- Experiments are limited to small-scale models and datasets (MLP, CNN, Transformer on 5k/1k samples), with a theoretical explanation only via a toy model.
- The ill-conditioned-valley explanation is proposed as plausible rather than proven.
- Analysis focuses on optimization, not generalization.
Experimental setup
Datasets: MNIST-5k CIFAR10-5k SST2-1k
Benchmarks: MNIST-5k training loss CIFAR10-5k training loss SST2-1k training loss
Metrics: training loss chi_k gradient alignment fraction top-k Hessian eigenvalues effective learning rate
Baselines: SGD GD SAM SGD with momentum Adam
Builds on: Gur-Ari et al. gradient descent happens in a tiny subspace Edge of Stability (Cohen et al.) Sharpness-Aware Minimization Adam self-stabilization
Closest papers by meaning
- Decomposing SGD Dynamics in Neural Networks: Teacher-Induced Spikes and Variance Inflation (NeurIPS 2026 · similarity 0.90)
- On the Relation Between the Sharpest Directions of DNN Loss and the SGD Step Length (ICLR 2019 · similarity 0.90)
- Catapults in SGD: spikes in the training loss and their impact on generalization through feature learning (ICML 2024 · similarity 0.88)
- SGD vs GD: Rank Deficiency in Linear Networks (NeurIPS 2024 · similarity 0.87)
- Implicit bias of SGD in $L_2$-regularized linear DNNs: One-way jumps from high to low rank (ICLR 2024 · similarity 0.87)
- A Unified Stability Analysis of SAM vs SGD: Role of Data Coherence and Emergence of Simplicity Bias (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.