🤖 AI

E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory

E2Former-V2는 등변 축 정렬 희소화(Equivariant Axis-Aligned Sparsification)와 온더플라이 등변 어텐션(On-the-Fly Equivariant Attention) 메커니즘을 도입하여 3D 원자계에 대한 예측 성능을 유지하면서도 20배의 TFLOPS 향상을 달성함으로써 주류 모델들의 계산 병목 현상을 극복하는 확장 가능한 등변 그래프 신경망 아키텍처입니다.

원저자: Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

게시일 2026-06-08

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신은 거대한 군중(원자)이 방 안에서 어떻게 움직이고 상호작용할지 예측하려고 노력하고 있다고 상상해 보세요. 화학과 물리학의 세계에서 이것을 "원자 단위 모델링(atomistic modeling)"이라고 부릅니다. 오랫동안 과학자들은 **등변 그래프 신경망(Equivariant Graph Neural Network, EGNN)**이라는 특별한 형태의 AI를 사용해 왔습니다. 이 네트워크는 방 전체를 회전시키더라도 그 안의 물리 법칙은 변하지 않는다는 점을 잘 이해하고 있기 때문에 매우 뛰어납습니다.

하지만 현재 버전의 이 네트워크들은 군중이 커질 때 서투르고 느리다는 중대한 문제를 가지고 있습니다.

다음은 논문의 해결책인 E2Former-V2를 일상적인 비유를 사용하여 쉽게 설명한 내용입니다.

문제점: "에지 중심(Edge-Centric)" 교통 체증

모든 사람이 이웃과 대화하고 싶어 하는 파티를 상상해 보세요.

기존 방식 (전통적인 EGNN): 이 시스템은 파티에서 일어나는 모든 개별 대화를 거대한 물리적 목록으로 만듭니다. 만약 1,000명이 있고 각자가 50명의 이웃과 대화한다면, 시스템은 수학적 계산을 수행하기 전에 거대한 화이트보드(메모리)에 50,000개의 별도 대화 노트를 적어야 합니다.
병목 현상: 파티 규모가 커질수록(원자가 많아질수록), 이 화이트보드는 너무 거대해져서 컴퓨터의 공간이 부족해집니다. 컴퓨터는 실제로 생각을 하기보다는 이 노트들을 옮기는 데 모든 시간을 허비하게 됩니다. 이것이 기존 모델들이 단백질과 같은 큰 분자를 시뮬레이션할 때 왜 멈추거나 속도가 급격히 느려지는 이유입니다.

해결책: E2Former-V2

저자들은 모든 대화 노트를 일일이 적는 대신, 과정을 빛의 속도로 빠르게 만드는 두 가지 영리한 기술을 구축했습니다.

기술 1: "마법 나침반" (등변 축 정렬 희소화/Equivariant Axis-Aligned Sparsification)

기존 시스템에서 원자 간의 상호작용을 계산하는 것은 모든 조각이 어떤 방향으로든 회전할 수 있는 퍼즐을 푸는 것과 같습니다. 이는 매우 복잡하고 많은 수학적 계산을 요구합니다.

비유: 당신이 복잡한 3D 퍼즐을 가지고 있다고 상상해 보세요. 퍼즐이 빙글빙글 도는 동안 문제를 풀려고 애쓰는 대신, "마법 나침반"을 사용하여 퍼즐 전체를 회전시켜 특정 조각 하나가 항상 위를 향하도록 만듭니다.
결과: 퍼즐이 정렬되면 대부분의 조각은 빈 공간이 되거나 0이 됩니다. 더 이상 그 빈 부분에 대해 수학 계산을 할 필요가 없습니다. 논문에서는 이를 EAAS라고 부릅니다. 이는 복잡하고 밀집된 계산을 단순한 "재색인(re-indexing)" 게임(단순히 조각들을 옮기는 작업)으로 바꾸어 놓으며, 훨씬 더 빠르게 작동합니다.

기술 2: "스트리밍 셰프" (온더플라이 등변 어텐션/On-the-Fly Equivariant Attention)

이것이 가장 큰 돌파구입니다.

