1 Introduction↩︎

Machine learning interatomic potentials (MLIPs) [1]–[3] have become an indispensable tool for atomistic simulation in chemistry and materials science [4]. They have enabled studies ranging from the thermodynamics of water [5] to the properties of amorphous silicon [6], large-scale materials screening [7], and even all-atom simulations of complex biomolecular assemblies such as the HIV capsid [8].

Motivated by these successes, the MLIP community has invested heavily in two primary directions of improvement. The first is new architectures, progressing from descriptor-based [3], [9]–[12] to deep learning approaches such as invariant [13] and equivariant graph neural networks [14]–[17] and attention-based architectures [18]–[21]. The second direction is new datasets, with large-scale quantum chemical databases for organic molecules [22]–[24], catalytic systems [25], and inorganic materials [26], [27] providing the corpora on which these architectures are trained and evaluated.

A third axis of improvement, the training techniques that dictate the dynamics of learning, has received comparatively little attention. Crucially, the choice and configuration of the optimizer used is largely static. The vast majority of MLIP training pipelines default to Adam [28] or one of its closely related variants such as AdamW [29]. Even at the scale of training universal potentials, such as the Universal Models for Atoms (UMA) [30] and SevenNet-Omni [31], Adam-family optimizers remain the default choice in reported training protocols.

Meanwhile, the scaling of large language models (LLMs) has renewed interest in optimizers that go beyond the diagonal gradient scaling of Adam by exploiting matrix structure in neural-network weight tensors. This broad direction includes explicit structured preconditioners such as Shampoo [32], SOAP [33], Kron [34], SPlus [35], and DyKAF [36]; orthogonalized matrix-update methods such as Muon [37], Scion [38], Gluon [39], and PolarGrad [40]; and hybrid adaptive schemes such as SOAP-Muon [41], COSMOS [42], Mousse [43], and Newton-Muon [44].

In this work, we focus on Muon and SOAP as prominent optimizers, together with SOAP-Muon as a hybrid variant combining elements of both. For LLMs, these methods have demonstrated faster convergence to competitive loss values [33], [37], [45]. However, their potential for MLIP training remains underexplored. [46] recently argue that optimizer choice is an overlooked factor in MLIP fine-tuning, but restrict their benchmark to optimizers based on diagonal preconditioning methods. [47] provided encouraging evidence for this line of inquiry, showing that Muon outperforms Adam as an ingredient for budget-conscious foundation potential training, though without extensive ablations.

We work towards filling this gap by benchmarking Muon, SOAP, and SOAP-Muon against AdamW on two systems of physical significance: liquid water [5], modeled with NequIP [14] and the solid acid electrolyte (cesium dihydrogen phosphate, CDP) [48], important for electrochemical energy technologies, modeled with Allegro [16]. We additionally examine optimizer behavior under varying levels of force supervision, including energy-only training. Computing force labels in density functional theory (DFT) incurs little additional cost beyond the energy evaluation, thanks to the Hellmann–Feynman theorem [49]. Higher levels of theory, however, do not share this benefit. For coupled cluster methods, force computation is costly [50]. This limitation is shared by diffusion Monte Carlo, where energy-only training is typically used due to the prohibitive cost of force evaluations [51]. Methods that are better suited for sparse-force regimes are thus desirable for training more accurate MLIPs on reference data beyond DFT. Furthermore, even when forces are available, some works have trained on only a subset to reduce the training time and memory requirements [52], highlighting sparse force supervision as a practically relevant regime.

Our contributions here are as follows:

• We integrate Muon, SOAP, and SOAP-Muon into the nequip MLIP framework [53] and benchmark them against AdamW across two systems. We find that SOAP and SOAP-Muon consistently improve energy and force accuracy and accelerate convergence, with SOAP showing the most robust behavior across systems and SOAP-Muon achieving the strongest results in selected settings. Muon provides gains over AdamW only in one system.

• We examine how these optimizers behave under reduced force supervision: SOAP and SOAP-Muon retain strong accuracy as force supervision is reduced; in particular, SOAP-Muon trained with \(50\%\) of force labels matches AdamW trained with \(100\%\), suggesting a path to reducing force-label requirements in regimes where force labels are expensive.

