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

Motivation

• Known: MLIP 연구는 주로 새로운 architecture(NequIP, Allegro 등 equivariant graph neural network)와 새로운 dataset 구축에 집중되어 왔으며, LLM 학습에서는 Shampoo, SOAP, Muon과 같은 matrix-structured optimizer들이 Adam 대비 더 빠른 수렴을 보여 최근 주목받고 있다.
• Gap: MLIP 학습 파이프라인은 여전히 대부분 Adam 계열 optimizer를 기본값으로 사용하며, optimizer 선택이라는 세 번째 설계 축은 거의 탐구되지 않았다. 기존 연구(Liu et al. 2026b)는 diagonal preconditioning 기반 optimizer만 벤치마킹했고, Koker et al.(2025)은 Muon의 가능성을 보였으나 충분한 ablation 없이 제한적으로 다루었다.
• Why: DFT 이상의 고급 이론 수준(coupled cluster, diffusion Monte Carlo)에서는 force label 계산 비용이 매우 크기 때문에, force-sparse 환경에서도 잘 작동하는 optimizer는 label efficiency를 높여 더 높은 수준의 참조 데이터를 활용한 MLIP 학습을 가능하게 하는 실질적 중요성을 가진다.
• Approach: liquid water(NequIP)와 solid acid electrolyte CsH2PO4(Allegro)라는 두 물리적으로 중요한 시스템에서 Muon, SOAP, SOAP-Muon을 AdamW와 비교하고, 전체 force supervision부터 energy-only training까지 다양한 force label 비율에서의 성능을 체계적으로 평가한다.

Achievement

1. Optimizer 성능 비교: SOAP와 SOAP-Muon이 energy 및 force accuracy를 일관되게 개선하고 수렴 속도를 가속화하며, SOAP가 가장 강건한 성능을 보이고 SOAP-Muon이 특정 설정에서 최고 성능을 달성함을 보였다. Muon은 한 시스템에서만 AdamW 대비 이득을 제공했다.
2. Force label 효율성: force supervision이 줄어들 때도 SOAP와 SOAP-Muon이 높은 정확도를 유지하며, 특히 SOAP-Muon이 50%의 force label만으로 AdamW의 100% force label 성능에 필적함을 보여 force label 요구량을 줄일 수 있는 가능성을 제시했다.
3. 물리적 타당성 검증: 학습된 MLIP이 ab initio 계산 및 실험적 관측값을 잘 재현함을 확인했으며, 특히 한 시스템에서 SOAP-Muon은 5%의 force label만으로도 물리적 신뢰도를 유지한 반면, 동일 조건의 AdamW 모델은 불안정해졌다.

How

• nequip framework 내에 Muon, SOAP, SOAP-Muon optimizer를 구현·통합
• diagonal preconditioning(Adam)과 대비되는 matrix-structured/Kronecker-structured preconditioning(SOAP, Shampoo 계열) 및 orthogonalized matrix-update(Muon) 방식을 적용
• NequIP(liquid water)와 Allegro(CsH2PO4) 두 시스템에서 full force supervision부터 부분적 force supervision, energy-only training까지 다양한 조건에서 test MAE, wall-clock convergence 비교
• radial distribution function(RDF) 및 MD 시뮬레이션 기반 물리적 observable(예: CDP 시스템의 물성)을 통해 학습된 모델의 물리적 충실도(fidelity) 검증

Originality

• MLIP 학습에서 그동안 거의 탐구되지 않았던 optimizer라는 설계 축을 체계적으로 조명
• 최근 LLM 학습에서 각광받는 matrix-structured optimizer(Muon, SOAP, SOAP-Muon)를 MLIP 도메인에 최초로 광범위하게 적용 및 벤치마킹
• 특히 force-sparse regime(coupled cluster, diffusion Monte Carlo 등 force 계산이 비싼 고급 이론 수준을 겨냥)에서의 optimizer 성능을 탐구한 점이 실용적으로 참신함

Limitation & Further Study

• 두 개의 시스템(liquid water, CsH2PO4)과 두 architecture(NequIP, Allegro)에 한정된 실험으로, 더 다양한 화학 시스템·universal potential 규모로의 일반화 가능성은 추가 검증이 필요
• Muon이 한 시스템에서만 이득을 보인 이유에 대한 심층적 분석(예: 시스템별 loss landscape 차이)이 부족
• hyperparameter tuning의 공정성(각 optimizer별 튜닝 노력의 동등성)에 대한 논의가 제한적일 수 있음
• 향후 연구로 더 큰 규모의 foundation model 학습이나 다양한 원소 조합에 대한 optimizer의 스케일링 거동 분석이 필요

Evaluation

총평: MLIP 학습에서 그동안 간과되었던 optimizer 선택이라는 설계 축을 체계적으로 조명하고, 특히 force label이 부족한 실용적 상황에서 SOAP-Muon과 같은 matrix-structured optimizer가 실질적인 label efficiency 향상을 가져올 수 있음을 보인 의미 있는 실증 연구이다.