How Inner Optimizers Shape Replica Scaling
Summary
该论文研究了分布式LLM训练中内部优化器如何影响副本扩展(replica scaling),识别出伪梯度协方差结构与噪声尺度两个关键因素,并从理论和实证上证明基于白化的预条件优化器能实现最强的worker扩展性能。
View Cached Full Text
Cached at: 10/15/26, 06:50 AM
# How Inner Optimizers Shape Replica Scaling
Source: [https://arxiv.org/html/2610.19064](https://arxiv.org/html/2610.19064)
Eugene Belilovsky Stefan Horoi Edouard Oyallon Benjamin Thérien
###### 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 \(for8−968\-96replicas\), with performance relative to DP improving as parameters and data are scaled\.
## 1Introduction
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](https://arxiv.org/html/2610.19064#bib.bib5);[Douillard et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib6);[Liu et al\., 2024](https://arxiv.org/html/2610.19064#bib.bib17)\)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](https://arxiv.org/html/2610.19064#bib.bib3)\)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](https://arxiv.org/html/2610.19064#bib.bib3)\)\. Recent empirical results show that modifying either the inner optimizer\([Thérien et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib27)\)or the outer aggregation rule\([Horoi et al\., 2026](https://arxiv.org/html/2610.19064#bib.bib9)\)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](https://arxiv.org/html/2610.19064#bib.bib11)\)\. Much of this work and theory\([Stich, 2019](https://arxiv.org/html/2610.19064#bib.bib26);[Karimireddy et al\., 2020](https://arxiv.org/html/2610.19064#bib.bib11);[Yu et al\., 2019](https://arxiv.org/html/2610.19064#bib.bib32)\)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\\beta\_\{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](https://arxiv.org/html/2610.19064#bib.bib9)\), using an Adam inner optimizer, can outperform the much more computationally expensive MuLoCo\([Thérien et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib27)\)\.
- •Our results demonstrate that using a SOAP\([Vyas et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib29)\)or SPlus\([Frans et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib7)\)inner optimizer can lead to very effective replica scaling , obtaining significantly improved performance relative to existing methods from the literature \(as much as6%6\\%less degradation relative to DP at 96 replicas\)\.
## 2Related Work
##### Local\-step methods and DiLoCo\.
Local SGD\([Stich, 2019](https://arxiv.org/html/2610.19064#bib.bib26)\)and federated averaging\([McMahan et al\., 2017](https://arxiv.org/html/2610.19064#bib.bib22)\)reduce communication by taking several local steps between averaging events\. DiLoCo\([Douillard et al\., 2023](https://arxiv.org/html/2610.19064#bib.bib5)\)adapts this to LLM pretraining by utilizing a large number of inner steps, AdamW\([Loshchilov & Hutter, 2019](https://arxiv.org/html/2610.19064#bib.bib18)\)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](https://arxiv.org/html/2610.19064#bib.bib17)\)and overlap or shard the exchange of information between replicas\([Douillard et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib6)\)\.[Charles et al\. \(2025\)](https://arxiv.org/html/2610.19064#bib.bib3)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 asRRgrows\. 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](https://arxiv.org/html/2610.19064#bib.bib24), SPARTA;\)or Top\-kksparsification SparseLoCo\([Sarfi et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib25)\)\. IsoLoco[Horoi et al\. \(2026\)](https://arxiv.org/html/2610.19064#bib.bib9)views DiLoCo’s outer step as task\-arithmetic model merging and replaces plain pseudo\-gradient averaging with Iso\-C merging\([Marczak et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib20)\)\. For each weight matrix, it takes the SVD of the averaged pseudo\-gradient and replaces all singular values with their mean, givingσ¯UV⊤\\bar\{\\sigma\}UV^\{\\top\}\. It then applies the usual Nesterov momentum outer step\. This was shown to improve worker scaling over DiLoCo\([Horoi et al\., 2026](https://arxiv.org/html/2610.19064#bib.bib9)\)\.
##### Small\-batch training and adaptive\-optimizer scaling rules\.
Scaling rules relate optimizer hyperparameters to batch size\([Malladi et al\., 2022](https://arxiv.org/html/2610.19064#bib.bib19)\); most directly,[Marek et al\. \(2025\)](https://arxiv.org/html/2610.19064#bib.bib21)show that Adam trains stably at very small batches providedβ2\\beta\_\{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](https://arxiv.org/html/2610.19064#S4.SS4)transfers theβ2\\beta\_\{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](https://arxiv.org/html/2610.19064#bib.bib8)\)preconditions with Kronecker\-factored curvature estimates; SOAP\([Vyas et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib29)\)runs Adam in Shampoo’s eigenbasis; Muon\([Jordan et al\., 2024](https://arxiv.org/html/2610.19064#bib.bib10)\)orthogonalizes matrix updates with no accumulated second\-moment statistics\. SPlus\([Frans et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib7)\)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](https://arxiv.org/html/2610.19064#bib.bib31)\), 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](https://arxiv.org/html/2610.19064#bib.bib13)\)\. In the DiLoCo setting, MuLoCo\([Thérien et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib27)\)uses Muon as the inner optimizer and reports more directionally consistent pseudo\-gradients asRRgrows, while IsoLoCo\([Horoi et al\., 2026](https://arxiv.org/html/2610.19064#bib.bib9)\)treats outer aggregation as model merging and isotropizes the aggregated pseudo\-gradient\. Section[4](https://arxiv.org/html/2610.19064#S4)develops a complementary criterion: whether the stochastic update geometry emitted by each replica remains well conditioned under averaging\.
## 3Background: the DiLoCo Framework
The methods in this paper can be seen as an instance of Algorithm[1](https://arxiv.org/html/2610.19064#alg1)which is a generalization of local optimization from[Reddi et al\. \(2021\)](https://arxiv.org/html/2610.19064#bib.bib23);[Douillard et al\. \(2023\)](https://arxiv.org/html/2610.19064#bib.bib5)that utilizes both inner \(on the local gradients\) and outer optimization \(on the pseudo\-gradients\)\.RRreplicas start each round from the shared parameters, each runsHHsteps of an inner optimizer on its own data stream\. The resulting displacements𝚫r\\bm\{\\Delta\}\_\{r\}are uniformly averaged, and an outer optimizer applies the aggregate𝚫¯\\bar\{\\bm\{\\Delta\}\}to the shared parameters from the start of the round\. DiLoCo\([Douillard et al\., 2023](https://arxiv.org/html/2610.19064#bib.bib5)\)instantiates the inner optimizer as AdamW and the outer as Nesterov momentum\. For example MuLoCo\([Thérien et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib27)\)replaces the inner optimizer with Muon\. IsoLoCo\([Horoi et al\., 2026](https://arxiv.org/html/2610.19064#bib.bib9)\)isotropizes𝚫¯\\bar\{\\bm\{\\Delta\}\}before the outer step\. The inner optimizer statessrs\_\{r\}are never communicated: each is estimated locally from its replica’s data stream\.
Algorithm 1The DiLoCo framework\. All methods studied here are instances of this algorithm\.1:initial parameters
𝜽\(0\)\\bm\{\\theta\}^\{\(0\)\}; replicas
r∈\{1,…,R\}r\\in\\\{1,\\dots,R\\\}; inner steps
HH; inner optimizerInnerwith per\-replica state
srs\_\{r\}; outer optimizerOuterwith state
SS; global batch size
BB;
2:forouter round
t=0,…,T−1t=0,\\dots,T\-1do
3:for allreplicas
rrin paralleldo
4:
𝜽r←𝜽\(t\)\\bm\{\\theta\}\_\{r\}\\leftarrow\\bm\{\\theta\}^\{\(t\)\}
5:for
h=1,…,Hh=1,\\dots,Hdo
6:
𝒈←∇𝜽ℒ\(𝜽r,xr,h\)\\bm\{g\}\\leftarrow\\nabla\_\{\\bm\{\\theta\}\}\\mathcal\{L\}\(\\bm\{\\theta\}\_\{r\};x\_\{r,h\}\)⊳\\trianglerightlocal minibatch:B/RB/Rtokens
7:
\(𝜽r,sr\)←Inner\(𝜽r,𝒈,sr\)\(\\bm\{\\theta\}\_\{r\},\\,s\_\{r\}\)\\leftarrow\\textsc\{Inner\}\(\\bm\{\\theta\}\_\{r\},\\,\\bm\{g\},\\,s\_\{r\}\)⊳\\trianglerightsrs\_\{r\}is never synchronized
8:endfor
9:
𝚫r←𝜽\(t\)−𝜽r\\bm\{\\Delta\}\_\{r\}\\leftarrow\\bm\{\\theta\}^\{\(t\)\}\-\\bm\{\\theta\}\_\{r\}
10:endfor
11:
𝚫¯←1R∑r𝚫r\\bar\{\\bm\{\\Delta\}\}\\leftarrow\\tfrac\{1\}\{R\}\\sum\_\{r\}\\bm\{\\Delta\}\_\{r\}⊳\\trianglerightuniform average
12:
\(𝜽\(t\+1\),S\)←Outer\(𝜽\(t\),𝚫¯,S\)\(\\bm\{\\theta\}^\{\(t\+1\)\},\\,S\)\\leftarrow\\textsc\{Outer\}\(\\bm\{\\theta\}^\{\(t\)\},\\,\\bar\{\\bm\{\\Delta\}\},\\,S\)⊳\\trianglerightNesterov momentum in DiLoCo
13:endfor
## 4Inner 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](https://arxiv.org/html/2610.19064#S4.SS1)introduces the optimizers considered in this work, Section[4\.2](https://arxiv.org/html/2610.19064#S4.SS2)separates the update error \(relative to non\-local steps with the same base optimizer\) into amean mismatchterm and astochastic errorterm\. It shows how whitening controls the largest directional variance, and how matching covariance structures across replicas makes uniform averaging optimal\. Section[4\.4](https://arxiv.org/html/2610.19064#S4.SS4)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](https://arxiv.org/html/2610.19064#S4.SS4)bounds the Adam mean mismatch caused by fresh\-gradient second moments and explains how appropriately adjustingβ2\\beta\_\{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\.1Background 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\\bm\{G\}\\in\\mathbb\{R\}^\{m\\times n\}\. An optimizer returns a direction𝑼\\bm\{U\}and an updated state after observing𝑮\\bm\{G\}; its parameter displacement is𝚫=ηvec\(𝑼\)\\bm\{\\Delta\}=\\eta\\,\\operatorname\{vec\}\(\\bm\{U\}\)\. We use subscriptst−1t\-1andttfor the incoming and updated optimizer states, respectively, and suppress the step index on the gradient and returned direction\. The first\-moment proposal is
𝑴t=β1𝑴t−1\+\(1−β1\)𝑮\.\\bm\{M\}\_\{t\}=\\beta\_\{1\}\\bm\{M\}\_\{t\-1\}\+\(1\-\\beta\_\{1\}\)\\bm\{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γ\\gammaof Muon is shown explicitly\. All decay coefficients lie in\[0,1\)\[0,1\); numerical damping is denoted byϵ\>0\\epsilon\>0\.
##### Adam and Shampoo\.
Adam uses an entrywise second moment in parameter coordinates\([Kingma & Ba, 2015](https://arxiv.org/html/2610.19064#bib.bib14);[Loshchilov & Hutter, 2019](https://arxiv.org/html/2610.19064#bib.bib18)\):
𝑽A,t=β2𝑽A,t−1\+\(1−β2\)𝑮⊙2,𝑼Adam=𝑴t⊘\(𝑽A,t\+ϵ\)\.\\bm\{V\}\_\{A,t\}=\\beta\_\{2\}\\bm\{V\}\_\{A,t\-1\}\+\(1\-\\beta\_\{2\}\)\\bm\{G\}^\{\\odot 2\},\\qquad\\bm\{U\}\_\{\\rm Adam\}=\\bm\{M\}\_\{t\}\\oslash\(\\sqrt\{\\bm\{V\}\_\{A,t\}\}\+\\epsilon\)\.\(2\)WhereUUis the update\. Here,⊙\\odotand⊘\\oslashdenote elementwise multiplication and division, respectively, with𝑮⊙2\\bm\{G\}^\{\\odot 2\}representing the elementwise square\. Shampoo instead maintains left and right Gram factors\([Gupta et al\., 2018](https://arxiv.org/html/2610.19064#bib.bib8)\):
𝑳t=β3𝑳t−1\+\(1−β3\)𝑮𝑮⊤,𝗥t=β3𝗥t−1\+\(1−β3\)𝑮⊤𝑮,𝑼Shampoo=\(𝑳t\+ϵ𝑰\)−1/4𝑴t\(𝗥t\+ϵ𝑰\)−1/4\.\\begin\{split\}\\bm\{L\}\_\{t\}&=\\beta\_\{3\}\\bm\{L\}\_\{t\-1\}\+\(1\-\\beta\_\{3\}\)\\bm\{G\}\\bm\{G\}^\{\\top\},\\qquad\\bm\{\\mathsf\{R\}\}\_\{t\}=\\beta\_\{3\}\\bm\{\\mathsf\{R\}\}\_\{t\-1\}\+\(1\-\\beta\_\{3\}\)\\bm\{G\}^\{\\top\}\\bm\{G\},\\\\ \\bm\{U\}\_\{\\rm Shampoo\}&=\(\\bm\{L\}\_\{t\}\+\\epsilon\\bm\{I\}\)^\{\-1/4\}\\bm\{M\}\_\{t\}\(\\bm\{\\mathsf\{R\}\}\_\{t\}\+\\epsilon\\bm\{I\}\)^\{\-1/4\}\.\\end\{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](https://arxiv.org/html/2610.19064#bib.bib29)\)\. 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](https://arxiv.org/html/2610.19064#bib.bib16)\)\.
##### SOAP and SPlus\.
Let𝑸L,𝑸R\\bm\{Q\}\_\{L\},\\bm\{Q\}\_\{R\}be orthogonal eigenvector matrices obtained from the historical Gram factors above\. Their action is the matrix rotation
𝑮~=𝑸L⊤𝑮𝑸R,𝑴~t=𝑸L⊤𝑴t𝑸R\.\\widetilde\{\\bm\{G\}\}=\\bm\{Q\}\_\{L\}^\{\\top\}\\bm\{G\}\\bm\{Q\}\_\{R\},\\qquad\\widetilde\{\\bm\{M\}\}\_\{t\}=\\bm\{Q\}\_\{L\}^\{\\top\}\\bm\{M\}\_\{t\}\\bm\{Q\}\_\{R\}\.\(4\)SOAP \(ShampoO with Adam in the Preconditioner’s eigenbasis\) performs Adam\-style normalization in this product basis\([Vyas et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib29)\):
𝑽S,t=β2𝑽S,t−1\+\(1−β2\)𝑮~⊙2,𝑼SOAP=𝑸L\[𝑴~t⊘\(𝑽S,t\+ϵ\)\]𝑸R⊤\.\\begin\{split\}\\bm\{V\}\_\{S,t\}&=\\beta\_\{2\}\\bm\{V\}\_\{S,t\-1\}\+\(1\-\\beta\_\{2\}\)\\widetilde\{\\bm\{G\}\}^\{\\odot 2\},\\\\ \\bm\{U\}\_\{\\rm SOAP\}&=\\bm\{Q\}\_\{L\}\\big\[\\widetilde\{\\bm\{M\}\}\_\{t\}\\oslash\(\\sqrt\{\\bm\{V\}\_\{S,t\}\}\+\\epsilon\)\\big\]\\bm\{Q\}\_\{R\}^\{\\top\}\.\\end\{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](https://arxiv.org/html/2610.19064#bib.bib7)\)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\.\\bm\{U\}\_\{\\rm SPlus\}=\\bm\{Q\}\_\{L\}\\sign\(\\widetilde\{\\bm\{M\}\}\_\{t\}\)\\bm\{Q\}\_\{R\}^\{\\top\},\\qquad\\sign\(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](https://arxiv.org/html/2610.19064#bib.bib29);[Frans et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib7);[Lin et al\., 2026](https://arxiv.org/html/2610.19064#bib.bib16)\)\. Section[4\.2](https://arxiv.org/html/2610.19064#S4.SS2)gives the corresponding covariance identities under explicit assumptions and explains their impact on pseudo\-gradient aggregation\.
##### Muon\.
For a singular value decomposition𝑴t=𝑼M𝚺M𝑽M⊤\\bm\{M\}\_\{t\}=\\bm\{U\}\_\{M\}\\bm\{\\Sigma\}\_\{M\}\\bm\{V\}\_\{M\}^\{\\top\}on the nonzero singular subspace, ideal Muon returns
𝑼Muon=γpolar\(𝑴t\)=γ𝑼M𝑽M⊤\.\\bm\{U\}\_\{\\rm Muon\}=\\gamma\\polar\(\\bm\{M\}\_\{t\}\)=\\gamma\\bm\{U\}\_\{M\}\\bm\{V\}\_\{M\}^\{\\top\}\.\(7\)It replaces nonzero singular values by a common value\([Jordan et al\., 2024](https://arxiv.org/html/2610.19064#bib.bib10)\)\. In matched singular\-vector coordinates,𝑼M⊤𝑴t𝑽M=𝚺M\\bm\{U\}\_\{M\}^\{\\top\}\\bm\{M\}\_\{t\}\\bm\{V\}\_\{M\}=\\bm\{\\Sigma\}\_\{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\.2Impact 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 mismatchand \(2\)stochastic error\. The covariance analysis in this subsection addresses the stochastic term; Section[4\.4](https://arxiv.org/html/2610.19064#S4.SS4)then addresses the mean mismatch\. Following the pseudo\-gradient interpretation of client displacements in[Reddi et al\. \(2021\)](https://arxiv.org/html/2610.19064#bib.bib23)and[Wang et al\. \(2024\)](https://arxiv.org/html/2610.19064#bib.bib30), 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\\bm\{\\Delta\}\_\{r\}\\in\\mathbb\{R\}^\{d\}be the completeHH\-step displacement returned by replicarr\. 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\)\.\\bm\{\\Delta\}\_\{r\}=\\bm\{\\mu\}\_\{r\}\+\\bm\{\\varepsilon\}\_\{r\},\\qquad\\bm\{\\mu\}\_\{r\}=\\mathbb\{E\}\[\\bm\{\\Delta\}\_\{r\}\],\\qquad\\mathbb\{E\}\[\\bm\{\\varepsilon\}\_\{r\}\]=0,\\qquad\\bm\{C\}\_\{r\}=\\operatorname\{Cov\}\(\\bm\{\\Delta\}\_\{r\}\)\.\(8\)Assume finite second moments,𝔼\[‖Δr‖22\]<∞\\mathbb\{E\}\[\\\|\\Delta\_\{r\}\\\|\_\{2\}^\{2\}\]<\\inftyfor allrr, and zero cross\-replica noise covariances,Cov\(εr,εs\)=𝟎\\operatorname\{Cov\}\(\\varepsilon\_\{r\},\\varepsilon\_\{s\}\)=\\bm\{0\}for allr≠sr\\neq 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\\bar\{\\bm\{\\Delta\}\}=R^\{\-1\}\\sum\_\{r\}\\bm\{\\Delta\}\_\{r\}and𝝁¯=R−1∑r𝝁r\\bar\{\\bm\{\\mu\}\}=R^\{\-1\}\\sum\_\{r\}\\bm\{\\mu\}\_\{r\}\. Note that for the typical setting studied in[Douillard et al\. \(2023\)](https://arxiv.org/html/2610.19064#bib.bib5);[Thérien et al\. \(2025\)](https://arxiv.org/html/2610.19064#bib.bib27)with i\.i\.d\. data and identical incoming optimizer states,𝝁r\\bm\{\\mu\}\_\{r\}is the same for all replicas\.
We now consider𝒅HDP\\bm\{d\}\_\{H\}^\{\\rm DP\}: the displacement of the*same inner optimizer*run forHH*synchronized*steps and same global batch size, in expectation\. We allow a deterministic scalar calibration factorαR\\alpha\_\{R\}, and define
𝒃R=αR𝝁¯−𝒅HDP,𝑪agg=Cov\(𝚫¯\)=1R2∑r𝑪r\.\\bm\{b\}\_\{R\}=\\alpha\_\{R\}\\bar\{\\bm\{\\mu\}\}\-\\bm\{d\}\_\{H\}^\{\\rm DP\},\\qquad\\bm\{C\}\_\{\\rm agg\}=\\operatorname\{Cov\}\(\\bar\{\\bm\{\\Delta\}\}\)=\\frac\{1\}\{R^\{2\}\}\\sum\_\{r\}\\bm\{C\}\_\{r\}\.\(9\)ForH=1H=1,αR\\alpha\_\{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𝐪\\bm\{q\},
𝔼\[\(𝒒⊤\(αR𝚫¯−𝒅HDP\)\)2\]⏟error relative to the DP mean=\(𝒒⊤𝒃R\)2⏟mean mismatch\+αR2𝒒⊤𝑪agg𝒒⏟stochastic error\.\\underbrace\{\\mathbb\{E\}\\\!\\left\[\\big\(\\bm\{q\}^\{\\top\}\(\\alpha\_\{R\}\\bar\{\\bm\{\\Delta\}\}\-\\bm\{d\}\_\{H\}^\{\\rm DP\}\)\\big\)^\{2\}\\right\]\}\_\{\\text\{error relative to the DP mean\}\}=\\underbrace\{\(\\bm\{q\}^\{\\top\}\\bm\{b\}\_\{R\}\)^\{2\}\}\_\{\\text\{mean mismatch\}\}\+\\underbrace\{\\alpha\_\{R\}^\{2\}\\bm\{q\}^\{\\top\}\\bm\{C\}\_\{\\rm agg\}\\bm\{q\}\}\_\{\\text\{stochastic error\}\}\.\(10\)Consequently, the largest directional error is
λmax\(𝒃R𝒃R⊤\+αR2𝑪agg\)≤‖𝒃R‖22\+αR2λmax\(𝑪agg\)\.\\lambda\_\{\\max\}\(\\bm\{b\}\_\{R\}\\bm\{b\}\_\{R\}^\{\\top\}\+\\alpha\_\{R\}^\{2\}\\bm\{C\}\_\{\\rm agg\}\)\\leq\\\|\\bm\{b\}\_\{R\}\\\|\_\{2\}^\{2\}\+\\alpha\_\{R\}^\{2\}\\lambda\_\{\\max\}\(\\bm\{C\}\_\{\\rm agg\}\)\.\(11\)
###### Proof\.
Subtract and addαR𝝁¯\\alpha\_\{R\}\\bar\{\\bm\{\\mu\}\}\. Using the Rayleigh quotient gives \([11](https://arxiv.org/html/2610.19064#S4.E11)\)\. ∎
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](https://arxiv.org/html/2610.19064#S4.SS4)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𝑷\\bm\{P\}whitens a covariance𝚺\\bm\{\\Sigma\}if𝑷𝚺𝑷⊤=𝑰\\bm\{P\}\\bm\{\\Sigma\}\\bm\{P\}^\{\\top\}=\\bm\{I\}\([Kessy et al\., 2018](https://arxiv.org/html/2610.19064#bib.bib12)\); a prescribed scalar multiple gives a spherical covariance at another scale\. For any covariance𝑪\\bm\{C\}of dimensionddand traceτ\\tau\(fixed total variance\), the following inequality holds:
λmax\(𝑪\)≥τ/d,with equality precisely when𝑪=\(τ/d\)𝑰\.\\lambda\_\{\\max\}\(\\bm\{C\}\)\\geq\\tau/d,\\quad\\text\{with equality precisely when \}\\bm\{C\}=\(\\tau/d\)\\bm\{I\}\.\(12\)Thus, isotropy minimizes the*largest*directional variance at fixed total variance, precisely the covariance quantity appearing in \([11](https://arxiv.org/html/2610.19064#S4.E11)\)\. Moreover, a bound on the maximum directional variance of each replica carries over to the aggregate:
λmax\(𝑪agg\)≤1R2∑rλmax\(𝑪r\)\.\\lambda\_\{\\max\}\(\\bm\{C\}\_\{\\rm agg\}\)\\leq\\frac\{1\}\{R^\{2\}\}\\sum\_\{r\}\\lambda\_\{\\max\}\(\\bm\{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\.\\sum\_\{r=1\}^\{R\}w\_\{r\}^\{2\}\\,\\bm\{q\}^\{\\top\}\\bm\{C\}\_\{r\}\\bm\{q\},\\qquad\\sum\_\{r\}w\_\{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](https://arxiv.org/html/2610.19064#bib.bib1)\); they do so in every direction when𝑪1=⋯=𝑪R\\bm\{C\}\_\{1\}=\\cdots=\\bm\{C\}\_\{R\}\. Whitening each replica toward the same spherical targetν𝑰\\nu\\bm\{I\}promotes precisely this agreement, since‖𝑪i−𝑪j‖≤‖𝑪i−ν𝑰‖\+‖𝑪j−ν𝑰‖\.\\\|\\bm\{C\}\_\{i\}\-\\bm\{C\}\_\{j\}\\\|\\leq\\\|\\bm\{C\}\_\{i\}\-\\nu\\bm\{I\}\\\|\+\\\|\\bm\{C\}\_\{j\}\-\\nu\\bm\{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\\bm\{b\}\_\{R\}, which is analyzed separately in Section[4\.4](https://arxiv.org/html/2610.19064#S4.SS4)\.
### 4\.3Impact 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](https://arxiv.org/html/2610.19064#bib.bib29)\), the whitening motivation of[Frans et al\. \(2025\)](https://arxiv.org/html/2610.19064#bib.bib7), and the covariance\-estimation analysis of[Lin et al\. \(2026\)](https://arxiv.org/html/2610.19064#bib.bib16)\.
For this comparison, consider a random matrix proposal𝑴=𝑴t∈ℝm×n\\bm\{M\}=\\bm\{M\}\_\{t\}\\in\\mathbb\{R\}^\{m\\times n\}, with𝚺=Cov\(vec\(𝑴\)\)≻0\\bm\{\\Sigma\}=\\operatorname\{Cov\}\(\\operatorname\{vec\}\(\\bm\{M\}\)\)\\succ 0\. Let𝑼\\bm\{U\}denote the resulting update matrix, and defined=mnd=mnand𝑪=Cov\(vec\(𝑼\)\)\\bm\{C\}=\\operatorname\{Cov\}\(\\operatorname\{vec\}\(\\bm\{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\(𝑴\)\.\\bm\{Q\}=\\bm\{Q\}\_\{R\}\\otimes\\bm\{Q\}\_\{L\},\\qquad\\operatorname\{vec\}\(\\bm\{Q\}\_\{L\}^\{\\top\}\\bm\{M\}\\bm\{Q\}\_\{R\}\)=\\bm\{Q\}^\{\\top\}\\operatorname\{vec\}\(\\bm\{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𝑴\\bm\{M\}is divided by its true standard deviation\. Writing𝑫=diag\(Σ11,…,Σdd\)\\bm\{D\}=\\operatorname\{diag\}\(\\Sigma\_\{11\},\\ldots,\\Sigma\_\{dd\}\), this gives
vec\(𝑼\)=𝑫−1/2vec\(𝑴\),𝑪=𝑫−1/2𝚺𝑫−1/2\.\\operatorname\{vec\}\(\\bm\{U\}\)=\\bm\{D\}^\{\-1/2\}\\operatorname\{vec\}\(\\bm\{M\}\),\\qquad\\bm\{C\}=\\bm\{D\}^\{\-1/2\}\\bm\{\\Sigma\}\\bm\{D\}^\{\-1/2\}\.Every diagonal entry of𝑪\\bm\{C\}is one, butCij=Σij/ΣiiΣjjC\_\{ij\}=\\Sigma\_\{ij\}/\\sqrt\{\\Sigma\_\{ii\}\\Sigma\_\{jj\}\}\. Thus, even with access to the true coordinate variances, Adam’s normalization preserves the correlations present in𝑴\\bm\{M\}\.
##### Shampoo\.
Consider the idealized case where the factors equal the population Gram matrices of a zero\-mean input:𝑳=𝔼\[𝑴𝑴⊤\]\\bm\{L\}=\\mathbb\{E\}\[\\bm\{M\}\\bm\{M\}^\{\\top\}\]and𝗥=𝔼\[𝑴⊤𝑴\]\\bm\{\\mathsf\{R\}\}=\\mathbb\{E\}\[\\bm\{M\}^\{\\top\}\\bm\{M\}\]\. Assume the covariance is separable,𝚺=\(𝗥⊗𝑳\)/tr\(𝑳\)\\bm\{\\Sigma\}=\(\\bm\{\\mathsf\{R\}\}\\otimes\\bm\{L\}\)/\\operatorname\{tr\}\(\\bm\{L\}\), with positive\-definite factors\. The canonical Shampoo transformation then gives:
𝑼=𝑳−1/4𝑴𝗥−1/4,𝑪=𝗥1/2⊗𝑳1/2tr\(𝑳\)\.\\bm\{U\}=\\bm\{L\}^\{\-1/4\}\\bm\{M\}\\bm\{\\mathsf\{R\}\}^\{\-1/4\},\\qquad\\bm\{C\}=\\frac\{\\bm\{\\mathsf\{R\}\}^\{1/2\}\\otimes\\bm\{L\}^\{1/2\}\}\{\\operatorname\{tr\}\(\\bm\{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\)](https://arxiv.org/html/2610.19064#bib.bib16)\.
##### SOAP\.
Suppose the fixed rotation𝑴~t=𝑸L⊤𝑴t𝑸R\\widetilde\{\\bm\{M\}\}\_\{t\}=\\bm\{Q\}\_\{L\}^\{\\top\}\\bm\{M\}\_\{t\}\\bm\{Q\}\_\{R\}gives uncorrelated coordinates, and let𝑽S,t\\bm\{V\}\_\{S,t\}denote their true variances:\(𝑽S,t\)ij=Var\(\(𝑴~t\)ij\)\(\\bm\{V\}\_\{S,t\}\)\_\{ij\}=\\operatorname\{Var\}\(\(\\widetilde\{\\bm\{M\}\}\_\{t\}\)\_\{ij\}\)\. Dividing each rotated entry by its true standard deviation and rotating back gives
𝑼=𝑸L\[𝑴~t⊘𝑽S,t\]𝑸R⊤,𝑪=𝑰\.\\bm\{U\}=\\bm\{Q\}\_\{L\}\[\\widetilde\{\\bm\{M\}\}\_\{t\}\\oslash\\sqrt\{\\bm\{V\}\_\{S,t\}\}\]\\bm\{Q\}\_\{R\}^\{\\top\},\\qquad\\bm\{C\}=\\bm\{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𝑸\\bm\{Q\}decorrelates the proposal coordinates, and their variances𝑽S,t\\bm\{V\}\_\{S,t\}are known exactly\.
##### SPlus\.
Supposevec\(𝑴\)\\operatorname\{vec\}\(\\bm\{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\+1and−1\-1values\. Therefore
𝑼=𝑸Lsign\(𝑴~t\)𝑸R⊤,𝑪=𝑰\.\\bm\{U\}=\\bm\{Q\}\_\{L\}\\operatorname\{sign\}\(\\widetilde\{\\bm\{M\}\}\_\{t\}\)\\bm\{Q\}\_\{R\}^\{\\top\},\\qquad\\bm\{C\}=\\bm\{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\\bm\{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\\bm\{Q\}\_\{L\},\\bm\{Q\}\_\{R\}before samplingvec\(𝑴t\)∼𝒩\(𝟎,𝚺\)\\operatorname\{vec\}\(\\bm\{M\}\_\{t\}\)\\sim\\mathcal\{N\}\(\\bm\{0\},\\bm\{\\Sigma\}\), and let𝓒\\bm\{\\mathcal\{C\}\}be the correlation matrix of the entries of𝑴~t=𝑸L⊤𝑴t𝑸R\\widetilde\{\\bm\{M\}\}\_\{t\}=\\bm\{Q\}\_\{L\}^\{\\top\}\\bm\{M\}\_\{t\}\\bm\{Q\}\_\{R\}\. To characterize how SPlus transforms these correlations, we apply the Gaussian arcsine law\([Dirksen et al\., 2022](https://arxiv.org/html/2610.19064#bib.bib4)\), which gives the covariance after taking signs\. Rotating back yields
𝑪SPlus=𝑸\[2πarcsin\(𝓒\)\]𝑸⊤,\|2πarcsin\(ρ\)\|≤\|ρ\|\.\\bm\{C\}\_\{\\rm SPlus\}=\\bm\{Q\}\\left\[\\frac\{2\}\{\\pi\}\\arcsin\(\\bm\{\\mathcal\{C\}\}\)\\right\]\\bm\{Q\}^\{\\top\},\\qquad\\left\|\\frac\{2\}\{\\pi\}\\arcsin\(\\rho\)\\right\|\\leq\|\\rho\|\.\(13\)Herearcsin\\arcsinacts 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=𝑸𝓒𝑸⊤\\bm\{C\}\_\{\\rm SOAP\}=\\bm\{Q\}\\bm\{\\mathcal\{C\}\}\\bm\{Q\}^\{\\top\}: 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\\displaystyle\\\|\\bm\{C\}\_\{\\rm SPlus\}\-\\bm\{I\}\\\|\_\{F\}^\{2\}=∑i≠j\[2πarcsin\(𝒞ij\)\]2\\displaystyle=\\sum\_\{i\\neq j\}\\left\[\\frac\{2\}\{\\pi\}\\arcsin\(\\mathcal\{C\}\_\{ij\}\)\\right\]^\{2\}≤∑i≠j𝒞ij2=‖𝑪SOAP−𝑰‖F2\.\\displaystyle\\leq\\sum\_\{i\\neq j\}\\mathcal\{C\}\_\{ij\}^\{2\}=\\\|\\bm\{C\}\_\{\\rm SOAP\}\-\\bm\{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𝑴\\bm\{M\}is full rank almost surely, and its distribution is unchanged by any orthogonal transformation𝑴↦𝑶𝑴\\bm\{M\}\\mapsto\\bm\{O\}\\bm\{M\}whenm≥nm\\geq n\(or𝑴↦𝑴𝑶\\bm\{M\}\\mapsto\\bm\{M\}\\bm\{O\}whenm<nm<n\)\. For the exact polar transformation with a fixed scaleγ\>0\\gamma\>0,
𝑼=γpolar\(𝑴\),𝑪=γ2max\(m,n\)𝑰\.\\bm\{U\}=\\gamma\\operatorname\{polar\}\(\\bm\{M\}\),\\qquad\\bm\{C\}=\\frac\{\\gamma^\{2\}\}\{\\max\(m,n\)\}\\bm\{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](https://arxiv.org/html/2610.19064#bib.bib27)\)\. 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](https://arxiv.org/html/2610.19064#S4.E10)\) and \([11](https://arxiv.org/html/2610.19064#S4.E11)\) apply to the complete displacement after any number of inner optimization stepsHH\. However, if𝚫r,h\\bm\{\\Delta\}\_\{r,h\}denotes the displacement at thehh\-th local step, the covariance of replicarr’s complete displacement is given by
𝑪r=∑h=1HCov\(𝚫r,h\)\+∑h≠kCov\(𝚫r,h,𝚫r,k\)\.\\bm\{C\}\_\{r\}=\\sum\_\{h=1\}^\{H\}\\operatorname\{Cov\}\(\\bm\{\\Delta\}\_\{r,h\}\)\+\\sum\_\{h\\neq k\}\\operatorname\{Cov\}\(\\bm\{\\Delta\}\_\{r,h\},\\bm\{\\Delta\}\_\{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\.4Single\-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](https://arxiv.org/html/2610.19064#S4.SS2)\. This eliminates multi\-step local drift, allowing us to study the remaining mean mismatch𝒃R\\bm\{b\}\_\{R\}more easily\. At fixed global batch sizeBB, let𝒈^r≔vec\(𝑮r\)\\widehat\{\\bm\{g\}\}\_\{r\}\\coloneqq\\operatorname\{vec\}\(\\bm\{G\}\_\{r\}\)denote the independent and identically distributed replica gradients, and define
𝒈≔𝔼\[𝒈^r\],𝒈¯≔1R∑r𝒈^r,𝚺g≔Cov\(𝒈¯\)\.\\bm\{g\}\\coloneqq\\mathbb\{E\}\[\\widehat\{\\bm\{g\}\}\_\{r\}\],\\qquad\\bar\{\\bm\{g\}\}\\coloneqq\\frac\{1\}\{R\}\\sum\_\{r\}\\widehat\{\\bm\{g\}\}\_\{r\},\\qquad\\bm\{\\Sigma\}\_\{g\}\\coloneqq\\operatorname\{Cov\}\(\\bar\{\\bm\{g\}\}\)\.\(14\)ThenCov\(𝒈^r\)=R𝚺g\\operatorname\{Cov\}\(\\widehat\{\\bm\{g\}\}\_\{r\}\)=R\\bm\{\\Sigma\}\_\{g\}\. For IID examples,𝚺g=𝛀/B\\bm\{\\Sigma\}\_\{g\}=\\bm\{\\Omega\}/B, where𝛀\\bm\{\\Omega\}is the per\-example covariance\. LetT\(𝒈^,s\)T\(\\widehat\{\\bm\{g\}\};s\)denote the parameter\-update map that takes a gradient𝒈^\\hat\{\\bm\{g\}\}and incoming optimizer statessand returns the corresponding parameter displacement\. Then
𝚫¯=1R∑rT\(𝒈^r,s\),𝒅1DP=𝔼\[T\(𝒈¯,s\)\],𝒃R=αR𝔼\[𝚫¯\]−𝒅1DP\.\\bar\{\\bm\{\\Delta\}\}=\\frac\{1\}\{R\}\\sum\_\{r\}T\(\\widehat\{\\bm\{g\}\}\_\{r\};s\),\\qquad\\bm\{d\}\_\{1\}^\{\\rm DP\}=\\mathbb\{E\}\[T\(\\bar\{\\bm\{g\}\};s\)\],\\qquad\\bm\{b\}\_\{R\}=\\alpha\_\{R\}\\mathbb\{E\}\[\\bar\{\\bm\{\\Delta\}\}\]\-\\bm\{d\}\_\{1\}^\{\\rm DP\}\.\(15\)In this single\-step setting, DP applies the optimizer transformationTTto the averaged gradient𝒈¯\\bar\{\\bm\{g\}\}, whereas DiLoCo appliesTTindependently to each replica’s gradient before averaging the resulting updates\. While an affineTTcommutes 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](https://arxiv.org/html/2610.19064#S4.E1)\)\. The average of the replica momenta is exactly the DP momentum from the same incoming state\. Letm~r,t,j\\widetilde\{m\}\_\{r,t,j\}denote entryjjofvec\(𝑴~r,t\)\\operatorname\{vec\}\(\\widetilde\{\\bm\{M\}\}\_\{r,t\}\), forj=1,…,d=mnj=1,\\ldots,d=mn, and assume
m~r,t,j=σj\(aj\+RZr,j\),Zr,j∼iid𝒩\(0,1\),σj\>0\.\\widetilde\{m\}\_\{r,t,j\}=\\sigma\_\{j\}\(a\_\{j\}\+\\sqrt\{R\}\\,Z\_\{r,j\}\),\\qquad Z\_\{r,j\}\\stackrel\{\{\\scriptstyle\\rm iid\}\}\{\{\\sim\}\}\\mathcal\{N\}\(0,1\),\\quad\\sigma\_\{j\}\>0\.\(16\)Thusaja\_\{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](https://arxiv.org/html/2610.19064#S4.E16)\), set𝐚=\(aj\)j=1d\\bm\{a\}=\(a\_\{j\}\)\_\{j=1\}^\{d\},ψ\(t\)=2Φ\(t\)−1\\psi\(t\)=2\\Phi\(t\)\-1, andc=2/πc=\\sqrt\{2/\\pi\}\. WithαR=R\\alpha\_\{R\}=\\sqrt\{R\},
𝒃R=𝑸\[Rψ\(𝒂/R\)−ψ\(𝒂\)\],‖𝒃R‖22≤c236\(1−1R\)2∑jaj6,Cov\(R𝚫¯\)=𝑸Diag\(1−ψ\(aj/R\)2\)𝑸⊤⪯𝑰\.\\begin\{split\}\\bm\{b\}\_\{R\}&=\\bm\{Q\}\\big\[\\sqrt\{R\}\\,\\psi\(\\bm\{a\}/\\sqrt\{R\}\)\-\\psi\(\\bm\{a\}\)\\big\],\\\\ \\\|\\bm\{b\}\_\{R\}\\\|\_\{2\}^\{2\}&\\leq\\frac\{c^\{2\}\}\{36\}\\left\(1\-\\frac\{1\}\{R\}\\right\)^\{2\}\\sum\_\{j\}a\_\{j\}^\{6\},\\\\ \\operatorname\{Cov\}\(\\sqrt\{R\}\\,\\bar\{\\bm\{\\Delta\}\}\)&=\\bm\{Q\}\\operatorname\{Diag\}\\big\(1\-\\psi\(a\_\{j\}/\\sqrt\{R\}\)^\{2\}\\big\)\\bm\{Q\}^\{\\top\}\\preceq\\bm\{I\}\.\\end\{split\}\(17\)Hereψ\\psiacts entrywise\. In particular, every directional error in \([10](https://arxiv.org/html/2610.19064#S4.E10)\) is at most1\+c236\(1−R−1\)2∑jaj61\+\\frac\{c^\{2\}\}\{36\}\(1\-R^\{\-1\}\)^\{2\}\\sum\_\{j\}a\_\{j\}^\{6\}\.
The proof is given in Appendix[A\.2](https://arxiv.org/html/2610.19064#A1.SS2)\. For a small vector𝒂\\bm\{a\}, the mean mismatch is third order and its squared contribution is sixth order\. At𝒂=0\\bm\{a\}=0, the calibrated local aggregate and DP have the same mean and covariance, although their distributions differ\. TheR\\sqrt\{R\}calibration makes the stochastic term order one; it does not retain an additional1/R1/Rvariance gain at fixed global batch\.
##### Mean mismatch for Adam
We now bound the mean mismatch𝒃R\\bm\{b\}\_\{R\}for Adam\. At fixed global batch size, the local gradient variance isRRtimes that of the DP average\. Fresh\-gradient second moments therefore change the normalization even from a common incoming state\. Proposition[3](https://arxiv.org/html/2610.19064#Thmproposition3)bounds the resultingmean\-update error; Appendix[A\.3](https://arxiv.org/html/2610.19064#A1.SS3)gives the second\-moment identities and the Gaussian variance\-differentiation argument\([Voigtlaender, 2021](https://arxiv.org/html/2610.19064#bib.bib28)\)\.
###### Proposition 3\(Second\-moment memory controls a one\-step mean bound\)\.
In a fixed parameter coordinatejj, let the incoming Adam state be\(m,v\)\(m,v\), withv\>0v\>0and0<β2<10<\\beta\_\{2\}<1\. Omitϵ\\epsilon, 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\)=\\frac\{\\beta\_\{1\}m\+\(1\-\\beta\_\{1\}\)x\}\{\\sqrt\{\\beta\_\{2\}v\+\(1\-\\beta\_\{2\}\)x^\{2\}\}\}\.\(18\)Supposeg^r∼𝒩\(g,Rσ2\)\\widehat\{g\}\_\{r\}\\sim\\mathcal\{N\}\(g,R\\sigma^\{2\}\)independently, sog¯∼𝒩\(g,σ2\)\\bar\{g\}\\sim\\mathcal\{N\}\(g,\\sigma^\{2\}\)\. WithαR=1\\alpha\_\{R\}=1, writebR,jb\_\{R,j\}for this coordinate of the bias𝐛R\\bm\{b\}\_\{R\}from \([15](https://arxiv.org/html/2610.19064#S4.E15)\)\. It obeys
\|bR,j\|≤\(1−β2\)\(R−1\)σ22\(β2v\)3/2\(\|β1m\|\+3\(1−β1\)\|g\|\)\.\|b\_\{R,j\}\|\\leq\\frac\{\(1\-\\beta\_\{2\}\)\(R\-1\)\\sigma^\{2\}\}\{2\(\\beta\_\{2\}v\)^\{3/2\}\}\\left\(\|\\beta\_\{1\}m\|\+3\(1\-\\beta\_\{1\}\)\|g\|\\right\)\.\(19\)For SOAP with a basis fixed before the fresh gradients, the same bound applies to the corresponding coordinate of𝐐⊤𝐛R\\bm\{Q\}^\{\\top\}\\bm\{b\}\_\{R\}\.
This is a finite\-noise bound, proved in Appendix[A\.3\.1](https://arxiv.org/html/2610.19064#A1.SS3.SSS1)\. For fixedβ2\\beta\_\{2\}, the upper bound grows linearly withR−1R\-1, motivating an increase inβ2\\beta\_\{2\}as the number of replicas grows\. Adapting the scaling rule of[Marek et al\. \(2025\)](https://arxiv.org/html/2610.19064#bib.bib21)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\)\.\\beta\_\{2\}\(R\)=\\beta\_\{2\}\(R\_\{0\}\)^\{R\_\{0\}/R\},\\qquad 1\-\\beta\_\{2\}\(R\)=O\(R^\{\-1\}\)\.\(20\)Ifm,v,g,σm,v,g,\\sigmaare held fixed andβ2\(R\)\\beta\_\{2\}\(R\)follows \([20](https://arxiv.org/html/2610.19064#S4.E20)\), withβ2\(R\)\\beta\_\{2\}\(R\)bounded away from zero, the right\-hand side remains bounded rather than growing linearly inRR\. For eachRR, both routes in this comparison use the*same*β2\(R\)\\beta\_\{2\}\(R\)\.Scaling second\-moment memory therefore controls a bound on the Adam mean mismatch asRRgrows\.Summing the squared coordinate bounds controls‖𝒃R‖22\\\|\\bm\{b\}\_\{R\}\\\|\_\{2\}^\{2\}, and hence the mean term\(𝒒⊤𝒃R\)2\(\\bm\{q\}^\{\\top\}\\bm\{b\}\_\{R\}\)^\{2\}for every unit𝒒\\bm\{q\}in \([10](https://arxiv.org/html/2610.19064#S4.E10)\)\. The same conclusion holds after SOAP’s orthogonal rotation back to parameter coordinates\. This conditional one\-step result motivates the matched\-β2\\beta\_\{2\}experiments in Section[5\.2](https://arxiv.org/html/2610.19064#S5.SS2)\.
##### 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\\beta\_\{2\}, smaller local batches produce noisier gradients, which can increase fluctuations in the estimated second momentv^\\widehat\{v\}\. For SOAP, these fluctuations inv^\\widehat\{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\.
## 5Empirical Evaluation
We first compare replica scaling across inner optimizers \(Sec\.[5\.1](https://arxiv.org/html/2610.19064#S5.SS1)\), then isolate the effect of second\-moment memory on performance \(Sec\.[5\.2](https://arxiv.org/html/2610.19064#S5.SS2)\), and finally examine the update distributions behind these results \(Sec\.[5\.3](https://arxiv.org/html/2610.19064#S5.SS3)\)\. The analysis connects the observed scaling to the two terms in Equation \([10](https://arxiv.org/html/2610.19064#S4.E10)\): mean mismatch and stochastic aggregation error\.
##### Experimental setup\.
Our main experiments train 178M\- and 512M\-parameter LLaMA models on DCLM\([Li et al\., 2024](https://arxiv.org/html/2610.19064#bib.bib15)\)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 takeH=30H=30local steps between aggregations and all local optimization methods use a Nesterov SGD outer optimizer\. Appendix[B](https://arxiv.org/html/2610.19064#A2)records the architecture, training budgets, and result\-selection protocol\.
### 5\.1Comparing 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](https://arxiv.org/html/2610.19064#S5.F2)and[3](https://arxiv.org/html/2610.19064#S5.F3)show both absolute validation loss and the percentage change from each optimizer’s own DP reference\([Thérien et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib27)\)\. 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](https://arxiv.org/html/2610.19064#S5.F2)\), we observe, as in[Thérien et al\. \(2025\)](https://arxiv.org/html/2610.19064#bib.bib27), 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 largeRR\. AtR=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](https://arxiv.org/html/2610.19064#S5.SS3)examines how this ordering relates to the covariance of the returned updates\. Interestingly, at 178M,R=8R=8SPlus actually improves performance over its DP baseline\. This is likely because atR\>1R\>1it can mimic a majority vote\-based technique\([Bernstein et al\., 2018](https://arxiv.org/html/2610.19064#bib.bib2)\)\.
Figure 2:178M replica scaling\.R=1R=1reports DP performance, whileR\>1R\>1reports the performance of local optimization with synchronization intervalH=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](https://arxiv.org/html/2610.19064#S5.F3), we observe a similar pattern\. SOAP is, once again, the strongest DP reference and has the lower loss throughR=32R=32\. SPlus overtakes it atR=64R=64andR=96R=96\. AtR=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=1reports DP performance, whileR\>1R\>1reports the performance of local optimization with synchronization intervalH=30H=30\. Left: final validation loss\. Right: change in loss from the same optimizer’s own DP \(R=1\) baseline\. As in Fig\.[2](https://arxiv.org/html/2610.19064#S5.F2), 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 usesR=1R=1; other runs useH=30H=30\. Best in each column is bold\.
Table 2:512M final validation loss\.DP usesR=1R=1; other runs useH=30H=30\. Best in each column is bold\.
##### 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](https://arxiv.org/html/2610.19064#bib.bib27);[Charles et al\., 2025](https://arxiv.org/html/2610.19064#bib.bib3)\)\. In Figure[1](https://arxiv.org/html/2610.19064#S1), 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 atH=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](https://arxiv.org/html/2610.19064#S4)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](https://arxiv.org/html/2610.19064#S5.T3)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](https://arxiv.org/html/2610.19064#S4.SS4)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](https://arxiv.org/html/2610.19064#bib.bib26);[Karimireddy et al\., 2020](https://arxiv.org/html/2610.19064#bib.bib11)\)\.
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 for6,7756,775steps\. The DP,R=8R=8, andR=64R=64methods are directly tuned, while theR=16,32R=16,32configurations use interpolated hyperparameters rather than separate full sweeps\. The best\-performing method in each column is bolded\.Interestingly, Muon is less robust here than atH=30H=30and is worse than AdamW at the larger replica counts\. Its strongH=30H=30results 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\.2Second\-moment memory and replica scaling
Following Sec\.[4\.4](https://arxiv.org/html/2610.19064#S4.SS4), we next study a practical consequence of smaller local batches for Adam\-based methods\. Equation \([A\.6](https://arxiv.org/html/2610.19064#A1.E6)\) shows how local gradient variance enters the second\-moment discrepancy\. The token\-half\-life rule in Equation \([20](https://arxiv.org/html/2610.19064#S4.E20)\) increasesβ2\\beta\_\{2\}asRRgrows, 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](https://arxiv.org/html/2610.19064#bib.bib27);[Horoi et al\., 2026](https://arxiv.org/html/2610.19064#bib.bib9)\)\. This rule is already incorporated in the main results in Sec\.[5\.1](https://arxiv.org/html/2610.19064#S5.SS1)\. 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\\beta\_\{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\\beta\_\{2\}when scaling R[Charles et al\. \(2025\)](https://arxiv.org/html/2610.19064#bib.bib3);[Horoi et al\. \(2026\)](https://arxiv.org/html/2610.19064#bib.bib9);[Sarfi et al\. \(2025\)](https://arxiv.org/html/2610.19064#bib.bib25)
To empirically measure the importance of second moment estimation highlighted in Sec\.[4\.4](https://arxiv.org/html/2610.19064#S4.SS4), we compared the performance of optimizers using a fixedβ2\\beta\_\{2\}and tuning it at every replica scale\. Figure[5](https://arxiv.org/html/2610.19064#S5.F5)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\\beta\_\{2\}as replicas increase, but IsoLoCo and DiLoCo benefit more\. Underβ2\\beta\_\{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\\beta\_\{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\\beta\_\{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\\beta\_\{2\}is well tuned\.Figure 5:Tuningβ2\\beta\_\{2\}for larger model and replica\.We ablate using fixedβ2\\beta\_\{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](https://arxiv.org/html/2610.19064#S5.F5)\. Both methods use the publishedR=8R=8hyperparameters without retuning\. At 512M andR=64R=64, changing onlyβ2\\beta\_\{2\}reduces DiLoCo loss from2\.8842\.884to2\.8492\.849and IsoLoCo loss from2\.7342\.734to2\.7022\.702, approximately1\.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 maintainLLandRR, however, scalable recipes for SOAP have been studied\([Khona et al\., 2026](https://arxiv.org/html/2610.19064#bib.bib13)\), 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\\beta\_\{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\.3Analysis: 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=1R𝑪r,h,ah=\(dtr\(𝑪~h2\)tr\(𝑪~h\)2−1\)1/2\.\\widetilde\{\\bm\{C\}\}\_\{h\}:=\\frac\{1\}\{R^\{2\}\}\\sum\_\{r=1\}^\{R\}\\bm\{C\}\_\{r,h\},\\qquad a\_\{h\}=\\left\(\\frac\{d\\,\\operatorname\{tr\}\(\\widetilde\{\\bm\{C\}\}\_\{h\}^\{2\}\)\}\{\\operatorname\{tr\}\(\\widetilde\{\\bm\{C\}\}\_\{h\}\)^\{2\}\}\-1\\right\)^\{1/2\}\.\(21\)Hereddis the probed dimension;ah=0a\_\{h\}=0means isotropy, and larger values indicate greater anisotropy, independently of overall noise magnitude\. The constructed𝑪~h\\widetilde\{\\bm\{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](https://arxiv.org/html/2610.19064#A3.SS1)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](https://arxiv.org/html/2610.19064#S5.T5)below\.
Table 4:Update\-noise anisotropy atR=64R=64, round 75, using all 64 replicas and local batch four\. Lower is more isotropic\.
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\.
We observe thatcovariance 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 from3535for AdamW to0\.430\.43for 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\)](https://arxiv.org/html/2610.19064#bib.bib27)\. 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](https://arxiv.org/html/2610.19064#S4.E10)\) 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 runH=30H=30steps 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\\alpha\_\{R\}and a leading noise direction𝒒^\\widehat\{\\bm\{q\}\}; 24 additional repetitions estimate the directional mean mismatch\(𝒒^⊤𝒃R\)2\(\\widehat\{\\bm\{q\}\}^\{\\top\}\\bm\{b\}\_\{R\}\)^\{2\}and stochastic errorαR2𝒒^⊤𝑪agg𝒒^\\alpha\_\{R\}^\{2\}\\widehat\{\\bm\{q\}\}^\{\\top\}\\bm\{C\}\_\{\\rm agg\}\\widehat\{\\bm\{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](https://arxiv.org/html/2610.19064#A3.SS2)defines the estimators, including the paired finite\-sample correction\.
We observe*the mean\-mismatch term is only7\.57\.5–19%19\\%of the stochastic term on average along these directions\.*This supports controlling the stochastic term in Equation \([10](https://arxiv.org/html/2610.19064#S4.E10)\) 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](https://arxiv.org/html/2610.19064#S4.E11)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𝒒^\\widehat\{\\bm\{q\}\}that maximizes sample variance over 64 paired one\-step update differences\. With𝓗\\bm\{\\mathcal\{H\}\}the loss Hessian on 64 held\-out sequences, Table[6](https://arxiv.org/html/2610.19064#S5.T6)reports
κ\(𝒒^\)=𝒒^⊤𝓗𝒒^116∑j=116𝒗j⊤𝓗𝒗j,\\kappa\(\\widehat\{\\bm\{q\}\}\)=\\frac\{\\widehat\{\\bm\{q\}\}^\{\\top\}\\bm\{\\mathcal\{H\}\}\\widehat\{\\bm\{q\}\}\}\{\\frac\{1\}\{16\}\\sum\_\{j=1\}^\{16\}\\bm\{v\}\_\{j\}^\{\\top\}\\bm\{\\mathcal\{H\}\}\\bm\{v\}\_\{j\}\},where the𝒗j\\bm\{v\}\_\{j\}are independent random unit vectors in the same six\-matrix subspace\. A ratio of100100means100100times 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 batchb=256/Rb=256/R\.Appendix[C\.3](https://arxiv.org/html/2610.19064#A3.SS3)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\.
## 6Conclusion
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](https://doi.org/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*, pp\. 560–569, 2018\.URL[https://proceedings\.mlr\.press/v80/bernstein18a\.html](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](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](https://doi.org/10.1214/22-AOS2239)\.URL[https://arxiv\.org/abs/2104\.01280](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](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](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](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*, pp\. 1842–1850, 2018\.URL[https://proceedings\.mlr\.press/v80/gupta18a\.html](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](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/](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*, pp\. 5132–5143, 2020\.URL[https://proceedings\.mlr\.press/v119/karimireddy20a\.html](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](https://doi.org/10.1080/00031305.2016.1277159)\.URL[https://arxiv\.org/abs/1512\.00809](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](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](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](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](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](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](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](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](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](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*, pp\. 1273–1282, 2017\.URL[https://proceedings\.mlr\.press/v54/mcmahan17a\.html](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](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](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](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](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](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](https://doi.org/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](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](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](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*, pp\. 7184–7193, 2019\.URL[https://proceedings\.mlr\.press/v97/yu19d\.html](https://proceedings.mlr.press/v97/yu19d.html)\.
Appendix[A](https://arxiv.org/html/2610.19064#A1)collects the supporting theoretical results and proofs\. Appendix[B](https://arxiv.org/html/2610.19064#A2)records the training setup, and Appendix[C](https://arxiv.org/html/2610.19064#A3)gives the analysis protocols and estimators\.
## Appendix ASupporting theoretical results
### A\.1Muon covariance under rotational proposals
Letp=min\(m,n\)p=\\min\(m,n\)anddmax=max\(m,n\)d\_\{\\max\}=\\max\(m,n\)\. If a full\-rank proposal has SVD𝑴=𝑼~diag\(σi\)𝑽~⊤\\bm\{M\}=\\widetilde\{\\bm\{U\}\}\\operatorname\{diag\}\(\\sigma\_\{i\}\)\\widetilde\{\\bm\{V\}\}^\{\\top\}, the ideal Muon update with prescribed matrix\-size scaleγ\>0\\gamma\>0is
𝑼Muon=γpolar\(𝑴\)=γ𝑼~𝑽~⊤\.\\bm\{U\}\_\{\\rm Muon\}=\\gamma\\,\\operatorname\{polar\}\(\\bm\{M\}\)=\\gamma\\widetilde\{\\bm\{U\}\}\\widetilde\{\\bm\{V\}\}^\{\\top\}\.\(A\.1\)The polar map discards the realized singular values and keeps the polar factor\. In particular,polar\(a𝑴\)=polar\(𝑴\)\\operatorname\{polar\}\(a\\bm\{M\}\)=\\operatorname\{polar\}\(\\bm\{M\}\)for everya\>0a\>0, so a positive multiplicative change to the entire proposal leaves the update unchanged\.
###### Proposition 4\(Muon covariance under rotational proposals\)\.
Suppose𝐌r\\bm\{M\}\_\{r\}is full rank almost surely\. Ifm≥nm\\geq n, suppose𝐎𝐌r=𝑑𝐌r\\bm\{O\}\\bm\{M\}\_\{r\}\\overset\{d\}\{=\}\\bm\{M\}\_\{r\}for every𝐎∈O\(m\)\\bm\{O\}\\in\\mathrm\{O\}\(m\); ifm<nm<n, suppose𝐌r𝐎=𝑑𝐌r\\bm\{M\}\_\{r\}\\bm\{O\}\\overset\{d\}\{=\}\\bm\{M\}\_\{r\}for every𝐎∈O\(n\)\\bm\{O\}\\in\\mathrm\{O\}\(n\)\. Either one\-sided condition suffices whenm=nm=n\. Then
𝔼\[𝑼Muon,r\]=𝟎,Cov\(vec\(𝑼Muon,r\)\)=γ2dmax𝑰mn\.\\mathbb\{E\}\[\\bm\{U\}\_\{\{\\rm Muon\},r\}\]=\\bm\{0\},\\qquad\\operatorname\{Cov\}\(\\operatorname\{vec\}\(\\bm\{U\}\_\{\{\\rm Muon\},r\}\)\)=\\frac\{\\gamma^\{2\}\}\{d\_\{\\max\}\}\\bm\{I\}\_\{mn\}\.\(A\.2\)Consequently, independent replicas have a uniformly averaged update covarianceγ2𝐈mn/\(dmaxR\)\\gamma^\{2\}\\bm\{I\}\_\{mn\}/\(d\_\{\\max\}R\), even when their proposal singular\-value distributions differ\.
###### Proof\.
Considerm≥nm\\geq nand write the columns of𝑸r=polar\(𝑴r\)\\bm\{Q\}\_\{r\}=\\operatorname\{polar\}\(\\bm\{M\}\_\{r\}\)as𝒒r,1,…,𝒒r,n\\bm\{q\}\_\{r,1\},\\ldots,\\bm\{q\}\_\{r,n\}\. Equivariance of the polar map and left invariance of the proposal law imply𝑶𝑸r=𝑑𝑸r\\bm\{O\}\\bm\{Q\}\_\{r\}\\overset\{d\}\{=\}\\bm\{Q\}\_\{r\}for every𝑶∈O\(m\)\\bm\{O\}\\in\\mathrm\{O\}\(m\)\. Taking𝑶=−𝑰\\bm\{O\}=\-\\bm\{I\}gives𝔼\[𝑸r\]=𝟎\\mathbb\{E\}\[\\bm\{Q\}\_\{r\}\]=\\bm\{0\}\. Moreover,𝔼\[𝒒r,j𝒒r,k⊤\]\\mathbb\{E\}\[\\bm\{q\}\_\{r,j\}\\bm\{q\}\_\{r,k\}^\{\\top\}\]commutes with every orthogonal matrix and must therefore equalajk𝑰ma\_\{jk\}\\bm\{I\}\_\{m\}\. Since𝑸r⊤𝑸r=𝑰n\\bm\{Q\}\_\{r\}^\{\\top\}\\bm\{Q\}\_\{r\}=\\bm\{I\}\_\{n\},
majk=𝔼\[𝒒r,k⊤𝒒r,j\]=𝟏\{j=k\}\.ma\_\{jk\}=\\mathbb\{E\}\[\\bm\{q\}\_\{r,k\}^\{\\top\}\\bm\{q\}\_\{r,j\}\]=\\mathbf\{1\}\_\{\\\{j=k\\\}\}\.\(A\.3\)The column blocks of𝔼\[vec\(𝑸r\)vec\(𝑸r\)⊤\]\\mathbb\{E\}\[\\operatorname\{vec\}\(\\bm\{Q\}\_\{r\}\)\\operatorname\{vec\}\(\\bm\{Q\}\_\{r\}\)^\{\\top\}\]are consequently𝟏\{j=k\}𝑰m/m\\mathbf\{1\}\_\{\\\{j=k\\\}\}\\bm\{I\}\_\{m\}/m, which proves Equation \([A\.2](https://arxiv.org/html/2610.19064#A1.E2)\) becausedmax=md\_\{\\max\}=m\. The proof form<nm<napplies the same argument to the rows, using right invariance and𝑸r𝑸r⊤=𝑰m\\bm\{Q\}\_\{r\}\\bm\{Q\}\_\{r\}^\{\\top\}=\\bm\{I\}\_\{m\}\. ∎
### A\.2Proof of Proposition[2](https://arxiv.org/html/2610.19064#Thmproposition2)\(SPlus to DDP comparison\)
###### Proof of Proposition[2](https://arxiv.org/html/2610.19064#Thmproposition2)\.
Writeψ\(a\)=2Φ\(a\)−1\\psi\(a\)=2\\Phi\(a\)\-1andc=2/πc=\\sqrt\{2/\\pi\}\. For one coordinate, setXr=σ\(a\+RZr\)X\_\{r\}=\\sigma\(a\+\\sqrt\{R\}Z\_\{r\}\)andX¯=R−1∑rXr\\bar\{X\}=R^\{\-1\}\\sum\_\{r\}X\_\{r\}\. Then
Pr\(Xr\>0\)=Φ\(a/R\),𝔼\[sign\(Xr\)\]=ψ\(a/R\),𝔼\[sign\(X¯\)\]=ψ\(a\)\.\\Pr\(X\_\{r\}\>0\)=\\Phi\(a/\\sqrt\{R\}\),\\quad\\mathbb\{E\}\[\\sign\(X\_\{r\}\)\]=\\psi\(a/\\sqrt\{R\}\),\\quad\\mathbb\{E\}\[\\sign\(\\bar\{X\}\)\]=\\psi\(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\. Fora≥0a\\geq 0,
Rψ\(a/R\)−ψ\(a\)=c∫0a\(e−u2/\(2R\)−e−u2/2\)du\.\\sqrt\{R\}\\psi\(a/\\sqrt\{R\}\)\-\\psi\(a\)=c\\int\_\{0\}^\{a\}\\left\(e^\{\-u^\{2\}/\(2R\)\}\-e^\{\-u^\{2\}/2\}\\right\)du\.\(A\.4\)The integrand is nonnegative and at mostu2\(1−R−1\)/2u^\{2\}\(1\-R^\{\-1\}\)/2, by the unit Lipschitz bound fore−xe^\{\-x\}onx≥0x\\geq 0\. Integration and odd symmetry give
\|Rψ\(a/R\)−ψ\(a\)\|≤c6\(1−R−1\)\|a\|3\.\\left\|\\sqrt\{R\}\\psi\(a/\\sqrt\{R\}\)\-\\psi\(a\)\\right\|\\leq\\frac\{c\}\{6\}\(1\-R^\{\-1\}\)\|a\|^\{3\}\.Squaring and summing over coordinates is valid because𝑸\\bm\{Q\}is orthogonal\. Each sign has variance1−ψ\(aj/R\)21\-\\psi\(a\_\{j\}/\\sqrt\{R\}\)^\{2\}\. Independence of projected coordinates and of replicas gives the covariance in \([17](https://arxiv.org/html/2610.19064#S4.E17)\)\. Finally use Proposition[1](https://arxiv.org/html/2610.19064#Thmproposition1)\. ∎
The DDP covariance is𝑸Diag\(1−ψ\(aj\)2\)𝑸⊤\\bm\{Q\}\\operatorname\{Diag\}\(1\-\\psi\(a\_\{j\}\)^\{2\}\)\\bm\{Q\}^\{\\top\}\. ForR≥1R\\geq 1the calibrated local covariance minus this matrix is positive semidefinite, with operator norm at mostc2maxjaj2c^\{2\}\\max\_\{j\}a\_\{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\.3Adam/SOAP second moments and the mean comparison
Condition on common incoming optimizer states and parameters, as in Section[4\.4](https://arxiv.org/html/2610.19064#S4.SS4)\. Letttdenote 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−𝒈¯\)⊤\\bm\{S\}\_\{R\}=R^\{\-1\}\\sum\_\{r\}\(\\widehat\{\\bm\{g\}\}\_\{r\}\-\\bar\{\\bm\{g\}\}\)\(\\widehat\{\\bm\{g\}\}\_\{r\}\-\\bar\{\\bm\{g\}\}\)^\{\\top\}\. For the second moments in \([2](https://arxiv.org/html/2610.19064#S4.E2)\) and \([5](https://arxiv.org/html/2610.19064#S4.E5)\),
1R∑r𝒗A,r,t−𝒗A,DDP,t=\(1−β2\)diag\(𝑺R\),1R∑r𝒗S,r,t−𝒗S,DDP,t=\(1−β2\)diag\(𝑸⊤𝑺R𝑸\)\.\\begin\{split\}\\frac\{1\}\{R\}\\sum\_\{r\}\\bm\{v\}\_\{A,r,t\}\-\\bm\{v\}\_\{A,\\rm DDP,t\}&=\(1\-\\beta\_\{2\}\)\\operatorname\{diag\}\(\\bm\{S\}\_\{R\}\),\\\\ \\frac\{1\}\{R\}\\sum\_\{r\}\\bm\{v\}\_\{S,r,t\}\-\\bm\{v\}\_\{S,\\rm DDP,t\}&=\(1\-\\beta\_\{2\}\)\\operatorname\{diag\}\(\\bm\{Q\}^\{\\top\}\\bm\{S\}\_\{R\}\\bm\{Q\}\)\.\\end\{split\}\(A\.5\)The second identity assumes a common basis fixed before the fresh gradients;𝒗A,𝒗S\\bm\{v\}\_\{A\},\\bm\{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\)=R𝚺g\\operatorname\{Cov\}\(\\widehat\{\\bm\{g\}\}\_\{r\}\)=R\\bm\{\\Sigma\}\_\{g\}and the replica gradients are independent, so𝚺g\\bm\{\\Sigma\}\_\{g\}is the covariance of the DDP\-averaged gradient\. At fixedBB, taking expectations in \([A\.5](https://arxiv.org/html/2610.19064#A1.E5)\) gives
𝔼\[1R∑r𝒗A,r,t−𝒗A,DDP,t\]=\(1−β2\)\(R−1\)diag\(𝚺g\)\.\\mathbb\{E\}\\left\[\\frac\{1\}\{R\}\\sum\_\{r\}\\bm\{v\}\_\{A,r,t\}\-\\bm\{v\}\_\{A,\\rm DDP,t\}\\right\]=\(1\-\\beta\_\{2\}\)\(R\-1\)\\operatorname\{diag\}\(\\bm\{\\Sigma\}\_\{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\)\.\\frac\{1\}\{R\}\\sum\_\{r\}\\widehat\{\\bm\{g\}\}\_\{r\}^\{\\odot 2\}\-\\bar\{\\bm\{g\}\}^\{\\odot 2\}=\\frac\{1\}\{R\}\\sum\_\{r\}\(\\widehat\{\\bm\{g\}\}\_\{r\}\-\\bar\{\\bm\{g\}\}\)^\{\\odot 2\}=\\operatorname\{diag\}\(\\bm\{S\}\_\{R\}\)\.Multiplying by1−β21\-\\beta\_\{2\}proves the Adam identity in \([A\.5](https://arxiv.org/html/2610.19064#A1.E5)\); applying it to𝑸⊤𝒈^r\\bm\{Q\}^\{\\top\}\\widehat\{\\bm\{g\}\}\_\{r\}proves the SOAP identity when𝑸\\bm\{Q\}is fixed and common\. Under \([14](https://arxiv.org/html/2610.19064#S4.E14)\),𝔼\[𝑺R\]=\(R−1\)𝚺g\\mathbb\{E\}\[\\bm\{S\}\_\{R\}\]=\(R\-1\)\\bm\{\\Sigma\}\_\{g\}, which proves \([A\.6](https://arxiv.org/html/2610.19064#A1.E6)\)\.
For a scalar smooth mapffandXt=g\+tσZX\_\{t\}=g\+\\sqrt\{t\}\\,\\sigma Z, Gaussian variance differentiation\([Voigtlaender, 2021](https://arxiv.org/html/2610.19064#bib.bib28)\)gives the exact identity
ddt𝔼\[f\(Xt\)\]=σ22𝔼\[f′′\(Xt\)\],𝔼\[f\(XR\)\]−𝔼\[f\(X1\)\]=σ22∫1R𝔼\[f′′\(Xt\)\]𝑑t\.\\frac\{d\}\{dt\}\\mathbb\{E\}\[f\(X\_\{t\}\)\]=\\frac\{\\sigma^\{2\}\}\{2\}\\mathbb\{E\}\[f^\{\\prime\\prime\}\(X\_\{t\}\)\],\\qquad\\mathbb\{E\}\[f\(X\_\{R\}\)\]\-\\mathbb\{E\}\[f\(X\_\{1\}\)\]=\\frac\{\\sigma^\{2\}\}\{2\}\\int\_\{1\}^\{R\}\\mathbb\{E\}\[f^\{\\prime\\prime\}\(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\.1Finite\-noise mean bound and second\-moment memory
###### Proof of Proposition[3](https://arxiv.org/html/2610.19064#Thmproposition3)\.
For this proof only, abbreviatep=β1mp=\\beta\_\{1\}m,k=1−β1k=1\-\\beta\_\{1\},a0=β2v\>0a\_\{0\}=\\beta\_\{2\}v\>0, andb0=1−β2\>0b\_\{0\}=1\-\\beta\_\{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\\,\\ell\(x\),\\qquad h\(x\)=\(a\_\{0\}\+b\_\{0\}x^\{2\}\)^\{\-1/2\},\\quad\\ell\(x\)=x\(a\_\{0\}\+b\_\{0\}x^\{2\}\)^\{\-1/2\}\.Direct differentiation gives
h′′\(x\)=b0\(2b0x2−a0\)\(a0\+b0x2\)5/2,ℓ′′′\(x\)=−3a0b0\(a0−4b0x2\)\(a0\+b0x2\)7/2\.h^\{\\prime\\prime\}\(x\)=\\frac\{b\_\{0\}\(2b\_\{0\}x^\{2\}\-a\_\{0\}\)\}\{\(a\_\{0\}\+b\_\{0\}x^\{2\}\)^\{5/2\}\},\\qquad\\ell^\{\\prime\\prime\\prime\}\(x\)=\\frac\{\-3a\_\{0\}b\_\{0\}\(a\_\{0\}\-4b\_\{0\}x^\{2\}\)\}\{\(a\_\{0\}\+b\_\{0\}x^\{2\}\)^\{7/2\}\}\.\(A\.8\)Fory≥0y\\geq 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 ofy=1/2y=1/2, and the second on either side ofy=1/4y=1/4; their interior positive extrema have value less than one\. Substitutingy=b0x2/a0y=b\_\{0\}x^\{2\}/a\_\{0\}gives
‖h′′‖∞≤b0/a03/2,‖ℓ′′′‖∞≤3b0/a03/2\.\\\|h^\{\\prime\\prime\}\\\|\_\{\\infty\}\\leq b\_\{0\}/a\_\{0\}^\{3/2\},\\qquad\\\|\\ell^\{\\prime\\prime\\prime\}\\\|\_\{\\infty\}\\leq 3b\_\{0\}/a\_\{0\}^\{3/2\}\.Sinceℓ′′\\ell^\{\\prime\\prime\}is odd andZZis symmetric,𝔼\[ℓ′′\(tσZ\)\]=0\\mathbb\{E\}\[\\ell^\{\\prime\\prime\}\(\\sqrt\{t\}\\sigma Z\)\]=0\. The mean\-value bound therefore implies
\|𝔼\[ℓ′′\(g\+tσZ\)\]\|≤3b0\|g\|/a03/2\.\\big\|\\mathbb\{E\}\[\\ell^\{\\prime\\prime\}\(g\+\\sqrt\{t\}\\sigma Z\)\]\\big\|\\leq 3b\_\{0\}\|g\|/a\_\{0\}^\{3/2\}\.Thus, uniformly int\>0t\>0,
\|𝔼\[f′′\(g\+tσZ\)\]\|≤b0a03/2\(\|p\|\+3k\|g\|\)\.\|\\mathbb\{E\}\[f^\{\\prime\\prime\}\(g\+\\sqrt\{t\}\\sigma Z\)\]\|\\leq\\frac\{b\_\{0\}\}\{a\_\{0\}^\{3/2\}\}\\bigl\(\|p\|\+3k\|g\|\\bigr\)\.Integrating \([A\.7](https://arxiv.org/html/2610.19064#A1.E7)\) over1≤t≤R1\\leq t\\leq Rgives \([19](https://arxiv.org/html/2610.19064#S4.E19)\)\. 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\\beta\_\{2\}per local step\. Keeping its decay per fixed number of local examples unchanged giveslogβ2\(R\)/\(B/R\)=logβ2\(R0\)/\(B/R0\)\\log\\beta\_\{2\}\(R\)/\(B/R\)=\\log\\beta\_\{2\}\(R\_\{0\}\)/\(B/R\_\{0\}\), proving \([20](https://arxiv.org/html/2610.19064#S4.E20)\)\. Writec0=−R0logβ2\(R0\)\>0c\_\{0\}=\-R\_\{0\}\\log\\beta\_\{2\}\(R\_\{0\}\)\>0\. Then1−β2\(R\)=1−e−c0/R≤c0/R1\-\\beta\_\{2\}\(R\)=1\-e^\{\-c\_\{0\}/R\}\\leq c\_\{0\}/R, so\(R−1\)\(1−β2\(R\)\)≤c0\(R\-1\)\(1\-\\beta\_\{2\}\(R\)\)\\leq 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\\beta\_\{2\}\.
### A\.4Basis 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](https://arxiv.org/html/2610.19064#S4.E10)\)\.
##### 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\.1Scale errors and correlations left in the estimated basis
Let workerrrhave a centered matrix proposal𝑴r∈ℝm×n\\bm\{M\}\_\{r\}\\in\\mathbb\{R\}^\{m\\times n\}, and set𝒙r=vec\(𝑴r\)\\bm\{x\}\_\{r\}=\\operatorname\{vec\}\(\\bm\{M\}\_\{r\}\)and𝚺r=Cov\(𝒙r\)≻𝟎\\bm\{\\Sigma\}\_\{r\}=\\operatorname\{Cov\}\(\\bm\{x\}\_\{r\}\)\\succ\\bm\{0\}\. Let𝑸^L,r\\widehat\{\\bm\{Q\}\}\_\{L,r\}and𝑸^R,r\\widehat\{\\bm\{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\.\\bm\{Z\}\_\{r\}=\\widehat\{\\bm\{Q\}\}\_\{L,r\}^\{\\top\}\\bm\{M\}\_\{r\}\\widehat\{\\bm\{Q\}\}\_\{R,r\},\\qquad\\bm\{Q\}\_\{r\}=\\widehat\{\\bm\{Q\}\}\_\{R,r\}\\otimes\\widehat\{\\bm\{Q\}\}\_\{L,r\},\\qquad\\operatorname\{vec\}\(\\bm\{Z\}\_\{r\}\)=\\bm\{Q\}\_\{r\}^\{\\top\}\\bm\{x\}\_\{r\}\.\(A\.9\)No Kronecker assumption is imposed on𝚺r\\bm\{\\Sigma\}\_\{r\}here: the Kronecker matrix𝑸r\\bm\{Q\}\_\{r\}represents the two\-sided rotation, not a factorization of the full covariance\.
To distinguish scale error from an imperfect basis, let𝚪r\\bm\{\\Gamma\}\_\{r\}be the true covariance in the estimated basis,𝑫r\\bm\{D\}\_\{r\}its diagonal of coordinate variances, and𝓒r\\bm\{\\mathcal\{C\}\}\_\{r\}the correlation matrix remaining after exact variance normalization:
𝚪r=𝑸r⊤𝚺r𝑸r,𝑫r=Diag\(𝒗r\),vrj=\(𝚪r\)jj,𝓒r=𝑫r−1/2𝚪r𝑫r−1/2\.\\bm\{\\Gamma\}\_\{r\}=\\bm\{Q\}\_\{r\}^\{\\top\}\\bm\{\\Sigma\}\_\{r\}\\bm\{Q\}\_\{r\},\\qquad\\bm\{D\}\_\{r\}=\\operatorname\{Diag\}\(\\bm\{v\}\_\{r\}\),\\quad v\_\{rj\}=\(\\bm\{\\Gamma\}\_\{r\}\)\_\{jj\},\\qquad\\bm\{\\mathcal\{C\}\}\_\{r\}=\\bm\{D\}\_\{r\}^\{\-1/2\}\\bm\{\\Gamma\}\_\{r\}\\bm\{D\}\_\{r\}^\{\-1/2\}\.\(A\.10\)Thus𝓒r\\bm\{\\mathcal\{C\}\}\_\{r\}has unit diagonal; its off\-diagonal entries measure correlations that the basis has not removed\. The actual normalization uses frozen variance estimatesv^rj\\widehat\{v\}\_\{rj\}\. Write their relative errors aserje\_\{rj\}and collect the resulting inverse\-scale gains in𝑨r\\bm\{A\}\_\{r\}:
v^rj=vrj\(1\+erj\)\>0,𝑫^r=Diag\(𝒗^r\),𝑨r=Diag\(\(1\+erj\)−1/2\)\.\\widehat\{v\}\_\{rj\}=v\_\{rj\}\(1\+e\_\{rj\}\)\>0,\\qquad\\widehat\{\\bm\{D\}\}\_\{r\}=\\operatorname\{Diag\}\(\\widehat\{\\bm\{v\}\}\_\{r\}\),\\qquad\\bm\{A\}\_\{r\}=\\operatorname\{Diag\}\\big\(\(1\+e\_\{rj\}\)^\{\-1/2\}\\big\)\.\(A\.11\)If𝑽^r\\widehat\{\\bm\{V\}\}\_\{r\}is them×nm\\times nreshape of𝒗^r\\widehat\{\\bm\{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\\displaystyle\\bm\{U\}\_\{r\}^\{\\rm SOAP\}=𝑸^L,r\(𝒁r⊘𝑽^r\)𝑸^R,r⊤,\\displaystyle=\\widehat\{\\bm\{Q\}\}\_\{L,r\}\\left\(\\bm\{Z\}\_\{r\}\\oslash\\sqrt\{\\widehat\{\\bm\{V\}\}\_\{r\}\}\\right\)\\widehat\{\\bm\{Q\}\}\_\{R,r\}^\{\\top\},vec\(𝑼rSOAP\)\\displaystyle\\operatorname\{vec\}\(\\bm\{U\}\_\{r\}^\{\\rm SOAP\}\)=𝑸r𝑫^r−1/2𝑸r⊤𝒙r,\\displaystyle=\\bm\{Q\}\_\{r\}\\widehat\{\\bm\{D\}\}\_\{r\}^\{\-1/2\}\\bm\{Q\}\_\{r\}^\{\\top\}\\bm\{x\}\_\{r\},\(A\.12\)𝑼rSPlus\\displaystyle\\bm\{U\}\_\{r\}^\{\\rm SPlus\}=𝑸^L,rsign\(𝒁r\)𝑸^R,r⊤,\\displaystyle=\\widehat\{\\bm\{Q\}\}\_\{L,r\}\\sign\(\\bm\{Z\}\_\{r\}\)\\widehat\{\\bm\{Q\}\}\_\{R,r\}^\{\\top\},vec\(𝑼rSPlus\)\\displaystyle\\operatorname\{vec\}\(\\bm\{U\}\_\{r\}^\{\\rm SPlus\}\)=𝑸rsign\(𝑸r⊤𝒙r\)\.\\displaystyle=\\bm\{Q\}\_\{r\}\\sign\(\\bm\{Q\}\_\{r\}^\{\\top\}\\bm\{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\)\)\\bm\{C\}\_\{r\}^\{\\rm SOAP\}=\\operatorname\{Cov\}\(\\operatorname\{vec\}\(\\bm\{U\}\_\{r\}^\{\\rm SOAP\}\)\)\. For every centered proposal with finite covariance,
𝑪rSOAP=𝑸r𝑨r𝓒r𝑨r𝑸r⊤,\\bm\{C\}\_\{r\}^\{\\rm SOAP\}=\\bm\{Q\}\_\{r\}\\bm\{A\}\_\{r\}\\bm\{\\mathcal\{C\}\}\_\{r\}\\bm\{A\}\_\{r\}\\bm\{Q\}\_\{r\}^\{\\top\},\(A\.14\)and the whitening error decomposes exactly as
‖𝑪rSOAP−𝑰‖F2=∑j\(erj1\+erj\)2⏟coordinate\-scale error\+∑i≠j𝒞r,ij2\(1\+eri\)\(1\+erj\)⏟residual correlations after normalization\.\\\|\\bm\{C\}\_\{r\}^\{\\rm SOAP\}\-\\bm\{I\}\\\|\_\{F\}^\{2\}=\\underbrace\{\\sum\_\{j\}\\left\(\\frac\{e\_\{rj\}\}\{1\+e\_\{rj\}\}\\right\)^\{2\}\}\_\{\\text\{coordinate\-scale error\}\}\+\\underbrace\{\\sum\_\{i\\neq j\}\\frac\{\\mathcal\{C\}\_\{r,ij\}^\{2\}\}\{\(1\+e\_\{ri\}\)\(1\+e\_\{rj\}\)\}\}\_\{\\text\{residual correlations after normalization\}\}\.\(A\.15\)If additionally𝐱r\\bm\{x\}\_\{r\}is centered Gaussian, then
𝑪rSPlus\\displaystyle\\bm\{C\}\_\{r\}^\{\\rm SPlus\}=𝑸r\[2πarcsin\(𝓒r\)\]𝑸r⊤,\\displaystyle=\\bm\{Q\}\_\{r\}\\left\[\\frac\{2\}\{\\pi\}\\arcsin\(\\bm\{\\mathcal\{C\}\}\_\{r\}\)\\right\]\\bm\{Q\}\_\{r\}^\{\\top\},\(A\.16\)‖𝑪rSPlus−𝑰‖F2\\displaystyle\\\|\\bm\{C\}\_\{r\}^\{\\rm SPlus\}\-\\bm\{I\}\\\|\_\{F\}^\{2\}=∑i≠j\[2πarcsin\(𝒞r,ij\)\]2≤‖𝓒r−𝑰‖F2\.\\displaystyle=\\sum\_\{i\\neq j\}\\left\[\\frac\{2\}\{\\pi\}\\arcsin\(\\mathcal\{C\}\_\{r,ij\}\)\\right\]^\{2\}\\leq\\\|\\bm\{\\mathcal\{C\}\}\_\{r\}\-\\bm\{I\}\\\|\_\{F\}^\{2\}\.\(A\.17\)Herearcsin\\arcsinacts entrywise\. The same sign identities hold for the positive scale\-mixture model𝐱r=τr𝐲r\\bm\{x\}\_\{r\}=\\sqrt\{\\tau\_\{r\}\}\\bm\{y\}\_\{r\}, where𝐲r∼𝒩\(𝟎,𝚺r\)\\bm\{y\}\_\{r\}\\sim\\mathcal\{N\}\(\\bm\{0\},\\bm\{\\Sigma\}\_\{r\}\),τr\>0\\tau\_\{r\}\>0is independent of𝐲r\\bm\{y\}\_\{r\}, and𝔼\[τr\]=1\\mathbb\{E\}\[\\tau\_\{r\}\]=1\.
###### Proof\.
Vectorization gives Equations \([A\.12](https://arxiv.org/html/2610.19064#A1.E12)\) and \([A\.13](https://arxiv.org/html/2610.19064#A1.E13)\)\. Since the maps and denominators are frozen, the SOAP covariance is the linear covariance sandwich\. Substitute𝚪r=𝑫r1/2𝓒r𝑫r1/2\\bm\{\\Gamma\}\_\{r\}=\\bm\{D\}\_\{r\}^\{1/2\}\\bm\{\\mathcal\{C\}\}\_\{r\}\\bm\{D\}\_\{r\}^\{1/2\}and𝑫^r−1/2𝑫r1/2=𝑨r\\widehat\{\\bm\{D\}\}\_\{r\}^\{\-1/2\}\\bm\{D\}\_\{r\}^\{1/2\}=\\bm\{A\}\_\{r\}to obtain Equation \([A\.14](https://arxiv.org/html/2610.19064#A1.E14)\)\. 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\)\\mathcal\{C\}\_\{r,ij\}/\\sqrt\{\(1\+e\_\{ri\}\)\(1\+e\_\{rj\}\)\}\. Summing their squared differences from the identity proves Equation \([A\.15](https://arxiv.org/html/2610.19064#A1.E15)\)\.
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ρ\\rhohave sign disagreement probabilityarccos\(ρ\)/π\\arccos\(\\rho\)/\\pi: represent them as projections of an isotropic two\-dimensional Gaussian onto two unit directions with anglearccos\(ρ\)\\arccos\(\\rho\)and use its uniform angular distribution\. Their sign product therefore has expectation1−2arccos\(ρ\)/π=\(2/π\)arcsin\(ρ\)1\-2\\arccos\(\\rho\)/\\pi=\(2/\\pi\)\\arcsin\(\\rho\)\. This is the Gaussian arcsine identity used in one\-bit covariance estimation\([Dirksen et al\., 2022](https://arxiv.org/html/2610.19064#bib.bib4)\), and proves Equation \([A\.16](https://arxiv.org/html/2610.19064#A1.E16)\)\. For0≤t≤10\\leq t\\leq 1, convexity ofarcsin\(t\)\\arcsin\(t\)places it below the chord joining its endpoint values, so\(2/π\)arcsin\(t\)≤t\(2/\\pi\)\\arcsin\(t\)\\leq t\. Oddness extends the absolute\-value inequality to\[−1,1\]\[\-1,1\]\. Squaring the off\-diagonal entries proves Equation \([A\.17](https://arxiv.org/html/2610.19064#A1.E17)\)\. Finally, multiplication by the positive scalarτr\\sqrt\{\\tau\_\{r\}\}leaves every sign unchanged; independence and𝔼\[τr\]=1\\mathbb\{E\}\[\\tau\_\{r\}\]=1preserve 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=𝑰\\bm\{\\mathcal\{C\}\}\_\{r\}=\\bm\{I\}, SPlus has identity covariance for every worker, including workers using different bases and different input variances\.
For an explicit bound, ifmaxj\|erj\|≤ϵr<1\\max\_\{j\}\|e\_\{rj\}\|\\leq\\epsilon\_\{r\}<1, then
‖𝑪rSOAP−𝑰‖F2≤‖𝒆r‖22\+‖𝓒r−𝑰‖F2\(1−ϵr\)2\.\\\|\\bm\{C\}\_\{r\}^\{\\rm SOAP\}\-\\bm\{I\}\\\|\_\{F\}^\{2\}\\leq\\frac\{\\\|\\bm\{e\}\_\{r\}\\\|\_\{2\}^\{2\}\+\\\|\\bm\{\\mathcal\{C\}\}\_\{r\}\-\\bm\{I\}\\\|\_\{F\}^\{2\}\}\{\(1\-\\epsilon\_\{r\}\)^\{2\}\}\.\(A\.18\)This follows by bounding each denominator in Equation \([A\.15](https://arxiv.org/html/2610.19064#A1.E15)\) 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:
‖𝑸r𝑩𝑸r⊤−𝑰‖=‖𝑩−𝑰‖\.\\\|\\bm\{Q\}\_\{r\}\\bm\{B\}\\bm\{Q\}\_\{r\}^\{\\top\}\-\\bm\{I\}\\\|=\\\|\\bm\{B\}\-\\bm\{I\}\\\|\.\(A\.19\)Orthogonal multiplication preserves singular values, proving Equation \([A\.19](https://arxiv.org/html/2610.19064#A1.E19)\)\. Rotation back can change the directions of residual anisotropy, but cannot increase this norm of its error\.
## Appendix BTraining setup and hyperparameters
##### Models and data\.
The experiments use decoder\-only LLaMA models trained on DCLM\([Li et al\., 2024](https://arxiv.org/html/2610.19064#bib.bib15)\), 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 byRR\.
Setting178M512MTransformer layers912Hidden width10241536Attention heads / key\-value heads8/88/812/1212/12Feed\-forward width26885440ActivationSwiGLUSwiGLUSequence length20482048Global tokens per inner step524,2882,097,152Global sequences per inner step2561024Training tokens \(approximately\)3\.55B10\.25BTable 7:Architecture and training scales\. The global batch is measured per inner step, not per communication round\.
##### Aggregation and comparisons\.
Unless stated otherwise, replicas takeH=30H=30local steps and use a Nesterov SGD outer optimizer\. TheH=1H=1controls 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](https://arxiv.org/html/2610.19064#A3)\.
##### 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 CAnalysis 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\.1Conditional covariance replay
This appendix spells out the covariance measurement used in Table[5](https://arxiv.org/html/2610.19064#S5.T5)\. 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\.1Conditional replay protocol
At a communication boundary, all replicas share the synchronized parameters𝑾\\bm\{W\}, while replicarrhas its own saved optimizer statesrs\_\{r\}\. Table[5](https://arxiv.org/html/2610.19064#S5.T5)uses all 64 saved replica states\. For each replica, we run2n2nrollouts:n=16n=16\. Before every rollout, we restore the same pair\(𝑾,sr\)\(\\bm\{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 afterh∈\{1,30\}h\\in\\\{1,30\\\}local steps,
𝒖r,j\(h\)=concatℓ∈ℒvec\(𝑾ℓ−𝑾r,j,ℓ\(h\)\),j=1,…,2n,\\bm\{u\}\_\{r,j\}^\{\(h\)\}=\\mathop\{\\operatorname\{concat\}\}\_\{\\ell\\in\\mathcal\{L\}\}\\operatorname\{vec\}\\\!\\left\(\\bm\{W\}\_\{\\ell\}\-\\bm\{W\}\_\{r,j,\\ell\}^\{\(h\)\}\\right\),\\qquad j=1,\\ldots,2n,\(C\.1\)whereℒ\\mathcal\{L\}contains attentionq\_projand MLPdown\_projat layers 0, 4, and 8\. For the 178M model, this concatenatesd=11,403,264d=11\{,\}403\{,\}264coordinates\. Theh=1h=1andh=30h=30endpoints come from the same rollout and therefore have coupled finite\-sample errors\.
The per\-replica batch is four sequences forR=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 fixedrrandhh, write
𝝁r,h=𝔼\[𝒖r,j\(h\)\],𝑪r,h=Cov\(𝒖r,j\(h\)\)\.\\bm\{\\mu\}\_\{r,h\}=\\mathbb\{E\}\[\\bm\{u\}\_\{r,j\}^\{\(h\)\}\],\\qquad\\bm\{C\}\_\{r,h\}=\\operatorname\{Cov\}\(\\bm\{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\.\\bm\{z\}\_\{r,i\}^\{\(h\)\}=\\frac\{\\bm\{u\}\_\{r,2i\-1\}^\{\(h\)\}\-\\bm\{u\}\_\{r,2i\}^\{\(h\)\}\}\{\\sqrt\{2\}\},\\qquad i=1,\\ldots,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\.\\mathbb\{E\}\[\\bm\{z\}\_\{r,i\}^\{\(h\)\}\]=\\bm\{0\},\\qquad\\operatorname\{Cov\}\(\\bm\{z\}\_\{r,i\}^\{\(h\)\}\)=\\tfrac\{1\}\{2\}\(\\bm\{C\}\_\{r,h\}\+\\bm\{C\}\_\{r,h\}\)=\\bm\{C\}\_\{r,h\}\.\(C\.4\)The subtraction therefore removes the unknown mean, while1/21/\\sqrt\{2\}restores the original covariance scale\.
#### C\.1\.2Trace estimators from full\-coordinate inner products
Fix a horizon and suppresshhin the notation\. Define the Gram entries
Gri,sj=𝒛r,i⊤𝒛s,j\.G\_\{ri,sj\}=\\bm\{z\}\_\{r,i\}^\{\\top\}\\bm\{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 fromnnpaired differences per replica\.
For each replica, define
tr\\displaystyle t\_\{r\}=1n∑i=1nGri,ri,\\displaystyle=\\frac\{1\}\{n\}\\sum\_\{i=1\}^\{n\}G\_\{ri,ri\},\(C\.6\)br\\displaystyle b\_\{r\}=1n\(n−1\)∑i≠jGri,riGrj,rj,\\displaystyle=\\frac\{1\}\{n\(n\-1\)\}\\sum\_\{i\\neq j\}G\_\{ri,ri\}G\_\{rj,rj\},\(C\.7\)qrr\\displaystyle q\_\{rr\}=1n\(n−1\)∑i≠jGri,rj2\.\\displaystyle=\\frac\{1\}\{n\(n\-1\)\}\\sum\_\{i\\neq j\}G\_\{ri,rj\}^\{2\}\.\(C\.8\)Under the IID\-rollout idealization, for independent paired differences,
𝔼\[tr\]=tr\(𝑪r\),𝔼\[br\]=tr\(𝑪r\)2,𝔼\[qrr\]=tr\(𝑪r2\)\.\\mathbb\{E\}\[t\_\{r\}\]=\\operatorname\{tr\}\(\\bm\{C\}\_\{r\}\),\\qquad\\mathbb\{E\}\[b\_\{r\}\]=\\operatorname\{tr\}\(\\bm\{C\}\_\{r\}\)^\{2\},\\qquad\\mathbb\{E\}\[q\_\{rr\}\]=\\operatorname\{tr\}\(\\bm\{C\}\_\{r\}^\{2\}\)\.\(C\.9\)The last identity follows from
𝔼\[\(𝒛⊤𝒛′\)2\]=tr\(𝔼\[𝒛𝒛⊤\]𝔼\[𝒛′𝒛′⊤\]\)=tr\(𝑪2\)\\mathbb\{E\}\[\(\\bm\{z\}^\{\\top\}\\bm\{z\}^\{\\prime\}\)^\{2\}\]=\\operatorname\{tr\}\\\!\\left\(\\mathbb\{E\}\[\\bm\{z\}\\bm\{z\}^\{\\top\}\]\\mathbb\{E\}\[\\bm\{z\}^\{\\prime\}\\bm\{z\}^\{\\prime\\top\}\]\\right\)=\\operatorname\{tr\}\(\\bm\{C\}^\{2\}\)\(C\.10\)for independent𝒛\\bm\{z\}and𝒛′\\bm\{z\}^\{\\prime\}\. Usingi≠ji\\neq jis important: including the diagonal would add fourth\-moment self\-products and bias the estimate upward\.
For different replicasr≠sr\\neq s, the cross term
qrs=1n2∑i,jGri,sj2q\_\{rs\}=\\frac\{1\}\{n^\{2\}\}\\sum\_\{i,j\}G\_\{ri,sj\}^\{2\}\(C\.11\)estimatestr\(𝑪r𝑪s\)\\operatorname\{tr\}\(\\bm\{C\}\_\{r\}\\bm\{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\.3Scaled sum of local covariances
The probe constructs the scaled sum
𝑪~=1R2∑r=1R𝑪r\.\\widetilde\{\\bm\{C\}\}=\\frac\{1\}\{R^\{2\}\}\\sum\_\{r=1\}^\{R\}\\bm\{C\}\_\{r\}\.\(C\.12\)For the actual replica average𝒖¯=R−1∑r𝒖r\\bar\{\\bm\{u\}\}=R^\{\-1\}\\sum\_\{r\}\\bm\{u\}\_\{r\}, this equalsCov\(𝒖¯\)\\operatorname\{Cov\}\(\\bar\{\\bm\{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𝑪~\\widetilde\{\\bm\{C\}\}, it uses
T^\\displaystyle\\widehat\{T\}=1R2∑rtr,\\displaystyle=\\frac\{1\}\{R^\{2\}\}\\sum\_\{r\}t\_\{r\},𝔼\[T^\]\\displaystyle\\mathbb\{E\}\[\\widehat\{T\}\]=tr\(𝑪~\),\\displaystyle=\\operatorname\{tr\}\(\\widetilde\{\\bm\{C\}\}\),\(C\.13\)S^\\displaystyle\\widehat\{S\}=1R4∑r,sqrs,\\displaystyle=\\frac\{1\}\{R^\{4\}\}\\sum\_\{r,s\}q\_\{rs\},𝔼\[S^\]\\displaystyle\\mathbb\{E\}\[\\widehat\{S\}\]=tr\(𝑪~2\),\\displaystyle=\\operatorname\{tr\}\(\\widetilde\{\\bm\{C\}\}^\{2\}\),\(C\.14\)B^\\displaystyle\\widehat\{B\}=1R4\(∑rbr\+∑r≠strts\),\\displaystyle=\\frac\{1\}\{R^\{4\}\}\\left\(\\sum\_\{r\}b\_\{r\}\+\\sum\_\{r\\neq s\}t\_\{r\}t\_\{s\}\\right\),𝔼\[B^\]\\displaystyle\\mathbb\{E\}\[\\widehat\{B\}\]=tr\(𝑪~\)2\.\\displaystyle=\\operatorname\{tr\}\(\\widetilde\{\\bm\{C\}\}\)^\{2\}\.\(C\.15\)The separate estimatorB^\\widehat\{B\}avoids squaring the noisy estimateT^\\widehat\{T\}\. The reported anisotropy is
a^=max\{0,dS^B^−1\}\.\\widehat\{a\}=\\sqrt\{\\max\\\!\\left\\\{0,\\frac\{d\\widehat\{S\}\}\{\\widehat\{B\}\}\-1\\right\\\}\}\.\(C\.16\)Under the IID idealization,T^\\widehat\{T\},S^\\widehat\{S\}, andB^\\widehat\{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\\lambda\_\{1\},\\ldots,\\lambda\_\{d\}are the eigenvalues of𝑪~\\widetilde\{\\bm\{C\}\}, the population quantity is
a2=d∑iλi2\(∑iλi\)2−1,reffd=11\+a2,a^\{2\}=\\frac\{d\\sum\_\{i\}\\lambda\_\{i\}^\{2\}\}\{\(\\sum\_\{i\}\\lambda\_\{i\}\)^\{2\}\}\-1,\\qquad\\frac\{r\_\{\\mathrm\{eff\}\}\}\{d\}=\\frac\{1\}\{1\+a^\{2\}\},\(C\.17\)wherereff=\(∑iλi\)2/∑iλi2r\_\{\\mathrm\{eff\}\}=\(\\sum\_\{i\}\\lambda\_\{i\}\)^\{2\}/\\sum\_\{i\}\\lambda\_\{i\}^\{2\}is the participation\-ratio effective rank\. Thusaameasures covariance shape independently of scale; it is undefined when the covariance is identically zero\. For example,diag\(5,5\)\\operatorname\{diag\}\(5,5\)hasa=0a=0, whereasdiag\(9,1\)\\operatorname\{diag\}\(9,1\)hasa=0\.8a=0\.8\. Multiplying either matrix by a positive scalar leavesaaunchanged\.
#### C\.1\.4Scope 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, then\(n−1\)n\(n\-1\)andn2n^\{2\}sums reuse the samennvectors; their terms are not independent observations\. Finally, interpreting Equation \([C\.12](https://arxiv.org/html/2610.19064#A3.E12)\) as the covariance of the actual average requires vanishing conditional cross\-replica noise covariances\.
Loweraaindicates a more isotropic𝑪~\\widetilde\{\\bm\{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\.2DP\-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\.
Letjjindex one of the 63 hidden matrices, vectorized in parameter coordinates\. Write𝒙j,k\\bm\{x\}\_\{j,k\}for its averaged local displacement and𝒚j,k\\bm\{y\}\_\{j,k\}for its DP displacement in repetitionkk\. On the 16 fitting repetitions atR=8R=8, define
α^j=⟨𝒙¯j,fit,𝒚¯j,fit⟩‖𝒙¯j,fit‖2,\\widehat\{\\alpha\}\_\{j\}=\\frac\{\\langle\\overline\{\\bm\{x\}\}\_\{j,\\rm fit\},\\overline\{\\bm\{y\}\}\_\{j,\\rm fit\}\\rangle\}\{\\left\\lVert\\overline\{\\bm\{x\}\}\_\{j,\\rm fit\}\\right\\rVert^\{2\}\},and take𝒒^j\\widehat\{\\bm\{q\}\}\_\{j\}to be the leading eigenvector of the sample covariance ofα^j𝒙j,k\\widehat\{\\alpha\}\_\{j\}\\bm\{x\}\_\{j,k\}on those same fitting repetitions\. The reference mean is𝒅j=𝔼\[𝒚j,k\]\\bm\{d\}\_\{j\}=\\mathbb\{E\}\[\\bm\{y\}\_\{j,k\}\], and the mean discrepancy is𝒃j=α^j𝔼\[𝒙j,k\]−𝒅j\\bm\{b\}\_\{j\}=\\widehat\{\\alpha\}\_\{j\}\\mathbb\{E\}\[\\bm\{x\}\_\{j,k\}\]\-\\bm\{d\}\_\{j\}, conditional on the fitting sample\. The decomposition in Equation \([10](https://arxiv.org/html/2610.19064#S4.E10)\) therefore applies to each matrix with these fitted quantities held fixed\.
We use the remainingK=24K=24independent repetitions for evaluation\. Set𝒛j,k=α^j𝒙j,k−𝒚j,k\\bm\{z\}\_\{j,k\}=\\widehat\{\\alpha\}\_\{j\}\\bm\{x\}\_\{j,k\}\-\\bm\{y\}\_\{j,k\},tj,k=𝒒^j⊤𝒛j,kt\_\{j,k\}=\\widehat\{\\bm\{q\}\}\_\{j\}^\{\\top\}\\bm\{z\}\_\{j,k\}, andpj,k=α^j𝒒^j⊤𝒙j,kp\_\{j,k\}=\\widehat\{\\alpha\}\_\{j\}\\widehat\{\\bm\{q\}\}\_\{j\}^\{\\top\}\\bm\{x\}\_\{j,k\}\. Withs2s^\{2\}denoting the unbiased sample variance on these repetitions, the unnormalized estimates are
B^j=t¯j2−s2\(tj,k\)K,N^j=s2\(pj,k\)\.\\widehat\{B\}\_\{j\}=\\overline\{t\}\_\{j\}^\{\\,2\}\-\\frac\{s^\{2\}\(t\_\{j,k\}\)\}\{K\},\\qquad\\widehat\{N\}\_\{j\}=s^\{2\}\(p\_\{j,k\}\)\.Conditional on the fit and independent evaluation repetitions,B^j\\widehat\{B\}\_\{j\}is unbiased for\(𝒒^j⊤𝒃j\)2\(\\widehat\{\\bm\{q\}\}\_\{j\}^\{\\top\}\\bm\{b\}\_\{j\}\)^\{2\}andN^j\\widehat\{N\}\_\{j\}estimatesα^j2𝒒^j⊤𝑪j𝒒^j\\widehat\{\\alpha\}\_\{j\}^\{2\}\\widehat\{\\bm\{q\}\}\_\{j\}^\{\\top\}\\bm\{C\}\_\{j\}\\widehat\{\\bm\{q\}\}\_\{j\}, where𝑪j=Cov\(𝒙j,k\)\\bm\{C\}\_\{j\}=\\operatorname\{Cov\}\(\\bm\{x\}\_\{j,k\}\)\. Crucially, the correction uses the variance of the*paired*difference𝒛j,k\\bm\{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](https://arxiv.org/html/2610.19064#S5.T5)areB^j/‖𝒅^j‖2\\widehat\{B\}\_\{j\}/\\left\\lVert\\widehat\{\\bm\{d\}\}\_\{j\}\\right\\rVert^\{2\}andN^j/‖𝒅^j‖2\\widehat\{N\}\_\{j\}/\\left\\lVert\\widehat\{\\bm\{d\}\}\_\{j\}\\right\\rVert^\{2\}, with𝒅^j=K−1∑k𝒚j,k\\widehat\{\\bm\{d\}\}\_\{j\}=K^\{\-1\}\\sum\_\{k\}\\bm\{y\}\_\{j,k\}, expressed in units of10−610^\{\-6\}\. The third column is100B^j/N^j100\\widehat\{B\}\_\{j\}/\\widehat\{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\.\\widehat\{B\}\_\{j\}^\{\\mathrm\{tr\}\}=\\left\\lVert\\overline\{\\bm\{z\}\}\_\{j\}\\right\\rVert^\{2\}\-\\frac\{1\}\{K\(K\-1\)\}\\sum\_\{k\}\\left\\lVert\\bm\{z\}\_\{j,k\}\-\\overline\{\\bm\{z\}\}\_\{j\}\\right\\rVert^\{2\},\\qquad\\widehat\{N\}\_\{j\}^\{\\mathrm\{tr\}\}=\\frac\{\\widehat\{\\alpha\}\_\{j\}^\{2\}\}\{K\-1\}\\sum\_\{k\}\\left\\lVert\\bm\{x\}\_\{j,k\}\-\\overline\{\\bm\{x\}\}\_\{j\}\\right\\rVert^\{2\}\.They estimate‖𝒃j‖2\\left\\lVert\\bm\{b\}\_\{j\}\\right\\rVert^\{2\}andα^j2tr\(𝑪j\)\\widehat\{\\alpha\}\_\{j\}^\{2\}\\operatorname\{tr\}\(\\bm\{C\}\_\{j\}\), respectively\. The same DP\-mean normalizer and equal per\-matrix averaging apply; the mean/noise ratio is computed per matrix before averaging\. AtR=8,H=30R=8,H=30, these average ratios are6\.2%6\.2\\%for SPlus,2\.8%2\.8\\%for SOAP,10%10\\%for Muon, and6\.7%6\.7\\%for AdamW\.
### C\.3Curvature 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\\bm\{z\}\_\{i\}=\(\\bm\{u\}\_\{i\}\-\\bm\{u\}\_\{i\}^\{\\prime\}\)/\\sqrt\{2\}, using fresh local batches ofb=256/Rb=256/Rsequences\. Under the independent\-replay model these differences are centered and have the local\-update covariance\. For Table[6](https://arxiv.org/html/2610.19064#S5.T6), use the first 64 of the 128 paired samples to select
𝒒^=argmax‖𝒒‖=1164∑i=164\(𝒒⊤𝒛i\)2\.\\widehat\{\\bm\{q\}\}=\\underset\{\\left\\lVert\\bm\{q\}\\right\\rVert=1\}\{\\arg\\max\}\\;\\frac\{1\}\{64\}\\sum\_\{i=1\}^\{64\}\(\\bm\{q\}^\{\\top\}\\bm\{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\\bm\{v\}\_\{1\},\\ldots,\\bm\{v\}\_\{16\}\. Report
κ\(𝒒^\)=𝒒^⊤𝓗𝒒^116∑j=116𝒗j⊤𝓗𝒗j\.\\kappa\(\\widehat\{\\bm\{q\}\}\)=\\frac\{\\widehat\{\\bm\{q\}\}^\{\\top\}\\bm\{\\mathcal\{H\}\}\\widehat\{\\bm\{q\}\}\}\{\\frac\{1\}\{16\}\\sum\_\{j=1\}^\{16\}\\bm\{v\}\_\{j\}^\{\\top\}\\bm\{\\mathcal\{H\}\}\\bm\{v\}\_\{j\}\}\.The numerator measures curvature along a direction selected for high noise variance, not maximum curvature\. The denominator estimatestr\(𝓗\)/d\\operatorname\{tr\}\(\\bm\{\\mathcal\{H\}\}\)/din 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, notH=30H=30aggregate measurements\. The ratio removes the noise magnitude and does not estimatetr\(𝓗𝑪\)\\operatorname\{tr\}\(\\bm\{\\mathcal\{H\}\}\\bm\{C\}\)\.Similar Articles
An Exact Generate - Transform Decomposition of Small-LLM Team Scaling Across Orchestration Architectures
This paper introduces an exact generate–transform decomposition to analyze how small 7–9B LLM teams scale across eight orchestration architectures, finding that scaling returns are sharply task-dependent — Proposer-Critic wins on arithmetic benchmarks but no architecture dominates across all tasks.
Spectral-Driven Adaptive-Rank Optimizer for Large Language Model Reinforcement Learning with Verifiable Rewards
论文研究了LLM在可验证奖励强化学习(RLVR)下的优化几何特性,发现梯度能量会逐渐集中到不断收缩的有效子空间中,进而提出自适应秩优化器Ares,根据梯度谱能量分布动态调整更新秩,在多个模型规模和RLVR任务上取得优于主流优化器的表现。
The Ephemeral Cost of Switching Optimizers Across LLM Training Stages
该研究系统考察了在预训练与后训练阶段切换 AdamW 和 Muon 优化器的代价,发现最终性能主要由两个优化器各自独立的效果解释,切换优化器几乎没有额外的性能惩罚,后训练优化器的选择可基本独立于预训练优化器。
@AlphaSignalAI: You can now boost any LLM's accuracy 2-10x without training it. Most teams improve model accuracy by fine-tuning or swa…
OptiLLM is an open-source proxy that boosts any LLM's accuracy 2-10x by adding extra compute at inference time, using techniques like multi-agent cross-verification and Monte Carlo tree search.
SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales
This paper improves higher-order optimizers SOAP and Muon for large-scale LLM pretraining, addressing instabilities at large batch sizes and introducing a layer-wise distributed optimizer compatible with Megatron-LM. Experiments show they consistently outperform AdamW at billion-parameter scales.