Technical Summary: Cosine-Gated Adam-Decay (CGAD) for Decoupled DiLoCo

1. Problem Statement

In asynchronous Decoupled DiLoCo (DDiLoCo) systems, workers perform multiple local optimization steps before exchanging pseudo-gradients with a global parameter server. A critical challenge in this architecture is staleness: the pseudo-gradients applied at outer round $t$ may have been computed $\tau$ rounds earlier due to network latency, straggler workers, or asynchrony budgets.

The standard outer optimizer for DiLoCo is SGD with Nesterov momentum. This method applies a uniform update step to every incoming pseudo-gradient without conditioning on its age ($\tau$). As $\tau$ increases, the momentum buffer accumulates a stale view of the loss landscape. Consequently, the optimizer takes steps that overshoot directions which no longer descend the current loss, leading to instability or divergence, particularly as model scale increases. Existing theoretical analyses of asynchronous momentum-SGD suggest that convergence rates degrade with a factor proportional to $\tau_{max}^2$ (the square of the maximum delay).

2. Methodology: Cosine-Gated Adam-Decay (CGAD)

The authors propose Cosine-Gated Adam-Decay (CGAD), a drop-in replacement for the Nesterov outer optimizer that explicitly conditions updates on the age of the pseudo-gradient.

Core Mechanism

CGAD modifies the standard Adam update by scaling each incoming pseudo-gradient $g_f$ for fragment $f$ by a factor $\sigma(\tau)$ before it enters Adam's first- and second-moment buffers:
$$\sigma(\tau) = \gamma(\tau) e^{-\alpha \tau}$$

Where:

• Exponential Decay ($e^{-\alpha \tau}$): Models the information decay of stale gradients. As $\tau$ grows, the contribution of the gradient to the momentum buffers is exponentially reduced.
• Cosine Gate ($\gamma(\tau)$): A smooth cutoff function defined as:
  $$\gamma(\tau) = \frac{1}{2} \left(1 + \cos\left(\frac{\pi \tau}{\tau_{cut}}\right)\right) \cdot \mathbb{I}[\tau \le \tau_{cut}]$$
  This gate equals 1 for fresh gradients, smoothly decays to 0 at a chosen cutoff $\tau_{cut}$, and strictly zeroes out contributions for $\tau > \tau_{cut}$.

Key Design Choices

1. Pre-Buffer Gating: The scaling $\sigma(\tau)$ is applied to the gradient before it updates the first moment ($m$) and second moment ($v$) buffers. This prevents stale gradients from inflating the second moment ($v$), which would otherwise attenuate subsequent fresh updates.
2. Per-Fragment Variant (PA-CGAD): For partial-sync schedulers (e.g., Streaming DiLoCo), the method tracks the age of individual fragments ($a_f$) and gates updates based on $\max(\tau, a_f)$, allowing a unified staleness model for both network delay and scheduling delays.
3. Hyperparameters: The method introduces two hyperparameters, $\alpha$ (decay rate) and $\tau_{cut}$ (cutoff), with defaults ($\alpha=0.2, \tau_{cut}=32$) that transfer across model scales.

3. Theoretical Contributions

The paper provides a non-asymptotic convergence bound for an idealized gated-adaptive update on smooth non-convex objectives.

• Staleness-Bias Independence: The convergence bound's staleness-bias term depends on the gating hyperparameter $\alpha$ (specifically $O(1/e^\alpha)$) rather than the realized maximum delay $\tau_{max}$.
• Contrast with Standard Analysis: Standard analyses of asynchronous momentum-SGD yield a complexity term scaling with $\tau_{max}^2$. By replacing this with a term dependent on $\alpha$, the gating mechanism theoretically decouples the convergence penalty from the worst-case delay distribution, provided $\tau_{cut} \ge 1/\alpha$.
• Scope: The authors note this is a motivation for the gating mechanism rather than a full convergence theorem for empirical Adam, as it assumes bounded normalized steps rather than deriving them from empirical statistics.

4. Experimental Results

Experiments were conducted on Llama-style language model pretraining at three scales: 25 M, 1 B, and 7 B parameters, using controlled delays up to $\tau=16$.

Stability and Convergence

• 1 B Scale: At $\tau=16$, the standard Nesterov recipe diverged (loss $\approx 319$), while CGAD achieved a loss of 7.44, comparable to its $\tau=0$ performance (7.48). The Nesterov method lost convergence entirely at $\tau \ge 8$.
• 25 M Scale: CGAD and other staleness-aware baselines (Adam-Decay, SDM, Delayed Nesterov) all remained stable, while Nesterov diverged at $\tau \ge 4$. At this scale, the differences between staleness-aware methods were marginal.

Scale Insurance and Robustness

The most significant finding emerges at the 7 B scale under memory-constrained conditions (bf16 weights, 8-bit AdamW):

• Seed-to-Seed Variance: The closest baseline, Adam-Decay (CGAD without the cosine cutoff), exhibited a 27-fold increase in seed-to-seed standard deviation ($\sigma$) as the model scaled from 25 M to 7 B (from 0.07 to 1.89 at $\tau=8$). In contrast, CGAD's $\sigma$ remained flat (approx. 0.38).
• Single-Shot Risk: Using "mean + $\sigma$" as a proxy for single-shot deployment risk, Adam-Decay's risk at 7 B / $\tau=8$ (13.19) crossed the chance-level loss (10.83) by 2.4 nats, implying a high probability of failure in a single production run. CGAD's risk (11.28) stayed well below this threshold.
• Mechanism of Failure: The authors attribute Adam-Decay's failure at 7 B to the accumulation of squared rounding noise in the second-moment buffer ($v$) from stale gradients that are never fully zeroed out. The cosine cutoff in CGAD prevents this by strictly zeroing contributions beyond $\tau_{cut}$.

Tuning Sensitivity

• Nesterov vs. CGAD: Even when Nesterov was re-tuned (removing momentum entirely), it could not match the performance of a tuned CGAD configuration at high delays. Tuned Nesterov required different hyperparameters for different delay regimes, whereas CGAD maintained robustness across the full range.
• Hyperparameter Sensitivity: The optimal $\alpha$ shifts slightly with scale (0.2 for 25 M, 0.4 for 1 B at $\tau=8$), but the default $\alpha=0.2$ remains robust across the full delay spectrum.

5. Significance and Claims

The paper claims that CGAD is the most robust scaling choice for asynchronous DiLoCo training, specifically addressing the brittleness of Nesterov momentum under large controlled delays.

• Scale Insurance: The primary contribution is the cosine cutoff, which acts as "scale insurance." While exponential decay alone (Adam-Decay) suffices for smaller models or full-precision settings, the hard cutoff is critical for frontier-scale models (7 B+) operating with compressed optimizer states (8-bit/bf16). It prevents the accumulation of stale variance that leads to catastrophic seed-to-seed variance and single-shot failure.
• Drop-in Compatibility: CGAD requires only a one-line modification to existing Adam implementations and adds two hyperparameters with defaults that transfer across scales.
• Practical Impact: For frontier-scale asynchronous training, the method offers the difference between a single training run that successfully ships and one that fails (lands above chance loss).

The authors remain modest regarding the theoretical scope, acknowledging that the convergence bound applies to an idealized update and that the empirical superiority is demonstrated specifically within the DiLoCo family and Llama-style pretraining contexts. They note that the inner-step count ($H=8$) used in experiments is lower than production DiLoCo settings ($H=100-500$), though sensitivity studies suggest the ranking persists at higher $H$.