PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks

Yuanzhi Li, Colin Wei, Tengyu Ma

arXiv:1907.04595v2cs.LGstat.ML

TL;DR

The paper asks why small initial learning rates can optimize training faster yet generalize worse after annealing. It analyzes learning order in a two-layer network, proves a generalization advantage for large initial learning rates with annealing, and validates the mechanism using modified CIFAR-10 data. The analysis also motivates activation-noise mitigation.

  • Problem

    Existing explanations leave limited understanding of why large-learning-rate models can overtake small-learning-rate models in generalization only after annealing.

  • Method

    The paper studies learning order in a two-layer neural network, using a constructed distribution with distinct easy-to-generalize and hard-to-generalize pattern types.

  • Results

    Large initial learning rates with annealing provably generalize better than small learning rates in the constructed setting, with the pattern supported by modified CIFAR-10 experiments.

  • Takeaways & Limitations

    Learning-rate scale and SGD noise can influence which patterns are learned first, and activation noise provides a theoretically supported and empirically demonstrated mitigation strategy.

  • Takeaways & Limitations

    The theoretical analysis uses a simplified two-layer setting and spherical Gaussian SGD noise, leaving the corresponding real-world minibatch-noise analysis for future work.

Abstract

from arXiv · show

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.

1 Introduction

The paper explains why large initial learning rates can improve generalization after annealing through the order in which neural networks learn different pattern types. Its theory and experiments connect this learning-order effect to non-convexity, SGD noise, and performance on modified CIFAR-10.

  • Empirical phenomenon: In Figure 1, the small-learning-rate model leads before annealing, whereas the large-learning-rate model visibly outperforms it after annealing.The figure compares CIFAR-10 train and validation accuracy for WideResNet models initialized with learning rates 0.1 and 0.01.
  • Motivation and contribution: The paper proposes learning order—the rates at which a model learns different example types—as a theoretical explanation for learning-rate-dependent generalization.This training-time property contrasts with analyses focused mainly on post-training classifier complexity or output stability.
  • Theoretical results: Large learning rates with annealing provably generalize better in a constructed two-layer-network setting by initially learning hard-to-generalize, easy-to-fit patterns.The model uses an effective sample size of 0.8N for these patterns before annealing.
  • Theoretical results: After annealing, the large-learning-rate model learns easy-to-generalize, hard-to-fit patterns correctly from the remaining 0.2N examples.This learning sequence preserves performance on both pattern types in the constructed setting.
  • Theoretical results: Small initial learning rates quickly learn easy-to-generalize, hard-to-fit patterns, but use only 0.2N examples to learn hard-to-generalize patterns and perform relatively worse on them at test time.Examples containing both pattern types are initially ignored for the hard-to-generalize patterns.
  • Empirical validation: The CIFAR-10 experiment adds an artificial memorizable patch, which small-learning-rate models exploit early while large-learning-rate models defer it until after annealing.The authors report that this validates the theoretical learning-order effect and motivates injecting activation noise as a mitigation strategy.

2 Setup and Notations

The setup defines a two-component data distribution, a ReLU two-layer network, and the training objective used to study learning order and generalization.

  • Data distribution: The data distribution combines P, a linearly separable but high-noise component, with Q, a low-noise component supported on three directions and readily memorized by neural networks.P requires relatively large sample complexity, while Q is not linearly separable and requires nonlinear modeling.
  • Data distribution: Q assigns negative labels to z − ζ and z + ζ and a positive label to z, with ζ orthogonal to z and much smaller in norm.The construction uses x2 = α(z + bζ), where b is uniformly chosen from {−1, 1}.
  • Data distribution: Training examples are drawn independently from the distribution, with p and q denoting empirical fractions associated with the two component equations.The dataset contains N examples {(x^(1), y^(1)), …, (x^(N), y^(N))}.
  • Two-layer neural network model: The model is a ReLU two-layer network with first-layer weights U and second-layer weights u, whose output is expressed using activation patterns and element-wise products.The first layer has m × 2d weights, and the paper also decomposes U into matrices operating on the P and Q coordinates.
  • Training objective: The paper trains the network with logistic loss, averages losses over the dataset or subsets, and adds an ℓ2 regularizer on U while fixing the randomly initialized second-layer weights.The regularized objective is written as the training loss plus λ/2 times the squared Frobenius norm of U.

3 Main Results

