Transformers Discover Molecular Structure Without Graph Priors

기술 요약: 물리적 사전 지식 없이 분자 구조를 발견하는 트랜스포머(Transformers)

문제 정의
컴퓨팅 모델링은 물리적 시스템을 연구하는 데 필수적이지만, 전통적인 ab initio(제1원리) 방법론은 대규모 시스템에 적용하기에는 비용이 지나치게 많이 드는 경우가 많습니다. 머신러닝(ML)은 비용 효율적인 대안으로 부상했으나, 물리 과학 분야에서의 적용은 일반적으로 모델 아키텍처에 강력한 물리적 귀납적 편향(예: 기하학적 국소성, 그래프 구조, 특정 대칭성)을 직접 내재화하는 방식에 의존해 왔습니다. 이는 근본적인 질문을 제기합니다. 즉, 물리 법칙을 명시적으로 인코딩해야 하는 범위는 어디까지이며, 데이터만으로 자율적으로 발견할 수 있는 범위는 어디까지인가 하는 점입니다. 트랜스포머와 같은 도메인 불가지론적(domain-agnostic) 아키텍처가 언어와 시각 분야를 혁신했음에도 불구하고, 원자 모델링에서의 채택은 이들이 명시적인 기하학적 가이드 없이는 복잡한 물리적 상호작용을 회복할 수 없다는 가정(특히 역사적으로 작은 과학 데이터셋을 고려할 때)으로 인해 제한되어 왔습니다.

방법론
저자들은 상호작용 그래프나 수동으로 정의된 쌍별 포텐셜(pairwise potentials)과 같은 도메인 특화된 사전 지식 없이, 오직 원자론적 데이터만을 학습했을 때 표준적인 수정 없는 트랜스포머 아키텍처가 근본적인 물리적 구조를 회복할 수 있는지 조사합니다.

• 모델 아키텍처: 본 연구는 가공되지 않은 데카르트 좌표, 원자 유형, 그리고 글로벌 토큰(전하, 스핀)을 입력으로 받아 분자 에너지와 원자당 힘을 예측하는 그래프 프리(graph-free) 트랜스포머를 활용합니다. 핵심 메커니즘은 셀프 어텐션(self-attention) 연산이며, 여기서 어텐션 점수($XW_Q(XW_K)^T$)는 학습된 원자 간 상호작용 강도의 대리 지표 역할을 합니다.
• 데이터셋: 주요 데이터셋은 OMol25(Open Molecules 2025)로, $\omega$B97M-V/def2-TZVPD 수준의 이론으로 계산된 에너지 및 힘 라벨을 포함하며 다양한 화학 구조(생체 분자, 금속 착물, 전해질)를 포괄하는 방대한 데이터셋입니다.
• 학습 파이프라인: 모델은 타겟의 이산화(discretized) 버전에 대한 사전 학습(pre-training)과 연속값에 대한 미세 조정(fine-tuning)을 따르는 2단계 학습 과정을 채택합니다.
• 분석 기법:
  • 어텐션 분석: 저자들은 어텐션 점수를 원자 간 거리의 함수로 분석하여, 모델이 물리 법칙(예: 정전기력)과 일치하는 감쇠 패턴(decay patterns)을 학습하는지 확인합니다.
  • 적응형 컷오프(Adaptive Cutoffs): 각 원자의 전체 어텐션 중 90%를 포착하는 데 필요한 거리를 기준으로 학습된 "유효 상호작용 반경"($R_i$)을 정의하고, 이를 국소적 희소성(local sparsity) 및 물리적 프록시(밀도 행렬 및 정전기 단극자 프록시)와 비교합니다.
  • 스케일링 법칙(Scaling Laws): 모델의 크기를 다양하게(5M ~ 306M 파라미터, 사전 학습 시 최대 1B까지) 조절하고 에포크(epoch) 수를 달리하여 학습함으로써, 성능이 컴퓨팅 및 데이터와 함께 예측 가능한 거듭제곱 법칙(power-law) 관계를 따르는지 검증하는 뉴럴 스케일링 분석을 수행합니다.
  • 합성 실험: 학습 메커니즘을 격리하기 위해, 에너지 라벨이 단순한 쌍별 포텐셜($E(r) \propto 1/r^\alpha$)로 정의된 합성 데이터셋을 생성하여 모델이 특정 지수 $\alpha$를 회복할 수 있는지 테스트합니다.

