High dimensional theory of two-phase optimizers

TL;DR

This paper analyzes high-dimensional two-phase optimizers, especially LA-DiLoCo, revealing their noise tradeoffs, hyperparameter effects, and potential for acceleration, offering insights into their advantages over traditional methods.

Contribution

It provides a theoretical analysis of LA-DiLoCo, a two-phase optimizer, highlighting its noise characteristics, hyperparameter tuning, and acceleration potential with momentum.

Findings

LA-DiLoCo offers a different signal-noise tradeoff than SGD.

Multi-worker LA-DiLoCo generates more noise, but hyperparameters can mitigate this.

Stacking momentum operators can accelerate convergence, especially with Nesterov momentum.

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.