기술 요약: 플라토닉 트랜스포머(Platonic Transformers): 등변성(Equivariance)을 위한 견고한 선택

문제 정의

트랜스포머는 유연성과 효율성 덕분에 딥러닝 분야에서 광범위한 성공을 거두었으나, 과학적 도메인(물리학, 화학) 및 컴퓨터 비전(3D 포인트 클라우드)에 널리 퍼져 있는 기하학적 대칭성에 대한 내재적 귀납 편향(inductive biases)이 부족합니다. 이러한 대칭성을 준수하는 기존의 등변성 방법들은 군 합성곱(group convolutions)이나 클레브슈-고르단(Clebsch-Gordan) 텐서 곱과 같은 복잡한 연산에 의존하는 경우가 많습니다. 이러한 접근 방식은 상당한 계산 오버헤드와 구조적 복잡성을 초래하여, 표준 트랜스포머 아키텍처가 가진 속도, 유연성, 계산 효율성을 저해합니다. 핵심 과제는 표준 트랜스포머의 속도, 유연성, 계산 효율성을 희생하지 않으면서 강력한 기하학적 귀납 편향을 트랜스포머 프레임워크에 통합하는 것입니다.

방법론: 플라토닉 트로스포머(The Platonic Transformer)

저자들은 표준 트랜스포머의 기본 어텐션 메커즘이나 계산 그래프를 변경하지 않고도, 연속적인 평행 이동(translations)과 이산적인 회전 반사(roto-reflections, 플라토닉 다면체의 대칭성)에 대한 등변성을 달 Achieve하는 프레임워크인 플라토닉 트랜스포머를 제안합니다.

핵심 메커니즘

1. 참조 프레임 리프팅(Reference Frame Lifting): 입력 특징을 단일 글로벌 좌표계에서 처리하는 대신, 이 방법은 입력 특징을 이산 대칭군 $G \subset O(n)$(예: 정사면체, 정팔면체 또는 정이십면체의 회전 대칭군)에 대한 함수로 리프팅합니다. 특징은 $G$의 원소들로 정의된 일련의 참조 프레임들에 상대적으로 표현됩니다.
2. 프레임 의존적 RoPE(Frame-Dependent RoPE): 이 방법은 회전 위치 임베딩(Rotary Position Embeddings, RoPE)을 일반화합니다. 표준 RoPE가 평행 이동 등변성을 보장하는 반면, 플라토닉 트랜스포머는 각 참조 프레임 $R \in G$에 대해 상대적으로 위치 인코딩을 재정의합니다. 토큰의 위치가 $p_i$일 때, 위치는 각 프레임에 대해 $p_i(R) = R^{-1}p_i$로 투영됩니다. 그 후 RoPE 연산자가 이러한 상대적 위치에 적용됩니다.
3. 가중치 공유 스킴(Weight-Sharing Scheme): 모든 선형 레이어(Query, Key, Value 투영 및 MLP)는 그룹 $G$에 대해 등변성을 갖도록 제한됩니다. 수학적으로 이는 선형 레이어를 **군 합성곱(group convolutions)**으로 변환합니다. 평탄화된 벡터 표현에서, 이는 가중치 행렬이 학습 가능한 커널에 의해 결정되는 블록 행렬이 되는 구조화된 가중치 공유 스킴으로 나타납니다. 이는 파라미터 수를 제약 없는 레이어에 비해 $|G|$ 배만큼 줄이면서 동일한 계산 그래프 구조를 유지합니다.
4. 병렬 어텐션 헤드(Parallel Attention Heads): 어텐션 메커니즘은 모든 참조 프레임에 걸쳐 병렬로 작동합니다. 각 프레임은 사실상 별개의 어텐션 헤드 역할을 합니다. 원시 어텐션 점수는 $s_{ij}(R) = q_i(R)^\top \rho_\Omega(p_j(R) - p_i(R)) k_j(R)$로 계산됩니다(여기서 $\rho_\Omega$는 RoPE 회전 행렬). 소프트맥스는 프레임별로 적용되어 모델이 방향 의존적인 어텐션 패턴을 학습할 수 있게 합니다.

이론적 통찰

• 동적 군 합성곱(Dynamic Group Convolution): 저자들은 RoPE 기반 어텐션(소프트맥스가 생략되거나 선형 어텐션 설정인 경우)과 동적 군 합성곱 사이의 공식적인 동등성을 입증했습니다. 쿼리 벡터는 희소 푸리에 기저 확장(sparse Fourier basis expansion)의 계수 역할을 하여, 온 더 플라이(on the fly)로 콘텐츠 인식 기하학 필터를 합성합니다.
• 선형 복잡도: 이러한 합성곱 관점은 이 메커니즘이 토큰 수에 선형적인 $O(N)$ 시간 복로 계산될 수 있음을 보여주며, 이는 표준 어텐션의 이차 복잡도 $O(N^2)$에 대한 확장 가능한 대안을 제공합니다.

주요 기여

