How Inner Optimizers Shape Replica Scaling

Summary (Overview)

• Core Problem: The paper addresses replica scaling in distributed LLM training—the phenomenon where increasing the number of model replicas at a fixed token budget leads to monotonically decreasing performance relative to non-local data-parallel (DP) training.

• Key Insight: The choice of inner optimizer (the local optimizer used between synchronization rounds in DiLoCo-style methods) critically impacts replica scaling through two mechanisms: (1) the covariance structure of pseudo-gradients and (2) the impact of batch size on pseudo-gradient noise scale.

• Main Finding: Whitening-based preconditioning is a powerful desideratum for inner optimizers. SPlus (a sign-based optimizer in a Shampoo eigenbasis) achieves the strongest replica scaling (as low as +2.1% degradation relative to DP at 96 replicas), while SOAP offers the best absolute DP performance.

• Practical Contribution: The paper shows that appropriately scaling $\beta_2$ (second-moment EMA coefficient) as replicas grow—following the rule $\beta_2(R) = \beta_2(R_0)^{R_0/R}$—can make the computationally cheaper IsoLoCo method outperform the more expensive MuLoCo.

• Theoretical Framework: The authors decompose the error relative to DP into a mean mismatch term and a stochastic error term, providing explicit bounds and covariance analyses for Adam, Shampoo, SOAP, SPlus, and Muon.

Introduction and Theoretical Foundation

The computational scale of language-model training increasingly exceeds single-site capacity. Training across geographically separated clusters introduces communication bottlenecks that local optimization methods address by allowing replicas to optimize independently for several steps between parameter or pseudo-gradient averages. However, a central challenge is replica scaling: at a fixed token budget, increasing replicas degrades final performance relative to DP training.

The DiLoCo Framework

The paper generalizes the DiLoCo framework (Algorithm 1):

• R replicas start each round from shared parameters
• Each runs H steps of an inner optimizer on its own data stream
• Resulting displacements $\Delta_r$ are uniformly averaged
• An outer optimizer (typically Nesterov momentum) applies the aggregate $\bar{\Delta}$

Key notation: Inner optimizer states $s_r$ are never communicated—each is estimated locally from its replica's data stream.

Why Adaptive Optimizers Create Divergence

Classical analyses of local SGD attribute degradation to client drift over multiple local steps. However, nonlinear adaptive inner optimizers need not satisfy the commutativity property that makes averaging updates equivalent to updating with averaged gradients. A discrepancy arises even on the first step with identical starting states because transforming each local gradient before averaging differs from transforming the averaged gradient.

The Error Decomposition

The paper decomposes each replica's displacement into:

$$\boldsymbol{\Delta}_r = \boldsymbol{\mu}_r + \varepsilon_r, \quad \boldsymbol{\mu}_r = \mathbb{E}[\boldsymbol{\Delta}_r], \quad \mathbb{E}[\varepsilon_r] = 0, \quad \boldsymbol{C}_r = \text{Cov}(\boldsymbol{\Delta}_r)$$

Proposition 1 (Directional error decomposition):

$$\underbrace{\mathbb{E}\left[\left(\boldsymbol{q}^\top (\alpha_R \bar{\boldsymbol{\Delta}} - \boldsymbol{d}_H^{\text{DP}})\right)^2\right]}_{\text{error relative to DP mean}} = \underbrace{(\boldsymbol{q}^\top \boldsymbol{b}_R)^2}_{\text{mean mismatch}} + \underbrace{\alpha_R^2 \boldsymbol{q}^\top \boldsymbol{C}_{\text{agg}} \boldsymbol{q}}_{\text{stochastic error}}$$

where $\boldsymbol{b}_R = \alpha_R \bar{\boldsymbol{\mu}} - \boldsymbol{d}_H^{\text{DP}}$ and $\boldsymbol{C}_{\text{agg}} = \frac{1}{R^2}\sum_r \boldsymbol{C}_r$.

Methodology

Optimizers Analyzed

The paper analyzes five optimizers using unified notation:

Adam: Entrywise second moment in parameter coordinates:

$$\boldsymbol{V}_{A,t} = \beta_2 \boldsymbol{V}_{A,t-1} + (1-\beta_2)\boldsymbol{G}^{\odot 2}, \quad \boldsymbol{U}_{\text{Adam}} = \boldsymbol{M}_t \oslash (\sqrt{\boldsymbol{V}_{A,t}} + \epsilon)$$

Shampoo: Maintains left/right Gram factors:

$$\boldsymbol{U}_{\text{Shampoo}} = (\boldsymbol{L}_t + \epsilon\boldsymbol{I})^{-1/4}\boldsymbol{M}_t(\boldsymbol{R}_t + \epsilon\boldsymbol{I})^{-1/4}$$

SOAP: Adam in Shampoo's eigenbasis (rotation → Adam normalization → rotate back):

$$\boldsymbol{U}_{\text{SOAP}} = \boldsymbol{Q}_L[\widetilde{\boldsymbol{M}}_t \oslash (\sqrt{\boldsymbol{V}_{S,t}} + \epsilon)]\boldsymbol{Q}_R^\top$$

SPlus: Signs momentum in Shampoo's eigenbasis (no second moment):

$$\boldsymbol{U}_{\text{SPlus}} = \boldsymbol{Q}_L \text{sign}(\widetilde{\boldsymbol{M}}_t)\boldsymbol{Q}_R^\top$$

Muon: Polar decomposition of momentum:

$$\boldsymbol{U}_{\text{Muon}} = \gamma \text{polar}(\boldsymbol{M}_t) = \gamma \boldsymbol{U}_M \boldsymbol{V}_M^\top$$

Theoretical Analysis of Covariance Structure

