[Paper Review] Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks

Yuanzhi Li, Colin Wei+1 more|arXiv (Cornell University)|Jul 10, 2019

Stochastic Gradient Optimization Techniques40 references124 citations

TL;DR

The paper provides a theoretical and empirical explanation for why a large initial learning rate with subsequent annealing generalizes better than a small initial learning rate, by analyzing learning order in a two-type-pattern setting and validating with CIFAR-10 patches.

ABSTRACT

Stochastic gradient descent with a large initial learning rate is widely used for training modern neural net architectures. Although a small initial learning rate allows for faster training and better test performance initially, the large learning rate achieves better generalization soon after the learning rate is annealed. Towards explaining this phenomenon, we devise a setting in which we can prove that a two layer network trained with large initial learning rate and annealing provably generalizes better than the same network trained with a small learning rate from the start. The key insight in our analysis is that the order of learning different types of patterns is crucial: because the small learning rate model first memorizes easy-to-generalize, hard-to-fit patterns, it generalizes worse on hard-to-generalize, easier-to-fit patterns than its large learning rate counterpart. This concept translates to a larger-scale setting: we demonstrate that one can add a small patch to CIFAR-10 images that is immediately memorizable by a model with small initial learning rate, but ignored by the model with large learning rate until after annealing. Our experiments show that this causes the small learning rate model's accuracy on unmodified images to suffer, as it relies too much on the patch early on.

Motivation & Objective

• Motivate understanding of why large initial learning rate (LR) followed by annealing improves generalization compared to small LR from start.
• Propose a simple, two-pattern data distribution to study learning order effects in two-layer networks.
• Show that learning order impacts generalization through theoretical results and a regression that mirrors practical observations.

Proposed method

• Define a two-layer ReLU network with a specific U decomposition to separately handle two data components (P: easy-to-generalize, hard-to-fit; Q: easy-to-fit, hard-to-generalize).
• Construct a data distribution with two pattern types and fixed fractions p and q of samples containing each type.
• Use SGD with spherical Gaussian noise and a two-phase learning-rate schedule (large LR then anneal) to analyze learning dynamics.
• Derive informal theorems comparing large-LR with annealing vs. small-LR in terms of pattern learning order and generalization.
• Decompose network outputs into components g_t(x) on the Q-patterns and r_t(x) on the P-patterns to track learning progress.

Experimental results

Research questions

• RQ1 Does a large initial learning rate with annealing generalize better than a small initial learning rate for a two-pattern data distribution?
• RQ2 How does the order in which a network learns different pattern types affect final generalization?
• RQ3 Can the learning-order phenomenon be observed in practical settings via controlled experiments (e.g., memorizable patches on CIFAR-10)?

Key findings

• For a constructed dataset, a two-layer network with large initial LR followed by annealing learns hard-to-generalize, easy-to-fit patterns first and only learns easy-to-generalize, hard-to-fit patterns after annealing.
• A small initial LR quickly learns easy-to-generalize, hard-to-fit patterns and then overfits to them, leading to worse generalization on hard-to-generalize patterns after training.
• The final test error for the large-LR-then-anneal method is smaller by a factor related to p (roughly O(p) in the analysis) than the small-LR method.
• The paper provides a lower bound showing the small-LR method can achieve better training loss but worse test error due to memorization bias on certain pattern components.
• A mitigation strategy inspired by the analysis—adding noise before activations that is reduced at a chosen epoch—can match large-LR guarantees and improve robustness.