기존 방식 (뷔페): 기존 시스템은 누군가 음식을 먹기 전에 모든 뷔페 음식(50,000개의 대화 노트)을 준비하여 테이블 위에 올려둡니다. 만약 테이블이 너무 작으면 음식이 넘쳐흐르고 파티는 중단됩니다.
새로운 방식 (스트리밍 셰프): 새로운 시스템인 E2Former-V2는 고객 앞에서 한 번에 한 접시씩 요리하여 즉시 서빙하는 셰프처럼 행동합니다.
- 노트를 따로 적지 않습니다.
- 전체 목록을 저장하지 않습니다.
- 두 사람 사이의 상호작용을 계산하고, 그 결과를 최종 점수에 더한 뒤, 다음 쌍을 위해 즉시 세부 사항을 잊어버립니다.
결과: 이것을 "온더플라이(On-the-Fly)" 처리라고 합니다. 거대한 화이트보드를 만들지 않기 때문에 메모리를 거의 사용하지 않습니다. 이 방식은 모든 것을 컴퓨터의 초고속 내부 캐시(SRAM)에 담아내어, 기존 방식보다 20배 더 빠르게 만듭니다.

그들이 실제로 달성한 성과

저자들은 이 새로운 시스템을 두 가지 주요 항목에 대해 테스트했습니다.

작은 분자 (SPICE): 약물과 유사한 분자들을 테스트했습니다. 새로운 시스템은 기존의 가장 우수한 모델들과 대등한 정확도를 보이면서도 훨씬 빨랐습니다.
거대 분자 (OMol25): 거대한 화학 구조들을 테스트했습니다.
- 한계: 기존 모델들은 10,000개 이상의 원자를 가진 시스템을 시뮬레이션할 때 메모리 부족으로 충돌(crash)했습니다.
- 승리: E2Former-V2는 충돌 없이 100,000개의 원자를 가진 시스템을 성공적으로 시뮬레이션했습니다.

언급된 실제 물리 테스트

저자들은 단순히 숫자만 돌린 것이 아니라, 실제 물리학을 시뮬레이션했습니다.

액체 물: 물 분자 박스를 시뮬레이션했습니다. 새 모델은 물 분자들이 배열되는 방식(방사 분포 함수/radial distribution function)을 거의 완벽하게 예측하여, 이전 모델들보다 실제 실험 결과와 더 잘 일치하는 모습을 보였습니다.
단백질: 물속에 떠 있는 거대한 단백질(약 30,000개 원자)을 시뮬레이션했습니다. 이 모델은 단백질을 안정적으로 유지하면서 실제 실험 데이터와 일치하는 형태로 접히도록(folding) 도왔으며, 이는 복잡한 생물학적 구조를 다룰 수 있음을 증명합니다.

요약

E2Former-V2는 모든 대화를 거대한 화이트보드에 적는 느리고 답답한 시스템에서, 스마트한 셰프가 대화를 하나씩 즉석에서 요리하고 서빙하는 시스템으로 업그레이드한 것과 같습니다. 이를 통해 과학자들은 표준 컴퓨터를 사용하여 메모리 부족 문제 없이 거대한 단백질과 같은 거대 분자 시스템을 시뮬레이션하면서도 물리적 정확성을 유지할 수 있습니다.

기술 요약: E2Former-V2: 선형 활성화 메모리를 이용한 온더플라이(On-the-Fly) 등변량 어텐션

1. 문제 정의

등변 그래프 신경망(Equivariant Graph Neural Networks, EGNNs)은 회전 및 병진 대칭성을 보존하는 능력 덕분에 3D 원자계 모델링의 표준이 되었습니다. 그러나 주류 아키텍처(예: MACE, EquiformerV2, eSCN)는 에지 중심(edge-centric) 설계에 기인한 결정적인 효율성 병목 현상을 겪고 있습니다. 이러한 모델들은 모든 에지에 대해 기하학적 특징을 명시적으로 구축하거나 조밀한 텐서 곱을 수행하며, 이는 다음과 같은 결과를 초래합니다:

메모리 복잡도: $O(|E|)$ 의 활성화 메모리 복잡도를 가집며, 여기서 $|E| \approx kN$ 입니다 ( $k$ 는 평균 이웃 수). 대규모 시스템의 경우, 이는 GPU 용량을 초과하는 거대한 메모리 풋프린트를 생성합니다.
하드웨어 비효율성: 에지 텐서를 명시적으로 실체화(materialization)하는 과정은 고대역폭 메모리(HBM)와 온칩 SRAM 사이의 데이터 이동을 강제하며, 이는 처리량을 제한하는 메모리 벽(memory wall)을 형성합니다.
확장성 한계: 시스템 크기( $N$ )가 증가함에 따라 지연 시간(latency)이 불량하게 스케일링되어, 일반적인 하드웨어에서 대규모 분자 역학(MD) 시뮬레이션을 수행하는 것을 계산적으로 불가능하게 만듭니다.

