SOAP-Muon: Hybrid Optimizer for MLIPs

• SOAP-Muon is a hybrid optimizer that combines SOAP’s adaptive preconditioning with a Muon-inspired orthogonalization step for selected parameters.
• The method updates matrix-valued parameters by projecting gradients into a rotating eigenbasis before selectively applying orthogonalization and normalization.
• It demonstrates enhanced energy and force accuracy with improved label efficiency in MLIPs, outperforming AdamW under sparse supervision.

SOAP-Muon is a hybrid matrix-structured optimizer that combines SOAP’s Shampoo-eigenbasis adaptive preconditioning with a Muon-inspired orthogonalization and normalization step on selected parameter groups. In the current optimizer literature, it is presented as a practical method for training machine-learning interatomic potentials (MLIPs), specifically NequIP and Allegro models in the nequip framework, where SOAP and SOAP-Muon emerged as robust and consistently strong methods and the gains were especially pronounced under partial force supervision (Harari et al., 2 Jul 2026). Its conceptual lineage runs through SOAP, which interprets Shampoo as AdamW in a rotating preconditioner eigenbasis, and through Muon, which replaces a matrix momentum direction by an approximate orthogonal factor (Vyas et al., 2024).

1. Conceptual position in the optimizer landscape

SOAP-Muon occupies the intersection of two matrix-aware optimization ideas. The first is SOAP, or “ShampoO with Adam in the Preconditioner’s eigenbasis,” which uses left and right second-order statistics to define a rotating orthogonal basis and then performs Adam-style first- and second-moment adaptation in that basis (Vyas et al., 2024). The second is Muon, an orthogonalized matrix-update optimizer in which a Nesterov-style momentum direction is transformed by NewtonSchulz5, an approximation to the orthogonal factor $UV^\top$ of an SVD $G = U \Sigma V^\top$, before the parameter update (Harari et al., 2 Jul 2026).

Within that design space, SOAP-Muon is not described as a completely new preconditioner. Rather, it extends SOAP by adding a Muon-inspired orthogonalization step after the standard SOAP step. The paper’s phrasing is explicit: “SOAP-Muon extends SOAP by adding a Muon-inspired orthogonalization step after the standard SOAP step.” In practical terms, the hybridization is selective rather than uniform: parameters deemed suitable for Muon-style treatment receive an orthogonalized SOAP-Muon update, while the remaining parameters receive standard SOAP updates without orthogonalization (Harari et al., 2 Jul 2026).

This selective construction is central to the method’s identity. SOAP-Muon is therefore best understood not as “SOAP everywhere plus Muon everywhere,” but as a parameter-structural hybrid in which adaptive preconditioning is global in intent while orthogonalization is applied only where the tensor structure and layer role appear compatible with it (Harari et al., 2 Jul 2026).

2. Algorithmic construction

The paper places SOAP-Muon inside a generic preconditioned matrix-update viewpoint. For a parameter vector $w$ with gradient $g$, the generic form is

$w \leftarrow w - \eta P^{-1} g,$

and for a matrix-valued parameter $W$ with gradient $G \in \mathbb{R}^{m \times n}$, the structured analogue is

$W \leftarrow W - \eta L^{-p} G R^{-p},$

with left and right preconditioners $L$ and $R$ and preconditioning power $G = U \Sigma V^\top$0 (Harari et al., 2 Jul 2026).

SOAP supplies the rotating coordinate system. The gradient is projected into the Shampoo eigenbasis, Adam first and second moments are accumulated in that transformed basis, and the normalized direction is projected back to the original parameter space. The general update used for SOAP and SOAP-Muon begins with the current gradient $G = U \Sigma V^\top$1, projects it to the Shampoo eigenbasis, forms

$G = U \Sigma V^\top$2

$G = U \Sigma V^\top$3

computes the bias-corrected step size

$G = U \Sigma V^\top$4

and applies elementwise variance normalization

$G = U \Sigma V^\top$5

The resulting direction is then back-projected to the original parameter space before the final update (Harari et al., 2 Jul 2026).

SOAP-Muon enters when orthogonalization is enabled. If $G = U \Sigma V^\top$6, the back-projected direction is passed through

$G = U \Sigma V^\top$7

which yields a Muon-style orthogonalization. If orthogonalization is enabled with a nonzero singular-value power, the paper states that one computes an SVD and applies a singular-value-power transform; in practice, the study used $G = U \Sigma V^\top$8, where $G = U \Sigma V^\top$9 gives pure orthogonalization and $w$0 requires full SVD. If normalization is enabled, the update is further rescaled by

$w$1

Standard SOAP corresponds to ortho = False, normalize = False, whereas SOAP-Muon corresponds to ortho = True, normalize = True, but only on the designated Muon-style parameter group (Harari et al., 2 Jul 2026).

The architectural template comes from SOAP’s original analysis. SOAP was motivated by the formal observation that Shampoo with exponent $w$2 is equivalent to running Adafactor in the eigenbasis of Shampoo’s preconditioner, which in turn suggests replacing Adafactor by AdamW in that basis. This eigenbasis interpretation is the piece that SOAP-Muon inherits directly; the orthogonalization is an added layer on top of that rotated adaptive update (Vyas et al., 2024).

