[논문 리뷰] Bayesian Graph Neural Networks for Molecular Property Prediction

George S. Lamb, Brooks Paige|arXiv (Cornell University)|2020. 11. 25.

Machine Learning in Materials Science참고 문헌 30인용 수 9

한 줄 요약

이 논문은 QM9 데이터셋에서 방향성 메시지 전파 신경망(D-MPNN)의 메시지 전파 및 리더아웃 레이어에서의 불확실성 정량화를 활용하여 분자의 물리적 성질 예측을 위한 베이지안 그래프 신경망(GNN)을 제안한다. 전체 베이지안 GNN은 표준 딥러닝 기준선에 비해 예측 정확도, 校정성(calibration) 및 후속 분자의 탐색 작업 성능을 크게 향상시킨다.

ABSTRACT

Graph neural networks for molecular property prediction are frequently underspecified by data and fail to generalise to new scaffolds at test time. A potential solution is Bayesian learning, which can capture our uncertainty in the model parameters. This study benchmarks a set of Bayesian methods applied to a directed MPNN, using the QM9 regression dataset. We find that capturing uncertainty in both readout and message passing parameters yields enhanced predictive accuracy, calibration, and performance on a downstream molecular search task.

연구 동기 및 목표

• 학습 데이터에서의 과적합으로 인해 GNN의 분자의 분포 외 scaffolds에 대한 일반화 능력이 떨어지는 문제를 해결하기 위해.
• 특히 회귀 설정에서 분자의 물리적 성질 예측의 모델 校정성 및 불확실성 정량화를 향상시키기 위해.
• GNN의 메시지 전파 및 리더아웃 구성 요소에 대한 베이지안 추론의 영향을 평가하기 위해.
• 실제 분자의 탐색 작업에서 베이지안 방법의 성능을 벤치마킹하여 실용적 유용성을 평가하기 위해.
• 재현성 및 향후 연구를 위해 오픈소스 PyTorch 코드를 제공하기 위해.

제안 방법

• 분자의 표현 학습을 위한 기본 GNN 아키텍처로 방향성 메시지 전파 신경망(D-MPNN)을 사용한다.
• 여덟 가지 베이지안 추론 방법(MAP, GP, MC Dropout(DropR, DropA), SWAG, SGLD, BBP, DUN)을 적용하여 불확실성 정량화를 수행한다.
• 특히 애로카틱 불확실성(aleatoric uncertainty)이 명시적으로 모델링되지 않은 경우를 고려해 후행(post-hoc) t-분포 우도를 사용하여 校정성을 향상시킨다.
• 사후 근사(posterior approximation)에서 샘플링을 통해 앙상블 기반 예측 분포를 구현하여 불확실성 추정을 수행한다.
• 모델 파라미터의 사후 분포를 근사하기 위해 변분 추론 또는 MCMC 기반 방법을 사용하여 엔드 투 엔드로 모델을 훈련시킨다.
• 신뢰도 다이어그램을 활용해 불확실성의 정확도를 평가하며, 12개의 QM9 회귀 과제에서 평균 순위와 오차가 큰 면적(miscalibration area)을 핵심 평가 지표로 사용한다.

실험 결과

연구 질문

• RQ1 메시지 전파 및 리더아웃 레이어에서의 베이지안 학습이 분자의 물리적 성질 예측의 정확도 및 校정성 향상에 기여하는가?
• RQ2 다양한 베이지안 추론 방법은 QM9 회귀 과제에서 불확실성 정량화 및 성능 측면에서 어떻게 비교되는가?
• RQ3 베이지안 GNN은 후속 분자의 합성 또는 탐색 작업 성능 향상에 기여하는가?
• RQ4 애로카틱 불확실성을 모델링할 경우, 특히 정규분포와 t-분포 우도를 사용했을 때 校정성에 어떤 영향을 미치는가?
• RQ5 불확실성 정량화는 예측된 분자의 scaffolds에 대한 일반화 능력에 어떤 영향을 미치는가?

주요 결과

• 메시지 전파 및 리더아웃 레이어 양쪽에서 불확실성 정량화를 수행한 전체 베이지안 GNN은 QM9에서 평균 순위 1.95 ± 0.40을 기록하여 표준 MAP 훈련(4.08 ± 0.16)에 비해 유의미하게 뛰어난 성능을 보였다.
• 전체 GNN에 MC 드롭아웃을 적용한 DropA는 분자의 탐색 점수 77.52 ± 0.77%를 기록하여 MAP(72.22 ± 0.57%) 및 기타 베이지안 방법보다 뛰어난 성능을 보였다.
• 후행 t-분포 우도의 사용은 校정성을 향상시켜 정규분포 우도에서 관찰된 병적인 과소신뢰(underconfidence)를 감소시켰다.
• 베이지안 리더아웃(GP, DropR 등)을 적용한 모델는 단지 베이지안 메시지 전파만 적용한 모델보다 더 나은 校정성을 보였으며, 이는 리더아웃 불확실성이 예측 신뢰도에 더 중요한 영향을 미친다는 것을 시사한다.
• SGLD 및 BBP는 높은 불확실성 추정치를 보였지만 예측 정확도가 떨어져, MCMC 및 변분 방법이 이 설정에서는 항상 잘 일반화되지 않는다는 것을 시사한다.
• DUN 방법은 대각 + 저랭크 공분산을 사용하여 평균 순위(4.40 ± 0.25)와 탐색 점수(76.02 ± 1.09)에서 뛰어난 성능을 보였으며, 이는 불확실성 정량화의 강건성(robustness)을 나타낸다.