Source-linked AI summary
The Break-Even Point on Optimization Trajectories of Deep Neural Networks
Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, Krzysztof Geras
TL;DR
The paper asks how early SGD choices shape the later optimization trajectory, addressing limited understanding of optimization geometry and generalization links. It models and tests a break-even point, finding that large early learning rates reduce gradient noise and improve conditioning, including in batch-normalized networks. The authors identify connecting these effects to generalization as future work.
Problem
The connection between DNN optimization and generalization is not fully understood, including why large initial learning rates and batch normalization affect training and generalization.
Method
The paper develops a simplified break-even-point model and empirically studies how early SGD learning rates and batch sizes affect K and H along training trajectories.
Results
Large early learning rates reduce gradient variance, improve covariance-of-gradients conditioning, and steer training toward better-conditioned loss-surface regions, including with batch normalization.
Takeaways & Limitations
Key properties of the loss surface and gradient noise are strongly influenced by SGD hyperparameters used in the early training phase.
Takeaways & Limitations
The conjectures apply only to learning rates and batch sizes that guarantee training convergence, and the paper leaves closer Hessian investigation for future work.
Abstract
from arXiv · showhide
The early phase of training of deep neural networks is critical for their final performance. In this work, we study how the hyperparameters of stochastic gradient descent (SGD) used in the early phase of training affect the rest of the optimization trajectory. We argue for the existence of the "break-even" point on this trajectory, beyond which the curvature of the loss surface and noise in the gradient are implicitly regularized by SGD. In particular, we demonstrate on multiple classification tasks that using a large learning rate in the initial phase of training reduces the variance of the gradient, and improves the conditioning of the covariance of gradients. These effects are beneficial from the optimization perspective and become visible after the break-even point. Complementing prior work, we also show that using a low learning rate results in bad conditioning of the loss surface even for a neural network with batch normalization layers. In short, our work shows that key properties of the loss surface are strongly influenced by SGD in the early phase of training. We argue that studying the impact of the identified effects on generalization is a promising future direction.
1 INTRODUCTION
The introduction frames early SGD hyperparameters as determinants of later optimization geometry. It proposes a break-even point and argues that larger early learning rates reduce gradient noise and improve conditioning, including with batch normalization.
- Motivation: The paper studies how early-training SGD hyperparameters affect the entire optimization trajectory through gradient covariance K and Hessian H.K captures mini-batch gradient noise, while H captures local loss-surface curvature.
- Empirical Illustration: Figure 1 compares η = 0.01 and η = 0.001 trajectories on CIFAR-10, showing that lower η reaches a region with larger λ1K after the break-even point at matched training accuracy.Points encode UMAP embeddings of test predictions; background color encodes the spectral norm of K.
- The Break-Even Point: The break-even point is where gradient descent begins oscillating along the loss surface’s most curved direction.The paper models early curvature as changing monotonically along the trajectory and presents empirical evidence for this point in actual DNNs.
- Main Conjectures: Early SGD hyperparameters are conjectured to control the spectral norms and conditioning of both K and H after the break-even point.These properties describe gradient noise and local curvature along the subsequent optimization trajectory.
- Main Conjectures: A larger learning rate before the break-even point reduces K’s spectral norm, thereby reducing mini-batch gradient variance.The paper links lower gradient variance to improved convergence speed from the optimization perspective.
- Batch Normalization: The analysis extends to batch-normalized networks, where a large learning rate is necessary to reach better-conditioned loss-surface regions rather than batch normalization alone producing that condition.This complements the paper’s broader claim that early SGD choices shape later optimization geometry.
2 RELATED WORK
Related work connects optimization trajectories, early critical periods, gradient covariance, Hessian curvature, and implicit regularization to DNN optimization and generalization. The paper builds on these lines by focusing on how early SGD dynamics shape the full trajectory.
- Implicit Regularization: Prior studies report that large initial learning rates can improve generalization, while parameter norms and whole optimization trajectories are used to study this connection.Trajectory-based analyses complement bounds based on properties of the final minimum.
- Learning Dynamics: Earlier work identifies critical early learning periods in which example shuffling, regularization, or other interventions can affect the eventual solution.These studies motivate examining the early phase rather than only the final trained model.
- Gradient Covariance and Hessian: The Hessian measures local loss-surface curvature, and its largest eigenvalues can grow rapidly during early training.Prior work also studies how optimization hyperparameters influence Hessian properties at the final minimum.
- Gradient Covariance and Hessian: The Hessian can be decomposed into two terms, with uncentered gradient covariance G dominant at least near the end of training.This connects curvature analyses to gradient-statistics analyses.
- Gradient Covariance and Hessian: The covariance of gradients K captures variation and geometry across examples and has been related to generalization error.Related measures include cosine alignment between gradients computed on different examples.
- Learning Dynamics: The paper’s theoretical analysis is motivated by learning-dynamics studies reporting negative cosine similarity between consecutive SGD steps on classification benchmarks.It also draws on prior dynamical-stability analyses of neural-network training.
3 THE BREAK-EVEN POINT AND THE TWO CONJECTURES ABOUT SGD
The paper models a break-even point where SGD stability changes along the most curved direction, then conjectures that early learning rate and batch size shape curvature, gradient-noise variance, and conditioning along the trajectory.
- The Break-Even Point: The theoretical model assumes a quadratic loss along the top Hessian eigenvector, monotonic curvature behavior, and a batch size much smaller than the training set.These assumptions are used to generate empirically testable predictions about early training.
- The Break-Even Point: A larger learning rate or smaller batch size is predicted to reach the break-even point earlier and determine the Hessian and gradient-covariance properties there.The prediction is made for trajectories initialized in a stable region and is formally stated for the model before empirical testing.
- The Break-Even Point: The break-even point is the first trajectory iteration where the stability condition becomes equality, and it is unique by definition.The model uses this point to separate the early phase from subsequent optimization behavior.
- The Two Conjectures About SGD: After the break-even point, the model predicts that SGD avoids regions where the largest eigenvalues of H or K exceed their break-even values.This post-break-even restriction follows from maintaining stability along the leading Hessian direction.
- The Two Conjectures About SGD: The variance-reduction conjecture states that larger learning rates or smaller batches yield smaller maximum spectral values of the gradient covariance K.Lower λ1(K) can reduce the L2 distance between full-batch and mini-batch gradients.
- The Two Conjectures About SGD: The pre-conditioning conjecture states that larger learning rates or smaller batches increase the smallest non-zero eigenvalues of K and H while reducing their maximum traces.The conjectures apply only to learning rates and batch sizes that guarantee convergence, and K has at most N − 1 non-zero eigenvalues.
4 EXPERIMENTS
Experiments support the break-even framework across architectures and classification tasks: early SGD hyperparameters shape gradient covariance and loss-surface curvature throughout training. Higher learning rates generally reduce the largest covariance eigenvalue and improve conditioning, including in batch-normalized networks.
- Experimental setup: Experiments span CIFAR-10, IMDB, ImageNet, and MNLI using CNN, ResNet, LSTM, DenseNet, and BERT architectures.Additional results cover MLPs, momentum, and learning-rate decay.
- Covariance and curvature: Lower learning rates reach regions with larger λ1 K and worse conditioning of both K and H, including more spectral outliers.These effects are observed at comparable training accuracy and when the spectra attain their largest values.
- Covariance and curvature: Higher learning rates reduce the maximum λ1 K along the optimization trajectory, while larger batch sizes increase it.The reported trend supports variance reduction and pre-conditioning effects of SGD.
- Covariance and curvature: For ResNet-32 and SimpleCNN, smaller learning rates or larger batch sizes coincide with larger maximum λ1 H.The Hessian trend complements the covariance results.
- Broader validation: The same conjectured learning-rate effects hold for BERT on MNLI, DenseNet on ImageNet, and networks with batch normalization.In batch-normalized networks, a high learning rate is necessary to improve loss-surface conditioning relative to the corresponding network without batch normalization.
5 CONCLUSION
The paper argues that SGD's early-phase hyperparameters shape the spectral norms and conditioning of gradient covariance K and Hessian H after the break-even point. A large initial learning rate steers training toward better-conditioned loss-surface regions, while links to generalization remain future work.
- 5 CONCLUSION: Theoretical and empirical analysis identifies a break-even point after which early SGD hyperparameters control the spectral norms and conditioning of K and H.K describes mini-batch gradient noise, while H describes local loss-surface curvature.
- 5 CONCLUSION: A large initial learning rate steers training toward better-conditioned regions of the loss surface, benefiting optimization.
- 5 CONCLUSION: The paper proposes connecting these SGD-controlled measures to deep-network generalization as a future research direction.
- 5 CONCLUSION: Figure 8 visualizes conditioning of K, Tr(K), and the spectral norm of H across analogous training trajectories.
B PROOFS
The proofs formalize how SGD stability along the most curved direction determines the break-even point and how learning rate and batch size affect when it is reached. They also establish corresponding curvature bounds for stable and unstable initialization settings.
- B PROOFS: Theorem 1 defines the break-even point through the first trajectory iteration where the stability condition becomes equal to 1.
- B PROOFS: The derived threshold values of λ1 H are smaller for a larger η or a smaller S.
- B PROOFS: The proof rewrites the stability inequality using gradient variance, with s(t) defined through the variance of Hessian-related quantities.
- B PROOFS: Theorem 2 extends the argument to training initialized in an unstable region along the leading Hessian eigenvector.
- B PROOFS: The proof compares trajectories with different learning rates or batch sizes to derive their relative break-even behavior.
C APPROXIMATING THE EIGENSPACE OF K AND H
This section develops computational approximations for the eigenspace of K using mini-batch gradients and Gram matrices, avoiding the infeasible quadratic cost of explicitly forming K. It also describes related Hessian estimates and a correlation check between full-data and mini-batch spectra.
- C APPROXIMATING THE EIGENSPACE OF K AND H: A small subset of CIFAR-10 suffices to approximate the largest Hessian eigenvalues, and the experiments use approximately 5% of the dataset.
- C APPROXIMATING THE EIGENSPACE OF K AND H: The trace Tr(K) is related to average gradient variance through the normalized trace of the gradient covariance.
- C APPROXIMATING THE EIGENSPACE OF K AND H: Explicit eigenspace computation for K is infeasible at realistic network sizes because its cost is quadratic in the number of parameters.
- C APPROXIMATING THE EIGENSPACE OF K AND H: The method samples L mini-batch gradients of size M, forms the Gram matrix K_M, and computes its eigenspace with SVD.
- C APPROXIMATING THE EIGENSPACE OF K AND H: On CIFAR-10, the largest eigenvalue of K computed with M = 1 and M = 128 shows a strong correlation when using the whole training set and 10% subsampling.
D EXPERIMENTAL DETAILS FOR SEC. 4.2
The experiments cover convolutional, recurrent, multilayer-perceptron, transformer, and DenseNet models across CIFAR-10, IMDB, FashionMNIST, MNLI, and ImageNet. Training schedules, batch sizes, learning-rate settings, repetitions, and preprocessing are specified for each setup.
- D EXPERIMENTAL DETAILS FOR SEC. 4.2: ResNet-32 is trained on CIFAR-10 for 200 epochs with batch size 128 and weight decay 0.0001.
- D EXPERIMENTAL DETAILS FOR SEC. 4.2: SimpleCNN on CIFAR-10 uses four convolutional layers with 32 and 64 filters, max pooling, and ReLU activations.
- D EXPERIMENTAL DETAILS FOR SEC. 4.2: BERT-base is trained on MNLI for 20 epochs with batch size 32 and three initialization and shuffling seeds.
- D EXPERIMENTAL DETAILS FOR SEC. 4.2: The FashionMNIST MLP has hidden-layer sizes 300 and 100, batch size 64, and a 200-epoch training schedule.
- D EXPERIMENTAL DETAILS FOR SEC. 4.2: The IMDB LSTM uses a 20,000-word vocabulary, sequence length 80, and 100 training epochs, with settings varied by learning rate or batch size.
- D EXPERIMENTAL DETAILS FOR SEC. 4.2: DenseNet-121 is trained on ImageNet for 10 epochs with batch size 32, centered and normalized images, and no data augmentation.
E ADDITIONAL EXPERIMENTS FOR SEC. 4.2.
Additional experiments extend the paper’s findings across architectures, datasets, and learning-rate or batch-size settings. The FashionMNIST MLP experiments report that all conclusions carry over to that setting.
- Additional experiments: Additional metrics cover SimpleCNN and ResNet-32 on CIFAR-10, including training accuracy, validation accuracy, and conditioning-related quantities.The reported metrics include Tr(K) and λ1 H across the corresponding figures.
- Additional experiments: LSTM experiments on IMDB report accuracy, validation accuracy, and conditioning-related metrics under different learning rates and batch sizes.The figures include Tr(K) and λ1 among the additional metrics.
- Additional experiments: BERT on MNLI and DenseNet on ImageNet are evaluated using training accuracy, validation accuracy, and Tr(K).These experiments broaden the reported measurements beyond CIFAR-10 models.
- Additional experiments: The MLP experiments on FashionMNIST report that all conclusions carry over to this setting.Results are presented through separate metric and accuracy figures.
F ADDITIONAL EXPERIMENTS FOR SGD WITH MOMENTUM
The momentum experiments test whether the paper’s conjectures extend beyond ordinary SGD. They report analogous effects for varying momentum and confirm both conjectures with momentum held constant.
- Varying momentum: Higher momentum has an analogous effect to higher learning rate or smaller SGD batch size.The experiment varies β over 0.1, 0.5, and 0.9 using SimpleCNN on CIFAR-10.
- Fixed momentum: Both conjectures generalize when momentum is fixed at 0.9 and learning rates are varied.The learning rate 0.1 diverged, so only 0.01 and 0.001 were included.
- Fixed momentum: The fixed-momentum experiment therefore reports results only for the two non-divergent learning rates, 0.01 and 0.001.The setting uses SimpleCNN on CIFAR-10.
G ADDITIONAL EXPERIMENTS FOR SGD WITH LEARNING RATE DECAY
Learning-rate decay experiments examine whether the conjectures persist under a schedule. The reported results indicate that both conjectures generalize when the learning rate is divided by 10 after the 100th epoch.
- Learning-rate decay: Both conjectures generalize to a learning-rate schedule in which the learning rate is divided by 10 after the 100th epoch.The experiment uses SimpleCNN on CIFAR-10.
- Learning-rate decay: Changing the learning rate under the schedule does not change the relative ordering of the maximum λ1 K.The corresponding results are reported in Figures 24 and 25.
- Learning-rate decay: Figure 25 reports training and validation accuracy for the learning-rate-decay experiment.Figure 24 reports the associated variance-reduction and pre-conditioning metrics.
H ADDITIONAL EXPERIMENTS FOR SIMPLECNN-BN
Additional SimpleCNN-BN experiments vary batch size to further examine conditioning of the loss surface. Overall, decreasing batch size has effects similar to increasing learning rate, with an exception at batch size 10.
- SimpleCNN-BN: Decreasing batch size has a similar effect on the studied metrics as increasing the learning rate.The experiments repeat the Section 4.3 setting while varying batch size.
- SimpleCNN-BN: For the smallest batch size of 10, the maximum values of ∥g∥, ∥g5∥, and λ∗ K are smaller than with larger batch sizes.This is the stated exception to the overall batch-size pattern.
- SimpleCNN-BN: Figure 26 summarizes conditioning-related metrics for SimpleCNN-BN across different batch sizes.The figure concerns the evolution of various metrics quantifying loss-surface conditioning.