PaperlayerView exact revision on arXiv

Source-linked AI summary

On the Origin of Implicit Regularization in Stochastic Gradient Descent

Samuel L. Smith, Benoit Dherin, David G. T. Barrett, Soham De

arXiv:2101.12176v1cs.LGstat.ML

TL;DR

The paper addresses why finite, moderately large learning rates can improve SGD generalization despite lower training loss at smaller rates. It proves that shuffled SGD's mean iterate follows gradient flow on a modified loss with an implicit minibatch-gradient regularizer, and finds that explicitly including this term can enhance test accuracy.

  • Problem

    SGD's higher test accuracy at moderately large learning rates remains poorly understood because convergence bounds do not explain cases where smaller rates achieve lower training losses.

  • Method

    The paper uses backward-error analysis for randomly shuffled SGD, averaging over minibatch orders to derive a modified loss containing the full-batch loss and an implicit regularizer.

  • Results

    The mean SGD iterate stays close to gradient flow on the modified loss, whose regularization scale is proportional to learning rate divided by batch size under small-batch, sufficiently diverse gradients.

  • Takeaways & Limitations

    Explicitly including SGD's implicit regularizer in the training loss can enhance test accuracy when the learning rate is small.

  • Takeaways & Limitations

    The backward-error analysis assumes small finite learning rates and does not capture chaotic or oscillatory dynamics near instability, where the modified-loss Taylor series may not converge.

Abstract

from arXiv · show

For infinitesimal learning rates, stochastic gradient descent (SGD) follows the path of gradient flow on the full batch loss function. However moderately large learning rates can achieve higher test accuracies, and this generalization benefit is not explained by convergence bounds, since the learning rate which maximizes test accuracy is often larger than the learning rate which minimizes training loss. To interpret this phenomenon we prove that for SGD with random shuffling, the mean SGD iterate also stays close to the path of gradient flow if the learning rate is small and finite, but on a modified loss. This modified loss is composed of the original loss function and an implicit regularizer, which penalizes the norms of the minibatch gradients. Under mild assumptions, when the batch size is small the scale of the implicit regularization term is proportional to the ratio of the learning rate to the batch size. We verify empirically that explicitly including the implicit regularizer in the loss can enhance the test accuracy when the learning rate is small.

1 INTRODUCTION

The paper explains finite-learning-rate generalization in SGD by analyzing the mean iterate under random shuffling as gradient flow on a modified loss. This loss adds an implicit regularizer on minibatch-gradient norms, whose scale grows with learning rate and decreases with batch size.

  • Moderately large learning rates can improve SGD test accuracy even when smaller rates achieve lower training losses, a gap not explained by convergence bounds.
  • For small finite learning rates, the mean shuffled-SGD iterate after one epoch stays close to gradient flow on a modified loss eCSGD.The averaging is over all possible minibatch orders, with each training example sampled once per epoch.
  • eCSGD combines the full-batch loss with an implicit regularizer that penalizes the mean squared norms of minibatch gradients.Its scale is proportional to the learning rate.
  • When batch size is small and gradients are sufficiently diverse, the expected SGD regularization magnitude is proportional to the learning-rate-to-batch-size ratio.
  • Explicitly adding the SGD implicit regularizer to the training loss can enhance test accuracy at small learning rates and helps close the gap with larger-rate training.
  • Unlike GD's modified loss, SGD's regularizer penalizes both large full-batch gradients and non-uniform minibatch gradients.The latter term measures squared deviations of minibatch gradients from the full-batch gradient.

2 A BACKWARD ERROR ANALYSIS OF STOCHASTIC GRADIENT DESCENT

