DMFF：共建生产级可微分力场计算引擎，让复杂力场参数优化不再难

开发精度高、可迁移性好的分子力场模型一直是微尺度科学计算领域的核心问题之一。近期，DeepModeling社区的开发者们发起了可微分分子力场（Differentiable Molecular Force Field, DMFF）项目，希望通过可微分编程框架技术打造全新的生产级力场计算引擎，以帮助解决参数调优困难、复杂力场计算等开发过程中的诸多痛点问题。

项目背景

准确而高效地描述原子、分子、粗粒化粒子间的相互作用是多尺度建模的重要问题之一。尽管以DeePMD为代表的神经网络方法在材料模拟等领域取得了较大的成功，基于物理的分子力场在生物和有机高分子等领域的模拟中仍然扮演着关键作用。然而，这些力场的开发至今依然保持着“小农作坊”的模式——传统力场参数的更新和维护缺乏自动化流程，往往需要大量有经验人员的人工干预。同时，新一代力场（如AMOEBA、MPID等多极矩可极化力场）的函数形式复杂：对多数力场开发人员而言，对此类新力场进行快速的分子动力学（MD）验证也并非易事。因此，分子力场的开发严重依赖于某几个课题组或商业公司，开发过程专业性强且开发速度慢。造成此问题的重要原因之一是，宏观拟合目标（如密度、自由能等性质）和先进力场本身的计算过程相当复杂，其对力场参数或原子坐标的梯度的解析求解困难，这给梯度下降优化算法和分子动力学的实现造成了较大障碍。

区别于传统编程范式，在人工智能领域中新发展起来的可微分编程技术使得开发者能聚焦于计算流本身，而将对参数的微分和相应的优化完全交给计算机自动进行。DMFF是完全基于此技术搭建的分子力场开发框架：通过DMFF，我们可以将各种物理力场和模拟计算流视为人工智能模型的自然延伸，以第一性原理或实验数据为输入，采用先进机器学习算法（而非人类经验）反向矫正、优化物理模型参数。由此，我们期望摆脱传统力场研究中手动调参和手动写MD代码的小作坊模式，并为自动化可持续的力场开发工作流奠定基础。

我们关注到，对于上述问题，尽管已有一些初步的尝试，但大多是在概念验证的阶段，且从项目设计的角度来看会在走向实用的路上遇到较大瓶颈。因此，值得强调的是，DMFF从项目开始起，便有一个明确的定位：成为一个适配多场景、兼具灵活性和高性能的生产级力场开发解决方案。

DeepModeling核心开发者们在DP系列方法、分子动力学引擎、高性能优化等方面经验丰富，并深入了解学界、业界的需求，这为实现DMFF的目标定位奠定了基础。但是，一个好的基础设施类项目需要整个领域来推动。因此，我们希望能有更多的开发者和感兴趣的用户加入到DMFF项目中来，共建生产级可微分力场计算引擎，让复杂力场参数优化不再难。

目前，DMFF项目已经以LGPLv3协议开源在DeepModeling社区中：

https://github.com/deepmodeling/DMFF

项目介绍

DMFF可以通过下载源码并使用pip进行安装。

git clone https://github.com/deepmodeling/DMFF.git
cd DMFF
pip install .

如下图所示，DMFF主要由基于OpenMM的「前端」模块和基于可微分编程框架JAX的「后端」组成，并设计了用户友好的API，方便用户快速进行力场参数解析和力场计算。

在「前端」模块中，DMFF通过Hamiltonian类（继承自openmm.ForceField类）解析xml格式的力场文件，并根据拓扑信息对输入分子进行参数化，最终再通过调用「后端」模块中用JAX实现的不同能量项（键长、键角、二面角、多极相互作用等）的计算核心代码返回相应的势函数。以使用GAFF2力场计算有机分子的能量为例：

import jax
import jax.numpy as jnp
import openmm.app as app
import openmm.unit as unit
from dmff import Hamiltonian, NeighborList

# load molecule
app.Topology.loadBondDefinitions("lig-top.xml")
pdb = app.PDBFile("lig.pdb")
ff = Hamiltonian("gaff-2.11.xml", "lig-prm.xml")
bondfunc, anglefunc, torsionfunc, nbfunc = ff.createPotential(pdb.topology)

