A recipe for scalable attention-based ML potentials: unlocking long-range accuracy with all-to-all node attention

Motivation

• Known: 기존 MLIP은 SchNet, GemNet, NequIP, MACE, eSEN, Equiformer 등 물리 기반 inductive bias(대칭성, locality, equivariance)를 강하게 반영하여 발전해왔으며, 장거리 상호작용(electrostatics, dispersion 등)을 다루기 위해 charge equilibration, Ewald/PME/FMM 등 명시적 물리항을 추가하는 것이 일반적이다.
• Gap: 모델과 데이터 규모가 커질수록(예: OMol25 같은 O(100 million) 샘플 데이터셋) 기존의 명시적 물리 기반 장거리 처리 방식이 대규모 이질적 데이터셋에 효과적으로 확장되는지는 검증되지 않은 열린 문제였다.
• Why: 생체분자, 전해질 등 장거리 상호작용이 중요한 대형 시스템에 대해 정확하고 안정적인 장시간 MD 시뮬레이션이 가능한 확장성 있는 MLIP을 제공하는 것은 실험적으로 검증 가능한 물성(밀도, 기화열 등) 예측을 위해 중요하다.
• Approach: 이 논문은 rotational symmetry, locality, long-range interaction과 같은 inductive bias들이 충분한 데이터와 모델 규모 하에서는 학습 가능하다는 가설 하에, 국소 neighbor attention과 전역 all-to-all node attention 두 단계로 구성된 prior-light 아키텍처(AllScAIP)를 제안하고 이를 다양한 규모에서 ablation한다.

Achievement

Figure 5. OMol25 energy error vs. efficiency throughput (inference). Energy / Atom MAE (meV, ↓) vs. throughput (ns/day,

1. AllScAIP 모델 제안: EScAIP 기반 neighborhood self-attention과 all-to-all node self-attention을 결합하여 energy-conserving하며 O(100 million) 샘플까지 확장 가능한 attention 기반 MLIP을 설계함.
2. Inductive bias의 규모 의존성 규명: 소규모 데이터/모델에서는 LAE, ERoPE 등 기하학적 inductive bias가 성능에 유용하지만, 데이터와 모델 규모가 커지면 그 이점이 감소하거나 역전되는 반면 all-to-all node attention은 장거리 정보 처리에 지속적으로 중요함을 실험적으로 입증함.
3. State-of-the-art 성능 달성: 분자 시스템(OMol25)에서 energy/force 정확도 및 다양한 물리 기반 평가에서 SOTA를 달성했고, 재료(OMat24) 및 촉매(OC20) 데이터셋에서도 경쟁력 있는 성능을 보임.
4. 실제 물리적 관측치 재현: OMol25로 학습된 AllScAIP이 안정적인 장시간 MD 시뮬레이션을 수행하여 밀도, 기화열 등 실험적 관측치를 정확히 재현함을 보임.

How

Figure 3. Illustration of the geometric encoding used in AllScAIP. (a) Legendre Angular Encoding (LAE) (b) Euclidean rot

• Neighborhood self-attention: EScAIP 모델을 기반으로 각 center atom이 최대 k개의 neighbor와 self token을 모아 (#nodes)×(k+1)×d_model 텐서를 구성하고 표준 multi-head self-attention(MHSA)을 적용하며, center→neighbor와 neighbor→center 양방향 pass를 수행하고 distance 기반 envelope로 padding/mask 처리.
• All-to-all node self-attention: 국소 정보 처리 이후, 노드 스트림 전체를 1×(#nodes)×d_model로 패킹하여 MHSA를 적용함으로써 그래프 전체에 걸친 정보 전파(global mixing)를 O(N^2) 비용으로 수행.
• Legendre Angular Encoding (LAE): neighborhood attention에 고차 방향 정보를 제공하는 compact rotation-aware edge encoding.
• Euclidean Rotary Position Encoding (ERoPE): node attention에 거리 정보를 주입하는 isotropic distance-only encoding (Frank et al. 2024 기반).
• 두 attention 단계 모두 컴퓨터 비전/언어 분야에서 고도로 최적화된 off-the-shelf multi-head self-attention CUDA kernel(xFormers 등)을 재사용하여 효율적으로 구현.
• 모델 크기와 데이터 규모(OMol25, OMat24, OC20)를 스케일링하며 LAE, ERoPE, node attention 각 요소에 대한 ablation 실험 수행.

Originality

• 물리 기반 명시적 long-range 항(전하 예측, Ewald/PME/FMM 등) 없이 순수 데이터 기반 all-to-all node attention만으로 장거리 상호작용을 포착하려는 접근은 기존 MLIP 연구와 차별화됨.
• inductive bias가 데이터/모델 규모에 따라 "학습 가능"해진다는 관점을 체계적으로 검증한 ablation 설계가 독창적이며, 이는 최근 대규모 foundation model 스케일링 논의를 MLIP 분야에 적용한 시도로 볼 수 있음.
• 기존에 CV/NLP에서 고도로 최적화된 attention CUDA kernel을 그대로 재사용하여 MLIP의 확장성을 확보한 엔지니어링적 실용성도 특징적임.

Limitation & Further Study

• all-to-all node attention은 O(N^2) 비용을 가지므로 매우 큰 시스템(수만~수십만 원자)에 대한 확장성에는 한계가 있을 수 있으며, 이에 대한 throughput/memory 분석(Fig 4)이 필요함.
• 명시적 물리항(전하, 다중극자 등) 없이 학습된 표현이 물리적 해석 가능성(interpretability) 측면에서 부족할 수 있으며, 극단적인 장거리 시스템(예: 매우 큰 전해질 시스템)에서의 일반화 성능은 추가 검증이 필요함.
• ablation 결과가 특정 데이터셋(OMol25, OMat24, OC20) 범위 내에서의 관찰이므로, 더 다양한 화학 공간이나 극단적 규모에서도 동일한 경향이 유지되는지 후속 연구가 필요함.

Evaluation

총평: 명시적 물리 기반 장거리 항 없이 all-to-all attention만으로 장거리 상호작용을 데이터 기반으로 학습할 수 있음을 대규모 실험으로 설득력 있게 보인 실용적이고 중요한 연구이며, MLIP 설계에서 scale에 따른 inductive bias 재고를 촉구하는 의미 있는 기여이다.