Bayesian Neural Networks versus deep ensembles for uncertainty quantification in machine learning interatomic potentials

기술 요약: 머신러닝 원자 간 포텐셜의 불확실성 정량화를 위한 베이지안 신경망과 딥 앙상블의 비교

문제 정의
머신러킹 원자 간 포텐셜(MLIP)은 밀도 범함수 이론(DFT)과 같은 ab initio 방법의 계산 효율적인 대안으로서 원자 단위 모델링에 필수적이다. 그러나 표준 딥러닝 아키텍처는 결정론적이며 내재적인 예측 불확실성 정량화 메커니즘이 부족하다. 이러한 한계는 모델이 분포 외(out-of-distribution) 데이터에 자주 적용되는 재료 과학 분야에서 매우 중요하다. 신뢰할 수 있는 불확실성 추정치가 없다면 시뮬레이션은 오도하는 결과를 초래할 위험이 있으며, 모델 개선을 위한 능동 학습 전략은 지침을 얻지 못하게 된다. 딥 앙상블(DE)과 베이지안 신경망(BNN)이 불확실성을 해결하기 위해 확립된 전략이지만, 특히 에너지와 힘의 공동 예측(joint energy and force predictions)이 변하는 데이터 레짐(regime) 전반에서 MLIP의 특정 맥락 내에서의 상대적 성능, 즉 비교 성능에 대한 체계적인 평가가 필요하다.

방법론
저자들은 ænet-PyTorch 프레임워크 내에서 두 가지 주요 불확실성 정량화 접근 방식을 구현하고 비교하였다:

1. 딥 앙상블 (DE): Lakshminarayanan 등의 공식을 따라, 서로 다른 무작위 시드로 초기화된 5개의 독립적인 신경망을 훈련했다. 불확실성은 앙상블 예측의 표본 표준 편차로 정량화되었다.
2. 변분 베이지안 신경망 (VBNN): 저자들은 Pyro 확률 프로그래밍 프레임워크를 사용하여 변분 추론(VI)을 통한 BNN을 구현했다. 두 가지 특정 변분 가이드(variational guides)가 테스트되었다:
   • 국소 재매개변수화 기법 (LRT): 그래디언트 분산을 줄이기 위해 Bayes by Backprop 알고리プト와 LRT를 사용하여 최적화된, 완전히 인수 분해된 가우시안 사후 분포를 가진 평균장 근사(mean-field approximation)를 활용한다.
   • 방사형 가이드 (RAD): 가중치 간의 글로벌 상관관계를 포착하기 위해 방사형 변환을 도입한 비평균장(non-mean-field) 접근 방식으로, ELBO의 표준 몬테카를로 샘플링을 통해 최적화된다.

실험 설정
연구는 다음 두 가지 뚜렷한 데이터셋에 대해 이 방법들을 평가했다:

• TiO2 데이터셋: 다양한 단계의 이산화티타늄을 나타내는 7,815개의 구조. 모델은 데이터 효율성을 평가하기 위해 전체 데이터셋(High-Data regime)과 20% 부분 집합(Low-Data regime) 모두에서 훈련되었다. 에너지 전용 및 에너지-힘 공동 훈련이 모두 수행되었다.
• QM7 데이터셋: 화학적으로 다양한 시스템에 대한 일반화 가능성을 테스트하기 위해 사용된 7,165개의 작은 유기 분자. 모델은 원자화 에너지만을 사용하여 high-data regime에서 훈련되었다.

평가 지표
성능은 다음과 같은 종합적인 지표 세트를 통해 평가되었다:

• 정확도: 평균 절대 오차(MAE), 평균 제곱근 오차(RMSE), 음의 로그 가능도(NLL).
• 불확실성 품질: 제곱근 평균 교정 오차(RMSCE), 샤프니스(sharpness, 평균 예측 불확실성), 불확실성과 오차 사이의 스피어만 순위 상관관계($r_s$), 그리고 고불확실성 예측과 고오차 영역 사이의 정렬을 측정하는 "중첩 점수(Overlap Score)".

주요 결과

1. 예측 정확도:

   • **High-Data regime (TiO2)**에서 DE는 가장 낮은 MAE (6 meV/atom)와 RMSE (9 meV/atom), 그리고 가장 낮은 NLL (-4.2)을 달성하며 VBNN보다 일관되게 우수한 성능을 보였으며, 이는 우수한 정확도와 교정 성능을 나타낸다.
   • Low-Data regime에서 DE는 정확도(MAE 35 meV/atom) 측면에서 여전히 우위를 유지했다. 그러나 LRT는 경쟁력 있는 정확도(MAE 54 meV/atom)와 DE보다 낮은 RMSE (259 meV/atom)를 보여, 이상치에 대한 더 나은 강건성을 시사했다. RAD는 저데이터 환경에서 성능이 저조했다.
   • QM7 데이터셋에서는 역전 현상이 발생했다: BNN(LRT 및 RAD)이 DE(15.2 meV/atom)보다 낮은 에너지 RMSE(13.0–13.3 meV/atom)를 달성했으며, 이는 변분 사전 확률(variational prior)이 화학적으로 다양한 공간에서 유익한 규제 역할을 제공함을 시사한다.

