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

2026-05-27

AI Summary

This work addresses the inefficiency in communication and instability in convergence caused by outer-loop momentum accumulation in distributed optimization. To mitigate these issues, the authors propose a periodic momentum restarting mechanism that periodically resets the outer-loop momentum between communication rounds. This approach leverages phase cancellation to eliminate stale momentum information while preserving inner-loop optimization progress, thereby expanding the stable regime of outer-loop learning rates and momentum coefficients. Through analysis based on a linearized squared-loss model and empirical neural tangent kernel (NTK) theory—integrated within distributed frameworks such as DiLoCo—the study theoretically establishes a mode contraction effect induced by restarting. Experimental results demonstrate that the proposed method significantly enlarges the stability range of outer-loop hyperparameters and enhances overall optimization efficiency in large language model pretraining.

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.

Problem

• outer-momentum restarting
• distributed optimization
• two-phase optimization
• communication efficiency
• momentum memory

Innovation

• outer-momentum restarting
• distributed optimization
• phase cancellation
• communication efficiency
• two-phase optimization