图神经网络势函数：让分子动力学加速又精准

写这个系列写到了第63篇，今天把分子动力学（MD）里的图神经网络势函数这件事讲透。说句实话，被传统力场逼到去写机器学习势的人，都知道一个核心痛点：体系一换，力场参数就得手搓；温度一高，二面角势就开始飘；想研究反应路径，固定原子成键的模型压根没法断键。图神经网络势函数要解决的正是这个——用数据拟合一个能自动适应化学环境的势能面，让模拟在保留第一性原理精度的同时，跑出经验力场才有的速度。

这篇文章我尽量不绕弯子，按“为什么需要它→核心原理→完整落地实操→常见问题排查→进阶方向”的顺序写。不管你是做计算化学、材料模拟，还是刚从深度学习这边转过来，只要能看懂原子坐标和能量，这篇文章里的坑和建议都能直接用上。

1. 为什么分子动力学需要图神经网络势函数

1.1 传统力场绕不开的精度天花板

经典分子动力学的力场，本质上是一张参数表加一套固定函数形式。键伸缩用谐振项，角弯用谐振项，二面角用余弦展开，非键作用再叠上范德华和库仑项。这套做法几十年来非常成功，但它的天花板也很明显：函数形式一旦定死，能描述出的势能面形状就有限。

最典型的三个问题，第一是断键和成键。化学反应过程中键的生成与断裂是连续变化的，经典力场里键是否“存在”通常是写死的，反应坐标走到一半时就不知道怎么描述能量了。第二是极化与电荷转移。固定电荷模型没法描述原子环境变化后电荷重新分布带来的能量变化；可极化力场能改善一些，但参数拟合难度和代价会成倍上涨。第三是参数迁移性。同一套力场参数在这个分子上拟合得很好，换一个含杂原子、配位中心或者特殊官能团的体系，经常出现某些原子类型的相互作用参数根本查不到，需要人工去猜。

做过几个月力场拟合的人都会有一种感觉：调参数是在“补窟窿”。而机器学习势函数的思路完全不同，它不预设键项、角项、二面角项的公式，直接学“结构→能量→力”的映射。只要训练数据覆盖面够，它天然能处理断键、极化、电荷转移这些复杂效应。

1.2 机器学习势函数在做同一件什么事

机器学习势函数的思想并不复杂：构造一个可微函数，输入所有原子的元素类型和坐标，输出体系的总能量；能量对坐标求负梯度，就得到每个原子上的力向量。训练阶段，我们用第一性原理计算出的能量和力作为标签，做回归拟合。推理阶段，模型不再调用电子结构计算，而是直接给出能量和力。

这里有一个关键点值得先说清楚：机器学习势函数的“精度”上限，完全由训练数据决定。它不是黑魔法，更像是压缩了训练集覆盖范围内第一性原理势能面的高精度插值。数据里有反应路径，它才能预测反应路径；数据里只有平衡态附近的小振动，那描述高能碰撞/断键基本就是外推，早晚要出问题。

在模型结构的选择上，大致有三条路线：第一代是基于人工设计描述符加神经网络；第二代是图神经网络；第三代是带等变约束的图神经网络，也叫等变消息传递模型。这三条路线的本质区别在于：分子结构用什么方式表示，以及旋转对称性以什么形式进入模型。而图神经网络之所以成为主流，是因为“原子和化学键”天然就是一张图。

1.3 图结构本身比手工描述符更接近化学直觉

早期机器学习势函数常用手工设计的不变描述符，比如把中心原子周围的径向分布和角度分布编码成一组对称函数，再送入普通全连接网络。问题在于，这些描述符是人设计出来的，距离范围、光滑性、角度截断方式都要试；换一个原子密度差异很大的体系，描述符的选择可能又要重新来。而且很多手工描述符为了维持旋转不变性，会丢失大量方向信息。

图神经网络把分子或晶体看成一张图：原子是节点，截断半径内的近邻原子是边，边特征包含相对距离、方向等信息。在图这样的结构里，节点特征通过消息传递一层层更新，相当于模型自己去学习一段局部化学环境应该被压缩成什么样的数值表示。这比固定公式更灵活，也更容易迁移到不同分子体系。

