PaperlayerView exact revision on arXiv

Source-linked AI summary

How Inner Optimizers Shape Replica Scaling

Eugene Belilovsky, Stefan Horoi, Edouard Oyallon, Benjamin Thérien

arXiv:2610.19064v1cs.LGcs.AI

TL;DR

At a fixed token budget, adding replicas can worsen performance relative to data-parallel training, and the role of inner optimizers in that scaling remains incompletely understood. The paper analyzes pseudogradient covariance and second-moment estimation across inner optimizers. It reports strong worker scaling for whitening-oriented optimizers, with up to 6% less degradation relative to DP at 96 replicas for SOAP or SPlus than existing methods.

  • Problem

    At a fixed token budget, increasing replicas often degrades performance relative to non-local data-parallel training, while how inner optimizers contribute to replica scaling remains poorly understood.

    Source: §1 · p. 1

  • Method

    The paper analyzes adaptive and preconditioned inner optimizers through pseudogradient covariance and second-moment estimation, linking optimizer choice to discrepancy from the DP update direction.

    Source: §1 · p. 2

  • Results

    Whitening-oriented optimizers achieve the strongest worker scaling, and SOAP or SPlus yields up to 6% less degradation relative to DP at 96 replicas than existing methods.

    Source: §1 · pp. 1–2

  • Takeaways & Limitations

    Whitening pseudogradient covariance is a promising inner-optimizer desideratum for replica scaling, though SOAP can have stronger absolute performance despite poorer worker scaling.

    Source: §1 · pp. 1–2

    Local whitening does not control the mean-mismatch term, which the paper analyzes separately.

    Source: §4.2 · p. 6

Abstract

from arXiv · show

Training large language models (LLMs) across geographically distributed and poorly interconnected compute pools or data centers has become a topic of significant interest. This has led to increased study of local optimization methods such as DiLoco for LLMs, especially those that use local (or inner) adaptive or preconditioned optimizers, which are critical for high-performance LLM training. However, a central challenge remains: it is well known that increasing the number of model replicas (replica scaling) while fixing the total training budget leads to monotonically decreasing performance. Existing work has empirically demonstrated that the choice of inner optimizer (Muon v.s. AdamW) can significantly improve this trend, but offers only limited explanation of the performance difference. In this work, we consider a range of inner-optimizers and identify two important factors that impact replica scaling: (1) the covariance structure of the pseudo-gradients and (2) the impact of batch size on the pseudogradient's noise scale. Based on these insights, we theoretically and empirically demonstrate that whitening-based pre-conditioning is a powerful inner-optimizer desideratum. Our empirical evaluation extends the results to reveal that whitening-based inner-optimizers achieve strong replica scaling (for 8-96 replicas), with performance relative to DP improving as parameters and data are scaled.

1 INTRODUCTION

The paper examines why replica scaling can degrade performance at a fixed token budget and how inner optimizers affect that trend. It connects optimizer choice to pseudogradient covariance and second-moment estimation, finding strong scaling with whitening-oriented optimizers.

1 INTRODUCTION

  • At a fixed token budget, adding replicas often degrades performance relative to non-local data-parallel training.The paper frames this replica-scaling challenge in the setting of geographically separated clusters with poor connectivity and communication bottlenecks.Source: §1 · p. 1
  • Optimizers that whiten pseudogradient covariance achieve the strongest worker scaling; SPlus is the best proxy for whitening.SOAP’s stronger absolute performance sometimes yields the best result despite poorer worker scaling, and all inner optimizers show a strong scaling trend relative to DP.Source: §1 · p. 1
  • Nonlinear adaptive optimizers can create a discrepancy on the first step because transforming local gradients before averaging differs from transforming the averaged gradient.This can occur even with identical starting states and identically distributed data, motivating analysis of how inner optimizers transform gradient noise.Source: §1 · p. 2
  • The paper identifies pseudogradient covariance as important and links inner-optimizer choice to maximum discrepancy from the DP update direction.The study also examines how second-moment estimation behaves at low batch sizes.Source: §1 · p. 2
  • Appropriately scaling β2 for second-moment optimizer states can let Adam-inner IsoLoCo outperform the more computationally expensive MuLoCo.The paper validates the scaling finding through theoretical analysis and experiments.Source: §1 · p. 2
  • At 96 replicas, SOAP or SPlus achieves up to 6% less degradation relative to DP than existing methods.The reported improvement is measured as degradation relative to DP.Source: §1 · p. 2

