[논문 리뷰] Equivariant Graph Attention Networks for Molecular Property Prediction

Tuan Le, Frank Noé외 1명|arXiv (Cornell University)|2022. 02. 20.

Computational Drug Discovery Methods인용 수 6

한 줄 요약

이 논문은 분자의 3차원 기하학적 구조 정보를 유지하기 위해 회전 및 이동에 대해 등변성을 보장하는 새로운 GNN 아키텍처인 등변 그래프 어텐션 네트워크(EQGAT)를 제안한다. 이는 카르테시안 좌표와 공간-콘텐츠 인식 어텐션 메커니즘을 활용하여 회전 및 이동에 대해 등변적인 분자 표현 학습을 달성한다. 소분자에 대한 양자성질 예측에서 최고 성능을 기록하며, 단백질-리간드 결합 친화도와 같은 대규모 생물분자 과제에서도 경쟁력 있는 결과를 내며, 기존의 등변 GNN보다 더 빠른 추론 속도를 보이는 계산 효율성이 뛰어난 모델이다.

ABSTRACT

Learning and reasoning about 3D molecular structures with varying size is an emerging and important challenge in machine learning and especially in drug discovery. Equivariant Graph Neural Networks (GNNs) can simultaneously leverage the geometric and relational detail of the problem domain and are known to learn expressive representations through the propagation of information between nodes leveraging higher-order representations to faithfully express the geometry of the data, such as directionality in their intermediate layers. In this work, we propose an equivariant GNN that operates with Cartesian coordinates to incorporate directionality and we implement a novel attention mechanism, acting as a content and spatial dependent filter when propagating information between nodes. We demonstrate the efficacy of our architecture on predicting quantum mechanical properties of small molecules and its benefit on problems that concern macromolecular structures such as protein complexes.

연구 동기 및 목표

• 분자의 그래프에서 3차원 기하학적 구조 정보를 유지하기 위해 회전 및 이동에 대해 등변성을 보장하는 GNN을 개발한다.
• 새로운 어텐션 메커니즘을 통해 방향성(공간) 및 원자성(콘텐츠) 정보를 메시지 전파에 통합하여 분자성질 예측 성능을 향상시킨다.
• 소분자 외에도 단백질과 단백질-리간드 복합체를 포함한 대규모 분자 체계에서의 확장 가능하고 효율적인 학습을 가능하게 한다.
• 낮은 계산 비용을 유지하면서 다양한 분자 예측 벤치마크에서 최고 수준 또는 경쟁 가능한 성능을 달성한다.

제안 방법

• 모델은 직접 카르테시안 좌표를 사용하여, 메시지 전파 프레임워크 내에서 명시적인 변환 규칙을 통해 3차원 회전 및 이동에 대해 등변성을 확보한다.
• 새로운 그래프 어텐션 메커니즘이 상대적 위치(공간)와 노드 특징(콘텐츠)에 모두 의존하는 엣지 어텐션 가중치를 계산하여, 기하학적 인식이 가능한 적응형 메시지 전파를 가능하게 한다.
• 기하학적 정확성을 유지하기 위해 스칼라 및 벡터 표현(타입-0 및 타입-1 텐서)을 사용하여 텐서 필드 네트워크에서 사용하는 복잡한 기저 계산이 필요 없도록 한다.
• 메시지 전파 과정은 공간 대칭성에 대한 변환 행동을 유지하면서 기능을 전파하는 등변 업데이트 규칙을 사용한다.
• 메시지 전파 신경망(MPNN) 프레임워크 내에 구현되며, 표현 학습을 향상시키기 위해 어텐션 기반 집계를 사용한다.
• 기저 기반 등변 GNN과 비교해 기저 생성을 피하고 카르테시안 공간에서 직접 연산함으로써 FLOPs와 메모리 사용량을 줄였다.

실험 결과

연구 질문

• RQ1 공간 및 콘텐츠 정보를 동시에 인코딩하는 어텐션 메커니즘을 갖춘 GNN이 기존 GNN보다 더 나은 분자성질 예측 성능을 낼 수 있는가?
• RQ2 명시적인 등변성을 갖는 카르테시안 공간에서의 직접 구현이 기저 기반 텐서 필드 네트워크보다 성능과 효율성 면에서 뛰어나게 되는가?
• RQ3 등변 GNN이 단백질과 단백질-리간드 복합체와 같은 대규모 생물분자 체계에 효과적으로 확장될 수 있는가?
• RQ4 다양한 분자 벤치마크에서 EQGAT은 PaiNN 및 E(n)-GNN과 같은 최고 수준의 등변 GNN과 비교해 성능과 추론 속도 면에서 어떻게 비교되는가?

주요 결과

• EQGAT은 소분자에 대한 양자역학적 성질 예측에서 최고 성능을 기록하며, 이 벤치마크에서 이전 GNN들을 모두 앞서며 최고의 성능을 보였다.
• LBA(리간드 결합 친화도) 데이터셋에서 EQGAT은 테스트된 모든 GNN 아키텍처 중에서 가장 우수한 성능을 기록했으며, 전체적으로는 3D CNN 모델이 최고 성능을 냈음에도 불구하고 그러한 성과를 달성했다.
• EQGAT은 뛰어난 계산 효율성을 보이며, LBA에서 23.67ms, PSR에서 42.79ms, RSR에서 54.80ms의 추론 시간을 기록하여 모든 벤치마크에서 PaiNN 및 E(n)-GNN을 앞서는 성능을 보였다.
• 모델의 추론 속도 우월성은 카르테시안 공간에서의 직접적 구현 덕분이며, PaiNN에서 사용하는 메모리 집약적인 벡터 특징 표현 방식을 피하기 때문이다.
• PaiNN의 생고분자 벤치마크에서의 열악한 성능은 3차원 벡터 특징을 저장하는 데 소요되는 높은 메모리 사용량 때문일 수 있으며, 이는 그래프 크기가 커질수록 성능 저하를 야기한다.
• 저자들은 별도의 리간드 및 단백질 인코더를 사용할 경우 성능 향상이 가능할 수 있으며, 향후 연구의 유망한 방향으로 제안했다.