Learning Inter-Atomic Potentials without Explicit Equivariance

1. 문제 정의 (Problem)

분자 시뮬레이션 (신약 개발, 신소재 설계 등) 에 있어 정확한 원자 간 퍼텐셜 (Inter-Atomic Potentials, IP) 을 학습하는 것은 필수적입니다. 현재 가장 정밀한 방법인 밀도 범함수 이론 (DFT) 은 계산 비용이 시스템 크기의 세제곱에 비례하여 증가하므로 대규모 또는 장기간 시뮬레이션에 적용하기 어렵습니다. 이를 해결하기 위해 머신러닝 기반 원자 간 퍼텐셜 (MLIP) 이 개발되었으나, 기존 최첨단 모델들은 회전 - 병진 대칭성 (Rotational-Translational Symmetry, SE(3) equivariance) 을 보장하기 위해 특수한 공변적 (Equivariant) 신경망 아키텍처 (예: 구면 조화 함수, 공변적 메시지 전달) 를 사용합니다.

이러한 명시적 공변성 설계는 다음과 같은 한계가 있습니다:

• 유연성 저하: 아키텍처가 복잡해지고 표현력이 제한될 수 있음.
• 계산 효율성 및 확장성 문제: 구면 조화 함수 등 계산 집약적인 메커니즘으로 인해 대규모 모델이나 데이터셋에서의 확장성이 떨어짐.
• 데이터 활용도: 충분한 데이터가 존재할 때에도 복잡한 인덕티브 바이어스 (Inductive Bias) 를 강제로 부과하는 것이 항상 최선은 아님.

2. 방법론 (Methodology)

저자들은 TransIP (Transformer-based Inter-Atomic Potentials) 라는 새로운 학습 패러다임을 제안합니다. 이는 아키텍처에 명시적인 공변성 제약 (SO(3) 제약) 을 가하지 않고, 잠재 공간 (Embedding Space) 에서 공변성을 학습하는 방식입니다.

핵심 구성 요소:

1. 제약 없는 Transformer 백본 (Unconstrained Transformer Backbone):

   • 분자를 가변 길이의 토큰 시퀀스 (각 원자가 하나의 토큰) 로 모델링합니다.
   • 기존 GNN 과 달리 거리 컷오프에 의존하지 않고, 자기 주의 (Self-Attention) 메커니즘을 사용하여 분자 내 모든 원자를 처리합니다.
   • 전하 (Charge) 와 스핀 (Spin) 과 같은 전역 분자 속성을 학습 가능한 임베딩을 통해 각 레이어에 바이어스로 주입합니다.

2. 학습된 변환 네트워크 (Learned Transformation Network, $T_\tau$):

   • 입력 도메인에서의 군 작용 (예: 회전 $g \in SO(3)$) 이 잠재 공간의 임베딩 벡터에 어떻게 작용하는지 학습하는 MLP 를 도입합니다.
   • 수식: $T_\tau(\phi(g), f(m)) = \text{MLP}_\tau([\phi(g), f(m)])$
   • 이는 명시적인 $\rho(g)$ (출력 공간의 군 표현) 를 정의하지 않고, 신경망을 통해 군 작용을 암시적으로 학습하게 합니다.

3. 잠재 공변성 대비 손실 함수 (Contrastive Objective for Latent Equivariance):

   • 모델이 회전된 분자 입력 $f(\phi(g)(m))$ 을 처리했을 때, 변환 네트워크 $T_\tau$ 를 통해 원래 임베딩 $f(m)$ 을 회전시킨 결과와 일치하도록 유도합니다.
   • 손실 함수 ($L_{leq}$): $\| f(\phi(g)(m)) - T_\tau(\phi(g), f(m)) \|^2$
   • 이 손실 함수는 데이터 증강 (Data Augmentation) 없이도 모델이 회전 불변성/공변성을 학습하도록 가이드합니다.

4. 종합 학습 목적 함수:

   • 에너지 예측 오차 ($L_E$), 힘 예측 오차 ($L_F$), 잠재 공변성 오차 ($L_{leq}$) 를 가중치 합으로 결합하여 단일 스테이지로 학습합니다.

3. 주요 기여 (Key Contributions)

