[논문 리뷰] Uncertainty Estimation for Molecules: Desiderata and Methods

Tom Wollschläger, Nicholas Gao외 3명|arXiv (Cornell University)|2023. 06. 20.

Machine Learning in Materials Science인용 수 5

한 줄 요약

이 논문은 그래프 신경망(GNN) 기반 분자의 힘장에 대해 신뢰할 수 있는 불확실성 추정을 제공하는 가우시안 프로세스 기반 방법인 로컬라이저드 뉴럴 커널(LNK)을 제안한다. 이는 세 가지 물리 기반 및 세 가지 응용 중심의 요구사항을 충족한다. LNK는 높은 예측 정확도를 유지하면서도 단일 순방향 전파를 통해 신뢰할 수 있는 불확실성 정량화를 가능하게 하여, 비균형 상태 탐지에서 몬테카를로 드롭아웃 및 증거 기반 회귀보다 최대 2.5배와 2.1배 낮은 AUC-ROC 오차를 달성한다.

ABSTRACT

Graph Neural Networks (GNNs) are promising surrogates for quantum mechanical calculations as they establish unprecedented low errors on collections of molecular dynamics (MD) trajectories. Thanks to their fast inference times they promise to accelerate computational chemistry applications. Unfortunately, despite low in-distribution (ID) errors, such GNNs might be horribly wrong for out-of-distribution (OOD) samples. Uncertainty estimation (UE) may aid in such situations by communicating the model's certainty about its prediction. Here, we take a closer look at the problem and identify six key desiderata for UE in molecular force fields, three 'physics-informed' and three 'application-focused' ones. To overview the field, we survey existing methods from the field of UE and analyze how they fit to the set desiderata. By our analysis, we conclude that none of the previous works satisfies all criteria. To fill this gap, we propose Localized Neural Kernel (LNK) a Gaussian Process (GP)-based extension to existing GNNs satisfying the desiderata. In our extensive experimental evaluation, we test four different UE with three different backbones and two datasets. In out-of-equilibrium detection, we find LNK yielding up to 2.5 and 2.1 times lower errors in terms of AUC-ROC score than dropout or evidential regression-based methods while maintaining high predictive performance.

연구 동기 및 목표

• 표준 GNN이 실패하는 외부 분포(OOD) 예측에서 특히 신뢰할 수 있는 불확실성 추정의 필수적인 필요성을 해결한다.
• 효과적인 분자의 시스템에서의 불확실성 추정을 위해 세 가지 물리 기반(대칭성, 에너지 보존, 국소성)과 세 가지 응용 중심(속도, 해석 가능성, OOD 탐지)의 여섯 가지 핵심 요구사항을 규명한다.
• 분자의 머신러닝에서 기존의 불확실성 추정 방법을 조사하고, 이들 중 어느 것도 여섯 가지 요구사항을 모두 충족하지 못함을 보여준다.
• GNN 기반 모델에 통합되어 재학습 없이 단일 순방향 전파를 통해 불확실성 추정이 가능한, 새로운 방법인 LNK를 제안한다. 이는 모든 요구사항을 충족한다.
• 여러 백본과 불확실성 추정 방법을 사용하여 두 데이터셋에서 LNK를 실증적으로 평가하며, OOD 탐지 및 예측 성능에 중점을 둔다.

제안 방법

• 분자의 힘장에서의 불확실성 추정을 위한 여섯 가지 요구사항을 정의한다: 세 가지 물리 기반(유클리드 및 순열 대칭성에 대한 불변성, 커플-프리 힘에 의한 에너지 보존)과 세 가지 응용 중심(속도, 해석 가능성, OOD 탐지 능력).
• 분자의 그래프에 대해 국소화된 커널 함수를 학습하는 방식으로 불확실성을 모델링하는, 미분 가능하고 GP 기반의 GNN 확장인 로컬라이저드 뉴럴 커널(LNK)을 제안한다.
• 기존의 GNN 백본(예: DimeNet++)에 LNK를 통합하여, 백본을 재학습하지 않고도 불확실성 추정이 가능하게 하며, 단일 순방향 전파를 사용한다.
• 물리적 대칭성을 존중하고, 힘장이 구조적으로 보존되도록 보장하는, 미분 가능한 커널 메커니즘을 사용한다.
• 변동 기반 접근법을 통해 에너지(u_E)와 힘(u_F)에 대한 불확실성 정량화를 가능하게 하며, u_F는 커널의 추적값에서 계산된다.
• 에너지/힘 예측과 불확실성 校정을 통합한 확률적 손실을 사용하여 모델을 엔드 투 엔드로 학습함으로써, 잘 校정된 불확실성 추정을 보장한다.

실험 결과

연구 질문

• RQ1 분자의 힘장에서의 불확실성 추정을 위한 필수적인 물리적 및 응용 중심의 요구사항은 무엇인가?
• RQ2 기존의 분자의 머신러닝에서의 불확실성 추정 방법들은 이러한 요구사항을 어느 정도 충족하는가?
• RQ3 분자의 대칭성과 보존 법칙을 존중하면서도, 빠른 단일 순방향 전파를 통해 불확실성 추정이 가능한 GP 기반 방법을 설계할 수 있는가?
• RQ4 제안된 LNK 방법은 몬테카를로 드롭아웃, 증거 기반 회귀 등의 기존 불확실성 추정 기법과 비교해 OOD 탐지 및 예측 성능에서 어떻게 성능을 내는가?
• RQ5 GNN 백본에 LNK를 통합함으로써 높은 예측 정확도를 유지하면서도 불확실성 校정을 크게 향상시킬 수 있는가?

주요 결과

• 기존의 분자의 머신러닝에서의 불확실성 추정 방법 중 어느 것도 제안된 여섯 가지 요구사항을 모두 충족하지 못하며, 특히 물리적 일관성 또는 추론 속도에서 실패한다.
• QM7X 데이터셋에서 비균형 상태 탐지 시 LNK는 몬테카를로 드롭아웃 및 증거 기반 회귀보다 최대 2.5배와 2.1배 낮은 AUC-ROC 오차를 기록한다.
• LNK는 에너지 MAE 0.059 kcal/mol, 힘 MAE 0.007 kcal/mol의 높은 예측 성능를 유지하며, 1% 드롭아웃 비율에서도 최신 기술 수준의 모델과 유사한 성능을 보인다.
• MCD 및 앙상블과 같은 분산 기반 방법보다 4.6배 이상 빠르며, QM7X에서 샘플당 추론 시간이 59.3±2.68 ms로 증거 기반 회귀와 유사한 속도를 기록한다.
• 추론 시 드롭아웃을 사용하는 MCD-LNK의 변종은 드롭아웃 비율 증가에 따라 안정적인 성능을 보이며, 표준 MCD보다 OOD 탐지 및 예측 정확도에서 뛰어난 성능을 보인다.
• 앙상블은 힘에 대해 가장 높은 AUC-ROC(0.996)를 기록하지만, 학습 시간이 다섯 배 더 오래 걸리고 추론 시간이 274±8.25 ms로 훨씬 더 길어, 활성 학습이나 실시간 응용에는 실용적이지 않다.