arXiv · 2402.18824
Batch size invariant Adam
Abstract
We propose a batch size invariant version of Adam, for use in large-scale, distributed settings, in which the mini-batch is divided into micro-batches which are distributed among worker nodes. For the v term, standard Adam first computes the average over micro-batch gradients, then squares, while in the batch size invariant Adam proposed here, we first square the micro-batch gradients, then average. Previous work (e.g. Malladi et al. 2022) used an alternative approach that involved a square-root scaling of the learning rate, but this approach requires strong assumptions to work; in particular that the gradient variance dominates the square of the expected gradient. In contrast, the approach proposed here gives batch size invariance without this assumption. We confirm that in practice our scheme gives batch size invariance in a much larger range of scenarios than the previous approach.
Explore related subjects
Keep this discovery
Explore connections, maps & timelines
Xi Wang, Laurence Aitchison. 2024-02-29. Batch size invariant Adam. https://arxiv.org/abs/2402.18824
Cite the original work for its findings. Save a collection to share your selection of sources.