Cosine-Gated Adam-Decay: Drop-In Staleness-Aware Outer Optimization for Decoupled DiLoCo

arXiv:2605.09126v1 Announce Type: new Abstract: Asynchronous DiLoCo systems may receive pseudo-gradients computed several outer rounds earlier, yet the standard Nesterov outer optimizer does not explicitly condition its update on per-update age. This can make the outer momentum buffer brittle under large controlled delays. We propose Cosine Gated Adam Decay (CGAD), a simple, drop-in, age-aware outer optimizer that scales each incoming pseudo-gradient by sigma(au) = gamma(au) e^{-alphaau} before it enters Adam's first- and second-moment buffers; the exponential models information decay and the cosine gate gamma(au) smoothly zeroes contributions past a chosen cutoff. CGAD reduces to plain Adam at au=0, adds two hyperparameters whose defaults transfer across scales, and extends to partial-sync schedulers via a per-fragment age-aware variant (PA-CGAD). For an idealized gated-adaptive update on smooth non convex objectives, we prove a non-asymptotic convergence bound whose staleness-bias term depends on alpha alone, rather than on the realized maximum delay au_{max}; standard analyses of asynchronous momentum-SGD instead carry a au_{max}^2 factor. Empirically, on Llama style language model pretraining at 25M, 1B, and 7B parameters, CGAD trains stably across the controlled delays we sweep. The cosine cutoff acts as scale insurance: the closest baseline, Adam Decay (CGAD without the cutoff), is competitive at 25M but its seed-to-seed sigma at au=8 grows 27x from 25M to 7B, pushing its single-shot risk (mean + sigma) above the chance-level loss while CGAD's stays well below. The published Nesterov recipe is the least stable method on the full sweep.

Source: arXiv cs.LG | 2026-05-12