2 RELATED WORK

Related work establishes local optimization as a way to reduce communication, while documenting scaling challenges and several optimizer and aggregation variants. This paper situates its analysis around local state estimation, optimizer transforms, and small-batch second-moment scaling.

2 RELATED WORK

  • DiLoCo scaling laws document that adding replicas at a fixed token budget degrades quality, with optimal hyperparameters shifting as replica count grows.The paper complements those laws by tracing degradation to local state estimation and nonlinear optimizer transforms before aggregation.Source: §2 · p. 2
  • Local SGD and federated averaging reduce communication through multiple local steps between averaging events, while DiLoCo adapts this approach to LLM pretraining.DiLoCo uses AdamW internally and Nesterov momentum on aggregated pseudogradients; later work relaxes synchrony or overlaps and shards information exchange.Source: §2 · p. 2
  • Adam can train stably at very small batches when β2 keeps the second-moment half-life fixed in tokens rather than steps.The related-work discussion notes that replica scaling induces this small-batch regime for each DiLoCo replica, and transfers the rule to replica scaling.Source: §2 · pp. 2–3
  • Prior distributed-training methods vary inner and outer optimization, while optimizer comparisons report horizon- and batch-dependent rankings.Examples include MuLoCo’s Muon inner optimizer and IsoLoCo’s isotropization of the aggregated pseudogradient; a very-large-batch study reports a slight stabilized-SOAP edge over Muon with added memory and numerical-stability engineering.Source: §2 · p. 3

3 BACKGROUND: THE DILOCO FRAMEWORK

The DiLoCo framework combines independent local optimization with averaging and a shared outer update. Its inner optimizer state remains local to each replica.

3 BACKGROUND: THE DILOCO FRAMEWORK

  • Each round starts replicas from shared parameters, runs local inner-optimizer steps on separate data streams, averages their displacements, and applies an outer update.The inner optimizer states are estimated locally and are never communicated.Source: §3 · p. 3
  • DiLoCo uses AdamW as its inner optimizer and Nesterov momentum as its outer optimizer.MuLoCo replaces the inner optimizer with Muon, while IsoLoCo isotropizes the averaged displacement before the outer step.Source: §3 · p. 3
  • Algorithm 1 specifies initial parameters, replica count, inner-step count, inner optimizer and state, outer optimizer and state, and global batch size.Each replica processes B/R tokens per local minibatch, performs H inner steps, and contributes its displacement to the uniform average before the outer update.Source: §3 · p. 3

4 INNER OPTIMIZERS IMPACT ON REPLICA SCALING

The section analyzes how inner optimizers shape local-update means and covariances relative to DP updates. It argues that whitening can control stochastic error, while second-moment estimation and temporal covariance qualify the behavior of particular optimizers.

4.1 BACKGROUND ON POPULAR OPTIMIZERS

  • The optimizer framework makes update transformations explicit so their means and covariances can be compared; SPlus retains momentum and a historical basis while signing projected entries.The section includes popular adaptive and matrix-based optimizers, including Muon and Shampoo-family methods.Source: §4.1 · p. 4

4.2 IMPACT OF COVARIANCE STRUCTURE ON LOCAL OPTIMIZATION

  • Whitening local pseudo-gradients toward a common isotropic covariance controls stochastic error by bounding the aggregate’s largest directional variance.The analysis treats mean mismatch separately; whitening alone does not control it.Source: §4.2 · p. 6
  • Whitening replicas toward the same spherical covariance makes standard DiLoCo’s uniform averaging minimize stochastic directional variance when replica covariances match.Uniform weights minimize the quantity when replicas have equal directional variances, and in every direction when their covariances are identical.Source: §4.2 · p. 6