최근 트랜스포머(예: FlashAttention)의 발전은 스트리밍 및 타일 기반 실행을 통해 시퀀스 데이터의 유사한 병목 현상을 해결했지만, 이러한 원칙은 SO(3) 등변 아키텍처에서는 아직 탐구되지 않았습니다.

2. 방법론

E2Former-V2는 **대수적 희소화(algebraic sparsification)**와 하드웨어 인지적 노드 중심 실행 모델을 통합하여 이러한 문제를 해결합니다. 이 아키텍처는 두 가지 핵심 혁신으로 구성됩니다.

2.1. 등변 축 정렬 희소화 (Equivariant Axis-Aligned Sparsification, EAAS)

표준 SO(3) 텐서 수축(tensor contractions)은 클레브슈-고르단(Clebsch–Gordan) 계수에 대한 조밀한 합산을 포함하여 $O(L^6)$ 의 복잡도와 불규칙한 메모리 액세스 패턴을 유발합니다. E2Former-V2는 EAAS를 도입하여 이러한 조밀한 연산을 희소한 순열 기반 연산으로 변환합니다:

SO(3) $\to$ SO(2) 기저 변환: 이 방법은 전역 좌표계를 두 원자를 연결하는 로컬 벡터 $\vec{r}$ (축)에 정렬합니다. 이 정렬된 프레임 내에서 구면 조화 함수(solid spherical harmonics)는 최대로 희소해지며(Lemma 4.1 참조), 오직 $m=0$ 성분만이 0이 아닌 값을 가집니다.
재인덱싱 연산자(Re-indexing Operator): 회전을 텐서 곱을 통해 교환함으로써, 조밀한 수축은 결정론적인 재인덱싱 규칙(Eq. 13)으로 붕괴됩니다. 이 연산자는 입력 특징을 희소 블록 단위 선형 맵을 통해 출력 특징으로 매핑하여, 자기 지수(magnetic indices)에 대한 조밀한 합산의 필요성을 제거합니다.
결과: 이를 통해 컨볼루션 단계의 복잡도를 줄이고 표준 구현 대비 약 $1.5\times$ 의 속도 향상을 달성합니다.

2.2. 온더플라이(On-the-Fly) 등변량 어텐션

$O(|E|)$ 메모리 병목 현상을 제거하기 위해, 저자들은 커스텀 Triton 커널로 구현된 완전 퓨즈드(fully fused) 노드 중심 어텐션 메커니즘을 제안합니다:

노드 중심 인수 분해: 에지 특징을 실체화하는 대신, 어텐션 메커니즘은 계산을 분리합니다. 쿼리(Query)와 키(Key) 벡터는 SO(3)-불변 스칼라 투영(Eq. 6)을 사용하여 구축되는 반면, 등변 기하학 정보는 오직 밸류(Value) 경로를 통해서만 전파됩니다.
스트리밍 리덕션(Streaming Reduction): 커널은 이웃들을 스트리밍 방식으로 처리합니다. 이웃 $j \in \mathcal{N}(i)$ 에 대해 반복하며, 어텐션 점수를 계산하고 온칩 SRAM 내에서 온더플라이 방식으로 밸류 업데이트를 누적합니다.
에지 실체화 제거: $N \times K \times H \times D$ 형상의 중간 텐서가 HBM에 기록되지 않습니다. 커널은 온칩에서 실행 중인 통계치(max, 정규화, 밸류 누산기)를 유지하여, 활성화 메모리 복잡도를 $O(|E|)$ 에서 $O(|V|)$ 로 감소시킵니다.
온라인 소프트맥스(Online Softmax): 스트리밍 중 수치적 안정성을 보장하기 위해, 온라인 소프트맥스 공식을 사용하여 누산기를 점진적으로 업데이트합니다.

3. 주요 기여

