[논문 리뷰] SE(3) Equivariant Graph Neural Networks with Complete Local Frames

Weitao Du, He Zhang외 5명|arXiv (Cornell University)|2021. 10. 26.

Machine Learning in BioinformaticsBiochemistry, Genetics and Molecular Biology인용 수 22

한 줄 요약

이 논문은 크로스 곱을 통해 구성된 완전한 局소 프레임을 사용하여 기하 텐서의 충실하고 기준점이 없는 표현을 가능하게 하는 계산적으로 효율적인 SE(3)-동치성 그래프 신경망인 ClofNet을 제안한다. 고차원 기하 양을 정규직교 기저에 투영함으로써, 모델은 다체 역학 모델링과 분자 구조 생성에서 최고 성능을 달성하며, 뛰어난 샘플 효율성과 낮은 동치성 오차를 보인다.

ABSTRACT

Group equivariance (e.g. SE(3) equivariance) is a critical physical symmetry in science, from classical and quantum physics to computational biology. It enables robust and accurate prediction under arbitrary reference transformations. In light of this, great efforts have been put on encoding this symmetry into deep neural networks, which has been shown to improve the generalization performance and data efficiency for downstream tasks. Constructing an equivariant neural network generally brings high computational costs to ensure expressiveness. Therefore, how to better trade-off the expressiveness and computational efficiency plays a core role in the design of the equivariant deep learning models. In this paper, we propose a framework to construct SE(3) equivariant graph neural networks that can approximate the geometric quantities efficiently. Inspired by differential geometry and physics, we introduce equivariant local complete frames to graph neural networks, such that tensor information at given orders can be projected onto the frames. The local frame is constructed to form an orthonormal basis that avoids direction degeneration and ensure completeness. Since the frames are built only by cross product operations, our method is computationally efficient. We evaluate our method on two tasks: Newton mechanics modeling and equilibrium molecule conformation generation. Extensive experimental results demonstrate that our model achieves the best or competitive performance in two types of datasets.

연구 동기 및 목표

• 3D 다체 시스템을 위한 SE(3)-동치성 딥 러닝에서 표현력과 계산 효율성의 균형을 맞추는 도전 과제를 해결한다.
• 기존의 불완전한 반경 방향 프레임에 의존하는 동치성 GNN에서 발생하는 방향성 불확실성 문제를 해결한다.
• 기준점이 없는 방식으로 기하 양(예: 벡터, 비틀림 힘)을 손실 없는 방식으로 동치성 유지하면서 표현할 수 있도록 한다.
• 임의의 순서의 텐서를 지원하고, 고차원 리프팅이나 복잡한 텐서 연산 없이도 군의 동치성을 유지하는 프레임워크를 개발한다.
• 물리적 모델링과 분자 구조 생성 작업에서 뛰어난 일반화 능력과 데이터 효율성을 입증한다.

제안 방법

• 입자 벡터 간의 크로스 곱을 통해 3차원 공간에서 정규직교 기저를 형성하는 완전한 局소 프레임을 도입하여 완전성과 방향성 불확실성의 방지를 보장한다.
• 기하 텐서(스칼라, 벡터, 고차원)를 스칼라화 블록을 통해 이러한 局소 프레임에 투영하여 메시지 전파를 위한 스칼라 계수로 변환한다.
• 스칼라 계수에 대해 표준 그래프 신경망을 사용하여 메시지 전파를 수행하며, 프레임의 동치성 특성 덕분에 SE(3) 동치성이 유지된다.
• 크로스 곱 연산만을 사용하여 프레임을 구성함으로써, 클레브시-고르단 분해나 구면 조화 함수에 의존하는 방법에 비해 계산적으로 효율적이다.
• 크로스 곱을 반복 적용하고 정규직교성을 유지함으로써 고차원 텐서 입력을 처리할 수 있도록 프레임 구성 기법을 확장한다.
• 분자 구조 생성을 위한 스코어 기반 생성 모델링 설정에 프레임워크를 통합하고, 샘플링을 위해 안내된 랭그비안 역동법을 사용한다.

실험 결과

연구 질문

• RQ1 고차원 리프팅이나 텐서 분해의 높은 비용을 피할 수 있는 계산적으로 효율적인 SE(3)-동치성 GNN을 설계할 수 있는가?
• RQ2 완전한 局소 프레임을 사용할 경우, 반경 방향 프레임이나 불완전한 프레임 방법에 비해 표현력과 정확도가 향상되는가?
• RQ3 제안된 방법이 물리 시스템의 분포 외부 또는 저데이터 환경에서 얼마나 잘 일반화되는가?
• RQ4 다체 역학 모델링과 분자 구조 생성에서 모두 최고 성능을 달성하면서도 강력한 동치성을 유지할 수 있는가?
• RQ5 크로스 곱을 통한 프레임 구성 방식이 다른 동치성 기반 기능 공간 표현에 비해 효율성과 안정성 면에서 어떻게 비교되는가?

주요 결과

• POS 데이터셋에서 ClofNet은 보간 MSE 0.138과 외삽 MSE 2.502로 가장 뛰어난 성능을 보이며, Radial Field(7.665)와 EGNN(6.449)를 크게 앞서 간다.
• 낮은 동치성 오차($\Delta_{EQ} = 4.13 \cdot 10^{-6}$)를 유지하여 SE(3) 동치성의 강력한 수치적 보존을 입증한다.
• GEOM-QM9 및 GEOM-Drugs 데이터셋에서 ClofNet은 저 RMSD를 기록하며 고품질의 분자 구조를 생성하며, 최적으로 정렬된 구조의 시각화 결과에서 이를 확인할 수 있다.
• 학습 데이터의 0.4%만으로도 경쟁적인 성능를 달성하여 뛰어난 샘플 효율성을 입증한다.
• 제거 실험을 통해 완전한 局소 프레임의 사용이 토크력과 같은 복잡한 기하 양을 모델링하는 데 핵심적임을 확인한다.
• 고비용의 텐서 곱 분해를 피하기 위해 프레임 구성에 크로스 곱 연산만을 사용함으로써, 프레임워크의 계산 효율성이 검증되었다.