4.3 IMPACT OF OPTIMIZER TRANSFORMATIONS ON UPDATE COVARIANCE

  • Adam’s coordinate normalization preserves input correlations, canonical Shampoo does not generally whiten, and Muon is isotropic only under strong distributional symmetry assumptions.Canonical Shampoo’s output covariance condition number is the square root of the input’s; Muon’s polar transform does not generally guarantee isotropy.Source: §4.3 · pp. 6–7
  • SOAP can whiten update covariance when the input covariance has the assumed Kronecker structure, the rotation decorrelates coordinates, and their variances are known exactly.The result concerns idealized conditions and exact variance knowledge.Source: §4.3 · p. 7
  • Under imperfect decorrelation, SPlus’s sign transform makes update covariance at least as close to identity as SOAP’s, without coordinatewise variance estimates.This comparison assumes the same rotation and exact coordinate standard deviations for SOAP.Source: §4.3 · p. 7
  • Whitening individual local updates does not by itself whiten the complete pseudo-gradient, because temporal covariances can introduce anisotropy.Approximate isotropy can be maintained when the aggregate contribution of temporal covariances is sufficiently small.Source: §4.3 · p. 8

4.4 SINGLE-STEP MEAN MISMATCH AND SECOND-MOMENT ESTIMATION

  • Scaling β2 with replica count keeps the stated one-step Adam mean-mismatch bound bounded rather than growing linearly with replicas.This conclusion holds under the stated fixed-parameter assumptions and when β2(R) stays bounded away from zero.Source: §4.4 · p. 9

5 EMPIRICAL EVALUATION

The experiments compare inner optimizers across model and replica scales, then examine second-moment tuning and update-distribution properties. They report stronger replica scaling for whitening-based methods and relate scaling patterns to covariance structure and second-moment estimation.

5.2 SECOND-MOMENT MEMORY AND REPLICA SCALING

  • Experiments trained 178M- and 512M-parameter LLaMA models on DCLM, evaluated on a held-out 100M-token validation split, using respective global batches of 524K and 2M tokens and training budgets of approximately 3.55B and 10.25B tokens.Unless stated otherwise, replicas took H = 30 local steps between aggregations, and all local optimization methods used a Nesterov SGD outer optimizer.Source: §5 · p. 10

5.1 COMPARING INNER OPTIMIZERS IN REPLICA SCALING

  • At 178M and 512M, SPlus degrades less than SOAP, Muon, and AdamW at high replica counts, while SOAP has stronger DP performance; SPlus’s loss change is +3.1% at R = 64 and +2.1% at R = 96.At 512M, SPlus overtakes SOAP at R = 64 and R = 96; the figures show final validation loss and change from each optimizer’s own DP baseline.Source: §5.1 · pp. 10–11, Figure 2 · p. 10, Figure 3 · p. 11
  • Replica-scaling performance improves with model and data scale across SOAP, SPlus, and Muon, consistent with prior observations for AdamW and Muon.The paper compares results across 178M and 512M scales.Source: §5.1 · p. 11
  • At H = 1, SOAP and SPlus remain substantially more robust than Adam, Muon, and Shampoo, and optimizer choice and replica count affect performance more than the number of local steps.The experiment uses a single local update and pure parameter averaging to isolate the adaptive optimizer’s update from multi-step drift.Source: §5.1 · pp. 11–12

5.2 SECOND-MOMENT MEMORY AND REPLICA SCALING

  • At 512M and R = 64, changing only β2 reduces DiLoCo loss from 2.884 to 2.849 and IsoLoCo loss from 2.734 to 2.702, approximately 1.2% in each case.Tuning β2 benefits DiLoCo and IsoLoCo, and IsoLoCo can outperform MuLoCo after this control despite using an AdamW inner optimizer.Source: §5.2 · pp. 12–13, Figure 5 · p. 13