The paper derives covariance identities for each optimizer under idealized assumptions:

Key result for SPlus vs. SOAP under imperfect decorrelation: Using the Gaussian arcsine law:

$$\boldsymbol{C}_{\text{SPlus}} = \boldsymbol{Q}\left[\frac{2}{\pi}\arcsin(\boldsymbol{\mathcal{C}})\right]\boldsymbol{Q}^\top, \quad \left|\frac{2}{\pi}\arcsin(\rho)\right| \leq |\rho|$$

This shows SPlus contracts residual correlations without requiring variance estimates, giving $\|\boldsymbol{C}_{\text{SPlus}} - \boldsymbol{I}\|_F^2 \leq \|\boldsymbol{C}_{\text{SOAP}} - \boldsymbol{I}\|_F^2$.

Mean Mismatch Bounds

Proposition 2 (SPlus): With $\alpha_R = \sqrt{R}$, the mean mismatch satisfies:

$$\|\boldsymbol{b}_R\|_2^2 \leq \frac{c^2}{36}\left(1 - \frac{1}{R}\right)^2 \sum_j a_j^6$$

where $c = \sqrt{2/\pi}$ and $a_j$ is the mean divided by standard deviation of the DP momentum coordinate. The mismatch is third order in the signal-to-noise ratio.

Proposition 3 (Adam): The coordinate-wise mean mismatch is bounded by:

$$\left|b_{R,j}\right| \leq \frac{(1-\beta_2)(R-1)\sigma^2}{2(\beta_2 v)^{3/2}}\left(|\beta_1 m| + 3(1-\beta_1)|g|\right)$$

This grows linearly with R−1 for fixed $\beta_2$, motivating the scaling rule:

$$\beta_2(R) = \beta_2(R_0)^{R_0/R}, \qquad 1-\beta_2(R) = O(R^{-1})$$

Empirical Validation / Results

Experimental Setup

• Models: 178M and 512M parameter LLaMA models
• Data: DCLM, sequence length 2048
• Global batch: 524K tokens (178M), 2M tokens (512M)
• Local steps: $H = 30$ between aggregations
• Outer optimizer: Nesterov SGD

Main Results: Replica Scaling

178M scale (Table 1):

At R=64: loss change is +3.1% for SPlus, +4.8% for SOAP, +8.5% for Muon, +11% for AdamW.

512M scale (Table 2):

At R=96: SPlus has +2.1% degradation vs. +4.4% for SOAP, +7.7% for Muon, +12% for AdamW.

Single-Step Results (H=1)

The H=1 experiments (Table 3) isolate the effect of the optimizer's nonlinear transformation from multi-step drift. Results mirror the multi-step setting, confirming that the optimizer's update transformation itself—not just parameter drift—drives replica scaling degradation.

Effect of $\beta_2$ Scaling

Tuning $\beta_2$ according to the token-half-life rule significantly improves performance:

• At 512M and R=64: DiLoCo loss drops from 2.884 to 2.849; IsoLoCo from 2.734 to 2.702 (~1.2% improvement)
• IsoLoCo with tuned $\beta_2$ outperforms MuLoCo at large worker counts, despite using the cheaper AdamW inner optimizer

Empirical Covariance Analysis

Table 4: Update-noise anisotropy ($a_h = 0$ means isotropy):

Table 5: Mean mismatch vs. stochastic error:

Table 6: Curvature along high-noise directions (relative to random directions):

High noise along high-curvature directions is particularly damaging; whitening redistributes error more uniformly, explaining why SOAP and SPlus excel.

Theoretical and Practical Implications

Theoretical Contributions

1. Unified error decomposition: Separates replica scaling degradation into mean mismatch (deterministic) and stochastic error (covariance-driven) components.

2. Whitening as a design principle: The paper establishes that whitening-based preconditioning controls the worst-direction aggregate noise and makes uniform averaging optimal when replicas share similar covariance structures.

3. $\beta_2$ scaling rule: Transfers the small-batch training rule of Marek et al. (2025) to the replica scaling setting, providing a principled way to maintain second-moment EMA half-life in tokens.

Practical Implications

1. Optimizer selection guidance:

   • SOAP is preferred at smaller replica counts (stronger DP performance)
   • SPlus is preferred at larger replica counts (best scaling, minimal degradation)
   • Both achieve <3% degradation at 50+ replicas

2. Memory considerations: SOAP and SPlus require more memory than Muon (maintaining L and R factors), but scalable recipes exist (Khona et al., 2026) and cross-data-center settings can shard optimizer states across chips.

3. IsoLoCo as a practical candidate: With appropriate $\beta_2$ tuning, IsoLoCo (nearly the same overhead as DiLoCo) outperforms the more expensive MuLoCo, making it ideal for memory-constrained settings.

4. The H=1 insight: Replica scaling effects are visible even with a single local step, meaning the optimizer's nonlinear transformation matters more than multi-step parameter drift.

Conclusion

The paper demonstrates that replica scaling degradation in local optimization can be traced to two primary error sources: mean mismatch and stochastic error. Whitening-based inner optimizers (SOAP, SPlus) control the stochastic error by producing near-isotropic pseudo-gradient covariances, while appropriately scaling $\beta_2$ controls the mean mismatch for Adam-based methods.

Key findings:

• SPlus achieves the strongest replica scaling (+2.1% at 96 replicas vs. DP)
• SOAP offers the best DP performance and becomes competitive at scale
• $\beta_2$ scaling makes IsoLoCo competitive with or superior to MuLoCo
• Covariance structure empirically correlates with scaling performance across all optimizers

Future directions could include: scaling to larger models where whitening-based methods may benefit even more, developing memory-efficient implementations of SPlus, and extending the analysis to heterogeneous data distributions across replicas.