[论文解读] On the design space between molecular mechanics and machine learning force fields

Yuanqing Wang, Kenichiro Takaba+15 more|arXiv (Cornell University)|Sep 3, 2024

Machine Learning in Materials ScienceMaterials Science被引用 3

一句话总结

本文探索了分子力学（MM）与机器学习力场（MLFF）之间的设计空间，认为当前的MLFF虽然准确但速度过慢，难以广泛应用。为此，本文提出一种新一代MLFF架构，结合简单、可微分的操作与E(3)-等变归纳偏置，以实现高速与高精度的平衡，从而实现对生物分子体系的高效、物理上合理的模拟。

ABSTRACT

A force field as accurate as quantum mechanics (QM) and as fast as molecular mechanics (MM), with which one can simulate a biomolecular system efficiently enough and meaningfully enough to get quantitative insights, is among the most ardent dreams of biophysicists -- a dream, nevertheless, not to be fulfilled any time soon. Machine learning force fields (MLFFs) represent a meaningful endeavor towards this direction, where differentiable neural functions are parametrized to fit ab initio energies, and furthermore forces through automatic differentiation. We argue that, as of now, the utility of the MLFF models is no longer bottlenecked by accuracy but primarily by their speed (as well as stability and generalizability), as many recent variants, on limited chemical spaces, have long surpassed the chemical accuracy of $1$ kcal/mol -- the empirical threshold beyond which realistic chemical predictions are possible -- though still magnitudes slower than MM. Hoping to kindle explorations and designs of faster, albeit perhaps slightly less accurate MLFFs, in this review, we focus our attention on the design space (the speed-accuracy tradeoff) between MM and ML force fields. After a brief review of the building blocks of force fields of either kind, we discuss the desired properties and challenges now faced by the force field development community, survey the efforts to make MM force fields more accurate and ML force fields faster, envision what the next generation of MLFF might look like.

研究动机与目标

• 解决当前MLFF应用中的主要瓶颈——计算速度，尽管其精度很高。
• 通过探索MM与现代MLFF之间的设计空间，弥合传统分子力学（MM）与现代机器学习力场（MLFF）之间的差距。
• 识别MLFF开发中的关键挑战，包括速度、稳定性、泛化能力以及数据效率。
• 构想新一代MLFF，兼具高速与高精度，能够对复杂生物分子体系实现具有定量预测能力的模拟。
• 倡导开展社区范围的努力，创建高质量、多样化的数据集，以实现MLFF训练中对广阔化学空间的覆盖。

提出的方法

• 调研现有的MM与MLFF架构，识别在精度与效率之间实现平衡的功能形式与设计原则。
• 提出一种基于简单、可微分操作（如点积）的新MLFF架构，可普遍逼近E(3)-不变函数。
• 将物理归纳偏置（如旋转与平移不变性）融入神经网络架构，以确保平滑性与稳定性。
• 利用自动微分进行力的计算，并在通用张量加速框架（如PyTorch、JAX）中实现模型。
• 强调采用可扩展的、物理信息驱动的训练策略，包括课程学习与自适应批量处理，以提升数据效率。
• 倡导在大规模、高质量数据集上训练基础模型，以实现化学空间中的少样本或零样本泛化。

Figure 1: Overview of the design space between molecular mechanics (MM) and machine learning (ML) force fields.

实验结果

研究问题

• RQ1 现代力场中速度与精度之间的关键设计权衡是什么？如何实现优化？
• RQ2 如何显著提升机器学习力场的速度，同时保持低于1 kcal/mol的化学精度？
• RQ3 哪些架构选择与归纳偏置能够以最低计算成本实现对E(3)-不变能量函数的通用逼近？
• RQ4 社区驱动的高质量数据集是否能够实现MLFF在多样化化学空间中的泛化，而无需大量重新训练？
• RQ5 直接生成Boltzmann分布的生成建模在多大程度上可消除分子模拟中对显式力场的需求？

主要发现

• 当前的MLFF在有限化学空间上已超越1 kcal/mol的化学精度阈值，但其速度仍比分子力学力场慢几个数量级。
• MLFF采用的主要瓶颈是计算速度，而非精度，尤其在大规模生物分子模拟中尤为明显。
• 基于简单、可微分操作（如点积）的下一代MLFF架构，可实现对E(3)-不变函数的通用逼近，同时保持高速性能。
• 在模型架构中引入E(3)-等变归纳偏置，可确保物理一致性、平滑性与稳定性，且不损失表达能力。
• 社区范围内的高质量、多样化数据集的构建，对于实现在广阔化学空间中的少样本或零样本泛化至关重要。
• 新兴的生成模型（如Boltzmann生成器与扩散模型）可能最终使显式力场过时，通过一步直接从Boltzmann分布采样实现模拟。

Figure 2: Between MM and QM energies and forces, there is little correlation. Scatter plots and kernel density estimate (KDE) of: (a): MM energy ( $U_{\mathtt{MM}}$ , mean-subtracted) plotted against QM energy ( $U_{\mathtt{QM}}$ , mean-subtracted); (b): MM force magnitude ( $|F|_{\mathtt{MM}}$ ) pl