[논문 리뷰] Efficient and Accurate Physics-aware Multiplex Graph Neural Networks for 3D Small Molecules and Macromolecule Complexes

Shuo Zhang, Yang Liu외 1명|arXiv (Cornell University)|2022. 06. 06.

Machine Learning in Materials Science인용 수 6

한 줄 요약

이 논문은 3차원 소분자 및 고분자 복합체에서 스칼라 및 벡터성 성질을 효율적이고 정확하게 예측할 수 있도록, 국소적이고 비국소적 분자 상호작용을 별도의 메시지 전파 방식을 사용해 모델링하는 물리적 인지 능력을 갖춘 다중층 그래프 신경망인 PaxNet을 제안한다. PaxNet은 소분자에서 예측 오차를 15% 감소시키고 메모리 사용량을 73% 줄이며, 고분자 복합체에서는 추론 시간을 85% 감소시키고 메모리를 33% 줄여 최신 기준 모델보다 우수한 성능을 발휘한다.

ABSTRACT

Recent advances in applying Graph Neural Networks (GNNs) to molecular science have showcased the power of learning three-dimensional (3D) structure representations with GNNs. However, most existing GNNs suffer from the limitations of insufficient modeling of diverse interactions, computational expensive operations, and ignorance of vectorial values. Here, we tackle these limitations by proposing a novel GNN model, Physics-aware Multiplex Graph Neural Network (PaxNet), to efficiently and accurately learn the representations of 3D molecules for both small organic compounds and macromolecule complexes. PaxNet separates the modeling of local and non-local interactions inspired by molecular mechanics, and reduces the expensive angle-related computations. Besides scalar properties, PaxNet can also predict vectorial properties by learning an associated vector for each atom. To evaluate the performance of PaxNet, we compare it with state-of-the-art baselines in two tasks. On small molecule dataset for predicting quantum chemical properties, PaxNet reduces the prediction error by 15% and uses 73% less memory than the best baseline. On macromolecule dataset for predicting protein-ligand binding affinities, PaxNet outperforms the best baseline while reducing the memory consumption by 33% and the inference time by 85%. Thus, PaxNet provides a universal, robust and accurate method for large-scale machine learning of molecules. Our code is available at https://github.com/zetayue/Physics-aware-Multiplex-GNN.

연구 동기 및 목표

• 기존 GNN 모델이 다양한 분자 상호작용, 특히 3차원 구조에서의 상호작용을 모델링하는 데 한계가 있음을 해결하기 위해.
• 모든 상호작용을 동일하게 취급하는 전통적인 GNN과 달리 국소적 및 비국소적 상호작용 모델링을 분리하여 계산 비용을 줄이고, 비용이 많이 드는 각도 연산을 피하기 위해.
• 스칼라 성질 외에도 디폴트 모멘트 및 힘과 같은 벡터성 성질을 정확하게 예측할 수 있도록 하기 위해.
• 자원 제약 조건 하에서 단백질-리간드 시스템과 같은 고분자 복합체에 대해 효율성과 확장성을 향상시키기 위해.
• 대규모 분자 표현 학습을 위한 통합적이고 견고하며 정확한 프레임워크 개발하기 위해.

제안 방법

• 각 '플렉스'가 상호작용 유형(예: 국소 공유 결합, 비국소 비공유 결합)을 별도로 나타내는 다중층 그래프 구조를 사용하여, 각 상호작용 유형에 대해 별도의 메시지 전파를 수행한다.
• 국소(1- 및 2-호프) 및 비국소(전역) 메시지 전파를 분리하여 더 정확하고 효율적으로 상호작용을 모델링하는 물리적 인spired 설계를 도입한다.
• 원자 위치에서 유도된 기하학적 벡터와 메시지 가중치를 통합하여 디폴트 모멘트와 같은 벡터성 성질을 직접 예측한다.
• 다양한 상호작용 유형의 기여도를 동적으로 가중하기 위해 메시지 전파 방식 전역에 학습 가능한 어텐션 메커니즘을 적용한다.
• 각도 연산을 물리적 인지된 근사치로 대체하여 O(Nk²) 복잡도를 피하는 수정된 메시지 전파 방식을 사용한다.
• 각 원자에 벡터를 연결하고 메시지 전파를 통해 업데이트함으로써 스칼라 및 벡터성 분자 성질에 대한 엔드 투 엔드 학습을 지원한다.

실험 결과

연구 질문

• RQ1 국소 및 비국소 상호작용에 대해 별도의 메시지 전파를 적용하는 다중층 GNN 아키텍처가 3차원 분자 성질 예측 정확도를 향상시킬 수 있는가?
• RQ2 기존 GNN이 모든 상호작용을 동일하게 취급하는 것과 비교해, 상호작용 모델링을 분리함으로써 계산 비용은 어떻게 감소하는가?
• RQ3 스칼라 모델 뿐 아니라 디폴트 모멘트와 같은 벡터성 성질을 더 정밀하게 예측할 수 있는가?
• RQ4 대규모 고분자 복합체에서 PaxNet은 메모리 및 추론 시간을 크게 줄이면서도 높은 성능을 유지할 수 있는가?
• RQ5 어떤 구성 요소(예: 어텐션, 플렉스 간 통신, 다단계 메시지 전파)가 전체 모델 성능에 기여하는가?

주요 결과

• QM9 데이터셋에서 PaxNet은 디폴트 모멘트 예측의 평균 절대 오차(MAE)를 0.0108 D로 줄여 최고의 베이스라인(0.021 D)보다 50% 이상 우수하다.
• QM9에서 PaxNet은 최고의 베이스라인 모델 대비 15% 낮은 예측 오차와 73% 적은 메모리 사용량을 기록한다.
• PDBbind 데이터셋에서 PaxNet은 최고의 베이스라인 대비 메모리 소비를 33% 줄이고 추론 시간을 85% 감소시키며, 더 뛰어난 성능을 달성한다.
• 제거 실험 결과 어텐션 모듈을 제거하면 스칼라 목표에서 MAE가 최대 4% 증가하고, 벡터성 목표에서는 최대 3% 증가하여 그 중요성을 입증한다.
• 플렉스 간 통신 또는 메시지 전파 방식을 제거하면 스칼라 목표에서 최대 28% 성능 저하와 벡터성 목표에서 최대 102% 성능 저하가 발생하여 통합 설계의 필수성을 확인한다.
• 국소 메시지 전파에 이중 히트 기하학적 정보(θ)를 사용하면 U₀에서 성능이 64% 향상되고 μ에서 111% 향상되어 더 풍부한 기하학적 인코딩의 가치를 입증한다.