5.3 ANALYSIS: COVARIANCE STRUCTURE AND MEAN MISMATCH

  • Update anisotropy broadly follows replica-scaling performance: after one local step, a1 ranges from 35 for AdamW to 0.43 for SPlus, with Muon and SOAP between them.After 30 steps, SOAP and SPlus remain less anisotropic than Muon and AdamW; Muon is more isotropic than AdamW.Source: §5.3 · p. 14, Table 4 · p. 14
  • Along selected high-noise directions, mean mismatch averages only 7.5–19% of the stochastic term, supporting the analysis’s focus on stochastic aggregation error.The comparison is made along directions selected for high update variance.Source: §5.3 · p. 14
  • Adam and Muon have high curvature in their highest-noise directions as replicas scale, whereas SOAP and SPlus redistribute noise more uniformly through whitening.The paper identifies high noise along high-curvature directions as potentially damaging to loss.Source: §5.3 · pp. 14–15

6 CONCLUSION

The paper attributes differences from the DP pseudogradient to mean mismatch and stochastic error, and finds that inner-optimizer covariance properties align with replica-scaling performance. SPlus scales best relative to DP at higher worker counts, while SOAP achieves the best DP performance.

6 CONCLUSION

  • SPlus achieves the strongest performance at higher worker counts and the best overall scaling relative to DP, while SOAP has the best DP performance.SOAP also appears on trend to become the best inner optimizer at scale.Source: §6 · p. 15
  • Mean mismatch and stochastic error are the two primary sources of differences from the DP pseudogradient.The analysis connects optimizer covariance properties and these errors to observed performance trends.Source: §6 · p. 15
  • Measured mean-mismatch and stochastic-error quantities correlate directly with the observed optimizer performance trends.Table 6 reports loss curvature along a fitted high-noise direction relative to random directions; its ranges span replicas 0 and 1, not confidence intervals.Source: §6 · p. 15, Table 6 · p. 15

A SUPPORTING THEORETICAL RESULTS

Under rotational proposal assumptions, the ideal Muon update retains the proposal’s polar factor, and independent replicas have a uniform averaged update covariance. This covariance does not depend on differences in proposal singular-value distributions.

A.1 MUON COVARIANCE UNDER ROTATIONAL PROPOSALS

  • Muon’s polar update discards the proposal’s realized singular values and keeps its polar factor.For any positive scalar a, polar(aM) = polar(M), so positive rescaling of the whole proposal leaves the update unchanged.Source: §A.1 · p. 18
  • Independent replicas have averaged update covariance γ^2 I_mn/(dmax R), even when proposal singular-value distributions differ.The result assumes full-rank proposals almost surely and the stated one-sided rotational invariance condition.Source: §A.1 · p. 18
  • The covariance result follows from rotational invariance, which makes the polar-factor columns’ cross-covariances vanish and their diagonal covariances isotropic.The proof treats the m ≥ n case by left invariance and the m < n case by right invariance.Source: §A.1 · p. 18

A.2 PROOF OF PROPOSITION 2 (SPLUS TO DDP COMPARISON)

The proof compares the SPlus local update with signed DDP through its mean and covariance. It establishes a covariance bound while cautioning that first-order mean agreement does not imply lower variance than signed DDP.

A.2 PROOF OF PROPOSITION 2 (SPLUS TO DDP COMPARISON)

  • In the shared-state model, the expected local average equals the single-replica expectation, matching the signed DDP momentum-coordinate mean.The passage states this equality for the shared-state model, where momentum is affine before the sign operation.Source: §A.2 · p. 18
  • For R ≥ 1, the integrand is nonnegative and at most u^2(1 − R^-1)/2, by the unit Lipschitz bound for e^-x on x ≥ 0. The passage then applies integration and odd symmetry.Source: §A.2 · p. 18
  • The calibrated local covariance exceeds signed DDP covariance in positive-semidefinite order, with operator-norm difference at most c^2 max_j a_j^2.At zero signal, the covariances agree exactly.Source: §A.2 · p. 18

A.3 ADAM/SOAP SECOND MOMENTS AND THE MEAN COMPARISON

