SNOO: Step-K Nesterov Outer Optimizer - The Surprising Effectiveness of Nesterov Momentum Applied to Pseudo-Gradients

TL;DR

This work addresses the efficiency of two-loop optimization in large language model pre-training by showing that the surprising gains of DiLoCo stem from applying Nesterov momentum to pseudo-gradients. It introduces Step-$K$ Nesterov Outer Optimizer (SNOO), which wraps an inner optimizer and updates slow weights with gradient descent plus Nesterov momentum on the pseudo-gradient $s_t = w_t - \tilde{w}_{t,K}$. Across open-source and large-scale settings up to $1\times 10^{23}$ training FLOPs, SNOO yields compute-factor gains of about $1.5$–$2.5\times$, with larger gains for bigger models, and incurs negligible overhead due to infrequent outer updates and compatibility with sharding; it also exhibits implicit regularization, such as smaller weight norms and improved resilience to overfitting. These results suggest SNOO is a practical, scalable enhancement for efficient LLM pre-training, motivating further theoretical study of its convergence and regularization properties.

Abstract

The rapid development of large language models (LLMs) has driven the demand for more efficient optimization techniques. Among these, the Lookahead family of optimizers employs a two-loop framework, maintaining fast and slow sets of model weights. Multiple inner optimizer steps on the fast weights produce a trajectory - the pseudo-gradient - that is used to update the slow weights. DiLoCo, a notable example originally designed for distributed training, applies Nesterov momentum to the averaged pseudo-gradient from multiple workers, claiming to even outperform AdamW in a non-distributed setup. In this paper, we empirically show that DiLoCo's surprising effectiveness stems primarily from applying Nesterov momentum to the pseudo-gradient, which improves training in a non-distributed setting. We call this Lookahead variant the Step-$K$ Nesterov Outer Optimizer (SNOO). We demonstrate that SNOO achieves compute factor gains of 1.5 - 2.5$\times$ in a non-distributed setting up to a scale of 1e23 training FLOPs, with improvements that increase with model size. Because of its minimal compute and memory overhead and compatibility with model sharding, SNOO is a practical enhancement for a variety of inner optimizers, including AdamW and Muon.

Figures (7)

• Figure 1: Comparison of optimizer performance when varying the number of workers and number of inner steps.
• Figure 2: SNOO shows strong improvements across compute scales over the AdamW baseline on the C4 validation dataset. This figure plots NLL on a held-out validation set of C4 run using TorchTitan on OSS Llama-3 models.
• Figure 3: SNOO outperforms the AdamW baseline throughout training, exhibiting lower train and validation loss using a 300M dense transformer model. This figure plots NLL on training and held-out validation set of C4 run using TorchTitan on OSS Llama-3 models.
• Figure 4: SNOO outperforms the AdamW baseline across compute scales for dense models, with wins ranging from 1.35$\times$ through 2.75$\times$.
• Figure 5: SNOO outperforms the AdamW baseline across compute scales also for MoE models, with wins ranging from 1.2$\times$ through 2$\times$.