[논문 리뷰] Coherent energy and force uncertainty in deep learning force fields

Peter Bjørn Jørgensen, Jonas Busk외 2명|arXiv (Cornell University)|2023. 12. 07.

Machine Learning in Materials Science참고 문헌 16인용 수 8

한 줄 요약

이 논문은 상관된 노이즈 프로세스를 모델링하여 딥러닝 포스 필드의 에너지 및 힘의 애토로픽 불확실성을 연결하고, 힘의 분산 표현식을 도출하며, 등가적 메시지 전달 신경망 앙상블을 사용해 분자 데이터셋에서 일관된 불확실성을 검증한다.

ABSTRACT

In machine learning energy potentials for atomic systems, forces are commonly obtained as the negative derivative of the energy function with respect to atomic positions. To quantify aleatoric uncertainty in the predicted energies, a widely used modeling approach involves predicting both a mean and variance for each energy value. However, this model is not differentiable under the usual white noise assumption, so energy uncertainty does not naturally translate to force uncertainty. In this work we propose a machine learning potential energy model in which energy and force aleatoric uncertainty are linked through a spatially correlated noise process. We demonstrate our approach on an equivariant messages passing neural network potential trained on energies and forces on two out-of-equilibrium molecular datasets. Furthermore, we also show how to obtain epistemic uncertainties in this setting based on a Bayesian interpretation of deep ensemble models.

연구 동기 및 목표

• ML 포스 필드에서 에너지와 힘에 대한 일관된 불확실성 추정의 필요성에 대한 동기 제시.
• 에너지 애토로픽 불확실성을 공간적으로 상관된 프로세스를 통해 힘 불확실성과 연결하는 노이즈 모델 제안.
• 컬러 노이즈 하에서 에너지 및 힘 분산에 대한 닫힌 형식 표현 도출.
• 깊은 앙상블을 통한 인식적(에피스테믹) 불확실성의 검출 및 베이지안 해석 제시.
• 등가적 메시지 전달 신경망을 사용해 ANI-1x 및 Transition1x와 같은 비평형 분자 데이터에 대한 접근 평가

제안 방법

• 에너지 관측치를 E_obs = E_theta + rho_theta * eta로 모델링하되, eta는 미분 가능 autocorrelation을 가짐.
• 힘 분산 Var(-∂E_obs/∂r) = gamma_hat * rho_theta^2 + (∂rho_theta/∂r)^2 (식 8)로 도출.
• 노이즈 커널의 제곱 역 길이 스케일로 gamma_hat를 해석하고 이를 하이퍼파라미터로 학습하거나 설정.
• 분산 도출을 단순화하기 위해 Wide-Sense 정적 노이즈를 사용(부록 A).
• 포스터리어 샘플링을 통한 인식적 불확실성 확보를 위해 Bayesian 딥 앙상블 프레임워크 채택(부록 B).
• PaiNN 기반의 등가적 앙상블을 훈련하고 보정된 애토로픽/인식적 구성요소를 갖춘 상태에서 에너지 및 힘 MAE/RMSE, NLL, 보정 지표를 평가(표 1).

Figure 1: Ensemble model (with 10 instances) fitted to energy (left) and force (right) observations with correlated noise.

실험 결과

연구 질문

• RQ1 ML 포스 필드에서 에너지와 힘의 애토로픽 불확실성을 일관되게 연결할 수 있는가?
• RQ2 상관된(컬러링된) 노이즈 모델링이 에너지와 힘의 불확실성 추정에 어떤 영향을 미치는가?
• RQ3 이 일관된 프레임워크에서 인식적 불확실성은 딥 앙상블로 포착될 수 있는가?
• RQ4 제안된 방법이 비평형 분자 데이터 ANI-1x 및 Transition1x에서 일반적 baselines(바닐라 및 화이트 노이즈) 대비 어떻게 성능을 발휘하는가?
• RQ5 이 일관된 불확실성 접근법의 실용적 보정 및 예측 성능 이점은 무엇인가?

주요 결과

• 컬러 노이즈 에너지 모델은 식(8)에서 하이퍼파라미터 gamma_hat을 통해 에너지 불확실성과 결합된 닫힌 형식의 힘 분산을 얻는다.
• 힘 분산은 분산 출력의 제곱 기울기에 대한 항과 rho_theta^2에 비례하는 항을 포함하여 에너지/힘 불확실성을 일관되게 추정할 수 있게 한다.
• 깊은 베이지안 앙상블로 인식적 불확실성을 얻을 수 있으며 애토로픽과 인식적 기여를 구분한다(부록 B).
• ANI-1x 및 Transition1x에서 컬러/노이즈 기반 앙상블이 vanilla 및 화이트 노이즈 Baseline 대비 경쟁력 있는 MAE/RMSE를 보이며 불확실성 지표(NLL, ENCE, RZV, CV)가 개선되거나 보정된다(표 1).
• 검증 세트 보정으로 예측 정확성을 해치지 않으면서 불확실성 보정이 향상될 수 있다.

Figure 2: Reliability diagrams for models trained on the ANI-1x dataset.