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

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

Problem Statement
Machine Learning Interatomic Potentials (MLIPs) have become essential for atomistic simulations in chemistry and materials science, enabling studies from water thermodynamics to biomolecular assemblies. While the field has advanced significantly through new architectures (e.g., equivariant graph neural networks like NequIP and Allegro) and large-scale datasets, the training optimization landscape has remained static. The community predominantly defaults to Adam and its variants (e.g., AdamW) for training. This paper posits that the choice of optimizer is an overlooked yet impactful design axis. Specifically, the authors investigate whether matrix-structured optimizers, which exploit the tensor structure of neural network weights rather than relying on diagonal gradient scaling, can improve convergence speed and final accuracy, particularly in regimes where force labels are expensive or sparse.

Methodology
The authors implement and systematically benchmark three recently proposed matrix-structured optimizers against the standard AdamW baseline:

1. Muon: An orthogonalized matrix-update method using Newton–Schulz iterations to project gradients onto semi-orthogonal matrices.
2. SOAP: An optimizer that performs Adafactor-style updates in the eigenspace of the preconditioner, followed by a projection back to the original parameter space.
3. SOAP-Muon: A hybrid approach combining the SOAP update with a Muon-inspired orthogonalization step.

These optimizers were integrated into the nequip framework and applied to two distinct physical systems:

• Liquid Water: Modeled using the NequIP architecture.
• Solid Acid Electrolyte (CsH₂PO₄ or CDP): Modeled using the Allegro architecture.

The study employed a systematic hyperparameter tuning protocol to ensure fair comparisons. Crucially, the authors evaluated performance under varying levels of force supervision, ranging from full supervision (100% force labels) to energy-only training (0% force labels), and intermediate sparse regimes (75%, 50%, 10%, 5%). This setup tests the optimizers' ability to learn accurate potential energy surfaces (PES) when force calculations—often costly in high-level theories like coupled cluster or diffusion Monte Carlo—are limited.

Key Contributions

• Integration and Benchmarking: The authors successfully integrated Muon, SOAP, and SOAP-Muon into the nequip framework, handling the specific parameter reshaping required for e3nn layers (Linear and FullyConnectedTensorProduct) to apply matrix-structured updates.
• Performance Analysis: They demonstrated that these optimizers can substantially outperform AdamW in both convergence speed and final accuracy.
• Label Efficiency: The study highlights that SOAP and SOAP-Muon retain high accuracy even when force supervision is significantly reduced, suggesting a pathway to training accurate MLIPs with fewer expensive force labels.
• Physical Fidelity Validation: Beyond mean absolute error (MAE), the authors validated the physical correctness of the learned potentials by running molecular dynamics (MD) simulations and comparing observables (Radial Distribution Functions, Mean Squared Displacement, and diffusion activation energies) against ab initio reference data and experimental measurements.

Results

• Full Supervision: Under full energy and force supervision, SOAP and SOAP-Muon consistently outperformed AdamW.
  • On the CDP system (Allegro), SOAP achieved the best energy accuracy (9% reduction in MAE), while SOAP-Muon achieved the best force accuracy (14% reduction).
  • On the water system (NequIP), SOAP-Muon provided the best energy accuracy (24% reduction), and SOAP provided the best force accuracy (19% reduction).
  • Muon alone showed mixed results, underperforming AdamW on the water system but improving on CDP.
• Convergence Speed: The matrix-structured optimizers reached AdamW's minimum validation force MAE significantly faster. On CDP, SOAP was 4.9× faster; on water, it was 5.8× faster. This speedup occurred despite the higher computational cost per optimization step, implying a substantial reduction in the number of epochs required.
• Sparse Force Supervision: The advantages of matrix-structured optimizers became more pronounced as force labels were reduced.
  • Energy-Only Training: In the absence of force labels, SOAP and SOAP-Muon substantially outperformed AdamW. For example, on CDP, SOAP-Muon reduced energy and force MAE by ~47% and ~60% respectively compared to AdamW.
  • Sparse Labels: On CDP, SOAP-Muon trained with only 50% of force labels achieved accuracy comparable to AdamW trained with 100%.
  • Stability: A critical finding was observed at 5% force supervision on CDP. The AdamW model became catastrophically unstable during MD simulations, producing nonphysical results. In contrast, the SOAP-Muon model remained stable and accurately reproduced ab initio RDFs and MSDs.
• Physical Observables: Models trained with SOAP-Muon (even with only 5% force labels) reproduced experimental proton diffusion activation energies (0.42 eV) in good agreement with the experimental range (0.39–0.43 eV), whereas the unstable AdamW model failed to reach the linear diffusion regime.

Significance and Claims
The paper concludes that optimizer choice is a critical, yet often neglected, component of MLIP design. The authors argue that training techniques should be treated as a "first-class design axis" alongside architecture and dataset selection.

• Robustness: SOAP emerges as the most robust default optimizer across different systems and supervision levels. SOAP-Muon shows the strongest results in specific settings but requires more careful hyperparameter tuning (particularly regarding momentum coefficients and singular-value power) to ensure stability.
• Practical Impact: The ability of these optimizers to accelerate training and maintain accuracy with sparse force labels offers a practical solution for training high-fidelity MLIPs on reference data where force calculations are prohibitively expensive (e.g., beyond DFT).
• Limitations: The authors acknowledge that their work is restricted to two systems and two architectures. They note that while SOAP appears effective at the current scale, further systematic study is needed to confirm its efficacy in training universal foundation potentials. They also highlight that Muon's instability on the water system warrants further investigation, suggesting that the orthogonalization step may be a primary source of degradation in certain contexts.

The authors state that their code will be released in future versions of the open-source nequip framework, facilitating broader adoption of these optimization strategies.