所以图神经网络势函数迅速成为主流的底层逻辑很简单：用图来表达结构与相互作用，既有物理先验，又不限制具体函数形式。再往上走，加入等变约束，让模型对旋转操作有严格可预测的响应，方向敏感的性质也能一起学到。后面我会专门拆这部分原理。

2. 核心原理拆解：图构建、消息传递与力的解析梯度

2.1 把一帧分子变成一张图

在模型跑起来之前，第一步是建图。给定一帧结构，需要准备三样东西：节点特征、边索引、边特征。

节点特征通常是原子类型。我们把元素周期表里可能出现的元素放进一张可学习的嵌入表，比如氢对应索引1，碳对应索引6，氧对应索引8，模型训练过程中会为每个元素学到一个向量表示。之所以不用原子序数直接当数值特征，是因为元素之间没有平滑的数值关系，用可学习嵌入更通用。

边索引是邻居列表。对每个中心原子，在截断半径内搜索所有近邻原子，并把其中一对对应关系记为一条边。如果有周期性边界，还需要用最小镜像约定处理：只看最近的那个周期性镜像，避免同一个原子被复制出多条“假边”。

边特征至少要包含距离信息。常用的做法是把距离展开到一组光滑基函数上，比如径向基函数：

φ_k(r) = exp(-γ_k (r - μ_k)^2)

也就是一组中心位置不同的高斯函数，把连续的标量距离变成一个距离向量。这样模型就不再感知“原始距离数”，而是感知一个经过光滑编码的特征。为了让截断处的能量和力连续，还要乘上一个衰减函数，保证距离接近截断半径时边特征平滑归零。这一步非常关键，后面讲MD稳定性时还会提。

伪代码大致长这样：

# 构建分子图的伪代码 def make_graph(atom_types, positions, cell, cutoff): neighbor_list = find_neighbors_pbc(positions, cell, cutoff) n = len(atom_types) node_feat = embed_atom_type(atom_types) # (N, d_node) edge_index = neighbor_list.edge_index # (2, M) edge_vec = neighbor_list.relative_vectors # (M, 3) edge_dist = norm(edge_vec) # (M,) edge_feat = rbf(edge_dist) * cutoff_function(edge_dist) return Graph(node_feat, edge_index, edge_feat)

这步做得好不好，直接决定后面的模型上限。

2.2 邻居之间“传话”：消息传递的实质

图神经网络的核心操作叫消息传递。它的本质一句话就能概括：每个原子节点不断向邻居“打听”信息，再根据自己和邻居的状态更新自己的表示。

用公式表示就是：

h_i^(l+1) = h_i^(l) + Σ_{j∈N(i)} MLP([h_i^(l), h_j^(l), e_ij])

这里 h_i 是第 l 层节点i的特征，e_ij 是边特征，Σ 是邻居聚合操作。经过一层消息传递，中心原子就初步认识了自己的第一层邻居；经过 L 层，它就能“间接认识”距离 L×截断半径范围内的原子。这也是为什么层数和截断半径需要配套考虑：层数太少，长程信息进不来；层数太多，所有原子特征趋向平均，训练难度变大。

你可以把消息传递想象成一群人在一个圈子里传话：每传一圈，每个人掌握的信息就多了一圈邻居的消息。最终，模型从每个节点读出一个原子能量，再加总得到总能量：

E = Σ_i MLP_out(h_i)

这里每次加总的是“每原子能量”，而不是直接预测一个总能量。这么设计的好处是：模型天然具有原子总数的可扩展性，大规模体系也能用；同时，粒子间相互作用体现在加总前的节点状态里，而不是靠一个全连接层去硬拟合体系大小。

2.3 力不是另算的，而是能量对坐标的负梯度

很多刚上手的人会问：模型不是同时输出能量和力吗？严格来说不是。GNN势函数模型只输出能量，力的标准做法是自动微分：

F_i = -∂E / ∂r_i

训练时，我们对每一帧结构分别求能量和力损失。推理时，同样的操作也能一次性返回能量、力和原子坐标的梯度。这样做的好处是能量和力完全自洽：只要模型在坐标空间是光滑的，得到的力场就是保守力场，长时间MD里的总能量漂移能被控制住。

这里有一个实现上的小坑。如果训练时只算能量损失，模型输出的力会非常差；因为能量是标量，包含的信息量远小于 N 个原子的力向量。反过来，如果把力损失写好，就会自动促进能量在结构附近更光滑。所以在实际工程中，力损失几乎是必选的，而且权重通常要比能量损失更大。

