On the Origin of Implicit Regularization in Stochastic Gradient Descent

TL;DR

It is proved 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.

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.

Figures (7)

• Figure 1: (a) Explicitly including the implicit regularizer in the loss improves the test accuracy when training with small learning rates. (b) The optimal regularization coefficient $\lambda_{opt} = 2^{-6}$ is equal to the optimal learning rate $\epsilon_{opt} = 2^{-6}$. (c) Increasing either the learning rate $\epsilon$ or the regularization coefficient $\lambda$ reduces the value of the implicit regularization term $C_{reg}(\omega)$ at the end of training.
• Figure 2: (a) There is a clear generalization benefit to large learning rates when training on the original loss $C(\omega)$ with $\lambda=0$. (b) When we include the implicit regularizer explicitly in $C_{mod}(\omega)$ and set $\lambda = 2^{-6}$, the generalization benefit of large learning rates is diminished.
• Figure 3: (a) Different batch sizes achieve the same test accuracy if the ratio of the learning rate to the batch size $(\epsilon/B)$ is constant and $B$ is not too large. (b) The test accuracy is independent of the batch size if the ratio of the regularization coefficient to the batch size $(\lambda/B)$ is constant.
• Figure 4: (a) When training for 400 epochs, smaller values of $n$ are stable at larger bare learning rates $\alpha$, and this enables them to achieve higher test accuracies. (b) Similar conclusions hold when training for a fixed number of updates. (c) We show the test accuracy at the optimal learning rate for a range of epoch budgets. We find that smaller values of $n$ consistently achieve higher test accuracy.
• Figure 5: (a) When training on the original loss $C(\omega)$, the training accuracy achieved within 6400 epochs is maximized when the learning rate $\epsilon = 2^{-9}$. Smaller learning rates have not yet converged, while larger learning rates reach a plateau. When training on the modified loss $C_{mod}(\omega)$ with fixed learning rate $\epsilon = 2^{-9}$ and regularization coefficient $\lambda$, we achieve high training accuracies when $\lambda$ is small, but are unable to achieve high training accuracies when $\lambda$ is large. (b) We plot the value of the original loss $C(\omega)$ at the end of training. Remarkably, the training losses obtained when training on the original loss with a large learning rate are similar to the training losses achieved when training on the modified loss with small learning rate ($\epsilon = 2^{-9}$) and a large regularization coefficient $\lambda$.