[논문 리뷰] A new perspective on building efficient and expressive 3D equivariant graph neural networks

Weitao Du, Yuanqi Du외 6명|arXiv (Cornell University)|2023. 04. 07.

Machine Learning in Materials Science인용 수 7

한 줄 요약

이 논문은 3D 등변성 그래프 신경망의 효율적이고 표현력 있는 설계를 위한 새로운 이론적 프레임워크를 제안하며, 3D 동형사상의 계층적 구조를 통해 局소-전반적 기하 표현력을 분석한다. Local Substructure Encoding (LSE)와 Frame Transition Encoding (FTE)를 통합한 LEFTNet을 제안하여 분자 성질 예측 과제에서 최신 기술 수준을 넘어선 성능을 달성하며, 스칼라 성질의 경우 평균 20% 향상되고 벡터 성질의 경우 평균 16% 향상된다.

ABSTRACT

Geometric deep learning enables the encoding of physical symmetries in modeling 3D objects. Despite rapid progress in encoding 3D symmetries into Graph Neural Networks (GNNs), a comprehensive evaluation of the expressiveness of these networks through a local-to-global analysis lacks today. In this paper, we propose a local hierarchy of 3D isomorphism to evaluate the expressive power of equivariant GNNs and investigate the process of representing global geometric information from local patches. Our work leads to two crucial modules for designing expressive and efficient geometric GNNs; namely local substructure encoding (LSE) and frame transition encoding (FTE). To demonstrate the applicability of our theory, we propose LEFTNet which effectively implements these modules and achieves state-of-the-art performance on both scalar-valued and vector-valued molecular property prediction tasks. We further point out the design space for future developments of equivariant graph neural networks. Our codes are available at \url{https://github.com/yuanqidu/LeftNet}.

연구 동기 및 목표

• 3D 등변성 GNN의 표현력을 평가하기 위한 이론적 프레임워크를 구축하기 위해 기하 동형사상의 局소-전반적 분석을 수행한다.
• 표현력과 계산 효율성을 동시에 확보할 수 있는 핵심 아키텍처 구성요소인 LSE 및 FTE를 규명한다.
• 불변 기하 메시지와 등변성 집계를 조합하면 전반적 3D 기하 구조를 효과적으로 포착할 수 있음을 입증한다.
• 미래의 3D 기하 딥러닝 모델을 위한 모듈러 아키텍처 설계 가이드라인을 제공한다.

제안 방법

• GNN이 기하 구조를 포착하는 데 있어 표현력을 분석하기 위해 3D 동형사상의 계층적 구조—나무, 삼각형, 부분그래프—를 제안한다.
• 미분기하학의 스칼라화 기법을 활용하여 국소 3D 기하 패턴을 불변적으로 표현하는 Local Substructure Encoding (LSE) 기법을 도입한다.
• 메시지 전파 과정에서 국소 기준좌표계의 변환을 모델링하기 위해 Frame Transition Encoding (FTE)를 개발하여 전반적 기하 표현을 가능하게 한다.
• SE(3) 대칭성을 유지하기 위해 등변성 텐서 갱신과 함께 LSE 및 FTE를 통합한 모듈러 아키텍처인 LEFTNet을 설계한다.
• 고체체계의 고차원 잠재에너지 모델을 영감으로 삼아 텐서 곱 또는 합산을 사용한 기준좌표 기반 메시지 집계 메커니즘을 적용한다.
• 프레임 번들의 스칼라화를 적용하여 국소 기하 특징의 불변 인코딩을 가능하게 하여 LSE에 필수적인 요소를 확보한다.

실험 결과

연구 질문

• RQ1 불변 스칼라의 표현력은 3D 기하 패턴을 어떻게 인코딩하는가?
• RQ2 메시지 전파에서의 등변성은 전반적 3D 기하를 포착하는 데 진정으로 필수적인가?
• RQ3 어떻게 국소 기하 정보를 대칭성 유지 방식으로 체계적으로 집계하여 전반적 3D 기하 구조를 표현할 수 있는가?
• RQ4 3D 등변성 GNN에서 높은 표현력과 계산 효율성을 동시에 확보할 수 있는 아키텍처 구성요소는 무엇인가?

주요 결과

• LEFTNet은 QM9 데이터셋에서 12개의 분자 성질 예측 과제 전반에 걸쳐 평균 20% 향상되어 최신 기술 수준 성능을 달성한다.
• rMD17 데이터셋에서 힘 가중 손실이 100일 경우, 8개 분자 중 7개에서 모든 베이스라인 모델을 초월하며 평균 16% 향상된다.
• QM9 데이터셋의 두 데이터 분할 조건 하에서 12개 성질 중 10개에서 두 번째로 우수하거나 최고 성능을 기록하여 강건성과 일반화 능력을 입증한다.
• 제거 실험 결과에 따르면 LSE와 FTE 모두 성능 향상에 필수적이며, 각각이 최종 결과에 크게 기여하고 있음을 확인한다.
• 이 프레임워크는 분자 외에도 다양체 위의 기하 메esh 그래프로의 확장 가능성이 있으며, 접선 공간 기준좌표 정의를 통해 가능해진다.