Contrastive Regularization of Machine Learning Potentials

기술 요약: 머신러닝 포텐셜의 대조적 정규화 (Contrastive Regularization)

문제 정의
머신러닝 원자 간 포텐셜(MLIP)은 일반적으로 밀도 범함수 이론(DFT) 데이터로부터 에너지와 힘을 예측하기 위해 평균 제곱 오차(MSE)를 최소화하는 회귀 모델로서 학습된다. 이러한 모델들은 홀드아웃(held-out) 테스트 세트에서 화학적 정확도(chemical accuracy)를 달는 데 종종 성공하지만, 저자들은 이 모델들이 분자 역학이나 몬테카를로 시뮬레이션을 구동하는 생성기로 배치되었을 때 빈번하게 실패한다는 점을 입증한다.

핵심 문제는 구조적인 것이 아니라 통계적인 것이다. 고차원 구성 공간(configuration space)에서 유한한 학습 세트는 전체 공간의 극히 미미한 부분만을 커버한다. 표준 MSE 최소화는 데이터가 존재하는 영역에 대해서만 포텐셜을 제약한다. 따라서 표현력이 뛰어난 신경망은 필연적으로 방문되지 않은 나머지 구성 공간에 가짜 저에너지 극소점(spurious low-energy minima)을 형성하게 된다. 샘플러로 사용될 때, 궤적(trajectories)은 이러한 "유령(ghost)" 극소점으로 흘러 들어가며, 결과적으로 참조 데이터와는 판이하게 다른 비물리적인 구성과 열역학적 관측치(예: 자유 에너지 프로파일, 구조적 분포)를 초래한다. 이는 점 단위(pointwise)의 에너지 및 힘 예측이 정확함에도 불구하고 발생하는 현상이다.

방법론: 대조적 정규화 MSE (CRMSE)
이를 해결하기 위해 저자들은 모델의 암시적 볼츠만 분포와 타겟 분포 사이의 쿨백-라이블러(Kullback–Leibler, KL) 발산을 기반으로 한 대조적 항을 표준 MSE 손실 함수에 추가하는 **대조적 정규화 MSE (CRMSE)**를 도입한다.

주요 방법론적 구성 요소는 다음과 같다:

• 목적 함수: 전체 손실 함수는 $L_{CRMSE} = L_{MSE} + \lambda L_{CR}$이다. MSE 항은 학습 데이터에 대한 정확도를 유지하며, 대조적 항($L_{CR}$)은 확률 분포를 최적화한다.
• 대조 메커니즘: 대조적 항은 KL 발산의 최소화로부터 유도된다. 이는 변분 자유 에너지 목적 함수로서 작용한다. 즉, "양성(positive)" 샘플(학습 데이터)의 에너지는 낮추고, "음성(negative)" 샘플(모델의 현재 분포에서 추출된 구성)의 에너지는 높인다.
• 자기 일관적 샘플링 (Self-Consistent Sampling): 모델의 분포는 다루기 까다롭기 때문에(intractable), 음성 샘플은 비조정 랑주뱅 알고리즘(Unadjusted Langevin Algorithm, ULA)을 사용하여 생성된다. 저자들은 **지속적 대조 발산(Persistent Contrastive Divergence, PCD)**을 채택하여, 리플레이 버퍼(replay buffer)에 음성 구성을 유지하고 매 반복마다 몇 단계의 ULA를 수행하여 업데이트한다.
• 사후 학습(Post-Training) 특성: CRMSE는 이미 MSE로 최적화된 포텐셜에 적용되는 사후 학습 단계이다. 이는 새로운 ab initio 데이터를 요구하지 않으며 구조적 변경도 필요하지 않다. 모델은 시뮬레이션 중 미끄러져 들어가는 구성을 식별하고 그 에너지를 높임으로써 물리적 베이슨(basin) 주위에 보호 장벽을 형성함으로써 효과적으로 "자신의 실수를 스스로 레이블링"한다.