1. 아키텍처: 플라토닉 다면체 대칭군을 통한 원칙적인 가중치 공유 스킴을 통해 표준 트랜스포머에 기하학적 등변성을 통합하는 플라토닉 트랜스포머를 도입하였으며, 이를 통해 베이스라인과 동일한 아키텍처 및 계산 비용을 보존합니다.
2. 이론적 통합: RoPE 기반 어텐션이 동적 군 합성곱과 공식적으로 동등함을 입증하여, 모델이 적응형 기하학 필터를 학습할 수 있는 능력을 밝혀냈습니다.
3. 확장성: 대규모 3D 데이터에 적합하도록 효율적으로 확장 가능한 선형 시간 합성곱 변형을 제공합니다.
4. 실험적 검증: 다양한 벤치마크에 대한 종합적인 평가를 통해 기하학적 제약이 추가적인 계산 비용 없이 성능을 향상시킨다는 것을 보여주었습니다.

실험 결과

본 논문은 "비등변 작업(Non-Equivariant Tasks, 데이터에 정형화된 방향이 있는 경우)"과 "등변 작업(Equivariant Tasks, 데이터 방향이 임의적인 경우)" 모두에 대해 모델을 평가합니다.

• CIFAR-10 (2D 비전): 2D 회전 대칭(예: $C_6, D_3$)을 통합하는 것이 평행 이동 전용 베이스라인(RoPE가 적용된 ViT)보다 정확도를 높였습니다. 특히, 가중치 공유 덕분에 파라미터 수가 적은 가장 효율적인 구성이 최고의 성능을 보였는데, 이는 성능 향상이 용량 증가보다는 대칭성에 의한 규제(regularization)에서 기인함을 시사합니다.
• ScanObjectNN (3D 포인트 클라우드): 방향 가변성과 폐쇄(occlusion)가 심한 작업에서, 정사면체 그룹 변형이 베이스라인보다 어텐션에서는 2.43%, 합성곱에서는 1.54% 더 높은 성능을 보였습니다. 선형 시간 합성곱 변형은 경쟁력 있는 정확도를 유지하면서도 상당한 속도 향상을 보여주었습니다.
• QM9 (분자 생성 및 물성 예측):
  • 생성: 확산 트랜스포머(DiT) 내의 표준 트랜스포머를 플라토닉 트랜스포머(DiP)로 교체했을 때, ADiT, GeoLDM, EDM과 같은 베이스라인보다 더 높은 유효성과 현저히 개선된 분자 안정성을 가진 분자를 생성했습니다. 등변 변형(DiP-Tetrahedron)은 엄격한 기하학적 건전성 체크(PoseBusters)에서 완벽에 가까운 통과율을 달성했습니다.
  • 물성 예측: 표준 QM9 물성 예측 작업에서 경쟁력 있는 결과를 달성했습니다.
• OMol25 (대규모 힘/에너지 예측): 4백만 개의 구조를 가진 대규모 데이터셋에서, 플라토닉 트랜스포머(정사면체 및 정팔면체 변형)는 힘과 에너지 예측 정확도(MAE) 측면에서 최신 모델인 eSEN-sm 베이스라인을 크게 상회하는 동시에 2.6배 높은 학습 처리량을 달성했습니다. 대칭 그룹 크기를 키움으로써(정사면체에서 정팔면체로) 성능 손실을 최소화하면서 파라미터를 절반으로 줄였으며, 이는 가중치 공유 스킴의 효율성을 다시 한번 입증했습니다.
• ProteinMD: 모델은 단백질 역학 힘 예측에서 이전의 최첨단 방법들(G-Hyena 및 EGNN 포함)보다 우수한 성능을 보였습니다.

의의 및 주장

본 논문은 플라토닉 트랜스포머가 **대칭 인식(symmetry-awareness)**과 확장성(scalability) 사이의 오랜 트레이드오프를 해결한다고 주장합니다. 플라토닉 다면체의 기하학적 구조를 활용하여 가중치 공유를 유도함으로써, 이 방법은 트랜스포머의 핵심 계산 그래프를 수정하거나 계산 비용을 증가시키지 않고도 근사적인 $E(n)$ 등변성을 달성합니다.

저자들은 이 접근 방식이 원칙적인 등변성과 현대적인 확장성이 서로 배타적이지 않음을 보여준다고 강조합니다. 또한, 동적 군 합성곱과의 공식적인 연결은 모델이 적응형 기하학 필터를 학습할 수 있는 능력에 대한 이론적 토대를 제공하며, 이를 통해 매우 확장 가능한 선형 시간 변형을 가능하게 합니다. 결과적으로, 물리 법칙에 의해 지배되는 과학적 도메인(여기서 등변성은 '플라토닉 이상'임)을 위해 이러한 제약을 플라토닉 트랜스포머를 통해 강제하는 것이 표준 아키텍처나 복잡한 등변성 대안 모델보다 더 데이터 효율적이고, 견고하며, 물리적으로 타당한 모델을 만든다는 것을 시사합니다.