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

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Topollai, Kristi, Ma, Allan, Dimlioglu, Tolga, Tay, Sui Jiet, Choromanska, Anna
Format: Preprint
Veröffentlicht: 2026
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866918527996264448
author Topollai, Kristi
Ma, Allan
Dimlioglu, Tolga
Tay, Sui Jiet
Choromanska, Anna
author_facet Topollai, Kristi
Ma, Allan
Dimlioglu, Tolga
Tay, Sui Jiet
Choromanska, Anna
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2605_28585
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Outer-Momentum Restarting in High-Dimensional Two-Phase Optimization
Topollai, Kristi
Ma, Allan
Dimlioglu, Tolga
Tay, Sui Jiet
Choromanska, Anna
Machine Learning
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.
title Outer-Momentum Restarting in High-Dimensional Two-Phase Optimization
topic Machine Learning
url https://arxiv.org/abs/2605.28585