神经网络分子力场：3D等变深度学习实现又快又准的分子模拟

做分子模拟的人，大概都经历过这种纠结：想要精度就得跑量化计算，想要速度就得用经验力场。而神经网络分子力场这个方向，恰好就是奔着“又快又准”去的——用3D等变几何深度学习来学习分子体系的能量和受力，本质上是把第一性原理计算得到的势能面压缩进一个可微分的神经网络里。这几年我陆陆续续试过好几条技术路线，踩过不少坑，也看到它从“玩具分子”一步步走向真实材料模拟。这篇文章想把我对神经网络分子力场的理解、选型逻辑、实操流程和比较常见的翻车点都摊开讲一遍，适合刚接触等变模型、或者已经在跑传统消息传递模型但准备转向等变架构的读者参考。

1. 神经网络力场到底在解决什么问题

1.1 两难处境：量化计算太慢，经典力场太糙

模拟一个几百原子的体系，如果你要用密度泛函理论（DFT）算能量和受力，每一步自洽迭代动辄几十秒甚至几分钟，分子动力学一个纳秒需要成千上万步，时间上完全撑不住。经典力场倒是跑得飞快，但势函数是预先定好解析形式的，谐振子项、范德华项、库仑项的参数一旦写死，就很难描述键断裂、电荷转移、配位几何变化这类复杂化学过程。在我自己的项目里，最常见的场景是：体系里有配位键，力场参数库根本不覆盖这种元素组合，最后只能硬着头皮回去跑DFT。

神经网络力场的思路完全不同。它不规定势能面的解析形式，而是直接把原子坐标映射到能量和受力上，用大规模数据去拟合高阶相互作用。核心假设是分子的能量具备局部性——每个原子的贡献主要取决于它周围一小团邻居的几何构型。基于这个假设，模型可以做得非常快：一次前向推理，几百个原子的能量和力在毫秒级就能出来，比DFT快上好几个数量级，精度却能逼近训练数据所在级别。

1.2 题目里的三个关键词怎么理解

“3D等变几何深度学习”听上去绕，拆开其实很直白。3D指的是模型输入是三维空间中的原子坐标；几何深度学习意味着网络结构本身要尊重空间变换（旋转、平移、翻转）下的物理规律；分子表示指的是神经网络内部如何编码原子及其邻域环境。这三点凑在一起，构成了神经网络分子力场的核心设计约束。

一个合格的分子表示至少要满足四个条件：平移不变、旋转等变、局部平滑、对坐标可微。平移不变好办，把所有坐标减去质心就行。旋转等变稍微麻烦一点：体系转动之后，能量必须不变，而受力必须跟着一起转动，这就不是随便拿坐标塞进全连接层能搞定的了。局部平滑靠截断半径和衰减函数实现。可微是必须的，因为受力要从能量对坐标求负梯度得到。这四个条件缺一个，模型在真实分子动力学里就会出各种别扭的问题。

1.3 为什么传统分子描述符不够用

早期不少工作用对称函数、A汤键序参数这类手工描述符来喂神经网络。它们确实能保证旋转不变，但信息量和可学习性都受限。手工设计特征本质上是把人的先验固定死，分子里真正起作用的往往是多体关联、长程反馈这类复杂信号，靠几条预定义公式很难覆盖。等变深度学习的好处在于，它不需要你手工承担“把环境转换成不变特征”这件事，模型自己可以通过等变特征（向量、高阶张量）来编码方向信息，层层组合，最后在能量层完成对称性收缩。信息没有丢失，方向关系被保留，但最终标量能量对所有旋转都保持不变。

2. 为什么模型必须“3D等变”：从旋转测试说起

2.1 一个很多人忽略的检验：旋转一下输入

我有一次检查自己的模型，训练集误差已经收敛得很漂亮了，顺手把一批测试分子绕随机轴转了个角度再forward一遍，结果能量变了百分之十几。当时的模型把原子坐标直接映射成特征，表面对齐了训练分布，但物理规律完全没有学上。其实能量不该随旋转变化，因为它只依赖分子内部相对几何；力则应该跟着分子一起转，转之前指向右，转之后指向新的右。

这就是等变（equivariance）和不变（invariance）的区别。数学上，一个映射Φ满足Φ(g·x) = g·Φ(x)，叫等变；满足Φ(g·x) = Φ(x)，叫不变。分子力场里，标量能量是E(3)不变的，矢量力是SE(3)等变的，高阶响应（比如极化率张量）按更高阶表示变换。模型内部如果每一步都满足这个性质，输出自然满足，不需要靠数据去碰运气硬凑。

