SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales

# SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales
Source: [https://arxiv.org/html/2607.20548](https://arxiv.org/html/2607.20548)
Aditya VavreBoxiang WangDeyu FuHao WuMike ChrzanowskiBryan CatanzaroDheevatsa MudigereJeff PoolMichael LightstoneMohammad ShoeybiMostofa PatwaryNima TajbakhshTijmen Blankevoort NVIDIA \{mkhona, avavre, boxiangw, deyuf, skyw, mchrzanowski,bcatanzaro, dmudigere, jpool, mlightstone, mshoeybi, mpatwary, ntajbakhsh, tblankevoort\}@nvidia\.com

###### Abstract

Higher\-order optimizers such as Muon and SOAP offer faster convergence than AdamW, but their computational cost and numerical stability challenges have limited adoption at scale\. In this work, we adapt and enhance preconditioned gradient methods to overcome the practical challenges of large\-scale LLM pretraining\.

We first identify instabilities in SOAP at large batch sizes and propose algorithmic modifications including per\-step QR orthogonalization and improved preconditioning strategies that eliminate loss spikes and enable stable training in these regimes\. We then present a unified empirical study of SOAP, Muon, and AdamW using update\-RMS matching to ensure fair learning rate transfer across optimizers\. As part of this analysis, we empirically evaluate the orthogonalization quality of Muon\. Our experiments on multi\-billion\-parameter models trained on trillions of tokens reveal that SOAP and Muon consistently outperform AdamW at the scales we tested\. Notably, at batch sizes of up to 100M tokens for next\-token prediction, these optimizers maintain training stability and quality while AdamW degrades\.

To enable efficient training at large scale, we introduce a layer\-wise distributed optimizer compatible with Megatron\-LM\. Our implementation balances memory and hides communication while avoiding approximations to the optimizer computations, thus retaining their convergence benefits\. Additionally, we identify and build specific system\-level improvements to further accelerate our layer\-wise implementation\. To support the research community, we release a codebase that contains emerging algorithms for optimization:

https://github\.com/NVIDIA\-NeMo/Emerging\-Optimizers

## 1Introduction

The choice of the optimization algorithm remains at the center of frontier model training frameworks\. Indeed, it serves as the heart of all large\-scale model training, linking systems engineering with algorithms and model convergence\. From a systems perspective, the optimizer dictates the fundamental constraints of distributed execution: optimizer states typically consume significantly more memory than the model parameters themselves, heavily influencing model sharding and memory management strategies\. Further, an optimizer’s ability to maintain stability at extreme batch sizes dictates how far training can scale across large accelerator clusters without being bottlenecked by communication overhead due to the synchronous nature of training workloads\. Algorithmically, the optimizer governs data efficiency, convergence rates, and the generalization ability of the model\.

Despite this dual impact, the evolution of optimization in deep learning has been defined by a practical tension: the computational simplicity and scalability of element\-wise scalar optimizers, such as AdamW\[kingma2014adam,loshchilov2017decoupled\], RMSProp\[hinton2012neural\], and LaProp\[ziyin2020laprop\], have largely overshadowed curvature\-aware higher\-order methods\. While first\-order methods like AdamW dominate due to their ease of use, they inherently treat individual parameter element updates independently, ignoring the correlational structure between gradients and the operator nature of the corresponding neural network weights\. Conversely, higher\-order optimizers \(approximating second\-order information\) capture the geometry of the loss landscape\[george2018fast,martens2015optimizing,bollapragada2018progressive\], theoretically allowing for larger steps, faster convergence, and potentially higher accuracy, but suffer from a lack of scalability due to their complexity\.

Recently, tensor\-wise adaptive methods such as the Shampoo optimizer and modern variants like SOAP\[vyas2024soap\], Eigen\-corrected Shampoo\[eschenhagen2025purifying\], and KL\-Shampoo\[lin2025understanding\], along with spectral optimizers such as Muon and Scion\[pethick2025training\], have emerged as a middle ground, offering the benefits of second\-order methods with manageable computational and memory overhead\. However, these methods face scalability hurdles when applied to large frontier models like fine\-grained MoEs\.

The rest of this paper is organized around the formulation, empirical evaluation, and systems engineering required to scale these higher\-order optimizers\. Our core contributions are as follows:

- •Large Batch\-Size Scaling for MoEs with Muon: We use an update\-RMS matching framework to fairly transfer learning rates and compare AdamW, Muon, and SOAP optimizers\. We motivate the MoE large\-batch setting in Sec\.[3](https://arxiv.org/html/2607.20548#S3)and present empirical results in Sec\.[5\.3](https://arxiv.org/html/2607.20548#S5.SS3)and Sec\.[5\.5](https://arxiv.org/html/2607.20548#S5.SS5)\. Our evaluation on up to 72\-billion parameter MoE models demonstrates that while AdamW’s performance degrades beyond critical batch sizes, both Muon and SOAP maintain token\-efficient convergence at global batch sizes up to 100M tokens\.
- •Fixing Instability for SOAP’s Preconditioner Computation: We identify a critical instability in the standard implementation of SOAP111[https://github\.com/nikhilvyas/SOAP](https://github.com/nikhilvyas/SOAP)unique to the large\-batch regime, caused by a lag between the preconditioner and current gradient statistics\. Sec\.[5\.4](https://arxiv.org/html/2607.20548#S5.SS4)shows that enforcing the per\-step eigenbasis updates with current gradients and integrating a recently developed KL\-divergence\-based covariance estimation method for accumulating SOAP’s Kronecker factors\[lin2025understanding\]eliminates all training loss spikes and improves training stability, allowing for more confidence for scaling to larger batch and model sizes in production\.
- •Comparing SOAP and Muon: We compare the accuracy and stability of these optimizers in the large batch setting, finding that both Muon and SOAP outperform AdamW and KL\-SOAP has a slight edge over Muon\. We also discuss the limitations of our conclusions\.
- •Scalable Parallelization: We present an efficient, layer\-wise distributed implementation designed to enable scalability within the Megatron\-LM framework\. We briefly discuss the distributed optimizer constraints in Sec\.[3\.3](https://arxiv.org/html/2607.20548#S3.SS3)and present the implementation in Sec\.[6](https://arxiv.org/html/2607.20548#S6)\.
- •Emerging\-Optimizers: We contribute an open\-source library for optimizer research that includes implementations of all optimizers considered in this report as well as many experimental variants\.

## 2A Primer on Adam, Muon and SOAP

Adam\-style optimizers are the most common optimizers used in modern LLM pretraining, which extend stochastic gradient descent \(SGD\) with coordinate\-wise adaptive learning rates and decoupled weight decay\[kingma2014adam,loshchilov2017decoupled\]\. Here we provide a very brief overview of these methods and how they are related\.

Letgt∈ℝd=m​ng\_\{t\}\\in\\mathbb\{R\}^\{d=mn\}be the flattened gradient vectorGt∈ℝm×nG\_\{t\}\\in\\mathbb\{R\}^\{m\\times n\}of a parameter tensor at steptt\. AdamW maintains exponential moving averages \(EMA\) of the gradient and squared gradient for each parameter element independently withβ1,β2\\beta\_\{1\},\\beta\_\{2\}being the EMA timescale coefficients\.

mt=β1​mt−1\+\(1−β1\)​gt​vt=β2​vt−1\+\(1−β2\)​gt⊙gtm\_\{t\}=\\beta\_\{1\}m\_\{t\-1\}\+\(1\-\\beta\_\{1\}\)g\_\{t\}\\\\ v\_\{t\}=\\beta\_\{2\}v\_\{t\-1\}\+\(1\-\\beta\_\{2\}\)g\_\{t\}\\odot g\_\{t\}where⊙\\odotdenotes the element\-wise \(Hadamard\) product\. The preconditioned update direction for AdamW \(ignoring first\-moment momentummtm\_\{t\}and bias corrections for simplicity\) applies a diagonal preconditioner to the gradient:

ut=mt∘diag​\(1\(vt\)1/2\+ϵ\)u\_\{t\}=m\_\{t\}\\circ\\text\{diag\}\\left\(\\dfrac\{1\}\{\(v\_\{t\}\)^\{1/2\}\+\\epsilon\}\\right\)This element\-wise structure makes AdamW memory\-regular, and easy to shard in large distributed systems\. However, it also means that AdamW ignores the matrix or tensor structure of neural network weights: each coordinate is rescaled separately, but correlations across rows, columns, heads, experts, or projection subspaces are not modeled explicitly\.

A separate line of optimizers attempts to exploit this structure through preconditioning\. Shampoo\[gupta2018shampoo\]approximates second\-order information by maintaining Kronecker\-factored gradient statistics for each tensor dimension, allowing matrix\- or tensor\-shaped parameters to be updated using curvature\-aware transformations without forming a full dense preconditioner\. Shampoo maintains two separate symmetric covariance matrices for the rows and columns:

Lt\\displaystyle L\_\{t\}=β2​Lt−1\+\(1−β2\)​Gt​Gt⊤\\displaystyle=\\beta\_\{2\}L\_\{t\-1\}\+\(1\-\\beta\_\{2\}\)G\_\{t\}G\_\{t\}^\{\\top\}Rt\\displaystyle R\_\{t\}=β2​Rt−1\+\(1−β2\)​Gt⊤​Gt\\displaystyle=\\beta\_\{2\}R\_\{t\-1\}\+\(1\-\\beta\_\{2\}\)G\_\{t\}^\{\\top\}G\_\{t\}Shampoo’s update direction is then produced through a non\-diagonal kronecker\-factored preconditioner, derived from inverse powers of the kronecker factors:

ut=\(Rt−1/4⊗Lt−1/4\)​gtu\_\{t\}=\(R\_\{t\}^\{\-1/4\}\\otimes L\_\{t\}^\{\-1/4\}\)g\_\{t\}
SOAP \(Alg\.[2](https://arxiv.org/html/2607.20548#alg2)\) builds on Shampoo by combining Shampoo\-style preconditioning with Adam\-like elementwise adaptive updates in the preconditioner eigenbasis\[vyas2024soap\]:

ut=QL​Adam​\(QLT​mt​QR\)​QRTu\_\{t\}=Q\_\{L\}\\texttt\{Adam\}\(Q\_\{L\}^\{T\}m\_\{t\}Q\_\{R\}\)Q\_\{R\}^\{T\}\(1\)
Intuitively, SOAP first rotates gradients into a basis where their row and column correlations are approximately diagonalized \(i\.e\. the eigenbasis ofRtR\_\{t\}andLtL\_\{t\}\), applies an adaptive element\-wise update in that basis, and then rotates the update back out of the preconditioner eigenbasis\. This gives SOAP a closer connection to AdamW than classical Shampoo, since for identity\-valued eigenvectors, we exactly recover AdamW, while retaining the ability to exploit structure within weight matrices, possibly allowing for smoother hyperparameter transfer\. However both Shampoo and SOAP maintain full precision preconditioner and kronecker factor matrices requiring significantly more memory than AdamW\.

Muon \(Alg\.[1](https://arxiv.org/html/2607.20548#alg1)\) skips computing and storing preconditioners and follows a different strategy\. Rather than estimating preconditioners and full adaptive second\-moment statistics in a preconditioned basis, Muon applies momentum like AdamW and Shampoo, Eq\.[2](https://arxiv.org/html/2607.20548#S2.Ex1), and then orthogonalizes it\. To orthogonalizeMtM\_\{t\}and approximate its polar factor \(the closest orthogonal matrix, equivalent toU​V⊤UV^\{\\top\}from the singular value decompositionMt=U​Σ​V⊤M\_\{t\}=U\\Sigma V^\{\\top\}\), Muon employs Newton\-Schulz iterations\[pethick2025training,liu2025muon,shah2025practical,amsel2025polar,jordanmuon\]which use matrix polynomial iteration to approximately compute the polar factor\.

This can be viewed as a spectral update, since for 2D weight matrices, Muon changes the direction of the update so that its singular values are controlled, rather than applying independent coordinate\-wise rescaling like AdamW\. In simplified terms, SOAP/Shampoo\-style whitening and Muon\-style orthogonalization are closely related: both reduce the dominance of large singular directions and produce updates that better respect the operator structure of linear layers\. Notably, turning off EMA by settingβ1,β2=0\\beta\_\{1\},\\beta\_\{2\}=0in Shampoo mathematically reduces its preconditioned update directly to the polar factorU​V⊤UV^\{\\top\}, since\(𝐆𝐆𝐓\)−1/4​𝐆​\(𝐆𝐓​𝐆\)−1/4=U​V⊤\(\\mathbf\{GG^\{T\}\}\)^\{\-1/4\}\\,\\mathbf\{G\}\\,\(\\mathbf\{G^\{T\}G\}\)^\{\-1/4\}=UV^\{\\top\}, recovering the Muon update\. See Appendix[A\.1](https://arxiv.org/html/2607.20548#A1.SS1)for a detailed SVD interpretation of SOAP’s preconditioning and the intuitive relationships between all three optimizers\.

There are practical differences between using these optimizers at scale\. Here we summarize the qualitative differences\. AdamW is the easiest optimizer to scale because its state and computation are element\-wise and hence infinitely shardable\. SOAP is more expressive, but it must maintain and update Kronecker factors and eigenbases, making it more sensitive to numerical stability, preconditioner freshness, has a much larger memory footprint, and complicated matrix\-based distributed tensor layout for sharding\. Muon has a lower footprint than AdamW since it does not maintain second moments but it still requires full 2D matrix updates which faces the same matrix\-based sharding issues as SOAP\. Muon uses more computation in the optimizer step and this overhead depends on the quality and cost of orthogonalization which scales with number of Newton\-Schulz iterations\.

## 3A Primer on Optimization for Large Scale, Large Batch Pretraining

### 3\.1Batch Size Scaling for Mixture\-of\-Experts Models

Large\-batch scaling is particularly nuanced for Mixture\-of\-Experts \(MoE\) models because the majority of the model’s parameters reside in the expert linear layers \(often\>90%\>90\\%of the total parameter count\)\. Unlike dense models, these layers observe a significantly lowereffectivebatch size due to sparse top\-k routing\. Consider a global batch size ofBGlobalB\_\{\\text\{Global\}\}tokens trained with a top\-kkrouting strategy overNNtotal experts\. Assuming an ideally load\-balanced router, the effective batch sizeBeffexpertB^\{\\text\{expert\}\}\_\{\\text\{eff\}\}seen by any single expert is given by:

Beffexpert=BGlobal×kNB^\{\\text\{expert\}\}\_\{\\text\{eff\}\}=B\_\{\\text\{Global\}\}\\times\\frac\{k\}\{N\}\(2\)Sincek≪Nk\\ll Nin typical fine\-grained sparse MoE configurations \(e\.g\., top\-8 gating with 256 experts\[liu2024deepseek\], leading toBeffexpert/BGlobal=0\.03125\{B^\{\\text\{expert\}\}\_\{\\text\{eff\}\}\}/\{B\_\{\\text\{Global\}\}\}=0\.03125\), each expert sees only a small fraction of the global batch\. Consequently, when scaling to very large global batch sizes, the sparse expert parameters remain in a lower\-effective\-batch regime that is generally easier to optimize\. In contrast, the dense and shared parameters must process the full global batch\. As a result, increasing the global batch size primarily stresses the large\-batch tolerance of the dense components\. This makes MoE training highly reliant on optimization methods that can maintain stability and token efficiency specifically for these dense parameters operating at extreme batch sizes\[sun2024hunyuan\]\.

### 3\.2Batch Size–Learning Rate Scaling and Update Variance

Because our comparisons change global batch size by up to 4x, we need a learning\-rate transfer rule that keeps update scale comparable before attributing differences to the optimizer itself\. In the context of SGD, modifying the mini\-batch size requires a corresponding adjustment to the learning rate to maintain stable and transferrable optimization dynamics\. Let us consider a simple SGD update where the gradient estimatorg^B\\hat\{g\}\_\{B\}, computed over a mini\-batch of sizeBB, is given by:

g^B=1B​∑i=1B∇Li​\(θ\)\\hat\{g\}\_\{B\}=\\frac\{1\}\{B\}\\sum\_\{i=1\}^\{B\}\\nabla L\_\{i\}\(\\theta\)
The parameter update is defined asΔ​θ=η​g^B\\Delta\\theta=\\eta\\hat\{g\}\_\{B\}, whereη\\etarepresents the learning rate\. Assuming the individual sample gradients are independent, the covariance of the gradient estimator scales inversely with the batch size:

Cov​\(g^B\)=1B​Σ\\mathrm\{Cov\}\(\\hat\{g\}\_\{B\}\)=\\frac\{1\}\{B\}\\SigmawhereΣ\\Sigmais the covariance matrix of the per\-sample gradients\. Consequently, the covariance of the parameter update itself is:

Cov​\(Δ​θ\)=η2​1B​Σ\\mathrm\{Cov\}\(\\Delta\\theta\)=\\eta^\{2\}\\frac\{1\}\{B\}\\Sigma
When scaling the batch size from an initial sizeBBto a target sizeB′B^\{\\prime\}, a heuristic approach is to preserve the variance scale of the parameter updates\. We equate the covariances of the updates under both batch size regimes:

η2B​Σ=\(η′\)2B′​Σ\\frac\{\\eta^\{2\}\}\{B\}\\Sigma=\\frac\{\(\\eta^\{\\prime\}\)^\{2\}\}\{B^\{\\prime\}\}\\Sigma
Solving for the updated learning rateη′\\eta^\{\\prime\}yields the Square Root Scaling Rule:

η′=η​B′B\\eta^\{\\prime\}=\\eta\\sqrt\{\\frac\{B^\{\\prime\}\}\{B\}\}\(3\)
To ensure that the scale of the random fluctuations in parameter updates \(Δ​θ\\Delta\\theta\) remains similar when we switch from batch sizeBBtoB′B^\{\\prime\}, we adjust the learning rate by the square root of the ratio of the batch sizes\. The logic behind this rule is focused on noise or variance preservation\. Although the linear scaling rule\[goyal2017accurate\]has also been proposed, square\-root scaling is safer when making very large leaps in batch size, as it prevents the learning rate from exploding\. Several more rigorous treatments of this rule have been developed, such as in\[malladi2022sdes\]\.

### 3\.3Strategies for Enabling Distributed Training with Preconditioning

The gains from large batch training properties of Muon and SOAP can only be realized if the training system can preserve full matrix structure under tensor parallelism without additional exposed communication\. When tensor parallelism is applied, the weights of a linear layer are distributed across multiple GPUs\. Beyond the simplest method of preconditioning based solely on the weights owned by a given GPU, commonly referred to as “blocking”,Emerging\-Optimizersalso supports two distinct approaches for preconditioning on the entire layer’s weights:

- •Duplicated: The weights are initially all\-gathered to all GPUs within the Tensor Parallelism \(TP\) domain\. Subsequently, each GPU independently executes the Newton\-Schulz iteration\. This mode is preferable for smaller layers where communication is the limiting factor\.
- •Distributed: In each Newton\-Schulz iteration, the intermediate results of the first matrix multiplication are all\-reduced\. This mode is preferable for larger layers where computation is the limiting factor\.

In both the “duplicated” and “distributed” modes, the inputs are first normalized using the statistics of the entire layer’s weight, thereby ensuring mathematical equivalence to the scenario without tensor parallelism\.

Conventional distributed optimizer sharding approaches \(like ZeRO\-1\[zhao2023pytorch,rajbhandari2020zero\]\) are not directly appropriate for matrix\-based optimizers like Muon and SOAP\. If weights and optimizer states are evenly distributed among DP ranks, updates cannot be calculated based on the shard available on each rank\. Additional communication is needed to collect tensors for calculating the full update\. To overcome this, we developed a layer\-wise distributed optimizer, in which parameters of different layers are distributed to different DP ranks\. Each GPU has full layers’ worth of parameters so that the preconditioner can be calculated222[https://docs\.nvidia\.com/nemo/emerging\-optimizers/latest/primer/layerwise\-distributed\-optimizer\.html](https://docs.nvidia.com/nemo/emerging-optimizers/latest/primer/layerwise-distributed-optimizer.html)\.

## 4Background and Related Work

### 4\.1Shampoo and related optimizers

Shampoo\[gupta2018shampoo\]is an optimization algorithm that bridges the gap between first\-order methods like Adam and full\-matrix methods like Natural Gradient Descent and Adagrad\. At its core, Shampoo maintains a running estimate of the all\-dimensional correlations of the parameter gradients, but it does so efficiently by using a Kronecker\-factored approximation of the Fisher information or the empirical covariance matrix\. By treating gradients as tensors rather than flattened vectors, it preserves the spatial or structural correlations inherent in neural network layers\. Building upon this foundational work,\[anil2020scalable\]develops techniques to scale Shampoo to larger models, including “blocking" or restricting preconditioning to subsets of matrices and coupled iterations\[anil2020scalable,modoranu2026dash\]for root\-inverse computation that avoid eigensolvers\. These optimizations effectively mitigate the memory and computational bottlenecks traditionally associated with Kronecker\-factored second\-order methods\. By amortizing the cost of the preconditioner update across multiple steps and leveraging hardware\-accelerated matrix operations, these variants achieve a convergence rate superior to first\-order optimizers like Adam, while maintaining a comparable per\-step wall\-clock time\.

Modern variants of Shampoo have been proposed, such as Eigen\-decomposed Shampoo\[eschenhagen2025purifying\], KL\-Shampoo\[lin2025understanding\], SOAP\[vyas2024soap\], and DASH\[modoranu2026dash\]\. These algorithms improve the estimation of Shampoo’s preconditioner and thus improve convergence and stability\.

Spectral descent was first developed in\[carlson2015preconditioned,carlson2015stochastica,carlson2015stochasticb,flynn2017duality\], and a rigorous framework for spectral descent for all layer types in deep learning was established in\[pethick2025training\]\. In another line of work, Bernsteinet al\.\[bernstein2024modular,bernstein2024old,large2024scalable\]have shown that steepest descent under different norm constraints on parameters provides a picture that unifies element\-wise optimizers such as AdamW and matrix\-based optimizers such as Muon\. Several variations of Muon have recently been developed, such as Muon\+\[zhang2026muon\+\], Muon\+\+\[zhao2026towards\], AdaMuon\[si2025adamuon\], NorMuon\[li2025normuon\], Muon\-Hyperball\[wen2025hyperball\_part1\], Muon on the spectral sphere\[xie2026controlled\], NAMO\[zhang2026adam\], SPEL\[yang2026manifold\], Newton\-Muon\[du2026newton\], and Mousse\[zhang2026mousse\]\.

### 4\.2Systems to enable higher\-order optimizers

Conventional Fully Sharded Data Parallel \(FSDP\) systems restrict tensors to symmetric element\-wise shards\. While this is highly effective for standard element\-wise optimizers like AdamW, it fundamentally conflicts with higher\-order, structure\-aware optimizers\. These optimizers often require access to full 2D parameter matrices \(e\.g\., for computing Muon’s Newton\-Schulz matrix\-sign polynomial iteration\) or specific contiguous blocks\. To address this, veScale\-FSDP\[wang2026vescale\]introducesRaggedShard, a flexible sharding format that supports asymmetric and arbitrary data granularities\. This enables an execution pipeline where full parameters can be dynamically gathered to a root device for complex local computations, and the resulting updates asynchronously scattered back\. To support this flexibility without sacrificing scale, veScale\-FSDP pairsRaggedShardwith a structure\-aware planning algorithm and a zero\-copy distributed buffer, minimizing both communication overhead and memory fragmentation\. Canzona\[wang2026canzona\]builds an asynchronous framework that decouples the logical assignment of optimizer states from the physical distribution of model parameters\. For Data Parallelism \(DP\), it employs a static partitioning strategy that maintains matrix atomicity for the optimizer while preventing load imbalance\. For Tensor Parallelism \(TP\), Canzona utilizes an asynchronous pipeline with micro\-group scheduling to batch fragmented updates and hide matrix reconstruction overheads\. This approach masks the communication costs of higher\-order optimizers, reducing the latency of each step while preserving the throughput of established parallel architectures\.

## 5Pretraining Experiments with Muon and SOAP

Next, we describe our setup and follow up with experiments with Muon, SOAP and experiments comparing Muon and SOAP in the large batch regime on dense and MoE models\.

### 5\.1Model and Dataset

##### Model Architecture

We run all of our optimizer experiments on 8B dense GPT models\[brown2020language\]and 3B active/30B transformer\-only MoEs\[yang2025qwen3\], and finally scale up to 8B active/72B hybrid Mamba\-Transformer MoEs\[blakeman2025nemotron\]\.

##### Dataset

We use 1T\- and 3T\-token subsets of the full Nemotron\-3 dataset\[blakeman2025nemotron\]\.

##### Architecture Configurations

The core architectural hyperparameters for our dense and sparse models are detailed in Table[1](https://arxiv.org/html/2607.20548#S5.T1)\. For the attention mechanism, all models utilize a head dimension of 128 and Grouped\-Query Attention \(GQA\) to improve inference efficiency\.

Table 1:Core architectural hyperparameters for the evaluated models\. For MoE models, the FFN size denotes the hidden size of an individual expert\. All models use 32 heads with head dimension 128\.Table 2:MoE and routing configurations for the Mixture of Experts \(MoE\) models\.
##### Mixture of Experts \(MoE\) Configuration

For our sparse models, we employ a load\-balancing strategy utilizing a sigmoid score function combined with sequence\-level auxiliary loss \(seq\_aux\_loss\)\. By applying an auxiliary loss coefficient of1×10−41\\times 10^\{\-4\}, we effectively mitigate routing collapse—penalizing unbalanced token assignments and ensuring that active experts receive a relatively uniform distribution of tokens across the sequence\. Our MoE variants also leverage shared experts with overlapping computation\. For instance, the 8B Hybrid MoE scales up to 512 total experts with a top\-6 routing sparsity granularity and an additional shared expert intermediate size of 5376\.

Table 3:Learning rate schedules for the evaluated models\. All models utilize a global batch size of 3072 \(micro\-batch size of 1\), a constant weight decay of 0\.1, and a Warmup\-Stable\-Decay \(WSD\) schedule with a minus\-square\-root \(minus\_sqrt\) decay style\.
##### Optimization and Learning Rate Schedule

To train our models, the baselines use a standardized global batch size of 3072 samples, with a sequence length of 8192 tokens each \(leading to a global batch size of 25M tokens\) across all architectures\. We apply a uniform weight decay of 0\.1\. For learning rate progression, we adopt a Warmup\-Stable\-Decay \(WSD\) scheduling strategy\[hagele2024scaling\]\. Following a linear warmup phase, the learning rate is held constant during the stable phase before transitioning into an annealing phase\. Specifically, we utilize a minus\-square\-root \(minus\_sqrt\) decay style for the final WSD decay\.

### 5\.2Optimizer Comparison Protocol

A critical challenge in comparing optimizers is ensuring a fair hyperparameter configuration\. To address this, we utilize anupdate RMS matchingframework pioneered by Kimi\-Moonshot\[liu2025muon\]\. This technique allows us to transfer learning rates \(LRs\) effectively between optimizers by ensuring the root mean square \(RMS\) norm of the parameter updates is consistent across methods, reducing tuning compute\.

Under this RMS\-matching framework, SOAP and AdamW naturally have identical update RMS norm since rotation matrices do not change norm:

‖𝚫​𝐖‖F2\\displaystyle\\\|\\mathbf\{\\Delta W\}\\\|\_\{F\}^\{2\}=Tr​\(𝚫​𝐖T​𝚫​𝐖\)\\displaystyle=\\text\{Tr\}\\left\(\\mathbf\{\\Delta W\}^\{T\}\\mathbf\{\\Delta W\}\\right\)=Tr​\(\(𝐐L​𝐍𝐐RT\)T​\(𝐐L​𝐍𝐐RT\)\)\\displaystyle=\\text\{Tr\}\\left\(\(\\mathbf\{Q\}\_\{L\}\\mathbf\{N\}\\mathbf\{Q\}\_\{R\}^\{T\}\)^\{T\}\(\\mathbf\{Q\}\_\{L\}\\mathbf\{N\}\\mathbf\{Q\}\_\{R\}^\{T\}\)\\right\)=Tr​\(\(𝐐RT​𝐐R\)​𝐍T​\(𝐐LT​𝐐L\)​𝐍\)\\displaystyle=\\text\{Tr\}\\left\(\(\\mathbf\{Q\}\_\{R\}^\{T\}\\mathbf\{Q\}\_\{R\}\)\\mathbf\{N\}^\{T\}\(\\mathbf\{Q\}\_\{L\}^\{T\}\\mathbf\{Q\}\_\{L\}\)\\mathbf\{N\}\\right\)Since𝐐L\\mathbf\{Q\}\_\{L\}and𝐐R\\mathbf\{Q\}\_\{R\}are orthonormal,𝐐LT​𝐐L\\displaystyle\\mathbf\{Q\}\_\{L\}^\{T\}\\mathbf\{Q\}\_\{L\}=𝐈,𝐐RT​𝐐R=𝐈,\\displaystyle=\\mathbf\{I\},\\mathbf\{Q\}\_\{R\}^\{T\}\\mathbf\{Q\}\_\{R\}=\\mathbf\{I\},‖𝚫​𝐖‖F2\\displaystyle\\\|\\mathbf\{\\Delta W\}\\\|\_\{F\}^\{2\}=‖𝐍‖F2\\displaystyle=\\\|\\mathbf\{N\}\\\|\_\{F\}^\{2\}
We explicitly apply Kimi’s update RMS matching framework to Muon \(Algorithm[1](https://arxiv.org/html/2607.20548#alg1)\), which gives us a correction factor that depends on the EMA damping factor of momentum\[kexuefm\-11307,li2024stochastic,kosson2023rotational\]:

1−β11\+β1≈0\.2\\sqrt\{\\frac\{1\-\\beta\_\{1\}\}\{1\+\\beta\_\{1\}\}\}\\approx 0\.2Notably, we observed that using Nesterov momentum\[sutskever2013importance,dozat2016incorporating\]did not improve convergence and skipped using it for both Muon and SOAP, and using simple exponential moving averages \(EMA\) for momentum\.

### 5\.3Muon Scales Better than AdamW at Large Batch Sizes for Pretraining

Prior research indicates that Muon and its variants significantly outperform AdamW as the global batch size scales\[pethick2025training\]\. While AdamW often hits a "critical batch size"\[mccandlish2018empirical\]beyond which increasing parallelism yields diminishing returns, token inefficiency, or instability, Muon’s orthogonalization and spectral descent allow it to maintain convergence at much larger batch sizes\.

![Refer to caption](https://arxiv.org/html/2607.20548v1/nanov3_3T_adam_baseline.png)Figure 1:Loss curves of Muon at large batches \(2x and 3x\) with NVIDIA\-Nemotron\-Nano\-V3\-30B\-A3B on a 3T\-token horizon compared against the original AdamW baseline of\[blakeman2025nemotron\]\. The inset shows the relative loss gaps as a function of tokens against the AdamW baseline\. Larger batches typically underperform and eventually exceed lower batch sizes, especially as the token horizon scales\.
![Refer to caption](https://arxiv.org/html/2607.20548v1/latent_8b_bf16_vs_sqrt2lr.png)Figure 2:Hybrid Mamba\-MoE 72B\-A8B trained on 1T tokens with Muon at 2x the batch size of the AdamW baseline \(25M tokens\)\. Inset shows that Muon at larger batches haslowerloss and no loss spikes relative to the smaller batch AdamW baseline\.

Our experiments confirm that Muon provides better stability and lower loss compared to AdamW in this setting\. At a baseline \(1×1\\times\) batch size, Muon achieves lower training loss and exhibits fewer spikes, Fig\.[2](https://arxiv.org/html/2607.20548#S5.F2)\. We have empirically found training loss to correlate to held\-out validation loss on our dataset and we include evaluation on downstream tasks for a subset of models in Table[5](https://arxiv.org/html/2607.20548#S5.T5)\. Crucially, as we scale the global batch size, AdamW fails to scale effectively, rapidly encountering instabilities and diminishing returns\. In contrast, Muon maintains robust convergence characteristics, effectively utilizing the increased data parallelism to speed up training\.

Algorithm 1Muon:MomentUmOrthogonalized byNewton\-Schulz with Weight Decay1:Hyperparameters:

ηt,β1,λt,n,\{\(ai,bi,ci\)\}i=1n,ϵ\\eta\_\{t\},\\,\\beta\_\{1\},\\,\\lambda\_\{t\},\\,n,\\,\\\{\(a\_\{i\},b\_\{i\},c\_\{i\}\)\\\}\_\{i=1\}^\{n\},\\,\\epsilon
2:Optimizer States:

𝐌t∈ℝin×out​\(Momentum buffer\)\\mathbf\{M\}\_\{t\}\\in\\mathbb\{R\}^\{\\text\{in\}\\times\\text\{out\}\}\\text\{ \(Momentum buffer\)\}
3:

𝐆t=∇𝐖L​\(𝐖t−1\)\\mathbf\{G\}\_\{t\}=\\nabla\_\{\\mathbf\{W\}\}L\(\\mathbf\{W\}\_\{t\-1\}\)Update Momentum Buffer with EMA:

4:

𝐌t←β1​𝐌t−1\+\(1−β1\)​𝐆t\\mathbf\{M\}\_\{t\}\\leftarrow\\beta\_\{1\}\\mathbf\{M\}\_\{t\-1\}\+\(1\-\\beta\_\{1\}\)\\mathbf\{G\}\_\{t\}Newton\-Schulz Orthogonalization:

5:

𝐗0←𝐌t/max⁡\(‖𝐌t‖F,ϵ\)\\mathbf\{X\}\_\{0\}\\leftarrow\\mathbf\{M\}\_\{t\}/\\max\(\\\|\\mathbf\{M\}\_\{t\}\\\|\_\{F\},\\epsilon\)⊳\\trianglerightInitialize with normalized momentum

6:for

i=1i=1to

nndo

7:

𝐀←𝐗i−1​𝐗i−1T\\mathbf\{A\}\\leftarrow\\mathbf\{X\}\_\{i\-1\}\\mathbf\{X\}\_\{i\-1\}^\{T\}
8:

𝐁←bi​𝐀\+ci​𝐀2\\mathbf\{B\}\\leftarrow b\_\{i\}\\mathbf\{A\}\+c\_\{i\}\\mathbf\{A\}^\{2\}
9:

𝐗i←ai​𝐗i−1\+𝐁𝐗i−1\\mathbf\{X\}\_\{i\}\\leftarrow a\_\{i\}\\mathbf\{X\}\_\{i\-1\}\+\\mathbf\{B\}\\mathbf\{X\}\_\{i\-1\}
10:endfor

11:

𝐎t←𝐗n\\mathbf\{O\}\_\{t\}\\leftarrow\\mathbf\{X\}\_\{n\}⊳\\trianglerightResulting orthogonalized update directionPost\-processing update direction normalization:

12:

𝚫​𝐖←𝐎t⋅max⁡\(in,out\)\\mathbf\{\\Delta W\}\\leftarrow\\mathbf\{O\}\_\{t\}\\cdot\\sqrt\{\\max\(\\text\{in\},\\text\{out\}\)\}⊳\\trianglerightKimi scaling for update RMS matchingDecoupled weight decay and parameter update:

13:

𝐖t←\(1−ηt​λt\)​𝐖t−1−ηt​𝚫​𝐖\\mathbf\{W\}\_\{t\}\\leftarrow\(1\-\\eta\_\{t\}\\lambda\_\{t\}\)\\mathbf\{W\}\_\{t\-1\}\-\\eta\_\{t\}\\mathbf\{\\Delta W\}

Table 4:Muon hyperparameters used in Algorithm[1](https://arxiv.org/html/2607.20548#alg1)\.We hypothesize that at extreme batch sizes, the gradient noise decreases, which can cause AdamW’s element\-wise second\-moment estimator \(𝐯t\\mathbf\{v\}\_\{t\}\) to become poorly calibrated or overly confident, leading to suboptimal step sizes\. In contrast to AdamW, Muon’s spectral preconditioning \(i\.e\. orthogonalization of the momentum matrix\) structurally normalizes the update directionsindependentof coordinate\-wise scaling, making it resilient to the variance fluctuations of large\-batch scaling\.

We also find in Fig\.[2](https://arxiv.org/html/2607.20548#S5.F2)that implementing abatch size ramp\-upphase early in training allows Muon to significantly outperform AdamW on both final training loss and downstream evaluation metrics, while using a larger batch size for the majority of training\. We note that the rampup schedule was used only with the larger batch pretraining experiments for both optimizers\. Following\[li2025minimax\], we adopt a ramp\-up schedule governed by a few simple rules\. We treat the language modeling loss as the signal dictating the optimal batch size\. Starting from a batch size of 256 samples \(2M tokens\), we double everySSsteps over a fixed ramp\-up horizon ofNNtotal steps, irrespective of the token horizon, until the target batch size is reached, Fig\.[3](https://arxiv.org/html/2607.20548#S5.F3)\. The learning rate warmup is tied directly to the batch size ramp\-up to approximately maintain the square\-root scaling ratio throughout, i\.e\.,η∝B\\eta\\propto\\sqrt\{B\}, as motivated by large\-batch training theory\[You2020Large,hoffer2017train,you2017large,krizhevsky2014one\]\. Specifically, the learning rate at each stage is tuned to preserve the variance scale of the parameter update as the batch size changes\. We note that we did not tune the batch size ramp\-up schedule as a function of model scale\[li2025minimax,mccandlish2018empirical\], where the critical batch size typically depends on the training loss; hence, larger models are expected to reach higher critical batch sizes earlier than smaller models\.

We evaluate on two architectures: NVIDIA\-Nemotron\-Nano\-V3, a hybrid attention\-mamba MoE model with 30B total and 3B active parameters trained on 3T tokens, and a hybrid LatentMoE model with Multi\-Token Prediction \(MTP\) layers\[gloeckle2024better,liu2024deepseek\]with 72B total and 8B active parameters trained on a 1T token horizon\. Table[5](https://arxiv.org/html/2607.20548#S5.T5)reports downstream evaluation results for these models\. For both architectures, Muon consistently matches or outperforms AdamW at baseline batch size and continues to improve as batch size scales, whereas AdamW does not \(Fig\.[10](https://arxiv.org/html/2607.20548#A3.F10)\)\. The gains are most pronounced in coding and commonsense reasoning\. We also note that while ablating the hybrid Mamba\-MoE architecture of Table[5](https://arxiv.org/html/2607.20548#S5.T5), we found an important architectural sensitivity: Muon is best applied selectively to dense linear projections, as we found that falling back to AdamW for Mamba’sconv1Dweights\[dao2024transformers\]yields better convergence\. These results validate that Muon’s large\-batch efficiency transfers to MoE architectures up to the≈5​e22\\approx 5e^\{22\}pretraining FLOPs scale\.

We further validate that these results hold with MXFP8 pretraining in Appendix Sec\.[B](https://arxiv.org/html/2607.20548#A2)\.

Table 5:Benchmark results across MMLU, Coding, Math, and Commonsense Reasoning for Muon and AdamW optimizers at varying batch sizes \(n×n\\timesGBS\) for the Nano\-V3 and 8B Hybrid MTP models\. 1×\\timesGBS is equivalent to 25xM tokens per batch\.![Refer to caption](https://arxiv.org/html/2607.20548v1/batch_schedule_50M.png)
![Refer to caption](https://arxiv.org/html/2607.20548v1/batch_schedule_75M.png)

Figure 3:The batch size ramp\-up schedules \(in blue\) used in pretraining experiments for final batch sizes of 50M and 75M tokens \(in red\)\.
### 5\.4Training Stability with SOAP

We observed significant empirical hurdles when scaling SOAP to large\-batch, large parameter count regimes\. Here we document the instability with SOAP before moving on to experiments that compare them\.

#### 5\.4\.1Diagnosing the "Slingshot" Instability

The reference SOAP implementation333[https://github\.com/nikhilvyas/SOAP](https://github.com/nikhilvyas/SOAP)mitigates the computational expense of QR decomposition by refreshing the eigenbasis infrequently \(e\.g\., every 10 steps\) but also simultaneously omits the current step’s gradient from the eigenbasis recalculation\. While harmless for smaller models and higher eigenbasis recomputation intervals, our large\-batch pretraining experiments reveal that this configuration induces severe early\-training instability\. As illustrated in Fig\.[4](https://arxiv.org/html/2607.20548#S5.F4), this manifests as oscillating weight gradient norms followed closely by spikes in the language modeling loss\.

We attribute this to a "stale preconditioner" effect\. Because the loss landscape changes rapidly during the early phases of large\-batch training, applying a preconditioning matrix derived from outdated gradient statistics causes the optimization trajectory to "slingshot"\. For the smaller 621M model, these oscillations are eventually damped, and the model converges, Fig\.[4](https://arxiv.org/html/2607.20548#S5.F4)\. However, as we scale up to an 8B dense model, this lag becomes catastrophic, leading to training divergence and significantly poorer quality compared to AdamW, Fig\.[5](https://arxiv.org/html/2607.20548#S5.F5)\.

Algorithm 2KL\-SOAP used in experiments, based on\[vyas2024soap,lin2025understanding\]1:Hyperparameters:

βkron,β1,β2,p,ηt,ϵ,F,λt\\beta\_\{\\text\{kron\}\},\\,\\beta\_\{1\},\\,\\beta\_\{2\},\\,p,\\,\\eta\_\{t\},\\,\\epsilon,\\,F,\\,\\lambda\_\{t\}
2:Optimizer States:

𝐦t,𝐯t∈ℝin×out;𝐋t,𝐐𝐋∈ℝin×in;𝐑t,𝐐𝐑∈ℝout×out\\mathbf\{m\}\_\{t\},\\,\\mathbf\{v\}\_\{t\}\\in\\mathbb\{R\}^\{\\text\{in\}\\times\\text\{out\}\};\\,\\mathbf\{L\}\_\{t\},\\mathbf\{Q\_\{L\}\}\\in\\mathbb\{R\}^\{\\text\{in\}\\times\\text\{in\}\};\\,\\mathbf\{R\}\_\{t\},\\mathbf\{Q\_\{R\}\}\\in\\mathbb\{R\}^\{\\text\{out\}\\times\\text\{out\}\}
3:

𝐆t=∇𝐖L​\(𝐖t−1\)\\mathbf\{G\}\_\{t\}=\\nabla\_\{\\mathbf\{W\}\}L\(\\mathbf\{W\}\_\{t\-1\}\)Accumulate Kronecker Factors \(KL\-Shampoo\):

4:

λL=diag​\(𝐐LT​𝐋t−1​𝐐L\)\\mathbf\{\\lambda\}\_\{L\}=\\text\{diag\}\(\\mathbf\{Q\}\_\{L\}^\{T\}\\mathbf\{L\}\_\{t\-1\}\\mathbf\{Q\}\_\{L\}\)⊳\\trianglerightApproximate eigenvalues

5:

λR=diag​\(𝐐RT​𝐑t−1​𝐐R\)\\mathbf\{\\lambda\}\_\{R\}=\\text\{diag\}\(\\mathbf\{Q\}\_\{R\}^\{T\}\\mathbf\{R\}\_\{t\-1\}\\mathbf\{Q\}\_\{R\}\)
6:

𝐋t←\(1−βkron\)​𝐋t−1\+βkronout​𝐆t​\(λR\+ϵ\)−p​𝐆tT\\mathbf\{L\}\_\{t\}\\leftarrow\(1\-\\beta\_\{\\text\{kron\}\}\)\\,\\mathbf\{L\}\_\{t\-1\}\+\\frac\{\\beta\_\{\\text\{kron\}\}\}\{\\text\{out\}\}\\,\\mathbf\{G\}\_\{t\}\\,\(\\lambda\_\{R\}\+\\epsilon\)^\{\-p\}\\,\\mathbf\{G\}\_\{t\}^\{T\}
7:

𝐑t←\(1−βkron\)​𝐑t−1\+βkronin​𝐆tT​\(λL\+ϵ\)−p​𝐆t\\mathbf\{R\}\_\{t\}\\leftarrow\(1\-\\beta\_\{\\text\{kron\}\}\)\\,\\mathbf\{R\}\_\{t\-1\}\+\\frac\{\\beta\_\{\\text\{kron\}\}\}\{\\text\{in\}\}\\,\\mathbf\{G\}\_\{t\}^\{T\}\\,\(\\lambda\_\{L\}\+\\epsilon\)^\{\-p\}\\,\\mathbf\{G\}\_\{t\}Eigenbasis Re\-computation

8:if

tmodF=0t\\bmod F=0then

9:

𝐦^t=𝐐L​𝐦t​𝐐RT\\hat\{\\mathbf\{m\}\}\_\{t\}=\\mathbf\{Q\}\_\{L\}\\,\\mathbf\{m\}\_\{t\}\\,\\mathbf\{Q\}\_\{R\}^\{T\}
10:

𝐐L​QR=​\(𝐋t​𝐐L\)\.Q\\mathbf\{Q\}\_\{L\}\\,\{\}\_\{=\}\\texttt\{QR\}\(\\mathbf\{L\}\_\{t\}\\,\\mathbf\{Q\}\_\{L\}\)\.Q
11:

𝐐R​QR=​\(𝐑t​𝐐R\)\.Q\\mathbf\{Q\}\_\{R\}\\,\{\}\_\{=\}\\texttt\{QR\}\(\\mathbf\{R\}\_\{t\}\\,\\mathbf\{Q\}\_\{R\}\)\.Q
12:

𝐦t=𝐐LT​𝐦^t​𝐐R\\mathbf\{m\}\_\{t\}=\\mathbf\{Q\}\_\{L\}^\{T\}\\,\\hat\{\\mathbf\{m\}\}\_\{t\}\\,\\mathbf\{Q\}\_\{R\}
13:endifRotate gradient into eigenbasis

14:

𝐆t′←𝐐^LT​𝐆t​𝐐^R\\mathbf\{G\}\_\{t\}^\{\\prime\}\\leftarrow\\hat\{\\mathbf\{Q\}\}\_\{L\}^\{T\}\\,\\mathbf\{G\}\_\{t\}\\,\\hat\{\\mathbf\{Q\}\}\_\{R\}Calculate Adam update

15:

𝐦t←\(1−β1\)​𝐦t−1\+β1​𝐆′\\mathbf\{m\}\_\{t\}\\leftarrow\(1\-\\beta\_\{1\}\)\\,\\mathbf\{m\}\_\{t\-1\}\+\\beta\_\{1\}\\,\\mathbf\{G\}^\{\\prime\}
16:

𝐯t←\(1−β2\)​𝐯t−1\+β2​𝐆t′⊙𝐆t′\\mathbf\{v\}\_\{t\}\\leftarrow\(1\-\\beta\_\{2\}\)\\,\\mathbf\{v\}\_\{t\-1\}\+\\beta\_\{2\}\\,\\mathbf\{G\}\_\{t\}^\{\\prime\}\\odot\\mathbf\{G\}\_\{t\}^\{\\prime\}
17:

𝐍←𝐦t𝐯t\+ϵ\\mathbf\{N\}\\leftarrow\\dfrac\{\\mathbf\{m\}\_\{t\}\}\{\\sqrt\{\\mathbf\{v\}\_\{t\}\}\+\\epsilon\}Rotate update out of eigenbasis

18:

𝚫​𝐖=𝐐L​𝐍​𝐐RT\\mathbf\{\\Delta W\}=\\mathbf\{Q\}\_\{L\}\\,\\mathbf\{N\}\\,\\mathbf\{Q\}\_\{R\}^\{T\}Weight decay and update

19:

𝐖t←\(1−ηt​λt\)​𝐖t−1−ηt​𝚫​𝐖\\mathbf\{W\}\_\{t\}\\leftarrow\(1\-\\eta\_\{t\}\\lambda\_\{t\}\)\\mathbf\{W\}\_\{t\-1\}\-\\eta\_\{t\}\\mathbf\{\\Delta W\}

Table 6:KL\-SOAP hyperparameters used in Algorithm[2](https://arxiv.org/html/2607.20548#alg2)\.![Refer to caption](https://arxiv.org/html/2607.20548v1/dist_soap_embed_bs_1536_lr_3.0e-3_loss_and_grad_norm.png)Figure 4:SOAP with a preconditioner that is computed based on stale gradient statistics leads to an oscillatory instability of loss spikes preceded by gradient spikes\. The experiment was performed on a dense 621M pure transformer model pretrained with a global batch size of 12\.5M tokens\.![Refer to caption](https://arxiv.org/html/2607.20548v1/8b_gpt_academicsoap_vs_regularsoap.png)Figure 5:The instability of Fig\.[4](https://arxiv.org/html/2607.20548#S5.F4)at a larger model size leads to training divergence in SOAP \(blue\) and large loss gaps relative to a baseline that uses current gradient statistics to compute preconditioners \(in orange\)\.
#### 5\.4\.2Mitigating Staleness: Real\-Time Eigenbasis Updates

We find that simply increasing the eigenbasis update frequency to every step is insufficient if the current iteration’s gradient is excluded\. It is the combination of per\-step QR updatesandthe inclusion of the current step’s gradient in the eigenbasis calculation that fully resolves the issue\. When modified to use real\-time gradient statistics, SOAP eliminates the loss spikes and achieves an LM loss on par with Muon, Fig\.[7](https://arxiv.org/html/2607.20548#S5.F7)\. We note that this implementation was present with SOAP and is not present in\[anil2020scalable,lin2025understanding\]and other works\.

#### 5\.4\.3Gradient Covariance Estimation via Kullback–Leibler Minimization

Even with real\-time updates, the accumulation of Kronecker factor matrices \(G​GTGG^\{T\}andGT​GG^\{T\}G\) can suffer from instabilities at scale\. As described in\[lin2025understanding\], we find that replacing the standard accumulation with the KL divergence\-regularized covariance estimation algorithm provides a critical secondary layer of stability\. As shown in Fig\.[6](https://arxiv.org/html/2607.20548#S5.F6), integrating KL\-Shampoo’s update rule further suppresses loss spikes and yields a lower final training loss on the Qwen\-3\-30B\-A3B architecture\.

![Refer to caption](https://arxiv.org/html/2607.20548v1/qwen3b_30b_1T_soap_vs_soap_nokl.png)Figure 6:The inclusion of KL\-divergence\-based gradient covariance estimation\[lin2025understanding\]in SOAP stabilizes training and results in lower training loss and fewer loss spikes\. Experiments were performed on a Qwen\-3\-30B\-A3B architecture\.This empirical improvement comes from better estimation of the gradient covariance, and the stability benefit comes from the spectral properties of the update\. As detailed in Appendix[A\.3](https://arxiv.org/html/2607.20548#A1.SS3), KL\-Shampoo reduces the condition number of the Kronecker factor matrices to the square root of the standard formulation:

κ​\(𝐒KL\-Shampoo\)=κ​\(𝐒Shampoo\)\.\\kappa\(\\mathbf\{S\}\_\{\\text\{KL\-Shampoo\}\}\)=\\sqrt\{\\kappa\(\\mathbf\{S\}\_\{\\text\{Shampoo\}\}\)\}\.\(4\)We hypothesize that this property improves the numerical stability of the eigenbasis recomputation, whether using QR or eigendecomposition\-based eigensolvers, safeguarding the training dynamics against floating\-point noise and preventing subsequent gradient explosions or loss spikes\.

### 5\.5Comparing Muon and SOAP

To evaluate the relative performance of Muon and SOAP, we use a controlled comparison\. In both configurations, we employ the simplest "blocking" method for tensor parallelism support\. Further, to ensure an apples\-to\-apples structural comparison, we disable Query\-Key\-Value \(QKV\) splitting in Muon\. By treating the fused QKV projection as a single matrix, matching the default operational mode of our implementation of SOAP, we isolate the algorithmic differences in preconditioning and orthogonalization from structural choices\.

We do a comparative analysis for MoEs using the Qwen\-3\-30B\-A3B architecture\. Fig\.[7](https://arxiv.org/html/2607.20548#S5.F7)shows the results of a large batch study comparing Muon and SOAP\. We also include a version of Muon we term Momentum Orthogonalized by Polar \(MOP\), which is Muon that replaces Newton\-Schulz orthogonalization with exact polar decomposition via SVD444This was first developed in[https://github\.com/KellerJordan/cifar10\-airbench/tree/master](https://github.com/KellerJordan/cifar10-airbench/tree/master)\. Fig\.[8](https://arxiv.org/html/2607.20548#S5.F8)presents the training trajectories for these optimizers under the same controlled tensor parallelism and non\-split QKV conditions\. Throughout the majority of the pretraining run, KL\-SOAP maintains a consistent, albeit slight, edge over Muon in the CE loss\.

![Refer to caption](https://arxiv.org/html/2607.20548v1/qwen3b_30b_1T_muon_vs_soap_by_gbs_1x.png)

![Refer to caption](https://arxiv.org/html/2607.20548v1/qwen3b_30b_1T_muon_vs_soap_by_gbs_2x.png)

![Refer to caption](https://arxiv.org/html/2607.20548v1/qwen3b_30b_1T_muon_vs_soap_by_gbs_4x.png)

Figure 7:Large\-batch comparative study of Muon and SOAP for the Qwen3\-30B\-A3B architecture \(Top: 1x, Middle: 2x, Bottom: 4x GBS\)\. Inset shows relative loss differences\. Note that these experiments did not use any batch\-size ramp\-up\.![Refer to caption](https://arxiv.org/html/2607.20548v1/qwen3_30b_mop_vs_soap_vs_muon.png)Figure 8:Qwen\-3\-30B\-A3B architecture pretrained with Muon, SOAP, and MOP with update\-RMS matching\. Both SOAP and MOP have slightly lower losses than Muon\. Inset shows relative loss differences to the Muon baseline\. The global batch size was fixed to 24M tokens for this experiment\.##### Tuning the "Epsilon"

Both SOAP and Muon expose anϵ\\epsilonhyperparameter, and recent work has shown that this seemingly small numerical constant can affect optimizer stability, transfer across scale, and the predictability of training dynamics\[everett2024scaling,wortsman2023small,yang2023tensor,epsilon\_nuisance\]\. In SOAP,ϵ\\epsilonplays a role analogous to theϵ\\epsilonterm in AdamW: it floors the denominator of the second\-moment estimate after the momentum has been rotated into the preconditioned eigenbasis and the eigenvalues in the KL\-Shampoo accumulation step\. In addition, because SOAP uses Shampoo\-style kronecker factors where the preconditioning is effectively applying inverse square powers to the kronecker factor eigenvalues\[vyas2024soap\], theϵ\\epsilonused for the Adam\-like computation in the preconditioner basis can be interpreted as setting a soft floor on smallest effective eigenvalue of the Shampoo preconditioner\. In Muon,ϵ\\epsilonhas a different role and is used directly in the normalization step before Newton\-Schulz orthogonalization, where the momentum matrix is scaled bymax⁡\(‖𝐌t‖F,ϵ\)\\max\(\\\|\\mathbf\{M\}\_\{t\}\\\|\_\{F\},\\epsilon\)\. Thus, Muon’sϵ\\epsilonlower\-bounds the Frobenius norm used for matrix normalization rather than flooring an element\-wise second\-moment estimate\. We therefore treatϵ\\epsilonas an optimizer\-specific hyperparameter: for SOAP it controls the numerical floor in preconditioned adaptive moments and factor inversions, while for Muon it controls the smallest normalization scale entering the Newton\-Schulz iteration\.

##### Limitations of our work

In this report, we do not tuneϵ\\epsilonsystematically\. We useϵ=10−8\\epsilon=10^\{\-8\}for SOAP and AdamW, andϵ=10−7\\epsilon=10^\{\-7\}for Muon\. Understanding how these values should scale with model size, optimizer parameterization, precision format, and training horizon is an important direction for subsequent work\.

## 6Implementation Details: Layer\-Wise Distributed Optimizer

While higher\-order optimizers like SOAP and Muon demonstrate better convergence and stability at larger batch sizes, realizing their empirical benefits in production environments requires the underlying training system to supply the optimizers with unfragmented 2D weight matrices without incurring too much exposed communication overhead\. This conflicts with the tensor sharding inherent to standard distributed training approaches \(e\.g\., ZeRO\[rajbhandari2020zero\]or FSDP\[zhao2023pytorch\]\)\. To overcome this system bottleneck, we developed a strategy native to the Megatron\-LM framework\[shoeybi2019megatron\]\.

This implementation outlines a distributed optimizer strategy designed to support matrix\-based optimizers by keeping parameter matrices fully intact, while still balancing memory and hiding communication bottlenecks\. Here we describe the main features of the implementation, which is available in Megatron\-LM555[https://github\.com/NVIDIA/Megatron\-LM/blob/main/megatron/core/optimizer/layer\_wise\_optimizer\.py](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/optimizer/layer_wise_optimizer.py)\.

1. 1\.Load Balancing: Instead of flattening and slicing individual tensors, entire parameter matrices are sorted by size and distributed across GPUs in a round robin sequence\. This ensures every GPU has a roughly equal memory load without fragmenting the matrices\.
2. 2\.Parameter Update: Each GPU acts independently to update only its assigned parameters\. It applies the optimizer math \(e\.g\., using Muon or SOAP for 2D matrices and AdamW for standard vectors\) to its specific layers, and then flattens the results into a buffer to be gathered in the next step’s forward pass\.
3. 3\.Overlapped Parameter All\-Gather: Parameter synchronization is chunked into sequential communication buckets that logically align with the model’s execution order \(leveraging Megatron’s DDP bucket infrastructure\)\. By grouping layers into these buckets, the system bounds message sizes and enables pipelining\. Instead of a single, monolithic blocking synchronization, the system issues an asynchronous all\-gather collective for each bucket\. As the forward pass computes activations using the current bucket’s weights, the GPUs simultaneously execute the all\-gather to fetch the updated parameter matrices for the upcoming bucket from their respective owners, effectively hiding the network latency behind the computation\. Keeping parameter matrices intact results in slightly uneven total parameter counts across GPUs, so standard uniform all\-gather collectives cannot be used\. Instead, the asynchronous synchronization employs a variable\-sized allgather\-V666[https://www\.mpich\.org/static/docs/v3\.2/www3/MPI\_Allgatherv\.html](https://www.mpich.org/static/docs/v3.2/www3/MPI_Allgatherv.html)collective\. This allows each GPU to synchronize its uniquely sized, contiguous buffer of updated layers to the rest of the ranks natively, avoiding the need for memory\-wasting padding or tensor fragmentation\.

## 7Conclusions and Future Work

In this report, we conducted large\-scale pretraining studies of the Muon and SOAP optimizers, identifying key best practices for maintaining stability at large batch sizes\. We have open\-sourced our implementations in Megatron\-LM andEmerging\-Optimizers\. Our findings demonstrate that both SOAP and Muon consistently outperform AdamW, scaling effectively to significantly larger batch sizes during pretraining\. Furthermore, the KL\-SOAP variant emerged as the most effective approach overall; therefore, in scenarios where memory footprint is not a limiting factor, we recommend KL\-SOAP over Muon\.

We summarize several algorithmic, system\-level, and empirical directions for advancing and scaling these research directions below\.

##### SOAP

To fully realize the benefits of SOAP at frontier model scales, it is essential to support additional structural and distributed features\. Specifically, extending our layer\-wise distributed optimizer to natively support Tensor Parallelism \(TP\) for SOAP, handling fused tensors, for example by implementing QKV splitting for attention, and further optimizing the throughput of optimizer matrix multiplications such as the KL\-Shampoo eigenbasis updates are areas of potential improvement\.

##### Optimizer\-Driven Aligned Memory Buffer Layouts

Currently, standard Distributed Data Parallel \(DDP\) implementations force parameters into rigid, 1D contiguous memory buffers\. While our layer\-wise algorithm successfully navigates this via bucketed, variable\-sized collectives, the underlying buffer allocation remains largely optimizer\-agnostic\. This creates inherent friction for higher\-order methods like SOAP and Muon, which operate strictly on 2D matrices and structured factorization states\. A direction for future work is to refactor the DDP buffer allocation engine to be natively optimizer\-driven\. Rather than having the DDP module flatten parameters without considering the structure and boundaries of network weights, the system should allow the optimizer to pre\-compute and dictate the layout of the parameter and gradient buffers\. By granting the optimizer control, the physical communication buckets will map to the structural shape of the optimizer computations, reducing communication overhead\. This optimization is being built in the Megatron\-LM framework777[https://github\.com/NVIDIA/Megatron\-LM/pull/4509](https://github.com/NVIDIA/Megatron-LM/pull/4509)and will replace gradient all\-reduces with reduce\-scatters and further optimize memory management and sharding\.

##### Optimizer\-Aware Architectural Co\-Design

We observed that applying orthogonalization to Mamba2’s Conv1D filters degraded accuracy and sometimes caused instability in the form of NaNs, highlighting that matrix\-based optimizers can be geometrically mismatched with certain architectures\. Further investigation is needed to map out how architectural components such as Multi\-head Latent Attention \(MLA\)\[liu2024deepseek\]or low\-rank adapters \(LoRA\) interact with full\-rank preconditioning assumptions, potentially leading to hybrid training recipes that selectively route parameters to AdamW or Muon or a variant based on their tensor shapes and operator functions\[bernstein2024modular\]\. Along these lines, recent work has developed a matrix\-based optimizer for native low\-rank factored training\[janson2026stabilizing\]\.

##### Large Batch Training and Scaling Laws for Batch Size

Recent empirical studies have established scaling laws to predict optimal batch sizes as a function of compute budget and training hyperparameters\[bi2024deepseek,li2025minimax,li2025predictable\], with theoretical foundations linking these scaling laws to the interplay between gradient noise reduction and optimization step efficiency\[islamov2026role\]\. While we do not explicitly derive or measure scaling laws in this work, we instead approach batch size primarily as a systems\-level constraint\. Specifically, we push the global batch size as high as possible, while maintaining quality, with the goal of maximizing GPU utilization by minimizing data\-parallel communication overhead, which is necessary for efficient pretraining at scale\. An optimal approach would combine these two points of view, for example, by using scaling laws to better tune the batch size ramp\-up and learning rate schedule and heuristics for changing the maximum batch size as token horizons vary\.

##### More Accurate Orthogonalization at Larger Scales

Scaling matrix\-based optimizers that perform orthogonalization to larger models and longer token horizons can introduce new numerical challenges that warrant further investigation\. The typical size of matrices increases as we scale models up, and so does the “heavy\-tail" of small singular values of these matrices’ gradients\. Since the optimizer step always occurs in FP32 single precision, this “heavy\-tail" of small singular values and their associated vectors is sensitive to numerical issues and can become increasingly susceptible to floating\-point noise, often being below FP32 machine precision \(2−23≈1\.19×10−72^\{\-23\}\\approx 1\.19\\times 10^\{\-7\}\)\. For Muon’s Newton\-Schulz iterations which occur in BF16, it is coarser \(2−7=0\.00781252^\{\-7\}=0\.0078125\) and small singular values often represent numerical noise rather than informative signal, creating an effective "noise floor" in FP32\. Truncated or regularized orthogonalization approaches and corresponding hardware\-friendly algorithms that implement them are open research directions, with notable recent work measuring spectral scaling laws for these quantities\[magakyan2026spectral\]\.

## References

## Appendix AAppendix

### A\.1An SVD perspective on SOAP’s preconditioner: Relating Muon, SOAP, and Shampoo

As has been studied before, singular value decomposition can give us insight into how Shampoo, SOAP, and Muon are related to each other\[shah2025practical\]\.

Let us examine the SVD of the averaged weight gradient:888An ill\-conditioned gradient will have singular values with a largemax/min\\max/\\minratio:max⁡\(Σ\)min⁡\(Σ\)\>\>1\\dfrac\{\\max\(\\Sigma\)\}\{\\min\(\\Sigma\)\}\>\>1

𝐆t=𝐔⋅𝚺⋅𝐕T\.\\mathbf\{G\}\_\{t\}=\\mathbf\{U\}\\cdot\\mathbf\{\\Sigma\}\\cdot\\mathbf\{V\}^\{T\}\.\(5\)
Now notice \(with EMAs turned off for simplicity\):

- •𝐐L\\mathbf\{Q\}\_\{L\}is the eigenvector matrix of𝐆t​𝐆tT\\mathbf\{G\}\_\{t\}\\,\\mathbf\{G\}\_\{t\}^\{T\}, which is𝐔\\mathbf\{U\}, since 𝐆t​𝐆tT=𝐔​𝚺2​𝐔T\.\\mathbf\{G\}\_\{t\}\\,\\mathbf\{G\}\_\{t\}^\{T\}=\\mathbf\{U\}\\,\\mathbf\{\\Sigma\}^\{2\}\\,\\mathbf\{U\}^\{T\}\.\(6\)
- •𝐐R\\mathbf\{Q\}\_\{R\}is the eigenvector matrix of𝐆tT​𝐆t\\mathbf\{G\}\_\{t\}^\{T\}\\,\\mathbf\{G\}\_\{t\}, which is𝐕\\mathbf\{V\}, since 𝐆tT​𝐆t=𝐕​𝚺2​𝐕T\.\\mathbf\{G\}\_\{t\}^\{T\}\\,\\mathbf\{G\}\_\{t\}=\\mathbf\{V\}\\,\\mathbf\{\\Sigma\}^\{2\}\\,\\mathbf\{V\}^\{T\}\.\(7\)

We have

𝐆tR=𝐐LT​𝐆t​𝐐R=\(𝐔T​𝐔\)⏟=𝐈in×in​𝚺​\(𝐕T​𝐕\)⏟=𝐈out×out=𝚺\\mathbf\{G\}^\{R\}\_\{t\}=\\mathbf\{Q\}\_\{L\}^\{T\}\\,\\mathbf\{G\}\_\{t\}\\,\\mathbf\{Q\}\_\{R\}=\\underbrace\{\\left\(\\mathbf\{U\}^\{T\}\\,\\mathbf\{U\}\\right\)\}\_\{=\\,\\mathbf\{I\}\_\{\\text\{in\}\\times\\text\{in\}\}\}\\,\\mathbf\{\\Sigma\}\\,\\underbrace\{\\left\(\\mathbf\{V\}^\{T\}\\,\\mathbf\{V\}\\right\)\}\_\{=\\,\\mathbf\{I\}\_\{\\text\{out\}\\times\\text\{out\}\}\}=\\mathbf\{\\Sigma\}Thus the gradient rotated to the eigenbasis of Shampoo’s preconditioner effectively diagonalizes it\. Now applying Adam to𝐆tR\\mathbf\{G\}^\{R\}\_\{t\}:

𝐍=Adam⁡\(𝐆t′\)=Adam⁡\(𝚺\)\.\\mathbf\{N\}=\\operatorname\{Adam\}\(\\mathbf\{G\}^\{\\prime\}\_\{t\}\)=\\operatorname\{Adam\}\(\\mathbf\{\\Sigma\}\)\.
The weight update is then given by:

Δ​𝐖=−η​𝐐L​𝐍​𝐐RT=−η​𝐔​𝐍​𝐕T=−η​𝐔⋅Adam⁡\(𝚺\)⋅𝐕T\.\\Delta\\mathbf\{W\}=\-\\eta\\,\\mathbf\{Q\}\_\{L\}\\,\\mathbf\{N\}\\,\\mathbf\{Q\}\_\{R\}^\{T\}=\-\\eta\\,\\mathbf\{U\}\\,\\mathbf\{N\}\\,\\mathbf\{V\}^\{T\}=\-\\eta\\,\\mathbf\{U\}\\cdot\\operatorname\{Adam\}\(\\mathbf\{\\Sigma\}\)\\cdot\\mathbf\{V\}^\{T\}\.
Considering the limit where Adam can be approximated by thesignfunction\[balles2018dissecting\]and when EMA is turned off, the eigenvector matrices perfectly diagonalize the gradient,sign​\(Σ\)=Idmin​\(in,out\)×min​\(in,out\)\\texttt\{sign\}\(\\Sigma\)=\\text\{Id\}\_\{\\text\{min\}\(\\text\{in\},\\text\{out\}\)\\times\\text\{min\}\(\\text\{in\},\\text\{out\}\)\}, the update is given by the polar factor:

Δ​𝐖=−η​𝐔⋅Adam⁡\(𝚺\)⋅𝐕T≈−η​𝐔⋅sign​\(𝚺\)⋅𝐕T=−η​𝐔​𝐕T\\Delta\\mathbf\{W\}=\-\\eta\\,\\mathbf\{U\}\\cdot\\operatorname\{Adam\}\(\\mathbf\{\\Sigma\}\)\\cdot\\mathbf\{V\}^\{T\}\\approx\-\\eta\\,\\mathbf\{U\}\\cdot\\texttt\{sign\}\(\\mathbf\{\\Sigma\}\)\\cdot\\mathbf\{V\}^\{T\}=\-\\eta\\,\\mathbf\{U\}\\,\\mathbf\{V\}^\{T\}\(8\)
Similarly, consider Shampoo with exponentp=−1/4p=\-1/4\[gupta2018shampoo\]:

Δ​𝐖=−η​\(𝐆𝐆𝐓\)−1/4​𝐆​\(𝐆𝐓​𝐆\)−1/4\.\\Delta\\mathbf\{W\}=\-\\eta\\mathbf\{\(GG^\{T\}\)\}^\{\-1/4\}\\,\\mathbf\{G\}\\,\\mathbf\{\(G^\{T\}G\)\}^\{\-1/4\}\.
Taking the−14\-\\frac\{1\}\{4\}power of each term from Eq\.[6](https://arxiv.org/html/2607.20548#A1.E6)and Eq\.[7](https://arxiv.org/html/2607.20548#A1.E7), we obtain

\(𝐆𝐆𝐓\)−1/4\\displaystyle\\mathbf\{\(GG^\{T\}\)\}^\{\-1/4\}=\(𝐔​𝚺2​𝐔T\)−1/4=𝐔​𝚺−1/2​𝐔T,\\displaystyle=\\bigl\(\\mathbf\{U\}\\,\\mathbf\{\\Sigma\}^\{2\}\\,\\mathbf\{U\}^\{T\}\\bigr\)^\{\-1/4\}=\\mathbf\{U\}\\,\\mathbf\{\\Sigma\}^\{\-1/2\}\\,\\mathbf\{U\}^\{T\},\(𝐆𝐓​𝐆\)−1/4\\displaystyle\\mathbf\{\(G^\{T\}G\)\}^\{\-1/4\}=\(𝐕​𝚺2​𝐕T\)−1/4=𝐕​𝚺−1/2​𝐕T\.\\displaystyle=\\bigl\(\\mathbf\{V\}\\,\\mathbf\{\\Sigma\}^\{2\}\\,\\mathbf\{V\}^\{T\}\\bigr\)^\{\-1/4\}=\\mathbf\{V\}\\,\\mathbf\{\\Sigma\}^\{\-1/2\}\\,\\mathbf\{V\}^\{T\}\.
Thus, substituting into the update expression for Shampoo,

\(𝐆𝐆𝐓\)−1/4​𝐆​\(𝐆𝐓​𝐆\)−1/4\\displaystyle\\mathbf\{\(GG^\{T\}\)\}^\{\-1/4\}\\,\\mathbf\{G\}\\,\\mathbf\{\(G^\{T\}G\)\}^\{\-1/4\}=\(𝐔​𝚺−1/2​𝐔T\)​\(𝐔​𝚺​𝐕T\)​\(𝐕​𝚺−1/2​𝐕T\)\\displaystyle=\\Bigl\(\\mathbf\{U\}\\,\\mathbf\{\\Sigma\}^\{\-1/2\}\\,\\mathbf\{U\}^\{T\}\\Bigr\)\\Bigl\(\\mathbf\{U\}\\,\\mathbf\{\\Sigma\}\\,\\mathbf\{V\}^\{T\}\\Bigr\)\\Bigl\(\\mathbf\{V\}\\,\\mathbf\{\\Sigma\}^\{\-1/2\}\\,\\mathbf\{V\}^\{T\}\\Bigr\)=𝐔​𝚺−1/2​𝐔T​𝐔⏟𝐈in×in​𝚺​𝐕T​𝐕⏟𝐈out×out​𝚺−1/2​𝐕T\\displaystyle=\\mathbf\{U\}\\,\\mathbf\{\\Sigma\}^\{\-1/2\}\\,\\underbrace\{\\mathbf\{U\}^\{T\}\\mathbf\{U\}\}\_\{\\mathbf\{I\}\_\{\\text\{in\}\\times\\text\{in\}\}\}\\,\\mathbf\{\\Sigma\}\\,\\underbrace\{\\mathbf\{V\}^\{T\}\\mathbf\{V\}\}\_\{\\mathbf\{I\}\_\{\\text\{out\}\\times\\text\{out\}\}\}\\,\\mathbf\{\\Sigma\}^\{\-1/2\}\\,\\mathbf\{V\}^\{T\}=𝐔​\(𝚺−1/2​𝚺​𝚺−1/2\)​𝐕T\\displaystyle=\\mathbf\{U\}\\,\(\\mathbf\{\\Sigma\}^\{\-1/2\}\\,\\mathbf\{\\Sigma\}\\,\\mathbf\{\\Sigma\}^\{\-1/2\}\\,\)\\mathbf\{V\}^\{T\}=𝐔​\(𝐈𝐝\)​𝐕T\\displaystyle=\\mathbf\{U\}\\,\(\\mathbf\{Id\}\)\\,\\mathbf\{V\}^\{T\}=𝐔​𝐕T\.\\displaystyle=\\mathbf\{U\}\\,\\mathbf\{V\}^\{T\}\.
Hence, Shampoo’s preconditioner is whitening the gradient, which matches what SOAP and Muon are doing in the simplified SVD view without momentum, eq\.[8](https://arxiv.org/html/2607.20548#A1.E8):

Δ​𝐖=−η​\(𝐆𝐆𝐓\)−1/4​𝐆​\(𝐆𝐓​𝐆\)−1/4=−η​𝐔​𝐕T\.\\Delta\\mathbf\{W\}=\-\\eta\\mathbf\{\(GG^\{T\}\)\}^\{\-1/4\}\\,\\mathbf\{G\}\\,\\mathbf\{\(G^\{T\}G\)\}^\{\-1/4\}=\-\\eta\\mathbf\{U\}\\,\\mathbf\{V\}^\{T\}\.

### A\.2An SVD Perspective on Shampoo’s Hessian: Whitening Curvature

For a weight parameterW∈ℛin×outW\\in\\mathcal\{R\}^\{\\text\{in\}\\times\\text\{out\}\}, the optimal Kronecker approximation to the Gauss\-Newton Hessian\[morwani2024new\]was shown to be related to Shampoo’s preconditioner:

H=𝔼​\[G​G⊤\]⊗𝔼​\[G⊤​G\]Tr⁡\(𝔼​\[G​G⊤\]\)H=\\frac\{\\mathbb\{E\}\\left\[G\\,G^\{\\top\}\\right\]\\otimes\\mathbb\{E\}\\left\[G^\{\\top\}\\,G\\right\]\}\{\\operatorname\{Tr\}\\left\(\\mathbb\{E\}\\left\[G\\,G^\{\\top\}\\right\]\\right\)\}
Substituting the spectral decompositions of the covariance matrices from Eqs\.[6](https://arxiv.org/html/2607.20548#A1.E6)and[7](https://arxiv.org/html/2607.20548#A1.E7)into the Kronecker approximation \(assuming the expectations are approximated by the current statistics\), we obtain:

𝐇∝\(𝐔​𝚺2​𝐔T\)⊗\(𝐕​𝚺2​𝐕T\)\.\\mathbf\{H\}\\propto\(\\mathbf\{U\}\\mathbf\{\\Sigma\}^\{2\}\\mathbf\{U\}^\{T\}\)\\otimes\(\\mathbf\{V\}\\mathbf\{\\Sigma\}^\{2\}\\mathbf\{V\}^\{T\}\)\.\(9\)
Using the mixed\-product property of the Kronecker product,\(𝐀⊗𝐁\)​\(𝐂⊗𝐃\)=\(𝐀𝐂\)⊗\(𝐁𝐃\)\(\\mathbf\{A\}\\otimes\\mathbf\{B\}\)\(\\mathbf\{C\}\\otimes\\mathbf\{D\}\)=\(\\mathbf\{AC\}\)\\otimes\(\\mathbf\{BD\}\), we can factorize the expression into an eigendecomposition of the Hessian:

𝐇∝\(𝐔⊗𝐕\)⏟Eigenvectors⋅\(𝚺2⊗𝚺2\)⏟Eigenvalues⋅\(𝐔⊤⊗𝐕⊤\)⏟Basis Change\.\\mathbf\{H\}\\propto\\underbrace\{\(\\mathbf\{U\}\\otimes\\mathbf\{V\}\)\}\_\{\\text\{Eigenvectors\}\}\\cdot\\underbrace\{\(\\mathbf\{\\Sigma\}^\{2\}\\otimes\\mathbf\{\\Sigma\}^\{2\}\)\}\_\{\\text\{Eigenvalues\}\}\\cdot\\underbrace\{\(\\mathbf\{U\}^\{\\top\}\\otimes\\mathbf\{V\}^\{\\top\}\)\}\_\{\\text\{Basis Change\}\}\.\(10\)
This reveals the spectral structure of the implicit Hessian used by Shampoo and SOAP:

- •Theeigenvectorsof Shampoo’s approximate Hessian are the Kronecker products of the left and right singular vectors of the gradient,𝐮i⊗𝐯j\\mathbf\{u\}\_\{i\}\\otimes\\mathbf\{v\}\_\{j\}\.
- •Theeigenvalues\(curvature\) are the products of the squared singular values,λi​j∝σi2​σj2\\lambda\_\{ij\}\\propto\\sigma\_\{i\}^\{2\}\\sigma\_\{j\}^\{2\}\.

Consequently, the preconditioning operation \(which effectively applies𝐇−1/4\\mathbf\{H\}^\{\-1/4\}in Shampoo or𝐇−1/2\\mathbf\{H\}^\{\-1/2\}in SOAP\) dampens the gradient updates primarily along directions corresponding to large singular values \(where curvatureσi2​σj2\\sigma\_\{i\}^\{2\}\\sigma\_\{j\}^\{2\}is high\) and amplifies them where singular values are small\.

### A\.3An SVD perspective on KL\-Shampoo: Improving Conditioning for Inversion and Decomposition

While Shampoo and SOAP update Kronecker factors using the marginal covariance statistics𝐆𝐆T\\mathbf\{G\}\\mathbf\{G\}^\{T\}and𝐆T​𝐆\\mathbf\{G\}^\{T\}\\mathbf\{G\}, KL\-Shampoo\[lin2025understanding\]introduces a coupled update rule where each factor is updated using the inverse of the other\. Here we analyze the spectral properties of this update using the singular value decomposition \(SVD\) of the gradient\.

Let the SVD of the gradient at stepttbe𝐆t=𝐔​𝚺​𝐕T\\mathbf\{G\}\_\{t\}=\\mathbf\{U\}\\mathbf\{\\Sigma\}\\mathbf\{V\}^\{T\}\. We assume the Kronecker factors𝐒a\\mathbf\{S\}\_\{a\}and𝐒b\\mathbf\{S\}\_\{b\}are approximately aligned with the gradient’s eigenbasis, such that:

𝐒a\\displaystyle\\mathbf\{S\}\_\{a\}≈𝐔​𝚲a​𝐔T\\displaystyle\\approx\\mathbf\{U\}\\mathbf\{\\Lambda\}\_\{a\}\\mathbf\{U\}^\{T\}𝐒b\\displaystyle\\mathbf\{S\}\_\{b\}≈𝐕​𝚲b​𝐕T\\displaystyle\\approx\\mathbf\{V\}\\mathbf\{\\Lambda\}\_\{b\}\\mathbf\{V\}^\{T\}where𝚲a\\mathbf\{\\Lambda\}\_\{a\}and𝚲b\\mathbf\{\\Lambda\}\_\{b\}represent the eigenvalues captured by the Kronecker factors\.

Recall the KL\-Shampoo update rule for𝐒a\\mathbf\{S\}\_\{a\}:

𝐒a←\(1−β\)​𝐒a\+β​\(𝐆𝐒b−1​𝐆T\)\.\\mathbf\{S\}\_\{a\}\\leftarrow\(1\-\\beta\)\\mathbf\{S\}\_\{a\}\+\\beta\\left\(\\mathbf\{G\}\\mathbf\{S\}\_\{b\}^\{\-1\}\\mathbf\{G\}^\{T\}\\right\)\.\(11\)
Substituting the spectral forms into the update term:

𝐆𝐒b−1​𝐆T\\displaystyle\\mathbf\{G\}\\mathbf\{S\}\_\{b\}^\{\-1\}\\mathbf\{G\}^\{T\}=\(𝐔​𝚺​𝐕T\)​\(𝐕​𝚲b−1​𝐕T\)​\(𝐕​𝚺​𝐔T\)\\displaystyle=\(\\mathbf\{U\}\\mathbf\{\\Sigma\}\\mathbf\{V\}^\{T\}\)\(\\mathbf\{V\}\\mathbf\{\\Lambda\}\_\{b\}^\{\-1\}\\mathbf\{V\}^\{T\}\)\(\\mathbf\{V\}\\mathbf\{\\Sigma\}\\mathbf\{U\}^\{T\}\)=𝐔​𝚺​\(𝐕T​𝐕\)​𝚲b−1​\(𝐕T​𝐕\)​𝚺​𝐔T\\displaystyle=\\mathbf\{U\}\\mathbf\{\\Sigma\}\(\\mathbf\{V\}^\{T\}\\mathbf\{V\}\)\\mathbf\{\\Lambda\}\_\{b\}^\{\-1\}\(\\mathbf\{V\}^\{T\}\\mathbf\{V\}\)\\mathbf\{\\Sigma\}\\mathbf\{U\}^\{T\}=𝐔​\(𝚺​𝚲b−1​𝚺\)​𝐔T\\displaystyle=\\mathbf\{U\}\(\\mathbf\{\\Sigma\}\\mathbf\{\\Lambda\}\_\{b\}^\{\-1\}\\mathbf\{\\Sigma\}\)\\mathbf\{U\}^\{T\}=𝐔​\(𝚺2​𝚲b−1\)​𝐔T\.\\displaystyle=\\mathbf\{U\}\(\\mathbf\{\\Sigma\}^\{2\}\\mathbf\{\\Lambda\}\_\{b\}^\{\-1\}\)\\mathbf\{U\}^\{T\}\.
At a fixed point where the preconditioning statistics stabilize \(𝐒n​e​w≈𝐒o​l​d\\mathbf\{S\}\_\{new\}\\approx\\mathbf\{S\}\_\{old\}\), the eigenvalues of the updated factor must match the eigenvalues of the update term\. Thus, for the diagonal entries:

𝚲a≈𝚺2​𝚲b−1\.\\mathbf\{\\Lambda\}\_\{a\}\\approx\\mathbf\{\\Sigma\}^\{2\}\\mathbf\{\\Lambda\}\_\{b\}^\{\-1\}\.\(12\)
By symmetry, the update for𝐒b\\mathbf\{S\}\_\{b\}yields𝚲b≈𝚺2​𝚲a−1\\mathbf\{\\Lambda\}\_\{b\}\\approx\\mathbf\{\\Sigma\}^\{2\}\\mathbf\{\\Lambda\}\_\{a\}^\{\-1\}\. Solving this system for the case where the spectra are balanced \(𝚲a≈𝚲b=𝚲\\mathbf\{\\Lambda\}\_\{a\}\\approx\\mathbf\{\\Lambda\}\_\{b\}=\\mathbf\{\\Lambda\}\):

𝚲\\displaystyle\\mathbf\{\\Lambda\}≈𝚺2​𝚲−1\\displaystyle\\approx\\mathbf\{\\Sigma\}^\{2\}\\mathbf\{\\Lambda\}^\{\-1\}𝚲2\\displaystyle\\mathbf\{\\Lambda\}^\{2\}≈𝚺2\\displaystyle\\approx\\mathbf\{\\Sigma\}^\{2\}𝚲\\displaystyle\\mathbf\{\\Lambda\}≈𝚺\.\\displaystyle\\approx\\mathbf\{\\Sigma\}\.
This result highlights the fundamental difference between standard Shampoo and KL\-Shampoo:

- •Shampoo:Kronecker factors track the covariance, so𝚲Shampoo≈𝚺2\\mathbf\{\\Lambda\}\_\{\\text\{Shampoo\}\}\\approx\\mathbf\{\\Sigma\}^\{2\}\.
- •KL\-Shampoo:Kronecker factors track the singular values, so𝚲KL\-Shampoo≈𝚺\\mathbf\{\\Lambda\}\_\{\\text\{KL\-Shampoo\}\}\\approx\\mathbf\{\\Sigma\}\.

This has an important advantage for numerical stability\. The condition number of the matrices being inverted \(or decomposed\) in KL\-Shampoo is the square root of that in standard Shampoo:

κ​\(𝐒KL\-Shampoo\)=σmaxσmin=κ​\(𝐒Shampoo\)\.\\kappa\(\\mathbf\{S\}\_\{\\text\{KL\-Shampoo\}\}\)=\\frac\{\\sigma\_\{\\max\}\}\{\\sigma\_\{\\min\}\}=\\sqrt\{\\kappa\(\\mathbf\{S\}\_\{\\text\{Shampoo\}\}\)\}\.\(13\)For ill\-conditioned gradients whereκ​\(𝐆\)\\kappa\(\\mathbf\{G\}\)is large, KL\-Shampoo operates on significantly better\-conditioned matrices, reducing the accumulation of numerical errors during the eigen\-decomposition steps\.

## Appendix BMXFP8 Pretraining

We use the MXFP8 recipe developed by NVIDIA for Blackwell\[mishra2025recipes\]together with Muon and pretrain NVIDIA\-Nemotron\-NanoV3 and the 8B Hybrid Mamba\-Attention LatentMoE for 1T tokens\. We find that MXFP8 has a small loss gap, while most downstream evaluations are higher with MXFP8 than BF16, Table[7](https://arxiv.org/html/2607.20548#A2.T7)\. Further, we find that replacing AdamW with Lion\[chen2023symbolic\]with\(β1,β2\)=\(0\.95,0\.98\)\(\\beta\_\{1\},\\beta\_\{2\}\)=\(0\.95,0\.98\)improves losses and evaluations \(Table[7](https://arxiv.org/html/2607.20548#A2.T7), last column\)\.

Table 7:Benchmark results across MMLU, Coding, Math, and Commonsense Reasoning comparing BF16 and MXFP8 precision formats for the NVIDIA\-Nemotron\-NanoV3\-1T and 8B Hybrid LatentMoE\-1T models\. All experiments were performed at2×2\\timesGBS of 25M tokens with Muon as the backbone optimizer\.![Refer to caption](https://arxiv.org/html/2607.20548v1/qwen3_30b_muon_split_vs_nosplit.png)Figure 9:Training loss of a Qwen\-3\-30B\-A3B architecture comparing Muon’s Newton\-Schulz with and without splitting the fused QKV tensor\. With split\-QKV, the model has lower training loss early in pretraining, but this difference disappears as we scale token counts\.
## Appendix CAdamW at Larger Batch Sizes

We find that the techniques that enabled larger batch pretraining with Muon did not help with AdamW, Fig\.[10](https://arxiv.org/html/2607.20548#A3.F10)\. We leave tuning the hyperparametersβ1,β2,ϵ\\beta\_\{1\},\\beta\_\{2\},\\epsilonof AdamW with batch size as future work\.

![Refer to caption](https://arxiv.org/html/2607.20548v1/nanov3_1T_adam_1x_vs_2x.png)Figure 10:NVIDIA\-Nemotron\-NanoV3 pretrained with AdamW for a 1T\-token horizon at 25M and 50M batch sizes\. Batch size ramp\-up and square\-root scaling were applied to the 50M batch size run\.
## Appendix DOptimizing Conv1d Filters with Muon

For hybrid architectures incorporating Mamba2 SSM layers, we found that excludingConv1Dc parameters from Muon and optimizing them with AdamW instead yielded a modest but consistent improvement in training and validation loss, as shown in Fig\.[12](https://arxiv.org/html/2607.20548#A5.F12)\. The Conv1D weights in Mamba2 encode localized temporal filters that are shared and re\-applied at every sequence position, making orthogonality constraints geometrically ill\-motivated compared to standard linear projections\.

### D\.1Splitting fused QKV tensors before orthogonalization

Fig\.[9](https://arxiv.org/html/2607.20548#A2.F9)shows Muon with and without splitting QKV tensors before orthogonalization\. While splitting achieves lower losses early in training, the losses at the end of the full 1T\-token horizon are identical\.

## Appendix ESOAP Eigenbasis Update Frequency

![Refer to caption](https://arxiv.org/html/2607.20548v1/nanov3_1T_muon_original_vs_better_recipe.png)Figure 11:Training loss of Nano\-V3 comparing Muon’s quintic\-coefficient optimized Newton\-Schulz recipe and the Polar Express recipe run for 16 iterations\.![Refer to caption](https://arxiv.org/html/2607.20548v1/nanov3_1T_muon_vs_adamw_conv1d.png)Figure 12:Training loss of Nano\-V3 comparing the application of Muon vs\. AdamW on the Mamba mixer Conv1D\. Using AdamW on the Conv1D layer leads to slightly lower \(≈0\.1%\\approx 0\.1\\%\) loss\.![Refer to caption](https://arxiv.org/html/2607.20548v1/8b_dense_1x_soap_eigh_vs_noeigh.png)Figure 13:Usingeighvs\.QRdecomposition to recompute SOAP’s eigenbasis has a negligible impact on large\-scale pretraining loss on a Qwen\-3\-30B\-A3B architecture; hence, we default to usingQRsince it is cheaper\.Fig\.[13](https://arxiv.org/html/2607.20548#A5.F13)is a comparison ofeighvs\.QRdecomposition to recompute the eigenbasis in SOAP\. We found it necessary to perform QR decomposition at every step to reduce the loss spikes\. This is likely due to the rapidly changing nature of the preconditioner basis early in training, which slower update schedules failed to capture\. With this frequency, there was negligible difference between the use ofQRoreigh\.