주요 결과

1. 물리적 구조의 자율적 회복: 트랜스포머는 고전 물리학을 모사하는 상호작용 강도를 자율적으로 학습합니다. 초기 레이어에서 어텐션은 거리($\sim 1/r^2$)에 따라 급격히 감소하며, 이는 정전기력과 일치합니다. 더 깊은 레이어에서는 감쇠가 더 완만해지며($\sim 1/r$), 이는 정전기 에너지 스케일링과 일치합니다.
2. 적응형 상호작용 패턴: 고정된 반경 컷오프를 사용하는 그래프 신경망(GNN)과 달리, 트랜스포머는 적응형 유효 상호작용 반경을 학습합니다. 희소한 환경에서 모델은 장거리 상호작용을 포착하기 위해 수용 영역(receptive field)을 확장하며, 밀도가 높은 영역에서는 국소적 상호작용에 집중합니다. 이렇게 학습된 컷오프는 물리 기반 모델에서 수동으로 선택되는 값들과 일치합니다.
3. 예측 가능한 스케일링: 모델은 예측 가능한 뉴럴 스케일링 법칙 행동을 보입니다. 성능은 데이터 및 컴퓨팅 증가에 따라 일관되게 향상되며, 다른 ML 분야에서 관찰된 것과 유사한 거듭제곱 법칙 관계를 따릅니다. 저자들은 더 작은 모델들로부터 도출된 스케일링 법칙을 사용하여 306M 파라미터 모델의 성능을 성공적으로 예측했습니다.
4. 합성 데이터 회복: 알려진 $1/r^\alpha$ 포텐셜을 가진 합성 데이터셋에서, 모델이 학습한 어텐션 지수는 실제 값과 강력한 상관관계($R^2 = 0.98$)를 보입니다. 데이터와 모델 크기가 증가함에 따라 학습된 지수는 실제 물리적 지수에 점점 더 근접합니다.
5. 화학적 및 분자 간 이해: 모델은 주기율표를 반영하는 의미 있는 원소 임베딩을 학습합니다(예: 비활성 기체의 클러스터링, 알칼리 금속의 순서화). 또한, 물 이량체(water dimer)의 분자 간 상호작용 에너지 곡선을 정확하게 재현하여 DFT 및 실험 데이터와 일치함을 보여주었으며, 고정 컷오프 모델에서 흔히 발생하는 급격한 절단 아티팩트(truncation artifacts)를 피합니다.
6. 경쟁적인 정확도: 동일한 컴퓨팅 예산 하에 물리 기반 GNN과 비교 학습된 1B 파라미터 트랜스포머는 에너지 및 힘 예측 작업에서, 특히 분포 외(out-of-distribution) 분자 조성에 대해서도 경쟁력 있는 정확도를 달달성했습니다.

의의 및 주장
본 논문은 범용 아키텍처, 특히 트랜스포머가 과학적 모델링을 위한 원칙적인 베이스라인 역할을 할 수 있다고 주장합니다. 결과는 많은 물리적 귀납적 편향들이—전통적으로 모델 아키텍처에 하드코딩되어야 한다고 여겨졌던—충분한 데이터와 컴퓨팅 자원이 확보될 경우 데이터로부터 직접 학습 가능하다는 것을 시사합니다.

저자들은 이 연구가 설계된 귀납적 편향과 학습 가능한 물리적 구조 사이의 경계를 명확히 한다고 주장합니다. 그들은 명시적인 사전 지식(priors)이 기본 요구 사항이 아니라, 경험적으로 필요할 때만 도입되어야 한다고 제안합니다. 도메인 불가지론적 모델이 복잡한 물리적 패턴(적응형 컷오프 및 정전기 스케일링 등)을 회복하고 예측 가능하게 스케일링될 수 있음을 입증함으로써, 본 연구는 더 넓은 딥러닝 문헌의 통찰력을 원자론적 모델링 커뮤니티로 전달하여 더욱 유연하고 일반화 가능한 과학적 모델을 가능하게 할 수 있음을 뒷받침합니다.