Local Adam and SOAP second moments differ from DDP because replica gradients scatter around their average. The appendix also gives a finite-noise mean bound and a β2 scaling rule that controls an upper bound without guaranteeing monotonic bias reduction.

A.3 ADAM/SOAP SECOND MOMENTS AND THE MEAN COMPARISON

  • Locally updated Adam and SOAP second moments already differ from DDP before parameter drift, because averaging squared replica gradients differs from squaring their average.The discrepancy is the replica-gradient scatter; the SOAP identity assumes a common basis fixed before fresh gradients.Source: §A.3 · p. 19
  • Under the stated batch model, the expected replica-gradient scatter is (R − 1)Σ_g.Here Σ_g is the covariance of the DDP-averaged gradient, with independent replica gradients and Cov(g_r) = RΣ_g at fixed B.Source: §A.3 · p. 19

A.3.1 FINITE-NOISE MEAN BOUND AND SECOND-MOMENT MEMORY

  • The finite-noise mean bound does not replace a random denominator with its expectation or require a small-noise limit.It follows by integrating the Gaussian variance identity over 1 ≤ t ≤ R.Source: §A.3.1 · p. 19
  • The token-half-life scaling rule raises β2 as replica count R grows, keeping the decay per fixed number of local examples unchanged.It ensures (R − 1)(1 − β2(R)) ≤ c0, controlling the upper bound with fixed incoming state; it does not force bias to zero or prove its magnitude decreases monotonically.Source: §A.3.1 · p. 20

A.4 BASIS AND SCALE ERRORS

The analysis separates two sources of error in SOAP’s whitening: imperfect decorrelation by the basis and inaccurate coordinate-scale estimates. Its covariance results apply to a frozen-state, idealized one-step optimizer core, not the full fresh-gradient SOAP step.

A.4 BASIS AND SCALE ERRORS

  • SOAP whitening can err both because its basis leaves correlations and because its estimated coordinate scales are inaccurate.The section analyzes these errors separately and their effects on covariance matching and the stochastic term in Equation (10).Source: §A.4 · p. 20
  • The covariance calculations assume centered proposals and bases and denominators fixed before the proposal is drawn.They describe an idealized one-step optimizer core, omitting a common learning-rate factor and numerical stabilizer.Source: §A.4 · p. 20
  • The frozen-state identities do not describe the full fresh-gradient Adam or SOAP step.Practical SOAP uses gradient second moments for denominators while the numerator contains momentum; refreshing state from the same gradient couples the map to its input.Source: §A.4 · p. 20

A.4.1 SCALE ERRORS AND CORRELATIONS LEFT IN THE ESTIMATED BASIS

The analysis distinguishes scale-estimation error from correlations left by the estimated basis in the covariance of the returned update. In the centered model, SPlus avoids the estimated-scale error, while its covariance still depends on residual correlations.

A.4.1 SCALE ERRORS AND CORRELATIONS LEFT IN THE ESTIMATED BASIS

  • SOAP’s whitening error separates into coordinate-scale error and residual correlations after normalization.These are distinct sources: either can remain even when the other is eliminated.Source: §A.4.1 · p. 20
  • For centered Gaussian proposals, the sign transform’s covariance follows the entrywise arcsine relation and also holds for positive scale mixtures.A positive scalar multiplier leaves signs unchanged, so the stated identity extends to the specified scale-mixture model.Source: §A.4.1 · p. 21
  • SPlus removes the estimated-scale term in the centered model, but its covariance still depends on correlations left by the basis.When the residual correlation matrix is identity, SPlus has identity covariance for every worker, including workers with different bases and input variances.Source: §A.4.1 · p. 21
  • Underestimating a variance can amplify SOAP’s whitening error through the reciprocal scale estimate.Rotating the result back preserves the error norm, although it can change the directions of residual anisotropy.Source: §A.4.1 · p. 21

B TRAINING SETUP AND HYPERPARAMETERS

The training setup uses decoder-only LLaMA models and holds global batch size fixed within each model scale, making local batch size depend on replica count.

