Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentials

Abstract: Machine learning interatomic potentials (MLIPs) have become a hallmark of AI for scientific simulation. While efforts on new architectures and datasets have led to increasingly accurate and general models, the choice of optimizer for training has largely remained unexplored, defaulting to Adam and its variants in the community. Here, we implement and systematically compare a class of recently proposed matrix-structured optimizers, including Muon, SOAP, and the hybrid SOAP-Muon, for training NequIP and Allegro MLIP models. We find that these optimizers can substantially outperform Adam in both convergence speed and final accuracy. SOAP and SOAP-Muon emerge as robust and consistently strong methods, while Muon only provides partial gains relative to Adam. The improvements are particularly pronounced under partial force supervision. Our results indicate that optimizer choice is an overlooked yet impactful design axis for MLIPs.

• The paper demonstrates that SOAP and SOAP-Muon optimizers significantly reduce energy and force MAE compared to AdamW.
• It shows that label-efficient training is achievable, with SOAP-Muon matching full-label performance at only 50% force supervision.
• The study reveals that matrix-structured optimizers accelerate wall-clock training times by up to 5.8× while maintaining physical fidelity in MD simulations.

Matrix-Structured Optimizers for Machine Learning Interatomic Potentials: An Expert Analysis

Context and Motivation

Recent progress in machine learning interatomic potentials (MLIPs) has been dominated by advances in two primary axes: model architectures—particularly E(3)-equivariant GNNs like NequIP and Allegro—and the assembly of large, diverse quantum chemical datasets. However, the optimization landscape that governs MLIP training has not received equivalent systematic scrutiny. Despite success in other domains with matrix-structured optimizers—including explicit Kronecker-factored preconditioners (e.g., Shampoo, SOAP) and matrix-orthogonalization methods (e.g., Muon)—MLIP training pipelines almost universally default to Adam or its variants. This work rigorously investigates the impact of optimizer choice, with a focus on SOAP, Muon, and the hybrid SOAP-Muon, on both convergence and final generalization in MLIP settings.

Methodology

The comparative study integrates Muon, SOAP, and SOAP-Muon into the NequIP framework and evaluates their efficacy relative to AdamW in two canonical MLIP scenarios: liquid water (modeled by NequIP) and the multicomponent solid acid electrolyte CsH₂PO₄ (CDP, modeled by Allegro). Critically, the analysis systematically ablates the force-label fraction, ranging from full (100%) to energy-only (0%) supervision, to quantify optimizer behavior in both data-rich and label-sparse regimes. Each optimizer's parameterization and application to e3nn layers is technically refined to ensure correct shape preservation and alignment with structured updates. For SOAP and SOAP-Muon, Shampoo-style preconditioning is applied with careful eigenbasis tracking; for Muon, orthogonalization is integrated for appropriate parameter subsets with auxiliary AdamW for others.

Models are evaluated via per-atom energy and force mean absolute error (MAE), wall-clock time to accuracy, and by examining physical fidelity through MD simulations, reporting observables such as radial distribution functions (RDF) and mean squared displacement (MSD). Convergence analysis is performed with hyperparameter tuning and multi-seed aggregation to ensure robustness.

Main Findings

Accuracy and Speed

Matrix-structured optimizers uniformly improve upon AdamW in both convergence rate and final test MAE under joint energy and force supervision. For the Allegro/CDP system:

• Energy MAE: SOAP reduces AdamW's error by 9%; SOAP-Muon achieves comparable results, while Muon offers only partial improvement.
• Force MAE: SOAP-Muon achieves a 14% reduction versus AdamW.

For NequIP/water:

• Energy MAE: SOAP-Muon reduces AdamW's error by 24%; SOAP achieves a notable improvement as well.

The reduction in training wall-clock time is substantial; for SOAP, the time to AdamW's minimum median force MAE is 4.9x (CDP) to 5.8x (water) faster. These speedups far outweigh the modest per-step computational overhead of the matrix structure, indicating improved loss surface conditioning.

Sparse Supervision

In the energy-only regime, matrix-structured optimizers yield dramatic improvements over AdamW. On CDP, SOAP-Muon reduces energy MAE by ~47% and force MAE by ~60%; on water, SOAP reduces these errors by 26% and 23% respectively. Notably, SOAP-Muon trained with 50% of force labels matches AdamW trained with a fully labeled set—demonstrating dramatic label efficiency gains.

At extremely low force-label fractions (e.g., 5%), AdamW-based models for CDP become dynamically unstable in MD, failing to reproduce physical observables, while SOAP-Muon preserves both physical fidelity and agreement with experimental activation energies.

Robustness and Tuning

Results on NequIP/water highlight a caveat: Muon and SOAP-Muon do not universally dominate. In this setting, Muon underperforms AdamW, and SOAP-Muon requires careful tuning (especially for momentum and SVD hyperparameters) to guarantee robust convergence. This instability is consistent with observations in physics-informed ML and underlines the potential pitfalls of matrix orthogonalization without additional regularization.

Physical Fidelity

Empirical gains in MAE translate to pronounced improvements in the faithfulness of the learned PES, particularly in label-sparse regimes. For CDP, SOAP-Muon maintains realistic RDFs and MSDs, matching AIMD benchmark and experimental activation energies ($E_a = 0.42$ eV vs. experimental $0.39$–$0.43$ eV) even with only 5% force supervision; AdamW diverges catastrophically. For water, both SOAP and AdamW produce stable and accurate MD trajectories, although SOAP achieves lower errors.

Implications and Future Prospects

This work demonstrates that optimizer selection is a critical, underexplored axis for MLIP training—on par with architecture and dataset design. Matrix-structured optimizers, especially SOAP and SOAP-Muon, offer meaningful advances in label efficiency, convergence rate, and physical reliability, especially in regimes of reduced force supervision where label cost is high (e.g., CCSD(T), QMC-based datasets).

From a theoretical standpoint, the results underscore the power of preconditioning by capturing local curvature structure, reducing anisotropy in the loss geometry, and enabling rapid convergence even at low supervision. Hybrid approaches (e.g., SOAP-Muon) further extend this by blending adaptive steps with orthogonalization for robustness. However, their stability is problem-dependent, demanding further research into adaptive spectral regularization and parameter group assignments.

As MLIPs scale toward universal foundation models and broader atomistic domains, training on ever more expensive and sparse reference data, matrix-structured optimizers will likely be essential for practical, label-efficient, and robust PES learning.

Conclusion

Systematic benchmarking reveals that SOAP and SOAP-Muon consistently outperform AdamW for MLIP training in both accuracy and efficiency, with gains amplified in data-sparse regimes. The optimizer should be treated as a first-order hyperparameter alongside the architecture and dataset, especially as the field transitions to universal and label-limited MLIPs. Future directions include extension to larger, cross-domain foundation models and robustification of hybrid methods for broad applicability.

Reference: "Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentials" (2607.02499).