2.2 从“坐标不可用”到“距离描述符”到“等变特征”

直接拿笛卡尔坐标喂网络是行不通的，因为坐标数值本身会随平移动。于是大家最早用距离矩阵——两个原子之间的距离天然旋转不变。问题在于距离只保留了径向信息，丢掉角度和二面角，而分子势能面对键角、扭转角非常敏感。想表达完整局部几何，至少要包含键长、键角、二面角的组合信息。沿这条路发展出了基于角和二面角的显式消息传递模型，代表性方法包括定向消息传递那一族，这类方法相对容易实现，但消息和特征混合了不同阶信息，表达力有一定天花板。

另一条路是直接在网络内部维护一组等变特征。我拿到一个位置的坐标差，把它投影到球谐函数基上。球谐函数Y_l^m的最大优点是旋转时有明确的变换规律，层与层之间用张量积把不同阶特征组合起来，信息逐层丰富，但对称性始终被精确控制。这类模型从不可约表示出发，把节点特征分解成多阶分量，每一阶按旋转群的不可约表示变换，整体严格等变。代表作品包括早期的等变网络模型，以及后来大批基于同一套数学框架的改进版本。

2.3 两个“笨办法”测试等变是否成立

实操层面，判断一个力场模型是否等变，不需要懂群论，做两个实验就行。

第一，旋转一致性测试。对同一个体系生成若干个随机旋转副本，每个副本单独推理，比较输出能量是否一致、输出力是否在旋转后被拉回到同一条参考线上。误差应该在数值精度范围内（比如机器精度乘以尺度，而不是百分之几的相对误差）。

第二，质量守恒测试。神经网络力场输出力的总和应当为零，因为平移自由度的合力不改变。如果模型对每个原子的受力都独立预测，很容易出现合力不为零的情况，动力学上会让整个体系自己加速飘走。等变架构天然不会出这个问题，这也是我坚持用等变模型而不是普通标量图网络的原因之一。

3. 核心架构怎么选：两代等变模型对比

3.1 三代模型演进的逻辑线

我还是想用一个非常粗略的谱系来梳理：第一代是手工描述符加神经网络，精度有限。第二代是以图神经网络为核心的消息传递模型，从只考虑距离发展到距离加角度、再加二面角，表达能力沿着“信息类型变多”这条线提升。第三代是严格等变表示模型，特征携带旋转信息，层内通过张量积更新，物理对称性从网络结构上被保证。

第二代里比较有代表性的有SchNet（只做距离卷积）、DimeNet（显式加入键角）、GemNet（进一步加入二面角）。它们的思路是把分子几何拆成显式的二体、三体、四体项，用数学公式组合成不变标量。优点是实现相对直观、训练稳定，缺点是不变标量做多了之后，高阶几何信息容易丢失，或者要靠特征维度去堆。

第三代以NequIP、Allegro、MACE等为代表。NequIP用球谐基和张量积维护等变特征，Allegro进一步把每个原子中心的等变信息在局部坐标系里标量化，速度和稳定性都有提升，MACE引入了多体消息，让每个原子读到的信息不只是周围邻居的一层聚合，还有更远尺度的有效关联。这几类方法在公开基准上的表现已经明显超过第二代模型，尤其是在力预测和能量守恒方面。

3.2 张量积为什么贵，贵在哪里

等变模型最核心的操作是张量积。假设我有一个l1阶特征和一个l2阶特征，张量积会生成从|l1-l2|到l1+l2的所有阶的分量。物理上这对应角动量耦合，化学里常见的s、p、d轨道杂化也能用它来描述。但计算上张量积需要按各个分量循环，复杂度随阶数增长很快。所以模型里l通道一般不会开太高，常见是l=0,1,2或者0到3。l=1只管向量信息，l=2能描述更精细的方向分布，但训练和推理都要多花一两倍时间。

我在实际中一般这样取舍：体系大小几百原子、需要长时间MD采样时，把l上限设为2，隐藏特征控制在128左右，速度优先；体系只有几十个原子、关键是要高精度势能面时，l开到3也没有问题。注意张量积使用的是三维旋转的Wigner-D矩阵，这部分在高精度下计算也比较慢，部分实现会预计算旋转矩阵并缓存，能省不少事。

3.3 局部坐标系技巧：既快又稳的秘密