3. Parameter grouping and implementation in equivariant MLIPs

The implementation studied for MLIPs is architecture-specific and depends heavily on tensor structure. The authors do not apply orthogonalization uniformly across all parameters. Instead, they use a grouping rule similar to Muon: matrix-like hidden-layer weights form the orthogonalized group, while embeddings, readouts, tensor-product weights, and other sensitive or non-matrix parameters remain in the non-orthogonalized group and receive standard SOAP updates (Harari et al., 2 Jul 2026).

In the paper’s own terminology, some tables label the non-orthogonalized group as “ADAM,” but within SOAP-Muon that group actually receives standard SOAP rather than AdamW. The resulting partition is as follows.

Model	Orthogonalized “Muon group”	Standard SOAP group
NequIP	edge-MLP weight matrices; e3nn.o3.Linear weights	type embedding; e3nn.o3.FullyConnectedTensorProduct weights; energy readout MLP weights; per-type energy scale and shift
Allegro	scalar/tensor embedding MLP weights; latent MLP weights; first-layer environment MLP weights	type embeddings; tensor-product weights; readout MLP weights

A further implementation subtlety arises from e3nn. Many equivariant weights are stored as flattened 1D parameter vectors even when they correspond to structured blocks in the forward pass. For e3nn.o3.Linear, each instruction has a multiplicity-mixing matrix $w$3; for FullyConnectedTensorProduct, each instruction has a multiplicity-mixing tensor $w$4. The study reshaped the Linear weights for structured optimization in Muon, SOAP, and SOAP-Muon, but explicitly did not slice and reshape the FullyConnectedTensorProduct weights into 3D tensors for structured optimization; that was left for future work (Harari et al., 2 Jul 2026).

The paper also notes that its implementation does not exactly match the official SOAP-Muon implementation. The study used separate $w$5 hyperparameters for the orthogonalized “Muon group” and the non-orthogonalized “Adam group,” whereas the official implementation shares momentum settings across all parameters. This makes the reported SOAP-Muon results specific to a practically tuned, parameter-grouped realization rather than to a single monolithic update rule (Harari et al., 2 Jul 2026).

4. Empirical behavior on NequIP and Allegro

SOAP-Muon was evaluated on two MLIP problems: NequIP on liquid water and Allegro on crystalline cesium dihydrogen phosphate (CDP). Hyperparameter sweeps used shortened schedules of 100 epochs for CDP and 500 epochs for water; final training used five random seeds—7, 42, 123, 2026, and 1618—with 1000 epochs for CDP and 2000 epochs for water, on NVIDIA A100-SXM4-80GB and H200 GPUs, with OpenEquivariance and cuEquivariance tensor-product kernels (Harari et al., 2 Jul 2026).

Under full energy+force supervision, SOAP-Muon was among the strongest methods but not uniformly dominant. For CDP with Allegro, it achieved energy MAE $w$6 meV/atom and force MAE $w$7 meV/Å, giving the best force accuracy among the four tested optimizers. For water with NequIP, it achieved energy MAE $w$8 meV/atom and force MAE $w$9 meV/Å, giving the best energy accuracy and essentially tying SOAP on force. Muon alone was markedly less reliable, especially on water, where it produced energy MAE $g$0 meV/atom and force MAE $g$1 meV/Å (Harari et al., 2 Jul 2026).

In energy-only training, SOAP-Muon’s strongest case appeared on CDP. There it achieved energy MAE $g$2 meV/atom and force MAE $g$3 meV/Å, which the paper identifies as the strongest overall result in that regime and summarizes as roughly 47% lower energy MAE and 60% lower force MAE than AdamW. On water, by contrast, SOAP remained better: SOAP-Muon achieved $g$4 meV/atom and $g$5 meV/Å, which still improved on AdamW but did not exceed SOAP (Harari et al., 2 Jul 2026).

The most prominent claim concerns label efficiency. Under partial force supervision, the paper states that “SOAP-Muon trained with 50% of force labels matches AdamW trained with 100%.” The clearest evidence is on CDP, where the normalized relative MAE values at 50% force supervision are approximately 1.01 for both energy and force relative to AdamW at 100% labels; at 100% force supervision the corresponding values improve to 0.93 for energy and 0.86 for force. The sparse-force regime is also where the physical-fidelity advantage becomes especially significant. At 5% force supervision on CDP, SOAP-Muon remained stable and accurately reproduced AIMD radial distribution functions and proton diffusion behavior, whereas the corresponding AdamW model was catastrophically unstable. The reported activation energies were $g$6 eV for SOAP-Muon with full energy+force supervision and $g$7 eV for SOAP-Muon with energy plus 5% force supervision, against an experimental range of $g$8–$g$9 eV (Harari et al., 2 Jul 2026).