• We demonstrate that the resulting MLIPs are physically faithful, reproducing corresponding ab initio calculations and experimental observables. Notably, in one of the systems studied, SOAP-Muon preserves this fidelity even when trained with only \(5\%\) of the force labels, while the corresponding AdamW model becomes unstable at the same level of force supervision.

2 Background↩︎

2.0.0.1 MLIP Training

MLIPs approximate the potential energy surface (PES) of an atomistic system by fitting a neural network with weights \(\theta\) to quantum mechanical reference data, most commonly computed using DFT. The network maps a set of atomic positions and chemical species \(\{\mathbf{r}_j,Z_j\}\) to a total energy \(E_\theta\), which decomposes into a sum of local atomic contributions \(\varepsilon_{i, \theta}\) to ensure size extensivity, \[E_\theta(\{\mathbf{r}_j, Z_j\}) = \sum_i \varepsilon_{i,\theta}(\{\mathbf{r}_j, Z_j\}_{j \in \mathcal{N}_i}),\] where \(\mathcal{N}_i\) denotes the local neighborhood of atom \(i\) within a cutoff radius. Forces are obtained as the negative gradient of the predicted energy with respect to atomic positions via automatic differentiation, \[\mathbf{F}_{i,\theta} = -\frac{\partial E_{\theta}}{\partial \mathbf{r}_i},\] enforcing energy conservation by construction. Training is performed by minimizing a weighted combination of energy and force errors over a dataset of \(N\) reference configurations \(\{(\{\mathbf{r}_j,Z_j\}_n, E_n, \mathbf{F}_n)\}_{n=1}^N\), where \(\mathbf{F}_n \in \mathbb{R}^{3N_\text{atoms}}\) collects the reference forces on all atoms in configuration \(n\), \[\begin{align} \label{eq:loss} \mathcal{L}(\theta) &= \frac{\lambda_E}{N}\sum_{n=1}^N \left(E_\theta(\{\mathbf{r}_j, Z_j\}_n) - E_n\right)^2 \nonumber\\ &+ \frac{\lambda_F}{3\sum_{n=1}^{N}N_\text{atoms}^{(n)}}\sum_{n=1}^N \left\|\mathbf{F}_\theta(\{\mathbf{r}_j, Z_j\}_n) - \mathbf{F}_n\right\|^2, \end{align}\tag{1}\] where \(\lambda_E\) and \(\lambda_F\) control the relative contribution of energy and force supervision. Including forces in the training loss substantially improves accuracy and stability [14]. In the force-sparse setting we study here, the force loss is computed over only a subset of configurations, with the limiting case \(\lambda_F = 0\) corresponding to energy-only training (see Appendix 8 for details).

2.0.0.2 NequIP and Allegro

NequIP [14] and Allegro [16] are \(E(3)\)-equivariant graph neural networks, in which internal features are geometric tensors that transform under rotations, reflections, and translations, yielding energy and force predictions that respect the fundamental 3D Euclidean symmetries of atomistic systems. Local information is mixed across the network through equivariant tensor products. NequIP propagates features via message-passing between neighboring atoms, while Allegro builds many-body representations through iterated tensor products of atom-pairwise features within a strictly local cutoff. Together, these architectures laid the foundations for the most widely used class of equivariant MLIPs, known for high accuracy and strong data efficiency attributed to the inductive bias of equivariance [14], [16]. Both are implemented within the nequip software framework [53].

2.0.0.3 Optimizers

Vanilla gradient descent updates parameters in the direction of steepest descent. However, the loss surface \(\mathcal{L}(\theta)\) is generally anisotropic, i.e., some directions change the loss more than others per unit step. Preconditioning addresses this inconsistency by rescaling the gradient to account for the local geometry of the loss, allowing larger steps in flat directions and smaller steps in steep ones. For a vector-valued parameter \(w\) with gradient \(g\), a preconditioned update takes the form \[w \leftarrow w - \eta P^{-1} g,\] where \(P\) is a positive definite matrix describing the local curvature and \(\eta\) is the learning rate. A natural choice is the Hessian, but forming, storing, and inverting the Hessian is prohibitively expensive at scale. Practical optimizers therefore choose tractable approximations for \(P\). Additionally, most optimizers incorporate momentum, an exponential moving average of past gradients, to smooth noisy updates and accelerate convergence.