需要补充的是，反向传播“能量对坐标的梯度”时，如果还要继续更新模型参数，那么输入坐标必须处于需要梯度的状态，并开启二阶梯度选项。很多初学者会忘掉这一步，结果前向算完一求梯度，参数就再也得不到更新。

2.4 不变模型还是等变模型，这里到底在争论什么

图神经网络势函数有两个路线：不变模型和等变模型。

不变模型只输出标量特征，能量天然满足旋转不变性；消息传递过程也只使用距离这类旋转不变特征。实现简单、训练稳定、计算量小。它的问题在于，很多方向依赖的信息需要在节点特征里表达，而纯标量特征对方向关系的编码比较弱。

等变模型允许节点特征按角动量分层，包含标量、向量、张量等不同部分；旋转分子时，这些特征会按对应的旋转矩阵一起变换。这样做的好处是，像原子力、偶极矩、极化率这类方向敏感量可以被模型结构直接承载，而不是靠后面硬凑。消息传递时，模型可以用球谐函数表达边的方向信息，能更好地捕捉取向和局部环境之间的关联。

普通分子体系的能量和力，用不变模型通常已经足够；但如果目标涉及极性分子、光谱性质、磁场响应，或者你想做带偶极/极化率的统一预测模型，等变路线会更合适。代价是等变模型的实现复杂度高，训练和推理的算力成本也更高。把这个选择放在“先跑通再谈扩展”的阶段会务实一点。

3. 完整落地实操：从数据采集到MD模拟

3.1 一套最小可行的端到端流水线

真正落地一套图神经网络势函数，流水线通常长这样：

1. 收集参考结构：比如从第一性原理MD轨迹、实验晶体结构、主动学习采样结果中收集一批原子构象。
2. 计算标签：用第一性原理方法计算每帧结构的能量和原子力。
3. 数据预处理：统一单位，统一截断半径，处理周期性盒子，划分训练集/验证集/测试集。
4. 构建图数据管道：把结构和标签转成节点特征、边索引、边特征、能量和力标签。
5. 训练模型：先小规模过拟合，确认代码没问题，再开完整训练，配合早停和模型平均。
6. 导出势函数：把训练好的模型包成一个输入坐标即输出能量和力的接口。
7. MD验证：先跑短时间NVT，观察温度、总能量漂移和结构合理性，再到目标工况下做正式模拟。

很多项目死在第5步之前，因为数据没整理好。数据质量永远比模型结构更重要。

3.2 数据、单位和切分：决定成败的前三步

先说单位。第一性原理软件输出的能量单位五花八门，有的用电子伏特，有的用哈特里；力也要对应统一。进入训练前，所有标签全部转成同一套单位，比如能量以 eV 为单位，力以 eV/Å 为单位。坐标也用 Å。这个看似基础的问题，会造成损失函数里量级混乱。

然后是能量的中心化。体系的绝对能量通常在几百甚至几千电子伏特量级，但模型更关心的是相对变化。常见做法是减去孤立原子能量参考，或者对每帧能量的每原子平均值做中心化。如果不做，损失会被一个“预测大数”主导，模型先拼命去平均能量，细节力反而学不到。

再然后是数据切分。这是最容易做错的一步。按帧随机切分是典型的错误，因为同一分子在不同帧之间高度相关，模型只要记住“这个骨架的原子类型组合”，测试集就能接近满分。正确做法是按“化学实体”切分：同一个分子、同一条聚合物链、同一块晶体结构的所有帧，都只出现在同一个集合里。对反应路径数据，甚至最好按反应过程整段切分，避免时间相邻构象跨集合泄漏。

考虑到主动学习循环，还可以补充一批远离平衡态的构象。比如对平衡结构做随机扰动、拉伸键、挤压角度，让模型见过高能区域，避免MD跑偏时直接外推到无约束区。

3.3 训练损失与超参数这样设比较稳

训练损失一般写成一个组合：

L = λ_E · MSE(E_pred, E_ref) + λ_F · MSE(F_pred, F_ref) + L2正则

