[Paper Review] On the Origin of Implicit Regularization in Stochastic Gradient Descent

Samuel Smith, Benoît Dherin+2 more|arXiv (Cornell University)|Jan 28, 2021

Stochastic Gradient Optimization Techniques31 references40 citations

TL;DR

The paper shows that SGD with small finite learning rates behaves like gradient flow on a modified loss that includes an implicit regularizer, and derives this modified loss for SGD by a backward error analysis accounting for minibatch structure.

ABSTRACT

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.

Motivation & Objective

• Motivate the unexplained generalization benefits of SGD with finite learning rates.
• Derive a modified loss for SGD that includes an implicit regularizer penalizing minibatch gradient norms.
• Explain differences between SGD and GD in terms of their implicit regularizers.
• Empirically validate that including the implicit regularizer in the loss can improve test accuracy.

Proposed method

• Use backward error analysis adapted to SGD with minibatch structure to derive a modified loss for the mean SGD iterate after one epoch.
• Show that the modified loss for SGD is C(ω) + (ε/4m) sum_{k=0}^{m-1} ||∇Ĉ_k(ω)||^2, where Ĉ_k are minibatch costs.
• Expand the relation between GD and SGD modified losses to compare gradients and batch-scale effects.
• Compute the expected SGD update after one epoch to identify the bias term arising from minibatch ordering.
• Demonstrate a linear scaling rule between learning rate and batch size within the modified loss framework.
• Provide empirical evidence that explicitly including the implicit regularizer can boost test accuracy.

Experimental results

Research questions

• RQ1 Does SGD with finite learning rates follow a gradient-flow path on a modified loss?
• RQ2 What is the form of the implicit regularizer arising in SGD due to minibatch structure?
• RQ3 How does the implicit regularizer scale with learning rate and batch size?
• RQ4 Can including the implicit regularizer in training loss improve generalization?
• RQ5 How do SGD and GD modified losses differ in their minima and trajectories?

Key findings

• The mean SGD iterate after one epoch stays close to the gradient-flow path of a modified loss.
• The modified SGD loss is detailed as C(ω) + (ε/4m) ∑_{k=0}^{m-1} ||∇Ĉ_k(ω)||^2.
• The implicit regularizer penalizes the mean squared norm of minibatch gradients, with scale ~ ε/(4m).
• If minibatch gradients are diverse, the implicit regularizer term scales with ε/B, explaining batch-size effects.
• Explicitly optimizing the modified loss can improve test accuracy for small learning rates.
• Experiments show the implicit regularizer can enhance test performance when included in the loss.