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

Kristi Topollai, Allan Ma, Tolga Dimlioglu, Sui Jiet Tay, Anna Ewa Choromanska

Keywords: optimization, distributed optimization, communication-efficient training, local updates, two-phase optimization, momentum methods

TL;DR: We study outer-momentum restarts for DiLoCo-style training and show, via NTK mode analysis and LLaMA pretraining, that they widen the stable range of outer optimizer hyperparameters.

Abstract:

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.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 140