[논문 리뷰] Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs

Saro Passaro, C. Lawrence Zitnick|arXiv (Cornell University)|2023. 02. 07.

Advanced Neural Network Applications인용 수 27

한 줄 요약

본 논문은 주 축을 엣지 방향에 맞춰 SO(3) 등가 컨볼루션을 효율적인 SO(2) 연산으로 변환하는 방법을 제시하고, 텐서 곱의 복잡도를 O(L^6)에서 O(L^3)로 감소시키며, Equivariant Spherical Channel Network (eSCN)을 도입하여 OC-20/OC-22 데이터셋에서 최첨단 결과를 달성한다.

ABSTRACT

Graph neural networks that model 3D data, such as point clouds or atoms, are typically desired to be $SO(3)$ equivariant, i.e., equivariant to 3D rotations. Unfortunately equivariant convolutions, which are a fundamental operation for equivariant networks, increase significantly in computational complexity as higher-order tensors are used. In this paper, we address this issue by reducing the $SO(3)$ convolutions or tensor products to mathematically equivalent convolutions in $SO(2)$ . This is accomplished by aligning the node embeddings' primary axis with the edge vectors, which sparsifies the tensor product and reduces the computational complexity from $O(L^6)$ to $O(L^3)$, where $L$ is the degree of the representation. We demonstrate the potential implications of this improvement by proposing the Equivariant Spherical Channel Network (eSCN), a graph neural network utilizing our novel approach to equivariant convolutions, which achieves state-of-the-art results on the large-scale OC-20 and OC-22 datasets.

연구 동기 및 목표

• SO(3)-등가 그래프 신경망의 계산 효율성을 향상시키기 위해 텐서 곱을 SO(2) 컨볼루션으로 축소한다.
• 실용적으로 더 높은 차수의 불변 표현(irreps)을 비용이 크게 들지 않게 가능하게 한다.
• 대규모 원자 데이터셋(OC-20/OC-22)에서 최첨단 성능을 시연한다.
• GNN의 메시지 패싱에 대한 SO(3)와 SO(2) 공식 간의 원칙적 연결 고리를 제공한다.

제안 방법

• 불변 표현의 주 축을 엣지 방향과 일치시켜 텐서 곱을 희소화한다.
• 회전된 SO(3) 컨볼루션은 원형 조화에 대한 등가 SO(2) 컨볼루션으로 계산될 수 있음을 보인다.
• 전체 Clebsch–Gordan 계수 계산을 피하고 복잡도를 O(L^3)로 감소시키는 계산적으로 효율적인 형식을 도출한다.
• SO(2) 형식을 활용한 메시지 패싱을 위한 Equivariant Spherical Channel Network (eSCN)을 제안한다.
• GPU 친화적인 재배열을 사용하여 m 채널과 차수에 걸친 컨볼루션을 병렬화한다.

실험 결과

연구 질문

• RQ1 SO(3) 등가 컨볼루션을 등가성을 잃지 않으면서 효율적인 SO(2) 컨볼루션으로 계산할 수 있는가?
• RQ2 축 정렬이 3D 데이터에 대해 GNN에서 더 높은 차수의 불변 표현을 실용적으로 사용할 수 있게 하는가?
• RQ3 대규모 오픈 촉매 데이터셋에서 eSCN의 성능이 최첨단 등가 GNN과 비교하여 어떤가?

주요 결과

• 축 정렬된 SO(2) 형식을 사용할 때 텐서 곱 계산이 O(L^6)에서 O(L^3)로 감소한다.
• 주 축을 엣지 방향과 맞추면 Y^(l) 항이 희소해져 계산 속도가 빨라지고 Clebsch–Gordan 계수의 필요성이 사라진다.
• eSCN은 OC-20과 OC-22에서 최첨단 결과를 달성하며, 힘 예측에서 OC-20에서 9%, OC-22에서 21%의 유의미한 개선과 느슨한 구조 예측에서 15% 개선을 보인다.
• 적절히 불변 표현을 회전시킬 때 SO(3) 등가 연산은 일반화된 SO(2) 컨볼루션의 집합으로 간주될 수 있다.