Allegro系列的方法给我印象最深的一点，是把全局等变特征转换到原子局部坐标系里做标量化。做法是：对每个原子中心，先根据它最稳定的几个邻居方向构建一个局部坐标系（比如两个键方向做Gram-Schmidt正交化），然后把所有等变向量投影到这个局部坐标系的三个轴方向上，得到一组标量。标量后续操作不会丢失旋转信息，因为原始向量已经在局部坐标系里被唯一表示过。

这样做的好处是：球谐张量积在整个网络主体里根本不需要，大部分操作是普通标量神经网络，速度提升非常明显。同时向量在不同结构之间被“对齐”到了局部坐标系，模型对结构的响应更平滑。代价是局部坐标系本身需要在形变中保持连续稳定，否则同一种构型在不同坐标系下投影会跳跃。实操时我吃过这个亏，分子在模拟中大幅扭转，参考键方向接近共线，局部坐标轴发生翻转，界面处能量噪声增加。解决办法是选多个参考向量做平均，同时对翻转情况加连续性惩罚。

3.4 架构选择对照表

4. 从零搭一个神经网络力场：数据、训练与评估

4.1 数据准备比模型更决定成败

神经网络力场本质上是在做从几何结构到能量的回归。最标准的训练数据来自DFT计算，每一帧包含原子种类、原子坐标、体系能量、每个原子受力。数据组织形式一般是扩展XYZ，原子符号和坐标照常写，注释行里放总能量和盒子信息，力可以放在每个原子坐标后面作为附加列，也可以用独立字段存。

数据切分是第一个大坑。绝大多数初学项目都按随机比例划分训练验证测试，这种方法带来的结果是同分子内不同构象容易被分到两个集合里，模型“见过”的构象家族其实涵盖了测试集，性能虚高得离谱。正确的做法是按分子骨架、分子类别或者反应路径来切分，保证测试集分子与训练集分子在结构上足够不同。比如分子甲的训练构象和分子乙的测试构象之间不会互相串场。

我在处理QM小分子基准数据时，习惯先把所有体系按化学骨架做聚类，再按7:1:2划训练验证测试；对于催化材料类数据，按吸附位点来分，绝不让同一个活性中心的同构型同时出现在训练和测试里。划分完之后还需要检查一个事：训练集和测试集原子种类的覆盖范围。稀有元素如果只出现在测试集，模型根本学不出来，需要在数据收集阶段就想办法平衡。

4.2 邻居构建、截断半径与平滑衰减

分子力场势能面可以近似看作局域相互作用叠加。邻居定义通常用截断半径r_c。r_c设太小，长程氢键、范德华作用被硬生生截断，力和能量在边界处突变；r_c设太大，每个原子的邻居数暴涨，图网络的消息聚合规模迅速膨胀，训练变慢。多数分子体系我选5埃，如果是含带电基团或氢键体系，至少要到6到7埃。对于周期性体系，邻居要用最小镜像约定来做，坐标差经过wrap处理，同时要小心原子在盒子边界附近时邻居重复计数的问题。

截断处不能简单地“一刀切”掉原子，否则势能面会不连续，分子动力学必然出问题。标准做法是把距离项乘一个衰减函数，常见CAuchy型或者余弦型：f_cut(r) = 0.5 * (cos(πr/r_c) + 1)，在r=0时值为1，r=r_c时值为0，一阶导数为0。这个平滑性直接影响力的连续性，也影响MD中的能量守恒。我见过有人贪省事直接硬截断，结果NVE模拟里体系总能量每步跳一个小台阶，原因就是截断边界处的受力不连续。

4.3 核心代码思路：能量回归，力用自动微分

神经网络力场的输出设计有一点和普通分类回归模型不同：网络前向出来的标量是总能量E，不是各原子力。推理时用自动微分求原子受力：

python复制import torch

coord = torch.tensor(...)  # (N, 3)，需要求梯度
coord.requires_grad_(True)

# 模型返回总能量（单位：eV/kcal/mol，取决于数据单位）
E = model(coord, species, edge_index)  # (1,)

F = -torch.autograd.grad(E.sum(), coord, create_graph=True)[0]

在训练时如果对每个batch都执行一次E.sum()的梯度求解，计算图会保留两层，内存消耗比普通前向多不少。因此一个很常见的做法是：训练阶段让模型同时在“能量模块”和“力模块”上输出，能量模块给出预测总能量，力模块要么直接输出显式力，要么也走autograd。显式力头的缺点是能量和力可能出现不一致，解决办法是加一个能量-力一致性约束，让网络输出力与energy梯度互相拉近；纯梯度方案则保证严格自洽，但算力更贵。

