北大提出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 也提供了一个更一般的启示：在科学生成任务中，模型结构应尽可能反映对象本身的生成机制。对于分子动力学而言，这意味着不仅要生成合理结构，还要建模结构随时间演化的随机过程。