[논문 리뷰] On Graph Neural Network Ensembles for Large-Scale Molecular Property Prediction

Edward Elson Kosasih, Joaquín Cabezas외 6명|arXiv (Cornell University)|2021. 06. 29.

Machine Learning in Materials Science인용 수 7

한 줄 요약

이 논문은 380만 개의 분자를 포함하는 대규모 분자 성질 예측을 향상시키기 위해, GIN-Virtual, 베이지안 GNN, DiffPool 향상된 GNN의 세 가지 그래프 신경망(GNN) 모델을 다수의 랜덤 초기화로 훈련한 앙상블을 제안한다. 이 앙상블은 기준 모델 대비 7.6%의 상대적 향상도를 기록하며, PCQM4M-LSC 데이터셋에서 평균 절대 오차(MAE)를 0.1290으로 낮추었고, 예측 오차와 예측의 불확실성 간에 강한 상관관계가 있음을 입증하였다(Pearson R = 0.5181). 이는 향후 주도적 학습(active learning) 응용에 활용 가능함을 시사한다.

ABSTRACT

In order to advance large-scale graph machine learning, the Open Graph Benchmark Large Scale Challenge (OGB-LSC) was proposed at the KDD Cup 2021. The PCQM4M-LSC dataset defines a molecular HOMO-LUMO property prediction task on about 3.8M graphs. In this short paper, we show our current work-in-progress solution which builds an ensemble of three graph neural networks models based on GIN, Bayesian Neural Networks and DiffPool. Our approach outperforms the provided baseline by 7.6%. Moreover, using uncertainty in our ensemble's prediction, we can identify molecules whose HOMO-LUMO gaps are harder to predict (with Pearson's correlation of 0.5181). We anticipate that this will facilitate active learning.

연구 동기 및 목표

• OGB-LSC PCQM4M-LSC 데이터셋에 포함된 380만 개의 분자를 대상으로 대규모 분자 성질 예측 성능을 향상시키기 위해.
• 실제 대규모 데이터셋에 적용 가능한 그래프 기반 기계학습의 확장성과 성능 도전 과제를 해결하기 위해.
• 앙상블 예측의 불확실성이 주어진 예측 오차가 높은 영역을 나타내는 데 사용될 수 있는지 조사하기 위해.
• 다양한 GNN 아키텍처와 훈련 실행을 결합한 확장 가능한 앙상블 프레임워크를 개발하여 회귀 성능 향상에 기여하기 위해.

제안 방법

• 다양한 GNN 약한 학습기(weak learners)를 앙성화: GIN-Virtual, 베이지안 GNN(BNN), DiffPool 향상된 GNN 각각을 서로 다른 초기화 세 번으로 훈련.
• 각 모델의 여러 실행 결과 평균을 취해 개별 앙상블(예: GIN-Virtual-ensemble)을 구성.
• 9개의 별도 실험 결과에서 유도된 예측을 가중 평균 방식으로 조합하여 최종 앙상블-all을 구성.
• GNN의 읽기 출력 레이어에 베이지안 역전파(Bayes-by-Backpropagation, BBP)를 적용하여 베이지안 선형 레이어를 통해 불확실성 추정을 도입.
• DiffPool를 활용해 그래프를 5개 노드로 압축하여 구조적 정보를 유지하면서 표현을 향상시키는 그래프 풀링을 구현.
• 9개의 앙상블 실험 결과에서의 예측 표준편차를 불확실성으로 정의하고, 이를 모델 신뢰도의 대체 지표로 사용.

실험 결과

연구 질문

• RQ1 다양한 GNN 아키텍처를 앙상블하여 개별 모델을 초월해 대규모 분자 그래프에서 회귀 성능을 향상시킬 수 있는가?
• RQ2 앙성화된 예측 불확실성은 HOMO-LUMO 갭 예측에서 실제 예측 오차와 상관관계가 있는가?
• RQ3 제안된 앙성화 방법이 PCQM4M-LSC 데이터셋에서 OGB-LSC 기준 모델을 상당한 폭으로 초월할 수 있는가?
• RQ4 불확실성 추정치가 예측이 어려운 분자를 식별함으로써 주도적 학습을 얼마나 효과적으로 이끌 수 있는가?

주요 결과

• 앙성화-all 모델은 검증 세트에서 평균 절대 오차(MAE) 0.1290을 기록하여 기준 모델의 MAE 0.1396 대비 7.6% 상대적 향상도를 달성하였다.
• 9개의 앙성화 실험 결과에서의 예측 표준편차로 정의된 불확실성 추정치는 실제 예측 오차와 피어슨 상관계수 0.5181을 보였다.
• GIN-Virtual-BNN-ensemble 버전은 개별 앙성화 유형 중에서 가장 낮은 검증 MAE 0.1329를 기록하였다.
• 앙성화-all 모델은 GIN-Virtual를 포함한 모든 개별 모델 및 기준 모델을 초월하여 성능을 냈으며, 이는 대회에서 가장 뛰어난 단일 모델이었다.
• 불확실성과 오차 간의 상관관계는 불확실성 추정치가 더 예측하기 어려운 분자를 식별하는 데 유용하며, 향후 주도적 학습 응용에 기여할 수 있음을 시사한다.
• 고메모리 클라우드 GPU 인스턴스를 사용해 약 300만 개의 분자로 학습을 수행하여, 대규모 분자 데이터셋에 대한 확장성 잠재력을 입증하였다.