구면 등변 그래프 트랜스포머 완전 가이드

초록

본 가이드는 3차원 분자·생체분자 시스템에 적용되는 구면 등변 그래프 신경망(EGNN)의 수학적 기초와 구현 방법을 체계적으로 정리한다. 그룹 표현, 구면 조화함수, 텐서곱·클레브시–고르단 분해, SO(3) 등변 커널 설계 과정을 상세히 설명하고, 이를 기반으로 Tensor Field Network와 SE(3)-Transformer의 메시지 패싱·어텐션 메커니즘을 단계별로 구현한다. 마지막으로 QM9 데이터셋을 이용한 화학적 특성 예측 파이프라인을 제시한다.

상세 분석

이 논문은 구면 등변 그래프 신경망(EGNN)의 핵심 이론을 물리학과 양자역학에서 차용한 군 이론으로 풀어낸다. 먼저 SO(3) 회전군의 불변성과 등변성을 정의하고, 그룹 표현 ρ(g) 를 통해 입력 텐서가 어떻게 변환되는지를 수식적으로 기술한다. 구면 텐서는 차수 l 에 따라 irreducible representation 으로 분류되며, 각 차수는 구면 조화함수 Y_l^m 로 전개된다. 텐서곱을 수행할 때는 클레브시–고르단(Clebsch–Gordan) 계수를 이용해 두 텐서를 새로운 차수 J 로 결합하고, |k−l| ≤ J ≤ k+l 범위 내 모든 가능한 J 를 고려한다. 이 과정은 등변 커널의 기저 함수를 구성하는 핵심이며, 방사형 함수 f(r) 와 결합해 최종 커널 K_{lk}(r,Ω)=∑{J}C{lk}^J f_J(r) Y_J(Ω) 형태가 된다.

커널 설계 단계에서는 (1) 회전 변환에 대한 제약식 ρ_l(g)K_{lk}=K_{lk}ρ_k(g) 를 만족하도록 기저를 선택하고, (2) 학습 가능한 방사형 파라미터를 신경망으로 파라미터화한다. 논문은 이러한 커널을 이용해 Tensor Field Network(TFN)의 메시지 패싱을 구현한다. TFN에서는 각 노드의 특징을 차수 l 텐서로 유지하면서, 이웃 노드와의 상호작용을 K_{lk} 로 가중합하고, 이후 선형 self‑interaction 과 채널 믹싱을 통해 차원을 조정한다.

SE(3)-Transformer는 TFN의 메시지 패싱을 어텐션 메커니즘으로 확장한다. 쿼리·키·밸류 텐서는 각각 다른 차수의 텐서로 인코딩되며, 어텐션 스코어는 ⟨Q_i, K_j⟩ 형태의 내적을 구해 회전 불변성을 확보한다. 여기서 내적은 클레브시–고르단 계수를 포함한 텐서곱을 통해 구현되며, 스코어는 소프트맥스 후 밸류 텐서와 결합돼 최종 노드 업데이트를 만든다. 다중 헤드 어텐션은 서로 다른 차수와 방사형 함수를 병렬로 학습하게 해 표현력을 크게 향상시킨다.

실험 섹션에서는 QM9 데이터셋을 사용해 분자 특성(에너지, 전하 등) 예측을 수행한다. 그래프는 원자를 노드, 원자 간 거리 벡터를 엣지 특징으로 구성하고, 초기 임베딩은 원자 번호와 전자 배치 정보를 포함한다. 모델은 3~5개의 등변 레이어와 다중 헤드 어텐션을 쌓아 학습하며, 데이터 증강 없이도 회전 불변성을 보장해 기존 비등변 GNN 대비 높은 정확도와 빠른 수렴을 보인다.

전체적으로 이 가이드는 복잡한 수학적 배경을 코드와 함께 단계별로 풀어내어, 연구자와 엔지니어가 직접 구현하고 확장할 수 있는 실용적인 로드맵을 제공한다. 특히, 텐서 차수 관리, 클레브시–고르단 계수 계산, 방사형 함수 파라미터화 등 구현 난관을 상세히 다루어, 구면 등변 모델을 처음 접하는 독자도 이해하고 적용할 수 있도록 설계되었다.