cs.LGMar 27, 2026

High dimensional theory of two-phase optimizers

Authors: Atish Agarwala

Organizations: Google DeepMind

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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, ν=2\nu=2 ). Convergence rate is improved for Nesterov momentum. At the critical ν=(1−βout)−1\nu=(1-\beta_{\rm out})^{-1} , Nesterov convergence rate becomes SS -invariant at large SS (right). Reduced system with k~t:=0\tilde{\mathbf{k}}_{t}:=0 every cycle has similar behavior to full dynamics with k~t\tilde{\mathbf{k}}_{t} preserved for larger SS , but shows faster convergence/better stability at smaller SS . All experiments have η=1\eta=1 , λ=0.2\lambda=0.2 , βin=0.9\beta_{\rm in}=0.9 , βout=0.8\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=3200D=3200 , Btot=64B_{tot}=64 , various learning rates). Keeping η\eta and ν\nu independent of number of workers RR leads to similar dynamics at smaller learning rates (top left) but causes learning curves for larger RR to diverge early at larger learning rates (top right). R−1/2R^{-1/2} scaling rule for η\eta , with νη\nu\eta fixed gives better correspondence across RR for all learning rates (bottom), but dynamics still becomes non-universal at larger η\eta (bottom right)
Figure 3 : Loss after one cycle of S=10S=10 steps as a function of η\eta and ν\nu for spiked spectrum (left). Batch fraction B/D=0.2B/D=0.2 , and spectrum has two eigenvalues λ0\lambda_{0} ( 99%99\% of eigenmodes) and 20λ020\lambda_{0} ( 1%1\% of eigenmodes). Optimal η∗(ν)\eta^{*}(\nu) decreases in ν\nu . Loss at η∗(ν)\eta^{*}(\nu) is minimized for ν>1\nu>1 (middle). Power law spectrum with exponent α=−1.5\alpha=-1.5 has optimal ν<1\nu<1 (right, B/D=0.005B/D=0.005 ).
Figure 4 : Loss curves for a two-layer MLP on MNIST, trained using LA-Diloco with various RR . When the inner and outer learning rates are kept fixed as a function of RR (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 RR .
Figure 5 : Per-worker dispersion E[∣∣ψt,k−θt∣∣2]{\rm E}[||\psi_{t,k}-\theta_{t}||^{2}] for a two-layer MLP on MNIST, trained using LA-Diloco with various RR . When the inner and outer learning rates are kept fixed as a function of RR , we see a large increase in dispersion. When we use the proposed scaling rule, the dispersions are similar across RR , reflecting a more similar noise process.
WorkloadAlgorithmValidation Metricην\eta\nuν\nuβout\beta_{\rm out}βin\beta_{\rm in}
fineweb_edu_lmNesterovPerplexity: 28.50.0007850.8660.8320.804
fineweb_edu_lmHeavy BallPerplexity: 29.30.0007850.8660.8320.804
imagenet_vit_s16NesterovError Rate: 0.2650.001891.290.8170.930
imagenet_vit_s16Heavy BallError Rate: 0.2690.001891.290.8170.930
Appendix
Table 1 : Validation metrics and key hyperparameters for best training runs with SLA on the AlgoPerf benchmark ( 5050 point Halton search). Nesterov momentum appears to perform better than heavy ball momentum in this setting.

Explore similar work

CardsList

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

    May 27, 2026Kristi Topollai, Allan Ma, Tolga Dimlioglu +2MomentumPeak Memory

    New York University, New York, USA

  2. Optimistic Dual Averaging Unifies Modern Optimizers

    May 11, 2026Thomas Pethick, Wanyun Xie, Roman Machacek +1Adaptive OptimizersFederated Averaging

    Independent Researcher · EPFL (LIONS) · University of Bern

  3. Factored Gossip DiLoCo: Reducing Blocking Communication in DiLoCo

    Jun 22, 2026Chamin Hewa Koneputugodage, Thalaiyasingam Ajanthan, Sameera Ramasinghe +7Distributed OptimizationCpu-Gpu Hybrid Designs

    1Pluralis Research