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

제목: Adam을 넘어: 더 빠르고 라벨 효율적인 머신러닝 원자 간 포텐셜 학습을 위한 SOAP와 Muon

문제 정의
머신러닝 원자 간 포텐셜(MLIP)은 물의 열역학부터 생체 분자 조립에 이르기까지 원자 단위 시뮬레이션을 가능하게 함으로써 화학 및 재료 과학 분야에서 필수적인 요소가 되었습니다. 이 분야는 새로운 아키텍처(예: NequIP 및 Allegro와 같은 등변 그래프 신경망)와 대규모 데이터셋을 통해 크게 발전해 왔으나, 학습 최적화 기술은 정체된 상태에 머물러 있습니다. 커뮤니티는 주로 Adam과 그 변형들(예: AdamW)을 기본값으로 사용합니다. 본 논문은 옵티마이저(optimizer)의 선택이 간과되었지만 영향력 있는 설계 축임을 주장합니다. 구체적으로, 저자들은 행렬 구조 기반의 옵티마이저가 대각 성분 기반의 그래디언트 스케일링에 의존하는 대신 신경망 가중치의 텐서 구조를 활용함으로써, 특히 힘(force) 라벨이 비싸거나 희소한 환경에서 수렴 속도와 최종 정확도를 개선할 수 있는지 조사합니다.

방법론
저자들은 최근 제안된 세 가지 행렬 구조 옵티마이자를 표준 AdamW 베이스라인과 비교하여 체계적으로 벤치마킹하고 구현했습니다:

1. Muon: 뉴턴-슐츠(Newton–Schulz) 반복법을 사용하여 그래디언트를 반직교 행렬(semi-orthogonal matrices)로 투영하는 직교화 행렬 업데이트 방식입니다.
2. SOAP: 프리컨디셔너(preconditioner)의 고유 공간(eigenspace)에서 Adafactor 스타일의 업데이트를 수행한 후, 원래의 파라미터 공간으로 다시 투영하는 방식입니다.
3. SOption-Muon: SOAP 업데이트와 Muon 방식의 직교화 단계를 결합한 하이브리드 접근 방식입니다.

이 옵티마이저들은 nequip 프레임워크에 통합되었으며, 행렬 구조 업데이트를 적용하기 위해 e3nn 레이어(Linear 및 FullyConnectedTensorProduct)에 필요한 특정 파라미터 리셰이핑(reshaping) 과정을 거쳐 두 가지 서로 다른 물리 시스템에 적용되었습니다:

• 액체 물(Liquid Water): NequIP 아키텍처를 사용하여 모델링되었습니다.
• 고체 산성 전해질(CsH₂PO₄ 또는 CDP): Allegro 아키텍처를 사용하여 모델링되었습니다.

본 연구는 공정한 비교를 위해 체계적인 하이퍼파라미터 튜닝 프로토콜을 채택했습니다. 결정적으로, 저자들은 힘 계산(높은 수준의 이론인 결합 클러스터(coupled cluster)나 확산 몬테카를로(diffusion Monte Carlo)와 같이 비용이 많이 드는 경우)이 제한적인 상황을 테스트하기 위해, 전체 감독(100% 힘 라벨)부터 에너지 전용 학습(0% 힘 라벨) 및 중간 단계의 희소 구간(75%, 50%, 10%, 5%)에 이르기까지 다양한 수준의 힘 감독 하에서 성능을 평가했습니다.

주요 기여

• 통합 및 벤치마킹: 저자들은 행렬 구조 업데이트를 적용하기 위한 특정 파라미터 리셰이핑을 처리하며 Muon, SOAP, SOAP-Muon을 nequip 프레임워크에 성공적으로 통합했습니다.
• 성능 분석: 이들이 행렬 구조 옵티마이저가 AdamW 베이스라인보다 수렴 속도와 최종 정확도 측면에서 실질적으로 우수할 수 있음을 입증했습니다.
• 라벨 효율성: SOAP와 SOAP-Muon은 힘 감독이 크게 감소하더라도 높은 정확도를 유지함을 보여주었으며, 이는 비용이 많이 드는 힘 라벨을 적게 사용하면서도 정확한 포텐셜 에너지 표면(PES)을 학습할 수 있는 경로를 제시합니다.
• 물리적 충실도 검증: 평균 절대 오차(MAE)를 넘어, 저자들은 학습된 포텐셜의 물리적 정확성을 검증하기 위해 분자 역학(MD) 시뮬레이션을 실행하고, 관측 가능한 물리량(방사 분포 함수(RDF), 평균 제곱 변위(MSD), 확산 활성화 에너지)을 ab initio 참조 데이터 및 실험 측정값과 비교했습니다.