The paper compares a noisy SGD schedule with a large initial learning rate and annealing against fixed small-rate SGD, proving that their learning order produces different generalization outcomes.

  • Training algorithms: The analyzed algorithm uses SGD with spherical Gaussian noise, treating this noise distribution as a simplification of minibatch SGD noise in real-world networks.The paper explicitly leaves analysis of the exact minibatch setting to future work.
  • Training algorithms: Algorithm L-S uses η1 initially, anneals to the much smaller η2 after a loss threshold, and then continues until a second training-loss target is reached.Algorithm S instead uses the fixed learning rate η2 until its stopping criterion.
  • Main results: The L-S algorithm achieves final generalization error smaller by a factor O(κ) = O(p^1/2) than the fixed-small-rate algorithm.The theorem statements establish the comparison under the paper’s initialization, noise, over-parameterization, and parameter-scaling assumptions.
  • Main results: The L-S model first learns P while large-rate noise prevents meaningful learning of Q, then memorizes Q after annealing.Its test performance is analyzed through the separate contributions of the P and Q components.
  • Main results: The small-rate algorithm quickly memorizes Q, learns P from too few effective examples, and consequently has test error and test loss at least Ω(p).Its classifier obtains only a small margin on P-only examples, producing the stated lower bound.

4 Characterization of Algorithm 1 (L-S)

The analysis of L-S separates its initial and post-annealing phases: the large rate learns P first, while the smaller post-annealing rate enables memorization of Q and the resulting generalization guarantee.

  • Phase I: initial learning rate η1: During Phase I, L-S reaches its annealing loss threshold within a bounded time while accumulating only a limited gradient signal.The phase is analyzed using changing neural tangent kernels induced by noisy activation patterns.
  • Phase I: initial learning rate η1: Before annealing, the Q component remains essentially unlearned because the large learning rate creates too much noise for meaningful memorization.This behavior holds through the analyzed 1/(η1λ)-scale time horizon.
  • Phase II: after annealing: After annealing to η2, the model converges further and memorizes the Q component.The post-annealing lemmas bound convergence and training losses on the relevant example subsets.
  • Generalization: Low loss on Q-containing examples supports generalization on Q, while the P classifier generalizes using the approximately (1 − q)N examples with P components.The P-side test-error bound uses standard Rademacher complexity tools.

5 Characterization of Algorithm 2 (S)

The fixed-small-rate algorithm reaches low training loss by memorizing Q early, leaving insufficient learning of P and producing a lower-bounded test error.

  • Training loss convergence: The small-rate algorithm converges to small training error before its P weights become large enough to produce a substantial margin on examples with x2 = 0.This convergence occurs within the time scale given by Lemma 5.1.
  • Lower bound on generalization error: The resulting solution has both test error and test loss at least Ω(p), despite attaining the prescribed small training loss.This is the lower bound established for Algorithm S.
  • Lower bound on generalization error: Under a small learning rate, early Q memorization suppresses learning from the P component of examples containing both patterns.The accumulated gradients on the relevant subset M2 remain small.
  • Lower bound on generalization error: The P classifier therefore overfits to only pN examples, which are too few to learn the d-dimensional P distribution under the paper’s parameter choice.The resulting margin is poor on P samples, so predictions are strongly influenced by noise.

6 Proof Sketches

The proof sketches decouple signal from noise and analyze optimization through changing convex functions that share a common target. They then show how learning-rate annealing alters activation stability and pattern learning.

  • 6.1 Proof Sketches: The analysis abstracts the loss as a function of the signal portion while decoupling noise in the weights.This decomposition supports the subsequent analysis of network dynamics.
  • 6.1 Proof Sketches: A common target solution is good for all changing kernel functions K_t, enabling convergence analysis across iterations.The target is used in the proof of Lemma 4.1.
  • 6.1 Proof Sketches: Gradient descent on differentiable convex functions can converge when the functions change over time but share the same solution.Theorem 6.2 supplies the general optimization statement used by the proofs.
  • 6.1 Proof Sketches: Before annealing, activation noise prevents the model from learning the Q component because it requires a specific activation-pattern configuration.After annealing, activations change little, allowing the Q component to be learned without forgetting the previously learned P component.
  • 6.1 Proof Sketches: For the small-learning-rate analysis, rapid loss convergence limits the signal that P components contribute to the weights.The proof decomposes the relevant weight component into data-span and remainder terms before establishing the generalization bound.
  • 6.1 Proof Sketches: The resulting classifier can be too heavily influenced by noise in x_1, preventing small generalization error.This completes the argument for the small-learning-rate theorem.