我个人的经验是：数据量不大、体系刚性较强时，纯梯度方案更好，训练稳定，MD不飘；数据量大、需要跑批量推理时，可以用显式力头加一致性损失来省算力，但一致性损失的权重需要仔细调，太大会让能量损失训不动，太小又等于白加。

4.4 损失函数：能量、力、应力的权重怎么配

最基本的损失项是能量误差和力误差。由于总能量随分子规模增大而增大，直接用总能量MSE来回归，不同分子之间的能量数量级差异会主导损失。稳妥做法是用每原子平均能量：E_per_atom = E/N_atoms，预测值和参考值都做同样的处理。力误差是真正的“局部”监督信号，它告诉模型每个原子受到的相互作用方向是否正确，对势能面的光滑程度影响最大。

我的默认配置是：

• 能量项：平均每原子能量MSE，权重设为1。
• 力项：平均每个原子三分量力MSE，权重设为100到1000。
• 应力项（模拟周期性体系时）：应力张量分量MSE，权重按盒子体积归一。

力权重为什么比能量大这么多？因为力是能量对坐标的导数，它对局部几何非常敏感，稍微有一点能量函数波动，力的误差就会被放大。如果不把力的重要性拉高，模型可能能量拟合得很平，受力却很抖，跑MD必然不稳定。在实际调参时我发现，多数情况下把力权重从100调到1000，力的MAE能下降明显，而能量误差不会恶化太多；但如果力权重超过2000，能量精度反而会变差。

4.5 训练配置与超参选择

训练配置遵循一套我反复在用的稳妥配方。优化器用AdamW，初始学习率1e-3，配合一个短warmup阶段（前几千步从1e-5线性升温到目标值），然后cosine衰减到1e-5。batch size一般设16到32，如果体系很大导致显存溢出，就梯度累积。混合精度（AMP）能显著提速，但等变模型的张量积部分在fp16下稳定性差一些，我建议先在fp32下跑通一个小周期，确认指标没爆，再切AMP。若是NequIP这类强等变训练，AMP带来的速度提升很明显，但遇到NaN要回退到fp32。

学习率warmup的作用，我理解是等变张量积初始化时特征范数可能偏大，直接大步长更新容易震荡。warmup给了模型先稳定特征尺度再优化的时间。隐藏维度一般64到128，交互层数2到3层。层数再多，信息半径变大，模型开始去拟合中等距离上的弱相互作用，训练会慢、也不会更准。

我自己的经验是先固定一组小超参数跑通pipeline，再扫描截断半径和力权重。模型架构之类的参数扫描性价比不高，真正影响精度的是数据覆盖和损失权重。

4.6 评估指标与“上了车再测测”

训练指标就看三个：能量MAE（meV/atom）、力MAE（meV/Å）、运行时长。但这三个指标只能说明拟合能力，不能说明模型的动力学行为。我的评估清单里一定包含：旋转一致性测试、NVE能量漂移测试、RDF（径向分布函数）对比、与DFT轨迹重叠性检查。

NVE测试非常关键：在恒定能量系综下，体系总能量在统计上应当恒定，只有数值噪声。如果模型预测的能量表面存在局部振荡，NVE总能量会随时间线性或随机漂移。我见过模型在测试集上的力MAE已经非常好看，但NVE跑50皮秒能量就掉了几个kcal/mol，这说明势能面上的精细振荡没被数据监督到，梯度方向大体对但幅度不稳定。NVE漂移量级控制在零点几kcal/mol/100ps以内，是我对“能不能用于MD”的硬指标。

5. 实战中的坑与排查：等变性破功、能量漂移与外推失稳

5.1 等变性测试失败：先查输入预处理再查球谐

等变模型原则上不会“主动”破功，但测试失败往往出在输入没对齐。最常见的原因是整批数据做标准化时不小心用了全局平移，比如减质心变成了减每个sample自己的质心，这会造成不同原子受到不同平移，破坏相对坐标。严格做法是每个分子为单位减质心，平移到原点，标准化只作用于整体平移和缩放，旋转保持不变。另一个隐蔽问题是坐标dtype和模型内部使用的Wigner矩阵dtype不一致时，旋转矩阵计算精度下降，旋转一致性测试在某些高角度下会出现几个meV级别的误差，这不是架构问题，是精度问题。

5.2 能量发散或者训练出现NaN

