# Smoothing DiLoCo with Primal Averaging for Faster Training of LLMs

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 45 references that are in the map.
- First posted on arXiv 18 Dec 2025.

## Builds on (most cited first, 28 of 28)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [On First-Order Meta-Learning Algorithms](https://mlsanity.com/paper/1803.02999/) (2018)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Lookahead Optimizer: k steps forward, 1 step back](https://mlsanity.com/paper/1907.08610/) (2019)
- [Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer](https://mlsanity.com/paper/2203.03466/) (2022)
- [Training Deep Learning Models with Norm-Constrained LMOs](https://mlsanity.com/paper/2502.07529/) (2025)
- [The Road Less Scheduled](https://mlsanity.com/paper/2405.15682/) (2024)
- [DiLoCo: Distributed Low-Communication Training of Language Models](https://mlsanity.com/paper/2311.08105/) (2023)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [Modular Duality in Deep Learning](https://mlsanity.com/paper/2410.21265/) (2024)
- [Exponential Moving Average of Weights in Deep Learning: Dynamics and Benefits](https://mlsanity.com/paper/2411.18704/) (2024)
- [Scalable Optimization in the Modular Norm](https://mlsanity.com/paper/2405.14813/) (2024)
- [The AdEMAMix Optimizer: Better, Faster, Older](https://mlsanity.com/paper/2409.03137/) (2024)
- [Optimal Linear Decay Learning Rate Schedules and Further Refinements](https://mlsanity.com/paper/2310.07831/) (2023)
- [Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunch](https://mlsanity.com/paper/2501.18512/) (2025)
- [Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo](https://mlsanity.com/paper/2503.09799/) (2025)
- [Asynchronous Local-SGD Training for Language Modeling](https://mlsanity.com/paper/2401.09135/) (2024)
- [Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner](https://mlsanity.com/paper/2506.03595/) (2025)
- [Momentum via Primal Averaging: Theoretical Insights and Learning Rate Schedules for Non-Convex Optimization](https://mlsanity.com/paper/2010.00406/) (2020)
- [Training trajectories, mini-batch losses and the curious role of the learning rate](https://mlsanity.com/paper/2301.02312/) (2023)
- [LaProp: Separating Momentum and Adaptivity in Adam](https://mlsanity.com/paper/2002.04839/) (2020)
- [MuLoCo: Muon is a practical inner optimizer for DiLoCo](https://mlsanity.com/paper/2505.23725/) (2025)
- [Understanding Outer Optimizers in Local SGD: Learning Rates, Momentum, and Acceleration](https://mlsanity.com/paper/2509.10439/) (2025)
- [SNOO: Step-K Nesterov Outer Optimizer - The Surprising Effectiveness of Nesterov Momentum Applied to Pseudo-Gradients](https://mlsanity.com/paper/2510.15830/) (2025)
- [On the Curved Geometry of Accelerated Optimization](https://mlsanity.com/paper/1812.04634/) (2018)

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

- [Optimistic Dual Averaging Unifies Modern Optimizers](https://mlsanity.com/paper/2605.11172/) (2026)
- [High dimensional theory of two-phase optimizers](https://mlsanity.com/paper/2603.26954/) (2026)
- [EMA-Nesterov: Stabilizing Nesterov's Lookahead for Accelerated Deep Learning Optimization](https://mlsanity.com/paper/2605.25395/) (2026)
- [Unifying Local Communications and Local Updates for LLM Pretraining](https://mlsanity.com/paper/2606.11081/) (2026)
- [Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining](https://mlsanity.com/paper/2608.28442/) (2026)

## Abstract (arXiv)

We propose Generalized Primal Averaging (GPA), an extension of Nesterov's method that unifies and generalizes recent averaging-based optimizers like single-worker DiLoCo and Schedule-Free, within a non-distributed setting. While DiLoCo relies on a memory-intensive two-loop structure to periodically aggregate pseudo-gradients using Nesterov momentum, GPA eliminates this complexity by decoupling Nesterov's interpolation constants to enable smooth iterate averaging at every step. Structurally, GPA resembles Schedule-Free but replaces uniform averaging with exponential moving averaging. Empirically, GPA consistently outperforms single-worker DiLoCo and AdamW with reduced memory overhead. GPA achieves speedups of 8.71%, 10.13%, and 9.58% over the AdamW baseline in terms of steps to reach target validation loss for Llama-160M, 1B, and 8B models, respectively. Similarly, on the ImageNet ViT workload, GPA achieves speedups of 7% and 25.5% in the small and large batch settings respectively. Furthermore, we prove that for any base optimizer with O(√T) regret, where T is the number of iterations, GPA matches or exceeds the original convergence guarantees depending on the interpolation constants.