Source-linked AI summary
On the Relation Between the Sharpest Directions of DNN Loss and the SGD Step Length
Stanisław Jastrzębski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, Amos Storkey
TL;DR
The paper addresses the limited understanding of curvature along SGD trajectories and its relationship to the sharpest Hessian directions. It empirically analyzes these dynamics and finds that SGD initially moves toward sharper regions, while reduced learning rates along sharpest directions can improve training speed and yield sharper, equally or better-generalizing solutions.
Problem
Curvature along the SGD trajectory and its interaction with sharpest Hessian directions remain poorly understood, despite links between final-minimum curvature and generalization.
Method
The paper analyzes Hessian eigenvalues and SGD steps along sharpest directions across neural-network training trajectories, then evaluates a reduced-learning-rate variant called Nudged-SGD.
Results
SGD initially visits increasingly sharp regions with peak curvature influenced by learning rate and batch size, while reduced sharp-direction learning rates often accelerate training and produce sharper regions that generalize as well or better.
Takeaways & Limitations
Sharpest-direction dynamics influence the regions SGD reaches, overall training speed, and final generalization capability.
Takeaways & Limitations
The effect on final generalization can be dataset dependent, and broader studies on more diverse datasets are left for future work.
Abstract
from arXiv · showhide
Stochastic Gradient Descent (SGD) based training of neural networks with a large learning rate or a small batch-size typically ends in well-generalizing, flat regions of the weight space, as indicated by small eigenvalues of the Hessian of the training loss. However, the curvature along the SGD trajectory is poorly understood. An empirical investigation shows that initially SGD visits increasingly sharp regions, reaching a maximum sharpness determined by both the learning rate and the batch-size of SGD. When studying the SGD dynamics in relation to the sharpest directions in this initial phase, we find that the SGD step is large compared to the curvature and commonly fails to minimize the loss along the sharpest directions. Furthermore, using a reduced learning rate along these directions can improve training speed while leading to both sharper and better generalizing solutions compared to vanilla SGD. In summary, our analysis of the dynamics of SGD in the subspace of the sharpest directions shows that they influence the regions that SGD steers to (where larger learning rate or smaller batch size result in wider regions visited), the overall training speed, and the generalization ability of the final model.
1 Introduction
The paper studies how SGD interacts with the sharpest Hessian directions throughout training, finding that curvature initially increases before stabilizing or declining and that these directions shape optimization and generalization.
- 1 Introduction: The study extends prior curvature–generalization work by analyzing the entire SGD trajectory rather than only the endpoint.Its motivation is the empirical relationship between final-minimum curvature and generalization.
- 1 Introduction: SGD initially visits increasingly sharp regions, reaching a peak curvature influenced by both learning rate and batch size before the largest eigenvalues stabilize or decrease.This trajectory-wide pattern is reported across the networks and datasets explored.
- 1 Introduction: Along the sharpest directions, SGD commonly takes steps large enough to cross the bowl-like minimum rather than approach it.The paper interprets this as a step size too large relative to the local curvature.
- 1 Introduction: Reducing the learning rate along sharpest directions can accelerate optimization while producing sharper regions that generalize as well as or better than vanilla SGD.The comparison uses vanilla SGD with the same small learning rate.
- 1 Introduction: Sharpest-direction dynamics influence the regions SGD reaches, training speed, and final generalization, with larger learning rates or smaller batches leading to wider visited regions.These are presented as practical consequences of the trajectory analysis.
2 Experimental setup and notation
The experiments examine Hessian curvature and sharpest-direction dynamics in neural networks trained with SGD, using several architectures and datasets while defining the spectral quantities used for analysis.
- 2 Experimental setup and notation: Experiments primarily use Resnet-32 and SimpleCNN on CIFAR-10, with additional studies of VGG-11 and a bidirectional LSTM on Penn Tree Bank.The setup includes standard augmentation, Resnet-32 L2 regularization, and LSTM dropout.
- 2 Experimental setup and notation: For SimpleCNN and Resnet-32 trained on CIFAR-10 with η = 0.1 and S = 128, the largest eigenvalues grow sharply at first and then evolve in an oscillatory-like manner.Training and test accuracy are provided alongside the curvature traces for reference.
- 2 Experimental setup and notation: The analysis tracks Hessian eigenvalues along SGD trajectories, where η and S denote learning rate and batch size, respectively.The largest eigenvalue λmax is the Hessian spectral norm, and the top eigenvectors define the sharpest directions.
- 2 Experimental setup and notation: Top Hessian eigenvalues are approximated with the Lanczos algorithm using approximately 5% of the training data.The full Hessian spectrum is considered computationally infeasible for reasonably large models.
3 A study of the Hessian along the SGD path
Along the SGD trajectory, curvature initially grows toward a peak shaped by learning rate and batch size, while SGD steps become large and aligned with sharpest directions. These dynamics influence the regions visited and motivate modifying optimization along those directions.
- 3.1 Largest eigenvalues of the Hessian along the SGD path: The largest Hessian eigenvalues initially grow, reach a maximum, and then generally decrease or stabilize during training.This growth can coincide with difficult optimization and, in full-batch training, instability and a sharp accuracy drop.
- 3.1 Largest eigenvalues of the Hessian along the SGD path: Higher learning rates or smaller batch sizes correlate with earlier, lower peaks in the Hessian spectral norm across the examined models and datasets.The observed pattern affects the full SGD trajectory, not only its endpoint.
- 3.2 Sharpest direction and SGD step: In the sharpest direction, the SGD step commonly overshoots the loss minimum, so increasing the step scale raises loss while smaller scales reduce it.For SimpleCNN, scales 2 and 4 increase loss by 2.1% and 11.1%; for Resnet-32, they increase loss by 3% and 13.1%.
- 3.2 Sharpest direction and SGD step: SGD gradients are substantially aligned with the top five sharpest directions, with maximum average cosine values roughly between 0.2 and 0.4 depending on learning rate and model.This alignment links sharpest-direction dynamics to overall training behavior.
- 3.3 How SGD steers to sharp regions in the beginning: Restricting or removing top-eigenvector updates changes the region reached: top-direction-only updates lower λmax but raise loss, whereas subtracting them finds sharper regions at similar loss to vanilla SGD.These comparisons show that early top-direction updates strongly influence the overall optimization path.
4 Optimizing faster while finding a good sharp region
NSGD reduces the learning rate along the sharpest Hessian directions, and experiments show that it can optimize faster while reaching sharper endpoints with comparable or better generalization than vanilla SGD.
- NSGD is motivated by evidence that SGD steps are highly aligned with sharpest directions yet commonly fail to minimize loss within those directions.The authors present NSGD as an investigative variant rather than a practical optimizer.
- NSGD uses η′ = γη along the top K Hessian eigenvectors while retaining the standard SGD gradient in the remaining directions.The experiments varied K from 1 to 20 and γ over {0.01, 0.1, 1, 5}.
- NSGD with K = 20 reached λmax ≈ 8 · 10^3 versus approximately 150 for baseline SGD and retained over 5% validation-accuracy advantage after 50 epochs.For SimpleCNN, the validation-accuracy advantage was 1%.
- Using γ < 1 can produce a significantly sharper final endpoint with slightly improved generalization compared with baseline SGD using the same η = 0.01.Using γ > 1 instead led to a wider endpoint.
- NSGD incurs approximately 2–3× more time per epoch than vanilla SGD because of the Lanczos iterations used to compute top eigenvectors.This overhead depends on the number of iterations required for convergence.
- Baseline SGD reached 86.4% test accuracy with ||H||F = 272, whereas NSGD with γ = 0.01 reached 87.0% with ||H||F = 1018.On SimpleCNN, the corresponding values were 86.6% and 191 for baseline SGD, versus 87.4% and 287 for NSGD.
- Across additional datasets, NSGD showed faster optimization and sharper regions, but final generalization was not always improved.The authors note that the generalization effect may depend on the dataset.
5 Related work
Prior work linked endpoint curvature with generalization, tracked Hessian evolution, and studied SGD dynamics; this paper instead emphasizes sharp-direction dynamics throughout training.
- Previous Hessian studies often focused on late training, although some reported initial spectral-norm growth and later reduction.
- Earlier studies commonly related wider endpoint minima to better generalization and connected width with batch size or the learning-rate-to-batch-size ratio.
- Related SGD-dynamics work examined optimizer trajectories, noise along the top eigenvector, or interpolation between consecutive parameter values.
- This paper differs by analyzing from the beginning of training whether SGD steps are too large relative to curvature along the sharpest directions.
6 Conclusions
The paper argues that SGD dynamics along sharpest directions shape trajectory curvature, optimization speed, and final generalization, motivating further study of these directions.
- SGD moves toward increasingly sharp regions until a level governed by learning rate and batch size, while its step is large and highly aligned with sharpest directions.
- Understanding optimization along sharpest directions may help explain neural-network generalization and inform optimizers tailored to neural networks.
A.1 Additional results on the evolution of the largest eigenvalues of the Hessian
Additional experiments show that Batch-Normalization partially stabilizes early full-batch training, while architecture, schedules, and momentum affect Hessian evolution.
- Batch-Normalization partially resolves instability during the early phase of full-batch training.
- The largest-eigenvalue evolution reported for the main experiments extends to VGG-11 and Batch-Normalized Resnet-32, with large absolute eigenvalues in inference-mode Resnet-32.
- The first-stage length of a learning-rate schedule changes how Hessian curvature and generalization performance evolve during later training.The schedule used η = 0.1 for L epochs before dropping to η = 0.01, with L in {10, 20, 40, 80}.
- Large momentum limits the maximum spectral norm of the Hessian throughout training on VGG11 with CIFAR10 and CIFAR100.
B Additional results for Sec. 3.2
Additional experiments extend the sharp-direction analysis across eigenvectors, architectures, learning rates, and longer training runs.
- Training and validation accuracy are reported for the SimpleCNN and Resnet-32 experiments from Fig. 6.
- The third and fourth Hessian eigenvector visualizations show loss surfaces that are much flatter than the sharpest-direction reference.
- Repeating the Resnet-32 and SimpleCNN experiments with η = 0.05 tests whether the observed results depend on learning rate.
- Replotting the experiment for 200 epochs extends the analysis beyond the early training phase.
C Additional results for Sec. 4
Additional Sec. 4 experiments examine NSGD across architectures, datasets, learning rates, and momentum, revealing faster optimization but dataset- and setting-dependent curvature and generalization effects.
- The SimpleCNN, CIFAR-100, and Fashion-MNIST experiments provide additional NSGD results, with γ = 5.0 diverging in the SimpleCNN and high-learning-rate settings.
- On Fashion-MNIST, final generalization is similar for γ < 1 and γ = 1.0, indicating dataset-dependent behavior.
- The spectral and Frobenius Hessian norms are tracked across different γ values under larger learning rates and momentum.
- NSGD improves training speed with a larger base learning rate or momentum, while final curvature and momentum-based generalization are not significantly affected.These experiments use η = 0.1 and momentum µ = 0.9.
E Comparing NSGD to Newton Method
NSGD leverages curvature but differs from Newton’s method because it does not generally choose a learning rate optimal for the local curvature.
- Unlike Newton’s method, NSGD is driven toward regions whose curvature is too large relative to its typical step.
- NSGD produces sharper loss-surface regions when reducing gradients along the top K directions permits optimization over sharper areas.The illustrated setting uses K = 5 and γ = 0.01.
- For a diagonal Hessian with five eigenvalues equal to 100, two equal to 1, and λ = 0, NSGD with γ = 0.01 and K = 5 matches the Newton update.
F Experiments on sentiment analysis dataset
Experiments on IMDB extend the Hessian-trajectory analysis to sentiment classification and test NSGD under controlled learning-rate, batch-size, and model-capacity settings.
- The IMDB experiments use a simple CNN model and increase selected layer widths by a factor of 2 to encourage overfitting.
- Learning rate and batch size limit the maximum spectral norm along the IMDB training trajectory, although curvature growth takes many epochs.The experiments test η ∈ {0.025, 0.05, 0.1} and S ∈ {2, 8, 32}.
- With η = 0.01, S = 8, and K = 1, NSGD using γ < 1 optimizes significantly faster and initially reaches a sharper region.The γ values tested are 0.1, 1.0, and 5.0, with experiments repeated over 3 seeds.
- Using γ < 1 does not produce a better-generalizing region in the IMDB setting.