EquiformerV3: 효율성, 표현력, 범용성을 갖춘 대규모 SE(3)-불변 그래프 어텐션 트랜스포머

3D 원자 구조를 다루는 AI 모델에서 물리적 대칭성(회전, 평행이동)을 유지하면서 연산 효율을 높이는 것은 소재 과학과 신약 개발의 핵심 과제이다. 이 논문은 기존 모델의 연산 병목을 해결하고 물리적 일관성을 강화하여, 더 빠르고 정확한 분자 시뮬레이션을 가능하게 하는 차세대 아키텍처를 제시한다.

3D 공간의 원자 시스템을 다룰 때는 시스템이 회전하거나 이동해도 물리적 성질이 변하지 않아야 한다는 SE(3)-Equivariance(불변성)가 매우 중요하다. 기존의 Equiformer 시리즈는 이를 위해 복잡한 수학적 도구인 Tensor Product를 사용해왔으나, 이는 차수가 높아질수록 연산량이 기하급수적으로 늘어나는 한계가 있었다.

EquiformerV3는 이 문제를 해결하기 위해 특징값들을 구체(S²) 그리드로 투영하여 연산하는 방식을 채택했다. 복잡한 텐서 곱셈을 구체 위에서의 단순한 요소별 곱셈(Element-wise multiplication)으로 변환함으로써 연산 복잡도를 O(L⁶)에서 O(L⁴)로 획기적으로 낮추었다.

또한, 최근 LLM에서 성능이 검증된 SwiGLU 구조를 이 구체 연산 체계에 맞게 변형한 SwiGLU-S²를 도입했다. 이는 모델이 원자들 사이의 복잡한 다체 상호작용을 더 잘 포착하게 만들면서도, 물리적 대칭성을 깨뜨리지 않도록 설계되었다. 결과적으로 더 적은 연산 자원으로도 더 정교한 물리 법칙을 학습할 수 있게 되었다.

전체 아키텍처는 EquiformerV2를 기반으로 하되, 효율성과 물리적 일관성을 극대화하는 세 가지 핵심 모듈을 통합했다. 먼저 Equivariant Merged Layer Normalization은 입력 특징값 x의 각 차수 L별 RMS를 구한 뒤, 이를 전체 차수에 대해 평균 낸 통합 σ를 사용하여 정규화한다. [각 차수의 제곱 평균 산출 → 전체 차수 평균 σ 계산 → x/σ 연산 → 차수 간 상대적 크기 유지 및 학습 안정화]

어텐션 메커니즘에는 Smooth Radius Cutoff를 도입했다. 원자 i와 j 사이의 거리 r_ij에 대해 envelope 함수를 정의하고, 이를 Softmax 연산 내부에 직접 통합했다. [원자 간 거리 입력 → envelope 함수 적용 → Softmax 분모/분자에 가중치 반영 → 원자가 컷오프 경계를 넘나들 때 급격한 수치 변화 방지 및 연속적인 에너지 평면 구현]

활성화 함수인 SwiGLU-S²는 입력 특징을 ToSphere 연산을 통해 구체 그리드로 변환한 후, 스칼라 게이팅 값과 그리드 특징 간의 곱셈을 수행한다. [이레프 특징 입력 → 구체 그리드 투영 → Sigmoid(스칼라)와 그리드 값 곱셈 → FromSphere로 복원 → 다체 상호작용 정보가 포함된 불변 특징 출력]

OC20 데이터셋의 S2EF-2M 작업에서 EquiformerV3는 기존 모델 대비 학습 효율을 최대 5.9배 향상시켰다. 특히 OMat24 벤치마크에서는 파라미터 수가 23배 더 많은 대형 모델(UMA-L)과 대등한 수준의 힘 예측 오차(MAE)를 기록하면서도 모델 크기는 훨씬 작게 유지했다.

Matbench Discovery 리더보드에서는 열전도율 예측 등 고차 미분이 필요한 작업에서 기존 SOTA 모델인 eSEN 대비 18%에서 31%의 성능 향상을 보였다. 또한, 전체적인 성능 지표인 CPS(Combined Performance Score)에서 최초로 0.9를 돌파하며 1위를 차지했다.

Ablation Study 결과, 제안된 SwiGLU-S² 활성화 함수와 Merged LN이 에너지 및 힘 예측 정확도를 높이는 데 결정적인 역할을 했음이 확인되었다. 특히 SwiGLU-S²를 사용했을 때 에너지 MAE가 약 41meV 감소하는 성과를 거두었다.

EquiformerV3는 SE(3)-Equivariant GNN의 고질적인 문제인 고차 이레프(Irreps) 연산 비용을 해결하기 위해 Fast Tensor Product 기법을 적극 활용한다. 핵심은 SO(3) 텐서 곱을 SO(2) 선형 레이어와 구체 그리드 상의 연산으로 분해하는 것이다. 이를 통해 Clebsch-Gordan 계수를 명시적으로 계산하는 비용을 회피한다.

아키텍처 측면에서는 Transformer의 FFN(Feedforward Network) 히든 레이어 크기를 4배 확장하여 모델 용량을 키웠으며, 이 과정에서 발생하는 연산 부담은 SwiGLU-S²의 효율적인 게이팅 메커니즘으로 상쇄했다. 또한, DeNS(Denoising Non-equilibrium Structures)라는 보조 학습 태스크를 사용하여 평형 상태가 아닌 구조에 대한 복원 능력을 길러 모델의 일반화 성능을 강화했다.

구현 상으로는 PyTorch의 torch.compile과 호환되도록 동적 셰이프 문제를 해결하고, 커스텀 CUDA 커널을 통해 eSCN 컨볼루션 연산을 최적화했다. 이는 단순히 알고리즘적 개선을 넘어 실제 하드웨어 가속기에서의 실행 효율을 극대화한 결과이다.

논문은 현재의 데이터 큐레이션 방식(예: 구조 필터링)이나 데이터 증강 기법이 Matbench Discovery와 같은 특정 벤치마크에서 성능 병목이 될 수 있음을 언급하며, 향후 연구 과제로 남겨두었다.