# Distributed Sign Momentum with Local Steps for Training Transformers

arXiv: https://arxiv.org/abs/2411.17866 · HTML: https://mlsanity.com/paper/2411.17866/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 49 references that are in the map.
- First posted on arXiv 26 Nov 2024.

## Builds on (most cited first, 30 of 33)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization](https://mlsanity.com/paper/2007.07481/) (2020)
- [signSGD: Compressed Optimisation for Non-Convex Problems](https://mlsanity.com/paper/1802.04434/) (2018)
- [Can Decentralized Algorithms Outperform Centralized Algorithms? A Case Study for Decentralized Parallel Stochastic Gradient Descent](https://mlsanity.com/paper/1705.09056/) (2017)
- [Symbolic Discovery of Optimization Algorithms](https://mlsanity.com/paper/2302.06675/) (2023)
- [Lookahead Optimizer: k steps forward, 1 step back](https://mlsanity.com/paper/1907.08610/) (2019)
- [Error Feedback Fixes SignSGD and other Gradient Compression Schemes](https://mlsanity.com/paper/1901.09847/) (2019)
- [Tighter Theory for Local SGD on Identical and Heterogeneous Data](https://mlsanity.com/paper/1909.04746/) (2019)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [Don't Use Large Mini-Batches, Use Local SGD](https://mlsanity.com/paper/1808.07217/) (2018)
- [On the Linear Speedup Analysis of Communication Efficient Momentum SGD for Distributed Non-Convex Optimization](https://mlsanity.com/paper/1905.03817/) (2019)
- [Stochastic Gradient Push for Distributed Deep Learning](https://mlsanity.com/paper/1811.10792/) (2018)
- [Dissecting Adam: The Sign, Magnitude and Variance of Stochastic Gradients](https://mlsanity.com/paper/1705.07774/) (2017)
- [SlowMo: Improving Communication-Efficient Distributed SGD with Slow Momentum](https://mlsanity.com/paper/1910.00643/) (2019)
- [Noise Is Not the Main Factor Behind the Gap Between SGD and Adam on Transformers, but Sign Descent Might Be](https://mlsanity.com/paper/2304.13960/) (2023)
- [Local SGD: Unified Theory and New Efficient Methods](https://mlsanity.com/paper/2011.02828/) (2020)
- [MARINA: Faster Non-Convex Distributed Learning with Compression](https://mlsanity.com/paper/2102.07845/) (2021)
- [Understanding the Acceleration Phenomenon via High-Resolution Differential Equations](https://mlsanity.com/paper/1810.08907/) (2018)
- [Collaborative Deep Learning in Fixed Topology Networks](https://mlsanity.com/paper/1706.07880/) (2017)
- [Lion Secretly Solves Constrained Optimization: As Lyapunov Predicts](https://mlsanity.com/paper/2310.05898/) (2023)
- [Stochastic Sign Descent Methods: New Algorithms and Better Theory](https://mlsanity.com/paper/1905.12938/) (2019)
- [CSER: Communication-efficient SGD with Error Reset](https://mlsanity.com/paper/2007.13221/) (2020)
- [Communication Efficient Distributed Training with Distributed Lion](https://mlsanity.com/paper/2404.00438/) (2024)
- [How to scale distributed deep learning?](https://mlsanity.com/paper/1611.04581/) (2016)

## Built on by (most cited first, 1 of 1)

- [Sharp High-Probability Rates for Nonlinear SGD under Heavy-Tailed Noise via Symmetrization](https://mlsanity.com/paper/2507.09093/) (2025)

## Abstract (arXiv)

Pre-training Transformer models is resource-intensive, and recent studies have shown that sign momentum is an efficient technique for training large-scale deep learning models, particularly Transformers. However, its application in distributed training remains underexplored. This paper investigates a novel communication-efficient distributed sign momentum method with multiple local steps, to cope with the scenarios where communicating at every step is prohibitive. Our proposed method allows for a broad class of base optimizers for local steps, and uses sign momentum in the global step, where momentum is generated from differences accumulated during local steps. For generic base optimizers, by approximating the sign operator with a randomized version that acts as a continuous analog in expectation, we present a general convergence analysis, which specializes to an O(1/√T) rate for a particular instance. When local step is stochastic gradient descent, we show an optimal O(1/T1/4) rate in terms of ℓ1 gradient norm for nonconvex smooth cost functions. We extensively evaluate our method on the pre-training of various sized GPT-2 models from scratch, and the empirical results show significant improvement compared to other distributed methods with multiple local steps.