High dimensional theory of two-phase optimizers

Atish Agarwala

cs.LGMar 27, 2026

High dimensional theory of two-phase optimizers

Authors: Atish Agarwala

Organizations: Google DeepMind

Abstract

The trend towards larger training setups has brought a renewed interest in partially asynchronous two-phase optimizers which optimize locally and then synchronize across workers. Additionally, recent work suggests that the one-worker version of one of these algorithms, DiLoCo, shows promising results as a (synchronous) optimizer. Motivated by these studies we present an analysis of LA-DiLoCo, a simple member of the DiLoCo family, on a high-dimensional linear regression problem. We show that the one-worker variant, LA, provides a different tradeoff between signal and noise than SGD, which is beneficial in many scenarios. We also show that the multi-worker version generates more noise than the single worker version, but that this additional noise generation can be ameliorated by appropriate choice of hyperparameters. We conclude with an analysis of SLA -- LA with momentum -- and show that stacking two momentum operators gives an opportunity for acceleration via a non-linear transformation of the "effective'' Hessian spectrum, which is maximized for Nesterov momentum. Altogether our results show that two-phase optimizers represent a fruitful new paradigm for understanding and improving training algorithms.

Figures & tables

Figure 1 : Eigenmodes damped by inner optimizer show decreasing per-step convergence rate with momentum-GD in the outer optimizer (left,

\nu=2

). Convergence rate is improved for Nesterov momentum. At the critical

\nu=(1-\beta_{\rm out})^{-1}

, Nesterov convergence rate becomes

S

-invariant at large

S

(right). Reduced system with

\tilde{\mathbf{k}}_{t}:=0

every cycle has similar behavior to full dynamics with

\tilde{\mathbf{k}}_{t}

preserved for larger

S

, but shows faster convergence/better stability at smaller

S

. All experiments have

\eta=1

\lambda=0.2

\beta_{\rm in}=0.9

\beta_{\rm out}=0.8

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2 : Averaged loss curves for LA diloco (solid lines) are well captured by theoretical model (dashed lines) across many settings with fixed total batch size (all panels,

D=3200

B_{tot}=64

, various learning rates). Keeping

\eta

and

\nu

independent of number of workers

R

leads to similar dynamics at smaller learning rates (top left) but causes learning curves for larger

R

to diverge early at larger learning rates (top right).

R^{-1/2}

scaling rule for

\eta

, with

\nu\eta

fixed gives better correspondence across

R

for all learning rates (bottom), but dynamics still becomes non-universal at larger

\eta

(bottom right)

Panel of Figure 3 — Figure 3 : Loss after one cycle of $S=10$ steps as a function of $\eta$ and $\nu$ for spiked spectrum (left). Batch fraction $B/D=0.2$ , and spectrum has two eigenvalues $\lambda_{0}$ ( $99\%$ of eigenmodes) and $20\lambda_{0}$ ( $1\%$ of eigenmodes). Optimal $\eta^{*}(\nu)$ decreases in $\nu$ . Loss at $\eta^{*}(\nu)$ is minimized for $\nu>1$ (middle). Power law spectrum with exponent $\alpha=-1.5$ has optimal $\nu<1$ (right, $B/D=0.005$ ).

Figure 4 : Loss curves for a two-layer MLP on MNIST, trained using LA-Diloco with various

R

. When the inner and outer learning rates are kept fixed as a function of

R

(left), we see no collapse. When we use the proposed scaling rule (right), we see good agreement except at early times or for higher values of

R

Figure 5 : Per-worker dispersion

{\rm E}[||\psi_{t,k}-\theta_{t}||^{2}]

for a two-layer MLP on MNIST, trained using LA-Diloco with various

R

. When the inner and outer learning rates are kept fixed as a function of

R

, we see a large increase in dispersion. When we use the proposed scaling rule, the dispersions are similar across

R

, reflecting a more similar noise process.

Workload	Algorithm	Validation Metric	$\eta\nu$	$\nu$	$\beta_{\rm out}$	$\beta_{\rm in}$
fineweb_edu_lm	Nesterov	Perplexity: 28.5	0.000785	0.866	0.832	0.804
fineweb_edu_lm	Heavy Ball	Perplexity: 29.3	0.000785	0.866	0.832	0.804
imagenet_vit_s16	Nesterov	Error Rate: 0.265	0.00189	1.29	0.817	0.930
imagenet_vit_s16	Heavy Ball	Error Rate: 0.269	0.00189	1.29	0.817	0.930

Appendix

Table 1 : Validation metrics and key hyperparameters for best training runs with SLA on the AlgoPerf benchmark (

50

point Halton search). Nesterov momentum appears to perform better than heavy ball momentum in this setting.

Explore similar work

May 27, 2026cs.LG

Outer-Momentum Restarting in High-Dimensional Two-Phase Optimization

Communication-efficient distributed optimizers such as DiLoCo reduce synchronization costs by letting workers perform many local updates before aggregating their progress with an outer momentum optimizer. Recent theory suggests that the outer optimizer acts on an effective spectrum induced by the inner optimization loop, and that the choice of outer momentum controls how progress from local updates is accumulated across communication rounds. We study periodic restarting of the outer momentum as a simple complementary mechanism for controlling this outer memory. In a linearized squared-loss model where prediction-space residuals evolve under the empirical NTK, we derive a mode-wise restart contraction showing that resets exploit phase cancellation by discarding stale momentum while preserving inner-loop progress. Toy experiments verify the predicted contraction behavior, and language-model pretraining experiments show that periodic restarts widen the stable range of outer learning rates and momentum values across communication periods.

Kristi Topollai, Allan Ma, Tolga Dimlioglu +2

New York University, New York, USA

Momentum Peak Memory

Ask PDF

May 11, 2026cs.LG

Optimistic Dual Averaging Unifies Modern Optimizers

We introduce SODA, a generalization of Optimistic Dual Averaging, which provides a common perspective on state-of-the-art optimizers like Muon, Lion, AdEMAMix and NAdam, showing that they can all be viewed as optimistic instances of this framework. Based on this framing, we propose a practical SODA wrapper for any base optimizer that eliminates weight decay tuning through a theoretically-grounded

1/k

decay schedule. Empirical results across various scales and training horizons show that SODA consistently improves performance without any additional hyperparameter tuning.

Thomas Pethick, Wanyun Xie, Roman Machacek +1

Independent Researcher · EPFL (LIONS) · University of Bern

Adaptive Optimizers Federated Averaging Taco

Ask PDF

Jun 22, 2026cs.LG

Factored Gossip DiLoCo: Reducing Blocking Communication in DiLoCo

To make large-scale distributed training practical outside high-bandwidth datacenters, we must reduce blocking, high-volume synchronization. While DiLoCo communicates infrequently, its outer synchronization remains bandwidth-heavy and brittle to stragglers and transient failures. We relax exact synchronization to approximate synchronization via mixing/gossip, which degrades gracefully under delays and communication failures. This allows us to factorize DiLoCo synchronization into a non-blocking mixing step that overlaps computation with no staleness, and a blocking mixing step that tightens worker agreement, yielding a tunable trade-off between compute utilization and optimization stability. On up to billion-parameter language models in low-bandwidth settings, our framework substantially improves compute utilization compared to DiLoCo, with training progress ranging from comparable to closely matching it, and is more robust to failures.

Chamin Hewa Koneputugodage, Thalaiyasingam Ajanthan, Sameera Ramasinghe +7

1Pluralis Research

Distributed Optimization Cpu-Gpu Hybrid Designs Gossip Data Centers

Ask PDF