Reproduce “SlowMo: Improving Communication-Efficient Distributed SGD with Slow Momentum” on OpenResearch

About this paper

Jianyu Wang, Vinayak Tantia, Nicolas Ballas, Michael Rabbat

The SLOWMO framework enhances communication-efficient distributed optimization algorithms by integrating a slow momentum mechanism, consistently improving training accuracy and generalization across tasks like image classification and neural machine translation with minimal additional communication overhead. It also provides the first theoretical convergence guarantees for related methods such as BMUF and Lookahead.

Key insights

• Integrating a 'slow momentum' update, performed only after periodic global synchronizations, effectively stabilizes and accelerates training for communication-efficient algorithms.
• The performance benefits of slow momentum largely stem from the update mechanism itself rather than strict global synchronization of the momentum buffer at every outer step, as indicated by a variant without exact averaging.
• The SLOWMO framework unifies and provides the first theoretical convergence guarantees for previously empirical nested-loop optimizers like Block-wise Model Update Filtering (BMUF) and the Lookahead optimizer in non-convex settings.

Summary by alphaXiv.