Paper deep dive
How Inner Optimizers Shape Replica Scaling
Eugene Belilovsky, Stefan Horoi, Edouard Oyallon, Benjamin Thérien
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 94%
Last extracted: 10/15/2026, 4:55:02 AM
Summary
This paper investigates how inner optimizers affect replica scaling in distributed large language model (LLM) training, specifically within the DiLoCo framework. It identifies that the covariance structure of pseudo-gradients and the impact of batch size on noise scale are critical factors. The authors demonstrate that whitening-based preconditioning (e.g., SPlus, SOAP) significantly improves replica scaling performance compared to other optimizers like AdamW or Muon by controlling directional variance and promoting covariance matching across replicas.
Entities (9)
Relation Signals (7)
SPlus → improves → Replica Scaling
confidence 95% · SPlus is the best proxy for whitening... using a SOAP or SPlus inner optimizer can lead to very effective replica scaling
DiLoCo → uses → AdamW
confidence 95% · DiLoCo (Douillard et al., 2023) instantiates the inner optimizer as AdamW
SOAP → improves → Replica Scaling
confidence 94% · using a SOAP (Vyas et al., 2025) or SPlus (Frans et al., 2025) inner optimizer can lead to very effective replica scaling
MuLoCo → uses → Muon
confidence 94% · MuLoCo (Thérien et al., 2025) uses Muon as the inner optimizer
Whitening → controls → Directional Variance
confidence 93% · whitening controls the largest directional variance, while also promoting covariance matching across replicas
Covariance Structure → impacts → Replica Scaling
confidence 92% · identify two important factors that impact replica scaling: (1) the covariance structure of the pseudo-gradients
IsoLoCo → outperforms → MuLoCo
confidence 90% · IsoLoCo method... can outperform the much more computationally expensive MuLoCo
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract: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.
Tags
Links
- Source: https://arxiv.org/abs/2610.19064v1
- Canonical: https://arxiv.org/abs/2610.19064v1
Trouble viewing inline? Open PDF directly →
Full Text
103,870 characters extracted from source content.
Expand or collapse full text
How Inner Optimizers Shape Replica Scaling Eugene Belilovsky Stefan Horoi Edouard Oyallon Benjamin Thérien Affiliation: Mila Concordia University Université de Montréal CNRS Abstract 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−968-96 replicas), with performance relative to DP improving as parameters and data are scaled. 1 Introduction The computational scale of language-model training increasingly exceeds the capacity of a single site. Training across geographically separated clusters can pool otherwise fragmented resources, but the connectivity between these clusters is poor, leading to significant communication bottlenecks. Local optimization methods address this constraint by allowing replicas to optimize independently for several steps between parameter or pseudo-gradient averages (Douillard et al., 2023; Douillard et al., 2025; Liu et al., 2024) often drastically reducing the frequency of communication. A central difficulty in practical local optimization is replica scaling: at a fixed token budget, increasing the number of replicas often degrades the final model performance (Charles et al., 2025) relative to its non-local data-parallel (DP) counterpart. This can limit the broad applicability of these methods. It has been shown that the replica scaling can improve with increasing model size (Charles et al., 2025). Recent empirical results show that modifying either the inner optimizer (Thérien et al., 2025) or the outer aggregation rule (Horoi et al., 2026) can further improve the scalability and efficiency of distributed training. However, the impact of inner optimizers on scaling is not well understood. More broadly, the primary sources of divergence from standard non-local training are not well understood. Figure 1: Local optimization performance as replicas are scaled. Each local optimizer is compared to its own DP baseline, to eliminate confounding replica scaling with absolute loss. Optimizers that whiten the pseudogradient covariance achieve the strongest worker scaling. SPlus is the best proxy for whitening. We note that SOAP’s stronger absolute performance sometimes leads to the best results despite its poorer worker scaling. Notably, we also observe a strong scaling trend relative to DP for all inner-optimizers. Classical analyses of local SGD and federated averaging study the discrepancy accumulated over multiple local steps, often with heterogeneous client data. The primary source of performance degradation is often attributed to client drift or other notions of divergence of model replicas along their trajectory (Karimireddy et al., 2020). Much of this work and theory (Stich, 2019; Karimireddy et al., 2020; Yu et al., 2019) focuses on SGD and momentum SGD as the inner optimizer, a single local step from common parameters and optimizer states that commute with averaging: averaging the updates gives the same result as updating with the averaged gradient. Nonlinear adaptive inner optimizers need not satisfy this property and are much less well understood in the context of this literature. Notably, a discrepancy can arise on the first step, even with identical starting states and identically distributed data, because transforming each local gradient before averaging differs from transforming the averaged gradient. Understanding replica scaling therefore requires examining how the inner optimizer transforms gradient noise, in addition to how parameters drift over successive local steps. In this work, we focus on a broad range of popular adaptive or preconditioner based optimizers and identify two critical properties that can greatly impact scaling: (a) the covariance structure of the pseudo-gradient induced by the local optimization (b) how the optimizers deal with statistical estimation of second moments in the low batch size setting. Our contributions can be summarized as follows: • We demonstrate, both empirically and theoretically, the importance of the pseudogradient covariance structure induced by the choice of inner optimizer. Specifically, we link the choice of inner optimizer to the maximum discrepancy with respect to the DP update direction. • We identify the importance of appropriately scaling β2 _2 when using second-moment optimizer states, validating our findings through theoretical analysis and experiments. This reveals that the recently introduced IsoLoCo method (Horoi et al., 2026), using an Adam inner optimizer, can outperform the much more computationally expensive MuLoCo (Thérien et al., 2025). • Our results demonstrate that using a SOAP (Vyas et al., 2025) or SPlus (Frans et al., 2025) inner optimizer can lead to very effective replica scaling , obtaining significantly improved performance relative to existing methods from the literature (as much as 6%6\% less degradation relative to DP at 96 replicas). 2 Related Work Local-step methods and DiLoCo. Local SGD (Stich, 2019) and federated averaging (McMahan et al., 2017) reduce communication by taking several local steps between averaging events. DiLoCo (Douillard et al., 2023) adapts this to LLM pretraining by utilizing a large number of inner steps, AdamW (Loshchilov & Hutter, 2019) as the inner optimizer, and SGD with Nesterov momentum on the aggregated pseudo-gradients as the outer optimizer. Follow-ups relax synchrony (Liu et al., 2024) and overlap or shard the exchange of information between replicas (Douillard et al., 2025). Charles et al. (2025) establish scaling laws for DiLoCo and document the central difficulty we study: adding replicas at a fixed token budget degrades quality, with the optimal hyperparameters shifting as R grows. Our analysis complements these scaling laws by tracing the degradation back to local state estimation and to the nonlinear optimizer transforms applied before aggregation. Local methods can also be complemented with local optimization by using Sparse parameter averaging (Reed & Baioumy, 2025, SPARTA;) or Top-k sparsification SparseLoCo (Sarfi et al., 2025). IsoLoco Horoi et al. (2026) views DiLoCo’s outer step as task-arithmetic model merging and replaces plain pseudo-gradient averaging with Iso-C merging (Marczak et al., 2025). For each weight matrix, it takes the SVD of the averaged pseudo-gradient and replaces all singular values with their mean, giving σ¯UV⊤ σUV . It then applies the usual Nesterov momentum outer step. This was shown to improve worker scaling over DiLoCo (Horoi et al., 2026). Small-batch training and adaptive-optimizer scaling rules. Scaling rules relate optimizer hyperparameters to batch size (Malladi et al., 2022); most directly, Marek et al. (2025) show that Adam trains stably at very small batches provided β2 _2 is set so that the second moment’s half-life is fixed in tokens rather than in steps. Replica scaling in DiLoCo induces exactly this small-batch regime for each individual replica. Section 4.4 transfers the β2 _2 scaling rule to replica scaling; to our knowledge this hasn’t been studied has not been studied in the local-step setting previously. Spectral and second-order optimizers. Shampoo (Gupta et al., 2018) preconditions with Kronecker-factored curvature estimates; SOAP (Vyas et al., 2025) runs Adam in Shampoo’s eigenbasis; Muon (Jordan et al., 2024) orthogonalizes matrix updates with no accumulated second-moment statistics. SPlus (Frans et al., 2025) retains Shampoo’s historical eigenbasis but signs momentum in that basis, avoiding cached eigenvalue magnitudes and SOAP’s additional coordinate-wise second moment. SPlus uses this instant-sign transformation independently in local DiLoCo replica bases; our focus is its distributed covariance and replica-scaling behavior. Controlled single-stream comparisons find the ranking horizon-dependent (Wen et al., 2025), while a recent multi-billion-parameter study reports a slight edge for a stabilized SOAP over Muon at very large batch, at the price of memory and numerical-stability engineering (Khona et al., 2026). In the DiLoCo setting, MuLoCo (Thérien et al., 2025) uses Muon as the inner optimizer and reports more directionally consistent pseudo-gradients as R grows, while IsoLoCo (Horoi et al., 2026) treats outer aggregation as model merging and isotropizes the aggregated pseudo-gradient. Section 4 develops a complementary criterion: whether the stochastic update geometry emitted by each replica remains well conditioned under averaging. 3 Background: the DiLoCo Framework The methods in this paper can be seen as an instance of Algorithm 1 which is a generalization of local optimization from Reddi et al. (2021); Douillard et al. (2023) that utilizes both inner (on the local gradients) and outer optimization (on the pseudo-gradients). R replicas start each round from the shared parameters, each runs H steps of an inner optimizer on its own data stream. The resulting displacements r _r are uniformly averaged, and an outer optimizer applies the aggregate ¯ to the shared parameters from the start of the round. DiLoCo (Douillard et al., 2023) instantiates the inner optimizer as AdamW and the outer as Nesterov momentum. For example MuLoCo (Thérien et al., 2025) replaces the inner optimizer with Muon. IsoLoCo (Horoi et al., 2026) isotropizes ¯ before the outer step. The inner optimizer states srs_r are never communicated: each is estimated locally from its replica’s data stream. Algorithm 1 The DiLoCo framework. All methods studied here are instances of this algorithm. 1: initial parameters (0) θ^(0); replicas r∈1,…,Rr∈\1,…,R\; inner steps H; inner optimizer Inner with per-replica state srs_r; outer optimizer Outer with state S; global batch size B; 2: for outer round t=0,…,T−1t=0,…,T-1 do 3: for all replicas r in parallel do 4: r←(t) θ_r← θ^(t) 5: for h=1,…,Hh=1,…,H do 6: ←∇ℒ(r,xr,h) g← _ θL( θ_r;x_r,h) ⊳ local minibatch: B/RB/R tokens 7: (r,sr)←Inner(r,,sr)( θ_r,\,s_r)← Inner( θ_r,\, g,\,s_r) ⊳ srs_r is never synchronized 8: end for 9: r←(t)−r _r← θ^(t)- θ_r 10: end for 11: ¯←1R∑r ← 1R _r _r ⊳ uniform average 12: ((t+1),S)←Outer((t),¯,S)( θ^(t+1),\,S)← Outer( θ^(t),\, ,\,S) ⊳ Nesterov momentum in DiLoCo 13: end for 4 Inner Optimizers Impact on Replica Scaling The transformations applied to local updates by local optimizers can have a significant impact on the final aggregated pseudo-gradient. We investigate how the choice of inner optimizer shapes the distribution of the resulting pseudo-gradients and their deviation from the corresponding non-local updates. Section 4.1 introduces the optimizers considered in this work, Section 4.2 separates the update error (relative to non-local steps with the same base optimizer) into a mean mismatch term and a stochastic error term. It shows how whitening controls the largest directional variance, and how matching covariance structures across replicas makes uniform averaging optimal. Section 4.4 specializes this framework to one step from common state: under the stated Gaussian model, SPlus admits a scalar calibration that matches the DP mean to first order in the global proposal’s signal-to-noise ratio. Finally, Section 4.4 bounds the Adam mean mismatch caused by fresh-gradient second moments and explains how appropriately adjusting β2 _2 limits that bound as local batches shrink. It then isolates a complementary advantage of SPlus: retaining the estimated basis without relying on coordinatewise scale estimates to control the covariance term. 4.1 Background on Popular Optimizers In this section we present a number of popular optimizers using a unified notation to facilitate their common analysis and comparison. In particular, we include recent matrix-based optimizers such as Muon and those in the Shampoo family. The optimizers differ in how they retain or remove the magnitudes of a momentum proposal before averaging across replicas. We write these transformations explicitly so their mean and covariance can be studied in the same framework. Consider a matrix parameter with stochastic gradient ∈ℝm×n G ^m× n. An optimizer returns a direction U and an updated state after observing G; its parameter displacement is =ηvec() =η\,vec( U). We use subscripts t−1t-1 and t for the incoming and updated optimizer states, respectively, and suppress the step index on the gradient and returned direction. The first-moment proposal is t=β1t−1+(1−β1). M_t= _1 M_t-1+(1- _1) G. (1) We compare the following matrix-update cores, using standard EMA momentum to isolate their normalizations. For simplicity we omit standard deterministic bias corrections and weight decay from the presentation. The matrix-size scale γ of Muon is shown explicitly. All decay coefficients lie in [0,1)[0,1); numerical damping is denoted by ϵ>0ε>0. Adam and Shampoo. Adam uses an entrywise second moment in parameter coordinates (Kingma & Ba, 2015; Loshchilov & Hutter, 2019): A,t=β2A,t−1+(1−β2)⊙2,Adam=t⊘(A,t+ϵ). V_A,t= _2 V_A,t-1+(1- _2) G 2, U_ Adam= M_t ( V_A,t+ε). (2) Where U is the update. Here, ⊙ and ⊘ denote elementwise multiplication and division, respectively, with ⊙2 G 2 representing the elementwise square. Shampoo instead maintains left and right Gram factors (Gupta et al., 2018): t=β3t−1+(1−β3)⊤,t=β3t−1+(1−β3)⊤,Shampoo=(t+ϵ)−1/4t(t+ϵ)−1/4. split L_t&= _3 L_t-1+(1- _3) G G , R_t= _3 R_t-1+(1- _3) G G,\\ U_ Shampoo&=( L_t+ε I)^-1/4 M_t( R_t+ε I)^-1/4. split (3) The above formulation is the canonical version of Shampoo. We note that some variants consider half-power instead of quarter (Vyas et al., 2025). We also note variants with grafting or additional normalizations are often used, which can impact the interpretation of the underlying preconditioner (Lin et al., 2026). SOAP and SPlus. Let L,R Q_L, Q_R be orthogonal eigenvector matrices obtained from the historical Gram factors above. Their action is the matrix rotation ~=L⊤R,~t=L⊤tR. G= Q_L G Q_R, M_t= Q_L M_t Q_R. (4) SOAP (ShampoO with Adam in the Preconditioner’s eigenbasis) performs Adam-style normalization in this product basis (Vyas et al., 2025): S,t=β2S,t−1+(1−β2)~⊙2,SOAP=L[~t⊘(S,t+ϵ)]R⊤. split V_S,t&= _2 V_S,t-1+(1- _2) G 2,\\ U_ SOAP&= Q_L [ M_t ( V_S,t+ε) ] Q_R . split (5) In this work, we consider a signed update in a Shampoo basis, a variant of the recently introduced SPlus (Frans et al., 2025) optimizer, which does not utilize iterate averaging. This optimizer retains the momentum and historical basis but signs each projected entry: SPlus=Lsign(~t)R⊤,sign(0)=0. U_ SPlus= Q_L ( M_t) Q_R , (0)=0. (6) This approach can be seen as a simplified version of SOAP that removes the dependence on the second moment. In the following sections, we will further illustrate how this optimizer naturally arises in the context of local optimization with replica scaling. SOAP’s rotated normalization and SPlus sign transform have a whitening interpretation (Vyas et al., 2025; Frans et al., 2025; Lin et al., 2026). Section 4.2 gives the corresponding covariance identities under explicit assumptions and explains their impact on pseudo-gradient aggregation. Muon. For a singular value decomposition t=MMM⊤ M_t= U_M _M V_M on the nonzero singular subspace, ideal Muon returns Muon=γpolar(t)=γMM⊤. U_ Muon=γ ( M_t)=γ U_M V_M . (7) It replaces nonzero singular values by a common value (Jordan et al., 2024). In matched singular-vector coordinates, M⊤tM=M U_M M_t V_M= _M, so at its core Muon also has the form “rotate, remove magnitudes, rotate back.” Unlike SPlus, its basis is determined by the current proposal, not by historical covariance statistics. In practice, Newton-Schulz is used to approximate the polar map. 4.2 Impact of Covariance Structure on Local Optimization We consider the general local optimization setting (under any local optimizer) and separate the error relative to the expected DP update with the same optimizer into two terms: (1) mean mismatch and (2) stochastic error. The covariance analysis in this subsection addresses the stochastic term; Section 4.4 then addresses the mean mismatch. Following the pseudo-gradient interpretation of client displacements in Reddi et al. (2021) and Wang et al. (2024), we treat each replica’s complete local update as a random vector. All expectations and covariances below are conditional on the start-of-round parameters and the incoming optimizer states. Let r∈ℝd _r ^d be the complete H-step displacement returned by replica r. We decompose each local displacement into its conditional mean and a zero-mean stochastic component: r=r+r,r=[r],[r]=0,r=Cov(r). _r= μ_r+ _r, μ_r=E[ _r], [ _r]=0, C_r=Cov( _r). (8) Assume finite second moments, [‖Δr‖22]<∞E[\| _r\|_2^2]<∞ for all r, and zero cross-replica noise covariances, Cov(εr,εs)=Cov( _r, _s)= 0 for all r≠sr≠ s. Independent local data streams, conditional on the start-of-round parameters and optimizer states, are sufficient to ensure zero cross-replica noise covariances. Set ¯=R−1∑r =R^-1 _r _r and ¯=R−1∑r μ=R^-1 _r μ_r. Note that for the typical setting studied in Douillard et al. (2023); Thérien et al. (2025) with i.i.d. data and identical incoming optimizer states, r μ_r is the same for all replicas. We now consider HDP d_H DP: the displacement of the same inner optimizer run for H synchronized steps and same global batch size, in expectation. We allow a deterministic scalar calibration factor αR _R, and define R=αR¯−HDP,agg=Cov(¯)=1R2∑r. b_R= _R μ- d_H DP, C_ agg=Cov( )= 1R^2 _r C_r. (9) For H=1H=1, αR _R can represent a step-size change, for example. We now consider the following decomposition of the error in any direction and the corresponding bound: Proposition 1 (Directional error decomposition). The conditional bias–variance decomposition gives, for every unit vector q, [(⊤(αR¯−HDP))2]⏟error relative to the DP mean=(⊤R)2⏟mean mismatch+αR2⊤agg⏟stochastic error. E\! [ ( q ( _R - d_H DP) )^2 ]_error relative to the DP mean= ( q b_R)^2_mean mismatch+ _R^2 q C_ agg q_stochastic error. (10) Consequently, the largest directional error is λmax(RR⊤+αR2agg)≤‖R‖22+αR2λmax(agg). _ ( b_R b_R + _R^2 C_ agg)≤\| b_R\|_2^2+ _R^2 _ ( C_ agg). (11) Proof. Subtract and add αR¯ _R μ. Using the Rayleigh quotient gives (11). ∎ The second term allows us to analyze the effect of inner optimizers that induce whitened updates. As we show below, whitening helps control the stochastic error by bounding the largest directional variance of the aggregated pseudo-gradient, while also promoting covariance matching across replicas. The mean mismatch term is more challenging to analyze; however in Section 4.4 we derive explicit expressions and bounds for several inner optimizers under a single-local-step setting. Moreover we will show in the empirical results that this term is generally dominated in the largest noise direction. Whitening and directional variance. A linear map P whitens a covariance if ⊤= P P = I (Kessy et al., 2018); a prescribed scalar multiple gives a spherical covariance at another scale. For any covariance C of dimension d and trace τ (fixed total variance), the following inequality holds: λmax()≥τ/d,with equality precisely when =(τ/d). _ ( C)≥τ/d, equality precisely when C=(τ/d) I. (12) Thus, isotropy minimizes the largest directional variance at fixed total variance, precisely the covariance quantity appearing in (11). Moreover, a bound on the maximum directional variance of each replica carries over to the aggregate: λmax(agg)≤1R2∑rλmax(r). _ ( C_ agg)≤ 1R^2 _r _ ( C_r). Consequently, whitening local pseudo-gradients toward a common isotropic covariance controls the stochastic error. Whitening and optimality of uniform averaging. Local update whitening has an additional, pre-aggregation benefit when inner optimizer trajectories produce different update distributions. For the stochastic term alone, a weighted aggregate has directional variance ∑r=1Rwr2⊤r,∑rwr=1. _r=1^Rw_r^2\, q C_r q, _rw_r=1. Uniform weights, which are used by default in standard DiLoCo settings, minimize this quantity when the replicas have equal directional variances (Aitken, 1935); they do so in every direction when 1=⋯=R C_1=·s= C_R. Whitening each replica toward the same spherical target νν I promotes precisely this agreement, since ‖i−j‖≤‖i−ν‖+‖j−ν‖.\| C_i- C_j\|≤\| C_i-ν I\|+\| C_j-ν I\|. Thus, local whitening both controls the worst-direction aggregate noise and makes the default uniform averaging well suited to the individual replica noise geometries. However, it does not control the mean mismatch R b_R, which is analyzed separately in Section 4.4. 4.3 Impact of Optimizer Transformations on Update Covariance In this subsection, we analyze how the transformations induced by different inner optimizers affect the covariance structure of their updates under idealized assumptions, with particular emphasis on their whitening properties. This comparison is motivated by the eigenbasis interpretation of SOAP (Vyas et al., 2025), the whitening motivation of Frans et al. (2025), and the covariance-estimation analysis of Lin et al. (2026). For this comparison, consider a random matrix proposal =t∈ℝm×n M= M_t ^m× n, with =Cov(vec())≻0 =Cov(vec( M)) 0. Let U denote the resulting update matrix, and define d=mnd=mn and =Cov(vec()) C=Cov(vec( U)). We treat optimizer bases and linear preconditioners as fixed independently of the sampled proposal, and define ideal normalization using the proposal’s centered covariance. Write =R⊗L,vec(L⊤R)=⊤vec(). Q= Q_R Q_L, ( Q_L M Q_R)= Q vec( M). We now examine how each optimizer transforms the covariance of the stochastic proposal and under which conditions the resulting update covariance is isotropic. Adam. Suppose each entry of M is divided by its true standard deviation. Writing =diag(Σ11,…,Σdd) D=diag( _11,…, _d), this gives vec()=−1/2vec(),=−1/2−1/2.vec( U)= D^-1/2vec( M), C= D^-1/2 D^-1/2. Every diagonal entry of C is one, but Cij=Σij/ΣiiΣjjC_ij= _ij/ _i _j. Thus, even with access to the true coordinate variances, Adam’s normalization preserves the correlations present in M. Shampoo. Consider the idealized case where the factors equal the population Gram matrices of a zero-mean input: =[⊤] L=E[ M M ] and =[⊤] R=E[ M M]. Assume the covariance is separable, =(⊗)/tr() =( R L)/tr( L), with positive-definite factors. The canonical Shampoo transformation then gives: =−1/4−1/4,=1/2⊗1/2tr(). U= L^-1/4 M R^-1/4, C= R^1/2 L^1/2tr( L). The condition number of the resulting covariance, i.e., the ratio of its largest to smallest eigenvalues, is the square root of that of the input covariance. Thus, canonical Shampoo does not generally whiten the updates, as also discussed in Lin et al. (2026). SOAP. Suppose the fixed rotation ~t=L⊤tR M_t= Q_L M_t Q_R gives uncorrelated coordinates, and let S,t V_S,t denote their true variances: (S,t)ij=Var((~t)ij)( V_S,t)_ij=Var(( M_t)_ij). Dividing each rotated entry by its true standard deviation and rotating back gives =L[~t⊘S,t]R⊤,=. U= Q_L[ M_t V_S,t] Q_R , C= I. The resulting variance is one in every unit direction, not only along the original coordinates. Thus, SOAP can whiten the update covariance when the input covariance has the assumed Kronecker structure, the rotation Q decorrelates the proposal coordinates, and their variances S,t V_S,t are known exactly. SPlus. Suppose vec()vec( M) is zero-mean Gaussian and the fixed rotation above gives uncorrelated entries. These entries are then independent, and their signs are independent, equally likely +1+1 and −1-1 values. Therefore =Lsign(~t)R⊤,=. U= Q_Lsign( M_t) Q_R , C= I. Unlike SOAP, this transformation does not require estimating the coordinate standard deviations. SPlus vs. SOAP Under Imperfect Decorrelation. The preceding analyses of SOAP and SPlus assume that the rotation decorrelates the entries of t M_t, with SOAP additionally having access to the true coordinate standard deviations. We can get further insight by considering an imperfect rotation: fix the estimated L,R Q_L, Q_R before sampling vec(t)∼(,)vec( M_t) ( 0, ), and let C be the correlation matrix of the entries of ~t=L⊤tR M_t= Q_L M_t Q_R. To characterize how SPlus transforms these correlations, we apply the Gaussian arcsine law (Dirksen et al., 2022), which gives the covariance after taking signs. Rotating back yields SPlus=[2πarcsin()]⊤,|2πarcsin(ρ)|≤|ρ|. C_ SPlus= Q [ 2π ( C) ] Q , | 2π (ρ) |≤|ρ|. (13) Here arcsin acts entrywise: the signs have unit variance in the rotated basis, and residual correlations contract, but need not vanish. With the same rotation and exact coordinate standard deviations, SOAP instead gives SOAP=⊤ C_ SOAP= Q C Q : coordinate normalization leaves these correlations unchanged. For both methods, the covariance in the rotated basis has unit diagonal. Its squared Frobenius distance to identity is therefore just the sum of the squared off-diagonal entries. Rotating back preserves this distance, yielding ‖SPlus−‖F2 \| C_ SPlus- I\|_F^2 =∑i≠j[2πarcsin(ij)]2 = _i≠ j [ 2π (C_ij) ]^2 ≤∑i≠jij2=‖SOAP−‖F2. ≤ _i≠ jC_ij^2=\| C_ SOAP- I\|_F^2. Under these assumptions, SPlus produces an update covariance at least as close to the identity in Frobenius norm as SOAP, even when the rotation does not perfectly decorrelate the proposal: its sign transformation contracts residual correlations without requiring coordinatewise variance estimates. Muon. Suppose M is full rank almost surely, and its distribution is unchanged by any orthogonal transformation ↦ M O M when m≥nm≥ n (or ↦ M M O when m<nm<n). For the exact polar transformation with a fixed scale γ>0γ>0, =γpolar(),=γ2max(m,n). U= ( M), C= γ^2 (m,n) I. Thus, Muon produces an isotropic update covariance under strong distributional symmetry assumptions, but its polar transformation does not generally guarantee isotropy. Relation to MuLoCo MuLoCo studies how Muon’s normalized matrix updates affect pseudo-gradient aggregation (Thérien et al., 2025). It argues that large singular values in individual update matrices contribute to misalignment during aggregation, an issue mitigated by using Muon as the inner optimizer. In contrast, our work examines the covariance structure of Muon’s updates, a distribution-level property that provides a complementary perspective and highlights potential limitations of Muon’s normalization. Furthermore, our error decomposition directly connects this covariance structure to the stochastic component of the error relative to the expected DP update. Beyond one local step. Equations (10) and (11) apply to the complete displacement after any number of inner optimization steps H. However, if r,h _r,h denotes the displacement at the h-th local step, the covariance of replica r’s complete displacement is given by r=∑h=1HCov(r,h)+∑h≠kCov(r,h,r,k). C_r= _h=1^HCov( _r,h)+ _h≠ kCov( _r,h, _r,k). Consequently, whitening individual (one-step) local updates does not by itself whiten the complete pseudo-gradient, as temporal covariances between successive updates can introduce anisotropy. Approximate isotropy of the complete pseudo-gradient can nevertheless be maintained when the aggregate contribution of these temporal covariances is sufficiently small. 4.4 Single-Step Mean Mismatch and Second-Moment Estimation We now consider the single-step setting (H=1H=1), where all replicas share the same initial parameters and optimizer states, and analyze the mean mismatch term of the decomposition introduced in Section 4.2. This eliminates multi-step local drift, allowing us to study the remaining mean mismatch R b_R more easily. At fixed global batch size B, let ^r≔vec(r) g_r ( G_r) denote the independent and identically distributed replica gradients, and define ≔[^r],¯≔1R∑r^r,g≔Cov(¯). g [ g_r], g 1R _r g_r, _g ( g). (14) Then Cov(^r)=RgCov( g_r)=R _g. For IID examples, g=/B _g= /B, where is the per-example covariance. Let T(^,s)T( g;s) denote the parameter-update map that takes a gradient g and incoming optimizer state s and returns the corresponding parameter displacement. Then ¯=1R∑rT(^r,s),1DP=[T(¯,s)],R=αR[¯]−1DP. = 1R _rT( g_r;s), d_1 DP=E[T( g;s)], b_R= _RE[ ]- d_1 DP. (15) In this single-step setting, DP applies the optimizer transformation T to the averaged gradient ¯ g, whereas DiLoCo applies T independently to each replica’s gradient before averaging the resulting updates. While an affine T commutes with averaging sample-by-sample, sharing the incoming optimizer states does not generally make an adaptive optimizer’s update transformation affine. We first analyze SPlus, whose sign-based transformation permits explicit expressions for the mean and covariance of its updates. SPlus: a direct bound on the mean term. Fix the shared SOAP basis before the step; this covers a step between refreshes or a lagged-basis implementation. Momentum is still updated by (1). The average of the replica momenta is exactly the DP momentum from the same incoming state. Let m~r,t,j m_r,t,j denote entry j of vec(~r,t)vec( M_r,t), for j=1,…,d=mnj=1,…,d=mn, and assume m~r,t,j=σj(aj+RZr,j),Zr,j∼iid(0,1),σj>0. m_r,t,j= _j(a_j+ R\,Z_r,j), Z_r,j iid N(0,1), _j>0. (16) Thus aja_j is the mean divided by the standard deviation of the DP momentum coordinate, not its replica-local SNR. We use unit step scale in the following proposition and restore other deterministic scales multiplicatively. Proposition 2 (SPlus mean and covariance relative to DP). Under (16), set =(aj)j=1d a=(a_j)_j=1^d, ψ(t)=2Φ(t)−1ψ(t)=2 (t)-1, and c=2/πc= 2/π. With αR=R _R= R, R=[Rψ(/R)−ψ()],‖R‖22≤c236(1−1R)2∑jaj6,Cov(R¯)=Diag(1−ψ(aj/R)2)⊤⪯. split b_R&= Q [ R\,ψ( a/ R)-ψ( a) ],\\ \| b_R\|_2^2&≤ c^236 (1- 1R )^2 _ja_j^6,\\ Cov( R\, )&= QDiag (1-ψ(a_j/ R)^2 ) Q I. split (17) Here ψ acts entrywise. In particular, every directional error in (10) is at most 1+c236(1−R−1)2∑jaj61+ c^236(1-R^-1)^2 _ja_j^6. The proof is given in Appendix A.2. For a small vector a, the mean mismatch is third order and its squared contribution is sixth order. At =0 a=0, the calibrated local aggregate and DP have the same mean and covariance, although their distributions differ. The R R calibration makes the stochastic term order one; it does not retain an additional 1/R1/R variance gain at fixed global batch. Mean mismatch for Adam We now bound the mean mismatch R b_R for Adam. At fixed global batch size, the local gradient variance is R times that of the DP average. Fresh-gradient second moments therefore change the normalization even from a common incoming state. Proposition 3 bounds the resulting mean-update error; Appendix A.3 gives the second-moment identities and the Gaussian variance-differentiation argument (Voigtlaender, 2021). Proposition 3 (Second-moment memory controls a one-step mean bound). In a fixed parameter coordinate j, let the incoming Adam state be (m,v)(m,v), with v>0v>0 and 0<β2<10< _2<1. Omit ϵε, bias correction, and the learning rate, and define the complete same-step scalar map T(x,m,v)=β1m+(1−β1)xβ2v+(1−β2)x2.T(x;m,v)= _1m+(1- _1)x _2v+(1- _2)x^2. (18) Suppose g^r∼(g,Rσ2) g_r (g,Rσ^2) independently, so g¯∼(g,σ2) g (g,σ^2). With αR=1 _R=1, write bR,jb_R,j for this coordinate of the bias R b_R from (15). It obeys |bR,j|≤(1−β2)(R−1)σ22(β2v)3/2(|β1m|+3(1−β1)|g|).|b_R,j|≤ (1- _2)(R-1)σ^22( _2v)^3/2 (| _1m|+3(1- _1)|g| ). (19) For SOAP with a basis fixed before the fresh gradients, the same bound applies to the corresponding coordinate of ⊤R Q b_R. This is a finite-noise bound, proved in Appendix A.3.1. For fixed β2 _2, the upper bound grows linearly with R−1R-1, motivating an increase in β2 _2 as the number of replicas grows. Adapting the scaling rule of Marek et al. (2025) to our setting, we can preserve the second-moment EMA’s half-life in terms of local training examples by setting: β2(R)=β2(R0)R0/R,1−β2(R)=O(R−1). _2(R)= _2(R_0)^R_0/R, 1- _2(R)=O(R^-1). (20) If m,v,g,σm,v,g,σ are held fixed and β2(R) _2(R) follows (20), with β2(R) _2(R) bounded away from zero, the right-hand side remains bounded rather than growing linearly in R. For each R, both routes in this comparison use the same β2(R) _2(R). Scaling second-moment memory therefore controls a bound on the Adam mean mismatch as R grows. Summing the squared coordinate bounds controls ‖R‖22\| b_R\|_2^2, and hence the mean term (⊤R)2( q b_R)^2 for every unit q in (10). The same conclusion holds after SOAP’s orthogonal rotation back to parameter coordinates. This conditional one-step result motivates the matched-β2 _2 experiments in Section 5.2. Scale-estimation error at smaller local batches. Beyond the mean mismatch analyzed above, smaller local batches can also increase the variability of second-moment estimates. Indeed, at fixed β2 _2, smaller local batches produce noisier gradients, which can increase fluctuations in the estimated second moment v v. For SOAP, these fluctuations in v v change the coordinatewise normalization and hence the update covariance. SPlus avoids this source of variability by taking signs in its estimated basis rather than dividing by estimated coordinate scales. 5 Empirical Evaluation We first compare replica scaling across inner optimizers (Sec. 5.1), then isolate the effect of second-moment memory on performance (Sec. 5.2), and finally examine the update distributions behind these results (Sec. 5.3). The analysis connects the observed scaling to the two terms in Equation (10): mean mismatch and stochastic aggregation error. Experimental setup. Our main experiments train 178M- and 512M-parameter LLaMA models on DCLM (Li et al., 2024) with sequence length 2048. We evaluate performance on a held-out 100M-token validation split. The 178M runs use a fixed global batch of 524K tokens and approximately 3.55B training tokens; the 512M runs use 2M tokens per global batch and 10.25B-token budget. Unless stated otherwise, replicas take H=30H=30 local steps between aggregations and all local optimization methods use a Nesterov SGD outer optimizer. Appendix B records the architecture, training budgets, and result-selection protocol. 5.1 Comparing Inner Optimizers in Replica Scaling We compare SOAP, SPlus, Muon, AdamW, and Shampoo as inner optimizers for DiLoCo while scaling the number of replcas. Figures 2 and 3 show both absolute validation loss and the percentage change from each optimizer’s own DP reference (Thérien et al., 2025). The latter allows making a direct comparison between inner optimizers as it relates to their impact on worker scaling without being confounded by absolute performance differences of the optimizers (e.g., accounting for the fact that Moun has lower absolute loss than AdamW). 178M Experiment Scale At 178M scale (Fig 2), we observe, as in Thérien et al. (2025), that Muon has better replica scaling than AdamW, not just DP performance. SOAP has both stronger DP performance and significantly improved replica scaling compared to Muon and Adam. SPlus has the strongest worker scaling, leading to the smallest increase over its own DP reference and the lowest observed loss at large R. At R=64R=64, the loss change is +3.1%+3.1\% for SPlus and +4.8%+4.8\% for SOAP, compared with +8.5%+8.5\% for Muon and +11%+11\% for AdamW. Shampoo improves on AdamW but remains behind SOAP and SPlus. As discussed in our work, it does not assure an approximately whitened update unlike SOAP and SPlus. Thus, both the normalization and the coordinates in which it is applied matter empirically. Section 5.3 examines how this ordering relates to the covariance of the returned updates. Interestingly, at 178M, R=8R=8 SPlus actually improves performance over its DP baseline. This is likely because at R>1R>1 it can mimic a majority vote-based technique (Bernstein et al., 2018). Figure 2: 178M replica scaling. R=1R=1 reports DP performance, while R>1R>1 reports the performance of local optimization with synchronization interval H=30H=30. Left: final validation loss. Right: change in loss from the same optimizer’s own DP (R=1) baseline. We observe that matrix-based optimizers have stronger worker scaling than AdamW and that SPlus achieves the strongest scaling of all relative to its DP baseline. 512M Experiment Scale In Figures 3, we observe a similar pattern. SOAP is, once again, the strongest DP reference and has the lower loss through R=32R=32. SPlus overtakes it at R=64R=64 and R=96R=96. At R=96R=96, SPlus’ los change is +2.1%+2.1\% above its own DP loss, versus +4.4%+4.4\% for SOAP, +7.7%+7.7\% for Muon, and +12%+12\% for AdamW. The crossover is consistent across model scales: SOAP’s DP performance is stronger, whereas SPlus degrades less as replicas are scaled and the local batch size shrinks. Both SPlus and SOAP show drastically better scaling than Adam and Muon, making them practical alternatives when training with high worker counts. Figure 3: 512M replica scaling. R=1R=1 reports DP performance, while R>1R>1 reports the performance of local optimization with synchronization interval H=30H=30. Left: final validation loss. Right: change in loss from the same optimizer’s own DP (R=1) baseline. As in Fig. 2, we observe that matrix-based optimizers have stronger worker scaling than AdamW and that SPlus achieves the strongest scaling of all relative to its DP baseline. Table 1: 178M final validation loss. DP uses R=1R=1; other runs use H=30H=30. Best in each column is bold. Replicas R Optimizer DP 8 16 32 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 Shampoo 2.842 2.881 2.918 2.976 3.069 Table 2: 512M final validation loss. DP uses R=1R=1; other runs use H=30H=30. Best in each column is bold. Replicas R Optimizer DP 8 16 32 64 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 Comparing results across 178M and 512M Prior work has demonstrated that the worker scaling performance of AdamW and Muon improves as data and model scale are increased (Thérien et al., 2025; Charles et al., 2025). In Figure 1, we compare the performance of SPlus, SOAP, AdamW, and Muon when scaling replicas. We also observe that scale improves replica scaling across all optimizers can significantly benefit from this including SOAP, SPlus, and Muon. Isolating the Effect of Local Steps: Replica-scaling effects are visible at H=1H=1, before multi-step parameter drift can occur! Unlike SGD, in adaptive optimizers, the divergence to standard DP training occurs even with a single step of local optimization (due to the non-linear map of gradients to update). Indeed, the arguments regarding covariance structure in Sec. 4 apply directly to this case. To isolate the effects of local drift when taking multiple steps from the impact of the adaptive optimizer update, we perform scaling experiments using a single local update. We also remove the impacts of the outer optimizer by using a standard SGD outer optimizer with a learning rate of 1 (equivalent to parameter averaging). Table 3 reports the results of this experiment. We observe similar results to the multi-step setting: Adam, Muon, and Shampoo have poorer replica scaling, while SOAP and SPlus remain substantially more robust. This is the setting in which Section 4.4 isolates the effect of applying the optimizer before gradient averaging. We also observe that performance (even in terms of scaling with respect to DP) is far more impacted by the choice of optimizer and the number of replicas than the number of local steps. This suggests the impact of the local optimizer on the update and subsequent aggregation can be more significant than the parameter drift often studied in local optimization (Stich, 2019; Karimireddy et al., 2020). Table 3: 178M performance using H=1 (single-step) local updates. We report the final validation loss for each method. To isolate the effect of the adaptive optimizer’s update, we use pure parameter averaging at each step and persistent local optimizer states. All methods train for 6,7756,775 steps. The DP, R=8R=8, and R=64R=64 methods are directly tuned, while the R=16,32R=16,32 configurations use interpolated hyperparameters rather than separate full sweeps. The best-performing method in each column is bolded. Inner optimizer DP R=8R=8 R=16R=16 R=32R=32 R=64R=64 SOAP 2.816 2.842 2.864 2.901 2.958 SPlus 2.853 2.837 2.853 2.873 2.912 Muon 2.858 2.979 3.073 3.230 3.417 AdamW 2.880 2.932 2.969 3.050 3.157 Shampoo 2.842 2.899 2.947 3.002 3.097 Interestingly, Muon is less robust here than at H=30H=30 and is worse than AdamW at the larger replica counts. Its strong H=30H=30 results therefore do not transfer unchanged to pure one-step averaging. We hypothesise that this may be due to differences in their outer optimizer, which could be more favourable for some methods and less for others. 5.2 Second-moment memory and replica scaling Following Sec. 4.4, we next study a practical consequence of smaller local batches for Adam-based methods. Equation (A.6) shows how local gradient variance enters the second-moment discrepancy. The token-half-life rule in Equation (20) increases β2 _2 as R grows, preserving the number of tokens represented by the running average. In this section, we study the impact of applying this and, in particular, evaluate how it impacts conclusions in prior work (Thérien et al., 2025; Horoi et al., 2026). This rule is already incorporated in the main results in Sec. 5.1. Here we evaluate both the Adam-based DiLoCo optimizer and IsoLoCo which is a modification of the outer optimizer. We incorporate it into our analysis as it is based on the Adam inner optimizer, which can benefit from appropriate β2 _2 tuning. We note that in the prior section’s experiments, to ensure the strongest baseline, this rule is already used by construction in all optimizers with a second moment. This section, on the other hand, emphasizes the improvement compared to existing literature, which does not scale β2 _2 when scaling R Charles et al. (2025); Horoi et al. (2026); Sarfi et al. (2025) To empirically measure the importance of second moment estimation highlighted in Sec. 4.4, we compared the performance of optimizers using a fixed β2 _2 and tuning it at every replica scale. Figure 5 reports the final validation loss of DiLoCo, IsoLoCo, and MuLoCo under these strategies. We note that while IsoLoCo and DiLoCo both apply AdamW to each tensor, MuLoCo only applies it to non-hidden tensors. We observe all optimizers improve from tuning β2 _2 as replicas increase, but IsoLoCo and DiLoCo benefit more. Under β2 _2 tuning, we find that IsoLoCo, despite using an AdamW inner optimizer, can even outperform MuLoCo at large worker counts. Figure 4: The effect of tuning AdamW’s β2 _2 in DiLoCo, IsoLoCo, and MuLoCo. We report the performance of each method at 178M scale when fixing (dashed lines) and tuning (solid lines) β2 _2 at each worker count. While IsoLoCo and DiLoCo both apply AdamW to each tensor, MuLoCo only applies it to non-hidden layers. We observe all optimizers improve from this change, but IsoLoCo and DiLoCo benefit significantly more, revealing that IsoLoCo can outperform MuLoCo at larger worker counts when β2 _2 is well tuned. Figure 5: Tuning β2 _2 for larger model and replica. We ablate using fixed β2 _2 vs scaling at high replica count (R=64) and 512M model scale. We observe that performance consistently improves for DiLoCo and IsoLoco. Notably IsoLoco outerperforms Muon inner optimizer (MuLoco). We also confirm that this scaling rule works well at 512M in Figure 5. Both methods use the published R=8R=8 hyperparameters without retuning. At 512M and R=64R=64, changing only β2 _2 reduces DiLoCo loss from 2.8842.884 to 2.8492.849 and IsoLoCo loss from 2.7342.734 to 2.7022.702, approximately 1.2%1.2\% in each case. Second-moment memory is thus an effective control on replica scaling, including for IsoLoCo, whose outer optimizer does not remove this inner AdamW sensitivity. Notably, after the relevant control, IsoLoCo, which uses an AdamW inner optimizer, can outperform MuLoCo, while only doing the spectral decomposition during the outer step. Discussion of Practical Implications Our results suggest that local optimization methods can be deployed across a large number of replicas (50+) while maintaining minimal performance degradation of <3%<3\% when utilizing SOAP or SPlus. Notably, these methods consume more memory than Muon due to the need to maintain L and R, however, scalable recipes for SOAP have been studied (Khona et al., 2026), demonstrating that it can perform well at very large scale. A particularly relevant setting for our work is cross-data center training where each replica has a considerable number of chips that can be used to shard the additional optimizer states and perform the optimizer step in parallel. At 178M and 512M scales, our results suggest that at a smaller number of replicas SOAP is the preferred inner optimizer, while at a large number of replicas, SPlus performs best. Another observation of our work is that the recently proposed IsoLoCo outer optimizer, which has nearly the same overhead as AdamW inner optimizer-based DiLoCo, can perform very effectively at a large number of replicas simply by adjusting its β2 _2 according to the local batch size. When doing so, we show that it outperforms MuLoCo. This can make it an excellent candidate for settings where each replica has limited memory or existing infrastructure and training recipe are already well designed for Adam. 5.3 Analysis: covariance structure and mean mismatch We empirically analyze the covariance structure and mean mismatch discussed in Section 4 for the different optimizer classes. Empirically Measuring Whitening Using the 178M training setting and beginning with checkpoints at outer round 75, we empirically estimate how close to whitened the covariance of the pseudo-gradient is for different inner optimizers. We perform 64 instantiations of inner training loop starting with the same checkpoint and inner optimizer state, thus only varying the data used. Subsequently we use pairs to construct an estimator for the following quantities: ~h:=1R2∑r=1Rr,h,ah=(dtr(~h2)tr(~h)2−1)1/2. C_h:= 1R^2 _r=1^R C_r,h, a_h= ( d\,tr( C_h^2)tr( C_h)^2-1 )^1/2. (21) Here d is the probed dimension; ah=0a_h=0 means isotropy, and larger values indicate greater anisotropy, independently of overall noise magnitude. The constructed ~h C_h equals the covariance of the averaged update when cross-worker noise covariances vanish, conditional on the starting parameters and optimizer states. Appendix C.1 gives the estimators and exact replay protocol. We chose this particular quantity as a proxy for isotropy as the trace and square trace have efficient estimators available in the statistical literature. The results are shown in Table 5 below. Table 4: Update-noise anisotropy at R=64R=64, round 75, using all 64 replicas and local batch four. Lower is more isotropic. Optimizer a1↓a_1\; a30↓a_30\; AdamW 35 11 Muon 3.8 3.3 SOAP 2.3 2.5 SPlus 0.43 2.4 Table 5: Mean mismatch is smaller than stochastic error in the detected high-noise directions, and particularly small for SPlus and SOAP compared with AdamW. Optimizer Mean mismatch (×10−6× 10^-6) Stochastic error (×10−6× 10^-6) Mean / noise (%\%) SPlus .044 .91 7.9 SOAP .12 1.4 7.5 Muon .11 1.1 14 AdamW 5.4 19 19 We observe that covariance structure broadly follows the scaling results. with optimizers that have better replica scaling having better covariance structure. The anisotropy differences are already pronounced after one local step, before parameter drift across replicas can affect subsequent gradients. a1a_1 ranges from 3535 for AdamW to 0.430.43 for SPlus, with Muon and SOAP between them. After 30 steps, SOAP and SPlus remain less anisotropic than Muon and AdamW, broadly following their replica-scaling performance. Interestingly, Muon is significantly more isotropic than Adam, helping to explain its improved scaling Thérien et al. (2025). The next measurements examine the curvature of high-variance noise directions, connecting this shape comparison to optimization. Noise dominates mean mismatch along high-noise directions. Finally, we compare the two terms of Equation (10) along directions selected for high update variance. For each optimizer, all replicas and its DP reference start from the same checkpoint weights and replica 0’s saved optimizer state. They run H=30H=30 steps on paired data: DP uses the union of the local batches at each step. We analyze the 63 hidden attention and MLP weight matrices separately, excluding embeddings, the output head, and other auxiliary parameters. For each matrix, 16 repetitions fit a scalar αR _R and a leading noise direction q; 24 additional repetitions estimate the directional mean mismatch (^⊤R)2( q b_R)^2 and stochastic error αR2^⊤agg _R^2 q C_ agg q. Both terms are divided by the squared norm of that matrix’s DP mean update. This per-matrix calibration applies the decomposition to each matrix separately; it is not a change to the training learning rates. Appendix C.2 defines the estimators, including the paired finite-sample correction. We observe the mean-mismatch term is only 7.57.5–19%19\% of the stochastic term on average along these directions. This supports controlling the stochastic term in Equation (10) which is the focus of our analysis in Section 4, and allow to better confirm that the covariance structure of SPlus and SOAP lead to their improved replica scaling. Extreme noise also lies in high-curvature directions. We have seen bounding the maximum noise direction of the covariance as in Equation 11 can reduce the maximum discrepancy to the DDP direction, in particular since the above empirical experiments suggest it dominates the bound on the error with respect to the DDP target. We next ask whether directions carrying extreme noise variance are directly responsible for poor performance of methods such as AdamW which do not bound them well. We do this by estimating these direction and then the underlying curvature. At each method’s own checkpoint for outer step 75, we fit a unit direction q that maximizes sample variance over 64 paired one-step update differences. With H the loss Hessian on 64 held-out sequences, Table 6 reports κ(^)=^⊤^116∑j=116j⊤j,κ( q)= q H q 116 _j=1^16 v_j H v_j, where the j v_j are independent random unit vectors in the same six-matrix subspace. A ratio of 100100 means 100100 times the average random-direction curvature, not noise variance. Table 6: Loss curvature along a fitted high-noise direction, relative to random directions. Ranges span replicas 0 and 1, not confidence intervals; local batch b=256/Rb=256/R. Curvature ratio κ(^)κ( q) Optimizer R=8R=8 R=64R=64 AdamW 8686–170170 460460–730730 Muon 6.46.4–7.77.7 4747–100100 SOAP 3.13.1–3.23.2 4.24.2–5.25.2 SPlus 1.5 1.5–1.6 1.6 1.5 1.5 Appendix C.3 details the estimator used. We observe that Adam and Muon have very high curvature in the directions with the most noise compared to average directions, in particular as we scale the replica size, while SOAP and SPlus are able to effectively redistribute noise in any particular direction due to their whitening. High noise along high curvature directions can be particularly problematic as small errors can be very damaging for the loss in high curvature regions. This gives an additional mechanism for understanding why whitening (which redistributes error more uniformly along all directions) can be particularly important in this setting. 6 Conclusion We have shown that differences to the DP pseudogradient can be seen to result from two primary sources of error: mean mismatch and stochastic error. Building on these insights, we showed that covariance structure in optimizers like AdamW is very important. In our empirical study, we establish the performance of different inner optimizers (SOAP, SPlus, Muon, AdamW), which have more and less appealing properties for reducing stochastic error according to our theory. We find that SPlus reaches the strongest performance at higher worker counts, and the best overall scaling with respect to DP, while SOAP reaches the best DP performance and seems on trend to become the best inner optimizer at scale. We then directly test our theoretical claims by calculating the mean mismatch and stochastic error for each optimizer, showing that these quantities are directly correlated with the observed performance trends. Acknowledgements We acknowledge support from NSERC Discovery RGPIN-2021-04104 and compute resources from Digital Research Alliance Canada. References Aitken (1935) A. C. Aitken. On least squares and linear combinations of observations. Proceedings of the Royal Society of Edinburgh, 55:42–48, 1935. doi:10.1017/S0370164600014346. Bernstein et al. (2018) Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Anima Anandkumar. signSGD: Compressed optimisation for non-convex problems. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, p. 560–569, 2018. URL https://proceedings.mlr.press/v80/bernstein18a.html. Charles et al. (2025) Zachary Charles, Gabriel Teston, Lucio Dery, Keith Rush, Nova Fallen, Zachary Garrett, Arthur Szlam, and Arthur Douillard. Communication-efficient language model training scales reliably and robustly: Scaling laws for DiLoCo. In Advances in Neural Information Processing Systems, 2025. URL https://arxiv.org/abs/2503.09799. Dirksen et al. (2022) Sjoerd Dirksen, Johannes Maly, and Holger Rauhut. Covariance estimation under one-bit quantization. The Annals of Statistics, 50(6):3538–3562, 2022. doi:10.1214/22-AOS2239. URL https://arxiv.org/abs/2104.01280. Douillard et al. (2023) Arthur Douillard, Qixuan Feng, Andrei A. Rusu, Rachita Chhaparia, Yani Donchev, Adhiguna Kuncoro, Marc’Aurelio Ranzato, Arthur Szlam, and Jiajun Shen. DiLoCo: Distributed low-communication training of language models. arXiv:2311.08105, 2023. URL https://arxiv.org/abs/2311.08105. Douillard et al. (2025) Arthur Douillard, Yanislav Donchev, Keith Rush, Satyen Kale, Zachary Charles, Zachary Garrett, Gabriel Teston, Dave Lacey, Ross McIlroy, Jiajun Shen, Alexandre Ramé, Arthur Szlam, Marc’Aurelio Ranzato, and Paul Barham. Streaming DiLoCo with overlapping communication: Towards a distributed free lunch. arXiv:2501.18512, 2025. URL https://arxiv.org/abs/2501.18512. Frans et al. (2025) Kevin Frans, Sergey Levine, and Pieter Abbeel. A stable whitening optimizer for efficient neural network training, 2025. URL https://arxiv.org/abs/2506.07254. Gupta et al. (2018) Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, p. 1842–1850, 2018. URL https://proceedings.mlr.press/v80/gupta18a.html. Horoi et al. (2026) Stefan Horoi, Benjamin Thérien, Guy Wolf, and Eugene Belilovsky. Can model merging improve aggregation in DiLoCo? arXiv:2607.03011, 2026. URL https://arxiv.org/abs/2607.03011. Jordan et al. (2024) Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks. Technical article, 2024. URL https://kellerjordan.github.io/posts/muon/. Karimireddy et al. (2020) Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, p. 5132–5143, 2020. URL https://proceedings.mlr.press/v119/karimireddy20a.html. Kessy et al. (2018) Agnan Kessy, Alex Lewin, and Korbinian Strimmer. Optimal whitening and decorrelation. The American Statistician, 72(4):309–314, 2018. doi:10.1080/00031305.2016.1277159. URL https://arxiv.org/abs/1512.00809. Khona et al. (2026) Mikail Khona, Aditya Vavre, Boxiang Wang, Deyu Fu, Hao Wu, Mike Chrzanowski, Bryan Catanzaro, Dheevatsa Mudigere, Jeff Pool, Michael Lightstone, Mohammad Shoeybi, Mostofa Patwary, Nima Tajbakhsh, and Tijmen Blankevoort. SOAP, muon, and beyond: Pushing LLM pretraining scales. arXiv:2607.20548, 2026. URL https://arxiv.org/abs/2607.20548. Kingma & Ba (2015) Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. URL https://arxiv.org/abs/1412.6980. Li et al. (2024) Jeffrey Li et al. DataComp-LM: In search of the next generation of training sets for language models. arXiv:2406.11794, 2024. URL https://arxiv.org/abs/2406.11794. Lin et al. (2026) Wu Lin, Scott C. Lowe, Felix Dangel, Runa Eschenhagen, Zikun Xu, and Roger B. Grosse. Understanding and improving shampoo and SOAP via Kullback–Leibler minimization. In International Conference on Learning Representations, 2026. URL https://arxiv.org/abs/2509.03378. Liu et al. (2024) Bo Liu et al. Asynchronous local-SGD training for language modeling. arXiv:2401.09135, 2024. URL https://arxiv.org/abs/2401.09135. Loshchilov & Hutter (2019) Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. URL https://arxiv.org/abs/1711.05101. Malladi et al. (2022) Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the SDEs and scaling rules for adaptive gradient algorithms. In Advances in Neural Information Processing Systems, 2022. URL https://arxiv.org/abs/2205.10287. Marczak et al. (2025) Daniel Marczak, Simone Magistri, Sebastian Cygert, Bartłomiej Twardowski, Andrew D. Bagdanov, and Joost van de Weijer. No task left behind: Isotropic model merging with common and task-specific subspaces. In Proceedings of the 42nd International Conference on Machine Learning, 2025. URL https://arxiv.org/abs/2502.04959. Marek et al. (2025) Martin Marek, Sanae Lotfi, Aditya Somasundaram, Andrew Gordon Wilson, and Micah Goldblum. Small batch size training for language models: When vanilla SGD works, and why gradient accumulation is wasteful, 2025. URL https://arxiv.org/abs/2507.07101. McMahan et al. (2017) H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Ag”uera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, volume 54 of Proceedings of Machine Learning Research, p. 1273–1282, 2017. URL https://proceedings.mlr.press/v54/mcmahan17a.html. Reddi et al. (2021) Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H. Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2021. URL https://arxiv.org/abs/2003.00295. Reed & Baioumy (2025) Matthew Reed and Mohamed Baioumy. SPARTA: Distributed training with sparse parameter averaging. EXO Labs, 2025. URL https://openreview.net/forum?id=stFPf3gzq1. See also the OpenReview preprint “Improving the Efficiency of Distributed Training using Sparse Parameter Averaging”. Sarfi et al. (2025) Amir Sarfi et al. Communication efficient LLM pre-training with SparseLoCo. arXiv:2508.15706, 2025. URL https://arxiv.org/abs/2508.15706. Stich (2019) Sebastian U. Stich. Local SGD converges fast and communicates little. In International Conference on Learning Representations, 2019. URL https://arxiv.org/abs/1805.09767. Thérien et al. (2025) Benjamin Thérien, Xiaolong Huang, Aaron Defazio, Irina Rish, and Eugene Belilovsky. MuLoCo: Muon is a practical inner optimizer for DiLoCo, 2025. URL https://arxiv.org/abs/2505.23725. Voigtlaender (2021) Felix Voigtlaender. A general version of Price’s theorem: A tool for bounding the expectation of nonlinear functions of Gaussian random vectors. Journal of Theoretical Probability, 34:1474–1485, 2021. doi:10.1007/s10959-020-01017-w. Vyas et al. (2025) Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. SOAP: Improving and stabilizing shampoo using adam. In International Conference on Learning Representations, 2025. URL https://arxiv.org/abs/2409.11321. Wang et al. (2024) Jianyu Wang, Rudrajit Das, Gauri Joshi, Satyen Kale, Zheng Xu, and Tong Zhang. On the unreasonable effectiveness of federated averaging with heterogeneous data. Transactions on Machine Learning Research, 2024. URL https://arxiv.org/abs/2206.04723. Wen et al. (2025) Kaiyue Wen, David Hall, Tengyu Ma, and Percy Liang. Fantastic pretraining optimizers and where to find them. arXiv:2509.02046, 2025. URL https://arxiv.org/abs/2509.02046. Yu et al. (2019) Hao Yu, Rong Jin, and Sen Yang. On the linear speedup analysis of communication efficient momentum SGD for distributed non-convex optimization. In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, p. 7184–7193, 2019. URL https://proceedings.mlr.press/v97/yu19d.html. Appendix A collects the supporting theoretical results and proofs. Appendix B records the training setup, and Appendix C gives the analysis protocols and estimators. Appendix A Supporting theoretical results A.1 Muon covariance under rotational proposals Let p=min(m,n)p= (m,n) and dmax=max(m,n)d_ = (m,n). If a full-rank proposal has SVD =~diag(σi)~⊤ M= Udiag( _i) V , the ideal Muon update with prescribed matrix-size scale γ>0γ>0 is Muon=γpolar()=γ~~⊤. U_ Muon=γ\,polar( M)=γ U V . (A.1) The polar map discards the realized singular values and keeps the polar factor. In particular, polar(a)=polar()polar(a M)=polar( M) for every a>0a>0, so a positive multiplicative change to the entire proposal leaves the update unchanged. Proposition 4 (Muon covariance under rotational proposals). Suppose r M_r is full rank almost surely. If m≥nm≥ n, suppose r=r O M_r d= M_r for every ∈O(m) O (m); if m<nm<n, suppose r=r M_r O d= M_r for every ∈O(n) O (n). Either one-sided condition suffices when m=nm=n. Then [Muon,r]=,Cov(vec(Muon,r))=γ2dmaxmn.E[ U_ Muon,r]= 0, (vec( U_ Muon,r))= γ^2d_ I_mn. (A.2) Consequently, independent replicas have a uniformly averaged update covariance γ2mn/(dmaxR)γ^2 I_mn/(d_ R), even when their proposal singular-value distributions differ. Proof. Consider m≥nm≥ n and write the columns of r=polar(r) Q_r=polar( M_r) as r,1,…,r,n q_r,1,…, q_r,n. Equivariance of the polar map and left invariance of the proposal law imply r=r O Q_r d= Q_r for every ∈O(m) O (m). Taking =− O=- I gives [r]=E[ Q_r]= 0. Moreover, [r,jr,k⊤]E[ q_r,j q_r,k ] commutes with every orthogonal matrix and must therefore equal ajkma_jk I_m. Since r⊤r=n Q_r Q_r= I_n, majk=[r,k⊤r,j]=j=k.ma_jk=E[ q_r,k q_r,j]=1_\j=k\. (A.3) The column blocks of [vec(r)vec(r)⊤]E[vec( Q_r)vec( Q_r) ] are consequently j=km/m1_\j=k\ I_m/m, which proves Equation (A.2) because dmax=md_ =m. The proof for m<nm<n applies the same argument to the rows, using right invariance and rr⊤=m Q_r Q_r = I_m. ∎ A.2 Proof of Proposition 2 (SPlus to DDP comparison) Proof of Proposition 2. Write ψ(a)=2Φ(a)−1ψ(a)=2 (a)-1 and c=2/πc= 2/π. For one coordinate, set Xr=σ(a+RZr)X_r=σ(a+ RZ_r) and X¯=R−1∑rXr X=R^-1 _rX_r. Then Pr(Xr>0)=Φ(a/R),[sign(Xr)]=ψ(a/R),[sign(X¯)]=ψ(a). (X_r>0)= (a/ R), [ (X_r)]=ψ(a/ R), [ ( X)]=ψ(a). In the shared-state model the last expression is exactly the mean of the signed DDP momentum coordinate, because momentum is affine before the sign. The expectation of the local average equals the single-replica expectation. For a≥0a≥ 0, Rψ(a/R)−ψ(a)=c∫0a(e−u2/(2R)−e−u2/2)du. Rψ(a/ R)-ψ(a)=c _0^a (e^-u^2/(2R)-e^-u^2/2 )du. (A.4) The integrand is nonnegative and at most u2(1−R−1)/2u^2(1-R^-1)/2, by the unit Lipschitz bound for e−xe^-x on x≥0x≥ 0. Integration and odd symmetry give |Rψ(a/R)−ψ(a)|≤c6(1−R−1)|a|3. | Rψ(a/ R)-ψ(a) |≤ c6(1-R^-1)|a|^3. Squaring and summing over coordinates is valid because Q is orthogonal. Each sign has variance 1−ψ(aj/R)21-ψ(a_j/ R)^2. Independence of projected coordinates and of replicas gives the covariance in (17). Finally use Proposition 1. ∎ The DDP covariance is Diag(1−ψ(aj)2)⊤ QDiag(1-ψ(a_j)^2) Q . For R≥1R≥ 1 the calibrated local covariance minus this matrix is positive semidefinite, with operator norm at most c2maxjaj2c^2 _ja_j^2. Thus mean agreement to first order should not be read as a variance improvement over signed DDP. At zero signal the covariances agree exactly. A.3 Adam/SOAP second moments and the mean comparison Condition on common incoming optimizer states and parameters, as in Section 4.4. Let t denote the fresh-gradient step. Averaging squared gradients is not the same as squaring their average: the difference is the scatter of the replica gradients. Define R=R−1∑r(^r−¯)(^r−¯)⊤ S_R=R^-1 _r( g_r- g)( g_r- g) . For the second moments in (2) and (5), 1R∑rA,r,t−A,DDP,t=(1−β2)diag(R),1R∑rS,r,t−S,DDP,t=(1−β2)diag(⊤R). split 1R _r v_A,r,t- v_A, DDP,t&=(1- _2)diag( S_R),\\ 1R _r v_S,r,t- v_S, DDP,t&=(1- _2)diag( Q S_R Q). split (A.5) The second identity assumes a common basis fixed before the fresh gradients; A,S v_A, v_S are the vectorized buffers. Thus locally updated second moments already differ from DDP before any parameter drift. Under the main-text batch model, Cov(^r)=RgCov( g_r)=R _g and the replica gradients are independent, so g _g is the covariance of the DDP-averaged gradient. At fixed B, taking expectations in (A.5) gives [1R∑rA,r,t−A,DDP,t]=(1−β2)(R−1)diag(g).E [ 1R _r v_A,r,t- v_A, DDP,t ]=(1- _2)(R-1)diag( _g). (A.6) SOAP has the same expected identity in its fixed basis. Expanding squares around the sample mean gives 1R∑r^r⊙2−¯⊙2=1R∑r(^r−¯)⊙2=diag(R). 1R _r g_r 2- g 2= 1R _r( g_r- g) 2=diag( S_R). Multiplying by 1−β21- _2 proves the Adam identity in (A.5); applying it to ⊤^r Q g_r proves the SOAP identity when Q is fixed and common. Under (14), [R]=(R−1)gE[ S_R]=(R-1) _g, which proves (A.6). For a scalar smooth map f and Xt=g+tσZX_t=g+ t\,σ Z, Gaussian variance differentiation (Voigtlaender, 2021) gives the exact identity dt[f(Xt)]=σ22[f′(Xt)],[f(XR)]−[f(X1)]=σ22∫1R[f′(Xt)]t. ddtE[f(X_t)]= σ^22E[f (X_t)], [f(X_R)]-E[f(X_1)]= σ^22 _1^RE[f (X_t)]\,dt. (A.7) Differentiation under the integral is justified, for example, by bounded first and second derivatives and Gaussian integrability; the Adam map below satisfies the required conditions when the incoming second moment is positive. A.3.1 Finite-noise mean bound and second-moment memory Proof of Proposition 3. For this proof only, abbreviate p=β1mp= _1m, k=1−β1k=1- _1, a0=β2v>0a_0= _2v>0, and b0=1−β2>0b_0=1- _2>0. Write f(x)=ph(x)+kℓ(x),h(x)=(a0+b0x2)−1/2,ℓ(x)=x(a0+b0x2)−1/2.f(x)=p\,h(x)+k\, (x), h(x)=(a_0+b_0x^2)^-1/2, (x)=x(a_0+b_0x^2)^-1/2. Direct differentiation gives h′(x)=b0(2b0x2−a0)(a0+b0x2)5/2,ℓ′(x)=−3a0b0(a0−4b0x2)(a0+b0x2)7/2.h (x)= b_0(2b_0x^2-a_0)(a_0+b_0x^2)^5/2, (x)= -3a_0b_0(a_0-4b_0x^2)(a_0+b_0x^2)^7/2. (A.8) For y≥0y≥ 0, the suprema of |2y−1|/(1+y)5/2|2y-1|/(1+y)^5/2 and |1−4y|/(1+y)7/2|1-4y|/(1+y)^7/2 are both one. The first follows by differentiation on either side of y=1/2y=1/2, and the second on either side of y=1/4y=1/4; their interior positive extrema have value less than one. Substituting y=b0x2/a0y=b_0x^2/a_0 gives ‖h′‖∞≤b0/a03/2,‖ℓ′‖∞≤3b0/a03/2.\|h \|_∞≤ b_0/a_0^3/2, \| \|_∞≤ 3b_0/a_0^3/2. Since ℓ′ is odd and Z is symmetric, [ℓ′(tσZ)]=0E[ ( tσ Z)]=0. The mean-value bound therefore implies |[ℓ′(g+tσZ)]|≤3b0|g|/a03/2. |E[ (g+ tσ Z)] |≤ 3b_0|g|/a_0^3/2. Thus, uniformly in t>0t>0, |[f′(g+tσZ)]|≤b0a03/2(|p|+3k|g|).|E[f (g+ tσ Z)]|≤ b_0a_0^3/2 (|p|+3k|g| ). Integrating (A.7) over 1≤t≤R1≤ t≤ R gives (19). This proof does not replace the random denominator by its expectation and does not require a small-noise limit. ∎ Scaling rule A local example or token is downweighted by β2 _2 per local step. Keeping its decay per fixed number of local examples unchanged gives logβ2(R)/(B/R)=logβ2(R0)/(B/R0) _2(R)/(B/R)= _2(R_0)/(B/R_0), proving (20). Write c0=−R0logβ2(R0)>0c_0=-R_0 _2(R_0)>0. Then 1−β2(R)=1−e−c0/R≤c0/R1- _2(R)=1-e^-c_0/R≤ c_0/R, so (R−1)(1−β2(R))≤c0(R-1)(1- _2(R))≤ c_0. This controls the upper bound with fixed incoming state; it neither forces the bias to zero nor proves that its actual magnitude decreases monotonically with β2 _2. A.4 Basis and scale errors SOAP whitening requires both a basis that reduces correlations and coordinate scales that normalize the remaining variances. We first separate these two sources of error in the covariance of the returned update, then show how they affect covariance matching and the stochastic term in Equation (10). Conditioning and scope. The covariance calculations below concern centered proposals with bases and denominators fixed before the proposal is drawn. They describe an idealized one-step optimizer core, omitting a common learning-rate factor and the numerical stabilizer. The ideal denominators estimate the centered variance of this proposal. In practical SOAP, the denominator instead uses gradient second moments, whereas the numerator contains momentum; refreshing either state from the same gradient also couples the map to its input. The identities below therefore describe the frozen-state model, not the full fresh-gradient Adam or SOAP step. A.4.1 Scale errors and correlations left in the estimated basis Let worker r have a centered matrix proposal r∈ℝm×n M_r ^m× n, and set r=vec(r) x_r=vec( M_r) and r=Cov(r)≻ _r=Cov( x_r) 0. Let ^L,r Q_L,r and ^R,r Q_R,r be the fixed estimated left and right bases. They are orthogonal matrices. Rotating into these bases gives the matrix proposal on which normalization or sign acts: r=^L,r⊤r^R,r,r=^R,r⊗^L,r,vec(r)=r⊤r. Z_r= Q_L,r M_r Q_R,r, Q_r= Q_R,r Q_L,r, ( Z_r)= Q_r x_r. (A.9) No Kronecker assumption is imposed on r _r here: the Kronecker matrix r Q_r represents the two-sided rotation, not a factorization of the full covariance. To distinguish scale error from an imperfect basis, let r _r be the true covariance in the estimated basis, r D_r its diagonal of coordinate variances, and r C_r the correlation matrix remaining after exact variance normalization: r=r⊤r,r=Diag(r),vrj=(r)jj,r=r−1/2r−1/2. _r= Q_r _r Q_r, D_r=Diag( v_r), v_rj=( _r)_j, C_r= D_r^-1/2 _r D_r^-1/2. (A.10) Thus r C_r has unit diagonal; its off-diagonal entries measure correlations that the basis has not removed. The actual normalization uses frozen variance estimates v^rj v_rj. Write their relative errors as erje_rj and collect the resulting inverse-scale gains in r A_r: v^rj=vrj(1+erj)>0,^r=Diag(^r),r=Diag((1+erj)−1/2). v_rj=v_rj(1+e_rj)>0, D_r=Diag( v_r), A_r=Diag ((1+e_rj)^-1/2 ). (A.11) If ^r V_r is the m×nm× n reshape of ^r v_r, the SOAP core divides by these estimated scales, whereas SPlus takes signs. Both then rotate the result back to the parameter coordinates: rSOAP U_r SOAP =^L,r(r⊘^r)^R,r⊤, = Q_L,r ( Z_r V_r ) Q_R,r , vec(rSOAP) ( U_r SOAP) =r^r−1/2r⊤r, = Q_r D_r^-1/2 Q_r x_r, (A.12) rSPlus U_r SPlus =^L,rsign(r)^R,r⊤, = Q_L,r ( Z_r) Q_R,r , vec(rSPlus) ( U_r SPlus) =rsign(r⊤r). = Q_r ( Q_r x_r). (A.13) Division, square roots, and signs in these matrix expressions act entrywise. Proposition 5 (Separate effects of the basis and coordinate scales). Under the preceding frozen-state assumptions, define rSOAP=Cov(vec(rSOAP)) C_r SOAP=Cov(vec( U_r SOAP)). For every centered proposal with finite covariance, rSOAP=rrrrr⊤, C_r SOAP= Q_r A_r C_r A_r Q_r , (A.14) and the whitening error decomposes exactly as ‖rSOAP−‖F2=∑j(erj1+erj)2⏟coordinate-scale error+∑i≠jr,ij2(1+eri)(1+erj)⏟residual correlations after normalization.\| C_r SOAP- I\|_F^2= _j ( e_rj1+e_rj )^2_coordinate-scale error+ _i≠ j C_r,ij^2(1+e_ri)(1+e_rj)_residual correlations after normalization. (A.15) If additionally r x_r is centered Gaussian, then rSPlus C_r SPlus =r[2πarcsin(r)]r⊤, = Q_r [ 2π ( C_r) ] Q_r , (A.16) ‖rSPlus−‖F2 \| C_r SPlus- I\|_F^2 =∑i≠j[2πarcsin(r,ij)]2≤‖r−‖F2. = _i≠ j [ 2π (C_r,ij) ]^2≤\| C_r- I\|_F^2. (A.17) Here arcsin acts entrywise. The same sign identities hold for the positive scale-mixture model r=τrr x_r= _r y_r, where r∼(,r) y_r ( 0, _r), τr>0 _r>0 is independent of r y_r, and [τr]=1E[ _r]=1. Proof. Vectorization gives Equations (A.12) and (A.13). Since the maps and denominators are frozen, the SOAP covariance is the linear covariance sandwich. Substitute r=r1/2rr1/2 _r= D_r^1/2 C_r D_r^1/2 and ^r−1/2r1/2=r D_r^-1/2 D_r^1/2= A_r to obtain Equation (A.14). Orthogonal conjugation preserves Frobenius norm. In the rotated frame, the diagonal covariance entries are (1+erj)−1(1+e_rj)^-1 and the off-diagonal entries are r,ij/(1+eri)(1+erj)C_r,ij/ (1+e_ri)(1+e_rj). Summing their squared differences from the identity proves Equation (A.15). For the sign map, each centered nondegenerate Gaussian coordinate has an equiprobable sign and hence mean zero and variance one. Two standardized Gaussian coordinates with correlation ρ have sign disagreement probability arccos(ρ)/π (ρ)/π: represent them as projections of an isotropic two-dimensional Gaussian onto two unit directions with angle arccos(ρ) (ρ) and use its uniform angular distribution. Their sign product therefore has expectation 1−2arccos(ρ)/π=(2/π)arcsin(ρ)1-2 (ρ)/π=(2/π) (ρ). This is the Gaussian arcsine identity used in one-bit covariance estimation (Dirksen et al., 2022), and proves Equation (A.16). For 0≤t≤10≤ t≤ 1, convexity of arcsin(t) (t) places it below the chord joining its endpoint values, so (2/π)arcsin(t)≤t(2/π) (t)≤ t. Oddness extends the absolute-value inequality to [−1,1][-1,1]. Squaring the off-diagonal entries proves Equation (A.17). Finally, multiplication by the positive scalar τr _r leaves every sign unchanged; independence and [τr]=1E[ _r]=1 preserve the stated input covariance. ∎ An accurate basis does not remove SOAP’s scale-estimation error. The first term remains when the basis decorrelates the proposal but the local scales are inaccurate; the second remains when exact scales are used in a basis that leaves correlations. SPlus removes the estimated-scale term in this centered model, while retaining dependence on those correlations. In particular, if r= C_r= I, SPlus has identity covariance for every worker, including workers using different bases and different input variances. For an explicit bound, if maxj|erj|≤ϵr<1 _j|e_rj|≤ _r<1, then ‖rSOAP−‖F2≤‖r‖22+‖r−‖F2(1−ϵr)2.\| C_r SOAP- I\|_F^2≤ \| e_r\|_2^2+\| C_r- I\|_F^2(1- _r)^2. (A.18) This follows by bounding each denominator in Equation (A.15) from below. Underestimating a variance can therefore amplify the whitening error through its reciprocal. Rotation back, by contrast, preserves the error in either operator or Frobenius norm: ‖rr⊤−‖=‖−‖.\| Q_r B Q_r - I\|=\| B- I\|. (A.19) Orthogonal multiplication preserves singular values, proving Equation (A.19). Rotation back can change the directions of residual anisotropy, but cannot increase this norm of its error. Appendix B Training setup and hyperparameters Models and data. The experiments use decoder-only LLaMA models trained on DCLM (Li et al., 2024), with sequence length 2048 and evaluation on a held-out 100M-token validation split. The model configuration files give the architecture below. Global batch size is held fixed within each model scale, so the local batch is the global batch divided by R. Setting 178M 512M Transformer layers 9 12 Hidden width 1024 1536 Attention heads / key-value heads 8/88/8 12/1212/12 Feed-forward width 2688 5440 Activation SwiGLU SwiGLU Sequence length 2048 2048 Global tokens per inner step 524,288 2,097,152 Global sequences per inner step 256 1024 Training tokens (approximately) 3.55B 10.25B Table 7: Architecture and training scales. The global batch is measured per inner step, not per communication round. Aggregation and comparisons. Unless stated otherwise, replicas take H=30H=30 local steps and use a Nesterov SGD outer optimizer. The H=1H=1 controls use uniform averaging without outer momentum. Percentage changes in validation loss use each optimizer’s own DP reference and are computed before rounding. The covariance, mean/noise, and curvature probes use the 178M outer-step-75 checkpoints; their replay settings and sample counts are specified in Appendix C. Hyperparameter sweep. We tuned the learning rate for each optimizer in the DP setting and the inner and outer learning rates and outer momentum for local training. For each optimizer we tuned R=8,64,96 each with at least 10 grid points extrapolating the HP for R=16 and 32. Appendix C Analysis protocols and estimators All three analyses use the 178M outer-step-75 checkpoints. The covariance and curvature probes use six matrices (attention query and MLP down-projection at layers 0, 4, and 8); the DP-referenced mean/noise comparison covers all 63 hidden attention and MLP matrices. Each optimizer uses its own trained checkpoint. The protocols below distinguish native saved states from the common-state replay used for comparison with DP. C.1 Conditional covariance replay This appendix spells out the covariance measurement used in Table 5. The probe estimates the minibatch-induced distribution of an optimizer’s local update conditional on one trained checkpoint. It does not treat successive updates from the training trajectory as independent covariance samples. C.1.1 Conditional replay protocol At a communication boundary, all replicas share the synchronized parameters W, while replica r has its own saved optimizer state srs_r. Table 5 uses all 64 saved replica states. For each replica, we run 2n2n rollouts: n=16n=16. Before every rollout, we restore the same pair (,sr)( W,s_r), then consume a new segment of the replica’s shuffled data stream. Within a rollout, the parameters and optimizer state evolve normally. We record the parameter displacement after h∈1,30h∈\1,30\ local steps, r,j(h)=concatℓ∈ℒvec(ℓ−r,j,ℓ(h)),j=1,…,2n, u_r,j^(h)= concat_ vec\! ( W_ - W_r,j, ^(h) ), j=1,…,2n, (C.1) where ℒL contains attention q_proj and MLP down_proj at layers 0, 4, and 8. For the 178M model, this concatenates d=11,403,264d=11,403,264 coordinates. The h=1h=1 and h=30h=30 endpoints come from the same rollout and therefore have coupled finite-sample errors. The per-replica batch is four sequences for R=64R=64, giving a global batch of 256 sequences. Each replay uses the optimizer configuration saved by the corresponding training run, including its replica-local momentum and second moments and its saved replica-specific or shared basis. For fixed r and h, write r,h=[r,j(h)],r,h=Cov(r,j(h)). μ_r,h=E[ u_r,j^(h)], C_r,h=Cov( u_r,j^(h)). (C.2) Two consecutive rollouts form one paired difference, r,i(h)=r,2i−1(h)−r,2i(h)2,i=1,…,n. z_r,i^(h)= u_r,2i-1^(h)- u_r,2i^(h) 2, i=1,…,n. (C.3) If the two rollouts are independent draws from the same conditional distribution, then [r,i(h)]=,Cov(r,i(h))=12(r,h+r,h)=r,h.E[ z_r,i^(h)]= 0, ( z_r,i^(h))= 12( C_r,h+ C_r,h)= C_r,h. (C.4) The subtraction therefore removes the unknown mean, while 1/21/ 2 restores the original covariance scale. C.1.2 Trace estimators from full-coordinate inner products Fix a horizon and suppress h in the notation. Define the Gram entries Gri,sj=r,i⊤s,j.G_ri,sj= z_r,i z_s,j. (C.5) The code computes these inner products over every selected coordinate; it does not use a random projection. It nevertheless estimates population moments from n paired differences per replica. For each replica, define tr t_r =1n∑i=1nGri,ri, = 1n _i=1^nG_ri,ri, (C.6) br b_r =1n(n−1)∑i≠jGri,riGrj,rj, = 1n(n-1) _i≠ jG_ri,riG_rj,rj, (C.7) qrr q_r =1n(n−1)∑i≠jGri,rj2. = 1n(n-1) _i≠ jG_ri,rj^2. (C.8) Under the IID-rollout idealization, for independent paired differences, [tr]=tr(r),[br]=tr(r)2,[qrr]=tr(r2).E[t_r]=tr( C_r), [b_r]=tr( C_r)^2, [q_r]=tr( C_r^2). (C.9) The last identity follows from [(⊤′)2]=tr([⊤][′′⊤])=tr(2)E[( z z )^2]=tr\! (E[ z z ]E[ z z ] )=tr( C^2) (C.10) for independent z and ′ z . Using i≠ji≠ j is important: including the diagonal would add fourth-moment self-products and bias the estimate upward. For different replicas r≠sr≠ s, the cross term qrs=1n2∑i,jGri,sj2q_rs= 1n^2 _i,jG_ri,sj^2 (C.11) estimates tr(rs)tr( C_r C_s) under the same IID and conditional replica-independence idealization. The quantity is the scale-dependent Frobenius overlap between two replicas’ covariance matrices, not a cross-covariance between their stochastic updates. C.1.3 Scaled sum of local covariances The probe constructs the scaled sum ~=1R2∑r=1Rr. C= 1R^2 _r=1^R C_r. (C.12) For the actual replica average ¯=R−1∑r u=R^-1 _r u_r, this equals Cov(¯)Cov( u) if the cross-replica noise covariances vanish conditional on the starting parameters and optimizer states. The local-replay probe does not measure those cross-covariances. To estimate the trace moments of ~ C, it uses T T =1R2∑rtr, = 1R^2 _rt_r, [T^] [ T] =tr(~), =tr( C), (C.13) S S =1R4∑r,sqrs, = 1R^4 _r,sq_rs, [S^] [ S] =tr(~2), =tr( C^2), (C.14) B B =1R4(∑rbr+∑r≠strts), = 1R^4 ( _rb_r+ _r≠ st_rt_s ), [B^] [ B] =tr(~)2. =tr( C)^2. (C.15) The separate estimator B B avoids squaring the noisy estimate T T. The reported anisotropy is a^=max0,dS^B^−1. a= \! \0, d S B-1 \. (C.16) Under the IID idealization, T T, S S, and B B are built from unbiased moment estimators, but their ratio and square root are not themselves unbiased. The clipping only handles a negative raw value caused by finite-sample error; it can therefore return zero even when the population covariance is anisotropic. If λ1,…,λd _1,…, _d are the eigenvalues of ~ C, the population quantity is a2=d∑iλi2(∑iλi)2−1,reffd=11+a2,a^2= d _i _i^2( _i _i)^2-1, r_effd= 11+a^2, (C.17) where reff=(∑iλi)2/∑iλi2r_eff=( _i _i)^2/ _i _i^2 is the participation-ratio effective rank. Thus a measures covariance shape independently of scale; it is undefined when the covariance is identically zero. For example, diag(5,5)diag(5,5) has a=0a=0, whereas diag(9,1)diag(9,1) has a=0.8a=0.8. Multiplying either matrix by a positive scalar leaves a unchanged. C.1.4 Scope and interpretation The estimator is conditional on one checkpoint and one set of replica-local states. Each method is evaluated on its own trained trajectory, so a cross-method difference can reflect both the update rule and the state reached by training; it is not a same-state causal ablation. Concatenating six matrices also mixes within-matrix anisotropy, scale imbalance between matrices, and cross-matrix covariance. The input streams are disjoint segments of one shuffled replica-local permutation, rather than independent samples drawn with replacement. The unbiasedness equalities above are therefore exact only under the stated IID idealization. Moreover, the n(n−1)n(n-1) and n2n^2 sums reuse the same n vectors; their terms are not independent observations. Finally, interpreting Equation (C.12) as the covariance of the actual average requires vanishing conditional cross-replica noise covariances. Lower a indicates a more isotropic ~ C, independently of total noise magnitude. For SOAP, the saved position relative to a basis-refresh step can also affect the measured covariance: the fresh gradient may enter both the basis and the proposal. At this single checkpoint, the probe characterizes update geometry but does not identify its causal contribution to final loss. C.2 DP-referenced directional mean and noise Each optimizer is evaluated at its own round-75 checkpoint. At the start of every repetition, the checkpoint weights and replica 0’s saved optimizer state are copied to every replica and to the DP reference. “Common state” means identical states within that replay, not an averaged optimizer state or a state shared across optimizers. Each replica consumes a persistent, shuffled data stream without reusing sequences across repetitions. In each step DP sees the union of the same local batches, evaluated at its own iterate. This preserves the pairing between local and DP data without reusing examples between the direction-fitting and evaluation samples. Let j index one of the 63 hidden matrices, vectorized in parameter coordinates. Write j,k x_j,k for its averaged local displacement and j,k y_j,k for its DP displacement in repetition k. On the 16 fitting repetitions at R=8R=8, define α^j=⟨¯j,fit,¯j,fit⟩‖¯j,fit‖2, α_j= x_j, fit, y_j, fit x_j, fit ^2, and take ^j q_j to be the leading eigenvector of the sample covariance of α^jj,k α_j x_j,k on those same fitting repetitions. The reference mean is j=[j,k] d_j=E[ y_j,k], and the mean discrepancy is j=α^j[j,k]−j b_j= α_jE[ x_j,k]- d_j, conditional on the fitting sample. The decomposition in Equation (10) therefore applies to each matrix with these fitted quantities held fixed. We use the remaining K=24K=24 independent repetitions for evaluation. Set j,k=α^jj,k−j,k z_j,k= α_j x_j,k- y_j,k, tj,k=^j⊤j,kt_j,k= q_j z_j,k, and pj,k=α^j^j⊤j,kp_j,k= α_j q_j x_j,k. With s2s^2 denoting the unbiased sample variance on these repetitions, the unnormalized estimates are B^j=t¯j 2−s2(tj,k)K,N^j=s2(pj,k). B_j= t_j^\,2- s^2(t_j,k)K, N_j=s^2(p_j,k). Conditional on the fit and independent evaluation repetitions, B^j B_j is unbiased for (^j⊤j)2( q_j b_j)^2 and N^j N_j estimates α^j2^j⊤j^j α_j^2 q_j C_j q_j, where j=Cov(j,k) C_j=Cov( x_j,k). Crucially, the correction uses the variance of the paired difference j,k z_j,k; it retains the local–DP cross-covariance. The stochastic term is the variance of a single aggregated update, not that variance divided by the number of evaluation repetitions. For each matrix, the first two numerical columns in Table 5 are B^j/‖^j‖2 B_j/ d_j ^2 and N^j/‖^j‖2 N_j/ d_j ^2, with ^j=K−1∑kj,k d_j=K^-1 _k y_j,k, expressed in units of 10−610^-6. The third column is 100B^j/N^j100 B_j/ N_j. Every column is averaged with equal weight over the 63 matrices. In particular, the reported average ratio is not the ratio of the two reported averages. Bias-corrected mismatch estimates can be negative and are not clipped before averaging. The estimated DP normalizer and the ratios introduce additional finite-sample uncertainty; the normalized entries are not claimed to be unbiased. The matrices share data and are not independent experimental repeats. Trace-based comparison. The corresponding whole-matrix estimates are B^jtr=‖¯j‖2−1K(K−1)∑k‖j,k−¯j‖2,N^jtr=α^j2K−1∑k‖j,k−¯j‖2. B_j^tr= z_j ^2- 1K(K-1) _k z_j,k- z_j ^2, N_j^tr= α_j^2K-1 _k x_j,k- x_j ^2. They estimate ‖j‖2 b_j ^2 and α^j2tr(j) α_j^2tr( C_j), respectively. The same DP-mean normalizer and equal per-matrix averaging apply; the mean/noise ratio is computed per matrix before averaging. At R=8,H=30R=8,H=30, these average ratios are 6.2%6.2\% for SPlus, 2.8%2.8\% for SOAP, 10%10\% for Muon, and 6.7%6.7\% for AdamW. C.3 Curvature along fitted high-noise directions At each method’s own checkpoint, restore the saved parameters and optimizer state before every one-step replay. Form 128 paired update differences i=(i−i′)/2 z_i=( u_i- u_i )/ 2, using fresh local batches of b=256/Rb=256/R sequences. Under the independent-replay model these differences are centered and have the local-update covariance. For Table 6, use the first 64 of the 128 paired samples to select ^=argmax‖=1164∑i=164(⊤i)2. q= q =1 \; 164 _i=1^64( q z_i)^2. Hessian–vector products on 64 held-out validation sequences evaluate curvature at the same checkpoint parameters. Draw 16 independent standard Gaussian vectors in the six-matrix subspace and normalize each to unit length, giving 1,…,16 v_1,…, v_16. Report κ(^)=^⊤^116∑j=116j⊤j.κ( q)= q H q 116 _j=1^16 v_j H v_j. The numerator measures curvature along a direction selected for high noise variance, not maximum curvature. The denominator estimates tr()/dtr( H)/d in this subspace. Each method uses its own trained checkpoint, so the Hessian and reference curvature are method-specific. Ranges span the two estimates from replicas 0 and 1; a single displayed value means both endpoints round to it. These are one-step local-update measurements, not H=30H=30 aggregate measurements. The ratio removes the noise magnitude and does not estimate tr()tr( H C).