스케일러블 시공간 SE(3) 확산을 통한 마이크로초 단백질 동역학 생성

초록

STAR‑MD는 시공간을 동시에 고려하는 인과적 트랜스포머와 SE(3)‑등변 확산 모델을 결합해, 대형 단백질의 마이크로초 수준 장기 궤적을 물리적으로 타당하게 생성한다. 공동 시공간 어텐션과 히스토리 노이즈 주입 기법으로 메모리와 연산을 효율화하고, ATLAS 벤치마크에서 기존 방법을 크게 앞선 구조 다양성·동역학 충실도·구조 유효성을 달성한다.

상세 분석

STAR‑MD는 기존 단백질 동역학 생성 모델이 직면한 세 가지 핵심 한계를 동시에 해결한다. 첫째, 공간‑시간 정보를 별도로 처리하던 기존 아키텍처는 복합적인 상호작용을 포착하지 못한다. STAR‑MD는 잔기‑프레임 쌍을 토큰으로 하는 공동 시공간 (S × T) 어텐션을 도입해, 특정 잔기의 움직임이 과거 다른 부위의 움직임과 어떻게 결합되는지를 직접 모델링한다. 2D‑RoPE를 이용한 잔기·프레임 위치 임베딩은 훈련 시 프레임 수를 초과하는 시계열에도 자연스럽게 외삽할 수 있게 한다.

둘째, 전통적인 비디오‑형 확산 모델은 모든 프레임을 동시에 디노이징하거나 고정된 윈도우만을 사용해 메모리·연산 비용이 급증한다. STAR‑MD는 블록‑인과 인과적 트랜스포머 구조를 채택해, 현재 프레임을 디노이징하면서 과거 정제된 프레임만을 KV‑캐시로 유지한다. 이는 훈련 시에도 동일한 인과적 패턴을 유지하도록 블록‑와이즈 어텐션을 설계함으로써, 한 번의 포워드 패스로 전체 시퀀스를 병렬 처리하게 한다.

셋째, 장기 롤아웃 시 누적 오류가 급격히 증가한다는 문제를 해결하기 위해 ‘컨텍스추얼 노이즈 퍼터베이션’ 기법을 도입했다. 훈련 단계에서 과거 프레임에 작은 전방 확산 노이즈를 추가해 모델이 손상된 히스토리를 보정하도록 학습시키고, 추론 시에도 동일한 수준의 노이즈를 삽입해 훈련‑추론 분포 차이를 최소화한다. 이와 더불어 물리적 시간 간격 Δt를 로그균등으로 샘플링하고 AdaLN을 통해 조건화함으로써, 프레임 수와 실제 시간 간의 비선형 관계를 학습한다.

이러한 설계 선택은 복잡도 측면에서도 의미가 있다. 기존 Pairformer‑기반 모델은 O(N³L) 의 입체 연산과 O(N²L) 메모리를 요구하지만, STAR‑MD는 단일 잔기 피처만을 캐시해 O(NL) 메모리와 O(N²L²) 연산으로 충분히 확장 가능하다. 실험에서는 ATLAS 데이터셋(다양한 크기·시간 스케일)에서 100 ns부터 1 µs까지의 궤적을 생성했으며, RMSD, TM‑score, Fréchet Distance 등 다중 지표에서 기존 AlphaFold‑2‑style, MDGen, ConfRover를 크게 앞섰다. 특히 1 µs 롤아웃에서 STAR‑MD는 구조 붕괴 없이 안정적인 궤적을 유지했으며, 동역학 충실도(시간 상관 함수)와 구조 다양성(클러스터 수) 모두에서 최고 성능을 기록했다.

요약하면, STAR‑MD는 (1) 공동 시공간 어텐션을 통한 비분리적 상호작용 모델링, (2) 인과적 블록‑와이즈 트랜스포머와 KV‑캐시를 이용한 메모리 효율화, (3) 히스토리 노이즈 주입을 통한 장기 롤아웃 안정성 확보라는 세 축을 결합해, 마이크로초 수준의 장기 단백질 동역학을 고품질·고효율로 생성한다는 점에서 현재 가장 진보된 데이터‑기반 MD 가속화 프레임워크라 할 수 있다.