米喜神经网络分子力场这几年在计算化学圈子里算是顶流话题，但真正能把SE(3)等变性和分子表示这层逻辑讲清楚的内容并不多。大部分入门资料要么堆数学公式，要么直接甩代码，读者看完依然不明白为什么要折腾等变图神经网络，也不清楚训练一个能跑的模型需要跨过哪些坑。这篇博文我想从一个实际做过的角度，把3D等变几何深度学习与分子表示结合的核心思路、实操细节和踩坑记录完整拆一遍，项目代号就叫“模拟项目X”。
在动笔之前，先给还没接触这个领域的朋友一个定位。传统分子力场用手工设计的键长、键角、二面角势函数描述分子能量，固定原子类型、固定电荷，遇到新骨架或者复杂化学环境就抓瞎。神经网络力场的思路完全不一样——用神经网络直接把原子坐标映射到体系能量，训练数据来自高精度量子化学计算（比如DFT或者波函数方法），理论上保留精度又大幅降低推理成本，让毫秒级甚至微秒级分子动力学模拟成为可能。
我接下来要写的，就是这个话题在各个层面如何落地，从为什么需要等变，到模型怎么选，再到具体训练和部署中真正遇到的问题。
1. 为什么分子表示必须满足3D等变性，这决定了模型能不能用
1.1 物理对称性不是可选项，而是硬约束
先想一个问题：你拿一个水分子，把整个分子在空间里旋转任意角度，能量变了吗？答案显然是没有。分子没变形、没碰壁、没发生任何实质变化，就是整体朝向换了个方向，物理体系完全一样，能量当然一模一样。这就是旋转不变性。同时，平动也不该影响能量。分子整体平移之后，能量也该保持不变。而原子受力方向则必须跟随旋转一起旋转，这叫等变性，也叫协变性。
一个神经网络力场如果不满足这些约束，会发生什么？同一个构象换个朝向，算出来的能量和受力就变了，模型直接失效，连最基础的旋转一致性都保证不了，做不了任何可靠的应用。所以3D几何深度学习用于分子表示，第一步就要把原子坐标的任意旋转和平动“处理掉”。
经常有人问，能不能在输入上做数据增强，比如每次训练时随机旋转分子，让模型“见得多”从而“学习到”旋转不变性？纸上谈兵可以，实操却不可靠。用数据增强近似等变性，样本效率和泛化能力都很差，模型很难在连续对称空间上真正学到精确的不变性。而且等变性是一种结构约束，体现在网络架构里之后，不需要用额外数据去反复教。就好比你雇人搬箱子，与其让他自己体会一百次怎么不闪腰，不如直接给他正确的搬运动作框架，一次就到位。
1.2 等变性与不变性的关系：从旋转坐标到旋转预测
再细化一点，等变性和不变性是不同级别的对称性要求。
- 标量任务，例如预测总能量、特定温度下的自由能差值，要求旋转不变即可。
- 矢量或张量任务，例如预测原子受力、偶极矩、极化率、应力张量，必须同时满足等变性。力要跟着分子一起转。
在实践中，最常见的是用能量-受力联合训练。先让网络输出体系总能量，这个能量是旋转不变的。然后利用受力是能量对坐标负梯度这一关系，通过自动微分把原子力算出来。这里有个非常优雅的地方：如果能量满足旋转不变性，那么力的等变性是自动成立的。这种做法也被大量力场模型采用，既保证了物理一致性，又不用额外设计复杂的高维等变网络去单独预测力。实际操作时，注意力得分、边特征、原子位置编码都要考虑旋转行为，不能随意设计。
1.3 从群论到网络设计：SE(3)等变几何深度学习在做什么
“SE(3)”指的是三维空间中的旋转群SO(3)和平移群的半直积，整体上描述的是空间中刚体运动。处理分子这一体系，就是在SE(3)对称性约束下做深度学习。
几何深度的具体落地做法是把分子的原子坐标、原子编号、可选的电荷信息以及三维盒子（周期边界条件）转成图数据，然后用图神经网络迭代更新节点特征（原子表示）、边特征（原子对信息）。在等变模型中，特征的每个通道不仅要存标量值，还要存矢量以及更高阶的不可约表示（irreps）。简单说，不同角动量阶数的特征在旋转变换下按各自的矩阵表示转换，网络通过层与层之间的耦合算子保证输出仍然守规矩地变换。
最直观的结果就是，你把输入分子转个角度，网络内部每个节点特征的“方向分量”也跟着转，而标量分量保持不变。最终经过张量积和可学习的径向权重收缩，得到旋转不变能量。整个过程不需要任何分子对齐、不需要把分子放入固定坐标系的预处理，网络天然就尊重物理对称性。
需要模型API调用？ 免费领10W Token，多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 分子表示方案的进化史：从原子到几何再到等变特征
2.1 从二维图到三维几何：为什么必须使用距离和角度
最早的分子表示是二维拓扑图。原子是节点，化学键是边，用RDKit之类的工具可以从SMILES或者分子图直接构造。但是二维图丢掉了太多空间信息——顺反异构、手性中心的绝对构型、空间位阻、分子在三维空间中的实际形状都没表达。对于预测能量、催化活性、结合亲和力这类任务，空间结构的影响往往比谁和谁连了一条键更大。
所以GNN家族中SchNet的出场更像是拐点。这个模型以原子类型作为初始节点特征，以原子对距离作为边特征，经过连续滤波卷积更新节点表示，最终将节点特征求和得到总能量。它不显式区分化学键类型，也不需要预定义键连关系，完全依赖距离信息。很多人第一次接触SchNet时觉得“这也能行？”但实际上它在QM9等数据集上已经取得了远超传统手工特征基线的精度，也让后续大量模型选择不依赖化学键信息，而是依赖几何。
只靠距离依然不够。分子中的角度信息，尤其是氢键方向、共轭体系的轨道取向，很多时候比距离更敏感。DimeNet这类模型就试图加入角度项，利用方向性消息传递编码三原子之间的夹角。从距离到角度，模型终于开始建立局部分子几何的完整图景。同理，如果再做二面角，就能区分顺反式构象。
2.2 等变神经网络的三种典型实现路线
真要动手实现一个等变网络，路径不止一条。主流方案我大体归纳为三类。
第一类是基于球谐函数和张量积构造的高阶等变消息传递网络。这是目前表现最强但理解门槛也最高的方向。代表是使用SO3 irreps的TFN框架，后续被NequIP、MACE、Allegro等模型发展壮大。每个隐藏节点的特征由多个不可约表示组成，定义了旋转阶数l，每个l对应一组矩阵和相应的特征通道。消息传递中通过张量积把边方向与节点irreps组合起来，再经过Clebsch-Gordan变换和可学习的径向网络收缩。这类模型组里对张量积的归一化方式非常敏感，特征通道分配也直接决定了参数大小和训练速度。使用者需要认真处理torch.nn.Module里的irreps定义，一个不匹配就会报维度错误，Debug起来相当难受。
第二类是隐式等变模型，不直接在内部构建irreps，而是保证输入输出的坐标行为正确，例如PaiNN。PaiNN节点上既有标量特征，也有显式的3D向量特征。消息传递中利用归一化方向向量和标量特征迭代产生新向量方向，并且保持向量跟随分子旋转而旋转。相比TFN家族的模型，PaiNN在实现上轻量很多，训练速度也快，在中等规模体系上表现非常稳，这是我的路线偏好。
第三类是纯标量特征加基于距离和角度的几何描述。这是严格意义上的不变模型，不是完全等变网络。训练时可以配合随机旋转数据增强弥补一部分旋转泛化损失，但除非任务极简单或者算力受限，否则不太推荐走这个方法，尤其在力和高精度能量预测这类任务上差距很大。
2.3 能量-受力联合输出：自动微分带来的“白嫖”优势
神经网络分子力场的输出往往是多任务的：总能量、原子受力、可选的应力张量、偶极矩、原子化能量等。力和能量之间天然存在梯度关系——力是能量负梯度。通过自动微分框架（比如PyTorch中的torch.autograd.grad）可以同时在一次前向中拿到能量和所有原子力。
这样做的第一重好处是精度一致：不单独训一个“力网络”，力直接继承能量的几何敏感度。第二重好处是信息共享：训练目标同时包含能量与力的监督，模型在某一个构型附近学到的几何敏感性更全面。
当然自动微分也是有成本的。分子的原子个数是N，能量是一个标量，对N组坐标求梯度，反传图占据的内存会增加不少。批处理大体系时，显存压力会明显上升。我的做法是尽量让批量大小控制在32到64之间的中等分子，但如果用大盒子跑几百个原子，批量就得进一步缩水。此外，由于反向传播本身要保留中间激活值，因此推荐在网络设计时尽量控制消息传递层数，例如4到6层就够了，不要盲目堆多层。
3. 训练神经网络分子力场的完整实操流程
3.1 选择基底数据：精度、覆盖度和分布均匀性比数量更重要
没有好的训练数据，再强的模型也没有意义。神经网络分子力场的训练数据一般来自DFT计算或更高级的量子化学方法（比如CCSD(T）校正后的数据）。数据集覆盖的化学空间越广、几何结构越多样，模型的迁移能力越强。
具体来说，第一要筛去难以收敛或虚频过大的结构，保证参考能量可靠。第二要对构象空间做聚类，避免动力学模拟时只覆盖了一小段势能面。第三要尽量包含非平衡态结构。这是很多人容易忽视的：如果你只用优化好的平衡结构训练网络，那么在分子动力学模拟时模型一旦偏离能量极小值区域，就会进入训练数据稀薄地带，给出完全不靠谱的预测。
常见做法是利用主动学习采样。比如先训练一个初版力场，跑几皮秒短程动力学，对轨迹中预测置信度低的结构调用DFT计算，然后把新结构和标签合并进训练集，再重训模型。这个过程循环两三轮，模型覆盖的势能面就会显著扩张。我在实际操作中发现，用多组不同温度（从300K到800K）的MD轨迹数据去扩充训练集，对力场在高温下的稳定性很有帮助，特别是在做热分解机理或者高温反应动力学模拟时。
3.2 数据归一化与输入特征：能量分布处理要小心
一个常见的细节是单位的统一和能量的归一化。DFT计算给出的绝对能量数值很大，而且不同泛函给出能量的基线值不同，直接回归绝对能量往往使模型收敛困难。实操中一般把参考能量减掉孤立原子能量，用原子化能或者“总能量相对参考构型差值”作为训练标签。不过要注意，如果后面要算活化能、反应焓，那么从绝对能量转换过程中任何一个常量偏差都会影响最终结果，所以需要把能量偏移量保存好。计算受力时，力的量纲是能量除以长度，常见单位是Hartree/Bohr或eV/Å，训练前要把能量单位与坐标单位的组合统一。
几何特征输入方面，虽然等变模型内部可以处理坐标，但坐标绝对值不能直接作为模型特征，因为平移等变要求不允许模型记住“绝对位置”。实际做法是边特征使用相对位移向量，节点特征使用元素嵌入（例如用周期表先验构造的初始嵌入），再在消息传递中组合出几何信息。某些模型还额外把周期性盒子的有关参数（晶格向量、体积）作为环境信息输入。对于周期体系，原子间的最近镜像邻居列表需要正确构建，否则边界处的相互作用会被算错。
3.3 模型训练策略：学习率、损失权重和验证划分
神经网络分子力场的损失一般写作：
L = w_E * MSE(E_pred - E_ref) + w_F * MSE(F_pred - F_ref)
力项的拟合难度通常远高于能量项。因为力的绝对值在数值上往往远小于能量，但有效信息量高，结构敏感性更强。参考很多公开发表的实现，力项的权重w_F一般设成1.0到100.0之间，能量项w_E通常在0.1到1.0之间。注意这是损失层面对量级差的补偿，不代表力的重要性只有这么一点。
训练中我习惯用带warm-up的Adam优化器，初始学习率依据模型规模在1e-4到5e-4之间。训练总轮数不一定追求超高epochs，因为模型参数越多，过拟合风险越大。尤其当训练集只有数千构型时，神经网络的容量必须控制住。早期我在小数据集上用过较大模型，验证误差持续下降但测试误差越来越差，这就是过拟合的开始。经验法则：数据量小于1万构型时，模型参数量尽量控制在几百万以内；数据量超过10万构型时，才适合放心使用大模型。
验证集的划分也值得强调。简单随机划分容易造成构型重叠度高，导致验证集误差虚低。更严格的做法是按体系划分，比如把某个分子完整拿出去做测试，考察外推性能。这能真实反映模型对未知化学环境的适应性。在分子动力学模拟和活性预测任务里，真正要关心的恰恰是这些完全没见过的新分子。
3.4 训练现场记录：一个中等规模体系的实证结果
这里以模拟项目X为例，目标是拟合一个包含约120个原子（有机分子加少量溶剂水）的凝聚相体系的能量和受力。数据来源是某稀土泛函中的杂化泛函级别DFT计算，共12800个构型，覆盖378K到620K的分子动力学采样。模型选用PaiNN构型，4层消息传递，节点特征宽128，标量特征与向量特征各占128与128通道。参数量大约250万，批大小16，在单张消费级GPU上训练了约36小时（50个epoch）。
最后的结果：验证集能量MAE约为0.35 kJ/mol（相对差值），受力MAE约为2.1 kJ/mol/Å。这个精度对于MD模拟来说可接受。跑了20ps的等温等压系综，密度、径向分布函数与直接做DFT-MD的结果吻合良好。相比DFT-MD的算力消耗（同体系量子化学方法单步耗时在小时量级），神经网络力场使得每个MD步长（0.5 fs）只需几十毫秒，整体加速约四个数量级。这就是神经网络分子力场真实的实用价值：不是精度完全替代DFT，而是在足够接近DFT精度的前提下，将模拟尺度和时长推进到一两个数量级以上。
4. 常见问题与排查技巧实录
4.1 训练Loss大幅震荡或爆炸
这个我碰到过好几次。典型症状是能量Loss前几十步下降正常，随后突然上涨几个数量级，甚至出现NaN。主要原因通常是坐标特征跨度过大、学习率过高、损失权重失衡，或者激活函数没有处理好数值稳定性。
排查顺序：先检查输入坐标和参考能量是否出现NaN或极端值；再确认损失权重是否过于悬殊；再降低学习率做一次短跑（比如100步）看曲线是否平滑。另外一个容易忽视的问题是对周期性体系，邻居搜索的截断半径若设置不连续，原子在边界处跨越时受力会发生突变，导致模型训练困难。建议截断半径设置时加上平滑衰减函数，并保证截断半径略小于盒子最短边长的一半。
4.2 预测的力方向对、大小偏小
常见原因是训练数据中平衡结构占比过高。力在任何势能面上的极小值附近都接近零，如果模型大量学平衡态结构，就会倾向低估偏离平衡态结构的力。解决方案是增加非平衡构型的采样比重，比如对最优结构进行热扰动、拉伸特定键长、随机旋转二面角后生成样本。
另一种可能是网络容量不够，消息传递层数偏少，导致对局部化学环境的区分能力不足。可以尝试把节点特征宽度增加1.5倍，或者增加一层消息传递，观察验证集力误差是否有明显下降。如果没有明显变化，说明问题不在容量，而在数据分布。
4.3 MD轨迹跑到一半结构崩坏
力场在长时间分子动力学模拟中崩溃是常见的无奈事件。原因往往是模型在轨迹中到达了一个远离训练分布的构象，给出错误预测，“反馈回路”又让系统走向更糟糕的结构。这就是外推失效。
让我给你几个实用处理办法：其一是做不确定性量化——训练多个模型（例如不同随机种子或不同数据子集），跑MD时同时预测多个模型输出的标准差，标准差超出阈值时停止模拟并补充高精度计算。其二是修改MD时间步。神经网络力场通常可以用1 fs甚至2 fs的时间步长（前提是体系不含特别轻的氢原子振动），但如果步长太大导致积分不稳定，宁可回到0.5 fs。其三是引入更保守的体系约束，例如对键长加SHAKE约束，把高频振动模冻结。当然最根本的解决之道还是扩充训练数据覆盖范围，让模型对势能面有更广的认知。
4.4 同一结构不同朝向能量不一致
理论上等变模型不该出现这种情况，但我实测过某款简化实现，确实出现过旋转后能量差了几个毫哈特里。原因排查结果令人意外：模型虽然用等变消息更新了节点特征，但边特征错误地使用了未旋转坐标下的距离，或者使用了某种与全局坐标绑定的归一化。也就是说，模型里存在一个非等变的“泄漏路径”。
排查方法是构造一个简单分子，对同一构象做多组随机旋转，计算能量和受力差异。如果差异超过阈值，检查输入是否真的只用了相对位移，检查模型内部任何归一化层的均值或方差是否依赖坐标绝对值，并检查是否有RemEmbedding这类位置编号被误用为坐标信息。经过严格排查和修复后，能量差异会趋近于机器精度级别（比如1e-6 eV以内）。
4.5 周期性体系受力不连续
对于固体、界面或者带溶剂的周期性盒子，邻居搜索使用最小镜像约定时，截断内部和外部的相互作用需要平滑过渡到零。直接截断会让受力在截断半径处不连续，分子动力学能量守恒被破坏。螺旋周期性条件下，推荐使用带平滑过渡函数的径向截断，并在模型边特征初始计算时就把cutoff效应编码进去。有些模型还在原子间距离矩阵计算后，对超出截断半径的边直接mask掉。但这些mask必须是硬性的，不能造成梯度的跳变。我建议做一次能量守恒测试：在微正则系综中跑几百步，观察总能量漂移。如果漂移超过一定范围，优先检查截断处理是否平滑。
5. 总结与后续扩展方向
说实话，这个领域的变化速度比我预期的还要快。在我刚开始做模拟项目X的时候，等变模型的选择还很少，社区里能参考的实操资料也有限。现在各类开源工具和预训练权重已经相当丰富，新入坑的研究者可以站在更高的起点上。不过核心问题依然没变：如何让模型在足够多化学空间内可靠地复现势能面。
我个人在过去几个月踩过的坑，基本都和数据分布、等变实现的细节以及训练稳定性有关。印象最深的教训是：模型架构再漂亮，训练数据覆盖不足照样会翻车，而覆盖充足但模型不守对称性，照样寸步难行。两者的配合才是这个方向的王道。
后续我会把模拟项目X的扩展方向往两个角度推。第一是研究主动学习采样中对不确定性更精确的估计方法，把训练集中在最有价值的构象上。第二是尝试结合等变扩散模型，将力场输出和构象分布统一起来，也许能为多尺度模拟提供一个更通用的接口。如果你正准备构建自己的神经网络分子力场，建议先从简单的不变模型或者轻量等变模型做起，用上述实验设计和评价方法把流程跑通，再逐步引入更复杂的架构。这个领域有门槛，但过了门槛之后带来的效率提升，确实让人上瘾。