[논문 리뷰] Equivariant Graph Neural Operator for Modeling 3D Dynamics

Minkai Xu, Jiaqi Han외 7명|arXiv (Cornell University)|2024. 01. 19.

Machine Learning in Materials ScienceMaterials Science인용 수 3

한 줄 요약

이 논문은 3D 물리적 동역학을 연속적인 궤적으로 모델링하기 위해 푸리에 기반 시간 컨볼루션을 갖는 SE(3)-동변성 신경 연산자로 구성된 새로운 방법인 등변성 그래프 신경 연산자(EGNO)를 제안한다. 시간에 대한 함수로서 동역학을 학습함으로써 EGNO는 입자 시뮬레이션, 인간 운동 캡처, 분자 동역학에서 뛰어난 성능을 달성하며, EGNN과 같은 기준 모델 대비 평균 52% 향상된 성능을 기록한다.

ABSTRACT

Modeling the complex three-dimensional (3D) dynamics of relational systems is an important problem in the natural sciences, with applications ranging from molecular simulations to particle mechanics. Machine learning methods have achieved good success by learning graph neural networks to model spatial interactions. However, these approaches do not faithfully capture temporal correlations since they only model next-step predictions. In this work, we propose Equivariant Graph Neural Operator (EGNO), a novel and principled method that directly models dynamics as trajectories instead of just next-step prediction. Different from existing methods, EGNO explicitly learns the temporal evolution of 3D dynamics where we formulate the dynamics as a function over time and learn neural operators to approximate it. To capture the temporal correlations while keeping the intrinsic SE(3)-equivariance, we develop equivariant temporal convolutions parameterized in the Fourier space and build EGNO by stacking the Fourier layers over equivariant networks. EGNO is the first operator learning framework that is capable of modeling solution dynamics functions over time while retaining 3D equivariance. Comprehensive experiments in multiple domains, including particle simulations, human motion capture, and molecular dynamics, demonstrate the significantly superior performance of EGNO against existing methods, thanks to the equivariant temporal modeling. Our code is available at https://github.com/MinkaiXu/egno.

연구 동기 및 목표

• 기존 그래프 신경망이 3D 물리적 동역학에서 장거리 시간적 상관관계를 포착하는 데 한계를 가진다는 문제를 해결한다.
• 다음 단계 예측의 제약을 극복하기 위해 전체 해 궤적을 시간에 대한 연속 함수로 모델링한다.
• 물리적 일반화에 매우 중요한 SE(3)-등변성을 유지하면서 3D 공간 동역학에서 시간 의존성을 학습한다.
• 재학습 없이도 임의의 시간 이산화에 대해 효율적이고 병렬적인 추론을 가능하게 하여 구현의 유연성을 높인다.
• 다양한 물리적 동역학 응용 분야에 적합한 일반화 가능한 프레임워크를 개발한다.

제안 방법

• 초기 상태에서 전체 궤적을 맵핑하는 함수로 3D 동역학을 정식화하여 신경 연산자를 통한 연속 시간 모델링을 가능하게 한다.
• 시간 모델링 중 SE(3)-등변성을 유지하기 위해 푸리에 도메인에서 매개변수화된 등변성 시간 컨볼루션 레이어를 도입한다.
• 푸리에 변환과 역푸리에 변환의 등변성을 활용하여 주파수 공간에서 회전 및 이동 불변성을 유지하는 커널 적분 연산자를 설계한다.
• 표준 EGNN 구성 요소와 푸리에 기반 시간 레이어를 스택하여 EGNO를 구축함으로써 기하학적 및 시간적 동역학의 엔드 투 엔드 학습을 가능하게 한다.
• 시간 임bedding을 사용하여 학습된 시간 단계 사이에서 보간함으로써 재학습 없이도 임의의 시간 해상도에서 제로샷 추론을 가능하게 한다.
• 기하학적 특징—노드 특징(h), 위치(x), 속도(v)—를 시간 컨볼루션에 통합하여 풍부한 동적 패턴을 포착한다.

Figure 1: Illustration of EGNO . EGNO blocks (green) can be built with any EGNN layers (blue) and the proposed equivariant temporal convolution layers (yellow). Consider discretizing the time window $\Delta T$ into $P$ points $\{\Delta t_{1},\dots,\Delta t_{P}\}$ . Given a current state ${\mathcal{G

실험 결과

연구 질문

• RQ1 신경 연산자 프레임워크는 연속적인 시간 3D 동역학을 모델링하는 동안 SE(3)-등변성을 유지할 수 있는가?
• RQ2 푸리에 도메인에서 시간 컨볼루션을 통합함으로써 자동회귀적 다음 단계 예측 대비 장거리 시간적 상관관계 모델링에 어떤 개선을 이룬다?
• RQ3 재학습 없이도 EGNO가 예측되지 않은 시간 이산화 단계에 얼마나 잘 일반화되는가?
• RQ4 다양한 기하학적 특징(h, x, v)이 EGNO의 시간 모델링 성능에 기여하는 정도는 어느 정도인가?
• RQ5 다양한 물리계에서 궤적 정확도와 일반화 능력 측면에서 기존 GNN 기반 모델과 비교해 EGNO는 어떠한가?

주요 결과

• 아스피린 분자 동역학 벤치마크에서 EGNO는 EGNN 대비 F-MSE 기준 36% 상대적 향상을 달성한다.
• 인간 운동 캡처(Mocap-Run) 데이터셋에서 EGNO는 EGNN 대비 예측 정확도 평균 52% 향상을 기록한다.
• 절단 실험 결과, 시간 컨볼루션에 속도(v)를 통합하는 것이 성능 향상에 크게 기여하며, h, x, v를 모두 사용하는 전체 EGNO 버전이 절단된 버전 대비 N-바디 시뮬레이션에서 최대 15% 높은 성능을 기록한다.
• EGNO는 재학습 없이도 시간 해상도를 두 배로 늘인 상태(2P)에서도 제로샷 일반화를 가능하게 하며, 저해상도 시간 단계 사이에서 보간함으로써 매끄럽고 정확한 궤적을 생성한다.
• 푸리에 모드 수(I)는 성능에 비단조화적인 영향을 미치며, 최적 성능은 I=2 또는 I=3에서 관찰되며, 이를 초과하면 과적합이 발생할 수 있다.
• EGNO는 자동회귀 모델이 순차적 추론이 필요한 것과 달리, 단일 순전파에서 임의의 시간 창에 걸친 미래 상태를 효율적으로 병렬 복원할 수 있다.

Figure 2: Ablation studies on the number of modes $I$ on N-body simulation and Mocap- Run datasets.