能量项用均方误差，力项也是均方误差，但要注意力的分量数量是3N，远多于能量一个标量。如果两个损失权重一样，力在梯度里的占比天然会压过能量。这不一定好：有些体系能量变化很关键，需要平衡。

我给一个从工程角度比较稳的起点：λ_E 取 0.1 到 1.0，λ_F 取 1.0 到 10.0，具体看能量和力的量级。先跑一个极小数据集，让训练集误差掉下去，再开始调权重。模型是否过拟合训练集，这一步一眼就能看出来。

常用超参数可以这样设：

训练时建议开梯度裁剪，梯度范数设到1.0。模型在早期容易出现尖峰损失，裁剪能避免参数一步飞出去。早停也很重要，验证集指标连续若干轮不降就停；最后可以留一份指数滑动平均权重，或者对最后几个checkpoint做平均，得到的模型用于MD通常更稳。

3.4 把训练好的模型接进MD引擎

模型训练完之后，必须把它包装成MD引擎能调用的“势函数”。典型接口是这样：

# 伪代码：势函数封装 class GNNPotential: def __init__(self, model, unit_convert): self.model = model self.unit_convert = unit_convert def compute(self, atom_types, positions, cell, pbc): # 内部重新建图、预测能量、自动微分得到力 energy, forces = self.model.predict(atom_types, positions, cell, pbc) # 单位转换为MD引擎所需单位 return self.unit_convert.energy(energy), self.unit_convert.force(forces)

接进MD引擎时，两件事最容易踩坑。第一，单位转换要一次做对；否则模拟温度直接翻倍或者减半。第二，周期性边界处理必须和训练时一致：模型训练用最小镜像和截断半径，MD引擎里的邻居列表也要用相同规则，否则同一个结构在训练和模拟中会得到不同输入。

性能上，纯Python逐帧调用模型，在几百原子的小体系里勉强够用，但超过几千原子会非常慢。如果目标是大体系或长模拟，建议把模型导出成C++可调用的格式，或者在MD引擎里对多条轨迹做批量并行。实测下来，一个100原子数量级的体系，单步调用毫秒级延时是常态；性能优化要和精度验证一起做，别等模拟跑起来再补救。

接入后的第一件事不是直接跑正式模拟，而是做“冒烟测试”：先跑100ps的NVT，看温度是否稳定在目标值附近，总能量漂移是否在可接受范围；再把模型输出的径向分布函数和从头算短轨迹对比。这一步能暴露大量建图和单位问题。

4. 常见问题与调参排坑实录

4.1 训练时能量不错，力误差却一直很高

这是我被问得最多的问题。模型能量拟合得很好，MAE在每原子毫电子伏特量级，但力误差高到不敢用。原因通常是三件事。

第一，力损失的权重太低。力分量数量巨大，做单点拟合很容易被能量均值主导。解决办法是把 λ_F 调大，比如从1.0调到10，看力误差是否明显下降。第二，训练数据的构象多样性不够。能量好拟合，力是高维信息，需要更多远离平衡态的构象提供约束。补充一些随机微扰结构、拉伸压缩结构，往往比单纯增加平衡态帧数更有效。第三，模型不够光滑。如果模型用了不光滑的激活函数或缺失截断衰减，能量在近邻点之间依然能拟合，导数会出现震荡。解决方法是检查截断衰减函数是否连续，是否在靠近截断处仍然可导。

4.2 MD跑到一半体系爆炸

模型在训练集上误差低，不代表MD能稳定跑。常见现象是模拟初始阶段正常，几百步之后温度飙升、原子重叠、能量陡降。这本质上是MD跑到了训练数据覆盖范围之外。

第一个排查点是截断半径与光滑性。如果截断处能量力不连续，系统会不停地从“截断边界”吸收或释放能量，表现为总能量漂移。解决方法是在图构建阶段就乘上平滑衰减函数，并确保在近邻距离接近截断半径时，力和一阶导数都趋于零。

第二个排查点是外推区域。模型没见过高能态、原子间距极近或极高曲率的构象，就可能出现非物理的“能量坑”。稳妥做法是做主动学习：先用模型跑一段MD，把高不确定性构象抽出来，加入下一轮第一性原理计算重新训练。也可以在高能构象上增加软排斥项，避免原子接近时出现过小的模型能量值。

