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

Machine learning interatomic potentials (MLIPs) are essential for modern atomistic simulations, yet their training pipelines almost exclusively rely on Adam or AdamW. This paper investigates whether matrix-structured optimizers—which exploit the structure of weight tensors to provide more effective preconditioning—can improve the training efficiency and accuracy of MLIPs compared to the standard AdamW baseline.

The authors integrate three matrix-structured optimizers—Muon, SOAP, and the hybrid SOAP-Muon—into the NequIP and Allegro MLIP frameworks. They systematically benchmark these methods against AdamW on two physically distinct systems: liquid water (modeled with NequIP) and solid-state cesium dihydrogen phosphate (CDP, modeled with Allegro). The study evaluates performance across varying levels of force supervision, ranging from full force labels to energy-only training, and assesses the physical fidelity of the resulting models through molecular dynamics (MD) simulations.

Matrix-structured optimizers consistently outperform AdamW in both convergence speed and final accuracy. SOAP and SOAP-Muon emerge as the most robust methods, with SOAP-Muon achieving the strongest results in selected settings. Notably, these optimizers provide significant wall-clock speedups, reaching target accuracy levels up to 5.8 times faster than AdamW.

Under sparse force supervision, the benefits of these optimizers become even more pronounced. SOAP-Muon trained with only 50% of force labels can match the accuracy of AdamW trained with 100% of force labels. Furthermore, in the extreme case of 5% force supervision, the SOAP-Muon model remains stable and physically faithful, whereas the AdamW model exhibits catastrophic instability. These results suggest that advanced optimization techniques can effectively compensate for the high cost of obtaining force labels in quantum chemical datasets.