[논문 리뷰] Fast, Expressive SE$(n)$ Equivariant Networks through Weight-Sharing in Position-Orientation Space

Erik J. Bekkers, Sharvaree Vadgama외 3명|arXiv (Cornell University)|2023. 10. 04.

Advanced Image and Video Retrieval Techniques인용 수 5

한 줄 요약

이 논문은 위치-방향성 공간 (R³×S²) 상에서 가중치 공유를 통해 빠르고 표현력 있는 SE(n) 동치 네트워크를 제안한다. 이는 SE(n) 변환 하에서 점-쌍의 동치류를 공식화함으로써 가중치 공유를 정의한다. 이러한 동치류를 유일하게 식별하는 기하적 특성에 기반해 소통 기능을 조건화함으로써, 원자 간 에너지 예측, N체계 궤도 예측, 확산 모델을 통한 분자 생성을 포함한 3D 포인트 클러스터 벤치마크에서 최고 수준의 정확도와 속도를 달성한다.

ABSTRACT

Based on the theory of homogeneous spaces we derive geometrically optimal edge attributes to be used within the flexible message-passing framework. We formalize the notion of weight sharing in convolutional networks as the sharing of message functions over point-pairs that should be treated equally. We define equivalence classes of point-pairs that are identical up to a transformation in the group and derive attributes that uniquely identify these classes. Weight sharing is then obtained by conditioning message functions on these attributes. As an application of the theory, we develop an efficient equivariant group convolutional network for processing 3D point clouds. The theory of homogeneous spaces tells us how to do group convolutions with feature maps over the homogeneous space of positions $\mathbb{R}^3$, position and orientations $\mathbb{R}^3 { imes} S^2$, and the group $SE(3)$ itself. Among these, $\mathbb{R}^3 { imes} S^2$ is an optimal choice due to the ability to represent directional information, which $\mathbb{R}^3$ methods cannot, and it significantly enhances computational efficiency compared to indexing features on the full $SE(3)$ group. We support this claim with state-of-the-art results -- in accuracy and speed -- on five different benchmarks in 2D and 3D, including interatomic potential energy prediction, trajectory forecasting in N-body systems, and generating molecules via equivariant diffusion models.

연구 동기 및 목표

• 기존의 컨볼루션 설정을 초월한 메시지 전달 네트워크에서의 가중치 공유를 공식화하는 것.
• 동차 공간 이론을 활용해 SE(n) 동치 메시지 기능에 대해 기하학적으로 최적의 특성을 도출하는 것.
• SE(n) 동치성을 유지하면서도 3D 포인트 클러스터 처리에 효율적이고 표현력 있는 아키텍처를 개발하는 것.
• 분자의 및 물리적 시스템을 포함한 다양한 3D 학습 벤치마크에서 최고 수준의 성능을 입증하는 것.

제안 방법

• SE(n) 변환 하에서 점-쌍의 동치류를 유일하게 식별하는 특성에 기반해 메시지 기능을 조건화함으로써 가중치 공유를 공식화하는 것.
• 동차 공간 이론을 활용해 기하학적 특성을 도출하고, 방향성 정보와 효율성 면에서 R³×S²가 최적의 특성 공간임을 규명하는 것.
• R³×S² 상에서 분리 가능한 군 컨벌루션을 구성함으로써 방향성과 채널에 대해 병렬 처리가 가능하게 하여 계산 효율성을 높이는 것.
• 구형 컨벌루션과 배치 채널 혼합을 활용해 추론 및 훈련 속도를 가속화하는 것.
• 상대 위치, 상대 방향, 상호작용 특성(예: 전하 곱)과 같은 불변 특성에 기반해 메시지 기능을 조건화하는 것.
• ConvNeXt 블록과 학습 가능한 구형 임베딩을 활용해 입력 및 출력 특성에 대해 완전히 컨벌루션 아키텍처를 구현하는 것.

실험 결과

연구 질문

• RQ1 메시지 전달 네트워크에서의 가중치 공유를 이동 동치성 컨벌루션을 초월해 전체 SE(n) 동치성으로 일반화할 수 있는가?
• RQ2 SE(n) 변환 하에서 점-쌍의 동치류를 유일하게 식별하는 기하학적 특성은 무엇인가?
• RQ3 왜 R³ 또는 전체 SE(3)에 비해 위치-방향성 공간 (R³×S²) 이 SE(3) 동치 학습에 더 효율적이고 표현력이 강한가?
• RQ4 순수하게 컨벌루션 기반의 SE(n)-동치 아키텍처가 다양한 3D 포인트 클러스터 벤치마크에서 최고 수준의 성능을 달성할 수 있는가?
• RQ5 기하학적 불변 특성에 기반해 메시지 기능을 조건화하면 동치 메시지 전달의 표현력과 효율성이 어떻게 향상되는가?

주요 결과

• 제안된 방법은 원자 간 잠재 에너지 예측에서 기존 방법을 뛰어넘는 정확도와 추론 속도를 달성하여 최고 수준의 성능을 기록한다.
• N체계 궤도 예측 작업에서는 기존의 동치 모델 대비 상당히 빠른 추론 시간을 기록하며 최고 수준의 결과를 도출한다.
• 확산 모델을 통한 분자 생성 작업에서는 높은 품질의 물리적으로 타당한 분자를 생성하며 샘플링 효율성이 향상된다.
• 위치-방향성 공간 (R³×S²) 의 사용은 전체 SE(3) 색인화에 비해 계산 비용을 감소시키면서도 방향 감도를 유지하고 성능을 향상시킨다.
• 이론적 프레임워크는 오직 기하학적 특성과 공유된 메시지 기능만을 사용해 SE(n)-동치 함수의 보편 근사가 가능함을 보여준다.
• 실험 결과는 유도된 특성에 기반해 메시지 기능을 조건화함으로써 모든 벤치마크에서 일반화 능력 향상과 더 빠른 수렴이 이루어짐을 확인한다.