# How Inner Optimizers Shape Replica Scaling

> Whitening-based inner optimizers like SPlus achieve near-optimal replica scaling by decorrelating pseudo-gradients, degrading only 2.1% versus data-parallel training at 96 replicas.

- **Source:** [arXiv](https://arxiv.org/abs/2610.19064)
- **Published:** 2026-10-17
- **Permalink:** https://picx.dev/p/VIJVfq
- **Whiteboard:** https://picx.dev/p/VIJVfq/image

## Summary

# 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:

| Optimizer | Resulting Covariance | Whitening? |
|-----------|---------------------|------------|
| Adam | $\boldsymbol{C} = \boldsymbol{D}^{-1/2}\boldsymbol{\Sigma}\boldsymbol{D}^{-1/2}$ (preserves correlations) | No |
| Shampoo | Condition number = square root of input's | No |
| SOAP | $\boldsymbol{C} = \boldsymbol{I}$ (with perfect decorrelation) | Yes |
| SPlus | $\boldsymbol{C} = \boldsymbol{I}$ (with perfect decorrelation) | Yes |
| Muon | $\boldsymbol{C} = \frac{\gamma^2}{\max(m,n)}\boldsymbol{I}$ (under rotational symmetry) | Conditional |

**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)**:

| Optimizer | DP | R=8 | R=16 | R=32 | R=64 |
|-----------|-----|------|-------|-------|-------|
| **SOAP** | **2.816** | 2.851 | 2.869 | 2.898 | 2.951 |
| **SPlus** | 2.853 | **2.851** | **2.867** | **2.897** | **2.942** |
| Muon | 2.858 | 2.887 | 2.928 | 3.000 | 3.099 |
| AdamW | 2.880 | 2.952 | 3.012 | 3.080 | 3.208 |

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)**:

| Optimizer | DP | R=8 | R=16 | R=32 | R=64 | R=96 |
|-----------|-----|------|-------|-------|-------|-------|
| **SOAP** | **2.563** | **2.579** | **2.598** | **2.618** | 2.649 | 2.676 |
| **SPlus** | 2.610 | 2.610 | 2.618 | **2.625** | **2.646** | **2.663** |
| Muon | 2.586 | 2.615 | 2.637 | 2.677 | 2.738 | 2.787 |
| AdamW | 2.646 | 2.691 | 2.726 | 2.785 | 2.881 | 2.958 |

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):

| Optimizer | $a_1 \downarrow$ | $a_{30} \downarrow$ |
|-----------|------------------|---------------------|
| AdamW | 35 | 11 |
| Muon | 3.8 | 3.3 |
| SOAP | 2.3 | 2.5 |
| SPlus | 0.43 | 2.4 |

**Table 5: Mean mismatch vs. stochastic error**:

| Optimizer | Mean mismatch (×10⁻⁶) | Stochastic error (×10⁻⁶) | Mean/noise (%) |
|-----------|------------------------|--------------------------|----------------|
| SPlus | 0.044 | 0.91 | 7.9 |
| SOAP | 0.12 | 1.4 | 7.5 |
| Muon | 0.11 | 1.1 | 14 |
| AdamW | 5.4 | 19 | 19 |

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

| Optimizer | κ(q̂) at R=8 | κ(q̂) at R=64 |
|-----------|-------------|--------------|
| AdamW | 86–170 | 460–730 |
| Muon | 6.4–7.7 | 47–100 |
| SOAP | 3.1–3.2 | 4.2–5.2 |
| SPlus | 1.5–1.6 | 1.5 |

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.

---

_Markdown view of https://picx.dev/p/VIJVfq, served by PicX — AI-generated visual whiteboard summaries of research papers._