7 Experiments

The experiments test activation-noise regularization and learning-order effects on WideResNet16 using patch-augmented CIFAR-10. Small learning rates memorize the patch early, while annealed methods delay patch memorization and retain higher clean-image accuracy.

  • 7 Experiments: Gaussian noise before every activation layer is annealed over time as a proposed regularizer for small learning rates.The experiment uses a WideResNet16 on CIFAR-10 without data augmentation.
  • 7 Experiments: A memorizable 7 × 7 pixel patch is added to a subset of CIFAR-10 images, with 20% unpatched, 16% patch-only, and 64% combined images.The patches have low variation but are not easily separable.
  • 7 Experiments: The small-learning-rate method picks up the patch very quickly, whereas the other methods memorize it only after annealing.Figure 3 compares clean validation images with patch-only images across epochs.
  • 7 Experiments: 70% clean-image validation accuracy after patch training contrasts with around 83% for a small-learning-rate model trained on the full dataset.The authors interpret this gap as evidence that the patched model learns CIFAR images from only a small fraction of the available data.

8 Conclusion

The paper concludes that learning order is central to generalization and that learning-rate scale can alter which patterns are learned first. It proposes activation noise as a mitigation, while leaving improved regularization algorithms for future work.

  • 8 Conclusion: The order in which a neural network fits different pattern types plays a crucial role in generalization.The paper supports this claim with a theoretical construction and modified CIFAR-10 experiments.
  • 8 Conclusion: Larger learning rates or more SGD noise bias the model toward generalizing kernels rather than memorizing kernels.This is presented as the mechanism revealed by the analysis.
  • 8 Conclusion: The CIFAR-10 experiments confirm that learning-rate scale can influence both learning order and generalization.The empirical setting uses artificially modified CIFAR-10 data.
  • 8 Conclusion: Injecting noise before activations is proposed as a mitigation strategy that works theoretically for the paper's construction and empirically.The authors identify designing better algorithms for regularizing learning order as future work.

B.1 Proof of Theorem 3.4

The proof of Theorem 3.4 argues that the large-learning-rate model learns one component while preserving the conditions needed to generalize on another. The argument combines margin properties, norm bounds, and post-annealing optimization.

  • B.1 Proof of Theorem 3.4: Low training error on the first mixture component forces the classifier to produce a large output on x_2 relative to its norm.This follows from the margin required for low loss.
  • B.1 Proof of Theorem 3.4: Small-norm x_2 examples force larger classifier outputs because the norm of x_2 is roughly uniform on [0, 1].The proof uses these examples to establish the required signal behavior.
  • B.1 Proof of Theorem 3.4: The proof checks that all examples in M2 are classified correctly and that the resulting classifier generalizes well on the corresponding test distribution.This is the stated goal of the Theorem 3.4 argument.
  • B.1 Proof of Theorem 3.4: The analysis bounds the network response on x_1 using the coefficient and noise decomposition together with a Frobenius-norm bound on the weights.The resulting bound is shown to remain controlled under the stated parameter choices.
  • B.1 Proof of Theorem 3.4: The combined classifier retains low loss on the relevant examples because the added component does not increase loss when its output has the appropriate sign.The proof combines this observation with earlier lemmas to complete the bound.

B.2 Proof of Theorem 3.5

The proof shows that small-learning-rate training can produce predictions dominated by a noisy low-rank component formed from examples in ¯ M2. This misalignment with the ground-truth direction yields a nontrivial test error.

  • Proof mechanism: A noisy vector α from the low-rank span of ¯ M2 examples can heavily influence the network’s predictions.With high probability, α does not align well with the ground-truth vector w⋆.
  • Probability argument: For randomly chosen x1, the proof establishes events occurring with probability Ω(1) that support the lower-bound argument.The argument uses the decomposition of x1 into a component along w⋆ and an independent orthogonal component.
  • Error construction: When the network output has the wrong sign on either x1 or −x1, one of the corresponding examples is misclassified.The sign comparison rt(x1) ≥ rt(−x1) holds with probability Ω(1) in the stated case.
  • Error construction: The resulting test error is at least Ω(p).This follows from the nonalignment of the learned prediction component with w⋆ and the paired sign argument.

C.4 Proof of Lemma 4.3

