[논문 리뷰] ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs

Limei Wang, Yi Liu외 3명|arXiv (Cornell University)|2022. 06. 17.

Machine Learning in Materials Science인용 수 27

한 줄 요약

tldr: ComENet은 3D 분자 그래프를 위한 완전하고 효율적인 1-hop 이웃 메시지 전달 체계를 도입하여 SE(3)-완전한 표현을 O(nk) 복잡도로 달성하고 대규모 데이터셋에서 강력한 실험적 성능을 발휘합니다.

ABSTRACT

Many real-world data can be modeled as 3D graphs, but learning representations that incorporates 3D information completely and efficiently is challenging. Existing methods either use partial 3D information, or suffer from excessive computational cost. To incorporate 3D information completely and efficiently, we propose a novel message passing scheme that operates within 1-hop neighborhood. Our method guarantees full completeness of 3D information on 3D graphs by achieving global and local completeness. Notably, we propose the important rotation angles to fulfill global completeness. Additionally, we show that our method is orders of magnitude faster than prior methods. We provide rigorous proof of completeness and analysis of time complexity for our methods. As molecules are in essence quantum systems, we build the \underline{com}plete and \underline{e}fficient graph neural network (ComENet) by combing quantum inspired basis functions and the proposed message passing scheme. Experimental results demonstrate the capability and efficiency of ComENet, especially on real-world datasets that are large in both numbers and sizes of graphs. Our code is publicly available as part of the DIG library (\url{https://github.com/divelab/DIG}).

연구 동기 및 목표

• 3D 정보의 완전한 통합이 분자 그래프의 정확한 학습을 위해 필요하다는 점을 동기화합니다.
• 3D 그래프 표현에 대한 기하학적 완전성(전역 및 지역)을 정의하고 보장합니다.
• 전역 완전성을 달성하기 위해 회전 각이 있는 효율적인 1-hop 메시지 전달 체인을 개발합니다.
• ComENet 내에서 양자에서 영감을 받은 특징들을 통합하여 3D 분자 학습을 개선합니다.
• 대규모 데이터셋(OC20, Molecule3D, QM9)에서 확장성과 경쟁력 있는 정확성을 입증합니다.

제안 방법

• SE(3) 하에서 3D 그래프의 완전성을 형식적으로 정의하고 가장자리의 회전 각을 통해 전역 완전성을 보여줍니다.
• 엣지당 1-hop 기하학적 특징(d, theta, phi, tau)을 계산하여 SE(3)-불변 표현을 형성합니다.
• 1-hop 이웃에서 작동하는 로컬 및 글로벌 컨볼루션 계층을 도입하여 O(nk) 복잡도를 달성합니다.
• 가장 가까운 이웃을 기반으로 한 회전 각 tau를 사용하여 3D 구성체를 고유하게 결정하고 구성체 식별을 가능하게 합니다.
• 양자에서 영감을 받은 기저 특징과 뉴럴 네트워크 기반 메시지 전달 프레임워크를 ComENet 아키텍처(임베딩 계층, 상호 작용 계층, 자기 원자 계층, 풀링) 내에 통합합니다.
• 완전한 메시지 전달을 양자에서 영감을 받은 특징과 함께 확장 가능하게 제시하고 시간 복잡성을 분석합니다; OC20, Molecule3D, QM9에서 평가합니다.

실험 결과

연구 질문

• RQ1 1-hop 이웃 정보만으로도 3D 분자 그래프 표현이 전역적으로 완전(정보 손실 없음)할 수 있는가?
• RQ2 회전 각을 어떻게 활용해 전역 완전성을 달성하고 구성체 수준의 판별을 가능하게 할 수 있는가?
• RQ3 완전한 3D GNN의 실제 시간 복잡도는 기존의 2-hop 방법과 비교해 어떠한가?
• RQ4 제안된 메시지 전달과 결합된 양자 영감 특징이 대규모 분자 데이터세트에서 경쟁력 있는 성능을 내는가?

주요 결과

• 1-hop 완전 기하학적 변환(d, theta, phi, tau)이 강하게 연결된 3D 그래프에 대해 전역 완전성을 보유함을 증명합니다.
• 1-hop 이웃과 회전 각을 사용하여 O(nk) 시간 복잡도를 달성하고 nk^2 또는 nk^3 방법보다 훨씬 빠릅니다.
• OC20, Molecule3D, QM9에서 최첨단 기법과 비교하여 경쟁적이거나 우수한 성능을 보이며, 일부 데이터셋에서 학습/추론 속도 향상을 6-10배 달성합니다.
• 회전 각이 OC20에서 구성체 수준의 판별을 가능하게 하고 중요한 구성 차이를 식별하게 합니다.
• 완전한 메시지 전달을 양자 영감 특징과 통합한 ComENet 아키텍처를 제안하여 확장 가능한 3D 분자 학습을 제공합니다.