大规模生物分子动力学的深度学习:哈佛大学研究在各种系统上扩展了一个大型、预训练的 Allegro 模型

哈佛大学研究在各种系统上扩展了一个大型、预训练的 Allegro 模型

计算生物学、化学和材料工程依赖于能够预测原子尺度上物质的时间演化。虽然量子力学统治着原子和电子在微小尺度上的振动、迁移和键解离,但主导可观察的物理和化学过程的现象通常发生在相当大的长度和更长的时间尺度上。创新的高度可并行化的架构,具备对超级计算机进行访问,并具备快速和高度准确的计算方法来捕捉量子相互作用,是连接这些尺度所必需的。目前的计算机方法无法探测真实物理和化学系统的结构复杂性,它们的可观察演化的持续时间对于原子级模拟来说太长。

在过去的二十年中,对机器学习插入原子势的研究有很多。从高精度参考数据中学习的能量和力被用来驱动机器学习插入原子势,其与原子数量呈线性关系。最早的尝试使用高斯过程或简单的神经网络与手动设计的描述符相结合。早期的机器学习插入原子势预测准确性较差,因为它们无法推广到训练中不存在的数据结构,导致不稳定的模拟无法在其他地方使用。

哈佛实验室的新研究表明,使用Allegro可以以SOTA精度对具有多达4400万个原子的生物分子系统进行建模。团队使用了一个大型的、预训练的Allegro模型,用于原子数量从23,000个的DHFR到91,000个的Factor IX,400,000个的纤维素,44,000,000个的HIV壳,以及超过100,000个的其他系统。使用了一个带有800万个权重的预训练Allegro模型,通过在100万个结构上进行训练,以SPICE数据集上的混合功能精度实现了仅26 meV/A的强迫误差。通过学习整个无机材料和有机分子的潜力,可以进行前所未有的大规模材料系统快速超级计算。这是一个非常庞大而强大的模型,有800万个权重。

为了进行主动学习以自动构建训练集,研究人员表明可以有效地量化深度等变模型对力和能量的预测的不确定性。由于等变模型是准确的,精度瓶颈现在在于训练机器学习插入原子势所需的量子电子结构计算。由于高斯混合模型可以在Allegro中轻松调整,将能够使用单一模型运行大规模的不确定性感知模拟,而不是使用集合模型。

Allegro是唯一能够超越传统的消息传递和基于变压器的设计的可扩展方法。在各种大型系统中,他们展示了每秒超过100步的最高速度,结果扩展到超过1亿个原子。即使在4400万个原子的HIV壳这样的大规模情况下,通常错误明显更多,模拟结果仍然稳定,可持续纳秒级。团队在整个生产过程中几乎没有遇到任何问题。

为了更好地理解巨大的生物分子系统的动力学以及蛋白质和药物之间的原子级相互作用,团队希望他们的工作能为生物化学和药物发现开辟新的途径。