The paper applies backward error analysis to finite-learning-rate SGD, showing that the mean iterate follows gradient flow on a modified loss with an implicit regularizer. The analysis explains how minibatch ordering creates this correction and evaluates its empirical relationship to generalization.

  • 2.1 AN INTRODUCTION TO BACKWARD ERROR ANALYSIS: Backward error analysis replaces finite-step dynamics with a continuous modified flow whose path matches discrete updates up to neglected higher-order terms.The modified flow equals the original flow as ϵ → 0 but differs when ϵ is finite.
  • 2.2 BACKWARD ERROR ANALYSIS AND STOCHASTIC GRADIENT DESCENT: A single epoch of SGD is equivalent to a single GD update with learning rate mϵ to first order, while the order-dependent second-order correction produces the SGD bias.The correction is random across minibatch orders, motivating the expectation over shuffles.
  • 2.2 BACKWARD ERROR ANALYSIS AND STOCHASTIC GRADIENT DESCENT: For SGD, averaging over all permutations of fixed minibatches yields a mean iterate that stays close after one epoch to gradient flow on the modified loss eCSGD(ω).The analysis holds batch composition fixed and averages only over minibatch order.
  • REMARKS ON THE ANALYSIS: The analysis assumes mϵ = Nϵ/B is small, so higher-order terms at O(m3ϵ3) may matter at practical learning rates and as dataset size increases.The authors report that experiments suggest the approximation explains most of the finite-learning-rate generalization benefit for Wide-ResNets trained on CIFAR-10.
  • 2.3 AN EMPIRICAL EVALUATION OF THE MODIFIED LOSS: Explicitly adding the implicit regularizer improves test accuracy at small learning rates, while matching optimal regularization and learning-rate coefficients can diminish the large-learning-rate generalization benefit.In Figure 1, λopt = 2−6 equals ϵopt = 2−6; in Figure 2, the test- and training-optimal learning rates coincide at 2−8 with λ = 2−6.

3 IMPLICIT REGULARIZATION AND THE BATCH SIZE

The paper derives how minibatch size shapes SGD’s implicit regularization and shows that, in the small-batch regime, the relevant scale is the learning-rate-to-batch-size ratio. Experiments support corresponding scaling rules for test accuracy, while very large batches approach GD-like behavior.

  • Analysis setup: The modified loss is derived by averaging SGD iterates over minibatch orderings, then extended to randomly shuffled non-overlapping minibatches.The analysis keeps batch composition fixed initially and later studies how the regularizer depends on batch size.
  • Implicit regularization: The SGD implicit regularizer contains a learning-rate term penalizing the full-batch gradient and an ϵ/B term penalizing per-example gradient covariance.The covariance contribution dominates when gradients are sufficiently diverse and Γ(ω)/B ≫ ||∇C(ω)||2.
  • Empirical scaling: Different batch sizes achieve similar test accuracy when ϵ/B is constant and the batch size is not too large.This is observed when training the Wide-ResNet on CIFAR-10 with the original loss.
  • Empirical scaling: Different batch sizes also achieve similar test accuracy when λ/B is constant while training on the explicitly modified loss.The experiment fixes ϵ = 2−9 and varies λ, B, and training for 10 million steps.
  • Large-batch regime: The linear scaling rule is expected to break down for very large batches because the full-batch-gradient regularization term dominates.In that regime, SGD’s bias becomes similar to the GD bias, and the optimal learning rate is expected to become batch-size independent.

4 FINITE LEARNING RATES AND STOCHASTIC DIFFERENTIAL EQUATIONS

The paper argues that SGD’s generalization benefit comes from finite-learning-rate implicit regularization rather than the temperature of an associated SDE. An n-step SGD experiment supports this distinction: increasing repeated updates suppresses regularization and lowers test accuracy.

  • SDE comparison: The SDE analogy assumes independently random minibatch sampling and converges to gradient flow as ϵ → 0, where its temperature T = ϵ/B also vanishes.With one pass through each example per epoch, the paper states that SGD has no first-order noise after one epoch.
  • n-step SGD: n-step SGD applies n sequential updates to one minibatch, with bare learning rate α and effective learning rate ϵ = nα.The next minibatch is sampled only after the n updates are completed.
  • SDE comparison: Standard SGD and n-step SGD agree through first order in ϵ but differ at O(ϵ2) and above.Thus, their dynamics can diverge at finite learning rates even if both share the same small-rate SDE limit.
  • n-step SGD: The implicit regularizer scales with α = ϵ/n, so holding ϵ fixed suppresses regularization as n increases.The case n = 1 recovers the paper’s standard SGD modified loss.
  • Experimental evidence: Smaller n values consistently achieve higher test accuracy because they remain stable at larger bare learning rates.Figure 4 reports this pattern across epoch budgets and also for a fixed number of updates.
  • Experimental evidence: 93.5% test accuracy for n = 1 fell to 88.8% for n = 16 on CIFAR-10, despite equal epoch budgets.The paper attributes the drop to a reduced largest stable bare learning rate and weaker implicit regularization.

5 DISCUSSION