• 단일 스테이지 학습 파이프라인: 사전 학습 - 미세 조정 (Pretraining-Finetuning) 프레임워크나 하드코딩된 공변적 레이어 없이, 일반 Transformer 를 사용하여 학습을 통해 SO(3) 공변성을 달성합니다.
• 아키텍처 무관 대비 손실 (Architecture-Agnostic Contrastive Loss): 제약 없는 모델의 잠재 공간에서 SO(3) 공변성을 촉진하는 새로운 손실 함수를 제안했습니다. 이는 기존 데이터 증강 기법보다 다양한 데이터셋과 모델 크기에서 더 우수한 확장성을 보입니다.
• 성능 입증: Open Molecules 25 (OMol25) 데이터셋 (소분자, 생체 분자 조각, 전해질 등 포함) 에서 TransIP 가 최신 공변적 베이스라인 (eSCN, GemNet-OC 등) 과 비교해 경쟁력 있는 성능을 보이며, 데이터 증강 기반 모델 대비 40~60% 의 성능 향상을 달성했습니다.

4. 실험 결과 (Results)

OMol25 데이터셋 (약 400 만 개의 분자) 을 기반으로 한 실험 결과는 다음과 같습니다.

• 데이터 크기 확장 (Data Scaling):
  • 데이터가 제한된 환경 (1M 개 분자) 에서 TransIP 는 데이터 증강을 사용한 TransAug 모델보다 힘 (Force) 예측 오차 (MAE) 에서 255 meV/Å 대 600 meV/Å 로 압도적인 우위를 보였습니다.
  • 데이터가 증가할수록 (4M 개) 두 모델 모두 성능이 향상되었으나, TransIP 는 여전히 힘 예측 정확도와 방향성 일치도 (Cosine Similarity) 에서 우위를 유지했습니다.
• 모델 크기 확장 (Model Scaling):
  • TransIP 는 모델 파라미터 수가 증가함에 따라 성능이 원활하게 향상되는 반면, TransAug 은 모델이 커질수록 오히려 성능이 저하되는 현상을 보였습니다. 이는 학습된 잠재 공변성이 대규모 모델에 더 효과적인 인덕티브 바이어스를 제공함을 시사합니다.
• 최첨단 모델 대비 성능:
  • 80 에포크 학습 후 TransIP-Large 모델은 eSCN (Small/Medium) 과 비교해 총 에너지 예측에서 경쟁력 있는 성능을 보였고, 전체 힘 예측 오차 (Force MAE) 에서도 유사한 수준을 달성했습니다.
• 추론 속도:
  • TransIP 는 공변적 모델 (eSEN) 보다 훨씬 빠른 추론 속도를 보입니다 (예: TransIP-Small 은 초당 16 만 개 원자 처리, eSEN 은 1 만 5 천 개).
• 학습된 공변성 분석:
  • 잠재 공간에서 회전된 입력에 대한 임베딩 변화가 단일 직교 변환 (Orthogonal Transformation) 으로 잘 설명됨을 확인했습니다. 이는 명시적 제약이 없어도 모델이 입력 회전과 잠재 공간의 직교 변환 간의 관계를 효과적으로 학습했음을 의미합니다.

5. 의의 및 결론 (Significance)

이 연구는 일반적인 Transformer 아키텍처가 명시적인 기하학적 제약 없이도 데이터로부터 공변성을 효과적으로 학습할 수 있음을 입증했습니다.

• 확장성: 복잡한 공변적 연산 (구면 조화 함수 등) 을 제거함으로써 하드웨어 효율성과 대규모 모델 확장성을 극대화했습니다.
• 효율성: 데이터 증강보다 학습된 잠재 공변성 (Learned Latent Equivariance) 이 더 강력한 인덕티브 바이어스로 작용하여, 특히 데이터가 부족하거나 모델이 클 때 더 빠른 수렴과 높은 정확도를 제공합니다.
• 미래 전망: 분자 역학 시뮬레이션 분야에서 더 크고 효율적인 모델을 구축할 수 있는 새로운 방향을 제시하며, "학습된 공변성"이 명시적 공변성이나 단순 데이터 증강의 강력한 대안이 될 수 있음을 보여줍니다.

결론적으로, TransIP 는 복잡한 기하학적 대칭성을 하드코딩하는 대신, 데이터와 손실 함수를 통해 이를 학습함으로써 정확성, 확장성, 계산 효율성을 모두 잡은 새로운 MLIP 패러다임을 제시합니다.