주요 결과
이 방법은 GNN-LF(Local Frame을 가진 그래프 신경망) 아키텍처를 사용하여 MD17 데이터셋의 에탄올과 아스피린 분자에 대해 검증되었다.

1. 표준 MSE의 실패: 표준 MSE로 학습된 모델은 홀드아웃 에너지와 힘에 대해 높은 정확도를 달성했다. 그러나 메트로폴리스-수정 랑주뱅 알고리즘(MALA) 시뮬레이션에 사용했을 때, 궤적은 물리적 베이슨을 빠르게 탈출하여 참조 에너지보다 수십 kcal/mol 낮은 가짜 극소점에 안착했다. 이는 비물리적인 원자 간 거리와 잘못된 에너지 분포를 초래했다.
2. CRMSE를 통한 교정:
   • 샘플링 안정성: CRMSE 사후 학습은 샘플러를 물리적 베이슨 내에 가두었다. 샘플링된 에너지 분포와 원자 간 거리 히스토그램은 DFT 참조 데이터와 거의 정량적인 일치를 회복했다.
   • 열역학적 관측치: 결정적으로, CRMSE는 디헤드럴 각도(에탄올의 H-C-C-O 및 아스피린의 에스테르 디헤드럴)에 대한 자유 에너지 프로파일을 DFT와 거의 정량적으로 일치하도록 복구했다. 이는 토션 장벽(torsional barriers)의 과소평가를 교정하였으며, 독립적인 시뮬레이션 체인 간의 분산(inter-chain dispersion)을 통계적 노이즈 수준 이내로 수렴시켰다.
   • 정확도 보존: 이러한 교정은 학습 매니폴드 상의 예측 정확도를 저하시키지 않았다. 홀드아웃 데이터에 대한 에너지 및 힘 오차는 여전히 화학적 정확도 범위 내에 머물렀다.
   • 강건성: 이 방법은 학습 세트를 950개에서 200개로 급격히 줄인 경우에도 효과적임을 보여주어, 데이터 희소성에 대한 개선된 데이터 효율성과 강건성을 입증했다.

의의 및 주장
본 논문은 신뢰할 수 있는 샘플링이 학습 데이터의 양보다는 모델이 생성하는 분포 자체를 학습시키는 것에 달려 있다고 주장한다. 저자들은 분포 수준의 학습이 단순한 정밀화 작업이 아닌 별개의 요구사항이라고 논한다.

• 별개의 요구사항: 본 연구는 점 단위의 정확도가 올바른 샘플링 분포를 보장하지 못한다는 점을 강조한다. 샘플링되어야 할 포텐셜은 반드시 자신의 샘플링 분포를 재현하도록 학습되어야 한다.
• 상보성: CRMSE는 능동 학습(Active Learning)과 상보적인 것으로 제시된다. 능동 학습이 오프-매니폴드(off-manifold) 영역의 정확도를 높이기 위해 새로운 ab initio 데이터를 주입한다면, CRMSE는 새로운 데이터 없이도 샘플링 분포의 올바른 지지 집합(support)을 복구하며, 에너지를 높임으로써 가짜 극소점을 억제한다.
• 일반성: 저자들은 이러한 교훈이 학습된 회귀 모델이 생성기로 배치되는 광범위한 AI 설정에 적용될 수 있다고 제안한다. 이 방법은 이미 표준 회 regression 지표에 의해 "잘 최적화된" 포텐셜을 저렴한 사후 학습 단계만을 사용하여, 추가적인 참조 계산 없이도 수리할 수 있다는 점에서 특히 중요하다.

저자들은 결론적으로, 생성기로 사용되는 대리 모델(surrogate models)의 경우, 훈련 목적 함수가 포텐셜에 의해 정의된 전체 확률 밀도를 고려해야 하며, 이는 표준 MSE 최소화 방식에서는 구조적으로 누락되어 있다는 점을 명시한다.