The paper distinguishes its finite-learning-rate explanation from variance-based SGD accounts and identifies a boundary where backward error analysis may fail. At very large learning rates, chaotic or oscillatory dynamics remain outside the analysis.

  • Variance-based explanations attribute improved generalization to greater SGD-iterate variance from larger learning rates or smaller batch sizes.
  • Backward error analysis applies only when the learning rate is small though finite, excluding rates near instability.At such rates, the modified loss is a learning-rate Taylor series that may not converge.
  • In the catapult regime, full-batch GD can oscillate along high-curvature directions, with the loss potentially increasing early in training.The cited work associates this regime with large learning rates close to divergence in wide networks trained on quadratic losses.
  • Whether backward error analysis fully explains small-batch SGD generalization, or chaotic effects also matter, remains open.

A THE EXPECTED NORM OF A MINIBATCH GRADIENT

The analysis centers on deviations of per-example gradients from the full-batch gradient. Their squared norms aggregate into Γ(ω), identified as the trace of the empirical covariance matrix of per-example gradients.

  • Xi measures each per-example gradient's deviation from the full-batch gradient ∇C(ω).The expectation is taken over all possible random shuffles of example indices.
  • Γ(ω) is the sum of squared norms ||∇Ci(ω) − ∇C(ω)||2 across training examples.
  • Γ(ω) equals the trace of the empirical covariance matrix of the per-example gradients.

B A BACKWARD ERROR ANALYSIS FOR N-STEP SGD

For n-step SGD, the analysis expands repeated minibatch updates and matches their expected evolution to a continuous modified flow. The resulting regularization strength depends on the bare learning rate rather than the effective rate when n exceeds one.

  • n-step SGD applies n updates to one minibatch before sampling the next, with effective learning rate ϵ = nα.
  • The n-step update is expanded in powers of α, retaining first- and second-order gradient terms plus O(n3α3).
  • Using ωi+1 = ωi − ϵ∇Ĉi(ωi) + O(ϵ2), the analysis reuses the earlier modified-flow derivation.
  • With ϵ = 2−9, small λ preserves high training accuracy, whereas large λ prevents high training accuracy; large-rate original-loss training has similar losses to small-rate, large-λ modified-loss training.
  • The expected update over batch orderings is matched to continuous modified flow by setting E(ωm) = ω(mϵ).
  • For n = 1, n-step SGD and SGD have the same modified loss; for n > 1, regularization scales with α = ϵ/n rather than ϵ.

C TRAINING LOSSES

The training-loss comparison shows that large learning rates and large regularization coefficients can produce similar training behavior under the original loss.

  • Large learning rates and large regularization coefficients achieve similar training accuracies and similar original losses.The comparison uses the original loss across learning rates and the modified loss at fixed ϵ = 2−9 across regularization coefficients.

D ADDITIONAL RESULTS ON FASHION-MNIST

Fashion-MNIST experiments with an over-parameterized fully connected network examine how learning rate, regularization, batch size, and n-step updates relate to test accuracy.

  • Model and setup: The Fashion-MNIST model is a highly over-parameterized three-layer fully connected MLP that reaches 100% training accuracy when the learning rate is not too large.The network uses width 4096 nonlinear layers with ReLU activations and a final linear softmax layer.
  • Regularization and batch size: Tuning explicit regularization has a similar influence on test accuracy to tuning the learning rate, but their optimal values differ by a factor of 8.The optimal learning rate is ϵ = 2^-5, while regularization reaches a higher maximum at λ = 2^-2.
  • Regularization and batch size: Test accuracy is primarily determined by ϵ/B across batch sizes when training on the original loss.This relationship was observed after 400 epochs.
  • Regularization and batch size: Test accuracy is primarily determined by λ/B across batch sizes when training on the modified loss, although smaller batches achieve slightly higher accuracies.The modified-loss comparison used ϵ = 2^-9 and 1.5 million training steps.
  • n-step SGD: For n-step SGD, test accuracy is governed by the bare learning rate α rather than the effective learning rate ϵ = nα.This pattern appears for both fixed epoch budgets and fixed numbers of parameter updates, contradicting the SDE analogy.
  • n-step SGD: The n-step results suggest that repeatedly applying updates from the same batch need not significantly degrade test accuracy, provided the largest stable bare learning rate is not exceeded.The authors speculate that near-orthogonal example gradients may explain this behavior in the over-parameterized model.
Loading 2101.12176v1…