JCTC | 深度学习分子模拟出现，传统分子动力学大势已去了吗？

背景

分子动力学模拟(MD)可以较为准确的模拟分子的运动，通常分子力场包含所模拟的各类原子及溶剂分子，并通过键长键角等来表示其作用力，在模拟分子构象、折叠等方面具有较高的准确性。但分子动力学模拟存在两个弊端，首先分子力场计算非常耗时，且参数拟合非常复杂。其次，分子动力学模拟很难在大的生物尺度上进行模拟。随着深度神经网络(DNN)体系结构的出现，机器学习(ML)变得特有吸引力，它使定义任意复杂的函数及其导数成为可能。DNNs提供了一种非常有前途的方法，在从更精确的方法获得的大规模数据库上进行训练后，在MD模拟中嵌入快速而准确的势能函数。DNNs的一个特别有趣的特征是它们可以学习多体相互作用，并预测系统的力和能量。TorchMD是一个从头构建的分子动力学代码，利用ML库PyTorch的语言。通过将MD中使用的键和非键扩展到任意复杂的DNN, TorchMD实现了快速的机器学习。TorchMD的两个关键点是，它是用PyTorch编写的，很容易集成其他ML PyTorch模型，如从头算神经网络(NNPs)和机器学习粗粒化。TorchMD在Lennard-Jones系统和生物分子系统上进行端到端可微分子模拟。本文介绍了TorchMD的功能，重点介绍了支持的功能形式和数据驱动DNN电位的有效拟合策略。

方法

2.1 TorchMD模拟和势能分析

TorchM不仅是一个标准的分子动力学代码，它提供NVT集成模拟，郎格文恒温器，初始原子速度是由麦克斯韦玻尔兹曼分布导出的，积分使用velocity Verlet算法。用反力场法对远距离静电场进行了近似计算。TorchMD也支持周期性系统的模拟。最小化是使用L-BFGS算法完成的。因为它是使用PyTorch数组用Python编写的，所以修改起来也非常简单，而且模拟可以在PyTorch支持的任何设备上运行(CPU、GPU、TPU)。然而，不像专门的MD代码，它不是为速度而设计的。TorchMD使用与经典MD代码一致的化学单位，如kcal/mol表示能量，K表示温度，g/mol表示质量，Å表示距离。TorchMD支持通过parmed读取AMBER力场参数。除此之外，为了更快地构建原型和开发，它实现了易于阅读的基于yaml的force-field格式。图1给出了模拟水盒子的YAML力场文件示例。目前，TorchMD缺失的功能包括氢键约束和邻居列表。TorchMD实现AMBER的基本功能形式。它提供了所有基本的AMBER内容：键长、键角、二面角、非键范德华能和静电能。

图1. YAML力场水分子示例

图片来源于JCTC

2.2 训练机器学习能力

TorchMD提供了一个完全可用的代码，用于在名为TorchMD-net的PyTorch中训练神经网络能力。目前使用的是基于SchNetbased的模型。通过提供一个简单的力计算器类或其他ML势，可以直接从非参数核方法导出力。将每一步的位置作为输入，并用选择的方法计算出外部能量和力。本文从SchNetPack中提取了特性化和原子层，但是重写了整个训练和推理部分。为了允许在多个GPU上进行训练，可使用PyTorch框架对网络进行训练。TorchMD也可以通过更改随机数生成器种子，将神经网络潜力分成一批来提高速度，从而同时运行一组相同的仿真，从而至少部分地恢复了优化的分子动力学代码的效率。对于QM9数据集，使用标准损失函数对能量进行均方误差训练模型。对于粗粒度模型，使用自下而上的“力匹配”方法进行训练，重点是从原子模拟中重现系统的热力学。

结果

3.1 全原子系统的模拟与性能

比较了TorchMD与高性能分子动力学代码ACEMD3的性能。在表1中可以看出，在测试系统上，TorchMD比在TITAN V NVIDIA GPU上运行的ACEMD3慢60倍左右。大部分的性能差异都可以归因于缺少非绑定交互的邻居列表，并且对于更大的系统来说这是不允许的，因为这对距离不能装进GPU内存中。尽管目前TorchMD实现的性能较低，但它的端到端可微性允许研究人员进行实验，这是传统更快的MD代码无法实现的。为了评估TorchMD实现AMBER力场的正确性，我们将其与OpenMM进行了比较，应用于离子、水盒子、小分子到全蛋白等14种不同的系统，从而测试所有不同的力场项。结果表明OpenMM与TorchMD的势能差均小于10 -3 kcal/mol。使用TorchMD对周期水箱进行了NVE模拟，模拟时间步长为1 ns，时间步长为1 fs。得到了1.1 ×10-5 K的平均值，显示出了良好的能量守恒。