Adam [28] approximates \(P\) with a diagonal matrix \(P = \operatorname{diag} (\hat{v}_t)^{1/2}\), where \(\hat{v}_t\) is a running estimate of the element-wise squared gradient. AdamW [29] decouples the weight decay from the Adam gradient update.

For a matrix-valued parameter \(W\) with gradient \(G \in \mathbb{R}^{m \times n}\), the full preconditioner on \(\text{vec}(W)\) would be an \(mn \times mn\) matrix, which is intractable in practice. An approximation is to assume Kronecker structure \(P \approx L^\mathsf{T}\otimes R\), which leads to the update \[W \leftarrow W - \eta L^{-p}GR^{-p},\] where \(L \in \mathbb{R}^{m \times m}\) and \(R \in \mathbb{R}^{n\times n}\) are left and right preconditioners, and \(p\) is the power of the preconditioner.

Shampoo [32] defines the left and right preconditioners using the accumulated uncentered row and column covariances of the gradient matrix, updating them as \(L \leftarrow L + GG^\mathsf{T}\) and \(R \leftarrow R + G^\mathsf{T}G\). The preconditioning power is set to \(p=1/4\). Later work by [54] showed that, when the accumulation is removed, Shampoo simplifies to a semi-orthogonal weight update. In this regime, the method can be interpreted as projecting the gradient matrix onto the nearest semi-orthogonal matrix under the Frobenius norm.

Muon [37] builds on this perspective, by first applying a Nesterov-style momentum update and then performing the orthogonalization step using efficient Newton–Schulz iterations.