등변 축 정렬 희소화 (EAAS): Wigner-6j 재결합(recoupling)과 SO(2) 축 정렬을 통합하여, 조밀한 SO(3) 텐서 수축을 희소한 재인덱싱 연산으로 변환하는 새로운 대수적 기법입니다.
온더플라이 등변량 어텐션: 에지 실체화를 피하고 SRAM 재사용을 극대화하여 표준 구현 대비 최대 20 $\times$ 높은 TFLOPS를 달성하는 완전 퓨즈드 Triton 커널입니다.
확장 가능한 성능: 단일 GPU에서 최대 100,000개의 원자를 처리할 수 있는 능력을 입증하였으며, 이는 이전의 등변 트랜스포머들이 OOM(Out-of-Memory) 오류로 인해 실패했던 영역입니다.
물리적 충실도: 액체 물 구조의 정확한 재현 및 AlphaFold 예측 단백질 구조(약 30k 원자)의 안정적인 정밀화(refinement)를 통해 분자 역학에서의 유효성을 검증했습니다.

4. 실험 결과

저자들은 SPICE(의약 화학) 및 OMol25(파운데이션 모델 규모) 데이터셋과 MD 시뮬레이션을 통해 E2Former-V2를 평가했습니다.

예측 정확도:
- SPICE에서 E2Former-V2는 대부분의 서브셋에서 가장 낮은 에너지 및 힘 MAE(Mean Absolute Error)를 달eric했으며, Dimers에 대한 Energy MAE에서 MACE-Large를 48% 상회했습니다.
- OMol25에서 Conservative 변형은 1.27 meV/atom의 종합 Energy MAE를 달성하여 eSEN-small과 대등한 성능을 보였으며, UMA-S(3.62 meV/atom)를 크게 앞질렀습니다.
- OC20 S2EF-2M에서 EquiformerV2 및 eSCN과 같은 강력한 베이스라인과 대등한 정확도를 달성했습니다.
효율성 및 확장성:
- 처리량(Throughput): 퓨즈드 커널은 나이브한 PyTorch 구현 대비 평균 20 $\times$ 빠른 TFLOPS를 달성했습니다.
- 메모리: 피크 메모리 사용량이 베이스라인 대비 16.9 $\times$ 에서 19.5 $\times$ 감소했습니다.
- 확장성: 이전 모델들(예: EquiformerV2, MACE-Large)이 10k~50k 원자 이상에서 OOM 실패를 겪은 반면, E2Former-V2는 100k 원자까지 안정적으로 확장되었습니다.
- 추론 속도: 100k 원자에서 E2Former-V2 (Conservative)는 0.12 steps/s에 도달하여 UMA-S보다 거의 3 $\times$ 빨랐습니다. Direct 설정에서는 작은 규모에서 Allegro 및 EquiformerV2보다 한 자릿수 더 빨랐습니다.
분자 역학 (Molecular Dynamics):
- 주기적 벌크 물(216개 분자) 시뮬레이션에서 E2Former-V2는 실험적 방사 분포 함수(RDF)와 밀접하게 일치함을 보여주었으며, MACE-OFF보다 우수한 성능을 보였습니다.
- 용매화된 단백질(약 30k 원자)에 대한 2 ns 시뮬레이션은 네이티브 구조에 대한 높은 TM-score를 유지하며 안정적인 역학을 입증했습니다.

5. 의의 및 주장

본 논문은 E2Former-V2가 범용 GPU에서 대규모 등변 트랜스포머의 실질적인 실행 가능성을 입증한다고 주장합니다. 실행 패러다임을 에지 중심에서 노드 중심으로 전환하고 하드웨어 인지적 스트리밍을 활용함으로써, 이 방법은 등변 모델의 확장을 제한해 온 근본적인 메모리 병목 현상을 해결했습니다.

저자들은 이것이 단순한 알고리즘적 개선이 아니라, 이론적 인수 분해(복잡도를 $O(|E|)$ 에서 $O(|V|)$ 로 감소)와 하드웨어 효율성 사이의 간극을 메우는 시스템 수준의 실현임을 강조합니다. 이 연구는 기존에는 계산적으로 불가능했던 대규모 분자 파운데이션 모델 및 장시간 MD 시뮬레이션이 이제 표준 GPU 인프라를 사용하여 접근 가능하다는 것을 시사합니다. 코드는 커뮤니티의 추가적인 채택을 돕기 위해 공개되었습니다.