The paper also reports a fast-convergence trend for matrix-structured optimizers, but the clearest quantified time-to-accuracy statements are given for SOAP rather than for SOAP-Muon specifically: on CDP, SOAP reaches AdamW’s minimum median validation force MAE 4.9× faster, and on water 5.8× faster. This suggests that SOAP-Muon participates in the same broader acceleration pattern, but the most explicit wall-clock claims in the study are attached to SOAP (Harari et al., 2 Jul 2026).

5. Hyperparameters, stability, and computational trade-offs

SOAP-Muon is presented as a high-upside but less forgiving alternative to SOAP. The paper is explicit that it required substantially smaller learning rates than the other optimizers, roughly one order of magnitude smaller. Learning rates that were optimal for AdamW, Muon, or SOAP often led to severe training instabilities when used with SOAP-Muon. Under full energy+force supervision, the reported best learning rates were $w \leftarrow w - \eta P^{-1} g,$0 for SOAP-Muon on both CDP/Allegro and water/NequIP, compared with $w \leftarrow w - \eta P^{-1} g,$1 or $w \leftarrow w - \eta P^{-1} g,$2 for the other optimizers depending on system (Harari et al., 2 Jul 2026).

The singular-value power $w \leftarrow w - \eta P^{-1} g,$3 is a second major knob. The official SOAP-Muon implementation defaults to $w \leftarrow w - \eta P^{-1} g,$4, which requires full SVD and therefore more compute. The cheaper Muon-style choice is $w \leftarrow w - \eta P^{-1} g,$5, enabling Newton–Schulz orthogonalization. The study found that the better setting was dataset-specific: $w \leftarrow w - \eta P^{-1} g,$6 for Allegro on CDP and $w \leftarrow w - \eta P^{-1} g,$7 for NequIP on water. The paper further notes, citing Vyas et al., that $w \leftarrow w - \eta P^{-1} g,$8 can cause dataset-dependent instability (Harari et al., 2 Jul 2026).

Momentum settings were likewise dataset- and group-dependent. The study searched over

$w \leftarrow w - \eta P^{-1} g,$9

The best settings were $W$0 for the Muon group and $W$1 for the non-orthogonalized group on Allegro-CDP, and $W$2 for both groups on NequIP-water. SOAP itself used the simpler default $W$3, $W$4, $W$5, with preconditioning frequency $W$6. SOAP-Muon also used preconditioner update frequency 10 and “full preconditioning,” meaning the preconditioner was applied along every mode of the weight tensor (Harari et al., 2 Jul 2026).

The trade-off against SOAP is therefore clear. SOAP-Muon adds orthogonalization and, depending on $W$7, potentially full SVD cost. The paper’s interpretation is that adaptive preconditioning is the robust ingredient, while orthogonalization can help in some settings but can also hurt if not tuned. Its discussion states that “the behavior across all three matrix-structured variants suggests that the orthogonalization step is the primary source of degradation, with adaptive preconditioning mitigating but not eliminating its effects.” For that reason, the paper’s practical default recommendation is SOAP rather than SOAP-Muon when minimal extra tuning is desired (Harari et al., 2 Jul 2026).

6. Theoretical status and broader developments

SOAP-Muon’s theory is indirect. SOAP itself has a formal conceptual foundation: the original SOAP paper shows that Shampoo with exponent $W$8 is equivalent to running Adafactor in Shampoo’s eigenbasis, which leads naturally to SOAP as AdamW in that rotated space. In large-batch language-model pretraining, SOAP reduced the number of iterations by over 40% and wall clock time by over 35% compared to AdamW, with approximately 20% improvements in both metrics compared to Shampoo (Vyas et al., 2024).

A later short note established a first convergence-rate analysis for generalized SOAP with arbitrary orthogonal projection matrices, provided they are conditionally independent of the current stochastic gradient. The result covers time-varying orthogonal projections constructed from stale information and shows that such bases need not come specifically from SOAP’s eigendecomposition. However, that paper is explicit that it does not analyze Muon itself and does not define or prove anything for a named SOAP-Muon hybrid (Li et al., 23 Apr 2026). This places SOAP-Muon in a partially supported but not fully resolved theoretical position. A plausible implication is that hybrids retaining the generalized SOAP form with orthogonal projections computed from past information may inherit part of that theory, whereas deeper hybrids that materially change the update rule remain outside the theorem’s scope (Li et al., 23 Apr 2026).

The optimizer family has continued to evolve beyond both SOAP and Muon. Muown, introduced as a drop-in replacement for Muon with explicit row-norm control, was reported to improve perplexity over Muon, SOAP, AdamW, and Lion across GPT-style pretraining runs from 124M to 2.7B parameters, while reducing sensitivity to weight decay and avoiding the spectral norm drift observed under Muon (Lion et al., 11 May 2026). In that broader context, SOAP-Muon appears less as a terminus than as a particularly instructive hybrid: it shows that combining Shampoo-style adaptive preconditioning with Muon-style orthogonalization can be highly effective, especially in label-efficient MLIP training, but it also makes clear that orthogonalized matrix optimization remains an active and unsettled design space (Harari et al., 2 Jul 2026).