[논문 리뷰] Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products

Shengjie Luo, T. Chen외 1명|arXiv (Cornell University)|2024. 01. 18.

Medical Imaging and Analysis인용 수 4

한 줄 요약

이 논문은 기하학적 표현의 기저가 되는 구면 조화함수를 2차원 푸리에 기저로 재표현함으로써, E(3)-동치성 연산의 계산 복잡도를 감소시켜 3차원 딥러닝에서의 효율성을 높이는 Gaunt 텐서 곱을 소개한다. 이는 기저가 되는 구면 조화함수를 2차원 푸리에 기저로 재표현함으로써 컨volution 정리와 빠른 푸리에 변환(FFT)을 활용하여, 기존의 O(L⁶)에서 O(L³)로 복잡도를 감소시킨다. 이로 인해 고차수의 불변 표현(irreducible representation, irrep) 연산을 효율적으로 수행할 수 있으며, 분자의 데이터셋에서 모델 성능이 향상된다.

ABSTRACT

Developing equivariant neural networks for the E(3) group plays an important role in modeling 3D data across real-world applications. Enforcing this equivariance primarily involves the tensor products of irreducible representations (irreps). However, the computational complexity of such operations increases significantly as higher-order tensors are used. In this work, we propose a systematic approach to substantially accelerate the computation of the tensor products of irreps. We mathematically connect the commonly used Clebsch-Gordan coefficients to the Gaunt coefficients, which are integrals of products of three spherical harmonics. Through Gaunt coefficients, the tensor product of irreps becomes equivalent to the multiplication between spherical functions represented by spherical harmonics. This perspective further allows us to change the basis for the equivariant operations from spherical harmonics to a 2D Fourier basis. Consequently, the multiplication between spherical functions represented by a 2D Fourier basis can be efficiently computed via the convolution theorem and Fast Fourier Transforms. This transformation reduces the complexity of full tensor products of irreps from $\mathcal{O}(L^6)$ to $\mathcal{O}(L^3)$, where $L$ is the max degree of irreps. Leveraging this approach, we introduce the Gaunt Tensor Product, which serves as a new method to construct efficient equivariant operations across different model architectures. Our experiments on the Open Catalyst Project and 3BPA datasets demonstrate both the increased efficiency and improved performance of our approach.

연구 동기 및 목표

• E(3)-동치성 딥러닝에서 불변 표현(irreps)의 전체 텐서 곱의 높은 계산 복잡도 문제를 해결함으로써, 실무에서 고차수의 불변 표현을 사용하는 데 제약를 줄이기 위해.
• 클레브시-고르단 계수와 가운트 계수 사이의 수학적 연결을 수립함으로써, 불변 표현 텐서 곱을 구면 함수의 곱셈으로 재해석할 수 있도록 하기 위해.
• 구면 조화함수에서 2차원 푸리에 기저로 기저를 전환하여 컨volution 정리를 활용하고, 빠른 푸리에 변환(FFT)을 적용함으로써 이러한 연산의 효율적 계산을 가능하게 하기 위해.
• 다양한 아키텍처와 응용 분야에 걸쳐 사용 가능한 일반적이고 효율적인 빌딩 블록인 Gaunt 텐서 곱을 개발하기 위해.
• 오픈 캐탈리스트 프로젝트 및 3BPA 데이터셋을 포함한 실제 3차원 기하학적 학습 벤치마크에서 더 높은 효율성과 향상된 성능을 입증하기 위해.

제안 방법

• 클레브시-고르단 계수와 가운트 계수 사이의 관계를 정량적으로 연결함으로써, 이는 세 개의 구면 조화함수의 적분으로 정의되며, 불변 표현 텐서 곱을 구면 함수의 곱셈으로 재해석할 수 있도록 한다.
• 구면 함수를 2차원 푸리에 기저로 재표현함으로써, 함수 간 곱셈을 컨볼루션 정리로 가속화할 수 있도록 한다.
• 빠른 푸리에 변환(FFT)을 적용하여 푸리에 도메인의 컨볼루션을 효율적으로 계산함으로써, 불변 표현 텐서 곱의 전체 복잡도를 감소시킨다.
• 기존의 텐서 곱 연산을 대체할 수 있는 새로운 미분 가능한 레이어인 Gaunt 텐서 곱을 유도함으로써, 동치성은 유지하면서도 계산 효율성을 향상시킨다.
• EquiformerV2와 MACE 등의 모델 아키텍처에 이 방법을 구현하여, 벤치마크 데이터셋에서 공정한 비교를 위한 표준 훈련 프rotocol를 사용한다.
• 불변 표현 기반의 동치성 네트워크의 보편 근사 능력을 유지하면서도 계산 비용을 극적으로 감소시킴을 보장한다.

실험 결과

연구 질문

• RQ1 E(3)-동치성 네트워크에서 불변 표현의 텐서 곱은 푸리에 도메인으로 문제를 재표현함으로써 가속화될 수 있는가?
• RQ2 클레브시-고르단 계수와 가운트 계수 사이에 수학적으로 엄밀한 연결 고리가 존재하는가? 이 연결 고리는 이러한 가속화를 가능하게 하는가?
• RQ3 구면 조화함수에서 2차원 푸리에 기저로 기저를 전환하면, FFT를 통해 효율적인 계산이 가능해지며, 동치성이 유지되는가?
• RQ4 제안된 Gaunt 텐서 곱은 고차수 불변 표현 설정에서 훈련 효율성과 모델 성능을 어느 정도 향상시키는가?
• RQ5 이 방법은 3차원 기하학적 학습 분야의 다양한 동치성 모델 아키텍처와 벤치마크 데이터셋으로 일반화될 수 있는가?

주요 결과

• Gaunt 계수와 2차원 푸리에 기저 변환을 활용함으로써, 불변 표현(irreps)의 전체 텐서 곱의 계산 복잡도가 O(L⁶)에서 O(L³)로 감소된다.
• Gaunt 텐서 곱을 통해 EquiformerV2에서 고차수 불변 표현(L=6까지)을 사용할 수 있게 되었으며, 이는 유의미하게 향상된 추론 효율성과 메모리 사용 감소를 가져왔다.
• 오픈 캐탈리스트 프로젝트 데이터셋에서, Gaunt 텐서 곱을 사용한 모델은 기준 모델 대비 더 빠른 훈련 및 추론 시간을 기록하며 최고 성능을 달성했다.
• 3BPA 데이터셋에서, 이 방법은 특히 분포 외(out-of-distribution) 및 다이헤드럴 슬라이스(dihedral-slice) 테스트 세트에서 뛰어난 일반화 성능을 보였으며, 더 부드럽고 정확한 잠재 에너지 표면을 나타내었다.
• 계산 비용을 감소시키면서도 모델 정확도를 유지하거나 향상시켰으며, EquiformerV2와 MACE를 포함한 여러 아키텍처에서 일관된 성능 향상이 관찰되었다.
• Gaunt 텐서 곱은 기존의 동치성 모델 설계와 호환되며, 표준 텐서 곱 레이어의 즉각적인 대체로 간편하게 통합될 수 있다.