2. 불확실성 정량화:

   • 교정 (Calibration): DE는 일반적으로 TiO2 실험에서 가장 잘 교정된 불확실성(가장 낮은 RMSCE)을 제공했다. Low-data regime에서 RAD는 놀랍게도 high-data 성능과 유사한 우수한 교정 성능(RMSCE 0.05)을 유지한 반면, LRT와 DE는 교정 성능이 저하되었다.
   • 샤프니스 (Sharpness): DE는 가장 샤프한(가장 확신에 찬) 예측을 생성했다. LRT는 특히 저데이터 환경에서 지나치게 넓은 불확실성 구간(높은 샤프니스 값)을 생성하는 경향이 있었다.
   • 상관관계 및 중첩: DE는 저데이터 regime에서 예측 불확실성과 실제 오차 사이의 가장 강력한 단조 상관관계($r_s = 0.45$)를 보여, LRT(0.28)와 RAD(0.19)를 능가했다. DE 또한 가장 높은 중첩 점수(55%)를 달성하여, 고오차 영역을 가장 효과적으로 식별했음을 나타냈다.

3. 힘 예측 (Force Predictions):

   • Low-data regime에서 VBNN(LRT 및 RAD)은 DE보다 유의미하게 낮은 힘 RMSE를 달성했는데, 이는 사전 확률이 제공하는 규제 때문으로 보인다. 그러나 힘의 교정은 모든 모델에서 불량했으며, DE만이 저데이터 regime에서 합리적인 교정을 보여주었다.

4. 계산 효율성:

   • 훈련: DE는 더 낮은 에폭당 비용을 갖지만 5개의 독립적인 모델을 훈련해야 한다. DE의 총 훈련 시간은 high-data regime에서는 LRT와 비슷했으나, low-data regime에서는 약간 더 높았다.
   • 추론: DE는 VBNN(0.64–0.71 s)에 비해 훨씬 빠른 추론 속도(0.03 s)를 보인다. BNN은 예측당 20개의 몬테카를로 샘플링을 필요로 하며, 힘 계산을 위한 자동 미분 과정이 포함되기 때문이다.

의의 및 결론
본 논문은 Deep Ensemble이 이론적으로 BNN보다 근거가 덜 탄탄할 수 있음에도 불구하고, 화학적 공간이 상대적으로 좁은 경우(TiO2) 정확도, 교정 및 추론 효율성 측면에서 경험적으로 변분 BNN보다 우수하다고 결론짓는다. 신뢰할 수 있는 불확실성 정량화가 우선순위이고 구현의 단순성을 원하는 경우 DE가 권장된다.

그러나 본 연구는 VBNN, 특히 LRT 변형이 경쟁력 있는 트레이드오프를 제공하며, 사전 확률이 효과적인 규제 역할을 하는 데이터 희소 또는 화학적 다양성이 높은 환경(QM7 사례)에서 더 우수할 수 있음을 강조한다. 저자들은 BNN이 무작위 초기화에 더 민감하며 견고성을 보장하기 위해 여러 번의 훈련 실행이 필요하다고 언급했다. 이 연구는 불확실성을 인지하는 MLIP를 선택하기 위한 실질적인 지침을 제공하며, DE와 VBNN 사이의 선택은 데이터 레짐과 대상 시스템의 화학적 다양성에 따라 달라져야 함을 시사한다.

기여

• 에너지-힘 공동 훈련을 지원하는, MLIP를 위한 GPU 가속 VBNN 구현을 ænet-PyTorch 프레임워크 내에 도입.
• 무기물(TiO2) 및 유기물(QM7) 데이터셋 모두에서 high- 및 low-data regime에 걸친 DE, LRT, RAD의 체계적인 비교.
• DE가 일반적으로 우수하지만, 화학적으로 다양한 데이터셋에서는 BNN의 상대적 성능이 역전될 수 있음을 입증하여 앙상블의 보편적 우월성에 도전함.
• 능동 학습 시나리오에서 불확실성 추정치의 유용성을 평가하기 위한 새로운 "중첩 점수(Overlap Score)" 지표 제공.