[Paper Review] The Break-Even Point on Optimization Trajectories of Deep Neural Networks

Stanisław Jastrzȩbski, Maciej Szymczak+5 more|arXiv (Cornell University)|Feb 21, 2020

Stochastic Gradient Optimization Techniques37 references33 citations

TL;DR

The paper investigates how SGD hyperparameters in the early training phase create a break-even point that implicit regularizes loss surface curvature and gradient noise, affecting the full optimization trajectory and potentially generalization.

ABSTRACT

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.

Motivation & Objective

• Motivate the study of the early training phase as critical for final performance in deep nets.
• Define and identify a break-even point on optimization trajectories where SGD stabilizes curvature and gradient noise.
• Hypothesize how initial learning rate and batch size control the covariance of gradients and Hessian conditioning across the trajectory.
• Demonstrate empirical evidence of variance reduction and pre-conditioning effects across multiple architectures and datasets.
• Extend analysis to networks with batch normalization to assess robustness of the findings.

Proposed method

• Model the early-training loss surface as dominated by the top Hessian direction and assume alignment between top Hessian and gradient-covariance eigenvectors.
• Define break-even point as the first SGD step where a stability condition is violated along the top Hessian direction.
• Conjecture that larger initial learning rate or smaller batch size reduce the maximum spectral norms of the gradient covariance and Hessian along the trajectory.
• Compute and analyze the spectral norms of the Hessian and the uncentered gradient covariance during early training to test conjectures.
• Empirically evaluate conjectures on multiple architectures and datasets, including BN and non-BN networks, and on large-scale tasks like BERT on MNLI and DenseNet on ImageNet.

Experimental results

Research questions

• RQ1 Does the early phase of SGD training determine the subsequent optimization trajectory via the break-even point?
• RQ2 How do initial SGD hyperparameters (learning rate, batch size) influence the spectral norms and conditioning of the gradient covariance and Hessian along the trajectory?
• RQ3 Do the proposed effects persist in networks with batch normalization and in large-scale models?
• RQ4 Can larger initial learning rates lead to better-conditioned regions of the loss surface after the break-even point?
• RQ5 What is the relationship between break-even dynamics and potential generalization improvements?

Key findings

• A break-even point exists where SGD dynamics along the most curved loss direction become unstable, after which the trajectory aligns with better-conditioned regions.
• Larger initial learning rates or smaller batch sizes reduce the maximum spectral norms of the gradient covariance and Hessian attained along training, indicating variance reduction and pre-conditioning effects.
• Higher learning rates and smaller batches lead to larger ratios of smallest non-zero eigenvalues to the top eigenvalues for K and H, implying improved conditioning.
• Networks with batch normalization still exhibit the break-even behavior, with high learning rates necessary to achieve improved conditioning compared to non-BN nets.
• The variance reduction and pre-conditioning effects are observed across SimpleCNN, ResNet-32, LSTM, BERT, DenseNet, and in larger scales like MNLI and ImageNet experiments.