北大提出ProAR：概率自回歸捕捉分子動力學軌跡中的構象不確定性

作者 | 論文團隊

編輯丨ScienceAI

AlphaFold 等結構預測模型已經顯著提升了我們預測蛋白質靜態結構的能力。但對許多生物過程而言，真正關鍵的并不只是某一個穩定構象，而是蛋白質如何在不同構象之間連續變化，以及這些變化中包含怎樣的不確定性。

分子動力學模擬（MD）正是研究這類動態過程的核心工具，但其計算代價高、可達時間尺度有限。近年來，深度生成模型開始嘗試學習 MD 數據并生成蛋白質軌跡；然而，現有方法往往將整段軌跡作為固定長度的高維時空對象聯合生成，難以廣泛探索自由能面，并靈活擴展到不同長度的動力學過程。

基于此，來自北京大學等機構的研究團隊提出了一種用于軌跡生成的概率自回歸框架——ProAR。ProAR 的核心貢獻，是將蛋白質 MD 軌跡建模為一個概率自回歸過程：模型沿時間逐步生成軌跡，并在每一步預測構象分布，而不是單一確定結構。通過這種設計，ProAR 使軌跡生成同時具備構象不確定性建模與可變長度生成能力。在 ATLAS 大規模蛋白質 MD 數據集上，ProAR 相比已有 SOTA 方法重建誤差降低約 7.5%，構象變化準確性平均提升約 25.8%。

相關論文以《ProAR: Probabilistic Autoregressive Modeling for Molecular Dynamics》為題發布，已被 AAAI 2026 錄用。

論文鏈接：https://ojs.aaai.org/index.php/AAAI/article/view/36974

項目鏈接：https://github.com/kevincheng7/ProAR

軌跡生成的難點，在于長度和不確定性

近年來，深度生成模型已經開始用于蛋白質構象和軌跡建模。例如，AlphaFlow、BioEmu 等方法關注從平衡分布中采樣構象；MDGEN、AlphaFolding 等方法進一步嘗試直接生成動力學軌跡。

不過，現有軌跡生成方法通常采用聯合建模高維時空表示的方式，在空間維度和時間維度上同時去噪，生成一段固定長度軌跡。這一路線帶來兩個關鍵限制。

第一，它與 MD 的序列演化結構并不完全一致。真實 MD 可以視為由隨機微分方程描述、并通過逐步積分得到的隨機過程。直接對整段軌跡進行聯合去噪，會顯著增加建模復雜度，也容易削弱模型對時間相關構象變化的表達能力。

第二，固定長度的非自回歸設計限制了推理階段的靈活性。實際應用中，研究者往往需要根據任務生成不同長度的軌跡，或者從已有構象繼續向前擴展。固定長度軌跡生成模型在這一點上天然受限。

另一方面，簡單地采用確定性自回歸也不夠。蛋白質動力學本質上包含熱漲落和多路徑構象變化。如果模型只預測一個確定的下一幀，往往會傾向于輸出平均化結果，導致構象多樣性不足，也難以充分覆蓋自由能面上的重要區域。

因此，一個更合適的建模目標應當同時滿足三點：

• 遵循 MD 逐步演化的序列結構；
• 顯式表達構象變化中的不確定性；
• 在長軌跡生成中控制誤差累積，保持軌跡穩定。

ProAR 正是圍繞這三個目標設計的。

ProAR 的關鍵轉變：預測一個分布，而不是預測一個點

圖 1：ProAR 的動機和方法論

ProAR 是 Probabilistic AutoRegressive Modeling 的縮寫。它將 MD 軌跡生成寫成自回歸問題：給定初始結構和條件信息，模型逐步學習并生成后續軌跡。

與普通自回歸模型不同，ProAR 不直接輸出一個確定的下一幀結構，而是將每個幀建模為一個結構化的多元高斯分布：

模型不僅預測該幀最可能出現的位置，即均值；還預測協方差，用于表達不同殘基區域的運動不確定性以及局部相關變化。隨后，模型從該分布中采樣得到具體構象。

這種設計有兩個直接作用。首先，它避免了確定性預測容易產生的平均化問題，使模型能夠生成多樣化構象。其次，協方差結構為蛋白質不同區域的動態差異提供了顯式表達能力，例如結構化區域和無序區域可能具有不同的運動幅度與相關模式。

雙網絡協同建模中間態與未來態

圖 2：ProAR 的架構設計

ProAR 使用兩個互補的網絡模塊：插值器（Interpolator）和預測器（Forecaster）。