表1. 不同方法以1 fs/ timestep对50,000步模拟的性能比较

表格来源于JCTC 

3.2 示例1——在QM9数据集上进行训练验证

为了验证TorchMD-Net的训练过程，本文在QM9基准数据集上进行了训练和性能评估。QM9由133,885个有机小分子组成，最多有9个C、O、N和F型重原子，展示了每个分子的一些热力学、能量和电子特性。我们对能量U0进行训练，由于数据集提示几何一致性检查失败，排除了3,054个分子。剩余的分子被分割成包含110,000个样本的训练集和包含6,541个样本(5%)的验证集，剩下14,290个样本供测试。通过使用PyTorch lightning，在多个GPU上使用带有学习率调度程序的Adam optimizer对网络进行训练。图2显示了QM9训练的配置文件示例。本文使用TorchMD-Net与不同数量的训练数据进行训练，在随机选择的5%数据集上进行验证并证明了训练的正确性。结果表明最佳表现为10 meV，略优于SchNet在QM9中的最佳表现。

图2. QM9训练的配置文件示例和QM9数据库学习曲线

图片来源于JCTC

3.3 示例2——点到点的可微模拟展示

分子动力学程序包中自动区分(AD)的可用性在ML应用之外是有益的。它能够计算所有数值运算的梯度，为灵敏度分析，力场优化和定向MD模拟以及高度复杂的约束和约束下的模拟打开了新的途径。为了演示这些功能，本示例从短MD轨迹推断出力场参数。首先，使用具有柔性键和角度的TIP3P水模型模拟了一个包含97个水分子和一个Na+ 和Cl-离子对的水盒子。在能量最小化和300 K的NVT平衡之后，在集合中以10 ps的速度进行模拟。模拟使用了1 fs的时间步长，9Å的截断半径和7.5Å的截断距离，每10步保存一次坐标和速度。接下来，中和系统的电荷，将其按0.01比例缩放以确保静电势的梯度不变。为了从MD轨迹推断q，首先初始化了积分器。然后，使用修改后的程序运行10个模拟步骤，并将此简短模拟的最终位置与相应的后续轨迹进行比较。由于TorchMD具有AD功能，因此该模拟器考虑到了周期性边界条件。其次使用Adam进行训练，图3显示了训练过程中训练损失和部分原子电荷的演变。1000次迭代之后，原始电荷的准确率达到了3％。

图3. 从短轨迹中推断部分原子电荷

图片来源于JCTC

3.4 示例3——粗粒化全原子系统

本文构建了两种粗粒化模型：一个仅基于α-碳原子，另一个基于α-碳和β-碳原子，如图4所示。

图4. 两种粗粒化模型

图片来源于JCTC

3.4.1 训练集

我们选择Chignolin的CLN025变体(序列YYDPETGTWY)，它在折叠时形成一个β型发夹结构。由于它的体积小(10个氨基酸)和快速折叠的特性，已通过MD进行了广泛的研究。训练数据是通过在GPUGRID.net分布式计算网络上使用ACEMD在显式溶剂中对蛋白质进行全原子模拟获得的。将包含一个奇诺林素链的系统放置在40Å的立体水盒子中，该水盒子包含1881个水分子和两个Na +离子。使用CHARMM22 力场和水的TIP3P模型在350 K下对该系统进行了模拟。使用Langevin积分器，其阻尼常数为0.1 ps-1。积分时间步长设置为4 fs，带有重氢原子，并且在所有氢重原子键项上均受完整约束。静电的计算是使用粒子网格Ewald进行的，其截断距离为9Å，网格间距为1Å。模拟时间为180μs，每100 ps保存一次力和坐标，总共获得1.8×106帧构象。

3.4.2 神经网络潜能训练

