[논문 리뷰] Reliable Graph Neural Networks for Drug Discovery Under Distributional Shift

Kehang Han, Balaji Lakshminarayanan외 1명|arXiv (Cornell University)|2021. 11. 25.

Computational Drug Discovery Methods참고 문헌 29인용 수 15

한 줄 요약

이 논문은 스펙트럴 정규화와 거리 인식 가우시안 프로세스 분류기를 통합하여 분포 이탈 하에서 약물 발굴을 위한 신뢰할 수 있는 그래프 신경망 GNN-SNGP를 제안한다. 새로운 CardioTox 벤치마크와 세 개의 기존 데이터셋에서 검증된 결과, 특히 분포 이탈 분자의 경우 과신 잡음 예측을 줄이며 높은 정확도를 유지하면서도 뛰어난 강건성을 입증하였다.

ABSTRACT

The concern of overconfident mis-predictions under distributional shift demands extensive reliability research on Graph Neural Networks used in critical tasks in drug discovery. Here we first introduce CardioTox, a real-world benchmark on drug cardio-toxicity to facilitate such efforts. Our exploratory study shows overconfident mis-predictions are often distant from training data. That leads us to develop distance-aware GNNs: GNN-SNGP. Through evaluation on CardioTox and three established benchmarks, we demonstrate GNN-SNGP's effectiveness in increasing distance-awareness, reducing overconfident mis-predictions and making better calibrated predictions without sacrificing accuracy performance. Our ablation study further reveals the representation learned by GNN-SNGP improves distance-preservation over its base architecture and is one major factor for improvements.

연구 동기 및 목표

• 약물 발굴에 사용되는 GNN에서 과신 잡음 예측의 심각한 신뢰성 문제를 해결함. 특히 분포 이탈 상황에서의 문제에 초점을 맞춤.
• 과신 잡음 예측은 일반적으로 훈련 데이터에서 구조적으로 거리가 먼 분자에서 자주 발생함을 확인함. 이는 거리 인식 능력의 부족을 시사함.
• 표현에서의 거리 유지 능력을 향상시키고 불확실성 인식 예측을 통합함으로써 모델의 신뢰성을 향상시키는 새로운 GNN 아키텍처인 GNN-SNGP를 개발함.
• 실제 약물 발굴 조건에서의 모델 신뢰성 평가를 위해 내재된 및 분포 이탈된 테스트 세트를 포함한 실제 세계 기반의 벤치마크인 CardioTox를 도입함.
• 향상된 거리 인식 능력이 보다 잘 校정된 예측과 과신된 거짓 음성 결과 감소에 기여함을 입증함. 이는 안전한 약물 개발에 매우 중요함.

제안 방법

• 불확실성 추정을 위해 스펙트럴 정규화(SNGP)와 무작위 특징 가우시안 프로세스(GP) 레이어를 결합한 엔드 투 엔드로 훈련 가능한 GNN인 GNN-SNGP를 제안함.
• 메시지 전달 레이어에 스펙트럴 정규화를 적용하여 잠재 공간 내 기하학적 거리 유지 능력을 향상시켜 거리 인식 능력을 강화함.
• 최종 밀집 레이어를 대체로, 훈련 데이터로부터의 거리 기반 예측 불확실성 모델링을 수행하는 GP 레이어를 도입함으로써 분포 이탈 분자에서의 과신을 감소시킴.
• 훈련 데이터로부터의 분자 거리를 정량화하고 평가를 위한 거리 기반 데이터 분할을 정의하기 위해 피ン거프린트 기반 그래프 거리(Rogers and Hahn, 2010)를 사용함.
• 표준 GNN 손실에 불확실성 정규화를 적용하여 정확도를 희생시키지 않은 채 신뢰할 수 있는 불확실성 정량화를 가능하게 함.
• GP 레이어와 거리 유지 표현의 기여도를 분리하기 위해 아블레이션 스터디를 수행함.

실험 결과

연구 질문

• RQ1 약물 독성 예측을 위한 GNN에서 과신 잡음 예측이 훈련 데이터에서 구조적으로 거리가 먼 분자에서 얼마나 자주 발생하는가?
• RQ2 약물 발굴에서 분포 이탈 상황에서 GNN 표현의 거리 인식 능력을 향상시키면 과신 잡음 예측을 줄일 수 있는가?
• RQ3 GNN-SNGP는 분포 이탈 데이터에서 보정 및 강건성 측면에서 기준 GNN 및 기타 불확실성 추정 방법과 비교해 어떻게 성능을 내는가?
• RQ4 거리 유지 표현과 GP 기반 분류기의 조합이 각각의 구성 요소보다 더 나은 불확실성 정량화를 이끌어내는가?
• RQ5 제안된 방법은 심장독성 외의 다양한 그래프 기반 약물 발굴 작업에 일반화 가능한가?

주요 결과

• GNN-SNGP는 테스트-OOD1 분할에서 최고의 AUROC(0.935)와 가장 낮은 ECE(0.002)를 기록하여, GNN 기준 모델 및 기타 변종보다 보다 뛰어난 보정 및 강건성 성능을 입증함.
• 가장 도전적인 테스트-OOD2 분할에서 GNN-SNGP는 과신 거짓 음성 결과(OFNs%)를 0.00%로 줄였으며, FPs + Dense 기준 모델의 5.71% 대비 유의미하게 향상됨.
• 아블레이션 스터디 결과, 거리 유지 GNN 인코더와 GP 레이어 모두 신뢰성 향상에 기여함을 확인함. 특히 GP 레이어가 과신을 감소시키는 데 뛰어난 효과를 보임.
• OOD 분할에서 ECE(최대 0.021)와 브리어 스코어(0.208)가 모두 향상되어 보다 보정된 불확실성 추정이 가능함을 시사함.
• 과신 예측을 줄이면서도 높은 정확도(AUROC > 0.89)를 유지함으로써, 신뢰성과 성능가 서로 상충되지 않음을 입증함.
• 모델은 일반화 능력이 뛰어남: molHIV, BBBP, BACE 데이터셋에서도 일관된 성능 향상을 기록하여 심장독성 외의 응용 분야에서도 넓은 적용 가능성을 확인함.