The proof constructs a target solution that preserves progress on ¯ M1 while enabling the network to learn patterns from Q after learning-rate annealing. It then uses convergence and approximation arguments to show that the algorithm reaches this target.

  • Target construction: After annealing, the network is shown to learn patterns from Q by constructing a common target for subsequent time steps.The convergence theorem is applied to this target after establishing its suitability.
  • Target construction: A target solution exists with good accuracy on ¯ M1 and without unlearning the network’s progress on M1.This is stated as the purpose of the proposition used in the proof.
  • Prediction control: Under a second-difference condition on gt and a loss bound, the predictions on z − ζ, z, and z + ζ are all O(·).The condition is |gt(z + ζ) + gt(z − ζ) − 2gt(z)| ≤ δ, with the loss bound specified in Proposition C.9.
  • Convergence control: The proof bounds activation-pattern changes and related network terms to control the evolution from the initial state toward the target.It uses bounds on indicator changes, noise terms, and the network output during the post-annealing phase.
  • Convergence control: Theorem 6.2 is then used to conclude convergence to the constructed target within the stated iteration bound.The proof combines the target’s properties with the convergence theorem and the initial-output bound.

D.2 Proof of Lemma 5.2

The proof bounds how much the network can learn from the P component of examples in M2 under the small-learning-rate regime. The key step is showing that large values of ρt cannot persist for many iterations.

  • Gradient control: When ρt is large, the total gradient norm must also be large.This relationship is established as the starting point for controlling the duration of large-ρt phases.
  • Iteration bound: The proof bounds the number of iterations for which ρt can be large.The resulting iteration count depends on the stated parameters, including r, ε′, and η2.
  • Learning limitation: Because ρt cannot remain large for many iterations, Wt obtains limited signal from the P component of examples in M2.This is the stated conclusion of Lemma 5.2.
  • Learning limitation: The decomposition of Wt separates contributions across prior gradient steps and is combined with Lemma 5.2 to derive the needed bound.The proof uses the decomposition and spectral-norm estimates for the data matrix.

E.2 Extension to two layer convolution network

The convolutional extension disentangles channel updates so that each channel behaves approximately separately. The resulting theorem transfers the two-layer conclusions with modified patch and dimension parameters.

  • Model extension: The convolutional model uses k channels, patch size d, and stride d/k, with channel-specific second-layer weights and shared hidden weights.The extension’s construction introduces channel patches while retaining shared hidden parameters.
  • Theoretical transfer: The convolutional setup assumes that z and ζ are supported only on the last d/k coordinates.This simplifying assumption is used in stating the main theorem.
  • Theoretical transfer: The general convolutional theorem replaces r by r/k and d by dk in Theorems 3.4 and 3.5 and their assumption.This is the stated conclusion of Theorem E.2.
  • Channel disentanglement: Convolutional channels can be analyzed as updating almost separately because a gradient from one channel has negligible later effect on other channels.This follows from the cross-channel disentangling lemma.

G.2 Additional Details on Patch-Augmented CIFAR-10

The patch-augmented CIFAR-10 experiment fixes a randomized clean/patch data split, constructs centered class-dependent patches, and compares three training configurations over 60 epochs.

  • Data split: The training split contains 10,000 clean images and 40,000 patched images, with patch-only and mixed examples generated during updates.The implementation yields approximately 20% clean, 16% patch-only, and 64% mixed updates.
  • Patch construction: Patches are centered and generated from a shared random vector plus class-dependent perturbations, with non-patch coordinates set to zero.Mixed images add the generated patch to a CIFAR example; patch-only images use class-dependent scalar multiples.
  • Training configurations: The experiment compares large-learning-rate annealing, fixed small learning rate, and fixed small learning rate with annealed activation noise.The large rate is 0.1 until annealing to 0.004 at epoch 30; the small rate is fixed at 0.004, and noise decays after epoch 30.
  • Evaluation: All models train for 60 epochs from the same dataset and patch assignments, with validation reported separately for patch-augmented and clean data.Table 2 distinguishes a mixed validation set containing patches from a clean validation set without patches.
  • Results: On clean images, the small-learning-rate patch model reaches 69.89% validation accuracy, while the large-learning-rate model reaches 87.61%.The small-learning-rate result is closer to its 65% accuracy on 10,000 clean images than to 83% on the full clean dataset; the large-learning-rate result is close to 90%.
Loading 1907.04595v2…