The Effect of Mini-Batch Noise on the Implicit Bias of Adam

TL;DR

This paper presents a theoretical framework analyzing how mini-batch noise affects Adam optimizer's implicit bias towards sharper or flatter minima, impacting generalization in multi-epoch training.

Contribution

It reveals how hyperparameters  and  influence implicit bias depending on batch size, guiding better hyperparameter choices for different batch regimes.

Findings

Higher  increases anti-regularization at large batch sizes, hurting generalization.

At smaller batch sizes, the dependence of anti-regularization on  reverses.

The default =0.9, =0.999 is optimal for small batches; larger batches benefit from  closer to .

Abstract

With limited high-quality data and growing compute, multi-epoch training is gaining back its importance across sub-areas of deep learning. Adam(W), versions of which are go-to optimizers for many tasks such as next token prediction, has two momentum hyperparameters $(\beta_1, \beta_2)$ controlling memory and one very important hyperparameter, batch size, controlling (in particular) the amount mini-batch noise. We introduce a theoretical framework to understand how mini-batch noise influences the implicit bias of memory in Adam (depending on $\beta_1$, $\beta_2$) towards sharper or flatter regions of the loss landscape, which is commonly observed to correlate with the generalization gap in multi-epoch training. We find that in the case of large batch sizes, higher $\beta_2$ increases the magnitude of anti-regularization by memory (hurting generalization), but as the batch size becomes smaller, the dependence of (anti-)regulariation on $\beta_2$ is reversed. A similar monotonicity shift (in the opposite direction) happens in $\beta_1$. In particular, the commonly "default" pair $(\beta_1, \beta_2) = (0.9, 0.999)$ is a good choice if batches are small; for larger batches, in many settings moving $\beta_1$ closer to $\beta_2$ is much better in terms of validation accuracy in multi-epoch training. Moreover, our theoretical derivations connect the scale of the batch size at which the shift happens to the scale of the critical batch size. We illustrate this effect in experiments with small-scale data in the about-to-overfit regime.