[논문 리뷰] 3D Infomax improves GNNs for Molecular Property Prediction

H. Stärk, Dominique Beaini외 5명|arXiv (Cornell University)|2021. 10. 08.

Machine Learning in Materials Science참고 문헌 53인용 수 85

한 줄 요약

이 논문은 3D Infomax를 도입하며 2D GNN이 암묵적으로 3D 분자 기하를 추론하도록 가르치고, 다운스트림 양자 특성 예측을 개선하고 데이터셋 간 일반화를 가능하게 한다.

ABSTRACT

Molecular property prediction is one of the fastest-growing applications of deep learning with critical real-world impacts. Including 3D molecular structure as input to learned models improves their performance for many molecular tasks. However, this information is infeasible to compute at the scale required by several real-world applications. We propose pre-training a model to reason about the geometry of molecules given only their 2D molecular graphs. Using methods from self-supervised learning, we maximize the mutual information between 3D summary vectors and the representations of a Graph Neural Network (GNN) such that they contain latent 3D information. During fine-tuning on molecules with unknown geometry, the GNN still generates implicit 3D information and can use it to improve downstream tasks. We show that 3D pre-training provides significant improvements for a wide range of properties, such as a 22% average MAE reduction on eight quantum mechanical properties. Moreover, the learned representations can be effectively transferred between datasets in different molecular spaces.

연구 동기 및 목표

• 명시적 3D 구조를 이용할 수 없는 대규모에서의 분자 특성 예측을 촉진한다.
• 2D 그래프로부터 잠재적 3D 정보를 인코딩하는 방식의 사전 학습 스킴을 개발한다.
• 학습된 3D 인식 표현이 다운스트림 양자 특성 예측을 향상시킴을 입증한다.
• 3D Infomax 표현이 서로 다른 분자 공간과 데이터세트에 걸쳐 일반화됨을 보여준다.
• 사전 학습 중 다중 콘포머를 사용하는 이점이 성능을 추가로 향상시키는지 탐구한다.

제안 방법

• 2D 분자 그래프에서 2D GNN f^a를 사전 학습하여 콘포머로부터 얻은 3D 표현 z^b와의 상호정보를 최대화해 z^a를 생성한다.
• 동일 분자에 대해 z^a와 z^b를 정렬하고 일치하지 않는 쌍을 분리하도록 NT-Xent 스타일의 대조 손실을 사용한다.
• 분자 내 모든 콘포머를 대조 손실(L^multi3D)을 계산할 때 양의 예제로 취급하여 다중 콘포머(c)를 포함한다.
• 3D 네트워크는 3D 좌표를 포인트 클라우드로 처리하고 거리 기반 MPNN을 통해 γ(d_uv) 형태의 사인파 특성 매핑으로 SE(3)-불변의 z^b를 인코딩한다.
• 2D 입력만 사용하여 다운스트림 작업에 대해 2D 네트워크를 미세조정하고, 암묵적으로 학습된 3D 정보를 활용해 예측을 향상시킨다.

실험 결과

연구 질문

• RQ1 2D GNN이 3D 콘포머와의 대조적 사전 학습을 통해 2D 분자 그래프로부터 잠재적 3D 정보를 인코딩하도록 학습할 수 있는가?
• RQ2 3D Infomax 사전 학습 모델이 2D 전용 사전 학습 기준선과 비교해 QM9 등 양자 특성 예측을 개선하는가?
• RQ3 학습된 표현이 서로 다른 크기와 구성을 가진 분자 공간 간에 전달 가능한가?
• RQ4 사전 학습 중 다중 콘포머를 활용하면 다운스트림 특성에서 추가 이득이 있는가?

주요 결과

• 3D Infomax는 양자 특성 예측에서 큰 향상을 달성하며 QM9 특성에서 평균 MAE를 약 22% 감소시켰다.
• 사전 학습된 표현은 서로 다른 분자 공간에 걸쳐 일반화되며 작은 분자에서 약물 유사 분자로의 전이 및 그 반대에서도 이점을 보인다.
• 사전 학습 중 다중 콘포머를 사용하면 추가 성능 이득이 있으며, 세 가지 콘포머를 넘으면 수익이 점점 감소한다.
• 3D Infomax는 평가된 데이터세트 전반에서 GraphCL 및 거리 예측 사전 학습과 같은 기존 SSL 기준선을 능가한다.
• 본 방법은 이전의 일부 사전 학습 방식과 달리 음의 전이를 나타내지 않는다.