[논문 리뷰] A benchmark study on reliable molecular supervised learning via Bayesian learning

Doyeong Hwang, Grace Lee외 3명|arXiv (Cornell University)|2020. 06. 12.

Computational Drug Discovery Methods참고 문헌 28인용 수 4

한 줄 요약

이 논문은 분자 그래프 신경망(GNN)에서 예측 신뢰도를 향상시키기 위한 베이지안 학습 방법을 벤치마킹하며, Stochastic Weight Averaging(SWA) 및 SWAG와 같은 기법이 다양한 GNN 아키텍처와 분자 성질 예측 과제에서 잘 校정된, 불확실성 인식 예측을 제공하는 것으로 밝혀졌다. 연구는 베이지안 방법이 과신을 감소시키고 신뢰할 수 있는 예측 확률을 제공함으로써 가상 스크리닝 성공률을 향상시킴을 보여주었다.

ABSTRACT

Virtual screening aims to find desirable compounds from chemical library by using computational methods. For this purpose with machine learning, model outputs that can be interpreted as predictive probability will be beneficial, in that a high prediction score corresponds to high probability of correctness. In this work, we present a study on the prediction performance and reliability of graph neural networks trained with the recently proposed Bayesian learning algorithms. Our work shows that Bayesian learning algorithms allow well-calibrated predictions for various GNN architectures and classification tasks. Also, we show the implications of reliable predictions on virtual screening, where Bayesian learning may lead to higher success in finding hit compounds.

연구 동기 및 목표

• 그래프 신경망에서 베이지안 딥 러닝을 활용한 분자 성질 예측의 신뢰도 평가.
• 특히 가상 스크리닝에서 분포 외(out-of-distribution, OOD) 화합물에 대해 과신하는 예측 문제 해결.
• 베이지안 추론 기법이 GNN 기반 분자 분류에서 교정성과 불확실성 추정을 어떻게 향상시키는지 평가.
• 약물 발견과 같은 긴급한 응용 분야를 위한 신뢰할 수 있는 분자 기계학습의 실용적 벤치마킹 제공.
• 재현 가능성과 계산 화학 및 가상 스크리닝 분야에서의 광범위한 채택을 위한 코드 공개.

제안 방법

• 예측 신뢰도를 정량적으로 평가하기 위해 예측 신뢰도와 실제 정확도 간의 괴리 정도를 측정하는 기대 교정 오차(Expected Calibration Error, ECE)를 사용.
• 모델 가중치의 사후 분포를 추정하기 위해 근사 추론 방법을 활용한 베이지안 학습: MAP, 딥 엔sembles, MC-Dropout, SGLD, SWA, SWAG.
• 기본 모델로 그래프 신경망(GIN, GAT, GatedGCN 등)을 사용하며, 입력 특징은 분자 그래프와 노드/에지 표현에서 유도된다.
• 베이지안 근사화를 위해 예측 분포는 다수의 샘플링된 가중치에 대한 예측 평균을 통해 계산되며, 이로써 예측 확률의 분산을 통한 불확실성 추정이 가능해진다.
• 모델은 네 개의 MoleculeNet 데이터셋(BACE, BBBP, HIV, Tox21)에서 스캐폴드 분할 데이터 분할(80:10:10)을 사용해 훈련되며, 각 기법에 맞게 초모수를 튜닝하였다.
• SWA와 SWAG는 순환 학습률 스케줄링과 함께 150~250 에포크 동안 가중치 평균을 적용하였으며, SWAG의 사후 공분산은 1.0으로 스케일링되었다.

실험 결과

연구 질문

• RQ1 다양한 아키텍처와 과제에서 베이지안 학습 기법이 GNN 기반 분자 성질 예측기의 교정성을 향상시킬 수 있는가?
• RQ2 다양한 베이지안 추론 기법은 분자 분류 벤치마크에서 예측 신뢰도와 성능 측면에서 어떻게 상호 비교되는가?
• RQ3 잘 교정된 예측이 거짓 양성 및 거짓 음성 수를 감소시켜 가상 스크리닝 성공률을 얼마나 향상시키는가?
• RQ4 불확실성 추정은 진짜 양성, 거짓 양성, 진짜 음성, 거짓 음성의 예측 행동에 어떤 영향을 미치는가?
• RQ5 SWA와 SWAG는 분자 성질 예측에서 높은 정확도와 낮은 ECE를 동시에 달성하기 위해 다른 베이지안 방법과 비교해 어떻게 성능을 내는가?

주요 결과

• SWA와 그 변종인 SWAG는 모든 GNN 아키텍처와 분자 성질 예측 과제에서 기대 교정 오차(Expected Calibration Error, ECE)가 가장 낮게 유지되어 뛰어난 교정성을 보였다.
• 베이지안 학습 기법의 적용으로 과신이 크게 감소하였으며, 특히 OOD 및 분포 외 샘플에서 실제 정확도 비율과 예측 확률이 매우 유사하게 유지되었다.
• 히스토GRAM 분석 결과, 진짜 양성은 대부분 높은 예측 확률을 할당받은 반면, 거짓 양성은 낮은 신뢰도로 예측되는 경향이 있었으며, 이는 신뢰할 수 있는 불확실성 정량화를 나타내었다.
• SWA와 SWAG를 사용한 모델은 BACE와 BBBP에서 AUROC(>0.85)와 F1-스코어(>0.75)가 높으면서도 ECE 값이 0.05 이하로 유지되어 높은 성능과 높은 신뢰도를 동시에 확보하였다.
• 연구 결과, MC-Dropout와 딥 엔셈블도 교정성을 향상시켰지만, 특히 클래스 불균형이 심한 Tox21 데이터셋에서 SWA와 SWAG만큼 일관되지는 않았다.
• https://github.com/AITRICS/mol_reliable_gnn 에 공개된 코드베이스는 재현 가능한 벤치마킹을 가능하게 하며, 코로나19와 같은 긴급한 약물 발견 응용 분야를 지원한다.