B TRAINING SETUP AND HYPERPARAMETERS

  • Experiments train decoder-only LLaMA models on DCLM and evaluate on a held-out 100M-token validation split.Training uses sequence length 2048.Source: B TRAINING SETUP AND HYPERPARAMETERS · p. 21
  • Within each model scale, the global batch is fixed, so local batch equals global batch divided by R.The table caption specifies that global batch is measured per inner step, not per communication round.Source: B TRAINING SETUP AND HYPERPARAMETERS · p. 21, Table 7 · p. 22
  • Unless stated otherwise, local training uses H = 30 steps and a Nesterov SGD outer optimizer; H = 1 controls use uniform averaging without outer momentum.Validation-loss percentage changes use each optimizer’s own DP reference and are computed before rounding.Source: B TRAINING SETUP AND HYPERPARAMETERS · p. 22
  • Learning rates and outer momentum are tuned for each optimizer, with sweeps at R = 8, 64, and 96.Each sweep uses at least 10 grid points; hyperparameters are extrapolated for R = 16 and 32.Source: B TRAINING SETUP AND HYPERPARAMETERS · p. 22

C ANALYSIS PROTOCOLS AND ESTIMATORS

The appendix describes conditional replay probes for update covariance, DP-referenced mean and noise, and curvature. These analyses use checkpoint-based measurements and state explicit assumptions and interpretation limits.

C.1 CONDITIONAL COVARIANCE REPLAY

  • The conditional covariance probe estimates minibatch-induced local-update variability at a trained checkpoint, not covariance across successive training-trajectory updates.It uses each optimizer’s own trained checkpoint.Source: §C.1 · p. 22, C ANALYSIS PROTOCOLS AND ESTIMATORS · p. 22

C.1.1 CONDITIONAL REPLAY PROTOCOL

  • For each of 64 saved replica states, the replay restores the same parameters and state before each rollout and records displacements after h = 1 or 30 steps.The protocol uses n = 16 paired rollouts per replica and horizon.Source: §C.1.1 · p. 22

C.1.2 TRACE ESTIMATORS FROM FULL-COORDINATE INNER PRODUCTS

  • The trace estimators use inner products across all selected coordinates, and exclude diagonal self-products to avoid upward bias from fourth-moment terms.The estimated covariance overlap between replicas is a scale-dependent Frobenius overlap, not cross-covariance between their stochastic updates.Source: §C.1.2 · p. 23

C.1.3 SCALED SUM OF LOCAL COVARIANCES

  • The scaled-sum analysis measures covariance shape independently of total noise magnitude, but its ratio and square-root estimates are not themselves unbiased.Finite-sample clipping can return zero even when the population covariance is anisotropic.Source: §C.1.3 · p. 24

C.1.4 SCOPE AND INTERPRETATION

  • Cross-method probe differences can reflect both update rules and training-reached states, so they are not same-state causal ablations.Interpreting the estimated covariance as that of the actual average also requires vanishing conditional cross-replica noise covariances; the single-checkpoint geometry probe does not identify causal contribution to final loss.Source: §C.1.4 · p. 24

C.2 DP-REFERENCED DIRECTIONAL MEAN AND NOISE

  • The DP-referenced analysis fits local-update scaling and a high-noise direction on fitting repetitions, then evaluates paired local–DP differences on independent repetitions.Its stochastic correction retains local–DP cross-covariance.Source: §C.2 · pp. 24–25
  • At R = 8 and H = 30, DP-referenced mean/noise ratios are 6.2% for SPlus, 2.8% for SOAP, 10% for Muon, and 6.7% for AdamW.These are equal-weight averages over 63 hidden matrices, with the ratio computed per matrix before averaging.Source: §C.2 · p. 25

C.3 CURVATURE ALONG FITTED HIGH-NOISE DIRECTIONS

  • The curvature probe selects directions with high noise variance and evaluates Hessian curvature at each method’s own checkpoint.It measures one-step local updates, not H = 30 aggregate updates, and does not estimate tr(HC).Source: §C.3 · p. 25
Loading 2610.19064v1…