[논문 리뷰] GemNet: Universal Directional Graph Neural Networks for Molecules

Johannes Gasteiger, Florian Becker외 1명|arXiv (Cornell University)|2021. 06. 02.

Computational Drug Discovery Methods인용 수 127

한 줄 요약

GemNet은 회전 불변/변환 가능 분자 예측을 위한 구면 표현의 보편성을 확립하고 방향성 엣지 임베딩을 갖춘 기하학적 두 홉 메시지 전달을 도입하며, 여러 MD 데이터셋에서 힘 예측의 최첨단 성능을 달성한다.

ABSTRACT

Effectively predicting molecular interactions has the potential to accelerate molecular dynamics by multiple orders of magnitude and thus revolutionize chemical simulations. Graph neural networks (GNNs) have recently shown great successes for this task, overtaking classical methods based on fixed molecular kernels. However, they still appear very limited from a theoretical perspective, since regular GNNs cannot distinguish certain types of graphs. In this work we close this gap between theory and practice. We show that GNNs with spherical representations are indeed universal approximators for predictions that are invariant to translation, and equivariant to permutation and rotation. We then discretize such GNNs via directed edge embeddings and two-hop message passing, and incorporate multiple structural improvements to arrive at the geometric message passing neural network (GemNet). We demonstrate the benefits of the proposed changes in multiple ablation studies. GemNet outperforms previous models on the COLL, MD17, and OC20 datasets by 34%, 41%, and 20%, respectively, and performs especially well on the most challenging molecules. Our implementation is available online.

연구 동기 및 목표

• 분자 그래프에 대한 일반 GNN의 표현력 격차를 해소하고 회전 불변 및 등가 예측을 가능하게 한다.
• 회전 불변 및 등가 작업에 대한 구면 표현의 보편성을 증명한다.
• 정확도와 샘플 효율성을 향상시키기 위해 기하학적 정보를 포함한 실용적이고 방향성 있는 간선 기반 두 홉 메시지 전달 프레임워크를 개발한다.

제안 방법

• SO(3)-불변 및 순열-등가 함수에 대한 구면 표현의 보편성을 증명한다.
• 구면 표현을 방향성 엣지 임베딩으로 이산화하고 엣지 간의 두 홉 메시지 전달을 모델링한다.
• φ와 θ를 포함한 원자 간 방향 및 각도를 사용한 구면 푸리에-베셀 기저를 활용한 기하학적 메시지 전달을 도입한다.
• 안정성과 계산 효율성을 개선하기 위해 대칭 메시지 전달과 효율적 바일리아인 계층을 통합한다.
• 표준 정규화 대신 미리 정해진 스케일링 팩터로 활성화 분산을 안정화한다.
• 두 홉 기하학적 메시지 전달, 한 홉 기하학적 메시지 전달, 원자 자기 상호작용의 세 가지 상호작용 형태를 가진 GemNet 아키텍처를 개발한다.

실험 결과

연구 질문

• RQ1 구면 표현이 분자 GNN에서 회전 불변 및 등가 예측에 대한 보편성을 달성할 수 있는가?
• RQ2 방향성(엣지 기반) 표현을 어떻게 이산화하여 분자를 위한 실용적이고 표현력 있는 GNN을 구현할 수 있는가?
• RQ3 대칭 메시징, 바일리니어 계층, 분산 안정화 등 어떤 아키텍처 혁신이 분자 동역학 데이터셋에서 에너지와 힘의 예측 정확도를 가장 크게 향상시키는가?
• RQ4 COLL, MD17, OC20에서 힘과 에너지 예측 측면에서 GemNet이 최첨단 모델과 어떻게 비교되는가?

주요 결과

• GemNet은 데이터세트 전반에서 상당한 오차 감소를 달성하며 COLL, MD17, OC20에서 이전 모델들보다 눈에 띄는 차이로 우수한 성능을 보였다.
• 방향 엣지 임베딩을 가진 두 홉 기하학적 메시지 전달이 성능을 크게 향상시키며, 특히 도전적인 비평면(non-planar) 분자에서 두드러진다.
• 대칭 메시지 전달, 바일리니어 계층, 활성화 분산 안정화가 정확도에 의미 있게 기여하며, 제거 연구를 통해 각 구성 요소의 이점이 확인된다.
• 직접 힘 예측은 훈련 및 추론을 가속화할 수 있으며, MAE의 트레이드오프는 데이터셋과 작업에 따라 달라진다.
• GemNet-T와 GemNet-Q 변형은 계산 비용과 정확도 사이에 서로 다른 트레이드오프를 제공하며, 더 복잡한 결합 데이터셋에서 GemNet-Q가 더 큰 이점을 제공한다.