SOAP [33] is motivated by the observation that Shampoo is equivalent to running Adafactor [55] in the eigenspace of the preconditioner. That is, the gradient is transformed as \(G_t'\leftarrow Q^\mathsf{T}_L G Q_R\) where \(Q_L\) and \(Q_R\) are the eigenvector matrices of the \(L\) and \(R\), respectively. SOAP replaces the Adafactor update in this eigenspace with an AdamW step, and then projects the result back to the original parameter space.

SOAP-Muon [41] extends SOAP by adding a Muon-inspired orthogonalization step after the standard SOAP step.

3 Experimental Setup↩︎

To assess the versatility of the studied optimizers, we evaluate them across different equivariant architectures and chemical environments. We consider two representative settings: homogeneous liquid water [5] modeled with NequIP, and solid-state CDP [48] modeled with Allegro.

Our chosen model configurations are based on previously reported parameters [14], [48]. To ensure a fair comparison across each model, dataset, and optimizer, we use a systematic hyperparameter tuning protocol (see Appendix 8 for details).

To probe the limits of each optimizer, we additionally consider reduced force supervision, including the energy-only regime. This setting evaluates how these optimizers perform with reduced data. Accordingly, we repeat the tuning protocol across training sets with varying fractions of force-labeled frames—\(75\%\), \(50\%\), \(10\%\), \(5\%\), and \(0\%\) (energy-only training)—while retaining all energy labels.

We evaluate the resulting models using per-atom energy and force mean absolute error (MAE), as well as training wall-clock time. However, energy and force error metrics alone do not fully capture the quality of an MLIP [56]. We therefore further assess the physical fidelity of the optimal PES identified by each optimizer by performing molecular dynamics (MD) simulations using the trained MLIPs and comparing resulting observables against experimental measurements and ab initio reference data.

4 Results↩︎

4.1 Full Supervision↩︎

4.1.0.1 Accuracy Gains

Under full energy and force supervision, the matrix-structured optimizers studied substantially outperform AdamW, with the magnitude and consistency of the gains depending on the system (Table ¿tbl:tab:main95test95metrics?, Figure 1). On CDP with Allegro, all three optimizers improve over AdamW. SOAP achieves the best energy accuracy, reducing the energy MAE by \(9\%\) whereas SOAP-Muon achieves the best force accuracy, reducing the MAE by \(14\%\). On water with NequIP, SOAP and SOAP-Muon again yield lower MAEs than AdamW. In this case, however, SOAP-Muon gives the best energy accuracy, reducing the MAE by \(24\%\) while SOAP provides the best force accuracy, reducing the MAE by \(19\%\). Muon alone underperforms the AdamW baseline, particularly for energy prediction.

4.1.0.2 Speedup

The accuracy improvements are accompanied by substantial reductions in time-to-accuracy. We compare the wall-clock time required for each optimizer to reach AdamW’s minimum median validation force MAE under full force supervision (Figure 2). On CDP, SOAP reaches that target \(4.9\times\) faster than AdamW, while on water it does so \(5.8\times\) faster. These gains are notable because matrix-structured methods incur additional computation per optimization step relative to AdamW. The wall-clock advantage therefore implies that they reduce the number of epochs needed by an even larger margin. In practice, improved conditioning more than compensates for the higher per-step cost, making SOAP attractive not only for final accuracy but also for training efficiency.

4.2 Reduced Supervision↩︎

4.2.0.1 Energy-Only Supervision

In the energy-only training regime, where forces are excluded from the training objective and therefore provide a stringent test of whether the learned PES captures physically meaningful gradients, the advantage of matrix-structured optimizers becomes more pronounced. On CDP, all three alternatives substantially outperform AdamW. SOAP-Muon achieves the strongest overall results, reducing the energy and force MAE by roughly \(47\%\) and \(60\%\). On water, SOAP performs best, reducing both energy and force errors relative to AdamW by \(26\%\) and \(23\%\), while Muon performs noticeably worse.

4.2.0.2 Sparse Force Supervision

The same ranking largely persists as the number of force-labeled configurations decreases, with the gains over AdamW becoming especially valuable in the low-label regime (Figure 1). On CDP, SOAP-Muon trained with only \(50\%\) force supervision achieves accuracy comparable to AdamW trained with the full force-labeled set, and SOAP shows a similar pattern. Both methods also remain clearly stronger than AdamW at \(10\%\) and \(5\%\) force supervision. This suggests that improved optimization can partially compensate for reduced access to force labels, lowering the amount of expensive force data needed to reach a given level of accuracy. On water, the trend is less regular, with some models trained at \(75\%\) and \(50\%\) force supervision often outperforming their respective models trained on \(100\%\). Nevertheless, SOAP is the most consistently competitive method across supervision levels and is the strongest single default in our experiments.

4.3 Physical Fidelity↩︎

To assess whether the optimizer-dependent differences in MAE translate into meaningful differences in the learned PESs, we evaluated the trained MLIPs in MD simulations.

4.3.0.1 CDP

We followed the MD protocol of [48]. All fully force-supervised models faithfully reproduce the ab initio MD (AIMD) results, matching the radial distribution functions (RDFs) and showing consistent mean squared displacement (MSD) behaviors across optimizers (Appendix Figure 7). The differences become much more pronounced in the sparse-label regime. Notably, at \(5\%\) force supervision, SOAP-Muon remains stable and accurately reproduces the AIMD RDFs and MSD. This is in sharp contrast to AdamW that exhibits catastrophic instability, with trajectories diverging almost immediately and producing nonphysical results under the same conditions, as shown in Figure 3 (the full set of RDFs and MSD is reported in Appendix Figure 9). Moreover, the \(5\%\) SOAP-Muon model reproduces experimental results, yielding a proton diffusion activation energy of \(E_a = 0.42\;\mathrm{eV}\) (Appendix Figure 9h), in good agreement with the experimental range of \(E_a = 0.39\text{--}0.43\;\mathrm{eV}\) [57]–[59].

4.3.0.2 Water

The fully force-supervised models also generate similar structural and dynamical observables across optimizers, with RDFs and MSDs broadly consistent with one another and with the experimental O-O RDF reference (Appendix Figure 8) [60]. For all sparsity percentages tested, both AdamW and SOAP maintained physical fidelity on RDFs and MSD (Appendix Figure 10).

5 Discussion↩︎

The two cases we study differ both in model architecture (NequIP vs.Allegro) and in chemical character (liquid water vs.multicomponent solid CDP). On CDP, all three matrix-based optimizers outperform AdamW, with a clear progression from AdamW to Muon, SOAP and SOAP-Muon. On water, SOAP and SOAP-Muon also improve upon AdamW, with SOAP having a slight advantage; by contrast, Muon performs consistently worse than AdamW across all force-label fractions tested. We note, however, that unlike CDP, the water case required additional tuning of SOAP-Muon (momentum coefficients and singular-value power; see Appendix 8 for details) to achieve these results. Taken together, SOAP and SOAP-Muon outperform AdamW across most metrics, training settings, and systems, with SOAP showing the most robust behavior across systems. For practitioners choosing a single default optimizer to use instead of AdamW without further tuning, our results suggest SOAP.

Muon’s underperformance on water is a notable result that warrants further investigation. 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. Similar concerns about the stability of Muon have been raised in the context of physics-informed neural networks [61], where step-size regulation along dominant spectral modes was proposed as a potential remedy.

5.0.0.1 Limitations

Our work is restricted to training on two systems and two equivariant architectures. Foundation potential training and fine-tuning are interesting open directions, and we expect SOAP to remain effective at the present scale of foundation models, though this requires a systematic study.

6 Conclusion↩︎

Across the systems and training settings studied, matrix-structured optimizers consistently outperform AdamW for MLIP training, with SOAP emerging as the most robust default with the advantage compounding under sparse force supervision. These accuracy gains translate into substantial wall-clock speedups despite per-step preconditioner overhead, and into faithful molecular-dynamics behavior including observables in good agreement with experiment. We argue that training techniques, exemplified by optimizer choice in this work, should be treated as a first-class design axis in MLIP training, alongside architecture and dataset, particularly as the field moves toward universal foundation potentials where label efficiency and convergence speed become increasingly important.

Code Availability↩︎

All optimizer code used in this work will be released in future versions of the open-source nequip framework and associated repositories, including github.com/mir-group/nequip and github.com/mir-group/allegro.

Acknowledgments↩︎

The authors gratefully acknowledge Menghang Wang for valuable insights and support regarding the CDP calculations. The computations in this paper were run on the FASRC Cannon cluster supported by the FAS Division of Science Research Computing Group at Harvard University. An award of computer time was provided by the INCITE program. This research used resources of the Oak Ridge Leadership Computing Facility, which is a DOE Office of Science User Facility supported under Contract DE-AC05-00OR22725. L.Z. was supported by the National Science Foundation Graduate Research Fellowship under Grant No. DGE-2140743. This work was supported by the National Science Foundation through the Harvard University Materials Research Science and Engineering Center Grant No. DMR-2011754.

Impact Statement↩︎

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

7 Optimizer Details↩︎

7.1 AdamW↩︎

We use the PyTorch torch.optim.AdamW implementation.

7.2 Muon↩︎

Algorithm 4 summarizes the Muon update for a matrix-valued parameter. The key operation is the orthogonalization of a momentum-augmented gradient matrix. In practice, this orthogonalization is approximated using the Newton–Schulz algorithm (denoted \(\mathrm{NewtonSchulz5}(\cdot)\) in Algorithm 4), an iterative matrix procedure for approximating the orthogonal factor \(UV^\top\) in the decomposition \(G = U\Sigma V^\mathsf{T}\), i.e., the matrix obtained by replacing the singular values of \(G\) with ones. Because this orthogonalization is defined for two-dimensional weight tensors (and 4-dimensional convolutional parameters), non-matrix parameters are assigned to an auxiliary AdamW optimizer (for simplicity will be referred to as Adam).  [37] further observed empirically that embedding and readout weights are also better optimized with Adam.

Our nequip implementation builds on the official Muon reference implementation and retains the original \(\beta\) momentum hyperparameter. To accommodate the split between parameters partitioned to Muon and those handled by the auxiliary Adam optimizer, we define two parameter groups, one for each optimizer. Broadly following the recommendations of Jordan et al. [37], we assign intermediate-layer linear weight matrices to Muon, while optimizing the type embeddings and readout layers with Adam. A complete breakdown of the parameter groups is provided in Table ¿tbl:tab:nequip-allegro-parameter-groups?.

To implement this split in NequIP, we must also account for the internal parameterization of e3nn layers. In particular, e3nn stores the weights of its equivariant operators as flattened one-dimensional parameter vectors, even when those weights act as structured matrices or tensors in the forward pass. Similar to the Nequix implementation [47], for the subset of e3nn’s Linear parameters assigned to Muon, we first recover the corresponding structured blocks from the flattened representation, perform the Muon update on those reshaped blocks, and then flatten and reassemble them into the original parameter vector. We formalize the instruction-wise action of the e3nn operators below.

7.2.0.1 e3nn and tensor-product weights.

We consider two e3nn.o3 layers: (1) Linear and (2) FullyConnectedTensorProduct (FCTP).

In Linear, each instruction connects an input irrep block to an output irrep block of the same irrep type. The learnable weight for instruction \(j\) is a multiplicity-mixing matrix \[W_j \in \mathbb{R}^{m^{\mathrm{in}}_j \times n^{\mathrm{out}}_j},\] which mixes copies of the irrep but does not act on the internal \(2\ell+1\)-dimensional irrep basis.

In FCTP, each instruction combines an irrep block from the first input, an irrep block from the second input, and an admissible output irrep block. The learnable weight for instruction \(j\) is a multiplicity-mixing tensor \[T_j \in \mathbb{R}^{m^{(1)}_j \times m^{(2)}_j \times n^{\mathrm{out}}_j},\] which mixes copies of the two input irreps and the output irrep but does not act directly on the internal \(2\ell+1\)-dimensional irrep bases; those couplings are fixed by the tensor-product Clebsch–Gordan structure.

Thus, Linear uses learned 2D multiplicity-mixing blocks, while FCTP uses learned 3D blocks over the two input multiplicities and the output multiplicity.

7.2.0.2 Flattened storage.

Although these operators act through structured blocks \(W_j\) or \(T_j\), e3nn stores all learnable weights as a single flattened parameter vector \[w \in \mathbb{R}^{P}, \qquad P = \sum_{j=1}^{J} |\mathbf{s}_j|,\] where \(|\mathbf{s}_j|\) is the number of entries in instruction \(j\)’s block. Equivalently, we define a slice specification \[\mathcal{S} = \{(\mathcal{I}_j,\mathbf{s}_j)\}_{j=1}^{J},\] with \(\mathcal{I}_j\) selecting the entries of \(w\) belonging to instruction \(j\), and \[B_j = \operatorname{reshape}(w_{\mathcal{I}_j}, \mathbf{s}_j),\] where \(B_j = W_j\) and \(\mathbf{s}_j = (m^{\mathrm{in}}_j,n^{\mathrm{out}}_j)\) for Linear, while \(B_j = T_j\) and \(\mathbf{s}_j = (m^{(1)}_j,m^{(2)}_j,n^{\mathrm{out}}_j)\) for FCTP. The forward pass uses the reshaped blocks \(B_j\). Adam-style elementwise optimizers can operate directly on the flattened parameter vector \(w\), because their updates do not depend on the structural interpretation of the parameter shape. By contrast, matrix-structured optimizers such as Muon, SOAP, and SOAP-Muon are shape-sensitive. To use their structured orthogonalization or preconditioning, the relevant slices of \(w\) must be reshaped into their instruction-wise matrix or tensor blocks, updated in that form, and then flattened back into the original parameter vector.

7.3 SOAP↩︎

Algorithm 5 presents the general SOAP update step, simplified by omitting the initialization steps and AdamW weight decay. For the standard SOAP setup, the orthogonalization flag \(\texttt{ortho}\) and normalization flag \(\texttt{normalize}\) are turned off. We use the default hyperparameters, \(\beta_1=0.95\), \(\beta_2=0.95\), Shampoo \(\beta=\beta_2=0.95\) and preconditioning frequency of \(10\) [33].

Unlike AdamW, whose elementwise updates are largely insensitive to how a parameter tensor is reshaped, SOAP uses Shampoo-style structured preconditioning along tensor modes. Consequently, the tensor shape assigned to each parameter block affects the preconditioner and the resulting update.

To leverage this preconditioning, the e3nn linear weights are unstacked and restored to their original two-dimensional forms, following the methodology described in Section 7.2. In the present study, the e3nn fully connected tensor product weights were not sliced and reshaped into three-dimensional tensors. We leave the implementation of such structural modifications for future investigations.

7.4 SOAP-Muon↩︎

Algorithm 5 presents the general SOAP update step with the orthogonalization flag \(\texttt{ortho}\) and the normalization flag \(\texttt{normalize}\) both enabled for SOAP-Muon (initialization and AdamW-style weight decay omitted for simplicity). The resulting update combines Muon-style orthogonalization with RMS normalization of the back-projected AdamW update direction. We use the default preconditioner update frequency of 10 and apply full preconditioning, in which the preconditioner is applied along every mode of the weight tensor. Throughout, the Shampoo preconditioner \(\beta\) was set at \(\beta=\beta_2\).

A key hyperparameter in this update is the singular-value power \(\rho\), which is applied to the singular-value matrix during the orthogonalization step. The official implementation uses the default value \(\rho=0.5\), which requires a full SVD and is therefore computationally more expensive. A cheaper Muon-style alternative is to set \(\rho=0\), which permits the use of Newton–Schulz iteration (as the transform reduces to \(UV^\mathsf{T}\)). However, [41] reported that using \(\rho=0\) can lead to dataset-dependent training instability. We therefore treat \(\rho=0.5\) as the safer default, while using \(\rho=0\) when training remains stable, since this enables the more efficient Newton–Schulz-based orthogonalization.

Finally, we use the parameter grouping scheme described in Appendix 7.2 to determine which parameters receive Muon-style orthogonalization, while the remaining parameters receive standard SOAP updates without the orthogonalization. This grouping also allows us to assign separate \(\beta_1,\beta_2\) hyperparameters to the orthogonalized Muon group and the non-orthogonalized Adam group. This differs from the official SOAP-Muon optimizer, which shares the same momentum settings across all parameters.

8 Training Setup↩︎

8.0.0.1 Sparse Force Training

Given a desired force sparsity level \(s \in [0,1]\), we retain \(N_{\mathrm{keep}} = \operatorname{round}(sN)\) frames with force labels, where \(N\) denotes the total number of frames in the corresponding dataset split. In practice, for each sparsity level, the force-labeled frames are sampled uniformly at random once at training initialization. We then define a binary mask \(m_n \in \{0,1\}\) by setting \(m_n = 1\) for the retained frames and \(m_n = 0\) for the remaining frames, whose force labels are discarded. The same sampled force-labeled subset is used across optimizers for a given data split and random seed.

Under this setting, the loss in Eq. 1 becomes \[\begin{align} \mathcal{L}(\theta) &= \frac{\lambda_E}{N}\sum_{n=1}^N \left(E_\theta(\{\mathbf{r}_j, Z_j\}_n) - E_n\right)^2 + \frac{\lambda_F}{3\sum_{n=1}^{N}m_nN_\text{atoms}^{(n)}}\sum_{n=1}^N m_n\left\|\mathbf{F}_\theta(\{\mathbf{r}_j, Z_j\}_n) - \mathbf{F}_n\right\|^2. \end{align}\] Thus, the force term is normalized only by the number of force components in the force-labeled configurations, rather than by all configurations in the split. The coefficients \(\lambda_E\) and \(\lambda_F\) are kept fixed across sparsity levels using the ratio tuned under full force supervision. The same masking procedure is applied to the training and validation subsets, whereas the test subset is left unmasked.

8.0.0.2 Model Architectures.

The model architectures used in this work are based on previously reported configurations for Allegro on CDP and NequIP on water, which were originally optimized using Adam. The architecture hyperparameters are presented in Table ¿tbl:tab:architecture95hyperparams?.

8.0.0.3 Hyperparameter Search

First, a grid search was performed for the learning rate and weight decay. For each optimizer and task, we first identified a candidate learning-rate order of magnitude \(10^{-p}\), and then swept over \(\{3\times10^{-(p+1)}, 10^{-p}, 3\times10^{-p}\}\) while weight decay was sampled across four values \(\{0, 10^{-5}, 10^{-4}, 10^{-3}\}\). For the water dataset with NequIP, which proved more sensitive to learning rate, we included an additional value of \(5\times10^{-p}\). In Muon and SOAP-Muon, the learning rate and weight decay for both parameter groups were fixed together, although they can be further optimized separately. With these parameters fixed, then, for the full force-supervision setting (\(100\%\)), the energy-force loss coefficient ratio was tuned by sweeping through \(\{1:1, 1:10, 1:100\}\). Final configurations were selected using a validation metric given by a weighted sum of the per-atom energy and force errors, with a \(10\):\(1\) weighting between the two components.. For NequIP on the water dataset, we used a reduce on plateau scheduler following the settings in Batzner et al. [14]. For Allegro on CDP, we opted for a different approach and used a cosine annealing scheduler. All hyperparameter sweeps were run with a single random seed and shortened training schedules: 100 epochs for CDP and 500 epochs for water.

8.0.0.4 Training Runs

The full training runs were conducted using five random seeds (7, 42, 123, 2026, and 1618). The best optimizer hyperparameters identified in the sweeps, and subsequently used in the training runs, are reported in Table ¿tbl:tab:hyperparameters?. Training was run for 1000 epochs for CDP and 2000 epochs for water. Early stopping was applied using the weighted validation objective. The minimum improvement thresholds were \(10^{-6}\) for water and \(10^{-5}\) for CDP, with patience values of 600 and 500 epochs, respectively. Training was conducted on NVIDIA A100-SXM4-80GB and H200 GPUs. To ensure consistency, all wall-clock time comparisons were derived exclusively from training runs performed on A100s without early stopping. All models were compiled [53] and accelerated using OpenEquivariance and cuEquivariance tensor-product kernels [62].

8.0.0.5 SOAP-Muon Stability

For the two model–system configurations considered in this work, we observed different sensitivities to the SOAP-Muon hyperparameters: the momentum parameters \((\beta_1,\beta_2)\) for each parameter group and the singular-value power \(\rho\) applied to the singular-value matrix. Allegro-CDP was relatively robust to these choices, whereas NequIP-water exhibited substantially greater sensitivity and benefited more from further tuning (Figure 6). For the momentum parameters, we searched over \((\beta_1,\beta_2)\in\{(0.9,0.95), (0.95,0.95), (0.95,0.98)\}\), following the combination of values reported in [33], [37], [41]. For the singular-value power, we considered \(\rho\in\{0,0.5\}\) following [41]. We used \(\rho=0\) whenever performance differences were small, in order to avoid the additional cost of the full SVD computation. For Allegro-CDP, we used \((\beta_1,\beta_2)=(0.95,0.98)\) for the Muon parameter group, \((0.9,0.95)\) for the Adam parameter group, and \(\rho=0\). For NequIP-water, we used \((\beta_1,\beta_2)=(0.95,0.95)\) for both the Muon and Adam parameter groups, together with \(\rho=0.5\). Since NequIP-water was more sensitive to these hyperparameters, we treat those settings as a practical default. However, broader evaluations across architectures and datasets are needed to establish more general hyperparameter recommendations. The dataset-specific instability observed for SOAP-Muon is consistent with [41], who found that optimizer stability varied across datasets: some datasets were stable without modification, whereas others required stability-enhancing adjustments. They further showed that datasets stable under the original optimizer were relatively insensitive to these modifications. Nevertheless, our conclusion is based on the limited set of experiments performed here. We note that SOAP-Muon required substantially smaller learning rates than the other optimizers. Learning rates that were optimal for the other optimizers often led to severe training instabilities when used with SOAP-Muon. Throughout our experiments, the best-performing SOAP-Muon learning rate was approximately one order of magnitude lower than those used for the other optimizers (Table ¿tbl:tab:hyperparameters?).

9 Full Results↩︎

References↩︎