对于粗粒化模拟，重要的是要提供一些固定势函数，以将动力学可以访问的空间限制为训练数据中采样的空间。本文将它们限制为键和排斥力。键可防止蛋白质聚合物链断裂，并且排斥力会在没有数据的非常近的原子距离处停止计算NNP。对于键合部分，使用了全原子训练数据来为每对键合类型构建距离直方图。计算了在距离间隔相等的间隔内适合选择的相应距离所花费的时间的比例，α碳之间的所有键的距离分别为3.55 Å和4.2 Å，而碳键之间的所有键数为1.3和1.8Å。通过α碳和β碳之间的所有键间距离分布经玻尔兹曼求逆得到自由能分布，并用于拟合平衡距离r0和各个谐振子常数k。类似地推导了非键合排斥项的先验电位。距离直方图由30个等距间隔的3Å和6Å之间的距离直方图构成，并用于拟合排斥电位的参数ϵ。在拟合电势曲线时，非线性曲线拟合使用SciPy软件包的Levenberg Marquardt方法执行。力的参数存储在YAML力场文件中。使用力匹配方法对网络进行了训练，其中将预测的力与训练集中的真实力进行比较。发现增加CACB模型的高斯函数的数量可以提供更高的模型稳定性，并防止在模拟过程中形成不合理的非物理结构。

3.4.3 NNP模拟

将力场和训练网络相结合，利用TorchMD对CA和CACB系统进行了模拟。将模拟的参数作为YAML格式的配置文件引入到TorchMD中，并具有指定的网络位置，嵌入和计算器。输入文件如图5所示。对于这两个粗粒化模型，都以1 fs的时间步长在350 K的条件下运行10 ns的模拟，每1000 fs保存一次输出。尽管模拟使用的时间步长很小，但粗粒系统的有效动力学要比全原子MD系统快得多。由于TorchMD可以轻松处理并行动力学，因此可同时运行十个模拟，其中五个从折叠状态开始，另外五个从展开状态开始。图6展示了通过时滞独立成分分析(TICA)获得的势能面，用于全原子基线模拟和通过TorchMD获得的粗粒度模拟。为了更好地诠释势能面，图7展示了具有代表性的轨迹的前2 ns的RMSD值的绘图。结果表明，两种模型的粗粒度模拟都能够获得折叠和展开的稳定构象。CA模型的能量表明，它捕获了折叠状态作为全局能量的最小值。CACB模型还可以正确检测全局最小值，但无法猜测展开区域的自由能。总体而言，模拟的稳定性不如CA模型稳定，并没能准确预测折叠状态下的最稳定结构。

图5. 输入文件的示例

图片来源于JCTC

图6. 全原子MD模拟的二维势能面(左)和两个粗粒度模型CA(中心)和CACB(右)

图片来源于JCTC

图7. CA模型和CACB模型的前2 ns轨迹的RMSD值和原始轨迹

图片来源于JCTC

结论

本文演示了TorchMD，这是一种基于PyTorch的分子动力学引擎，用于具有机器学习功能的生物分子模拟，并展示了几种可能的应用。包括从AMBER全原子模拟到端到端学习，以及最终的粗粒度神经网络在蛋白质折叠方面的潜力。此外蛋白质折叠构建NNP时，需要为其补充键和排斥的渐近分析势，以防止探索训练中未访问的构象。而在这些数据中NNP的预测是不可靠的。此外本文还展示了如何将蛋白质粗粒化为α-碳原子或α-碳和β-碳原子。当前，CA模型效果最好，但是未来的研究将表明哪种模型更适合于更多样化的目标。TorchMD参数的端到端可区分性是Open Force Field Initiative等项目可以利用的功能。此外，为了提高速度，计划促进OpenMM和ACEMD中潜在的机器学习功能的集成，并可能在将来开发插件来扩展对更多MD引擎的支持。同时，TorchMD可以通过促进ML和MD领域之间的实验，加快模型训练评估原型开发周期以及促进在分子模拟中采用基于数据的方法来发挥重要作用。

参考文献

Stefan Doerr, Maciej Majewski, Adrià Pérez, Andreas Krämer, Cecilia Clementi, Frank Noe, Toni Giorgino, and Gianni De Fabritiis, TorchMD: A Deep Learning Framework for Molecular Simulations, J. Chem. Theory Comput., 2021, ASAP. DOI: 10.1021/acs.jctc.0c01343.