원자 수준 생체분자 동역학을 위한 상태공간 모델 기반 ATMOS 프레임워크

본 논문은 원자 수준의 생체분자 동역학을 효율적으로 생성하기 위한 새로운 프레임워크 ATMOS(Atomic Trajectory Modeling with State Space Models)를 제안한다. 기존 MD 시뮬레이션은 물리적 정확도가 높지만 시간·자원 소모가 커 장기 시뮬레이션이 어려웠으며, 최근 딥러닝 기반 생성 모델은 정적 구조나 제한된 시간 범위만을 다루는 경우가 많았다. 이러한 격차를 메우기 위해 저자들은 상태공간 모델(SSM)을 핵심으로 채택하고, 이를 Pairformer 기반 전이 모듈과 확산 디코더와 결합한다. 먼저 문제 정의에서는 원자 좌표 시퀀스 X = {x₁,…,x_T}와 정적 화학적 컨텍스트 a를 조건으로 하는 확률 모델 p_θ(X|x₁,a)를 설정하고, 이를 autoregressive하게 factorize한다. SSM은 입력 신호 u(t)와 잠재 상태 h(t) 사이의 선형 ODE를 기반으로 하며, 이산화 후 h_t = A h_{t-1} + B u_t, y_t = C h_t 형태로 구현된다. 이 구조는 매 타임스텝 O(1) 연산으로 긴 시퀀스를 처리할 수 있어, 전통적인 어텐션 모델이 갖는 O(T²) 복잡도를 회피한다. ATMOS의 아키텍처는 크게 네 단계로 구성된다. (1) Context & Input Encoding: 화학적 속성 a를 임베딩해 초기 잠재 상태 h₀=(s₀,z₀)를 만든 뒤, 현재 좌표 x_t와 a를 기하학적 인코더 E_θ에 입력해 특징 v_t=(v_s_t, v_z_t)를 추출한다. (2) State Transition: Pairformer 변형 Φ_θ를 사용해 s_t와 z_t에 v_t와 시간 스텝 임베딩 τ_θ(Δt)를 더해 다음 잠재 상태 (s_{t+1}, z_{t+1})를 계산한다. Pairformer는 single-track와 pair-track 사이의 양방향 정보 흐름을 지원해 원자 간 장거리 상호작용을 효율적으로 모델링한다. (3) Decoding: 확산 기반 디코더 D_θ가 노이즈가 섞인 구조 ˜x(γ)와 진화된 잠재 상태 h_{t+1}을 조건으로 받아, 역확산 과정을 통해 깨끗한 좌표 ˆx_{t+1}를 복원한다. (4) Learning Objective: 실제 MD 궤적에 대해 로그우도 손실 L_SSM = -∑_{t} log p_θ(x_{t+1}|h_t, x_t, a) 를 최소화한다. 데이터 측면에서는 PDB의 정적 구조와 대규모 MD 데이터셋 mdCATH, MISATO를 활용한다. mdCATH은 단백질 단일체의 수천 나노초 시뮬레이션을, MISATO는 단백질‑리간드 복합체의 장기 동역학을 제공한다. 이러한 데이터는 ATMOS가 다양한 시스템에 일반화될 수 있음을 보장한다. 실험에서는 RMSD, 프레임 간 거리 분포, 에너지 보존, 그리고 프리시전-리콜 호출 등 여러 지표를 사용해 기존 모델(AlphaFlow, MDGen, TEMPO, ConfRover 등)과 비교한다. 결과는 ATMOS가 특히 장기 궤적(수천 프레임)에서 선형 시간 복잡도로 안정적인 샘플링을 수행하면서도, 물리적 에너지와 구조적 정확도 면에서 최고 수준의 성능을 보였음을 보여준다. 또한, 복합체 시스템에서 리간드 해리·결합 경로를 자연스럽게 재현해, 기존 모델이 놓치던 동역학적 세부 정보를 포착한다. 논문의 주요 기여는 다음과 같다. 1) SSM을 생체분자 궤적 생성에 적용해 물리적 마코프성을 유지하면서 연산 효율성을 확보했다. 2) Pairformer 기반 전이 함수를 도입해 원자 간 장거리 상호작용을 고차원 잠재 공간에서 효과적으로 학습했다. 3) 전체 원자 좌표를 직접 다루면서도 토큰화된 잠재 표현으로 메모리 요구량을 크게 줄였다. 4) 정적 구조와 대규모 MD 데이터를 통합 학습해 단일체와 복합체 모두에 대해 일반화 가능한 모델을 구축했다. 5) 확산 디코더와 결합해 물리적 유효성(에너지, 충돌 방지)을 보장하면서도 샘플링 속도를 크게 향상시켰다. 향후 연구 방향으로는 (a) 온도·압력 등 외부 조건을 명시적으로 모델링해 다양한 환경에서의 동역학을 예측, (b) 실험적 NMR·Cryo‑EM 데이터와의 혼합 학습을 통해 모델의 물리적 정확성을 더욱 강화, (c) 다중체 시스템(예: 단백질‑단백질 인터페이스) 및 대규모 복합체에 대한 스케일링, (d) 생성된 궤적을 약물 설계 파이프라인에 직접 연결해 리간드 결합 경로와 자유 에너지 프로파일을 빠르게 추정하는 응용을 제시한다. 전반적으로 ATMOS는 생체분자 동역학 모델링에 새로운 패러다임을 제시하며, 고성능 시뮬레이션과 딥러닝 기반 생성 모델의 장점을 융합한 중요한 진전이다.