神经网络力场训练中发散，多半不是优化器参数问题，而是输入数据存在极端构象。有些DFT取样会给出某些键断裂或原子非常接近的构象，距离接近0时球谐基值极端，消息传递数值爆炸。解决办法：训练前先清理数据，去掉原子间距明显小于共价半径之和的帧；或者在邻居特征计算时加一个安全下限，比如最小距离不取0而取0.05Å。还有一种情况是初始张量积的权重范围没设好，先用fan-in/fan-out这种经验初始化，再把等变层初始输出压缩一个小系数，能避免第一轮迭代就出现爆炸。

5.3 力与能量不一致

这个坑几乎每个用显式力头的人都会踩一次。训练时模型输出的力Field和能量Field各自拟合，单独看都还行，但MD里需要F=-∇E严格成立，否则能量守恒直接失效。应对办法：优先选择纯梯度方案，如果实在要显式力头，就在损失里加一致性正则项，并且在评估时不仅看力和能量的单独MAE，还要看“能量求得梯度”与“显式力输出”之间的差异均值。

注意：一致性正则项的权重不要一开始就拉满。先用小权重跑通，观察两个损失的比例稳定了再逐步增大，不然训练早期模型会为了满足一致性而牺牲能量精度，之后很难拉回来。

5.4 分布外化学空间失稳

神经网络力场的强项是插值，弱项是外推。训练集里没有的新官能团、新电荷态或者新的配位环境，模型预测出的势能面大概率是“合理地平滑乱走”。怎么判断模型“遇见了没见过的东西”？最简单办法是训练完保留所有训练样本的中间层特征，新样本进来时算一下它与训练样本特征分布的距离，距离太大就标记为低置信度。

主动学习是解决外推失稳的主力工具：在MD模拟中定期采样一批新构象，用集成模型算预测分歧度，把分歧最大的一批结构送去补做DFT，写进训练集再来一轮训练。我在催化表面吸附体系上试过这种迭代，前两轮精度提升特别快，到第三轮提升幅度会下降，说明数据覆盖已经比较饱和了。这个方法的一个隐藏副作用是训练集规模越来越大，周期性结构导致能量绝对基准不好对齐，需要额外处理参考态校正。

5.5 周期性体系、应力与长程相互作用

周期性与非周期体系的差别不只是邻居数变多，而是需要把盒子信息传进模型。邻居搜索要用最小镜像约定，势能面的形变要能响应盒子尺寸变化。如果要输出应力张量用于NPT系综模拟，模型还需要一个与盒子体积相关的输出通道，并且力的计算要考虑virial修正。不少等变模型没有直接给出应力输出，我用PBC数据训练时会额外做一个力头来预测应力，再和P=∂U/∂V的数值差分核对一致性。

长程相互作用是所有纯局域神经力场的命门。截断半径之外虽然不做图卷积，但静电和色散相互作用的尾巴还在，就靠硬截断会直接破坏体系总能量。一个折中办法是加一个后处理项：用经典长程库仑或者LJ势的贡献做校正，神经网络只拟合短程残差。在一线实际应用里，这种“神经网络+经典长程项”的混合框架，比重建一个大截断半径的模型要稳得多，也快得多。

6. 我踩过坑之后的几条心得

项目推进到中后期，我最大的体会是：等变模型真正要学的不只是“怎么算”，而是“怎么把物理约束嵌进去”。用户输入坐标，模型输出能量和力，看似一个回归问题，但中间牵着的几何对称性、能量-力自洽、局部性假设、长程近似每一条都是物理，哪一条松了，后面MD就会拿噪声教育你。

还有一点经验很实用：调试神经网络力场，不要只看单帧指标。把训练好的模型放到NVE模拟里跑，观察温度和能量轨迹。凡是训练指标好看但MD里不稳定，一定是模型学到了正确的平均势能面，但没有学到正确的曲率。力才是曲率的直接体现。这也解释了为什么力权重在训练里那么重。

对于刚入手的同学，我的建议是从第二代模型跑通pipeline，换到第三代等变模型时保留同一套数据、损失和评估流程，先做AB对比，不要一上来就换全套架构。保持数据固定是判断模型改进是否有效的前提。如果迭代到某个程度精度上不去了，优先去补数据和重新切分，而不是继续堆网络容量。

最后再分享一个小技巧：训练结束后，把模型在训练集和测试集上的特征分布都保存下来。以后每加一批新数据，都能快速判断它落在模型熟悉区域还是陌生区域。这套“置信度预警”机制，在神经网络力场这种高风险替代DFT的场景里，可能比提升0.1 meV的MAE还重要。