Scalable Spatio-Temporal SE(3) Diffusion for Long-Horizon Protein Dynamics (STAR-MD)

ByteDance Seed | arXiv: 2602.02128 | 2026年2月

🎯 Joint 時空間アテンション搭載の SE(3) 等変オートリグレッシブ拡散モデルで、タンパク質 MD 軌跡をマイクロ秒スケールまで安定生成。ATLAS 全指標 SOTA を達成する。

① 背景と課題

古典 MD はフェムト秒積分ステップが必要で μs〜ms スケールへのアクセスが計算コスト上困難。既存の深層学習 MD 手法は固定コンテキスト窓に制約され長時間拡張で時間的一貫性が失われる問題があった。

AlphaFolding/MDGen は固定コンテキスト窓—長時間拡張で過去記憶を切り捨て時間的一貫性喪失

ConfRover は O(N³L) 計算量・O(N²L) KV キャッシュで大規模・長時間に非スケーラブル

→ Joint S×T アテンションで空間-時間を統合、Block-causal 訓練で長時間一貫性を確保

② 手法: CausalDiffusionTransformer

② 手法: 3本柱の訓練テクニック

① Block-causal Attention

clean フレーム + noisy フレームを 2 倍長 concat し全フレームを 1 forward で並列訓練—長時間一貫性確保

② Contextual Noise Perturbation

履歴フレームに τ~U[0, 0.1] のノイズを付加—自己予測誤差への頑健性

③ Continuous-time Conditioning

物理ストライド Δt を LogUniform[0.01, 10] ns でサンプリングし AdaLN で条件付け—推論時の長時間外挿を可能に

③ 本研究で示したこと（要点）

④ 主な結果 (a) ATLAS 100ns 比較

④ 主な結果 (b) Recall / VAMP-2 比較

手法	Recall↑	tICA↓	VAMP-2差↓
MDGen	0.28	-	-
ConfRover	0.36	-	0.47
STAR-MD	0.54	0.17	0.07
MD Oracle	1.00	0.00	0.00

VAMP-2 差 0.07 — MD オラクルに最も近い値

④ 主な結果 (c) KVキャッシュスケーラビリティ

④ 主な結果 (d) 1μs長時間外挿

⑤ テイクホームメッセージ

Joint 時空間の理論的優位性
Mori-Zwanzig 形式により空間相関除去が非分離時空間結合を必要とすることを証明—理論保証付き設計。

ATLAS SOTA
全指標で ConfRover/MDGen/AlphaFolding を上回り 1 μs 外挿でも安定。現時点で最良のタンパク質 MD サロゲート。

限界
残基レベル表現で全原子詳細なし—FEP への直接適用は困難。ATLAS 100 ns 訓練でミリ秒スケール事象は未経験。

lib/md への統合価値
μs スケールコンフォメーションアンサンブルを高速生成し RMSDAnalyzer/HBondAnalyzer の入力として使用可能。

先行手法との比較

手法	長時間安定	KV キャッシュ	SOTA
AlphaFolding	✗	O(N²L)	JSD 0.59
MDGen	△	O(N²L)	R 0.28
ConfRover	△	O(N²L)	JSD 0.52
STAR-MD	✓ 1μs	O(NL)	JSD 0.43

本研究のインパクト