常见的一个误区是“数据里加误差大的构象不如多加点平衡态”。恰恰相反，活性高、形变大、接近解离态的构象，才是稳定MD的保障。

4.3 测试集指标漂亮，换个体系却完全不能用

这种“假成功”多半源自分组不当或数据偏移。随机切分时，同一分子的几千个构象被分到训练和测试两个集合，模型在训练时已经见过几乎一样骨架的力数据，测试集和训练集的分布高度重叠，指标自然好看。

但真实任务通常是：用A分子训练，跑到B分子上；或对某一类晶体结构训练，再预测另一类掺杂结构。测试集只能代表“训练集附近”的分布。如果想让泛化性更可信，需要按完整的化学实体切分，甚至刻意留出与训练集元素/官能团差异较大的体系作外部测试。

另一个隐蔽问题是元素嵌入表。图神经网络通常只能处理训练时见过的元素类型，一旦MD中遇到新元素，嵌入表直接失效。解决思路有几种：在训练数据中加入尽可能多的相关元素；将原子序数做成连续可外推特征；或者干脆把模型视为“专用势”，只在指定元素范围内使用。

4.4 问题速查表：症状、原因、操作

这张表我建议打印出来贴在工位旁边。绝大多数GNN势函数项目的前两次失败，都能在这里找到原因。

5. 进阶方向与多年实操下来的建议

5.1 一个主干多个输出：应力、偶极、电荷也能预测

图神经网络不只是预测能量和力。共享一个消息传递主干，在上面接多个输出头，可以同时学原子电荷、偶极矩、极化率、应力张量等。多任务学习还有一个额外好处：不同标签相当于不同的正则约束，会使中间表示更稳健。

应力张量对固体和应变模拟特别有用。要预测应力，模型必须感知周期盒形状，而且要对晶格参数求导；实现时需要把盒子向量也放入模型输入。偶极矩和极化率则需要等变特征；如果用不变模型硬拼，很难满足方向对称性。我的经验是，多任务输出要从项目开始就设计，而不是后面再加；因为标签类型会直接影响图特征的定义和消息传递结构。

5.2 局部GNN的天花板：长程相互作用怎么补

普通GNN势函数的信息交换只在截断半径内进行，适用于短程共价相互作用占主导的体系。但在离子液体、极性溶剂、带电缺陷这类长程静电显著的系统里，纯局部GNN会漏掉大量能量贡献。

常见的补救方法是“混合模型”：用GNN预测原子电荷，再显式加一个周期库仑项；或者把截断外长程作用通过网格方法并进模型，但整体拉长训练和推理成本。运行时也要注意，MD引擎如果本身有Ewald求和，势函数模块和引擎的静电处理不能重复计算，否则会双算。

5.3 预训练基础模型值得追，但别指望开箱即用

近两年出现了在大规模分子数据上预训练的基础势能模型，覆盖元素范围广，可以直接在通用分子上做初步筛选。这对小课题组很有吸引力，但也别抱着“开箱即用”的心态。基础模型在通用分布上表现好，在特定反应、特定缺陷或特定元素组合上仍需要微调。微调时建议全部层都参与更新，只微调输出头往往会丢掉底层结构表示。

如果算力有限，我的建议是先用中小规模模型加主动学习跑通自己的体系，评估收益比后再考虑要不要上更大规模的预训练模型。一味追大模型，反而会把项目拖进“数据清洗、环境配置、算力申请”的泥潭。

5.4 最想写给新手的一句话

如果让我给一个刚开始做GNN势函数的人提一条最核心的建议，那就是：在数据采集上花70%的精力，模型结构只花20%，剩下的10%留给MD验证。我见过太多人把时间耗在调网络宽度、层数和激活函数上，最后发现问题出在数据划分有泄漏、单位没统一、截断不连续。一个足够好的数据集配上中等规模的GNN，效果远超过一个“豪华模型”配劣质数据。

如果这篇写到这你心里已经有了一个完整的工作流，那就对了。从建图、消息传递到自动微分求力，从数据切分、训练到MD接驳，环环相扣。每一个模块单独看不难，难的是把它们串成一条稳定、可靠、可重复的流水线。这条流水线一旦跑通，你会发现分子动力学进入了另一个量级：第一性原理的精度，配上接近经典力场的速度，很多以前不敢想的长时程模拟都能提上日程了。