——简介——
如何高效采样蛋白质的多种结构一直是计算结构生物学领域的一大难点。通常此类问题通过分子模拟（包括MC采样和MD采样）来解决，同时配合各种增强采样算法提高效率。近年来，得益于GPU算力的进步，各种基于深度学习的构象生成算法被提出，但大都直接生成离散的构象系综，无法生成具有时间顺序的构象变化轨迹。2024年9月，MIT CSAIL研究所Tommi Jaakkola课题组的Bowen Jing等人开发了基于流匹配（flow matching）的基于关键帧的分子动力学模拟轨迹生成模型MDGen [1]。该模型可以基于给定的单一构象生成类似于分子动力学模拟轨迹的具有时间信息的连续变化构象系综，同时也能对于动态轨迹进行结构补全（即inpainting）。
——方法——
MDGen利用3000条三肽和1900多条单链蛋白的分子动力学模拟轨迹（各100 ns）进行训练，其关键在于对分子动力学模拟轨迹的特征工程，以便利用神经网络模型进行建模，基本逻辑如下图1所示。类似于已有的大量蛋白质结构深度学习模型，MDGen将单个残基作为token，但在sequence维度之外增加一个时间维度。作者并没有直接使用蛋白质SE(3)模型常用的residue frame进行残基建模，而是使用残基相对于关键帧的位移和旋转信息进行建模，具体的tokenize方法如下所示：
![]()
在给定K个关键帧的情况下，每个残基使用7K+14维的特征向量表示。同时，作者使用了目前生成模型领域较为常用的流匹配模型，具体的网络架构使用了在时序深度学习领域表现较好的Scalable Interpolant Transformer（SiT）网络架构。考虑到MD轨迹中的帧数可能很大（如100 ns的轨迹按10 ps一帧进行保存，则会产生10,000帧），作者将SiT架构中的时序attention部分更换为Hyena架构，用于以更低计算复杂度代替传统的Transformer。作者将训练好的模型应用在了4个下游任务上：前向模拟（Forward Simulation）、路径预测（Interpolation）、轨迹插帧（Upsampling）、动态修补（Inpainting）。其具体原理如图1所示：![]()
图1. MDGen的应用场景（左）和实现方法（右）作者随后对MDGen的四种下游任务进行了测试。对于Forward Simulation任务，作者使用第一帧作为参考，训练模型向后生成10 ns的轨迹。在生成时，模型对每个肽段迭代生成100 ns的轨迹，并和真实轨迹进行比较。作者比较的标准主要有：（1）生成轨迹与MD轨迹的结构分布是否匹配，（2）是否能捕捉动态内容，（3）是否能在更短的时间内遍历采样空间。作者采用对扭转角和TICA的降维分量计算Jensen-Shannon散度（JSD）的方法来衡量结构分布的匹配程度，结果如图2所示：![]()
图 2. Forward Simulation的采样结果。（A）MD（橙）与模型生成（蓝）的6种扭转角的分布。（B、C）TICA降维得到的自由能景观分布。（D）通过MSM构建的亚稳态相关性。（E）沿TICA的最慢独立分量的实际去相关时间。（F）扭转角的去相关时间相关性（蓝），不同亚稳态对之间的通量矩阵相关性（橙）。（G）MD和模型生成的丙氨酸四肽构象对比。作者通过马尔可夫状态模型（MSM）分别得到了真实轨迹和生成轨迹的亚稳态分布，并发现模型很好地捕捉到了状态的相对排序，且很少遗漏重要状态或对稀有状态赋予较高的权重。此外，通过比较轨迹中每个扭转角的去相关时间，可以发现生成轨迹和MD轨迹具有较高的一致性（图2F）。作者认为这一结果说明了模型能够区分缓慢和快速弛豫的扭转势垒。最后，作者表明模型能够在约60秒的GPU时间内生成等效于100 ns的轨迹，而MD模拟则需要约3小时的GPU时间，强调了模型生成相比于MD模拟的速度优势。 对于给定前后构象的变构过程采样任务（interpolation / transition path sampling），作者在训练时给定前后构象为关键帧，并使模型产生中间的1 ns构象。在测试时，作者挑选轨迹中两个肽段分离最远的状态，并产生1000个中间构象。最终的评估结果如图3所示：![]()
作者对生成的轨迹和MD轨迹分别构建MSM亚稳态，并将100 ns的轨迹MSM亚稳态结果作为参考标准。如图3所示，生成的路径比从任何少于100 ns的重复MD MSM中采样的路径具有更高的可能性，而100 ns正是参考MD模拟的长度。此外，MDGen的路径集合在访问状态的分布上与参考MD MSM的JSD表现最佳，并且具有最高比例的有效非零概率路径。 对于轨迹插帧（upsampling）任务，作者将10 ps间隔的轨迹插帧至100 fs。通过比较扭转角的自相关函数：![]()
的分布（该函数自变量为间隔时间Δt），可以评估的模型性能，具体如图4所示：![]()
图 4. 左图：扭转角的自相关函数（直线代表100 fs-timestep MD，小点代表100 fs-timestep生成结果，大点代表下采样的10 ps-timestep结果）。右图：自相关函数对对数时间尺度的负导数。作者随后探讨了基于动力学轨迹的补全任务，即在轨迹中遮去一部分结构，让模型进行修补（inpainting）。不同于常规的蛋白设计，作者评估了针对四肽（遮去中间两个残基，保留末端残基）的toy model，在给定观测到的残基构象转变的情况下，设计在相应马尔可夫状态的肽段残基类型，结果如表1所示：![]()
表 1. MDGen和baseline模型在序列恢复率上的比较此外，作者展示了MDGen在较大蛋白的系综生成任务上的效果，并和已有的系综生成模型AlphaFlow进行比较。在此任务中，作者使用ATLAS数据集的轨迹进行针对单链蛋白的训练。此时由于残基数目大大增加，作者生成250帧、时间步为400 ps的样本，使单个样本能够模拟100 ns的ATLAS参考轨迹，具体生成性能如表2：![]()
表 2. MDGen在单链蛋白上的结构系综生成性能
可见，MDGen精确度介于AlphaFlow和MSA下采样（MSA subsampling）之间，但每生成一个结构的速度比二者快几个数量级。图5是一个案例蛋白的MD轨迹系综与MDGen生成系综的可视化结果，其中，相似的RMSF图像也说明MDGen除了能较好的复现结构，也能较好的复现结构的某些动力学性质。![]()
图 5. （上）MDGen生成的6UOF_A蛋白和MD模拟的系综（下）RMSF数据
相比已有的蛋白构象系综生成模型，MDGen能够直接生成MD轨迹。但就目前而言，MDGen只能对短轨迹的小肽进行生成，同时仅对扭转角、降维数据等评价指标进行了分析，尚未提到生成的构象在能量上是否合理，是否存在一些原子碰撞等。同时，作者在文中提到，MDGen受限于关键帧，无法直接从序列开始采样，且依赖关键帧的准确度。笔者认为MDGen的主要贡献如下：1. 将类似于视频生成的训练流程和网络架构思路引入到MD轨迹中，并很好的和结构深度学习结合起来。2. 提出了4种下游任务，启发了这类模型的应用前景。3. 此类任务此前缺乏很好的评测流程，作者提出的很多评测项目都具有相当程度的创新性与合理性。[1] Jing B, Stärk H, Jaakkola T, et al. Generative Modeling of Molecular Dynamics Trajectories[J]. arXiv preprint arXiv:2409.17808, 2024.