插值器的任務是：給定兩個端點構象，預測中間時刻的構象分布。具體來說，它以當前幀和未來某個端點幀為條件，輸出中間幀的均值和稀疏結構化協方差，并通過負對數似然與結構損失共同訓練。

預測器的任務是：在給定當前歷史結構的條件下預測未來幀。它采用 corruption-refinement 機制：先對歷史構象施加與預測跨度相關的隨機擾動，再基于當前結構進行一次性精化預測。

兩個網絡共享基于 SE (3) 等變模塊的幾何骨干，結合 IPA 與 EGNN，并利用 ESM-2 嵌入以及時間步、氨基酸序列、二級結構和殘基索引等特征，為模型提供序列、結構和時間信息。

長軌跡能否穩定，取決于如何抑制自回歸漂移

自回歸生成的一個核心挑戰是誤差累積。即使每一步預測誤差很小，長時間遞推后也可能導致軌跡偏離真實分布。對于分子結構而言，這種漂移還可能帶來不合理的鍵長、原子碰撞或結構扭曲。

ProAR 為此設計了anti-drifting sampling（抗漂移采樣）策略。在每個自回歸循環中，預測器先給出遠端未來幀的預測；插值器再根據當前幀和該未來幀生成中間幀；隨后，預測器利用更接近目標時刻的中間幀繼續精化未來幀。隨著上下文逐漸接近目標時刻，未來幀預測被迭代修正，從而減輕自回歸過程中的漂移。

這一機制使 ProAR 能夠在保留概率采樣能力的同時，穩定生成任意長度的軌跡。

實驗驗證：更穩定的長軌跡，更準確的構象變化

ProAR 在 ATLAS 數據集上評估。ATLAS 包含約 1300 個不同大小和拓撲結構的蛋白質，每個蛋白質具有三條獨立的 100 ns MD 模擬軌跡。實驗覆蓋長軌跡生成、構象采樣和構象插值三個任務。

1. 長軌跡生成：250 幀預測中保持更低誤差

在長軌跡生成任務中，ProAR 從初始結構出發生成 250 幀軌跡，并與 100 ns 參考 MD 軌跡比較。以逐幀 Cα-RMSE 衡量，ProAR 在所有評估長度上均優于 MDGEN；在 250 幀處，誤差從 3.813 ? 降至 3.529 ?，降低約 7.5%。

圖 3：ProAR 生成的軌跡可視化

考慮到 MD 軌跡本身具有隨機性，論文進一步在 PCA 空間中比較生成軌跡與參考軌跡的構象位移分布。結果顯示，ProAR 在四項位移指標上均優于 MDGEN，平均提升約 25.8%，說明它不只是降低單幀誤差，也更準確地捕捉了構象變化的幅度與分布。

2. 構象采樣：面向軌跡訓練，也能逼近平衡分布

圖 4：ProAR 的構象采樣性能

雖然 ProAR 的主要目標是生成時間相關的 MD 軌跡，它在時間無關構象采樣中也表現穩健。與 AlphaFlow、CONFDIFF 等專門的構象采樣模型相比，ProAR 在 7 項指標中的 5 項取得最佳結果，說明概率自回歸軌跡生成也能夠較好覆蓋 MD 模擬對應的平衡構象分布。

圖 5：ProAR 軌跡生成情況

3. 構象插值：在兩個狀態之間恢復連續過渡

圖 6：ProAR 的構象插值性能

在構象插值任務中，ProAR 連接起始構象與遠端構象，生成中間過渡路徑。82 個 ATLAS 測試案例顯示，生成路徑在 PCA 空間中呈現平滑、定向的狀態轉變趨勢，并與參考 MD 軌跡較為一致。這表明插值器不僅能補全中間結構，也能學習構象轉變中的動態規律。

ProAR 的意義：把分子動力學的生成機制寫進模型

ProAR 的貢獻不只是提高了若干指標。更核心的是，它為蛋白質動力學軌跡生成提出了一種更貼合問題底層結構的建模范式：

用自回歸機制對應 MD 的逐步演化，用概率分布表達構象不確定性，用 anti-drifting sampling 控制長程遞推中的誤差累積。

這一區別使 ProAR 同時具備三方面能力：生成任意長度軌跡、捕捉時間相關的構象多樣性、并在長軌跡生成中保持穩定。

對 AI for Science 研究而言，ProAR 也提供了一個更一般的啟示：在科學生成任務中，模型結構應盡可能反映對象本身的生成機制。對于分子動力學而言，這意味著不僅要生成合理結構，還要建模結構隨時間演化的隨機過程。