Why (and When) does Local SGD Generalize Better than SGD?

Xinran Gu, Kaifeng Lyu, Longbo Huang, Sanjeev Arora

Keywords: local SGD, SDE, regularization, implicit bias, deep learning theory, optimization, distributed training

TL;DR: We derive a Stochastic Differential Equation (SDE) that captures the long-term behavior of Local SGD and provide a theoretical explanation why Local SGD generalizes better than SGD.

Abstract: Local SGD is a communication-efficient variant of SGD for large-scale training, where multiple GPUs perform SGD independently and average the model parameters periodically. It has been recently observed that Local SGD can not only achieve the design goal of reducing the communication overhead but also lead to higher test accuracy than the corresponding SGD baseline (Lin et al., 2020b), though the training regimes for this to happen are still in debate (Ortiz et al., 2021). This paper aims to understand why (and when) Local SGD generalizes better based on Stochastic Differential Equation (SDE) approximation. The main contributions of this paper include (i) the derivation of an SDE that captures the long-term behavior of Local SGD with a small learning rate, after approaching the manifold of minima, (ii) a comparison between the SDEs of Local SGD and SGD, showing that Local SGD induces a stronger drift term that can result in a stronger effect of regularization, e.g., a faster reduction of sharpness, and (iii) empirical evidence validating that having small learning rate and long enough training time enables the generalization improvement over SGD but removing either of the two conditions leads to no improvement.

Community Implementations: [ 2 code implementations](https://www.catalyzex.com/paper/arxiv:2303.01215/code)