결과

• 전체 감독: 에너지 및 힘에 대한 전체 감독 하에서 SOAP와 SOAP-Muon은 일관되적으로 AdamW보다 우수한 성능을 보였습니다.
  • CDP 시스템(Allegro)에서 SOAP는 최고의 에너지 정확도(MAE 9% 감소)를 달了했고, SOAP-Muon은 최고의 힘 정확도(14% 감소)를 달성했습니다.
  • 물 시스템(Nequip)에서 SOAP-Muon은 최고의 에너지 정확도(24% 감소)를, SOAP는 최고의 힘 정확도(19% 감소)를 제공했습니다.
  • Muon 단독으로는 물 시스템에서는 AdamW보다 낮은 성능을 보였으나, CDP에서는 개선된 결과를 보였습니다.
• 수렴 속도: 행렬 구조 옵티마이저들은 AdamW의 최소 검증 힘 MAE에 훨씬 빠르게 도달했습니다. CDP의 경우 SOAP는 4.9배 빨랐고, 물의 경우 5.8배 빨랐습니다. 이러한 속도 향상은 최적화 단계당 계산 비용이 더 높음에도 불구하고 발생했으며, 이는 필요한 에포크(epoch) 수가 실질적으로 감소했음을 의미합니다.
• 희소 힘 감독: 행렬 구조 옵티마이저의 이점은 힘 라벨이 줄어듦에 따라 더욱 두드러졌습니다.
  • 에너지 전용 학습: 힘 라벨이 없는 경우, SOAP와 SOAP-Muon은 AdamW를 크게 압도했습니다. 예를 들어, CDP에서 SOAP-Muon은 AdamW 대비 에너지 및 힘 MAE를 각각 약 47%와 60% 감소시켰습니다.
  • 희소 라벨: CDP에서 50%의 힘 라벨만 사용하여 학습한 SOAP-Muon은 100%의 힘 라벨을 사용한 AdamW와 대등한 정확도를 달성했습니다.
  • 안정성: CDP의 5% 힘 감독 조건에서 결정적인 발견이 있었습니다. AdamW 모델은 MD 시뮬레이션 중에 심각한 불안정성을 보이며 비물리적인 결과를 생성했습니다. 반면, SOAP-Muon 모델은 안정성을 유지하며 ab initio RDF 및 MSD를 정확하게 재현했습니다.
• 물리적 관측량: SOAP-Muon으로 학습된 모델(단 5%의 힘 라벨만 사용했을 때도)은 실험적 범위(0.39–0.43 eV)와 잘 일치하는 양성자 확산 활성화 에너지(0.42 eV)를 재현한 반면, 불안정한 AdamW 모델은 선형 확산 영역에 도달하는 데 실패했습니다.

의의 및 주장
본 논문은 옵티마이저의 선택이 아키텍처 및 데이터셋 선택과 함께 중요한, 그러나 종종 간과되는 "제1순위 설계 축"으로 취급되어야 한다고 결론짓습니다.

• 강건성(Robustness): SOAP는 다양한 시스템과 감독 수준에서 가장 견고한 기본 옵티마이저로 나타났습니다. SOAP-Muon은 특정 설정에서 가장 강력한 결과를 보여주지만, 안정성을 확보하기 위해 모멘텀 계수 및 특이값 거듭제곱(singular-value power)에 대한 더 세심한 하이퍼파라미터 튜닝이 필요합니다.
• 실질적 영향: 희소한 힘 라벨을 사용하여 높은 충실도의 MLIP를 학습할 수 있는 능력은, 계산 비용이 매우 높은(예: DFT를 넘어서는 수준) 참조 데이터를 사용하는 경우에 실질적인 해결책을 제공합니다.
• 한계점: 저자들은 본 연구가 두 가지 시스템과 두 가지 아키텍처로 제한되어 있음을 인정합니다. 또한 SOAP가 현재 규모에서는 효과적인 것으로 보이지만, 범용 파운데이션 포텐셜 학습에 대한 효용성을 확인하기 위해서는 추가적인 체계적 연구가 필요하다고 언급했습니다. 또한 Muon이 물 시스템에서 보이는 불안정성에 대해 추가 조사가 필요함을 강조하며, 이는 직교화 단계가 특정 맥락에서 성능 저하의 주요 원인이 될 수 있음을 시사합니다.

저자들은 자신들의 코드가 향-후 nequip 오픈 소스 프레임워크의 차기 버전에 공개되어, 이러한 최적화 전략의 광범위한 채택을 촉진할 것이라고 밝혔습니다.