[논문 리뷰] Scalable Spatio-Temporal SE(3) Diffusion for Long-Horizon Protein Dynamics

Nima Shoghi, Yuxuan Liu외 4명|arXiv (Cornell University)|2026. 02. 02.

Protein Structure and Dynamics인용 수 0

한 줄 요약

STAR-MD는 SE(3)-등가적인 자기회귀 확산 모델로, 결합된 시공간 주의를 통해 마이크로초 시간 척도에 걸친 물리적으로 그럴듯한 단백질 궤적을 생성하며, 이전 모델들보다 긴 시간대의 역학에서 우수합니다. ATLAS에서 안정적이고 고품질의 롤아웃을 달성하고 긴 시점으로의 일반화도 가능하다는 것을 보여줍니다.

ABSTRACT

Molecular dynamics (MD) simulations remain the gold standard for studying protein dynamics, but their computational cost limits access to biologically relevant timescales. Recent generative models have shown promise in accelerating simulations, yet they struggle with long-horizon generation due to architectural constraints, error accumulation, and inadequate modeling of spatio-temporal dynamics. We present STAR-MD (Spatio-Temporal Autoregressive Rollout for Molecular Dynamics), a scalable SE(3)-equivariant diffusion model that generates physically plausible protein trajectories over microsecond timescales. Our key innovation is a causal diffusion transformer with joint spatio-temporal attention that efficiently captures complex space-time dependencies while avoiding the memory bottlenecks of existing methods. On the standard ATLAS benchmark, STAR-MD achieves state-of-the-art performance across all metrics--substantially improving conformational coverage, structural validity, and dynamic fidelity compared to previous methods. STAR-MD successfully extrapolates to generate stable microsecond-scale trajectories where baseline methods fail catastrophically, maintaining high structural quality throughout the extended rollout. Our comprehensive evaluation reveals severe limitations in current models for long-horizon generation, while demonstrating that STAR-MD's joint spatio-temporal modeling enables robust dynamics simulation at biologically relevant timescales, paving the way for accelerated exploration of protein function.

연구 동기 및 목표

• 나노초를 넘어서는 장기적 단백질 궤적 생성의 필요성을 촉진하고 기존 모델의 한계를 다룬다.
• 비분리된 시공간 역학을 포착하기 위해 결합된 시공간 주의를 갖춘 확장 가능한 자기회귀 SE(3) 확산 모델을 개발한다.
• 블록-인과관계 주의(block-causal attention) 및 contextual noise perturbation과 같은 학습 기법을 통해 안정적인 장기 롤아웃을 달성한다.
• ATLAS 100 ns 궤적에서의 최첨단 성능과 240 ns 및 1 μs 시점으로의 견고한 외삽을 입증한다.
• 비마르코프 유관 기억 효과를 위한 Mori-Zwanzig 형식을 통한 이론적 정당화를 제공한다.

제안 방법

• 잔차당 이동과 회전을 나타내는 SE(3) 위에 자기회귀 확산 프레임워크를 사용한다.
• 비분리 의존성을 포착하기 위해 잔차 프레임 토큰에 대해 결합된 시공간 주의를 적용한다.
• long-horizon 생성을 가능하게 하는 KV 캐싱을 통한 블록-확산 스타일의 인과적 학습을 사용한다.
• 롤아웃의 오차 누적을 완화하기 위해 학습 및 추론 시 맥락적 노이즈 섭동(contextual noise perturbation)을 적용한다.
• 다중 시간 척도를 다루기 위한 적응형 레이어 노멀라이제이션으로 연속 시간 조건화를 수행한다.
• 역사 의존성 및 비분리 기억 커널의 필요성을 정당화하기 위한 Mori-Zwanzig 형식의 이론적 근거를 제시한다.

실험 결과

연구 질문

• RQ1 STAR-MD가 대형 단백질의 마이크로초 시간 척도에서 물리적으로 그럴듯한 궤적을 생성할 수 있는가?
• RQ2 결합된 시공간 주의가 공간-후 시간 구조보다 장단 결합 시공간 역동을 더 잘 포착하는가?
• RQ3 100 ns 데이터로만 학습한 후 STAR-MD가 240 ns 및 1 μs와 같은 더 긴 시점으로 얼마나 잘 외삽하는가?
• RQ4 장기 롤아웃의 안정성과 구조적 타당성을 크게 향상시키는 아키텍처 및 학습 선택은 무엇인가?
• RQ5 STAR-MD가 구성적 커버리지, 타당성 및 역학적 충실도 측면에서 최첨단 궤적 모델과 어떻게 비교되는가?

주요 결과

• STAR-MD는 구성적 커버리지, 구조적 타당성 및 역동적 충실도 측면에서 ATLAS 100 ns 벤치마크에서 최첨단 성능을 달성한다.
• STAR-MD는 1 μs까지의 장기 롤아웃 동안 높은 구조적 타당성과 넓은 구성적 커버리지를 유지하며, 빠르게 성능이 저하되는 베이스라인보다 우수하다.
• 결합된 시공간 주의와 맥락적 노이즈를 포함한 자기회귀 확산은 견고한 롤아웃과 오차 누적 감소를 가져온다.
• STAR-MD는 240 ns 및 1 μs 궤적으로 외삽되며 경쟁력 있는 품질 지표를 보이는 반면, 대안은 더 빨리 악화된다.
• 비마르코프 coarse-grained 다이나믹스와 기억 효과의 필요성을 역사 의존 커널과의 연결로 이론적으로 제시하여 모델 설계를 정당화한다.