通过以上类似OpenMM的前端API即可以得到各力场项的计算函数。利用JAX对这些函数进行自动微分，可非常方便地实现力（能量关于坐标的负梯度）和力场参数导数的计算：

force_fn = jax.grad(nbfunc, argnums=0)
force = -force_fn(positions, box, pairs, ff.getGenerators()[-1].params)

pgrad_fn = jax.grad(nbfunc, argnums=3)
pgrad = pgrad_fn(positions, box, pairs, ff.getGenerators()[-1].params)

案例分析

有机小分子电荷模型的优化

水合自由能是评价有机小分子电荷模型准确度的重要指标，常用的基于键电荷校正的模型如AM1-BCC），其参数就针对水合自由能进行过调优。然而由于自由能对力场参数导数的计算过于复杂，传统的参数调优过程要么依赖于费时费力的数值微分，要么依赖于计算化学家的直觉。DMFF则可以大大简化这一过程。作为概念验证，开发者用FreeSolv数据集中的27个酯类小分子对相关的BCC参数进行了调优。使用简单的梯度下降算法，经过8轮迭代，便可将自由能计算值相比于实验值的均方根误差从1.03 kcal/mol降低至0.36 kcal/mol。

水的多极矩可极化力场

在先进水模型的开发中，我们采用长短程分离的方案，从水的局域分子结构出发预测电荷、色散系数等原子参数，并使用这些原子参数计算体系总能量。对DMFF给出的多极矩可极化力场计算函数进行简单的二次封装，就可以在类似AMOEBA和MPID的复杂力场的基础上，进一步添加原子参数（电荷和色散系数）的结构涨落等新功能。基于DMFF框架，可以快速实现该力场的自动微分和分子动力学模拟，并给出相当准确的模拟结果。

水分子二体相互作用能，力场模型与CCSD(T)比较：（左）固定原子参数计算结果（右）结构涨落原子参数计算结果；蓝：最好，黑：一般，红：最差

RNA的粗粒化力场

RNA的CG模型致力于复现实验测得的三维结构，以达成预测折叠构象的目的。我们希望CG模型的能量最低结构与实验结构尽可能相符，这同样意味着实验结构下体系所受的应力需要尽可能小。通过简单的梯度下降方法，我们可以对模型中数十个力场参数同时进行优化，以降低晶体结构下体系中所含的应力。得益于可微分框架的灵活性，这一流程可以简便地进行扩展，加入成百上千的晶体结构信息令力场更加完备。

晶体结构下平均受力随训练步数的变化趋势

未来发展

DMFF正处在早期快速迭代发展的过程中，有许多地方亟待完善，也有更多的可能性值得探索：

• DMFF将与dflow等项目一起，实现一些常用的力场拟合工作流（如拟合二面角、自由能、概率分布等），进一步完善和发展案例库；

• DMFF将持续开发，根据需求实现更多的力场函数形式；

• DMFF与先进分子动力学算法（例如增强采样、路径积分等）的更良好的结合；

• DMFF自身也将持续优化，完善文档/API等，提高用户的使用体验；

如何加入

欢迎你在DMFF的Github项目中撰写 Issues，发起 Discussions，甚至提交 Pull Requests。具体来说：

• 如果你是分子力场的硬核开发者，正在探索新的力场函数形式，那么欢迎你与开发者深入交流，丰富DMFF的力场计算功能。

• 如果你正致力于某一特定体系的模拟研究，又苦于没有合适的力场参数，你可以成为DMFF的天使用户，根据自己的需求使用DMFF搭建力场优化工作流，并从实际的需求出发给我们提出宝贵建议，推动DMFF发展的更加完善。

此外非常欢迎你加入DMFF用户交流群，一起在群里探讨DMFF的使用技巧：

（如果群二维码失效，可以添加深度势能小编微信：deeppotential，小编会手动拉你入群）

结语

工欲善其事，必先利其器。在深度学习浪潮所引领的可微分编程技术飞速发展的今天，我们看到了新的力场研发的范式。我们希望将力场开发变成一个工程化、自动化、可复现的过程，从而享受持续集成/开发以及开源精神所带来的种种红利。DMFF将与DeepModeling社区内的相关项目一起，推广并实践这场变革。

- End -

（如需转载图文请与公众号后台联系）

-------------------------------

推荐阅读