Skip to content

等变神经网络(Equivariant Neural Networks)深度教程

从对称性到力场:完整数学推导、纯 Python 实现,与八个化学/材料实验

这是一份专门写给材料与化学领域科研人员、以及「人工智能 × 化学」交叉方向研究生的等变神经网络系统教程。

在分子动力学与高通量计算化学的日常研究中,每一位研究者都会建立起强烈的物理直觉:无论一个分子在三维空间中如何整体平移或旋转,它的能量、水溶解度、生成焓和能隙等内禀热力学性质绝不会发生丝毫改变;而如果分子整体旋转了一个角度,其受到的原子间相互作用力、分子偶极矩或主应力张量,也必须在空间中精准地跟随旋转相同的角度。然而,当我们将原子坐标送入常规机器学习模型时,这个看似平凡的物理常识却常常遭遇前所未有的破坏。

本教程从"为什么分子转个身,常规模型的预测就会剧烈抖动"这一核心痛点切入,带领读者建立坚实而直观的几何物理图像。我们将深入探讨群与表示论的物理本质,推导三维旋转的 Rodrigues 公式,厘清距离等不变量在立体化学与手性识别上的理论极限,拆解 Clebsch–Gordan 耦合法则 1⊗1 = 0⊕1⊕2 背后对应标量、矢量与二阶张量的本质,并借助 Schur 引理揭示等变双线性映射的数学唯一性。在此基础上,我们将完整推导 Tensor Field Networks (TFN) 的不可约球谐卷积核与 E(n) 等变图网络(EGNN)的坐标更新等变性与质心守恒律,并给出每一层网络反向传播梯度的闭式解析解。

为了彻底消除黑盒工具库带来的认知迷雾,本教程摒弃了 NumPy 等第三方数值库,采用纯原生 Python 构建起一套能够逐行对照数学推导的教学代码,让公式中的每一个求导符号与权重矩阵清晰可见。最后,这套理论将在包含药物分子水溶解度(ESOL)、手性对映体识别、高聚物势能面力场拟合以及三维周期性晶体(PBC)在内的八个真实化学与材料实验中得到系统验证,其中包括两个极具启发意义的"等变模型并未胜出"的真实负结果,以此客观勾勒出等变神经网络的真实应用边界。

封面:转动分子,读数不动

这份教程的定位与视角

当前讨论等变神经网络的技术文献呈现出鲜明的两极分化:在理论一端,原始文献(如 Tensor Field Networks、NequIP、MACE)往往预设读者拥有高能物理或纯粹代数背景,公式推导严谨却缺少化学体系的具体物理映照;在工程一端,开源框架(如 e3nn、PyTorch Geometric)提供了高效实现,却将底层几何变换深度封装于抽象算子之后。本教程致力于在这两极之间架起桥梁:从真实的材料与化学问题出发,将群论推导落到实处,并通过不依赖任何数值黑盒的纯 Python 代码在真实分子与晶体数据上跑通全流程。

教程仅预设读者掌握基本的微积分、线性代数与 Python 编程基础。至于群表示论、Schur 引理、Clebsch–Gordan 系数以及 Wigner-D 矩阵等抽象数学概念,在首次出现时均会结合化学直觉(如原子轨道对称性、角动量耦合与立体化学)从零剖析,并附带可直接运行的数值验证代码。


目录

  1. 写在前面
  2. 背景:为什么"转一下分子,预测就变"是个硬伤
  3. 数学准备(一):群、作用与表示
  4. 数学准备(二):不变量能走多远
  5. 核心机制:张量积与 Clebsch–Gordan 分解
  6. 经典架构:从距离到高阶张量
  7. EGNN:完整推导与稳定性
  8. 反向传播:手写每一行
  9. 纯 Python 实现:公式 ↔ 代码
  10. 化学与材料里的应用图谱
  11. 八个实验
  12. 优缺点与使用建议
  13. 把结论收拢
  14. 附录 A:符号表
  15. 附录 B:参考文献

0. 写在前面

0.1 这本教程想解决什么

许多初涉分子机器学习的研究者往往会带着一个极为朴素的疑问进入这个领域:既然多层感知机(MLP)拥有强大的通用函数拟合能力,为什么我们在处理分子或晶体三维坐标时,不能简单地将原子的笛卡尔坐标展开成一长串一维向量直接塞进全连接网络,再通过在训练集中随机旋转分子做数据增强呢?为什么学术界总在强调"旋转不变性"、"E(3) 等变性"与"不可约表示张量通道"?

这个困惑的答案深植于物理真实性与数学假设空间之中。

从客观事实层面来看,经验性数据增强无法解决几何泛化问题。正如我们在实验三中观察到的真实测量结果:一个在固定朝向上训练至均方误差仅为 0.071 的坐标感知机,仅仅将待测分子在空间中整体旋转 180°,其预测误差便剧增至 0.954,性能骤降十余倍;而内嵌了对称性先验的等变网络在 0° 与 180° 朝向下的预测结果则严丝合缝,数值精度分毫不差。

从物理机制层面剖析,这并非数据量欠缺或正则化不足的工程缺陷,而是常规全连接网络的函数假设空间从根本上违背了客观规律。一个未经几何约束的通用多层感知机能够自由地学出类似"当羰基指向实验室东方时溶解度高、指向北方时则不溶"这样荒谬的假象关系。在真实的物理世界中,孤立体系的能量与宏观热力学性质绝不会与观测者人为选定的空间参考系产生绑定。等变神经网络的核心使命,正是依据第一性原理的对称性要求,在网络架构设计的源头将这类违背客观规律的荒谬函数从假设空间中一劳永逸地剔除。

为了将这一机制彻底讲清,本教程围绕三大支柱展开叙述。首先是层层递进的数学推导,我们不仅直观解释为什么球谐函数能作为角度基底,更会手把手展开计算矢量外积的 1 ⊗ 1 = 0 ⊕ 1 ⊕ 2 正交分解,证明为什么三维欧几里得等变卷积核在数学上必然收敛于"径向函数与 Clebsch–Gordan 投影矩阵的乘积",并解析推导每一层张量操作与坐标迭代的前向与反向闭式解。其次是透明无盲区的代码呈现,我们坚持在核心模块中剥离诸如 NumPy 这样具有隐式广播与黑盒优化的第三方数值库,让矩阵乘法还原为嵌套循环,让自动微分还原为显式链式法则,使读者能够将推导中的偏导符号与代码行一一精准对齐。最后是立足化学本质的实验论证,全书八项实验均运行于经过严格物理检验的数据之上,涵盖水溶解度基准、构象力场能量与受力拟合、手性识别乃至十类常见空间点阵晶体的周期性计算,并在实践中诚实剖析等变模型在某些化学拓扑主导任务中的局限性。

0.2 读者的背景与阅读导引

本教程适合具备基础微积分、线性代数和基本神经网络反向传播常识的材料与化学背景研究人员阅读。化学与材料学背景的研究者天然具备深厚的对称性直觉——在化学课程中学习分子点群、晶体空间群以及原子轨道波函数时所积累的空间想象力,正是理解等变网络的极佳土壤。熟悉量子化学或凝聚态物理的读者可以跳过过于基础的群论概念,将目光聚焦于探讨手性分子空间几何的第 3.4 节以及阐释空间反演宇称特性的第 4.6 节,这两处不仅是物理化学概念与深度学习交互的核心地带,也是初学者极易踩中逻辑陷阱的关键环节。

0.3 文件结构与运行环境

教程所附带的完整工程代码保持了极高的自包含性与可复现性。各模块职能明确,相互印证:

文件说明
等变神经网络教程.md本教程主文件,包含系统数学推导、16 张手绘机制插图与 10 张实验数据图表
code/enn.py核心算法库,涵盖群作用、张量积、EGNN 层、SchNet 层与手写反向传播优化器
code/geom3d.py晶格几何运算、周期性镜像邻接图构建、最小镜像约定及解析力场计算
code/chem3d.py基于 RDKit 的化学分子接口,处理构象生成、通用力场(UFF)标签提取及偶极矩计算
code/runner.py统一训练引擎、评估指标体系与实验数据持久化落盘机制
code/tests_enn.py包含 107 项严格数值断言的测试套件,覆盖逐参数数值梯度检验与等变性体检
code/demo_*.py八个独立的实验运行脚本,均支持单脚本闭环复现
code/make_figures.py数据可视化脚本,将实验产出的原生 JSON 文件绘制为矢量数据图表
figures/存储最终的数据图表文件与各实验对应的原生评测数值
images/存储用于直观阐释物理与网络机制的手绘风格示意图集及生成提示词

读者可以在准备好的 Python 环境中依次执行如下指令,体验从单元测试、快速验证到完整重现的探索过程:

bash
cd code

# 1) 执行底层数学检验:运行 107 项断言与上千项参数的解析梯度检查(耗时约 0.5 秒)
/Users/polyai/anaconda3/envs/deepL/bin/python tests_enn.py

# 2) 以快速模式轻量化体验全流程实验(约 3~5 分钟)
bash run_all.sh --quick

# 3) 执行完整规模的化学与材料基准评测(纯 Python 运算,约 30~40 分钟)
bash run_all.sh

# 4) 将各实验生成的 JSON 原始指标绘制为可视化图表
MPLCONFIGDIR=/tmp/mplcache /Users/polyai/anaconda3/envs/deepL/bin/python make_figures.py

当前运行环境部署于 /Users/polyai/anaconda3/envs/deepL(基于 Python 3.12),内置了 deepchem 2.8.0rdkit 2026.03.6torch 2.14.0numpy 2.5.3。需要说明的是,在该特定测试环境下并未预装重型图学习扩展包(如 dgltorch_geometric),本教程特意选取了零依赖的自研纯 Python 架构与通用的 scikit-learn 基线配合,确保教学演示与推导验证不受任何外部框架版本的束缚。

0.4 关于纯 Python 原生实现的初衷

本教程在核心算法库(包括 enn.pygeom3d.py 以及所有 demo 的模型层)中贯彻了彻底的无 NumPy 准则:三维空间向量直接采用基础列表表达,张量积运算展开为清晰的三重嵌套循环,而反向传播的偏导流动亦全部由显式代码逐项累加。这种取舍并非出于追求极客复古,而是为了实现数学公式与执行代码之间的绝对透明。当我们在理论推导中写下 L/a=Sym(G)btrG3b\partial \mathcal{L}/\partial \mathbf{a} = \mathrm{Sym}(G)\mathbf{b} - \frac{\mathrm{tr}G}{3}\mathbf{b} 时,读者能在代码的 ga_vec = mat_vec(Sym, B[c]) 处一目了然地看到抽象张量导数如何平滑落地,中间没有任何黑盒广播或隐式内存重整的干扰。

在整个工程中,仅有两处合理的工程边界打破了这一纯粹性:其一是与 DeepChem 交互时处理其固有的 ndarray 数据集接口,其二是调用 Matplotlib 渲染实验曲线图,这两处工具性接口均完全被隔离在核心模型运算之外。

这一设计换取教学透明度的直接代价是运行速度:纯 Python 解释执行在 CPU 上的效率比经过矢量化优化的 NumPy 慢一至两个数量级,比 GPU 加速的 PyTorch 慢两至三个数量级。我们在第 8.7 节对实际计算耗时提供了翔实的实测统计与工业化加速路线。若读者的首要目标是开展千万级构象的大规模材料发现项目,可直接对接第 9.9 节介绍的成熟生态库(如 e3nn、NequIP 或 MACE),而本教程提供的透明底层框架则将成为彻底吃透这些大型工具运行机制的坚实基石。

0.5 数据集的物理性质:真实观测与物理建模

在探讨后续实验结论之前,我们需要界定实验数据中不同标签的来源范畴,以确立合理的评估基准。在材料与化学研究中,数据往往分为实验测量所得的宏观性质与微观量子/经验势函数计算所得的物理场量。

教程中的水溶解度数据(ESOL)直接采纳了 MoleculeNet 评测集中的真实湿法实验测得的对数溶解度(logS\log S)。而对于手性分类标签、分子偶极矩向量以及聚合链与空间晶体的能量和原子受力,我们分别基于三维几何手性不变量、RDKit 提供的 Gasteiger 局部电荷积分、普适分子力场(UFF)以及短程屏蔽库仑-范德华对势进行严格解析计算。采用这些解析而明确的微观物理量作为基准,不仅确保了全套教程结果在任何算力环境下都能秒级无损复现,更重要的是它们具备严格遵守力学与电磁学定律的数学特征。无论是密度泛函理论(DFT)、高阶多体微扰论还是经典解析力场,物理量对空间变换的响应规律始终如一——等变网络从不关心标签数据的计算精度层级,它所约束并保护的,是标签在三维空间中作为标量、矢量或高阶张量所必须恪守的几何变换法则。

0.6 针对不同目标的研读路径

读者可以根据自身的研究阶段与时间预算,采取不同的研读策略。

对于希望在两小时内快速掌握核心概念的读者,建议直奔核心主线:通读第 1 章建立直觉,研读第 3.4 节掌握手性对映体对距离不变量的本质挑战,解析第 4.3 节感受张量积耦合的几何美感,再阅读第 6.2 节体会 EGNN 如何以简洁形式维持等变性。随后运行 tests_enn.py 并观察实验一与实验二,便能迅速建立起"空间对称性完全可以精确数值化度量"以及"纯距离模型存在无法逾越的手性盲区"这两个最关键的物理洞见。

对于正在从事分子或材料表征学习的科研人员,建议顺次精读全篇。重点应当放在探讨几何算子唯一性的第 4.4 节与第 5.4 节,这两处从代数上回答了为什么主流等变架构不约而同收敛于球谐卷积形式;继而仔细比对第 7 章反向传播推导与第 8 章的手写实现,甚至可以尝试亲手微调其中一处梯度分支,观察单元测试中梯度检验报错的敏感响应,从而对模型的几何动力学机制获得刻骨铭心的体会。

而对于以工程实现与代码跑通为首要目标的读者,则可以直接执行一键脚本 bash run_all.sh --quick,随后对照第 10 章的各节分析,结合实验生成的 JSON 与图表,逐一考察模型在不同化学任务下的表现与局限。


1. 背景:为什么"转一下分子,预测就变"是个硬伤

1.1 从一个具体的失败说起

设想我们在计算化学实验室中承担了一项常见任务:根据一批小分子的空间三维构象预测它们的水溶解度。我们从化合物库中获取了一批分子的标准 SDF 结构文件,每个分子都记录了一组三维笛卡尔坐标。初入深度学习领域的开发者常常会非常自然地搭建如下全连接神经网络:

y^=MLP(flatten(x1,x2,,xN,z1,,zN)),\hat{y} = \mathrm{MLP}\big(\mathrm{flatten}(x_1, x_2, \dots, x_N, z_1, \dots, z_N)\big),

在这套输入方案中,xiR3x_i\in\mathbb{R}^3 代表第 ii 个原子的空间坐标,而 ziz_i 则标识其原子序数。我们将所有原子的三维坐标与化学特征依次拼接成一长串一维特征张量,直接送入多层感知机中进行有监督回归。

同一个分子,两种朝向,两张完全不同的纸条

在训练集上,庞大的参数容量往往能让模型迅速收敛,在结构相似的验证样本上也表现尚可。然而,当我们将计算流程部署到真实场景时——例如采用不同的构象搜索算法生成了同一分子的另一个等价姿态,或者仅仅在数学上将整个分子的坐标系绕 zz 轴旋转 9090^\circ——令人费解的现象出现了:由于溶液中的分子处于持续的无序热运动中,不同朝向在物理化学本质上属于同一个分子态,但模型的溶解度预测输出却发生了剧烈的跳跃,预测偏差甚至会达到几个数量级。

这一失效绝非单纯增加神经元层数、引入 Dropout 正则化或扩充训练数据所能弥合的浅层缺陷。上述全连接网络在数学上本质上是一个对各输入分量赋予独立权重的无约束映射,它可以毫无阻碍地拟合出"当分子的长轴平行于 xx 轴时判定其易溶、而旋转至平行于 yy 轴时判定其难溶"的病态函数。然而,真实物理世界中并不存在依赖绝对坐标轴朝向的热力学势。空间直角坐标系纯粹是研究人员为了记录原子位置而在白纸上人为划定的数学辅助线,绝非分子内禀性质的一部分。当我们把包含分子与观察者在内的整个物理体系在空间中刚性旋转,真实的物理测量读数必然岿然不动。

1.2 化学体系的刚体几何对称性

将上述化学物理直觉凝练为严格的数学表述:假定一个分子体系由 NN 个原子构成,其空间构型完全由坐标矩阵 X=(x1,,xN)RN×3X = (x_1,\dots,x_N)^\top \in \mathbb{R}^{N\times 3} 所刻画。从牛顿力学与非相对论薛定谔方程出发,以下两类基础坐标变换绝不会改变体系的任何内部相互作用能与平衡态性质:

空间整体平移:xixi+t,空间整体刚性旋转:xiRxi,\text{空间整体平移:}\quad x_i \mapsto x_i + t, \qquad \text{空间整体刚性旋转:}\quad x_i \mapsto R\,x_i,

此处 tR3t\in\mathbb{R}^3 代表任意三维位移矢量,而 RSO(3)R\in SO(3) 则代表保距且保持手征性的特殊正交旋转矩阵(满足正交性 RR=IR^\top R = I 且行列式 detR=+1\det R = +1)。若将三维空间中的任意平移与旋转操作复合,我们便得到了描述三维刚体运动的半直积欧几里得群:

E(3)=R3O(3),E(3) = \mathbb{R}^3 \rtimes O(3),

值得化学与材料研究者高度警惕的是群 O(3)O(3) 与子群 SO(3)SO(3) 的微妙边界。完全正交群 O(3)O(3) 除了包含行列式为 +1+1 的连续纯旋转外,还囊括了行列式为 1-1 的不当旋转(即旋转与镜像反演的复合)。如果我们严格要求排除空间镜像操作,仅允许刚性转动与平移,所对应的结构被称为特殊欧几里得群 SE(3)SE(3)

在基础物理学与无机晶体学中,O(3)O(3) 往往被视作理所当然的空间对称群;但在生物医药与立体化学领域,O(3)O(3)SE(3)SE(3) 的差异却关乎生命现象的核心——即分子的手性(Chirality)。空间镜像并非在真实宏观三维物理空间中可以通过刚体连续转动实现的操作,左右手对映异构体在受体结合与药理活性上往往天差地别。在文献术语中,若一个模型被称为具备 E(3) 等变性,通常意味着它对包含镜像反演在内的所有变换均无差别地响应;而若我们的任务需要敏锐分辨对映异构体的不同性质,就必须将对称性约束明智地限定在保持空间手征性的 SE(3)SE(3)SO(3)SO(3) 之上,并在特征通道中显式引入能够感知镜像翻转的赝标量。

1.3 不变性、等变性与协变性的物理本质

在几何深度学习文献中,不变(Invariant)、等变(Equivariant)与协变(Covariant)经常出现,澄清三者的严谨边界对于构建正确的化学网络至关重要。

假定我们设计了一个计算模型 ff,它接收分子构象 XX 并输出某种性质。当空间变换群 GG(如旋转群 SO(3)SO(3))的元素 gg 作用于分子坐标矩阵形成变换后的构象 gXg\cdot X 时,模型的响应形式直接决定了它的几何物理属性。

如果模型的输出满足 f(gX)=f(X)f(g\cdot X) = f(X),即输出值在分子变换前后恒定不动,我们称该模型具备不变性(Invariance)。化学研究中的能量标量、水溶解度、自由能、形成焓、电子亲和能与最高占有轨道能级等热力学与量子标量,均属于天然的不变量。

如果模型的输出并不保持原样,而是以与输入变换精准同步的代数规律跟随改变,满足 f(gX)=ρ(g)f(X)f(g\cdot X) = \rho(g)\,f(X),其中 ρ(g)\rho(g) 为定义在输出空间上的群表示矩阵,我们称该模型具备等变性(Equivariance)。直观而言,这代表"分子在空间中怎么转,物理输出便按对应的法则怎么转"。分子的永久偶极矩矢量 μR3\mu\in\mathbb{R}^3、各个原子受到的牛顿相互作用力矢量 FiR3F_i\in\mathbb{R}^3 以及体系的激发跃迁偶极矩,都是典型的等变物理量。

至于协变性(Covariance),在多数近代理论文献中已被广泛用作等变性的等价词汇,其在张量分析与连续介质力学中尤其特指物质张量分量随着坐标基底的变换而协同变换的规律(如分子的极化率张量、NMR 化学屏蔽各向异性张量与晶体内部的 Cauchy 应力张量)。为了行文统一起见,本教程通篇采用现代几何学习的标准词汇"等变性",仅在援引历史经典文献术语时说明其协变内涵。

不变量与等变量:数字纹丝不动,箭头跟着转

极其关键的一点在于,不变性绝非等变性的对立面,而恰恰是等变性在平凡表示下的特例。当输出空间的维度退化为一维标量,且群作用表示退化为恒等映射 ρ(g)=1\rho(g)=1 时,等变条件自然恒等为不变条件。这意味着等变性是整个几何体系中更为根本的普适范式,只要我们构建起健全的等变特征传播网络,标量不变量便能通过自然的张量收缩在最后一层唾手可得。

1.4 数据增强难以替代等变性的根本成因

面对常规网络对旋转的脆弱性,材料与化学领域最常见的经验主义尝试便是借助旋转数据增强:通过在训练批次中对分子构象随机乘以旋转矩阵,试图迫使多层感知机在大量样本中自行"悟出"旋转不变的规律。

这种策略诚然能在一定程度上缓解过拟合,但正如我们在后续实验三中展现的定量结果,经过大范围随机旋转增强训练的坐标感知机,在面对未曾涉足的测试构象时,其旋转 180° 的预测误差依然会出现显著抬升。数据增强与内嵌等变性架构之间,横亘着两座不可跨越的鸿沟。

首当其冲的是计算与参数代价的鸿沟。三维空间旋转群 SO(3)SO(3) 是一个三维连续李群流形。在有限的训练采样中,有限离散角度的覆盖充其量只是在连续流形上打下一串孤立的采样点。数据增强迫使模型耗费海量的可学习参数与表征容量,去反复学习一套原本在代数上早已确知的一般几何规律;而真正的等变网络则是通过特殊的函数耦合与几何约束,直接在网络前向结构的拓扑中锁死对称性。不论面对何种连续旋转角度,等变网络的输出偏差在理论上仅受限于计算机底层的浮点舍入极限(通常在双精度下为 101610^{-16} 数量级),这使得网络全部的学习容量得以完全解放,专注挖掘细微复杂的电子结构与化学相互作用。

其次是分布外泛化保证的鸿沟。通过经验训练学到的局部平滑性无法外推至整个紧致流形,更无法应对诸如手性中心反演等非刚性对称断裂;而等变模型在设计上就不需要"见识"过旋转,因为它从数学机理上就天然免疫于任何坐标系的人为摆动。

1.5 等变神经网络在计算化学与材料技术栈中的演进坐标

从计算化学特征工程的发展史来看,研究者表征分子三维几何信息的范式经历了一场由粗入细、从外在投影到内禀对称的深刻进化:

从手工不变量到高阶等变张量

演化谱系的最底层是传统化学信息学与经典机器学习中的手工不变量描述符,例如库仑矩阵(Coulomb Matrix)、Behler–Parrinello 原子中心对称函数以及平滑原子重叠位置谱(SOAP)。这类方法通过高斯展宽、径向衰减与球谐功率谱积分等物理手段,手工将空间三维坐标压缩为一组天然平移与旋转不变的实数向量,再交由核岭回归或随机森林等算法拟合。其特征构造过程固定且缺乏端到端可学习性。

演化谱系的中间层是引入图卷积范式的标量距离图神经网络,其典型代表为深度张量网络(DTNN)与 SchNet。它们不再依赖繁重的人工特征构造,而是将成对原子间的三维笛卡尔距离作为标量,通过连续径向基函数与节点标量特征进行连续滤波卷积。这一层级的模型虽然实现了自适应的特征表征,但信息在原子间传递的媒介始终局限于各向同性的距离标量,无法直接感知键角张角与空间取向。随后的 DimeNet 与 GemNet 等网络在此基础上迈出了一步,显式地将三体键角与四体二面角引入消息传递,极大丰富了局部几何表达力。

演化谱系的现代前沿,则是以张量场网络(Tensor Field Networks)、SE(3)-Transformer、NequIP、Allegro 以及 MACE 为代表的高阶等变张量与多体展开网络。在这一架构范式中,节点间流动的不再仅仅是各向同性的标量数字,而是包含方向矢量、多极矩张量乃至更高阶球谐不可约表示的多通道几何场量。

这一演化长河中孕育了两个颠覆性的技术拐点:其一是将空间方向矢量作为可传递的特征直接引入网络内部,从而使隐层特征摆脱了单一标量的束缚;其二则是利用量子力学中处理角动量耦合的 Clebsch–Gordan 理论,严密实现高阶张量在旋转变换下的闭合耦合。

然而,作为务实的材料与化学科研人员,必须清醒地认识到架构的层级攀升并不必然等同于实际任务指标的全面碾压。越高阶的几何表征往往伴随着越高的计算开销、越复杂的梯度反传与对精细空间构象质量的强烈依赖。在本教程后续的实验五(ESOL 水溶解度测试)中,我们将呈现一个发人深省的客观事实:对于那些主要由二维分子拓扑与局部官能团主导的热力学宏观性质,基于拓扑连通性的简单分子指纹(如 ECFP)搭配随机森林,其表现与高深复杂的 3D 等变网络完全处于同一误差量级。清晰地认识到何时必须使用等变几何模型、何时只需采用轻量级不变量基线,恰恰是每一位成熟的交叉学科研究者最为关键的专业判断力。


2. 数学准备(一):群、作用与表示

2.1 群与对称群:从分子点群到连续欧几里得空间

对于材料与化学领域的研究人员而言,群论绝非冰冷抽象的纯代数,而是早已融入日常分子光谱、晶体衍射与分子轨道理论中的老朋友。在结构化学课程中,我们习惯于借助水分子具备的 C2vC_{2v} 点群、苯分子具备的 D6hD_{6h} 点群、八面体配合物具备的 OhO_h 点群,或是无机晶体中 230 个空间群来简化哈密顿矩阵,对分子振动模进行对称性分类并推导红外与拉曼光谱的选择定则。

从代数定义来看,一个 GG 由一个集合以及定义在其上的二元结合运算构成,要求集合内元素满足运算封闭性、结合律,且存在唯一的单位元与逆元。化学中常见的分立点群由有限个对称操作元素构成;而当我们将目光转向孤立的三维空间时,一个自由游弋在溶剂中的孤立分子或悬浮在真空中的团簇,不再受限于特定外加晶体场的离散对称操作,它身处的是拥有完全连续对称性的三维物理空间。

群作用:平移、旋转、镜像

在这个连续空间中,有几个连续李群扮演着最核心的统治角色。首先是由全体三维正交变换构成的旋转群 SO(3)SO(3),其元素均为行列式为 +1+1 的正交矩阵,代表所有可能的连续刚体旋转。若进一步将包含空间镜像(行列式为 1-1)的操作吸纳进来,便构成了完全三维正交群 O(3)O(3)。将三维实空间中的连续平移群 R3\mathbb{R}^3 与空间变换群 O(3)O(3) 取半直积,便构成了掌管三维绝对刚体位移的欧几里得群 E(3)E(3);若在此基础上排除可能翻转空间手征性的镜像反演,便得到了特殊欧几里得群 SE(3)SE(3)。材料与化学中熟悉的各类晶体空间群与分子分立点群,在数学本质上无非是这些连续大群的离散子群。

理论物理中有一条贯穿始终的黄金定律——Noether 定理指出,物理系统的每一种连续对称性,都必然对应着一条守恒定律。连续平移对称性孕育了动量守恒,而连续旋转对称性孕育了角动量守恒。在分子动力学模拟(MD)中,这种守恒性确保了孤立系统在时间演化下质心静止且不会产生虚假的整体角动量漂移;而在几何深度学习中,对称性同样意味着神圣的约束——它在数学层面上将神经网络漫无边际的函数搜索空间严密锁死在符合物理客观律的子空间之中。

2.2 旋转的显式写法:Rodrigues 公式的完整推导

在动手编写任何几何算法代码之前,我们必须首先解决一个底层基础问题:如何在计算机中用数值解析式精确表达一个围绕任意空间轴旋转的矩阵。假定给定了一个单位空间方向矢量 k^=(kx,ky,kz)\hat{k}=(k_x,k_y,k_z)^\top(满足 k^=1\|\hat{k}\|=1)以及一个标量旋转角 θ\theta,我们的目标是求出对应的 3×33\times3 旋转矩阵 R(θ)R(\theta)

推导的出发点建立在旋转变换的两条最基本物理性质上。首先,位于旋转轴本身的任意矢量在旋转过程中必须保持绝对静止,即 R(θ)k^=k^R(\theta)\hat{k}=\hat{k}。其次,刚体旋转必须严格保持任意矢量之间的长度与内积夹角不变,这意味着 R(θ)R(\theta) 是一个保距的正交矩阵;进一步地,由于物理旋转是从单位阵 R(0)=IR(0)=I 连续演变而来,其行列式恒等于 +1+1,且旋转运算满足单参数连续群的加法性质:

R(θ)=R(θ)=R(θ)1,dRdθθ=0=[k^]×,R(\theta)^\top = R(-\theta) = R(\theta)^{-1},\qquad \left.\frac{dR}{d\theta}\right|_{\theta=0} = [\hat{k}]_\times,

在此处,[k^]×[\hat{k}]_\times 被定义为与三维叉乘相对应的反对称矩阵(即李代数 so(3)\mathfrak{so}(3) 的基底形式):

[k^]×=(0kzkykz0kxkykx0),[k^]×v=k^×v.[\hat{k}]_\times = \begin{pmatrix} 0 & -k_z & k_y\\ k_z & 0 & -k_x \\ -k_y & k_x & 0\end{pmatrix}, \qquad [\hat{k}]_\times v = \hat{k}\times v .

这个导数之所以精确等于反对称矩阵 [k^]×[\hat{k}]_\times,具有非常直观的物理力学图像:当旋转角 θ\theta 无限微小时,一个空间矢量 vv 的旋转位移本质上是沿着与旋转轴 k^\hat{k} 及矢量 vv 同时垂直的切向方向迈出的一小步,而这一微元切向线位移在矢量几何中恰恰精确等于 θ(k^×v)\theta(\hat{k}\times v)

借助单参数李群的微分几何性质,旋转矩阵满足如下常系数一阶矩阵线性微分方程:

dRdθ=[k^]×R(θ),R(0)=I,\frac{dR}{d\theta} = [\hat{k}]_\times R(\theta),\qquad R(0)=I,

形式上,它与我们在标量微积分中求解 r=arr' = a r 完全相同,其数学解析解可以严谨地写作反对称矩阵的李代数指数映射:

R(θ)=exp(θ[k^]×)=I+θ[k^]×+θ22[k^]×2+θ36[k^]×3+.R(\theta) = \exp\big(\theta[\hat{k}]_\times\big) = I + \theta[\hat{k}]_\times + \frac{\theta^2}{2}[\hat{k}]_\times^2 + \frac{\theta^3}{6}[\hat{k}]_\times^3 + \cdots .

为了将这个无穷矩阵级数闭式求和,我们来考察反对称矩阵 [k^]×[\hat{k}]_\times 的各阶代数幂次规律。利用三维矢量双重叉乘展开恒等式 a×(a×v)=a(av)a2va\times(a\times v) = a(a\cdot v)-\|a\|^2 v,并在单位向量 k^=1\|\hat{k}\|=1 条件下展开,可以直接计算出它的幂次循环性质:

[k^]×2=k^k^I,[k^]×3=[k^]×.[\hat{k}]_\times^2 = \hat{k}\hat{k}^\top - I,\qquad [\hat{k}]_\times^3 = -[\hat{k}]_\times .

这一规律极为奇妙地表明,三阶及更高阶的幂次均可循环折叠回一阶与二阶形式。将其代入泰勒展开式中,并将奇数项与偶数项分别合并,我们便得到了惊艳的三角函数重组:

R(θ)=I+(θθ33!+θ55!)[k^]×+(θ22!θ44!+)[k^]×2=I+sinθ[k^]×+(1cosθ)[k^]×2.\begin{aligned} R(\theta) &= I + \Big(\theta - \frac{\theta^3}{3!} + \frac{\theta^5}{5!} - \cdots\Big)[\hat{k}]_\times + \Big(\frac{\theta^2}{2!} - \frac{\theta^4}{4!} + \cdots\Big)[\hat{k}]_\times^2 \\ &= I + \sin\theta\,[\hat{k}]_\times + (1-\cos\theta)\,[\hat{k}]_\times^2 . \end{aligned}

这便是力学与图形学中著名的 Rodrigues 旋转公式。将其逐项代入三维矩阵分量中展开,便可以得到可以直接编码的显式代数矩阵:

R(θ)=(c+kx2CkxkyCkzskxkzC+kyskykxC+kzsc+ky2CkykzCkxskzkxCkyskzkyC+kxsc+kz2C),c=cosθs=sinθC=1cosθR(\theta)= \begin{pmatrix} c+k_x^2C & k_xk_yC-k_zs & k_xk_zC+k_ys\\ k_yk_xC+k_zs & c+k_y^2C & k_yk_zC-k_xs\\ k_zk_xC-k_ys & k_zk_yC+k_xs & c+k_z^2C \end{pmatrix},\quad \begin{aligned}c&=\cos\theta\\ s&=\sin\theta\\ C&=1-\cos\theta\end{aligned}

在底层的工程实现中,这套严谨的代数映射仅需数行代码即可实现(见 enn.rotation_matrix)。需要特别提醒的是,在测试神经网络的几何等变性时,必须使用基于四元数均匀采样的随机旋转方法(见 enn.random_rotation)生成各项同性的旋转矩阵,而绝不能偷懒地只围绕单一 zz 轴进行旋转检验,否则极易掩盖沿对角线或横向轴发生的高阶张量实现缺陷。

2.3 群表示:原子轨道与张量特征分类的物理直觉

当我们将抽象的旋转群操作作用在具体的物理对象上时,**群表示(Group Representation)**的概念便应运而生。群表示是将抽象群元 gGg \in G 映射为作用于特定矢量空间 VV 上的可逆线性变换矩阵 ρ(g)GL(V)\rho(g) \in GL(V),并且保持群结构的同态性,满足 ρ(g1g2)=ρ(g1)ρ(g2)\rho(g_1g_2)=\rho(g_1)\rho(g_2)ρ(e)=I\rho(e)=I。直白地说,群表示赋予了我们精确描述"某种特定的物理特征在几何变换下究竟如何形变"的数学语言。

对于化学与材料研究者而言,理解等变神经网络的表示论不仅毫无门槛,而且能够唤醒深厚的专业共鸣。我们在大一无机化学与量子力学中早已熟记于心的原子轨道角动量量子数 ll,正是三维旋转群 SO(3)SO(3) 不可约表示最经典的物理实体化载体:

当我们讨论标量特征(如分子的体系能量、基态电荷分布、原子核电荷数或范德华半径)时,这些量在空间旋转下分毫不动。它对应的正是标量表示(平凡表示)ρ(g)=1\rho(g)=1,其空间维度为 11。在物理化学中,这完完全全对应着球对称、无节面的 ss 原子轨道(l=0l=0

当我们讨论三维矢量特征(如原子的瞬时笛卡尔位移、受到的牛顿相互作用力 FF 以及分子的第一电偶极矩 μ\mu)时,这些量在空间旋转矩阵 RR 作用下按照 vRv\vec{v}\mapsto R\vec{v} 协同旋转。它对应的正是维度为 33 的标准矢量表示。在物理化学中,这完完全全对应着相互正交的三重简并 pp 原子轨道(px,py,pzp_x, p_y, p_zl=1l=1

当我们进一步进阶到更高阶的物理张量(如分子的电子极化率张量 α\boldsymbol{\alpha}、核磁共振中的化学位移各向异性张量、分子四极矩以及固体材料中的晶体 Cauchy 应力张量)时,这些物理张量在空间旋转下的变换法则体现为 TRTRT \mapsto R T R^\top。若将其中具有球对称性的各向同性迹项(即标量部分)剥离,剩下的对称无迹二阶张量恰好拥有 5 个独立的自由度分量。在物理化学中,这严丝合缝地对应着五重简并的 dd 原子轨道(dz2,dx2y2,dxy,dxz,dyzd_{z^2}, d_{x^2-y^2}, d_{xy}, d_{xz}, d_{yz}l=2l=2

从表示论的严谨分类出发,旋转群 SO(3)SO(3) 的每一个不可约表示(Irreducible Representation,简称为 irrep)均由角动量阶数 l=0,1,2,l = 0, 1, 2, \dots 唯一定义,且每一个 ll 阶不可约表示空间的固有维度严格满足:

dimVl=2l+1,l=0,1,2,\dim V_l = 2l+1,\qquad l = 0,1,2,\dots

这深刻解释了为什么在等变神经网络中,我们绝不能将所有特征不加区分地一并塞入常规向量数组:网络的每一层特征通道都必须严格指明其所归属的角动量阶数 ll,因为标量(l=0l=0)、矢量(l=1l=1)与高阶无迹张量(l=2l=2)在三维旋转变换下的几何演化规律截然不同

三类特征:标量、向量、二阶张量

2.4 直和、不可约表示与球谐函数展开

在实际搭建神经网络时,我们通常希望隐层神经元既能记忆整体热力学状态(标量),又能传递局部相互作用方向(矢量),甚至感知局域环境的各向异性变形(二阶张量)。在群论框架下,这通过将不同阶数的不可约表示空间并置拼接来实现,称为直和表示(Direct Sum Representation)

V=V0V1V2V = V_0 \oplus V_1 \oplus V_2 \oplus\cdots

在化学类比中,这就像是将 ss 轨道、pp 轨道与 dd 轨道杂化混合,构成一个多级角动量特征池。网络结构中常见的通道配置如 010 \oplus 1 代表一段标量拼接一段矢量特征,而 0120 \oplus 1 \oplus 2 则代表涵盖标量、矢量与二阶张量的完整几何通道。

要把定义在三维球面上的空间取向信息转化为标准的不可约表示特征,最完美的数学工具莫过于实球谐函数(Real Spherical Harmonics) Yl:R3R2l+1Y_l:\mathbb{R}^3\to\mathbb{R}^{2l+1}。球谐函数的本质定义特征在于:当其自变量空间位置经过刚性旋转矩阵 RR 作用时,输出的 2l+12l+1 维函数向量在空间中以线性变换精确保持闭合:

Yl(Rr)=Dl(R)Yl(r),Y_l(R\,r) = D_l(R)\,Y_l(r),

这里的变换矩阵 Dl(R)D_l(R) 即为量子力学与多体微扰论中极为著名的 (2l+1)×(2l+1)(2l+1)\times(2l+1)Wigner-D 矩阵。对标量阶数 l=0l=0,球谐函数退化为常数且 D0(R)=1D_0(R)=1;对矢量阶数 l=1l=1,球谐函数的三个分量正比于沿球面的单位方向矢量本身:

Y1,m(r)1r(x,y,z)m,m=1,0,1.Y_{1,m}(r) \propto \frac{1}{\|r\|}(x,y,z)_m,\qquad m=-1,0,1 .

此时的 Wigner-D 矩阵精确还原为常规的三维笛卡尔旋转矩阵 D1(R)=RD_1(R)=R。这意味着,在基于方向的等变神经网络中,原子间的单位连心矢量 r^ij=(xixj)/xixj\hat{r}_{ij} = (x_i-x_j)/\|x_i-x_j\| 就是大自然免费馈赠给我们的标准 l=1l=1 阶球谐不可约表示输入。

球谐函数:球面上的花纹随球一起转

2.5 等变性的严格数学定义与推论

有了表示论作为基石,我们便可以给出等变映射的严密代数定义。

设群 GG 分别以群表示 ρX\rho_{\mathcal{X}}ρY\rho_{\mathcal{Y}} 作用于输入特征空间 X\mathcal{X} 与输出特征空间 Y\mathcal{Y},若映射算子 f:XYf:\mathcal{X}\to\mathcal{Y} 对群流形内的任意元素 gGg\in G 与输入空间内的任意状态 xXx\in\mathcal{X} 恒满足:

f(ρX(g)x)=ρY(g)f(x),f\big(\rho_{\mathcal{X}}(g)\,x\big) = \rho_{\mathcal{Y}}(g)\,f(x),

则称算子 ff 在群 GG 作用下是等变的。当输出空间上的表示为平凡单位表示(即 ρY(g)=1\rho_{\mathcal{Y}}(g)=1)时,该定义直接退化为不变性条件。

在这个形式化体系中,有三个推论对深度架构的搭建至关重要:

其一是等变算子的复合封闭性。假定两层网络结构 ffhh 分别具备等变性,则其前向级联复合函数 hfh\circ f 必然全局等变:

(hf)(ρX(g)x)=h(f(ρX(g)x))=h(ρM(g)f(x))=ρY(g)(hf)(x).(h\circ f)\big(\rho_{\mathcal{X}}(g)\,x\big) = h\Big(f\big(\rho_{\mathcal{X}}(g)\,x\big)\Big) = h\Big(\rho_{\mathcal{M}}(g)\,f(x)\Big) = \rho_{\mathcal{Y}}(g)\,(h\circ f)(x).

这条性质为深度等变网络的堆叠提供了理论通行证:我们只需精心确保网络内部每一个微小构建块(如卷积核、聚合池化、通道混合)满足局部的等变约束,整座高深复杂的多层神经网络大厦便天然且严格地满足全局等变,绝不需要在网络最终输出端进行任何修补。

其二是等变张量模长向标量不变量的降维投射。对于任何遵循矢量旋转变换的 l=1l=1 特征 vRvv \mapsto Rv,由于正交矩阵严格保距,其欧几里得范数必然恒等不变:Rv2=vRRv=vv=v2\|Rv\|^2 = v^\top R^\top R v = v^\top v = \|v\|^2。这一推论表明,在网络深层,我们随时可以通过计算高阶张量的内部模长或互内积,合法且无损地将等变几何特征坍缩为旋转不变的标量特征,作为最终能量或热力学性质的读出头。

其三是平凡表示特征的不变本质。若一个特征在所有群操作作用下均严格自恒等,则它必然驻留在平凡表示空间中。这一判据直接指引了早期的手工不变量特征工程——如果一个描述符试图作为纯粹的能量输入,它必须由且仅由群不变量构建而成。

2.6 等变神经网络的通用架构范式

将上述几节的几何骨架串联起来,无论学术界提出的前沿架构外观多么千变万化,几乎所有等变神经网络都可以统一归纳为如下通用流水线:

空间几何输入  基底投影  等变不可约特征  (012)  等变层堆叠×L  几何表征  不变量收缩  物理读出.\text{空间几何输入} \;\xrightarrow{\quad\text{基底投影}\quad}\; \text{等变不可约特征}\;\big(0\oplus1\oplus2\oplus\cdots\big) \;\xrightarrow{\quad\text{等变层堆叠} \times L\quad}\; \text{几何表征} \;\xrightarrow{\quad\text{不变量收缩}\quad}\; \text{物理读出}.

这一流水线由三个关键环节紧密扣合:

在输入投影端,模型利用分子几何的相对位移消除平移自由度,提取成对原子间距离构建 l=0l=0 标量通道,并利用单位方向向量与高阶球谐基底提取 l=1l=1l=2l=2 的初始等变通道;

在层间演进端,模型严守几何运算规则,仅允许标量加权、按不可约表示分块的线性映射、以及能够维持空间旋转变换闭合的 Clebsch–Gordan 张量积操作进行特征交互;

在末端读出端,模型根据物理预测目标灵活应对——若目标为力或偶极矩等空间场,则直接由等变特征头输出;若目标为能量或生成自由能等标量,则通过内积收缩或模长汇聚转化为纯粹不变量后求和输出。


3. 数学准备(二):不变量能走多远

在深入探索高阶等变张量的复杂机制之前,一个务实且极富启发性的问题摆在每一位化学研究者面前:既然我们在绝大多数材料与药物计算中,最关心的目标终归是能量、结合自由能、带隙与溶解度这些旋转不变的标量,为什么我们不干脆将所有几何构象信息在输入端全部压缩成"旋转平移不变量",然后直接喂给高度成熟、训练极其轻快的传统神经网络呢?

这绝非空想。事实上,在整个 2010 年代,以库仑矩阵、Behler–Parrinello 对称函数和 SOAP 描述符为代表的"不变量表征流派"几乎统治了计算化学与材料建模领域。这一方案拥有清晰的数学根基,却也在立体化学的前沿遭遇了一道难以逾越的天花板。本章将带领读者用严格的几何代数将这两面边界彻底厘清。

3.1 平移不变性等价于局域相对坐标依赖

我们首先从最简单的平移不变性切入。

命题(平移不变量的相对表述):设映射函数 f:RN×3Rf:\mathbb{R}^{N\times3}\to\mathbb{R} 对任意整体平移变换具备不变性,即对任意位移矢量 tR3t\in\mathbb{R}^3,恒有 f(X+1t)=f(X)f(X + \mathbf{1}t^\top) = f(X) 成立,则必然存在一个等价映射函数 f~\tilde f,使得 ff 可以完全表示为仅依赖于质点间相对位移的形式:

f(X)=f~({xix1}i=2N).f(X) = \tilde f\big(\{x_i - x_1\}_{i=2}^N\big).

数学推导:由于上述等式对三维空间中的任意平移矢量 tt 均普遍成立,我们不妨特意取 t=x1t = -x_1。此时,整个分子坐标矩阵 X+1tX + \mathbf{1}t^\top 中的每一个原子坐标均演变为相对于第一个原子的相对位移 xix1x_i - x_1,而第一个原子的自身相对位置自然退化为坐标原点 00。基于平移不变性,我们直接得出:

f(X)=f(X1x1)=f(0,x2x1,,xNx1).f(X) = f(X - \mathbf{1} x_1^\top) = f\big(0, x_2-x_1,\dots,x_N-x_1\big).

只需将等号右侧仅依赖于 N1N-1 个相对位移的函数形式重命名为 f~\tilde f,命题即证完毕。\square

这个结论从数学上彻底证明了为什么在几乎所有三维分子表征模型中,底层特征工程从不直接使用孤立的绝对坐标 xix_i,而是从相对位移差 rij=xixjr_{ij} = x_i - x_j 开始构建。它同时揭示了初学者常采用的工程小技巧(如在送入模型前将分子整体平移使质心归零)的局限性:质心归零固然在工程上顺手满足了平移不变性,但对于随后更为复杂的空间旋转,平移归零则完全无能为力。

3.2 旋转不变性等价于内积与距离矩阵

在消除了平移自由度后,我们进一步考察刚体旋转变换。首先从最简单的双原子相对体系出发。

命题(旋转不变量的内积完备性):设双矢量标量函数 g:R3×R3Rg:\mathbb{R}^3\times\mathbb{R}^3\to\mathbb{R} 在任意刚体旋转下具备不变性,即对任意三维旋转矩阵 RSO(3)R\in SO(3),恒有 g(Ra,Rb)=g(a,b)g(Ra,Rb)=g(a,b) 成立,则函数 g(a,b)g(a,b) 在数学上必然且只能表达为三维内积量 a2,b2\|a\|^2,\|b\|^2 以及 aba\cdot b 的多元连续函数。

几何证明:为了确立这一性质,我们需要依赖经典的刚体几何基本事实:在三维欧几里得空间中,对于任意两组具有相同基数的矢量集合 {ai}\{a_i\}{ai}\{a_i'\},当且仅当它们两两之间的几何内积完全对应相等(即 ai,aj=ai,aj\langle a_i,a_j\rangle=\langle a_i',a_j'\rangle)时,必然存在一个全局正交矩阵 RR 实现全等空间刚体对齐 ai=Raia_i'=Ra_i。其构造过程只需在两组矢量中分别选出极大线性无关组,利用 Gram–Schmidt 正交化建立标准正交基对应,并借助空间正交补完成全空间扩展。

基于此几何公理,假定存在两组输入 (a,b)(a,b)(a,b)(a',b'),只要它们两两之间的长度与互内积完全相同,根据构造必存在正交矩阵 RR 使得 (a,b)=(Ra,Rb)(a',b')=(Ra,Rb)。代入旋转不变性假设,即刻得出 g(a,b)=g(a,b)g(a,b)=g(a',b')。这证明了函数 gg 的取值在所有保持内积相等的几何等价轨道上恒为常数,因此它只能被内积组 a2,b2,ab\|a\|^2,\|b\|^2,a\cdot b 所唯一函数化。\square

将这一论证直接拓展到包含 NN 个原子的多体复杂体系,我们便自然得到了大小为 N×NN\times N空间 Gram 矩阵 Gij=xixjG_{ij} = x_i\cdot x_j。所有满足刚体旋转不变性的物理预测函数,其有效输入的数学信息等价类完全被 Gram 矩阵所囊括。若进一步结合平移不变性要求,以相对位移内积代替绝对坐标内积,这组几何等价类便直接精炼为化学中无处不在的分子成对距离矩阵(Distance Matrix)

Dij=xixj.D_{ij} = \|x_i-x_j\| .

3.3 距离矩阵与 Gram 矩阵的信息等价性

上述结论常常给初学者带来一种直觉上的错觉,误以为距离矩阵似乎只是空间几何结构的一种粗糙低维投影。事实并非如此,在经典多维标度分析(MDS)框架下,成对距离矩阵与相对几何内积在刚体意义下携带的信息是完全等价的。

命题(距离矩阵与去心 Gram 矩阵的双向恒等转换):设 XRN×3X \in \mathbb{R}^{N\times 3} 为分子坐标矩阵,G=XXG = XX^\top 为对应的 Gram 内积矩阵,DD 为成对欧几里得距离矩阵。则成对距离平方可以直接由 Gram 矩阵分量线性展开:

Dij2=Gii+Gjj2Gij.D_{ij}^2 = G_{ii} + G_{jj} - 2G_{ij}.

反之,若给定已知的全部分子距离矩阵 DD,引入经典中心化投影矩阵 H=I1N11H = I - \frac{1}{N}\mathbf{1}\mathbf{1}^\top,则去心后的 Gram 矩阵 G~\tilde G 可以直接通过距离平方矩阵的双边投影逆向精确恢复:

G~=12HD2H,其中 (D2)ij=Dij2\tilde G = -\frac{1}{2} H D^{\odot 2} H,\qquad \text{其中 } (D^{\odot2})_{ij}=D_{ij}^2

恢复出的 G~\tilde G 与对原始坐标做质心归零后计算出的内积矩阵 YYY Y^\top(其中 Y=XHY = XH)分毫不差。

解析推导:第一个方向直接利用向量欧几里得范数的代数展开:xixj2=xixi+xjxj2xixj=Gii+Gjj2Gij\|x_i-x_j\|^2 = x_i\cdot x_i + x_j\cdot x_j - 2x_i\cdot x_j = G_{ii}+G_{jj}-2G_{ij},关系显然成立。

反向恢复的推导则是距离几何的核心支柱。由于去心坐标矩阵 Y=XHY = XH 的各列矢量和为零(iyi=0\sum_i y_i = 0),成对距离在中心化坐标系中保持不变:Dij2=yiyj2=yi2+yj22yiyjD_{ij}^2 = \|y_i-y_j\|^2 = \|y_i\|^2+\|y_j\|^2-2y_i\cdot y_j。我们将双边中心化投影矩阵展开为显式求和形式:

G~ij=12(Dij21NkDik21NkDkj2+1N2k,lDkl2).\tilde G_{ij} = -\frac{1}{2}\Big(D_{ij}^2 - \frac{1}{N}\sum_k D_{ik}^2 - \frac{1}{N}\sum_k D_{kj}^2 + \frac{1}{N^2}\sum_{k,l}D_{kl}^2\Big).

Dij2D_{ij}^2 的展开式逐项代入上式括号中,可以观察到极为精妙的代数对消:所有包含 yi2\|y_i\|^2yj2\|y_j\|^2 的纯范数项,在均值减法与全局中心化求和中全部抵消为零,唯一保留下来的正是交叉项 2yiyj-2y_i\cdot y_j1/2-1/2 乘子下的还原,最终严谨得出 G~ij=yiyj\tilde G_{ij} = y_i\cdot y_j\square

这个经典的数学事实深刻揭示了:对于一个不考虑镜像反演的孤立几何体系,全部分子成对距离矩阵已经完整囊括了该体系在刚体平移与旋转意义下的全部空间构型几何信息。这意味着,单纯对于以能量预测为代表的标量回归任务,只要不需要辨识手征性,基于成对距离的不变量模型在理论表达能力上已经具备了坚实的充分性。

3.4 不变量的信息边界:手性分子的理论盲区

然而,致命的例外恰恰发生在材料与药物化学的核心地带——手性(Chirality)

虽然中心化 Gram 矩阵 G~\tilde G 在数学上可以由距离矩阵无损还原,并且由于三维物理空间的约束,矩阵 G~\tilde G 的代数秩最大为 3,必然可以通过正交特征分解还原出一组三维坐标解 XX。然而,这种还原在空间正交变换意义下是不唯一的。假定 XX 是与 G~\tilde G 相匹配的一组有效三维笛卡尔解,若我们选取任意行列式为 1-1 的正交反射镜像矩阵 OO(3)O \in O(3),其镜像坐标构型 Xmirror=XOX_{\text{mirror}} = X O^\top 将生成一模一样的 Gram 矩阵与距离矩阵:

XmirrorXmirror=(XO)(XO)=X(OO)X=XX=G~.X_{\text{mirror}} X_{\text{mirror}}^\top = (XO^\top)(XO^\top)^\top = X\,(O^\top O)\,X^\top = XX^\top = \tilde G .

在化学与分子生物学中,这正是**立体对映异构体(Enantiomers)**的几何写照。互为镜像反演的左旋与右旋分子(如 L-丙氨酸与 D-丙氨酸,或是沙利度胺的两种构型),其内部所有成对原子间的空间距离完全一致,原子距离多重集重合度高达 100%。

于是,我们不可避免地推导出几何机器学习领域最重要的一条不可能定理:

理论极限公理:在三维欧几里得空间中,不存在任何仅依赖于成对原子间距离(或任意标量内积函数)的计算架构,能够区分一对严格互为镜像反演的对映异构体。

这绝非网络参数量不足或优化算法不收敛的训练问题,而是底层的数学表征在输入端就遭遇了严格的信息湮灭。在后续实验二中,我们将对这一论断进行严苛的实验检验:将天然 L-丙氨酸与合成 D-丙氨酸输入基于纯距离构建的 SchNet 模型中,两者在浮点精度下的预测响应完全重合,分类准确率严格钉死在与随机瞎猜无异的 50.0%。

手性:距离模型分不出左右手

那么,化学家在判定有机分子立体手性中心(R/SR/S 构型)时所依赖的几何信息,究竟是如何超越距离矩阵的?答案在于引入空间赝标量(Pseudoscalar)

在立体化学的 Cahn–Ingold–Prelog (CIP) 优先规则中,我们围绕一个手性碳原子,依次选取其四个配位原子中的三个特征空间键矢量 u1,u2,u3R3u_1, u_2, u_3 \in \mathbb{R}^3。此时,表征这三个矢量所夹空间有向体积的数学工具,正是标量三重积(混合积)

χ=(u1×u2)u3.\chi = (u_1\times u_2)\cdot u_3 .

我们来考察这个物理量在不同几何操作下的变换法则。当物理体系在三维空间中经历正常刚体旋转 RSO(3)R \in SO(3) 时,由于三维空间中叉乘具有等变性(Ru1×Ru2=R(u1×u2)Ru_1 \times Ru_2 = R(u_1\times u_2)),三重积保持绝对不变:

χrot=(Ru1×Ru2)Ru3=R(u1×u2)Ru3=(u1×u2)u3=χ.\chi_{\text{rot}} = (Ru_1\times Ru_2)\cdot Ru_3 = R(u_1\times u_2)\cdot Ru_3 = (u_1\times u_2)\cdot u_3 = \chi.

然而,当物理体系遭遇空间镜像反演操作 MO(3)M \in O(3)(满足 detM=1\det M = -1)时,由于叉乘本质上是一个伴随赝矢量的构造,反演变换将额外激发出一个因变量符号:

χmirror=(Mu1×Mu2)Mu3=det(M)(u1×u2)u3=χ.\chi_{\text{mirror}} = (Mu_1\times Mu_2)\cdot Mu_3 = \det(M)\,(u_1\times u_2)\cdot u_3 = -\chi .

这种在纯旋转下保持不变、而在空间镜像反演下严格反号的物理代数量,被称为赝标量(或轴标量)。在粒子物理中,它对应着宇称反演破坏;而在计算化学中,它正是刻画手性中心手征方向的唯一微观几何源泉。一个等变神经网络是否具备识别立体化学对映体的能力,完全取决于其隐层通道与消息聚合过程中,是否显式开辟了容纳赝标量流动并对其赋予奇宇称保护的特征通道。

3.5 经典局部环境描述符:不变量范式的巅峰与启示

在"仅使用不变量"的表征哲学指引下,计算物理与量子化学界在深度学习爆发之前便发展出了一整套高度精致的局域几何描述符。它们遵循着相同的核心范式:围绕每一个目标原子中心,通过精巧的数学构造将局域邻居几何压缩为一组满足旋转与平移不变性的实数向量,再交给核岭回归(KRR)、高斯过程或随机森林拟合。

**库仑矩阵(Coulomb Matrix)**由 Rupp 等人于 2012 年提出,其构造直接汲取了量子化学全电子哈密顿量中原子核间静电斥力项的物理灵感:

Mij={12Zi2.4i=jZiZjrijijM_{ij} = \begin{cases} \frac{1}{2}Z_i^{2.4} & i=j\\ \frac{Z_iZ_j}{r_{ij}} & i\ne j\end{cases}

对角线上的项基于多项式唯象拟合了孤立原子的中性自由原子势能,非对角项则直接对应不同核电荷 Zi,ZjZ_i, Z_j 在距离 rijr_{ij} 下的库仑作用。库仑矩阵天然具备三维旋转与平移不变性,但它面临着严重的原子置换排列敏感性——分子的原子编号一旦人为重排,矩阵的行列布局便会彻底错位。在工程实践中往往必须借助矩阵特征值谱或行范数降序排序来勉强维持置换不变性。

**Behler–Parrinello 原子中心对称函数(ACSF)**则代表了材料多体势领域的另一座里程碑。它将局域相互作用分解为径向两体函数与角向三体函数的线性求和组合:

Girad=jieη(rijrs)2fc(rij),Giang=j<k21ζ(1+cosθijk)ζeη(rij2+rik2+rjk2)fc(rij)fc(rik)fc(rjk),G_i^{\text{rad}} = \sum_{j\ne i} e^{-\eta(r_{ij}-r_s)^2}f_c(r_{ij}),\qquad G_i^{\text{ang}} = \sum_{j<k}2^{1-\zeta}\big(1+\cos\theta_{ijk}\big)^{\zeta} e^{-\eta(r_{ij}^2+r_{ik}^2+r_{jk}^2)}f_c(r_{ij})f_c(r_{ik})f_c(r_{jk}),

其中 fc(r)f_c(r) 是一段在特定截断半径处平滑衰减至零的紧支撑截断函数。径向对称函数 GiradG_i^{\text{rad}} 借助多组高斯基底中心 rsr_s 统计"不同距离处原子的分布密度",而角向对称函数 GiangG_i^{\text{ang}} 则精细刻画了"配位原子围绕中心原子的特征张开键角 θijk\theta_{ijk}"。这套人工特征不仅满足了旋转、平移与排列的全对称性,而且成功驱动了高精度神经网络势函数在固态硅、碳及水溶液模拟中的广泛应用。

**平滑原子重叠位置谱(SOAP)**由 Bartók 与 Csányi 等人提出,它将中心原子周围的高斯弥散原子密度直接在正交径向基底与球谐函数展开面上投影,进而构造出旋转不变的功率谱(Power Spectrum):

pi,nnl=12l+1m=ll(cinlm)cinlm,p_{i,nn'l} = \frac{1}{\sqrt{2l+1}}\sum_{m=-l}^{l}\big(c_{inlm}\big)^*\,c_{in'lm},

其中 cinlmc_{inlm} 为局域原子密度在 ll 阶实球谐函数上的积分展开系数。分析 SOAP 描述符的数学结构会发现,其在特征层面上通过系数自身共轭相乘(ccc^* c)求和,实际上已经在萌芽状态实现了群表示论中的张量收缩。

三类几何特征:距离、角度、赝标量

然而,正因为 SOAP 在定义中直接在磁量子数 mm 上实施了均匀求和收缩,它将包含空间取向的所有相位信息在输入构建的第一步就全部"各向同性化"抵消了。如同前面探讨的距离矩阵一样,经典的 SOAP 谱同样属于纯正的标量不变量,内部完全缺失了能够感知镜像反演的负宇称奇通道,因而同样无法在无额外先验辅助下区分立体的对映异构体。

手工不变量:只量距离与角度

3.6 本章总结:不变量表征的潜能与天花板

系统审视不变量表征的历史脉络与数学机制,我们可以客观总结出以下三条根本结论:

对于只涉及旋转与平移的纯标量预测任务(如基态总能量、带隙、晶格形成能),只包含成对距离与键角的内积型特征在信息论意义上几乎是完备的。它们不存在几何信息的严重亏蚀,因而在大量非手性材料基准测试中,设计精良的不变量模型展现出了丝毫不逊色于复杂等变模型的强劲拟合性能。

然而,一旦面临具有空间手征性的立体化学问题,纯内积不变量模型在数学结构上遭遇了不可逾越的阻断。任何只测量空间距离与夹角的算法都无法分辨镜像构型,若要使模型能够响应手性不对称催化或手性药物靶向筛选,神经网络的内部通道必须接纳奇宇称赝标量的传递。

更为决定性的是,传统不变量模型彻底割裂了矢量与张量性质的自然流动路径。当科学计算的目标从静态标量能量拓展到动态受力矢量、外加电场下的极化响应张量或是分子的瞬时偶极矩时,只具备标量通道的网络无法在内部协同维持特征的空间朝向,只能退化为输出缺乏空间关联的孤立幅值。要赋予神经网络内生理解矢量力场与微观多极矩的物理本领,我们必须打破纯不变量的茧房,迈入等变张量积的核心殿堂。


4. 核心机制:张量积与 Clebsch–Gordan 分解

在完成了群表示论与不变量边界的数学准备之后,我们正式进入等变神经网络的核心动力舱。前面我们已经确立了一个基本原则:三维空间中的几何特征必须按照角动量阶数 ll 分门别类存储,且每一层网络的输出必须保持严格的旋转等变性。接下来的关键问题自然浮出水面——在深层网络的前向传播中,上一层提取的几何特征究竟应当如何相互组合,才能在严格维持空间旋转等变的前提下,孕育出下一层更丰富、表达力更强的高阶特征?

4.1 为什么必须耦合两个方向:空间矢量的代数运算陷阱

设想在分子的一个局域配位环境中,某中心原子同时感知到了两个相邻原子的空间配位方向矢量 aR3a \in \mathbb{R}^3bR3b \in \mathbb{R}^3(例如水分子中两个由氧指向氢的 OH\mathrm{O-H} 键方向矢量)。我们希望设计一个神经网络神经元,将这两个方向信息融合成新的隐层特征。在深度学习的常规工具箱中,我们有哪些看似直观、实则违背物理定律的选择?

我们首先审视几种常见的朴素方案。第一种方案是将两个矢量直接相加 a+ba+b。从几何等变性来看,线性叠加确实是等变的(R(a+b)=Ra+RbR(a+b) = Ra + Rb),但这种做法在物理上极度粗糙——它将两个局域键方向彻底坍缩为一个合成矢量,完全抹杀了两个化学键之间的特征夹角 θab\theta_{ab}。第二种方案是将两个矢量的分量逐项相乘 aba\odot b(即阿达马积)。这种做法在物理上是灾难性的:矢量的逐分量相乘根本不是旋转等变的。分量相乘的结果完全依赖于观察者碰巧选取了哪一套直角坐标轴,只要坐标轴在空间中稍微转动一个微小角度,逐分量相乘计算出的数值就会彻底失真,在物理上毫无合法性。第三种方案则是将两个矢量直接并排拼接成六维数组 [a;b]R6[a; b] \in \mathbb{R}^6。从表示论角度来看,这不过是将两个 l=1l=1 空间机械地并置(111 \oplus 1),两个矢量之间没有任何实质性的代数交互,模型依然无法获知它们彼此的相对几何关系。

那么,在三维欧几里得空间中,如何才能在两组矢量之间建立起在空间旋转下物理自洽的非线性交互?从经典矢量几何出发,两个三维矢量之间能够在旋转变换下保持良好几何规律的代数构造在数学上恰恰只有三种形式:

首先是两个矢量的标量点积 aba\cdot b。点积将两个矢量收缩为一个单一实数,它直接反映了键长与键角余弦,在空间旋转下纹丝不动,对应于角动量阶数 l=0l=0 的标量不变量。

其次是两个矢量的空间叉积 a×ba\times b。叉积生成了一个垂直于两矢量所在平面的新三维矢量,其模长刻画了平行四边形面积,方向表征了平面的法向取向。由于刚体旋转严格保持右手螺旋定则,该矢量随着坐标系的旋转协同旋转,对应于角动量阶数 l=1l=1 的新矢量。

最后是两个矢量的对称无迹外积 symtraceless(ab)\mathrm{symtraceless}(a\otimes b)。外积矩阵 aba b^\top 包含了两个矢量全部 9 个成对分量乘积,将其中的迹项剥离并进行对称化后,便构成了描述局域形变与各向异性四极场的二阶张量,在空间旋转下按照 TRTRT \mapsto R T R^\top 协同变换,对应于角动量阶数 l=2l=2

这三种几何组合的维度分别为 1、3、5,三者相加恰好精确等于外积矩阵的总自由度 1+3+5=91 + 3 + 5 = 9。在量子力学与群表示论中,这一优美的几何事实被称为角动量的张量积分解,写成群论语言便是:

11=0    1    2(维度平衡:3×3=9=1+3+5).\mathbf{1}\otimes\mathbf{1} = \mathbf{0}\;\oplus\;\mathbf{1}\;\oplus\;\mathbf{2} \qquad(\text{维度平衡:}3\times3 = 9 = 1+3+5).

4.2 球谐函数与 Wigner-D 矩阵的表示基础

在推导张量分解的显式展开之前,我们先将不可约表示的基底记号规范化。前面提到,ll 阶球谐函数 Yl:R3R2l+1Y_l:\mathbb{R}^3\to\mathbb{R}^{2l+1} 构成了旋转群 SO(3)SO(3) 不可约表示的标准正交函数基底,其在三维旋转矩阵 RR 下的变换由对应阶数的 Wigner-D 矩阵 Dl(R)D_l(R) 精确刻画:Yl(Rr)=Dl(R)Yl(r)Y_l(Rr) = D_l(R)Y_l(r)

对于标量阶数 l=0l=0Y0Y_0 表现为常数,而 D0(R)=1D_0(R)=1 代表平凡不变标量。

对于矢量阶数 l=1l=1,实球谐函数的三个分量在笛卡尔坐标系中正比于单位方向向量 r^=r/r\hat{r} = r/\|r\|,此时 Wigner-D 变换矩阵严格恒等于笛卡尔旋转矩阵本身:D1(R)=RD_1(R)=R

对于二阶张量阶数 l=2l=2Y2Y_2 包含了 5 个独立分量。在直角坐标系中,这 5 个分量可以通过三阶对称无迹矩阵的独立元素等价表达,而 D2(R)D_2(R) 作用于张量矩阵的几何实质即为标准的张量旋转投影:TRTRT\mapsto RTR^\top

需要向读者特别指出的是:五维矢量形式的 l=2l=2 球谐特征与 3×33\times3 的对称无迹矩阵,在代数本质上是同构的同一物理量的两种写法。一个对称的 3×33\times3 矩阵拥有 6 个独立自由度,施加无迹约束(trT=0\mathrm{tr}T = 0)后剥离掉一个球对称自由度,恰好剩余 5 个独立分量,严密吻合 2l+1=52l+1=5 的理论维度。这一同构在工程编码中展现出极大的便利性——我们完全无需在底层繁琐地维护球谐展开式的各种球坐标系数,只需在直角坐标系下直接构造 3×33\times3 无迹对称矩阵,并在计算完成后将其投射回无迹子空间即可。

4.3 1⊗1 = 0⊕1⊕2:张量积闭式正交分解的完整推导

现在,我们以极其透明的初等线性代数,手把手展开外积矩阵的全部推导过程。设给定两组三维矢量 a,bR3a, b \in \mathbb{R}^3,考虑由它们构成的张量外积矩阵 abR3×3a\otimes b \in \mathbb{R}^{3\times3},其矩阵分量定义为 (ab)ij=aibj(a\otimes b)_{ij} = a_i b_j。我们的目标是将其唯一且正交地分解为迹标量(l=0l=0)、反对称矩阵(l=1l=1)以及对称无迹张量(l=2l=2)三个正交子空间。

推导的第一步是将外积矩阵分解为对称矩阵 SS 与反对称矩阵 AA 的直和:

ab=12(ab+ba)对称部分 S+12(abba)反对称部分 A.a\otimes b = \underbrace{\frac{1}{2}\big(a\otimes b + b\otimes a\big)}_{\text{对称部分 }S} + \underbrace{\frac{1}{2}\big(a\otimes b - b\otimes a\big)}_{\text{反对称部分 }A}.

分析反对称部分 AA 可以发现,由于主对角线全为零且 Aji=AijA_{ji} = -A_{ij},该矩阵仅包含 3 个独立的代数自由度。借助三维空间中的 Levi-Civita 反对称张量 ϵijk\epsilon_{ijk},反对称矩阵的分量可以精确表示为 Aij=12(aibjajbi)=12ϵijk(a×b)kA_{ij} = \frac{1}{2}(a_i b_j - a_j b_i) = \frac{1}{2}\epsilon_{ijk}(a\times b)_k。这意味着矩阵 AA 与两矢量的空间叉乘矢量 a×ba\times b 存在严格的一对一伴随对应关系:

A=12(abba)=12[a×b]×.A = \frac{1}{2}\big(a\otimes b - b\otimes a\big) = \frac{1}{2}\,[\,a\times b\,]_\times .

此时,反对称矩阵就是 l=1l=1 矢量通道在矩阵形式下的化身。根据三维叉乘的张量等变公理,在任意旋转矩阵 RR 作用下,恒有 R(a×b)=(Ra)×(Rb)R(a\times b) = (Ra)\times(Rb)。这一性质保证了反对称部分在空间旋转变换下始终保持严格的矢量等变性。

推导的第二步是将对称矩阵 SS 进一步解耦为各向同性迹项与无迹各向异性项。对称矩阵 SS 的矩阵迹等于两矢量的内积:trS=iaibi=ab\mathrm{tr}S = \sum_i a_i b_i = a\cdot b。在旋转变换下,矩阵迹具备良好的不变性:tr(RSR)=tr(SRR)=trS\mathrm{tr}(RSR^\top) = \mathrm{tr}(SR^\top R) = \mathrm{tr}S。我们将对称矩阵减去其平均对角迹分量,定义残差张量 TT

T=Sab3I,其分量显式形式为Tij=12(aibj+ajbi)ab3δij.T = S - \frac{a\cdot b}{3}I,\qquad \text{其分量显式形式为}\quad T_{ij}=\frac{1}{2}(a_ib_j+a_jb_i) - \frac{a\cdot b}{3}\delta_{ij}.

显然,张量 TT 满足对称性 Tji=TijT_{ji} = T_{ij},且其对角线元素之和严格归零:trT=trS(ab)tr(I)/3=abab=0\mathrm{tr}T = \mathrm{tr}S - (a\cdot b)\mathrm{tr}(I)/3 = a\cdot b - a\cdot b = 0。这便是 l=2l=2 的对称无迹张量通道。在旋转变换下,它遵循 TRTRT \mapsto R T R^\top 的几何规律协同演化,具备完整的张量等变性。

将上述解耦过程完整合并,我们便得到了三维矢量外积张量的正交 Clebsch–Gordan 投影分解定理

  ab=ab3Il=0 标量迹    12[a×b]×l=1 矢量叉乘    symtraceless(ab)l=2 对称无迹张量  \boxed{\;a\otimes b = \underbrace{\frac{a\cdot b}{3}I}_{l=0\text{ 标量迹}} \;\oplus\;\underbrace{\frac{1}{2}[a\times b]_\times}_{l=1\text{ 矢量叉乘}} \;\oplus\;\underbrace{\mathrm{symtraceless}(a\otimes b)}_{l=2\text{ 对称无迹张量}}\;}

4.4 投影正交性与范数守恒定理

上述分解中的符号 \oplus 代表三个子空间在 Frobenius 内积意义下保持严格的正交性。为了定量验证这种正交隔离,我们计算三个分解分量的 Frobenius 范数平方(定义为 MF2=ijMij2\|M\|_F^2 = \sum_{ij} M_{ij}^2):

对于标量通道,ab3IF2=(ab)29tr(I)=(ab)23\|\frac{a\cdot b}{3}I\|_F^2 = \frac{(a\cdot b)^2}{9}\mathrm{tr}(I) = \frac{(a\cdot b)^2}{3}

对于矢量反对称通道,利用反对称叉乘矩阵的范数性质 [v]×F2=2v2\|[v]_\times\|_F^2 = 2\|v\|^2,得到 12[a×b]×F2=142a×b2=a×b22\|\frac{1}{2}[a\times b]_\times\|_F^2 = \frac{1}{4}\cdot 2\|a\times b\|^2 = \frac{\|a\times b\|^2}{2}

对于对称无迹张量通道,保留其张量范数 TF2\|T\|_F^2

由于三个张量分量两两正交(例如对称张量与反对称张量的内积恒等于零,无迹张量与单位阵的内积恒等于零),全外积矩阵的总范数精确等于各分量范数之和:

  a2b2=(ab)23+a×b22+TF2  \boxed{\;\|a\|^2\|b\|^2 = \frac{(a\cdot b)^2}{3} + \frac{\|a\times b\|^2}{2} + \|T\|_F^2\;}

左端正是外积矩阵自身的模长平方 abF2\|a\otimes b\|_F^2。这个恒等式被称为几何张量积的能量守恒式。在计算机中,这一解析恒等式为数值检验张量积代码的数学正确性提供了终极判据——在测试脚本 tests_enn.py 的首组单元测试中,我们通过在数十组任意随机向量上检验该式两端的差值,其数值误差稳定低于 101010^{-10},从而在底层确证了张量分解逻辑的绝对无误。

4.5 Schur 引理:为什么等变张量积具有数学唯一性

在推导完 111\otimes1 的三种通道后,许多研究者都会自然产生一个深层疑问:我们推导出的内积、叉积与对称外积,究竟只是人类碰巧构想出的一种启发式设计,还是说在更深邃的数学空间里,原本还隐藏着第四种、第五种更具魔力的高维等变组合?

群表示论的核心基石——Schur 引理,以无可辩驳的代数力量给出了终极裁决:不存在其他任何形式,上述构造在等变双线性映射的意义下是唯一的

Schur 引理(Irreducible Representation Lemma):设 VVWW 是紧致李群 GG 的两个有限维不可约复表示空间,T:VWT:V\to W 是一个与群作用算子全局交换的等变线性映射(即对任意群元 gGg\in G,恒满足 TρV(g)=ρW(g)TT\rho_V(g) = \rho_W(g)T)。则该映射只有两种可能:若表示 VVWW 不等价,则映射必为零映射 T=0T=0;若表示空间等价 VWV\cong W,则映射 TT 必然且只能是单位映射的纯标量倍数,即 T=λIT = \lambda IλC\lambda \in \mathbb{C})。

Schur 引理的证明思想极为优美:算子与群作用的交换性迫使其核空间 kerT\ker T 与像空间 imT\mathrm{im} T 成为群作用下的不变子空间。由于不可约表示在定义上不存在任何非平凡的不变子空间,kerT\ker T 只能退化为零空间或占据全空间。这直接封死了线性映射的自由度,将其坍缩为唯一的同构缩放。

将 Schur 引理应用于我们的神经网络特征设计:我们希望寻找一个能够保持三维旋转等变的双线性特征融合算子 B:Vl1×Vl2Vl3B: V_{l_1} \times V_{l_2} \to V_{l_3}。根据代数张量积的普适性质,这完全等价于寻找一个从张量积直和空间向目标表示空间的等变线性映射 B~:Vl1Vl2Vl3\tilde{B}: V_{l_1} \otimes V_{l_2} \to V_{l_3}

而张量积空间 Vl1Vl2V_{l_1} \otimes V_{l_2} 的不可约表示分解,由物理学中著名的 Clebsch–Gordan 级数定理严格限定:

Vl1Vl2    l=l1l2l1+l2Vl.V_{l_1}\otimes V_{l_2} \;\cong\; \bigoplus_{l=|l_1-l_2|}^{l_1+l_2} V_l .

对于两组 l1=l2=1l_1=l_2=1 矢量通道的耦合,级数的角动量求和范围限定在 11=0|1-1|=01+1=21+1=2 之间,分解结果严格收敛为 V0V1V2V_0\oplus V_1\oplus V_2。在这一分解中,每一个目标不可约表示 VlV_l 恰好只出现了一次(无重复简并重数)。

结合 Schur 引理,这意味着:从两个矢量输入生成目标表示 VlV_ll{0,1,2}l\in\{0,1,2\})的任意旋转等变双线性算子,在数学上唯一受限于对应的 Clebsch–Gordan 耦合路径,且每条耦合路径的可学习参数在自由度上只允许存在一个全局标量权重因子 wlw_l

B(a,b)=wl(CG 耦合通道)(a,b),wlR.B(a,b) = w_l \cdot \big(\text{CG 耦合通道}\big)(a,b),\qquad w_l\in\mathbb{R}.

在材料与深度学习交叉领域充斥着大量启发式经验调优的背景下,这个结论具有非同寻常的定海神针作用:等变张量积网络绝非某种经验主义的架构尝试,而是由对称性法则从底层唯一确立的必然形态。研究者可以完全放下寻找"其他神奇耦合机制"的焦虑,确信物理对称性已经为我们指明了唯一合法的数学路径。

4.6 空间反演宇称:负宇称赝标量的几何演生机制

在确立了旋转群 SO(3)SO(3) 的等变性之后,我们必须面对三维空间中另一个极其关键的操作——空间反演(Inversion)与镜像反射。

当完全正交群 O(3)O(3) 作用于物理世界时,几何对象除了用角动量阶数 ll 标识旋转规律外,还必须携带一个标识空间反演响应的宇称标签(Parity) p{+1,1}p\in\{+1, -1\}。在空间坐标全反演变换 xxx \mapsto -x 下,一个几何特征若保持符号不变,则具有偶宇称(p=+1p=+1,通常标记为 "e");若发生全符号翻转,则具有奇宇称(p=1p=-1,通常标记为 "o")。这正是现代等变库 e3nn 中将不可约通道规范命名为 Irrep("1o")Irrep("0e") 的深层物理根源。

我们来深入追踪前面推导出的三个张量通道在空间全反演下的宇称演化。在反演操作下,空间矢量作为标准极矢量(如坐标位移差),其变换规律为奇宇称:aaa \mapsto -abbb \mapsto -b。将其分别代入三个耦合通道中:

标量点积通道表现为 (a)(b)=ab(-a)\cdot(-b) = a\cdot b,宇称为偶(p=+1p=+1);

矢量叉乘通道表现为 (a)×(b)=a×b(-a)\times(-b) = a\times b,宇称同样为偶(p=+1p=+1);

对称无迹张量通道表现为 symtraceless((a)(b))=symtraceless(ab)\mathrm{symtraceless}\big((-a)\otimes(-b)\big) = \mathrm{symtraceless}(a\otimes b),宇称依然为偶(p=+1p=+1)。

这一代数推导展现出一个令许多初涉几何学习的学者大为震惊的客观物理事实:任意两个奇宇称极矢量通过双线性张量积耦合,其负号必然在分量相乘中成对相互抵消,因而产出的三个几何通道全都是偶宇称的。其中,矢量通道 a×ba\times b 在反演下未发生翻转,这在物理力学中被称为轴矢量(或假矢量,如角动量与磁场);而标量通道 aba\cdot b 更是标准的真标量。

这就从最深层的群表示论机理上彻底解开了第 3.4 节所揭示的手性谜题:常规的等变网络即便引入了张量积,只要其隐层通道全部由偶宇称特征构成,它在面对空间镜像对映体时依然是完全失明的一对一退化。要让模型演生出能够识别对映异构体的负宇称赝标量,特征耦合必须至少容纳三个奇宇称矢量的级联复合:

χ=(u1×u2)u3  xx  (u1×u2)(u3)=χ(宇称 p=1).\chi = (u_1\times u_2)\cdot u_3 \;\xrightarrow{\quad x\mapsto -x\quad}\; (-u_1\times -u_2)\cdot(-u_3) = -\chi \quad(\text{宇称 }p=-1).

对于致力于攻关手性催化与药物立体活性的研究团队而言,这是一个决定性的架构设计指南:若要在模型中内生支持手性判别,必须采用具备宇称感知(Parity-aware)的高阶网络(如开辟了专门奇宇称通道的 NequIP 或 MACE 架构),或者在节点的初始几何特征中显式植入局域立体结构的有向标量三重积。

两个箭头耦合成三样东西

4.7 张量积层的手写解析梯度与反向传播

为了在纯 Python 环境中落地并严格验证张量积算子,我们必须给出张量积层在所有特征通道上的闭式解析反向传播梯度。设张量积层接收通道数为 CC 的两组矢量特征 A={ac}c=1CA = \{a_c\}_{c=1}^CB={bc}c=1CB = \{b_c\}_{c=1}^C,并通过可学习的权重矩阵输出三类通道:

sk=c=1CWkcs(acbc),vk=c=1CWkcv(ac×bc),Tk=c=1CWkctsymtraceless(acbc).s_k=\sum_{c=1}^C W^s_{kc}(a_c\cdot b_c),\qquad v_k=\sum_{c=1}^C W^v_{kc}(a_c\times b_c),\qquad T_k=\sum_{c=1}^C W^t_{kc}\,\mathrm{symtraceless}(a_c\otimes b_c).

设来自下游损失函数的上游误差梯度分别为标量项 gsRKsg^s \in \mathbb{R}^{K_s}、矢量项 gvRKv×3g^v \in \mathbb{R}^{K_v\times 3} 以及二阶张量项 GRKt×3×3G \in \mathbb{R}^{K_t\times 3\times 3}。我们以完全解析的方式推导对第 cc 个输入矢量 aca_cbcb_c 的反向梯度传递方程。

对于标量通道 l=0l=0,利用内积求导公式 (ab)a=b\frac{\partial (a\cdot b)}{\partial a} = b。定义线性权重的上游聚合项 αc=kgksWkcs\alpha_c = \sum_k g^s_k W^s_{kc},则该路径对输入矢量的梯度贡献为:

Lac+=αcbc,Lbc+=αcac.\frac{\partial\mathcal{L}}{\partial a_c} \mathrel{+}= \alpha_c\,b_c, \qquad \frac{\partial\mathcal{L}}{\partial b_c} \mathrel{+}= \alpha_c\,a_c .

对于矢量通道 l=1l=1,利用标量三重积的循环微分恒等式 g(a×b)=a(b×g)=b(g×a)g\cdot(a\times b) = a\cdot(b\times g) = b\cdot(g\times a)。定义聚合矢量 βc=kgkvWkcvR3\beta_c = \sum_k g^v_k W^v_{kc} \in \mathbb{R}^3,利用反对称性调整求导顺序,梯度贡献解析表达为:

Lac+=bc×βc,Lbc+=βc×ac.\frac{\partial\mathcal{L}}{\partial a_c} \mathrel{+}= b_c\times\beta_c, \qquad \frac{\partial\mathcal{L}}{\partial b_c} \mathrel{+}= \beta_c\times a_c .

对于二阶张量通道 l=2l=2,这是整个代数反求中最容易产生细微混淆的环节。定义上游张量梯度的加权求和 Gc=kWkctGkR3×3G_c = \sum_k W^t_{kc} G_k \in \mathbb{R}^{3\times3}。考虑张量内积 Gc,Tc=ij(Gc)ij(Tc)ij\langle G_c, T_c\rangle = \sum_{ij} (G_c)_{ij} (T_c)_{ij}。将对称无迹张量 Tij=12(aibj+ajbi)ab3δijT_{ij} = \frac{1}{2}(a_i b_j + a_j b_i) - \frac{a\cdot b}{3}\delta_{ij} 代入求导:

Tijak=12(δikbj+δjkbi)bk3δij.\frac{\partial T_{ij}}{\partial a_k} = \frac{1}{2}\big(\delta_{ik}b_j + \delta_{jk}b_i\big) - \frac{b_k}{3}\delta_{ij}.

将其与上游张量矩阵 GcG_c 进行全收缩求和:

ij(Gc)ijTijak=12[j(Gc)kjbj+i(Gc)ikbi]bk3i(Gc)ii=(Sym(Gc)b)ktrGc3bk,\sum_{ij} (G_c)_{ij} \frac{\partial T_{ij}}{\partial a_k} = \frac{1}{2}\Big[\sum_j (G_c)_{kj} b_j + \sum_i (G_c)_{ik} b_i\Big] - \frac{b_k}{3}\sum_i (G_c)_{ii} = \big(\mathrm{Sym}(G_c)\,b\big)_k - \frac{\mathrm{tr}G_c}{3} b_k ,

其中 Sym(Gc)=12(Gc+Gc)\mathrm{Sym}(G_c) = \frac{1}{2}(G_c + G_c^\top) 为张量梯度的对称化分量。由此得出无迹张量反向梯度的闭式解:

Lac+=Sym(Gc)bctrGc3bc,Lbc+=Sym(Gc)actrGc3ac.\frac{\partial\mathcal{L}}{\partial a_c} \mathrel{+}= \mathrm{Sym}(G_c)\,b_c - \frac{\mathrm{tr}\,G_c}{3}\,b_c, \qquad \frac{\partial\mathcal{L}}{\partial b_c} \mathrel{+}= \mathrm{Sym}(G_c)\,a_c - \frac{\mathrm{tr}\,G_c}{3}\,a_c .

在此处需要向手写几何神经网络的读者敲响警钟:在对矢量 aa 求导时,式中与对称化梯度矩阵相乘的配对矢量必须是 bb;反之对 bb 求导时配对的才是 aa。在特征输入对称相等的测试用例中(a=ba=b),写反这一关系依然能够巧合通过测试;而一旦应用于不同矢量特征的常规层,微小的代数错位将引发剧烈的梯度漂移。在代码 TensorProduct11.backward 中,我们以逐行严谨的形式固化了这一求导规律。


5. 经典架构:从距离到高阶张量

掌握了张量积与 Clebsch–Gordan 分解这一核心发动机后,我们现在可以站上理论的高位,系统俯瞰现代计算化学中涌现出的经典几何架构。它们并非互不相干的巧合产物,而是同一套几何图网络范式在不同"几何阶数深度"上的具象实现。

模型家族树:从只用距离,到角度,到高阶等变张量

消息传递:邻居把自己的信息带方向地送过来

5.1 SchNet:只依赖空间距离的连续滤波卷积

SchNet 由 Schütt 等人于 2017 年提出,是纯不变量图卷积网络的开山之作。其核心物理设计极为清晰:在量子化学基态能量近似中,主要的相互作用随原子间距平滑衰减,因此模型应当将空间距离作为连续卷积核的载体。

模型首先将成对原子间距 rij=xixjr_{ij} = \|x_i - x_j\| 通过一组高斯径向基函数(RBF)在空间上离散展开:

ek(rij)=exp(γ(rijμk)2),k=1,,K,e_k(r_{ij}) = \exp\big(-\gamma(r_{ij}-\mu_k)^2\big),\qquad k=1,\dots,K,

其中高斯中心 μk\mu_k 均匀散布在所关注的相互作用截断半径区间 [rmin,rmax][r_{\min}, r_{\max}] 上。随后,展开后的高斯距离特征通过小型滤波多层感知机映射为连续卷积核,并对相邻节点的隐层特征实施阿达马加权卷积:

mij=MLPfilter(e(rij))RF,hi(l+1)=hi(l)+jimij(Whj(l)).\mathbf{m}_{ij} = \mathrm{MLP}_{\text{filter}}\big(e(r_{ij})\big)\in\mathbb{R}^F, \qquad h_i^{(l+1)} = h_i^{(l)} + \sum_{j\ne i} \mathbf{m}_{ij}\odot \big(W h_j^{(l)}\big).

SchNet 具备严格的平移与旋转不变性,这是因为整个消息传递过程完全被成对距离标量所主导。然而,各向同性的距离核也是一把双刃剑:SchNet 内部流动的消息完全缺失了空间取向维度。当两个局域环境互为手性反演时,它给出的卷积消息完全一致。更为深层的是,若直接将 SchNet 用于预测分子力场,原子的受力只能通过对体系总能量进行解析反向求导获得:

Fi=Exi=jiErijxixjrij=jiErijr^ij.F_i = -\frac{\partial E}{\partial x_i} = -\sum_{j\ne i}\frac{\partial E}{\partial r_{ij}}\frac{x_i - x_j}{r_{ij}} = -\sum_{j\ne i}\frac{\partial E}{\partial r_{ij}}\hat{r}_{ij}.

这一求导关系揭示了纯距离模型的受力本质:每个原子受到的合力被严格约束在成对连心线 r^ij\hat{r}_{ij} 的线性组合方向上。在以两体中心势为主的简单共价体系中这一近似十分优异,但面对具有强烈离域共轭取向、多中心键或复杂配位场的各向异性体系时,它将面临表达力受限的挑战。

5.2 DTNN:早期深度张量交互网络

在 SchNet 诞生之前,深度张量神经网络(DTNN)作为其前驱,在 DeepChem 等分子学习框架中留下了深刻的技术印记。DTNN 的核心思想同样聚焦于距离不变量,它通过双线性张量乘法将成对原子的核电荷特征与连续距离表示进行矩阵隐层交叉,以两两相互作用逐步累积中心原子环境。DTNN 与 SchNet 的演进关系,如同特征工程时代向现代端到端连续图卷积跨越的过渡桥梁,其数学表达已被完全包容于广义的距离标量传递网络范畴。

5.3 DimeNet:将定向三体键角引入消息图

为了打破纯两体距离对方向特征的遮蔽,DimeNet 由 Klicpera 等人于 2020 年提出。它在图结构设计上实现了一个质的飞跃:图网络的消息传递实体不再是孤立的原子节点,而是带有空间指向的有向键向量(Directed Edges)

在 DimeNet 中,消息沿着原子对 (j,i)(j, i) 向相邻原子对 (i,k)(i, k) 进行传递,从而在局域空间中自然构造出以原子 ii 为顶点的三体键角 θjik\theta_{jik}。在特征表示上,模型引入了理论物理中解径向薛定谔方程的利器——二维球面傅里叶–贝塞尔(Spherical Fourier–Bessel)基底:

mjiik=f(n,lcnljl ⁣(znlrcrik)Yl(θjik)),\mathbf{m}_{ji\to ik} = f\Big(\sum_{n,l} c_{nl} j_l\!\left(\frac{z_{nl}}{r_c}r_{ik}\right) Y_l(\theta_{jik})\Big),

其中 jlj_l 为球贝塞尔函数,znlz_{nl} 为其零点坐标,而 YlY_l 则是关于键角张角的连带勒让德基底。通过这种方式,三维空间中两根相交化学键的几何取向通过球谐展开深度嵌入到消息函数中。DimeNet 显著增强了对共价键角刚性与势能面弯曲振动的拟合精度,但其计算代价也随之跃升——局域邻域的遍历复杂度从节点级度数 O(Nk)O(N\cdot k) 攀升到了连通键对的二阶复杂度 O(Nk2)O(N\cdot k^2)

5.4 Tensor Field Networks:球谐卷积核的完整推导

现在,我们将迎来现代等变网络史上最厚重、最优雅的一座理论里程碑:Thomas 等人于 2018 年在 Tensor Field Networks (TFN) 中建立的等变卷积完整理论。该理论从第一性数学原理回答了一个终极问题:一个作用于连续三维物理场、能够实现不同阶数不可约表示平滑转化的连续卷积核,在形式上究竟受到了怎样的绝对约束?

设网络在空间位置 rjr_j 处具有输入特征场 finR2lin+1f_{\text{in}} \in \mathbb{R}^{2l_{\text{in}}+1},我们希望通过局域连续积分卷积核 K(rirj)K(r_i - r_j),生成位于中心点 rir_i 处的输出特征场 foutR2lout+1f_{\text{out}} \in \mathbb{R}^{2l_{\text{out}}+1}

fout(ri)=jK(rirj)fin(rj).f_{\text{out}}(r_i) = \sum_{j} K(r_i - r_j)\,f_{\text{in}}(r_j).

我们要求该积分变换对整个物理系统的三维空间刚体旋转严格等变:当空间坐标经历刚体旋转 RSO(3)R \in SO(3)(相对位移变换为 R(rirj)R(r_i - r_j),输入特征在内部变换为 Dlin(R)finD_{l_{\text{in}}}(R)f_{\text{in}})时,卷积核计算出的输出特征必须精确以 Dlout(R)foutD_{l_{\text{out}}}(R)f_{\text{out}} 的规律协同演化。

第一步:基于球谐正交基底解耦卷积核的空间径向与角向自由度。在三维连续空间中,任何平方可积的矩阵值连续核函数 K(r)K(\vec{r}),都可以唯一地展开为径向标量函数与球谐角向函数的正交多极展开:

Kmout,min(r)=l=0m=llRlmout,min(r)Ylm(r^),K_{m_{\text{out}},m_{\text{in}}}(\vec{r}) = \sum_{l=0}^\infty \sum_{m=-l}^l R_l^{m_{\text{out}},m_{\text{in}}}(\|\vec{r}\|)\,Y_{lm}(\hat{r}),

其中 r^=r/r\hat{r} = \vec{r}/\|\vec{r}\| 为空间单位方向矢量。

第二步:将旋转等变性转化为对径向标量核的代数约束。将上述多极展开式代入等变泛函方程 K(Rr)Dlin(R)=Dlout(R)K(r)K(R\vec{r}) D_{l_{\text{in}}}(R) = D_{l_{\text{out}}}(R) K(\vec{r}) 中。利用 Wigner-D 矩阵在球谐基底上的正交投影性质,并在不可约张量积空间中展开,根据前面推导出的 Schur 引理,从表示空间 VlinVlV_{l_{\text{in}}} \otimes V_l 到目标表示空间 VloutV_{l_{\text{out}}} 的等变映射在代数上完全唯一,且各张量分量之间的耦合系数必须完全正比于 Clebsch–Gordan 投影矩阵。

这意味着,角向多极基底的展开系数被群表示论的对称性先验完全锁死,卷积核的最一般通用形态在数学上被彻底定型为:

  Kmout,min(r)=l=loutlinlout+lin  Rl(r)任意可学习的径向连续函数  Cl,mout;lin,minlout由代数对称性唯一固定的 CG 耦合常数  Ylm(r^)  \boxed{\; K_{m_{\text{out}},m_{\text{in}}}(\vec{r}) = \sum_{l=|l_{\text{out}}-l_{\text{in}}|}^{l_{\text{out}}+l_{\text{in}}}\; \underbrace{R_l\big(\|\vec{r}\|\big)}_{\text{任意可学习的径向连续函数}}\; \underbrace{C^{\,l_{\text{out}}}_{\,l,m_{\text{out}};\,l_{\text{in}},m_{\text{in}}}}_{\text{由代数对称性唯一固定的 CG 耦合常数}}\;Y_{lm}(\hat{r})\;}

第三步:透视现代等变网络的双重设计真谛

其一是参数空间的精纯化。TFN 的卷积核在空间方向上的角向分布根本不是由梯度下降随机学出来的,它从一开始就被三维旋转群的不可约球谐表示完全确定;模型唯一需要学习的自由度,仅仅是随着原子间标量距离变化的径向函数 Rl(r)R_l(\|\vec{r}\|)(在实现中通常采用一组高斯 RBF 搭配轻量全连接网络)。这赋予了等变卷积网络无与伦比的归纳偏置、极致的参数利用率以及极高的数据学习效率。

其二是角动量角选择定则的约束代价。式中的级数求和范围 loutlinllout+lin|l_{\text{out}}-l_{\text{in}}| \le l \le l_{\text{out}}+l_{\text{in}} 是量子力学角动量选择定则的直接映射。这揭示了一个关键的算力规律:如果我们在高阶网络中试图实现从 lin=1l_{\text{in}}=1 矢量向 lout=2l_{\text{out}}=2 二阶张量的复杂特征转换,参与卷积的球谐基底必须同时囊括 l=1,2,3l=1, 2, 3 的高阶项。随着网络截断角量子数 lmaxl_{\max} 的提升,张量通道数目与 Clebsch–Gordan 耦合路径呈现立方级急剧膨胀,这也是为什么工业级高阶等变网络通常谨慎地将最高阶数截断在 lmax=2l_{\max}=2lmax=3l_{\max}=3 之内。

5.5 等变图神经网络(EGNN):坐标直接流动范式

虽然 TFN 在理论上达到了极高的严密性,但其底层引入的球谐函数、Clebsch–Gordan 系数表以及复杂的张量积索引,在工程落地与计算耗时上构筑了不小的壁垒。面对这一现实,Satorras 等人于 2021 年提出了极富创新意味的等变图神经网络(EGNN),走出了一条极为轻灵、极具性价比的极简路线。

EGNN 的核心构想可以凝练为一句话:既然显式维护高阶不可约表示张量如此昂贵,我们为何不干脆让分子的三维真实笛卡尔坐标本身,作为直接参与隐层传递与迭代的几何矢量载体?

在 EGNN 中,特征系统被解耦为节点的非标量特征(仅包含标量 hiRFh_i \in \mathbb{R}^F)与每个原子的空间几何坐标 xiR3x_i \in \mathbb{R}^3。其每一层的演化方程简洁明快:

mij=ϕe(hi,hj,xixj2)(局域各向同性边消息)hi=hi+ϕh(hi,  jN(i)mij)(节点标量环境迭代)xi=xi+jN(i)(xixj)ϕx(mij)(原子坐标显式几何更新)\begin{aligned} \mathbf{m}_{ij} &= \phi_e\big(h_i,\,h_j,\,\|x_i-x_j\|^2\big) &&\text{(局域各向同性边消息)}\\ h_i' &= h_i + \phi_h\Big(h_i,\;\sum_{j\in\mathcal{N}(i)}\mathbf{m}_{ij}\Big) &&\text{(节点标量环境迭代)}\\ x_i' &= x_i + \sum_{j\in\mathcal{N}(i)}\big(x_i-x_j\big)\,\phi_x(\mathbf{m}_{ij}) &&\text{(原子坐标显式几何更新)} \end{aligned}

在此结构中,ϕe,ϕh,ϕx\phi_e, \phi_h, \phi_x 均为标准的小型多层感知机。整个设计处处彰显着精妙的几何自洽:边消息仅由平移与旋转不变的欧几里得距离平方激发,标量更新只聚合各向同性消息,而最为核心的坐标更新网络 ϕx\phi_x 仅输出一个标量纯量,去动态加权成对相对位移差 xixjx_i - x_j。由于更新位移严格沿着空间相对矢量的方向展开,当坐标经历刚体平移与旋转时,变换矩阵可以直接穿透标量加权提取出来,从而以零球谐开销达成了完美的刚体等变性。我们在第 6 章将对 EGNN 展开详尽推导。

5.6 几何矢量感知机(GVP)与矢量神经元模型

在轻量级标量网络(EGNN/SchNet)与重型高阶张量网络(TFN)之间,存在着一个极具实用价值的折中地带:严格截断在标量与矢量通道(l=0l=1l=0 \oplus l=1)的门控矢量架构。以 Jing 等人提出的几何矢量感知机(GVP)、以及 Schütt 等人提出的 PaiNN 为典型代表。

这类模型放弃了构造计算昂贵的 l=2l=2 二阶张量通道,将特征池严格限制在标量 hRFh \in \mathbb{R}^F 与三维矢量集合 VRM×3V \in \mathbb{R}^{M\times 3}。为了在避免破坏矢量等变的前提下引入非线性能力,模型采用标量门控(Scalar Gating)机制:

(h,V)=(MLP(h,V),  σ(Wh+b)V).\big(h',\,V'\big) = \Big(\mathrm{MLP}\big(h, \|V\|\big),\;\sigma\big(W h + b\big)\odot V\Big).

在这里,多层感知机仅利用矢量的欧几里得模长(旋转不变量)来更新标量状态,而对于多通道矢量特征,则通过标量特征激活出的无量纲门控系数进行各向同性缩放。由于标量门控因子 sRs \in \mathbb{R} 在空间旋转下恒定不变,标量与矢量的乘积满足 s(Rv)=R(sv)s\cdot(Rv) = R(s\cdot v),等变性天然成立。GVP 架构凭借极其轻量的参数体积与优秀的数值鲁棒性,在 AlphaFold 系列的前沿蛋白质侧链建模与大分子结合表征中得到了广泛的应用。

5.7 NequIP、Allegro 与 MACE:现代高精度机器学习势能面

在今日计算化学与材料模拟的最前沿基准(如 rMD17、OC20 催化表面、Matbench 材料基准)上,主导高精度势能面拟合的是以 NequIP、Allegro 和 MACE 为代表的现代高阶张量与多体展开模型。它们代表了等变几何网络深水区的集大成者。

NequIP 坚守 TFN 严谨的高阶不可约张量积范式,但它洞察到了几何深度学习中一条反常识的物理规律:对于分子原子体系,深层堆叠的图卷积往往会因过度平滑而削弱局域化学敏锐度,真正高精度的势能面往往来自于"极浅的图层数(通常仅 2~4 层)搭配极高阶的张量通道(最高阶数可达 lmax=3l_{\max}=3)"。少层浅网配合完备的高阶角动量表示,赋予了 NequIP 惊人的小样本拟合精度,在数千甚至数百个 DFT 构象下即可训练出接近从头算级别的超高精度力场。

Allegro 则直面了图神经网络在大规模并行分子动力学模拟中的致命痛点。传统图卷积每增加一层,中心原子的有效相互作用截断半径便成倍向外扩张,导致分布式计算中的通信开销剧烈膨胀。Allegro 完全摒弃了跨节点的长程消息传递,创新性地将等变不可约张量积完全聚焦在每对原子局域局域化学环境的直接耦合上,实现了严格的物理局域截断,使得等变神经网络势函数首次能够在超级计算机上扩展至百万甚至千万原子级的大规模材料动态模拟。

MACE 由 Batatia 等人于 2022 年提出,它在高阶张量的基石上巧妙融入了类似于材料原子团簇展开(Cluster Expansion)的高阶多体物理思想。MACE 证明了通过在节点处对不可约张量进行连续自对称化张量积,可以在保持线性计算复杂度的前提下,端到端显式重构出涵盖四体、五体乃至更高阶的多体相互作用能量。这种多体展开机制让模型在极少的消息传递跳数内便能充分逼近复杂的电子关联与化学键断裂重组,成为当今第一性原理力场生成领域最为璀璨的标杆工具之一。

5.8 主流三维几何表征架构横向全景比对

为了让材料与计算化学领域的研究者在面对具体课题时能够一目了然地确立选型坐标,我们将上述各类经典几何架构的核心技术特征与应用甜点区凝练为横向比对矩阵:

模型架构几何输入信息源对称性等级典型物理输出形式每层渐近计算复杂度核心优势与典型应用领域
经典手工描述符(CM / SOAP / ACSF)成对距离、局域键角与多极展开纯不变量标量热力学量不依赖深层消息传递极小规模数据集、轻量快速基线拟合
SchNet / DTNN连续高斯展宽成对原子间距纯不变量标量能量与热力学性质节点局域度数 O(Nk)O(Nk)构象能量预测、早期连续滤波神经网络力场
DimeNet / GemNet径向距离搭配定向有向键角纯不变量标量性质、几何构象基态连通键对复杂度 O(Nk2)O(Nk^2)强共价键角敏感体系、平衡构象精确优化
GVP / PaiNN径向距离搭配隐层极矢量门控E(3)E(3) 刚体等变标量与空间极矢量场节点局域度数 O(Nk)O(Nk)蛋白质残基结构建模、大分子结合自由能评估
EGNN成对距离平方搭配显式坐标移动E(3)E(3) 刚体等变标量能量与原子空间力场节点局域度数 O(Nk)O(Nk)反应构象生成扩散模型、低计算预算平衡动力学
TFN / SE(3)-Transformer高阶实球谐基底搭配 Clebsch–GordanSE(3)SE(3) 刚体等变任意不可约张量多极矩场依赖通道数与截断阶数 O(NkC)O(Nk\cdot C)严格物理场拟合、电多极矩与介电响应建模
NequIP / Allegro / MACE高阶球谐张量搭配不可约多体展开SE(3)SE(3) 刚体等变高精度势能面与动力学力场随最高角动量立方级增长 O(Nklmax3)O(Nk\cdot l_{\max}^3)第一性原理高精度材料反应力场、纳秒级相变模拟

在梳理完这幅浩瀚的架构演化谱系后,我们将镜头拉回最为纯粹、兼具理论美感与工程极简性的模型——等变图神经网络(EGNN)。在接下来的第 6 章中,我们将以最朴素的代数推导,亲手揭示其等变性成立的严密机制与隐藏在训练底层的稳定性密码。


6. EGNN:完整推导与稳定性

在理解了基于球谐张量的高阶等变网络之后,本章将视角切换至另一条极具工业美感的轻量化路径——等变图神经网络(EGNN)。这一架构摒弃了高阶张量通道与 Clebsch–Gordan 系数的庞大矩阵运算,转而直接在实空间中利用分子的三维笛卡尔坐标充当几何载体。

我们将系统推导 EGNN 的空间等变性成立的严密代数条件,剖析其内在的牛顿第三定律作用力反作用力抵消机制(质心守恒律),并对两种不同的坐标更新策略进行微分几何层面的解析求导。最后,教程将展示在真实高分子聚合物势能面拟合中亲历的数值发散真实记录,并阐明工程调优与数值稳定的底层机制。

6.1 层定义与几何自洽结构

EGNN 层的输入由三部分构成:节点标量特征集合 h={hiRF}i=1Nh = \{h_i \in \mathbb{R}^F\}_{i=1}^N、原子空间笛卡尔坐标集合 x={xiR3}i=1Nx = \{x_i \in \mathbb{R}^3\}_{i=1}^N,以及定义原子连通关系的边集 E\mathcal{E}。经过该层变换后,模型同时输出更新后的标量特征 hiRFh_i' \in \mathbb{R}^F 与更新后的原子坐标 xiR3x_i' \in \mathbb{R}^3

其前向传播由三个相互交织的几何步骤构成:

mij=ϕe(hi,hj,xixj2)(各向同性边消息传递)hi=hi+ϕh(hi,  jN(i)mij)(节点标量环境更新)xi=xi+jN(i)(xixj)ϕx(mij)(坐标显式几何更新)\begin{aligned} \mathbf{m}_{ij} &= \phi_e\big(h_i,\,h_j,\,\|x_i-x_j\|^2\big) &&\text{(各向同性边消息传递)}\\ h_i' &= h_i + \phi_h\Big(h_i,\;\sum_{j\in\mathcal{N}(i)}\mathbf{m}_{ij}\Big) &&\text{(节点标量环境更新)}\\ x_i' &= x_i + \sum_{j\in\mathcal{N}(i)}\big(x_i-x_j\big)\,\phi_x(\mathbf{m}_{ij}) &&\text{(坐标显式几何更新)} \end{aligned}

在上述结构中,ϕe,ϕh,ϕx\phi_e, \phi_h, \phi_x 均由两层带有 SiLU 激活函数的紧凑多层感知机实现。

坐标更新:原子被邻居轻轻推一下,推的方向随整体旋转

分析这一架构,必须关注一个极其根本的设计约束:坐标更新子网络 ϕx\phi_x 的输出维度必须严格限制为一个标量(一维实数),而绝不能输出高维向量。在物理图像上,ϕx(mij)\phi_x(\mathbf{m}_{ij}) 扮演着邻居原子 jj 对中心原子 ii 施加的局域微观"无量纲推拉强度";如果允许 ϕx\phi_x 输出非标量矢量,其输出方向将与坐标系的选定发生未受约束的耦合,从而彻底摧毁整个层的旋转等变性。

6.2 旋转与平移等变性的严格证明

我们以形式化代数来严谨证明 EGNN 的等变性。

命题(EGNN 层的刚体等变性):设存在任意三维正交刚体旋转矩阵 RSO(3)R\in SO(3) 与任意三维平移矢量 tR3t\in\mathbb{R}^3。若输入特征经历了对应的空间刚体变换,满足标量特征保持不变(h~i=hi\tilde{h}_i = h_i)且空间坐标整体刚性位移(x~i=Rxi+t\tilde{x}_i = Rx_i + t),则 EGNN 层的前向输出必然严格满足标量输出不变(h~i=hi\tilde{h}_i' = h_i')且坐标输出协同刚体位移(x~i=Rxi+t\tilde{x}_i' = Rx_i' + t)。

证明:推导过程沿计算图自然分为三步:

首先考察局域边消息的不变性。在空间变换后,成对原子间的相对位移演变为 x~ix~j=(Rxi+t)(Rxj+t)=R(xixj)\tilde{x}_i - \tilde{x}_j = (Rx_i + t) - (Rx_j + t) = R(x_i - x_j)。由于旋转矩阵具备严格的正交保距性(RR=IR^\top R = I),变换后的成对欧几里得距离平方满足:

x~ix~j2=(xixj)RR(xixj)=xixj2.\|\tilde{x}_i - \tilde{x}_j\|^2 = (x_i - x_j)^\top R^\top R (x_i - x_j) = \|x_i - x_j\|^2.

平移矢量 tt 在坐标差分中完全对消,而旋转矩阵 RR 在内积中被正交性完全吸收。因此,边消息网络 ϕe\phi_e 的所有输入分量分毫不差,计算出的边消息具备严格的不变性:m~ij=mij\tilde{\mathbf{m}}_{ij} = \mathbf{m}_{ij}

其次考察节点标量特征的不变性。由于更新函数 ϕh\phi_h 仅依赖于输入标量 hih_i 与相邻边消息的代数求和 jmij\sum_j \mathbf{m}_{ij},而这两部分均已确立了对刚体变换的不变性,因此节点标量特征的演进自然严格不变,即 h~i=hi\tilde{h}_i' = h_i'

最后完成坐标更新等变性的核心闭环。记标量加权权重为 wij=ϕx(mij)w_{ij} = \phi_x(\mathbf{m}_{ij})。由于消息 mij\mathbf{m}_{ij} 保持不变,权重 wijw_{ij} 是一个与坐标系选择无关的纯标量。将变换后的坐标代入坐标更新方程中展开:

x~i=x~i+jN(i)(x~ix~j)wij=Rxi+t+jN(i)R(xixj)wij.\tilde{x}_i' = \tilde{x}_i + \sum_{j\in\mathcal{N}(i)} (\tilde{x}_i - \tilde{x}_j)\,w_{ij} = Rx_i + t + \sum_{j\in\mathcal{N}(i)} R(x_i - x_j)\,w_{ij}.

在此处,正是因为 wijw_{ij} 是纯标量,线性正交变换算子 RR 得以被直接从求和号内部提至外部:

x~i=R(xi+jN(i)(xixj)wij)+t=Rxi+t.\tilde{x}_i' = R\Big(x_i + \sum_{j\in\mathcal{N}(i)}(x_i - x_j)\,w_{ij}\Big) + t = R x_i' + t.

由此严格确立了坐标输出的刚体等变性。\square

这个优雅证明再次强调了深层原理:标量输出保证了算子与旋转操作的可交换性。EGNN 正是通过这一极其克制的代数设计,在彻底规避高阶球谐张量的前提下,达成了严密的等变保证。

6.3 离散动量守恒与质心绝对静止定理

除了外在的旋转等变性,EGNN 在物理世界中最具魅力的特征在于其天然内嵌了经典力学中的内力抵消与质心守恒律。

当我们对分子体系中所有原子的坐标位移量实施全空间全局求和时,一个优美的物理对称性浮现出来:

i=1NΔxi=i=1NjN(i)(xixj)wij=0.\sum_{i=1}^N \Delta x_i = \sum_{i=1}^N \sum_{j\in\mathcal{N}(i)} (x_i - x_j)\,w_{ij} = 0.

这一求和恒等于零的充分必要前提在于:标量位移权重关于成对原子节点是对称的,即 wij=wjiw_{ij} = w_{ji}。其物理机制在几何上与牛顿第三定律如出一辙:在无向分子图中,任意一对相互作用原子 (i,j)(i, j) 在双向遍历求和中必然成对出现。对于端点 ii,邻居 jj 施加的相对位移贡献为 +(xixj)wij+(x_i - x_j)w_{ij};而对于端点 jj,端点 ii 施加的相对位移贡献为 +(xjxi)wji+(x_j - x_i)w_{ji}。当且仅当 wij=wjiw_{ij} = w_{ji} 时,有 +(xjxi)wji=(xixj)wij+(x_j - x_i)w_{ji} = -(x_i - x_j)w_{ij},两者在全图累加中方能严格精确对消。

工程实现注记:如何确保 wij=wjiw_{ij} = w_{ji} 在标准有向消息传递中,若原子节点特征 hihjh_i \neq h_j,通过非对称网络 ϕe(hi,hj,dij2)\phi_e(h_i, h_j, d_{ij}^2) 计算出的单向权重通常并不天然相等(wijwjiw_{ij} \neq w_{ji})。为了确保严格的动量守恒,通常有两种工程实现手段:

  1. 消息函数显式对称化:在输入端强制采用对称特征,例如 ϕe(hi+hj,hihj,dij2)\phi_e(h_i + h_j, |h_i - h_j|, d_{ij}^2)
  2. 按无向边成对更新(本教程代码实现方式):在底层代码 [code/enn.py](file:///Users/polyai/Desktop/飞书/等变神经网络教程/code/enn.py#L805-L824) 中,算法显式遍历无向边集合 edges,为每条无向边计算唯一的标量标度 w=ϕx(me)w = \phi_x(m_e),随后同时执行 x[i] += w * (x[i] - x[j])x[j] -= w * (x[i] - x[j])。这种成对相反的更新范式从图遍历底层强行锁死了 wijwjiw_{ij} \equiv w_{ji},使得离散质心守恒在浮点精度下坚如磐石。

这给深度几何网络带来了不可估量的数值稳定性:多层 EGNN 的连续前向传递,仅仅是在分子内部连续调整各化学键的相对局域构象,而绝不会引起整个分子团簇在空间中产生任何虚假的整体平移漂移。在分子动力学与构象生成任务中,这一性质彻底根除了因质心漂移导致的空间距离分布整体失真。

6.4 坐标更新的两种范式与其解析求导

在将 EGNN 应用于分子动力学与力场建模时,坐标更新公式的具体数学形态存在两种截然不同的工程选择。

第一种方案是文献中最初采用的原始差分形式(写法 A)

Δxi=jN(i)(xixj)wij.\Delta x_i = \sum_{j\in\mathcal{N}(i)} (x_i - x_j)\,w_{ij}.

在这种形式中,位移矢量的大小直接正比于原子间的初始空间距离 xixj\|x_i - x_j\|。这在物理直觉上是反常的——相距越远的两个原子,其相互施加的空间位移反而越大;在数值训练中,这极易形成正反馈发散。

第二种方案是现代几何库普遍采纳的单位方向归一化形式(写法 B)

Δxi=jN(i)xixjxixj+ϵwij=jN(i)r^ijwij.\Delta x_i = \sum_{j\in\mathcal{N}(i)} \frac{x_i - x_j}{\|x_i - x_j\| + \epsilon}\,w_{ij} = \sum_{j\in\mathcal{N}(i)} \hat{r}_{ij}\,w_{ij}.

在写法 B 中,位移的几何方向被严格限制为空间单位矢量 r^ij\hat{r}_{ij},其实际位移幅度完全由标量网络输出的权重 wijw_{ij} 所精确调控。这极大地抑制了远距离非键相互作用引发的数值激变。

在手写反向传播时,正确推导写法 B 的坐标梯度链式法则是力场训练能否成功的关键。设来自上游的坐标梯度为 gi=L/xig_i = \partial\mathcal{L}/\partial x_i'。对于写法 A,坐标与权重的梯度形式较为直接:

Lxi+=giwij,Lwij=(gigj)(xixj).\frac{\partial\mathcal{L}}{\partial x_i} \mathrel{+}= g_i\,w_{ij},\qquad \frac{\partial\mathcal{L}}{\partial w_{ij}} = (g_i - g_j)\cdot(x_i - x_j).

而对于方向归一化的写法 B,求导必须深入单位矢量的微分结构内部。记相对位移矢量为 r=xixjr = x_i - x_j,距离标量为 d=rd = \|r\|,正则化单位向量为 u=r/(d+ϵ)u = r/(d+\epsilon)。求向量 uu 对位移分量 rlr_l 的一阶偏导数:

ukrl=δkld+ϵrkrld(d+ϵ)2.\frac{\partial u_k}{\partial r_l} = \frac{\delta_{kl}}{d+\epsilon} - \frac{r_k r_l}{d\,(d+\epsilon)^2}.

式中第一项来自分子端矢量 rr 的直接线性求导,第二项则来自分母中距离模长 d=mrm2d = \sqrt{\sum_m r_m^2} 对坐标分量的非线性连锁求导。将此求导矩阵与上游梯度 gg 缩并,我们得到写法 B 中雅可比转置乘积算子 JgJ^\top g 的显式闭式解:

(Jg)l=gld+ϵrl(rg)d(d+ϵ)2.\big(J^\top g\big)_l = \frac{g_l}{d+\epsilon} - \frac{r_l\,(r\cdot g)}{d\,(d+\epsilon)^2}.

在核心实现 EGNNLayer.backward 中,该公式正是分支变量 JgD 的数学原型。此处隐藏着一个极为隐蔽的工程暗坑:在计算标量权重梯度 L/wij\partial\mathcal{L}/\partial w_{ij} 时,参与内积运算的梯度矢量 gig_i 必须是刚刚穿透坐标输出层、未经任何其他边累加污染的纯净上游梯度;若直接使用已被其他邻接边就地累加过的坐标梯度数组,计算出的几何导数将在深层网络中彻底崩溃。

6.5 数值稳定性实践:真实发散案例与调优心法

教科书上的几何公式往往预设了无限平滑的数学理想态,但在真实的化学数据拟合中,数值发散往往如影随形。在本书第 10.6 节的高聚物势能面与原子力场联合拟合实验中,我们曾使用原始的写法 A 搭配常规的 10310^{-3} 学习率进行训练,记录下了一场极为经典的真实发散过程:

训练轮次 Loss 演进轨迹:
Epoch 1~5:  1.252  1.244  1.245  1.245  1.247 (看似收敛平稳)
Epoch 6~8:  1.314  3.688  18.935 (瞬间梯度雪崩)
最终测试能量 MAE: 303.6 (约合 1.8×10⁴ kcal/mol,彻底发散)

这场发散的物理机制极为清晰:在训练迭代中,坐标更新层在物理上是在真实移动分子的空间构象。一旦某个批次的梯度扰动导致标量权重 wijw_{ij} 出现偶发性放大,某些成对原子就会被强行推向极端极近的距离。当原子间距 d0d \to 0 时,高斯径向基函数与库仑型项的一阶导数急剧飙升,激发出天文数字级别的反传梯度,从而在"更近的原子距离 \to 更大的位移梯度 \to 更加病态的构象畸变"之间形成了致命的正反馈雪崩。

平息这场数值风暴的核心手段包含三条联动的工程法则:

首先是将坐标更新彻底切换为写法 B(方向归一化),将每一步的几何位移上限刚性约束在单位模长之内;

其次是将多层感知机中负责输出坐标位移权重 ϕx\phi_x最后一层线性权重与偏置实施零初始化(Zero Initialization)。这一极具巧思的设计确保了在训练伊始的前几个轮次内,坐标更新幅度恒等于零,分子构象保持完全静止,从而让负责提取化学环境的标量网络 ϕe\phi_eϕh\phi_h 率先学习到平稳合理的局域表征;

最后是引入无量纲的坐标位移缩放因子 CC(在文献中通常取为平均原子配位数的倒数 C=1/(M1)C = 1/(M-1)),并辅以平缓的学习率(如 3×1043\times10^{-4})。

在采纳这套稳定化组合拳后,相同数据集上的拟合表现发生了质的蜕变:

引入方向归一化与零初始化后的训练演进:
Loss 轨迹:  1.252  1.244  1.245  1.244  1.245 ... (极度平稳)
测试能量 MAE: 1.20 (恢复到优异的物理拟合精度)

而在完全剔除坐标更新的分支对照实验中,模型给出的测试 MAE 为 1.36。这表明:在小样本学习场景中,坐标动态演化固然能带来额外的几何自适应收益,但前提必须是以极其克制的数值稳定算子为其保驾护航

6.6 空间搬运原子与旋转局域张量的范式抉择

至此,我们可以从物理哲学的高度将 EGNN 与高阶张量网络(TFN/MACE)进行对照。

EGNN 的等变逻辑是"在空间中搬运原子"——它保持隐层特征纯粹为各向同性的标量,将所有的方向性信息全部具象化为空间三维坐标矩阵的动态重构。这种范式实现极为轻盈,计算开销保持在优雅的 O(Nk)O(Nk),并且天然适合充当扩散生成模型(Diffusion Models)中预测原子位移漂移项的骨架。

而 TFN 与 MACE 的等变逻辑则是"在每个原子上旋转一组多极矩箭头"——原子在空间中的坐标骨架可以保持不动,但各个原子节点携带着一整套不可约表示张量基底,几何信息的交互通过高阶多极展开的 Clebsch–Gordan 耦合实现。这种范式具有近乎无懈可击的高维角向分辨力,但其实现与计算开销也高出一个数量级。对于计算化学研究者而言,明确这两条路线的物理本质,便能依据具体体系对精细势能面或快速动力学生成的不同需求,做出最精准的架构抉择。


7. 反向传播:手写每一行

在现代深度学习中,PyTorch 与 JAX 等自动微分引擎(Autograd)已经普及为极其便利的基础设施。然而,在深入几何等变网络的核心推导时,手写反向传播依然是不可替代的训练

自动微分机制具有先天的盲目性:它会忠实地反向传播研究者写错的任何前向公式。在等变网络中,诸如二阶张量积的转置缩并顺序、双向边消息的反传归并、以及由能量梯度衍生出牛顿力场的二次反传,皆是常规自动求导极其容易悄然通过、却在物理机制上彻底错位的重灾区。唯有亲手推导梯度的解析表达式,并通过双精度有限差分的逐参数数值梯度检查(Gradient Check),才能在代数层面上确证每一行几何算子的绝对正确。

7.1 计算图设计约定与内存状态解耦

为了确保手写反向传播的机械严密性与多任务复用性,本教程在核心代码库中树立了三条坚实的软件工程准则:

在前向传播端,每一个模块的 forward 方法统一返回二元元组 (output, ctx)。其中 output 为传递给下一层的特征张量,而 ctx 则是一个专门打包的上下文状态字典,严格留存当前层在求导时所必需的中间量(包括输入状态、仿射变换净输入、激活前后值以及预先计算好的成对几何距离)。这种函数式无状态设计使得同一个神经网络实例可以被并发应用于不同批次的构象计算,彻底杜绝了内部缓存被意外覆盖的隐患。

在反向传播端,每一个模块的 backward(ctx, grad_out) 接收上游误差梯度 grad_out,计算并返回对本层输入的解析梯度,同时以显式累加+=)而非覆盖的方式将参数梯度沉淀至 Param.g 中。累加机制至关重要,它不仅支撑了单批次中多个分子构象梯度的平滑聚合,更为后续利用二次反向传播实现原子力场匹配(Force Matching)提供了底层的接口保证。

在优化器端,参数梯度的清零、全局范数裁剪以及一阶二阶动量更新全部交由独立的优化引擎处理,从而将物理求导的几何力学与随机梯度下降的数值更新彻底解耦。

7.2 全连接层与 SiLU 激活函数的链式法则

作为所有网络大厦的基础砖瓦,我们首先快速对齐标准仿射全连接层与其激活函数的求导方程。对于前向仿射变换 z=Wx+bz = Wx + b 以及随后施加的非线性激活映射 a=ϕ(z)a = \phi(z),设来自下游的误差敏感度向量为 ga=L/ag_a = \partial\mathcal{L}/\partial a。根据多元微积分链式法则,误差向前流动的基本方程组如下:

δ=gaϕ(z),LW=δx,Lb=δ,Lx=Wδ.\delta = g_a \odot \phi'(z),\qquad \frac{\partial\mathcal{L}}{\partial W} = \delta\,x^\top,\qquad \frac{\partial\mathcal{L}}{\partial b} = \delta,\qquad \frac{\partial\mathcal{L}}{\partial x} = W^\top\delta .

在几何图网络中,**Sigmoid 加权线性单元(SiLU,又称 Swish 激活函数)**的表现在实践中显著优于经典的 ReLU。SiLU 的数学形式定义为自变量与其自身 Sigmoid 响应的乘积:

ϕ(z)=zσ(z)=z1+ez.\phi(z) = z\,\sigma(z) = \frac{z}{1 + e^{-z}}.

求导展开可得其显式一阶导函数:

ϕ(z)=σ(z)+zσ(z)=σ(z)+zσ(z)(1σ(z))=σ(z)(1+z(1σ(z))).\phi'(z) = \sigma(z) + z\,\sigma'(z) = \sigma(z) + z\,\sigma(z)\big(1 - \sigma(z)\big) = \sigma(z)\Big(1 + z\big(1 - \sigma(z)\big)\Big).

与 ReLU 在原点处的粗暴硬折折角不同,SiLU 在原点附近保持了平滑连续的非零导数(ϕ(0)=0.5\phi'(0) = 0.5),这使得成对原子间极其微弱的几何摄动能够产生平滑连续的物理力响应;其唯一的计算代价仅仅是每次反向求导时需要额外执行一次 Sigmoid 运算。

7.3 张量积层的多通道解析反传回顾

在第 4.7 节中,我们已经对包含 l=0l=0 标量缩并、l=1l=1 矢量叉乘和 l=2l=2 对称无迹张量的张量积反传进行了严密推导。在此处需要从代码架构的角度重申其独特性:这三个几何通道的反向梯度在代数拓扑上是完全不可合并的

标量通道的反传由外积点乘所驱动,矢量通道的反传本质上是利用反对称伴随矩阵重新组织的空间叉乘,而二阶张量通道的反传则由对称化梯度张量与互补输入矢量的矩阵向量乘法主导。在编写底层代码时,必须对这三个通道开辟独立的分支函数并执行逐分量求导;任何试图强行使用统一密集矩阵乘法替代的做法,都会在数学上破坏坐标旋转等变性,并在双精度梯度检验中引发剧烈报警。

7.4 EGNN 层的全反向传播:五大路径系统拆解

EGNN 层的反向传播是全套纯 Python 代码库中工程拓扑最为繁复的章节。来自下游网络的反向梯度同时包含着节点标量特征敏感度 ghRN×Fg_h \in \mathbb{R}^{N\times F}原子坐标敏感度 gxRN×3g_x \in \mathbb{R}^{N\times 3} 两股洪流。要将误差正确分解并回传给输入坐标与节点状态,梯度逆流必须严格按以下五条路径梯次推进:

第一路径:从输出坐标梯度反推位移权重与方向项。为了彻底杜绝第 6.4 节指出的累加污染,我们必须在进入任何边循环之前,将进入本层的原生上游坐标梯度完全冻结并克隆出一份独立副本 gxupg_x^{\text{up}}。随后遍历图中的每一条空间连通边 (i,j)(i, j),分别计算坐标位移加权标量网络 ϕx\phi_x 的目标梯度 gwijg_{w_{ij}} 与空间位移方向向量的敏感度 gΔijg_{\Delta_{ij}}

gwij=(gxiupgxjup)Δij,gΔij=wij(gxiupgxjup).g_{w_{ij}} = \big(g_{x_i}^{\text{up}} - g_{x_j}^{\text{up}}\big)\cdot \Delta_{ij}, \qquad g_{\Delta_{ij}} = w_{ij}\big(g_{x_i}^{\text{up}} - g_{x_j}^{\text{up}}\big).

对于归一化写法 B,矢量敏感度 gΔijg_{\Delta_{ij}} 必须进一步穿透第 6.4 节推导出的雅可比变换算子 JgJ^\top g,最终沉淀为对原始相对位移矢量 r=xixjr = x_i - x_j 的基础梯度贡献。

第二路径:节点标量更新函数的反向传导。在标量更新方程 hi=hi+ϕh(hi,misum)h_i' = h_i + \phi_h(h_i, \mathbf{m}_i^{\text{sum}}) 中,来自上游的节点标量梯度 ghig_{h_i} 一方面直接通过残差连接无损回传给输入端,另一方面作为敏感度注入多层感知机 ϕh\phi_h 的反向求导中,解析激发出对本节点输入状态的中间导数以及对聚合消息和 misum\mathbf{m}_i^{\text{sum}} 的梯度 gmisumg_{\mathbf{m}_i^{\text{sum}}}

第三路径:邻居消息聚合算子的转置转置投影。在图网络的前向传播中,节点聚合 misum=jmij\mathbf{m}_i^{\text{sum}} = \sum_{j} \mathbf{m}_{ij} 是典型的离散散开(Scatter)算子。在微积分的反向传播中,散开算子的伴随转置操作正是聚集(Gather)。由于每条空间无向边 (i,j)(i, j) 在前向过程中同时为端点 ii 与端点 jj 贡献了局域消息,因此该边消息所承受的总梯度必然同时汇聚两个端点的聚合敏感度:

gmij+=gmisum+gmjsum.g_{\mathbf{m}_{ij}} \mathrel{+}= g_{\mathbf{m}_i^{\text{sum}}} + g_{\mathbf{m}_j^{\text{sum}}}.

第四路径:边消息神经网络的逆向分解。将全图汇聚后的边消息梯度 gmijg_{\mathbf{m}_{ij}} 灌入边神经网络 ϕe\phi_e 的反向传播引擎中,得到对该网络三个拼接输入分量的联合梯度:其中对端点节点特征 hih_ihjh_j 的偏导数直接分别累加至全局节点梯度数组 ghg_{h} 中,而对成对距离平方项的偏导数 gd2=L/xixj2g_{d^2} = \partial\mathcal{L}/\partial \|x_i - x_j\|^2 则被完整提取,留待最后的几何链式收尾。

第五路径:由成对几何间距向笛卡尔坐标的终极连锁反传。这是初学者在手写几何神经网络时最常遗漏的幽灵路径:成对原子间的空间距离平方项本身是坐标的非线性函数。应用初等微积分求导:

xixj2xi=2(xixj),xixj2xj=2(xixj).\frac{\partial\,\|x_i-x_j\|^2}{\partial x_i} = 2(x_i - x_j),\qquad \frac{\partial\,\|x_i-x_j\|^2}{\partial x_j} = -2(x_i - x_j).

因此,距离敏感度 gd2g_{d^2} 必须通过连心矢量差 2(xixj)2(x_i - x_j) 转化为终极的笛卡尔受力,并分别以相反符号追加回传给节点 xix_ixjx_j。如果漏掉这最后一步,神经网络在纯标量滤波参数上的梯度检验将呈现出巨大的相对偏差,而其他密集矩阵参数的梯度却表现正常——这种局部错位本身就是诊断几何求导遗漏的典型指南。

7.5 势能面力场匹配:二次反向传播的高阶技巧

在计算材料学与高分子动力学中,最核心的任务莫过于构建机器学习相互作用势(MLP)。根据分析力学,保守系统中每个原子受到的真实相互作用力矢量等于体系总势能对原子坐标的负空间梯度:

Fi=Exi.F_i = -\frac{\partial E}{\partial x_i}.

在训练机器学习力场时,单纯拟合单点能量 EE 远远不足以约束整个势能面的局部曲率;我们必须采用**能量与原子力联合拟合(Force Matching)**策略,最小化包含受力均方误差的联合损失函数:

L=LE+LF=1σE2(EE)2+wFNσF2i=1NFiFi2.\mathcal{L} = \mathcal{L}_E + \mathcal{L}_F = \frac{1}{\sigma_E^2}(E - E^*)^2 + \frac{w_F}{N\sigma_F^2}\sum_{i=1}^N \|F_i - F_i^*\|^2.

此时,一个重大的数学挑战浮现在面前:模型预测的力 FiF_i 本身就是网络能量输出对输入坐标 xx 的一阶导数,因此计算损失函数对网络内部权重参数 θ\theta 的梯度 LF/θ\partial\mathcal{L}_F/\partial\theta,在形式上涉及到对网络参数与空间坐标的二阶混合偏导数 2Exiθ\frac{\partial^2 E}{\partial x_i \partial\theta}

在没有自动微分框架辅助的纯 Python 底层,我们难道需要手工解析推导并编写整座网络所有神经元的二阶 Hessian 矩阵吗?答案是完全不需要。借助微积分的伴随敏感度法,我们可以运用一次极为巧妙的**二次反向传播(Dual Backward)**技术,以极简的代码代价优雅解决这一难题。

推导逻辑极为优美。记模型对第 ii 个原子的受力预测为 Fi=gi(θ)F_i = -g_i(\theta),其中 gi(θ)=E/xig_i(\theta) = \partial E/\partial x_i 为能量前向计算在坐标入口处的中间导数。应用多元复合求导法则:

LFθ=i=1NLFFiFiθ=i=1NLFFigi(θ)θ.\frac{\partial\mathcal{L}_F}{\partial\theta} = \sum_{i=1}^N \frac{\partial\mathcal{L}_F}{\partial F_i}\cdot \frac{\partial F_i}{\partial\theta} = -\sum_{i=1}^N \frac{\partial\mathcal{L}_F}{\partial F_i}\cdot \frac{\partial g_i(\theta)}{\partial\theta}.

关键的顿悟发生在这一刻:表达式 i(LFFi)θ(Exi)\sum_i \Big(-\frac{\partial\mathcal{L}_F}{\partial F_i}\Big)\cdot \frac{\partial}{\partial\theta}\big(\frac{\partial E}{\partial x_i}\big) 在数学结构上,精确等价于对完全相同的计算图重新执行一次标准的反向传播,唯一的改变是将其输出端的标量能量梯度入口置零,而将物理力的残差矢量作为虚拟的外生敏感度种子,直接灌注到原子的空间坐标入口

这个注入坐标入口的虚拟敏感度种子具有极其简单的闭式解析解:

seedxi=LFFi=2wFNσF2(FiFi).\text{seed}_{x_i} = -\frac{\partial\mathcal{L}_F}{\partial F_i} = \frac{2w_F}{N\sigma_F^2}\big(F_i^* - F_i\big).

由此,我们建立起一套极为高效的力场双反向训练范式,在代码中仅需寥寥数行:

python
# 第一次反向传播:由能量标量损失驱动,既累加能量项参数梯度,又顺手得到原子受力预测
gx = model.backward(state, [g_e] * n)
f_pred = [[-gx[i][c] for c in range(3)] for i in range(n)]

# 第二次反向传播:将受力残差向量构造成种子,直接灌入坐标入口反传力的损失
seed_x = [[2.0 * w_f * (f_pred[i][c] - f_true[i][c]) / (scale_f**2) for c in range(3)] for i in range(n)]
model.backward(state, [0.0] * n, gx_seed=seed_x)

两次反向传播完全共享同一次前向前向所缓存的几何上下文 state,参数梯度在底层自动实现线性和叠加。这种架构在计算开销上仅仅耗费了约两倍的单次反向耗时,便在纯 Python 环境中极其严谨、优雅地实现了第一性原理力场的精准匹配。

7.6 有限差分梯度检查的理论与判据

在完成所有手写梯度的编码后,如何确保我们的推导与代码没有隐藏任何微小漏洞?有限差分梯度检查(Numerical Gradient Checking)是唯一的法官。

梯度检查采用截断误差极小的对称中心差分格式

LθkL(θk+ε)L(θkε)2ε=Lθk+ε263Lθk3+O(ε4).\frac{\partial \mathcal{L}}{\partial \theta_k}\approx \frac{\mathcal{L}(\theta_k+\varepsilon)-\mathcal{L}(\theta_k-\varepsilon)}{2\varepsilon} = \frac{\partial\mathcal{L}}{\partial\theta_k} + \frac{\varepsilon^2}{6} \frac{\partial^3\mathcal{L}}{\partial\theta_k^3} + O(\varepsilon^4).

与单侧前向差分相比,中心差分由于奇数阶泰勒展开项完全对消,其截断截断误差严格收敛在 O(ε2)O(\varepsilon^2) 级别。在实际选择扰动步长 ε\varepsilon 时,必须在微积分截断误差与计算机浮点舍入误差之间达成严密平衡:若 ε\varepsilon 取值过大,高阶导数项将引发显著截断失真;若取值过小(如接近双精度机器精度 εmach1016\varepsilon_{\text{mach}} \approx 10^{-16}),微小分子差值除以微小分母将引发严重的有效数字灾难性丢失。在理论数值分析中,中心差分的最佳综合步长为 ε(εmach)1/36×106\varepsilon \approx (\varepsilon_{\text{mach}})^{1/3} \approx 6\times10^{-6}。我们在实现中统一设置自适应步长 ε=106max(1,θk)\varepsilon = 10^{-6}\cdot\max(1, |\theta_k|)

在衡量解析导数 ganalyticg_{\text{analytic}} 与差分导数 gnumericg_{\text{numeric}} 的吻合度时,必须使用相对误差指标

Relative Error=ganalyticgnumericmax(1,ganalytic,gnumeric).\text{Relative Error} = \frac{|g_{\text{analytic}} - g_{\text{numeric}}|}{\max(1,\,|g_{\text{analytic}}|,\,|g_{\text{numeric}}|)}.

分母中引入下限 11 巧妙规避了当参数导数本身极度趋近于零时引发的假阳性误报警。在双精度纯 Python 环境下,我们将单元测试套件的断言门槛设定为严格的 2×1042\times10^{-4}。全套测试集 tests_enn.py 包含多达 107 项严格断言,对包含张量积、高斯滤波核与坐标更新在内的上千个模型参数逐一进行扰动检验,从而在数值上百分之百确证了梯度引擎的绝对刚健。

7.7 研发历程中踩中并排查的四大典型暗坑

在构建这套无依赖等变网络底层的过程中,即便理论推导完全正确的学者,在工程转换时也难免遭遇暗礁。本节将四个在研发实战中真实踩中、唯有通过梯度检查与等变体验才最终揪出的真实 Bug 完整复盘,为后来者提供珍贵的实战病理图谱:

第一暗坑:坐标累加器的原地指针污染(相对误差高达 1.83)。在前向传播中,坐标更新是全图所有邻居边的叠加;而在最初编写反向传播代码时,为了追求运行效率,我直接将上游传入的梯度数组用作当层梯度的累加容器。当程序处理第一条边时,计算出的中间量被就地累加到了数组中;而当循环处理到同一节点的第二条相邻边时,其读取到的"上游梯度"实际上已经混入了前一条边产生的微变贡献,导致第二条边计算出的标量权重梯度 gwg_w 彻底谬以千里。典型病理特征:单个独立层的梯度检验完全通过,而一旦在多层网络中将两层 EGNN 级联堆叠,测试立刻崩溃报警,且层数越深偏差越呈几何级数放大。根治方案:在进入计算循环的第一行,强制执行深拷贝冻结独立的上游梯度镜像。

第二暗坑:l=2l=2 张量积通道的输入对偶颠倒(相对误差 1.39)。在推导二阶无迹张量对第一组输入矢量 aa 的反向传播梯度时,式中配准的伴随矢量必须是第二组输入 bb。在初版代码中,我一时疏忽顺手将其写成了输入矢量 aa 自身的二次投影。典型病理特征:在早期的单元测试中,我们往往图方便取同一个随机向量作为双输入(A=BA = B),此时该 Bug 被完美的代数对称性悄然掩盖,单元测试全绿通过;而一旦在真实网络中两个输入通道来自不同特征提取分支时,梯度检查立刻剧烈报警。根治方案:严密核对 A/BA/B 索引配对,并在单元测试中强制增加一组两个输入张量维度不同且数值独立的专项梯度断言。

第三暗坑:读出层激活值与仿射净输入的代数混淆。在对输出端的 Softplus 或 SiLU 激活层实施反向求导时,代码在从上下文字典 ctx 提取变量时误将激活后的输出值当作未经非线性的净输入 zz 送入求导函数。典型病理特征:单层梯度检验在某些初始权重下可以偶然通过,而在另一些权重大小时误差大幅漂移。根治方案:在底层严格建立具有明确命名空间的上下文结构元组 (x_input, z_affine, a_activated),杜绝依赖下标盲猜上下文。

第四暗坑(最致命的架构设计陷阱):输入相同导致的矢量通道零退化。在搭建等变偶极矩预测网络时,若直接将同一组球谐特征不加区分地作为张量积的两路输入,根据外积反对称性的公理,三维空间中同一个矢量的自身叉积恒等于零(a×a0a\times a \equiv 0)。这意味着网络中的 l=1l=1 矢量特征通道在前向传播的第一步就被数学公理彻底清零,模型只能被迫退化为输出恒等于零的常数解。更为致命的是,此时整个网络的损失函数依然能够平滑下降,训练流程看似完全正常,因为网络迅速学会了输出零向量这个最平庸的平凡解!在初版偶极矩实验中,我曾惊喜地观察到模型报出了方向误差为 "0.00°" 的完美假象——其本质原因仅仅是零向量的方向角在数学上未定义而被评估脚本直接略过。根治方案:在矢量张量积的两路输入端前置独立可学习的通道混合投影层(pre_mix),打破输入对称性,并在实验评估框架中开辟专门的"平凡退化预测比例"监控指标。

这四大暗坑的共同启示在于:在等变几何学习中,单纯观察训练损失曲线平稳下降是极具欺骗性的。唯有建立起包含逐参数梯度检验、随机旋转等变体检以及预测解退化检测在内的全方位防御体系,才能确信模型在代数与物理意义上的绝对健康。


8. 纯 Python 实现:公式 ↔ 代码

在计算化学与材料建模的日常工作中,研究者们早已习惯了各类高度封装的软件工具:从量子化学程序 Gaussian、VASP 到经典分子动力学包 LAMMPS、GROMACS,抑或是化学信息学库 RDKit 与 ASE。当目光转向几何深度学习时,业界主流的生产级框架如 e3nn、NequIP 或 MACE 同样构建在厚重的底层抽象之上——动态图构造、C++/CUDA 算子融合、隐藏在自动微分引擎内部的向量雅可比积。对于希望洞悉其物理本质的研究者而言,这种封装往往如同一堵高墙,将清晰优美的群论几何机制遮蔽在层层张量转置与广播规则之后。

为了彻底打通从黑板上的数学推导到计算机内存中的数值运算,我们在教程附带的代码库 enn.py 中确立了三条严苛的教学设计原则。代码的组织形式与本教程的章节目录完全对应,每一个函数都带有指向对应小节的详细注释,让理论与实现形成镜像般的互文关系。

第一条原则是彻底剥离数值库,拥抱纯原生 Python。在整个核心实现中,我们没有引入 NumPy、SciPy 或 PyTorch,三维空间坐标就是长度为 3 的浮点数原生列表,二维矩阵是简单的列表嵌套列表,而 Clebsch-Gordan 张量积则被拆解为透明的三层嵌套循环。这一选择绝非出于技术复古的偏好,而是为了践行一种极致的教学承诺:数学公式里写了哪一项,代码里就有一目了然的对应行。这里没有黑盒般的多维广播机制,没有难以追踪的隐式内存重排与转置,也没有编译器自动推断张量形状时的不确定性,算法的每一步时空复杂度与数据流向都赤裸可见。

第二条原则是由第一性原理出发纯手写反向传播。每一个几何层与张量运算模块的 backward 方法,都严格按照微积分链式法则手动计算输入梯度并显式累加参数梯度。这种坚持不仅消除了对自动微分引擎的心理依赖,更让力的解析计算——即能量对原子坐标的负梯度——成为一种触摸得到的物理操作,并与严格的数值有限差分梯度检查无缝咬合。

第三条原则是全流程的可复现性与确定性。所有的伪随机数生成均绑定至显式声明的随机数生成器实例,杜绝了模块级全局随机状态带来的隐蔽扰动。每一个实验脚本都会将原始数值以结构化 JSON 的形式固化落盘,文中的所有图表均由独立脚本直接读取这些 JSON 数据绘制而成,确保任何一位读者在个人笔记本电脑上都能在数分钟内重演推导出的每一个结论。

8.1 公式 ↔ 代码总表

为了方便读者在阅读理论推导与翻阅源码之间自由穿梭,下表将前述章节的核心几何算子、数学恒等式与 enn.pychem3d.py 及测试脚本中的具体实现进行了严格的坐标锚定:

公式位置代码
R=I+sinθ[k^]×+(1cosθ)[k^]×2R=I+\sin\theta[\hat k]_\times+(1-\cos\theta)[\hat k]_\times^2§2.2enn.rotation_matrix
四元数均匀采样§2.2enn.random_rotation
(Ru1×Ru2)Ru3=det(M)()(Ru_1\times Ru_2)\cdot Ru_3=\det(M)\,(\cdots)§3.4enn.reflect, tests_enn.test_geometry
Dij2=Gii+Gjj2GijD_{ij}^2=G_{ii}+G_{jj}-2G_{ij}§3.3enn.distance_matrix / gram_matrix
ek(r)=exp(γ(rμk)2)e_k(r)=\exp(-\gamma(r-\mu_k)^2)§5.1enn.rbf_expand / rbf_centers
T=symtraceless(ab)T=\mathrm{symtraceless}(a\otimes b)§4.3enn.sym_traceless_outer
a2b2=(ab)23+a×b22+TF2|a|^2|b|^2=\frac{(a\cdot b)^2}{3}+\frac{|a\times b|^2}{2}+|T|_F^2§4.4tests_enn.test_tensor_product_math
sk,vk,Tks_k,v_k,T_k 与前向等变§4.3enn.TensorProduct11.forward
/a=Sym(G)btrG3b\partial/\partial a=\mathrm{Sym}(G)b-\frac{\mathrm{tr}G}{3}b§4.7enn.TensorProduct11.backward
mij,hi,xi\mathbf m_{ij},h_i',x_i'§6.1enn.EGNNLayer.forward
JgJ^\top g(方向归一化的坐标梯度)§6.4enn.EGNNLayer.backwardJgD 分支
hi=jmijWhjh_i'=\sum_j\mathbf m_{ij}\odot Wh_j§5.1enn.SchNetLayer.forward
seedi=2wFσF2(FiFi)\text{seed}_i=\frac{2w_F}{\sigma_F^2}(F_i^*-F_i)§7.5demo_polymer_uff.train_epoch_ef
μ=iqi(rir质心)\mu=\sum_i q_i(r_i-r_{\text{质心}})§10.4chem3d._gasteiger_dipole
r=xi(xj+offset)r=x_i-(x_j+\text{offset})§10.7enn.EGNNLayer.forwardoffsets 分支

8.2 数据结构与邻接图表示

在纯 Python 的语境下,一个分子系统或周期性晶胞被凝练为一个基础的字典结构:

python
mol = {
    "x": [[x, y, z], ...],      # 原子三维笛卡尔坐标,单位 Å
    "z": [元素编号, ...],        # 映射到 0..K 的连续原子序数
    "edges": [(i, j), ...],     # 相互作用无向边索引
    "offsets": [[dx,dy,dz], ...] # 可选:周期性镜像超胞位移(孤立分子体系省略)
}

在将化学体系转化为图神经网络的输入时,边的构建策略直接反映了建模者对底层物理规律的取舍。第一种思路是依据传统化学键理论,利用共价键连接原子。这种图结构高度稀疏且完美契合分子的路易斯结构式与有机合成直觉,但它的致命软肋在于将模型对外部环境的感知局限在拓扑相连的原子骨架上,完全屏蔽了主导凝聚态物理与大分子构象的非共价相互作用——如范德华吸引、氢键网络以及长程静电堆积。

第二种思路则是基于欧氏空间的距离截断,将指定截断半径 rcr_c(例如 5.0 Å)之内的所有原子对都纳入交互边缘。这正如经典分子力学中对非键相互作用所设定的势能截断,使网络能够天然捕捉三维空间构型诱导的所有近程物理效应。在我们的实现中,默认采用基于欧氏距离的邻域截断方案,并辅以最大近邻数 kk 的上限约束(对应 edges_from_cutoff 函数)。将度数上限截断为常数 kk,成功将大分子系统的边数增长由灾难性的 O(N2)O(N^2) 压制到良性的 O(N)O(N),这是使几何深度学习模型能够扩展至大尺度凝聚态与生物大分子的工业标准范式。

8.3 Adam 优化器与几何坐标诱发的梯度裁剪

为了支撑模型的自主训练,我们手写实现了标准的 Adam 优化器,其核心迭代递推式由动量平滑、方差自适应与参数步长三步组成:

mt=β1mt1+(1β1)gt,vt=β2vt1+(1β2)gt2,θt=θt1ηmt/(1β1t)vt/(1β2t)+ε.m_t=\beta_1m_{t-1}+(1-\beta_1)g_t,\quad v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2,\quad \theta_t=\theta_{t-1}-\eta\frac{m_t/(1-\beta_1^t)}{\sqrt{v_t/(1-\beta_2^t)}+\varepsilon}.

公式中关于偏差修正项 1β1t1-\beta_1^t1β2t1-\beta_2^t 的推导极为直观:由于初态 m0m_0 设定为零向量,在展开历史递推后,mtm_t 的数学期望实质上等于 (1β1t)E[gt](1-\beta_1^t)\,\mathbb{E}[g_t],因此必须除以该衰减系数方能在训练初始阶段消除向零点收缩的冷启动偏差。

在常规深度学习中,梯度裁剪常被视作一种可有可无的调优技巧,但在几何深度学习——尤其是包含坐标动力学演化的架构中,梯度裁剪是保障数值系统不发生物理崩溃的绝对生命线。正如前文多次探讨的,当原子坐标作为连续变量直接参与训练时,一次不合理的参数更新极易将相邻原子推入不切实际的紧密贴合状态。在极端微小的核间距下,几何距离的微小扰动会引发 1/r1/r 乃至更高阶导数的剧烈飙升,导致梯度瞬间暴涨数个数量级。

针对这种病态激增,我们采用全局范数裁剪策略,即先将网络所有权重与偏置的梯度平铺为一个超维向量,仅在其全局 L2L_2 范数突破预设阈值时实施等比例整体收缩。相较于简单粗暴的逐参数独立截断,全局裁剪在物理上保留了能量超曲面上最陡下降方向的矢量完整性,避免因破坏梯度方向而导致优化路径在势能谷底发生灾难性的锯齿震荡。

8.4 权重初始化:为什么何恺明初始化在几何图上会导致能量爆炸

在卷积神经网络与普通 MLP 的经典教材中,何恺明初始化(He Initialization,权重方差设定为 2/nin2/n_{\text{in}})被奉为避免梯度消失的黄金律,其核心假设是信号仅穿过单一线性映射与 ReLU 激活。然而,当这一经验被直接移植到几何图网络中时,系统会迅速陷入灾难性的数值发散。

几何图神经网络具有两个显著不同于平面 CNN 的结构特征:其一是层间普遍嵌入的残差跳跃连接,信号在前向传递中进行直接相加;其二是消息传递阶段对局域配位原子特征的求和聚合操作。在分子和晶体体系中,原子的配位数各不相同——碳原子可能拥有 4 个共价配位,而在面心立方晶体中一个金属原子的配位近邻多达 12 个。对邻居特征的无偏求和会使得聚合特征的方差成倍扩张为平均配位数 k\langle k \rangle 的倍数。当邻域膨胀与残差累加两相叠加,特征方差随网络层数的增加呈指数级爆炸。在早期的原型试验中,直接使用 He 初始化的网络在未经任何训练的情况下,对寡聚体能量的预测值便不可控制地飘到了数百乃至上千 kcal/mol,首个训练批次产生的巨大力梯度当场将整个网络参数推入无底的发散深渊。

解决这一系统性失效的药方是重塑几何网络的方差分配:将内部全连接层的权重方差收缩至 1/nin1/n_{\text{in}},剔除针对普通激活函数预设的翻倍增益;更为关键的是,将最终能量读出层以及坐标更新网络的末层输出权重与偏置严格归零初始化。在零初始化的保护下,初生网络对分子体系的输出基准天然归零,各原子的更新位移也严格为零。分子在初识阶段保持绝对静止,这为网络各层留出了宝贵的平稳期,使其能够优先通过不变的消息通道学出合理的局部化学环境表征,再从容启动坐标的微调迭代。

8.5 训练循环与能量-力双反向传播

在训练执行层,runner.train_epoch 与面向力场训练的 demo_polymer_uff.train_epoch_ef 构成了算法的核心调度中枢。对于常规的分子性质标量预测任务,前向过程返回预测值与网络各层激活状态缓存,损失标量通过链式法则注入 backward 过程;在一个批次内对各样本梯度进行算术平均后,施加全局范数裁剪,最终由优化器执行参数推进。

而当训练目标扩展至物理力场拟合时,系统展现出了极具魅力的物理对称性。正如量子力学中的 Hellmann-Feynman 定理指明外力源于哈密顿量对核坐标的导数,我们在模型中无需额外架设预测力的并行分支,而是直接将原子受力定义为预测总能量对原子坐标的负解析梯度:

Fi=Exi.F_i = -\frac{\partial E}{\partial x_i}.

在这一框架下,力场的训练需要执行双重反向传播。系统首先完成前向计算并针对标量能量项进行常规反传;紧接着,计算预测力矢量与真实物理力矢量之间的残差,将其加权映射为坐标端口的虚拟种子梯度,再次反向传入网络。两次反向传播完全共享同一份前向计算图与激活状态,参数梯度在原位自然叠加。这种优雅的构造不仅将网络参数规模削减了一半,更从根本上赋予了预测力以严格的保守场性质,彻底规避了非保守力场在长时间分子动力学微正则系综(NVE)积分中必然诱发的能量系统性漂移。

8.6 纯 Python 实现的性能代价与工业级加速路径

纯 Python 实现以极致的直观性揭开了几何深度学习的数学面纱,但这种教学上的透明性也伴随着不可避免的性能妥协。在基准测试平台(单核 CPU、纯原生 Python 解释器)上针对不同尺度的真实体系实测耗时如下:

规模一次前向+反向
17 原子 / 88 条边约 20 毫秒
26 原子 / 138 条边约 28 毫秒
63 原子 / 345 条边约 73 毫秒
张量积层,8 通道约 264 微秒
张量积层,32 通道约 3946 微秒(O(C3)O(C^3)

从基准数据中可以看出,处理单一中小分子构型的前向与反向耗时约为 20 至 70 毫秒。对于包含上千构型的数据集而言,跑完数十轮训练通常需要数十分钟乃至数小时。这对于在个人工作站上进行交互式算法验证、机理教学与单元测试已完全足够,但若要支撑起数百万原子步长的高通量从头算分子动力学(AIMD)拟合,代码必须走向工业级工程落地。

从教学原型向生产级系统演进的工程路线非常清晰。首先,利用张量重排与稀疏矩阵分解,将我们在原生三层循环中实现的 Clebsch-Gordan 张量积替换为基于 PyTorch 与 CUDA 算子深度优化的 e3nn 或 NequIP/MACE 原语;其次,将原生的字典邻接表与图循环结构,全面迁移至 PyTorch Geometric(PyG)或 torch_scatter 所支撑的高性能 CSR/COO 压缩稀疏图结构;最后,引入基于 LMDB 内存映射数据库的高并发数据预取管道。在保持底层几何映射公式完全不变的前提下,这一套工程重构将带来三到四个数量级的吞吐飞跃,将模型无缝送入工业级科研战壕。


9. 化学与材料里的应用图谱

9.1 模型选型的第一法则:目标物理量的张量阶数与空间宇称

在面对具体的化学或材料问题时,许多初涉该领域的研究者容易陷入盲目追逐复杂模型的陷阱。事实上,选择模型架构的核心判据并不在于网络层数的深浅,而在于一个纯粹的物理本原问题:你的模型输出,究竟属于欧几里得空间下的哪一种几何张量?

根据物理量在三维旋转与空间镜像下的变换法则,化学和凝聚态物理中的宏观与微观物理量可以被严谨地划分为四个基本象限:

第一象限是标量(阶数 l=0l=0 的零阶张量)。这类物理量在空间旋转与镜像反转下数值完全守恒。化学体系中的势能曲面、体系总能量、生成焓、带隙、分子轨道最高占据轨道(HOMO)与最低未占轨道(LUMO)能级差、水溶解度、辛醇-水分配系数(logP\log P)、聚合物玻璃化转变温度(TgT_g)以及缺陷形成能均属于此类。对于标量预测,纯粹基于核间距的不变量网络(如 SchNet 或 DimeNet)在数学表达能力上已经足够闭环,等变网络在此类任务中的优势主要体现为极高的数据利用效率与更精准的局部几何敏感度。

第二象限是向量(阶数 l=1l=1 的一阶极矢量)。这类物理量在空间旋转下经历与坐标相同的坐标变换,并在反演操作下精确反号。原子核受力矢量、体系瞬时永久偶极矩、分子激发跃迁偶极矩、电荷位移矢量以及扩散漂移速度皆属于此类。预测这些性质在数学上必须引入至少具备 l=1l=1 表象通道的等变架构,否则模型连输出方向随空间旋转而同步偏转的基本物理规律都无法维持。

第三象限是高阶张量(阶数 l=2l=2 及以上的对称二阶张量)。这类量在空间转动下经历双重旋转变换,涵盖了分子的电子极化率张量、晶体内部的柯西应力张量与弹性应变张量、宏观相对介电常数张量、分子四极矩以及 NMR 固态核磁化学位移各向异性(CSA)屏蔽张量。对此类物理特性的高保真回归,必须依赖具备对称无迹张量运算能力的 l=2l=2 等变通道。

第四象限则是极具化学特色的赝标量与手性敏感量(l=0l=0^-,奇宇称标量)。例如分子的旋光度、圆二色谱(CD)吸收差值、对映体过量百分比(eeee 值)以及手性催化反应中的立体选择性能垒差。由于对映异构体之间的核间距集合分毫不差,所有的标准纯旋转不变网络在此类任务上均会彻底失效,模型内部必须显式集成具有奇宇称响应的几何通道。

9.2 分子基态性质预测:2D 拓扑与 3D 空间的博弈

在 MoleculeNet、QM9 以及各类 ADMET 分子性质预测基准上,几何深度学习经历了长达数年的高强度检验。经验表明,几何深度学习架构并非在所有分子性质预测任务中都能稳操胜券。

当所预测的宏观物理化学性质主要由分子的官能团组成与共轭连通拓扑决定时——典型的例子包括平衡常数、水溶解度(ESOL)以及部分生物膜穿透性,传统的二维分子拓扑指纹(如 Morgan/ECFP)配合梯度提升决策树(LightGBM、XGBoost)往往能在极低的算力开销下取得不逊于甚至超越复杂几何网络的精度。更为关键的是,三维几何模型必须强依赖于输入构型的质量;如果训练与推断所采用的三维构象仅由构象生成算法(如 ETKDG)粗暴估算而来,其与真实溶液环境中的热力学系综构象偏差反而会沦为破坏预测稳定性的有害噪声。

相反,当目标物理量对微观空间排列极度敏感时——例如分子的基态热力学势能、多构象异构体的相对稳定性、大环分子的应变能以及非共价络合超分子系统的结合能,引入严格旋转不变或等变的几何表征则是不可替代的正确道路。

9.3 机器学习力场与从头算分子动力学

这是等变神经网络彻底重塑物理化学与计算材料学研究范式的核心主战场。长期以来,理论化学家深陷于"高精度小体系"(从头算量子化学与 DFT)与"低精度大体系"(经验力场如 AMBER、CHARMM 或 UFF)的两难权衡中。以 NequIP、Allegro 与 MACE 为代表的等变机器学习势函数(ML-PES),在保有接近量子力学基准精度的同时,将势能曲面和原子受力的计算速度提升了五到八个数量级。

势能面上的力:箭头指向能量下降的方向

在这一领域,等变架构的胜利建立在三大物理支柱之上。首先,体系能量是平移旋转标量,而原子受力是严格跟随分子整体旋转的一阶矢量,模型内部高阶张量通道为空间局域受力的方向学习提供了严丝合缝的表示载体;其次,在训练策略上,将能量与力的复合损失联合优化,能够从单帧构象中榨取出数百倍于单一标量能量的丰富物理约束;最后也是最具决定性的一点,通过能量对核坐标的解析求导输出保守力场,天然满足辛几何积分的守恒律约束,彻底根除了长时间分子动力学轨迹在微正则系综中发生非物理发热或能量耗散的顽疾。

9.4 构象采样生成与等变扩散模型

在从头分子设计与药物虚拟筛选中,三维分子构象生成正全面转向基于扩散概率模型(Diffusion Models)的架构。在这一类模型中,网络的核心任务是充当连续时间朗之万去噪器:它接收叠加了各向同性高斯噪声的三维坐标云,并精准预测出将原子推回低能稳定盆地的去噪位移矢量。

在此类生成任务中,去噪网络的 SE(3) 等变性不再仅仅是一个关乎拟合精度的指标,而是一道关乎统计力学分布数学合法性的刚性门槛。设想如果去噪网络不具备空间等变性,那么同一个分子在初始空间朝向发生任意翻转后,去噪网络输出的概率漂移场就会产生异化变形,进而导致同一个拓扑分子在欧几里得空间的不同方位上采样出完全不同的构象能量分布。采用 EGNN 或等变张量网络作为扩散骨架,天然确保了生成概率密度在全局欧氏变换下的绝对各向同性。

9.5 化学反应历程与过渡态搜索

探索复杂的有机合成路径与催化反应机理,本质上是在高维势能面上搜寻连接反应物与产物的一阶鞍点(过渡态,Transition State)。传统的轻推弹性带(NEB)方法对过渡态的精确定位高度依赖成百上千次的 DFT 能量与核梯度重算,计算开销极其惊人。

等变神经网络在此类任务中展现出了非凡的外推能力。由于等变架构在几何原理上完全抹除了"反应物在空间中的全局朝向"这一无意义的冗余自由度,模型能够在极其有限的微观过渡态构型训练集上,迅速捕获化学键断裂与重组过程中的通用受力模式。更进一步,针对不对称催化中的对映选择性控制,由于手性产物对映体的能垒差往往微弱至几个 kJ/mol,唯有具备严格奇宇称感知能力的几何网络,才能精准捕捉手性催化剂与底物在空间立体位阻上的微妙差异。

9.6 聚合物大分子与构象熵的挑战

当等变模型步入聚合物材料世界时,它将迎面遭遇三道独特的理论与工程难关。首先是大尺度体系带来的长程空间耦合与计算复杂度,聚合物分子往往包含成千上万个原子,如何在维持长程拓扑与空间关联的同时控制图卷积的感受野至关重要;其次是极其浩瀚的构象熵,单根柔性高分子链在构象空间中存在难以计数的局域低能极小值亚稳态;最后则是长链大分子往往需要引入一维或三维周期性边界条件进行连续介质近似。

聚合物链上的局部环境

在本教程的实验六中,我们选取了聚乙二醇(PEG)与聚丙二醇(PPG)的低阶寡聚体链作为研究切入点。通过在寡聚体上训练等变势能面并向更高聚合度的长链进行外推评测,不仅能够严谨检验几何网络对局部重复单元微环境的学习泛化能力,也为研发高保真聚合物多尺度力场提供了极具价值的概念验证。

9.7 凝聚态晶体与材料工程

相较于孤立的孤立分子体系,凝聚态固态材料展现出了两个决定性的新特征:宏观无限延展的晶格周期性以及极其庞杂的晶体结构配位多样性。将等变网络推广至材料科学领域,必须妥善处理两项核心工程机制:

其一是周期性边界条件(PBC)与最小镜像约定的融入。在周期性体系中,任何一个原子都在与遍布无限个镜像超胞中的原子发生物理交互。模型中的每条边必须携带一个显式的超胞镜像平移偏移矢量,原子间位移与距离必须在最小镜像平移准则下计算;同时,对于跨越晶胞边界的原子对,必须在无向图邻接构建中施加严格的半计数去重,防止人工引入双倍相互作用。

其二是元素参考态能量的剥离与残差学习。在由不同化学元素构成的无机晶体数据库中,不同化学组成的绝对内能差异可高达数十乃至上百 eV/atom。若直接强行让神经网络回归如此悬殊的绝对能量,损失函数的梯度会被重元素孤立原子的核电荷主导。工业界的标准规范是预先减去各元素单质或孤立原子的参考基态能量,引导网络聚焦拟合表征化学键合与相变驱动力的形成能残差。

周期性边界:影子邻居与镜像偏移

9.8 结构生物学与蛋白质大分子设计

在结构生物学与计算大分子药物设计中,以 AlphaFold、RoseTTAFold 以及几何矢量感知网络(GVP-GNN)为代表的等变模型引发了深刻的技术革命。

与小分子体系不同,生物大分子包含成千上万个氨基酸残基,其结构尺度远远超出了高阶不可约张量积(lmax3l_{\max} \ge 3)的算力承受极限。因此,在该领域中,设计哲学往往在"极致的群论完备性"与"卓越的大尺度工程吞吐量"之间做出审慎权衡:以 GVP 为代表的架构通过将空间信息限制在 l=0l=0 标量与 l=1l=1 几何矢量的轻量混合流,配合局部残基刚体坐标系(Frenet 标架),成功在大尺度蛋白质复合物折叠、抗体-抗原结合亲和力预测以及从头蛋白质序列设计任务中找到了完美的平衡点。

9.9 开源生态全景:站在巨人的肩膀上

本教程的核心使命在于通过纯原生代码的推导与剖析,使读者对每一处几何变换知其然更知其所以然。然而,当读者未来走出教程、投身于真实的科研攻关或工业级生产研发时,拥抱成熟的工业级开源生态则是最富效率的选择:

在基础张量代数库层面,e3nn 确立了欧几里得群张量运算的工业标准,提供了涵盖球谐展开、不可约表示直和分解与全自动 Clebsch-Gordan 系数查询的完备工具箱;在顶尖势能面与动力学模拟领域,NequIPAllegroMACE 以及引入注意力机制的 EquiformerV2 构成了当今兼顾精度与外推能力的标杆力量;在图结构底层与分布式训练支撑上,PyTorch Geometric(PyG)与 DGL 提供了顶级的稀疏图并行算子;而对于无机材料与表面异相催化,Meta 与卡内基梅隆大学联袂打造的 FairChem(前身为 Open Catalyst Project,OCP)则沉淀了规模空前的从头算数据集与极具参考价值的基线模型套件。


10. 八个实验

10.0 实验总览

前面九章建立了从群论基础到张量积推导、从架构设计到手写反向传播的完整理论链条。现在是检验这条链条的时候了。本章安排了八个精心设计的数值实验,它们共享同一套纯 Python 代码库与统一的评判标准体系,每一个实验都聚焦回答一个具体而尖锐的科学或工程问题,并且都能够独立运行、独立复现。所有表格中报告的数字均直接来源于 figures/results_*.json 中的原始输出,读者可以逐项对照核验。

实验数据来源标签规模关键数字
一、等变性体检RDKit 生成的阿司匹林构象无(只做变换)21 原子 / 138 边等变模型 101510^{-15},坐标 MLP 10010^{0}
二、手性10 个手性分子的对映体对几何手性指标80 个样本0.500 vs 0.938
三、旋转鲁棒性PEG/PPG/PE 寡聚体UFF 能量24/9 构象0.071 → 0.954
四、偶极矩32 个有机分子Gasteiger 偶极矩68/27 构象38.6° vs 82.5°
五、ESOLMoleculeNet(DeepChem)实验水溶解度902/113/113MAE 0.655 vs 0.581
六、聚合物力场PEG/PPG/PE 寡聚体UFF 能量 + 力42 构象 / 17~63 原子力 MAE 65.7 kcal/mol/Å
七、晶体(PBC)十种真实晶体结构经典对势的每原子能量与力80 个畸变构型998 meV/atom、旋转漂移 101510^{-15}
八、数据效率与代价同上张量积 O(C3)O(C^3)、图规模近似线性

以完整模式跑完全部八个实验大约需要 30 至 40 分钟,若切换为快速模式则仅需 3 至 5 分钟。这一耗时并非源于实验设计的低效,而是纯 Python 解释器不可避免的性能代价(详见 §8.6)。后续每个小节都会标注对应的运行时间,供读者合理安排工作节奏。

10.1 实验一:等变性体检——把对称性变成可以测量的数字

这个实验追问的是整套理论中最根本的一个问题:当我们把同一张分子图喂给不同模型,然后对输入施加随机的旋转、平移与镜像变换之后,各模型的输出到底会偏移多少?

实验取阿司匹林的一个三维构象作为探针分子(21 个原子,在 5 Å 截断半径下生成 138 条交互边),对每个被测模型实施 16 次随机刚体变换——其中旋转矩阵由四元数在 SO(3)SO(3) 上按各向同性 Haar 测度严格均匀采样,平移向量取自 [5,5]3[-5,5]^3 的均匀分布——然后记录变换前后模型输出相对偏差的最大值。

实验结果(results_equivariance.json,运行 11 秒)清晰地呈现如下:

模型旋转偏差(最大)平移偏差镜像偏差参数量
EGNN(标量 l=0l=0 + 坐标更新)6.5×10156.5\times10^{-15}5.7×10155.7\times10^{-15}0031 301
EGNN(只用距离,不更新坐标)6.2×10156.2\times10^{-15}4.6×10154.6\times10^{-15}0026 945
SchNet(距离 RBF 不变量)1.2×10151.2\times10^{-15}8.9×10168.9\times10^{-16}0012 225
坐标 MLP(不等变反面教材)5.975.9712.5912.591.251.253 809
张量场网络(l=0l=1l=0\oplus l=1,向量输出)5.8×10155.8\times10^{-15}4.2×10154.2\times10^{-15}004 905(向量等变误差 6.9×10156.9\times10^{-15}

仅凭这张对照表便足以理解为什么等变性值得被刻入模型结构:三个严格等变的模型——EGNN、SchNet 与张量场网络——其偏差无一例外地落在 101510^{-15} 量级,恰好是 IEEE 754 双精度浮点算术的舍入噪声水平。与之形成鲜明对比的是,坐标 MLP 的旋转偏差高达 5.97,平移偏差更飙升至 12.59,它连"输出不该随输入朝向改变"这一最基本的物理要求都无法满足。值得注意的是,EGNN 的坐标输出通道的等变误差同样仅为 7.1×10157.1\times10^{-15} Å,证实坐标更新路径同样维持了结构化的精确等变性。

等变性体检

这里有一个值得单独警示的实现陷阱:由于我们的读出层与坐标更新层均采用零初始化策略(§8.4 已详细阐述其物理动机),在模型尚未被"唤醒"之前,其输出恒等于零。在这种退化状态下,"旋转前后输出是否一致"的等变性检验会蜕变为 0/00/0 的不定式——我在第一版代码中就因此得到了一个外表漂亮的"误差 0.000e+00"假通过。为了杜绝这种自欺欺人的伪验证,当前版本的 demo_equivariance.py 会在检验之前先将网络的最后一层权重随机化。

将这个实验迁移到你自己的模型上非常简单:每次修改完任何一层网络之后,都应当立即重复上述检验流程。务必确保旋转测试绕的是真正的随机轴(千万不要只绕 zz 轴,那几乎必然会遗漏实现错误),平移向量要包含非整数位移分量,而镜像变换则可以顺便验证你的模型究竟能否区分手性——这恰好引出了接下来的实验二。

10.2 实验二:手性——不变量的天花板

如果说实验一回答的是"等变模型是否真的等变",那么实验二则直击另一个更具化学意义的灵魂之问:纯粹依赖原子间距离的模型,究竟能不能区分一对对映体?

实验选取了 10 个含手性中心的有机分子,用 RDKit 分别为每个分子生成三维构象,再对每个构象取空间镜像反演,由此获得一对严格互为镜像的训练样本。标签采用几何手性指标

χ=iedimin/λ(u^1×u^2)u^3,\chi=\sum_i e^{-d_i^{\min}/\lambda}\,(\hat u_1\times\hat u_2)\cdot\hat u_3 ,

的正负号——这是一个在空间反演下必然变号的赝标量。随后,实验在三种不同的特征表征与模型组合之间进行严格对照:纯距离径向基函数(RBF)、距离 RBF 附加赝标量通道、以及完整的 E(3) 等变 EGNN。在数据划分上,实验严格按分子身份而非按样本随机切分——同一对映体必须全部落在训练集或测试集的同一侧,否则镜像对之间的信息泄漏将完全摧毁实验的诊断价值。

实验结果(results_chirality.json,运行 165 秒,40 轮训练)如下:

模型参数量训练准确率测试准确率
距离 RBF + MLP2 6890.5000.500
距离 + 赝标量 + MLP2 7530.9690.938
EGNN(E(3) 等变标量读出)13 4920.5000.500

在解读这张表之前,先来看一项更基础的预检查:将每一对对映体的全部原子间距离排序后逐项比较,最大差异精确为 0.000 Å。换言之,这两个样本在所有成对距离的意义下是完全等价的孪生体。在这一前提下,任何以原子间距离作为唯一信息源的函数——无论其参数化形式多么复杂——在数学上都必然对它们给出完全相同的输出。因此,距离模型和 EGNN 停在 0.500 的测试准确率(等价于随机猜测)绝非训练失败的征兆,而是一个不可逾越的数学定理的直接体现。

再将目光投向唯一成功的模型:它仅仅比纯距离版本多引入了一维赝标量特征通道,参数量仅增加了 64 个,测试准确率却从随机水平一跃至 0.938。这个模型在本质上几乎就是在直接"读取"那个赝标量通道的符号——而这恰恰是这组对照实验最深刻的启示:模型的特征空间中有没有奇宇称通道,才是它能否表征手性的数学分水岭

手性判别

对于化学背景的读者,这里有一条关乎日常研究决策的重要提醒:默认配置下的 E(3) 等变模型——包括原版 EGNN、SchNet 以及大多数纯距离型架构——在空间反演下都是宇称偶的,它们与纯距离模型一样,在原理上就无法分辨对映体。如果你正在处理的科学问题涉及立体化学判定(手性催化中的对映选择性、手性药物的药效差异、不对称合成的产物分布),就必须在特征体系中显式引入宇称为 p=1p=-1 的几何通道,或者转而采用 SE(3)SE(3) 等变并配合宇称感知机制的架构(如 NequIP 或 MACE 中的 parity 通道设计)。

10.3 实验三:训练时只见过一个朝向,测试时被旋转会怎样

这是化学机器学习中最常见、也最容易被忽视的失败模式:数据集中每个分子恰好只有一个三维朝向。当模型直接以原始笛卡尔坐标作为输入特征时,它到底要付出怎样的代价?

实验取 PEG、PPG 和直链烷烃寡聚体的 UFF 力场计算能量作为标签——这是一个严格旋转不变的真实物理量——按分子身份划分训练集与测试集,然后训练四个对照模型:EGNN、SchNet、原始朝向训练的坐标 MLP、以及在训练过程中施加随机旋转数据增强的坐标 MLP。在测试阶段,系统地将分子绕随机轴旋转 0°180°180°,全程监测预测误差的变化曲线。

实验结果(results_rotation.json,运行 99 秒,24 轮训练)形成了全教程中最干净利落的一组对照:

模型见过的分子:00^\circ180180^\circ没见过的分子:00^\circ180180^\circ
EGNN(等变)1.514 → 1.5142.136 → 2.136
SchNet(不变量)0.327 → 0.3270.591 → 0.591
坐标 MLP(不等变)0.071 → 0.9541.389 → 1.152
坐标 MLP + 旋转增强0.772 → 0.6750.839 → 1.053

坐标 MLP 在训练集所见的分子上将拟合误差压到了极低的 0.071,但它学到的映射中不可避免地混入了"空间朝向"这一与物理量毫无关联的虚假变量:当分子被旋转至 180°180° 后,误差飙涨了 13 倍,从 0.071 猛增至 0.954。而 EGNN 与 SchNet 的对应两列数字在 0°180°180°严丝合缝地重合——它们的误差曲线在旋转角度轴上呈现为完美的水平直线。

旋转鲁棒性

数据增强对照组提供了另一层耐人寻味的洞察:随机旋转增强确实有效地抑制了坐标 MLP 对朝向的敏感度(误差从 0.772 到 0.675,在旋转扫描中基本持平),但代价是整体拟合精度出现了显著退化——从原始的 0.071 骤升至 0.772。这一现象精确地印证了 §1.4 中的核心论断:数据增强的本质是"用算力和样本数量去统计逼近对称性",而等变架构则是"将对称性以公理的形式直接写进网络结构"。

最后,我们必须诚实地指出这组实验中一个看似"不完美"却蕴含深刻教训的细节:在绝对误差的维度上,SchNet 的表现优于 EGNN(0.327 对比 1.514)。其原因在于,这份寡聚体数据集的能量标签主要由局部键长与键角决定,基于距离不变量的 SchNet 已经拥有充分的表达能力;EGNN 额外引入的坐标更新路径在这个极小规模的数据集上反而增加了优化的难度。等变不等于更准——它保证的是对称性的数学精确性,而非预测的绝对精度。我们刻意保留这个"不完美"的实验结果,正是为了传递这条朴素但至关重要的实践智慧。

10.4 实验四:偶极矩——当输出是向量

前面三个实验的预测目标都是标量——对于标量而言,"不变性"就是全部要求。然而一旦模型的输出从标量跃升至矢量(偶极矩、受力、位移),问题就从"输出不该变"升级为"输出必须同步旋转"。这个实验正是要量化等变与非等变架构在矢量预测任务中的鸿沟。

实验选取了 32 个覆盖广泛化学空间的有机分子(醇、醚、卤代烃、芳香烃、腈、硫化物等),利用 RDKit 的 Gasteiger 电荷方案构造物理真值标签 μ=iqi(rir质心)\mu=\sum_i q_i\,(r_i-r_{\text{质心}}),然后训练两个对照模型:一个是具备 l=0l=1l=0\oplus l=1 等变通道的张量场网络,另一个是将质心化坐标拍平后直接输出三个标量分量的全连接 MLP(坐标预减质心已经手工消除了平移自由度,只剩旋转问题悬而未决)。每个模型在三个独立的随机种子下分别训练,结果报告为均值 ±\pm 标准差。

实验结果(results_dipole.json,运行约 9 分钟,80 轮 × 3 个种子)如下:

模型分量 MAE方向平均误差等变误差 f(Rx)Rf(x)/f(x)|f(Rx)-Rf(x)|/|f(x)|退化比例
张量场网络(l=0l=1l=0\oplus l=12.72±0.642.72\pm0.6438.6±16.9\mathbf{38.6^\circ\pm16.9^\circ}3.1×10113.1\times10^{-11}0
坐标 MLP(不等变)3.69±0.073.69\pm0.0782.5±3.482.5^\circ\pm3.4^\circ1.501.500

需要注意的参考基准是:完全随机方向的期望误差恰好是 90°90°

等变张量场网络的等变误差为 101110^{-11},意味着"旋转输入后,输出精确地跟随同步旋转"——几何承诺在数值上得到了完美兑现。坐标 MLP 的对应数字为 1.50,直白地说就是"旋转之后它的预测矢量跑到了与原始结果毫不相关的方向"。在方向精度上,等变网络的 38.6°38.6° 明显优于 MLP 的 82.5°82.5°(后者已经逼近随机基线 90°90°),但也必须坦言 38.6°38.6° 本身并未达到工程可用的水平——这是一个仅有 4617 个参数、不包含消息传递的微型模型在 68 个训练样本上学出的结果。它证明了机制的正确性,不代表工业级精度。

偶极矩向量

这一节还值得重温 §7.7 中记录的第四个隐蔽 bug:如果张量积的两路输入恰好喂入了同一组几何特征向量,那么叉积 a×a0a\times a\equiv0 恒成立,l=1l=1 通道将在整个训练过程中退化为零向量。更为阴险的是,模型在此退化状态下仍会"正常训练、损失平稳下降"——因为它学会了输出零向量这个平凡解。我在第一版实现中得到的"方向误差 0.00°"假阳性正是源于此:零向量的方向没有定义,被指标计算代码直接跳过了。当前版本的指标体系中专门增设了"退化比例"一列,将零向量退化输出严格计为 90°90° 误差。

10.5 实验五:DeepChem 上的 ESOL——以及一个诚实的负结果

在经历了前四个围绕等变性机制进行精密验证的实验之后,我们终于来到了一个更接近日常研究场景的问题:在真实的分子性质预测数据集上,三维等变模型是否真的比二维分子指纹更强?

实验采用 DeepChem 的 load_delaney() 接口载入 MoleculeNet 的 ESOL 水溶解度数据集(训练集 902 个分子、验证集与测试集各 113 个),数据划分采用骨架切分(scaffold split)——按 Murcko 骨架将分子分组,同一骨架下的分子只出现在同一划分中。这比随机划分严格得多,更接近药物发现中"以已知骨架预测全新骨架"的真实应用场景。实验对比了三个对手:1024 位 ECFP 圆形拓扑指纹配合随机森林(纯二维方法,完全不使用三维几何信息)、Rupp(2012)提出的库仑矩阵配合随机森林(三维旋转不变量描述符),以及我们的纯 Python EGNN(三维等变模型,在 250 个训练分子上训练 25 轮)。

实验结果(results_deepchem_esol.json,运行 167 秒)如下:

模型MAER2R^2Pearson rr说明
ECFP + 随机森林0.6280.3710.4882D 指纹,无几何
库仑矩阵 + 随机森林0.5810.4310.4353D 不变量描述符
EGNN(纯 Python,3D 等变)0.6550.3290.68713 492 参数,训练 130 秒
EGNN(只用 1/4 训练数据)0.6830.254数据量扫描

DeepChem ESOL

这组结果必须从两面解读。一方面,三维等变模型在这个基准上并没有取胜:无论是 MAE 还是 R2R^2,它都略逊于"库仑矩阵 + 随机森林"这一已有十余年历史的经典方法,更不必说它比 ECFP 方案慢了几个数量级。但另一方面,一个令人意外的亮点是 EGNN 取得了三者中最高的 Pearson 相关系数(0.687)——这意味着它对溶解度的相对趋势排序学得最为准确,只不过预测的绝对尺度与偏移存在系统性偏差(因此 R2R^2 吃亏)。

这个负结果的物理根源是清晰的:ESOL 的水溶解度标签在化学本质上主要由分子的官能团组成与拓扑连通性决定,三维空间构型对最终溶解度的贡献极为有限。与此同时,三维模型还额外承担了一层"构象生成误差"的负担——我们用 ETKDG 算法生成的气相构象与真实溶液环境中分子的热力学系综构象之间存在不可忽视的偏差。当"标签对几何不敏感"与"输入几何本身带噪"两个不利因素叠加,三维等变模型自然处于结构性的劣势。

对于正在做研究的你,这条结论的实用价值可以凝练为一句话:永远先跑二维基线。如果 ECFP 配合梯度提升树已经给出了令人满意的精度,那么引入等变网络所带来的边际收益很可能被工程复杂度与调参成本吃掉。反过来,如果任务的标签确实深度依赖三维几何信息——构象选择性、力场拟合、手性识别、晶体能量——二维基线必然会露出明显的力不从心,就像实验二中那样。

10.6 实验六:聚合物力场——能量与力

这个实验将教程的技术栈推向了力场拟合这一等变网络最具变革性的应用场景,追问两个核心问题:神经网络能否同时学习能量与力?力的监督信号从何而来?

实验取聚乙二醇(PEG)、聚丙二醇(PPG)以及直链烷烃寡聚体作为模型体系(重复单元数为 3 至 8,加氢后原子数覆盖 17 至 63),在 UFF 力场优化的平衡构象上施加高斯随机位移(σ=0.06\sigma=0.06 Å)生成远离极小值的非平衡构型,再由 RDKit 调用 UFF 势函数计算真值能量与每个原子的笛卡尔受力矢量。训练过程同时拟合两个目标——能量标量误差与力矢量误差,其中力通过 E/x\partial E/\partial x 经由第二次反向传播求得(具体推导详见 §7.5)。测试集由更长的链(重复单元 5\ge5)组成,构成了真正的"链长度外推"评估。

实验结果(results_polymer.json,运行约 9 分钟,20 轮训练)如下:

模型训练能量 MAE测试能量 MAE测试力 MAE旋转漂移
EGNN(等变,坐标可动)0.891.7365.66ΔE=0\Delta E=0ΔF=5.9×1017\Delta F=5.9\times10^{-17}
SchNet(不变量,坐标不动)0.170.6865.52ΔE=1.2×1015\Delta E=1.2\times10^{-15}ΔF=1.3×1014\Delta F=1.3\times10^{-14}

(能量单位为标准化后的无量纲值,乘以 σE=69.8\sigma_E=69.8 kcal/mol 即可还原为物理量纲;力的单位为 kcal/mol/Å,训练集力矢量的 RMS 幅值为 103.9 kcal/mol/Å。)

从这张表中可以提炼出三条层层递进的核心观察。

首先,力的学习无需搭建任何额外的并行网络——它天然就是标量能量的空间梯度。在代码层面,所需的操作仅仅是"再反向传播一次",将力矢量残差作为梯度种子注入坐标入口端。更为深远的物理意义在于,由能量梯度导出的力天然是保守场,在长时间分子动力学轨迹的辛积分中,能量守恒可以得到远优于非保守力场的维持。

其次,等变保证在力矢量上同样严格成立:将整条高分子链施加全局旋转后,能量的漂移量为精确的零(双精度零),力矢量的漂移幅度仅有 5.9×10175.9\times10^{-17} kcal/mol/Å。这意味着"原子受力跟随分子整体旋转而同步偏转"是由网络结构在数学上刚性保证的,而非由训练数据中统计学习而来。

第三,也是本实验中最需要坦诚交代的事实:在这份极小的训练集上,两个模型都没有真正学好力。二者的力 MAE 均徘徊在 65 kcal/mol/Å 左右,相对于训练集力 RMS(103.9)约为 63%——这几乎等价于简单地预测平均力矢量。根本原因在于训练构型数量严重不足(仅 24 个),且 σ=0.06\sigma=0.06 Å 的位移在高频键伸缩振动模式上产生了高达上百 kcal/mol/Å 的恢复力峰值,而纯 Python 的计算预算不允许我们扩展到足够庞大的构象采样集。这个实验展示的是力场训练的完整机制与物理正确性,而非工业级精度;真正的生产力场拟合,请使用 MACE 或 NequIP 级别的优化实现,配合数十万乃至数百万构象的高质量量子化学数据集。

聚合物能量与力

10.7 实验七:周期性晶体——把 PBC 接进等变网络

从孤立分子迈入凝聚态材料领域,一个绕不开的物理现实是:在周期性晶体中,每个原子都在与遍布无穷多个镜像超胞中的"影子邻居"同时发生相互作用。等变网络如何在不破坏自身对称性保证的前提下,优雅地处理这种周期性?

实验选取了十种覆盖七大晶体结构类型的真实无机化合物——NaCl(岩盐型)、CsCl(氯化铯型)、ZnS(闪锌矿型)、CaF₂(萤石型)、MgO(岩盐型)、TiO₂(金红石型)、Si(金刚石型)、Cu(fcc)、Al(fcc)以及 Fe(bcc)。对每种晶体施加随机几何畸变(σ=0.03\sigma=0.03 Å)与整体随机旋转,利用经典短程对势(Lennard-Jones + 屏蔽库仑)计算每原子能量与原子受力矢量。在图结构层面,每条边携带一个显式的镜像偏移矢量,网络内部严格按 r=xi(xj+offset)r=x_i-(x_j+\text{offset}) 计算所有几何量。在训练目标方面,每原子能量预先减去该元素在训练集中的统计平均值作为参考态基准,引导网络聚焦学习形成能残差。

实验结果(results_materials.json,运行约 4 分钟,30 轮训练)如下:

划分模型每原子能量 MAE力的 MAE
插值(10 种晶体都见过,新畸变构型)EGNN(等变,PBC)998.5 meV/atom19.7 eV/Å
插值SchNet(不变量,PBC)637.3 meV/atom28.3 eV/Å
外推(CsCl、TiO₂、Si 三种新晶体)EGNN2991 meV/atom41.4 eV/Å
外推SchNet2250 meV/atom58.2 eV/Å

(训练集的每原子残差能量标准差 σ=3.54\sigma=3.54 eV/atom,力矢量 RMS 为 27.3 eV/Å。)

在这组结果中,最具说服力的发现不在于预测精度的绝对水平,而是等变性保证在周期性边界条件下依然严丝合缝:将整个晶胞的晶格矩阵与原子坐标联合旋转后,能量漂移仅为 4.8×10154.8\times10^{-15} eV,力矢量漂移为 1.7×10151.7\times10^{-15} eV/Å。而这套处理周期性体系的代码与前面处理孤立分子的代码共享完全相同的 EGNNLayer,唯一的差异仅仅在于每条边多附加了一个镜像偏移量。这正是本教程贯穿始终反复强调的核心理念:把几何信息"用对",远比盲目加深网络层数更重要

此外还有两个值得深入讨论的现象:其一,从"插值"到"外推"的精度退化极为显著(998 → 2991 meV/atom),这完全符合材料机器学习领域的共识——跨越结构类型的外推是真正意义上的分布外泛化问题;其二,两个模型的力预测精度仍然偏弱(19.7 至 58.2 eV/Å,相对于 27.3 eV/Å 的 RMS 并不出色),这同样根源于数据规模的严重不足:仅凭 80 个畸变构型覆盖十种化学体系,指望模型学到通用的凝聚态势能面是不切实际的。

晶体 PBC

10.8 实验八:数据效率与计算代价

最后一个实验直面一个所有实践者都必须回答的经济学问题:等变性在精度上的收益,是否值得为其付出的额外计算代价?以及这个代价究竟服从怎样的增长规律?

实验以聚合物数据集为测试平台,将训练样本数量从 4 逐步增加至 16,同时比较 EGNN、SchNet 与坐标 MLP 三者在测试误差与训练耗时两个维度上的表现。此外,实验还单独剖析了两类关键的微观计算开销——张量积层的耗时随通道数的增长规律,以及完整前向加反向传播的耗时随图规模(边数)的缩放行为。

实验结果(results_efficiency.json,运行约 6 分钟)凝聚为以下三张表:

模型4 个样本8 个12 个16 个(测试 MAE / 训练耗时)
EGNN1.56 / 1.2 s1.14 / 3.4 s1.43 / 4.5 s1.14 / 6.5 s
SchNet1.17 / 0.6 s1.22 / 1.5 s1.12 / 2.1 s1.08 / 3.1 s
坐标 MLP1.37 / 0.1 s1.20 / 0.2 s1.19 / 0.3 s1.29 / 0.3 s

张量积层(一次前向 + 反向):

通道数2481632
耗时(微秒)20.268.9263.51008.23945.8

前向 + 反向随图规模:

规模17 原子 / 88 边26 原子 / 138 边63 原子 / 345 边
耗时(毫秒)19.427.370.3

数据效率与代价

这三张表共同勾勒出三条清晰的结论。第一条关乎统计:在如此微小的样本量下,测试误差的随机波动远大于模型间的系统性差异。EGNN 的测试 MAE 在 1.14 至 1.56 之间跳荡不定,SchNet 稳定在 1.08 至 1.22 的窄带中,坐标 MLP 则在 1.19 至 1.37 之间徘徊。这有力地证伪了"等变模型在小数据集上必然优于非等变模型"这一过于乐观的信念——等变架构的数据效率优势,需要在"数据确实蕴含几何信号、且样本规模足以让非等变模型开始系统性过拟合"的条件下方能显现,实验三正是这一临界状态的完美例证。

第二条关乎算力:训练耗时与样本数呈现清晰的线性关系。当样本数翻四倍时,EGNN 的训练时间从 1.2 秒增长至 6.5 秒(约 5.4 倍),SchNet 从 0.6 秒增至 3.1 秒,坐标 MLP 从 0.1 秒增至 0.3 秒。三者之间的绝对差距——EGNN 在纯 Python 实现下比坐标 MLP 慢约 17 倍——就是"等变性"在当前技术栈下的精确价格标签。

第三条则揭示了等变架构最根本的计算瓶颈:张量积的运算代价随通道数呈立方级增长。通道数每翻一倍,耗时大约增长 3.8 倍(理论上限 23=82^3=8 倍,实测因常数项和低阶项而折半)。这条陡峭的幂律曲线,正是理解"为什么高阶等变模型(lmax=3l_{\max}=3)计算代价极其昂贵"的最直接定量证据。而在图规模的维度上,总耗时与边数近似呈线性增长(88 条边对应 19.4 毫秒,345 条边对应 70.3 毫秒),这一友善的线性缩放特性也正是几何图神经网络能够扩展至包含数万个原子的大尺度体系的物理基础。


11. 优缺点与使用建议

走过了十二章的数学推导与八个实验的数值验证之后,这一章将暂时搁下公式,转而以一个实践者的视角来回答每一位化学与材料研究者心中最直接的问题:等变神经网络到底什么时候该用、用了能收获什么、又必须为之付出什么?下文中的每一条结论都附带了本教程中支撑它的具体实验编号与数值证据,供读者按图索骥。

11.1 等变架构的核心优势

等变模型最具根本性的优势在于,对称性在其中是不可违背的刚性公理约束,而非可以被统计噪声扰动的柔性归纳偏置。实验一以数值精度给出了最有力的佐证:EGNN 与 SchNet 在随机旋转下的输出偏差分别为 6.5×10156.5\times10^{-15}1.2×10151.2\times10^{-15}——这是双精度浮点运算不可逾越的舍入地板——而不具备等变结构的坐标 MLP 的同一指标高达 5.97。这意味着等变模型无需借助任何数据增强、正则化或后处理近似就能精确满足物理对称性,且这一保证完全独立于训练数据的分布、优化是否收敛以及测试时的分子朝向。

由此自然引出第二大优势:显著更高的样本效率。实验三提供了一组教科书级别的对照——在"训练时仅见过一个固定朝向"的严苛设定下,坐标 MLP 的预测误差从 0.071(原始朝向)飙升至 0.954(旋转 180°180°),涨幅高达 13 倍;而 EGNN 与 SchNet 的误差在 0°180°180° 上严格相等。等变架构通过在网络结构中内建旋转不变性,一次性抹消了"朝向"这个与物理量无关的冗余自由度,从而将本该消耗在学习旋转对称性上的那些宝贵训练样本节省下来,投入到学习真正的化学规律中。

第三大优势是对方向性物理量的天然输出能力。力、偶极矩、位移、振动模式等矢量或张量性质,只有等变模型才能在数学上正确给出。实验四量化了这种差距:等变张量场网络的平均方向误差为 38.6°38.6°,坐标 MLP 则高达 82.5°82.5°(随机基线为 90°90°);更关键的是,等变网络的等变误差为 3.1×10113.1\times10^{-11}——"旋转输入,输出精确同步旋转"——而坐标 MLP 的对应数字为 1.50,意味着旋转后它的预测方向已经面目全非。

第四大优势是天然的物理一致性保证。当力由能量的解析梯度给出时(实验六),能量守恒在辛积分的离散化下自动得到最优维护;坐标更新路径天然保持体系质心不动(§6.3)。对于需要进行纳秒乃至微秒级长时间分子动力学模拟的研究者而言,这种结构化的守恒律保证至关重要。

最后,手性识别虽不是默认功能,但只需极小的架构修改即可实现。实验二证明,只要在特征空间中引入一维奇宇称赝标量通道(参数量仅增加 64 个),测试集上的对映体识别准确率就能从随机猜测的 0.500 跃升至 0.938。反过来说,这也意味着默认的 E(3) 等变模型——包括原版 EGNN 和 SchNet——在设计上是宇称偶的,不加修改便无法区分对映体。对于涉及立体化学的研究任务,这是一个必须在项目初期就明确审视的架构决策。

11.2 不可回避的代价与局限

等变性不是免费的午餐,它附带着一系列真实的工程与方法学代价。

首当其冲的是计算开销的非均匀膨胀。实验八的微观基准测试揭示,张量积层的计算耗时随通道数按 O(C3)O(C^3) 立方增长——8 通道仅需 264 微秒,而 32 通道已飙升至 3946 微秒。球谐函数展开与 Clebsch-Gordan 系数的计算在高阶不可约表示(lmax=3l_{\max}=3 及以上)下会迅速消耗掉算力预算,这正是 NequIP 类顶级势能面模型通常将网络深度限制在 2 至 4 层的根本原因。

其次是实现复杂度高且调试错误极其隐蔽。§7.7 中剖析的四个真实 bug 揭示了一条令人警醒的规律:它们无一例外地能够逃过损失曲线的常规监测。梯度数值检查、等变性刚体变换检查以及退化通道检测是手写等变层时不可或缺的三道保险,缺少任何一道都可能让静默的数值错误在生产环境中造成灾难性后果。

第三个代价是优化过程的系统性不稳定。当坐标更新路径参与训练时,分子体系被转化为一个离散动力学系统——每一步参数更新都会真实地移动原子的空间位置。

等变的代价:数据、算力与实现复杂度

实验六中最初的实现版本在第六个训练轮次便发生了灾难性发散(损失从 1.25 骤升至 18.9,测试能量 MAE 从 1.2 跳至 303),直到引入方向归一化与调低学习率后才恢复稳定。这并非偶发事件,而是"坐标参与梯度下降"这一类架构的结构性通病。

第四个现实是等变性不等于更高的预测精度。实验五在 ESOL 水溶解度数据集上给出了一个清醒的负面案例:1024 位 ECFP 配合随机森林的 MAE 为 0.628,十年前的库仑矩阵方案为 0.581,而我们的三维等变 EGNN 为 0.655——三者处于同一量级。当标签主要由分子官能团拓扑决定、三维几何贡献有限时,等变模型额外承担的"构象生成误差"负担反而会将其推入不利位置。

与此密切相关的是对输入构象质量的高度敏感性。三维模型的预测随构象而变,而构象本身来自 ETKDG 等近似生成器或半经验几何优化。当标签与构象之间的物理耦合较弱时(许多生物活性终点恰恰如此),几何模型可能学到的不是化学规律,而是构象生成器的系统偏差。

最后是数据与算力的联合门槛。高阶等变模型(lmax=3l_{\max}=3)在公开的势能面基准上通常需要数万至数百万个高质量构象方能发挥参数优势;在仅有几百个样本的小数据场景中,它们丰富的表征能力几乎无从施展。

11.3 何时应当克制使用等变网络

综合前述实验证据,以下几类典型场景中,更简单的方法往往是更明智的选择。

当可用数据仅有几百条且任务目标为标量预测时,ECFP 或手工分子描述符配合梯度提升树(LightGBM、XGBoost)通常是更务实的起点。等变模型在如此微小的数据量上很难体现出结构化优势——实验五中将训练数据削减至四分之一后,EGNN 的 MAE 仅从 0.655 小幅退化至 0.683,但工程复杂度的增量却完全不成比例。

当标签与三维分子构象的物理耦合天然较弱时——大部分 ADMET 性质、水溶解度、部分毒性终点均属此类——二维图神经网络或拓扑指纹方案通常在精度与效率之间提供更优的帕累托前沿。

当研究目标要求模型提供可解释的特征重要性归因时,等变模型内部流动的高维不可约张量表示难以直接映射到化学家熟悉的"官能团贡献"或"分子片段效应"等直觉化概念。

当任务仅要求预测旋转不变的正值标量(如总能量或形成能)而完全不涉及方向性输出时,基于成对距离不变量的轻量级网络(SchNet 级别)往往已经足够胜任——实验三和实验六中 SchNet 的绝对精度均不逊于甚至优于 EGNN。

11.4 从零到一的工程检查清单

一旦你审慎评估后决定在自己的研究项目中部署等变模型,下面这份按序执行的检查清单将帮助你规避绝大多数可预见的返工与陷阱。

第一步是执行严格的等变性体检。将 demo_equivariance.py 的检验逻辑完整移植到你的模型上,分别对输入施加旋转、平移和镜像三类变换,验证输出偏差是否落在 101010^{-10} 以下。这里特别需要注意的是:旋转测试必须使用在 SO(3)SO(3)各向同性均匀采样的真正随机旋转矩阵(绕随机轴、取随机角度),绝不能仅绕 zz 轴旋转,否则几乎必然会漏掉特定坐标分量的实现错误。

第二步是进行逐参数的数值梯度检查。检查覆盖范围至少应当包括"接收两个不同输入向量的张量积层"、"坐标更新路径"以及"距离计算路径"这三类最容易出错的参数区域。

第三步是审查物理量纲与参考态处理。能量标签是否已减去元素参考能量?力矢量的单位是 eV/Å 还是 kcal/mol/Å?目标量是否经过了零均值单位方差的标准化?实验七表明,不做参考态剥离时,不同化学体系的绝对能量差可达数十 eV/atom,严重干扰优化;施加参考态校正后,模型的每原子残差能量误差即可降至 0.6 至 1.0 eV/atom 的合理区间。

第四步是固定数据划分策略。随机划分几乎总是系统性地高估模型性能。对于分子性质预测,按 Murcko 骨架进行 scaffold 划分是更严格的工业标准;对于材料与晶体任务,按结构类型或化学体系划分(如实验七的"外推"设置)才能真实反映模型的分布外泛化能力。

最后一步是调节训练过程的数值稳定性旋钮:方向归一化、坐标更新层的零初始化与缩放因子、基础学习率以及全局梯度范数裁剪阈值。在几何深度学习中,首要目标永远是先让训练稳定运行不发散,然后才在稳定的基础上追求精度的最大化。

11.5 十个常见误区

在阅读文献、复现代码和指导学生的过程中,我反复遇到以下十个关于等变网络的典型误解,将它们逐一辨析于此,希望能为后来者扫清认知盲区。

第一个误区是将"不变"与"等变"混为一谈。标量输出(能量、电荷)所需要的是不变性——函数值在变换下完全不变;矢量或张量输出(力、偶极矩、极化率)所需要的则是等变性——输出必须与输入经历完全相同的坐标变换。二者在数学定义上截然不同,适用的架构设计也大相径庭。

第二个误区是相信原子间距离矩阵包含了分子的全部结构信息。§3.4 已经通过严格的数学论证表明,距离矩阵——乃至由距离矩阵经多维标度(MDS)重建的坐标——会系统性地丢失手性信息:一对对映体的距离矩阵在数值上完全等价。

第三个误区是想当然地认为 E(3) 等变模型天然具备区分对映体的能力。事实恰恰相反:默认的 E(3) 等变架构(含 EGNN 和 SchNet)在空间反演下保持宇称偶,除非显式引入奇宇称(p=1p=-1)通道,否则它们在数学上与纯距离模型一样对手性完全盲视。

第四个误区是在测试旋转不变性时仅绕 zz 轴旋转。这种偷懒做法几乎必然遗漏涉及 xxyy 坐标分量交换的实现错误,给出一个虚假的"通过"结论。

第五个误区是将 He 初始化(方差 2/nin2/n_{\text{in}})直接用于包含残差连接与邻居求和的几何图网络,然后对训练过程中的损失爆炸感到困惑不解(§8.4)。

第六个误区是不做数值梯度检查就盲目信任手写的等变层反向传播。"训练曲线看起来正常下降"完全不能证明梯度计算正确——§7.7 中四个 bug 的共同特征正是在损失曲线上毫无异常。

第七个误区是忽视坐标更新路径的发散风险。§6.5 已经详细复盘了原始实现在真实聚合物数据上的灾难性崩溃过程,缺乏方向归一化与输出缩放保护的坐标更新是一枚随时可能引爆的定时炸弹。

第八个误区是对标签进行了标准化处理(零均值、单位方差),却在推理阶段忘记将预测值反标准化回物理量纲,导致 MSE 在数值上看起来极小而产生虚假的性能幻觉。

第九个误区是在构建周期性体系的无向图邻接时,遗漏了镜像原子对的半计数去重(少乘了 1/21/2 因子),导致系统能量中出现精确整数倍的人为偏差(§3 与实验七中 pbc_edges_undirected 的设计详解了这一机制)。

第十个误区是在仅有几百条数据、且标签与三维几何关系薄弱的场景中强行上马等变模型,然后不加辨别地得出"等变模型没有用处"或"等变模型效果惊人"的结论。实验五清楚地表明,在这种边界条件下,实验结论的方向很可能完全由数据划分策略与标签质量所主导,而非由模型架构本身决定。


12. 把结论收拢

如果要用最简练的语言概括这门技术的核心精神,我会写下这三段话。

对称性不是数据中一种可以被统计抽样逐渐逼近的经验规律,而是物理函数必须精确满足的公理约束。将它以数学定理的形式刻入模型的网络结构,远比寄望于从有限的训练样本中归纳学习来得可靠——本教程中所有等变模型的对称性误差无一例外地落在双精度浮点噪声量级(101510^{-15}),而所有非等变模型的对应指标则停留在 10110^{-1} 乃至 10010^{0} 的尺度。

"只用不变量"是一条有着清晰数学边界的技术路线:它足以表达任意旋转不变函数,表达能力本身并不逊色,但它在定理层面无法区分镜像结构,也无法直接给出具有方向性的物理量。是否要越过这条边界,取决于你面对的科学问题是否涉及手性判别与矢量或张量输出——这是一个纯粹由物理需求驱动的决策,而非一个关乎模型复杂度的审美偏好。

一旦越过这条边界,代价是真实而不可回避的:张量积运算的 O(C3)O(C^3) 立方增长、坐标更新路径固有的数值不稳定性、以及远高于常规深度学习的调试难度与实现复杂度。这些代价在数据充裕、物理量对三维几何高度敏感的任务(力场拟合、过渡态搜索、手性识别、晶体势能面)中是值得的;在数据稀疏、标签与空间构象弱耦合的场景中则不值得。培养判断"值不值得"的能力,比记住 Clebsch-Gordan 系数的递推公式重要得多。

最后,我最想留给每一位化学与材料背景读者的一句话是:等变网络中的每一个数学设计,都能在你已经熟知的化学直觉中找到天然的对应物。角量子数 l=0l=0 对应的是不变标量(能量、电荷),l=1l=1 对应的是空间矢量(偶极矩、原子受力),l=2l=2 对应的是对称二阶张量(极化率、应力);Clebsch-Gordan 张量积的本质就是"两种不同阶数的几何量如何在对称性约束下耦合成第三种";宇称就是手性;坐标更新就是"让原子在由消息网络计算出的虚拟力场下互相推一把"。如果你能将这些数学符号与化学直觉一一对准,那么公式就不再是需要死记硬背的抽象负担,而只是你早已理解的物理图景的一种精确书写方式。


附录 A:符号表

下表汇总了本教程正文中反复出现的核心数学符号及其物理含义,供读者在阅读推导过程中随时查阅。

符号含义
NN原子数
xiR3x_i\in\mathbb{R}^3ii 个原子的坐标
hiRFh_i\in\mathbb{R}^Fii 个原子的标量特征
ziz_i原子类型(元素编号)
RSO(3)R\in SO(3)旋转矩阵
GG群(SO(3)SO(3)O(3)O(3)E(3)E(3)SE(3)SE(3) 等)
ρ(g)\rho(g)群元素 gg 的表示矩阵
ll不可约表示的阶数(角量子数式标号),维数 2l+12l+1
YlY_lll 阶(实)球谐函数
Dl(R)D_l(R)ll 阶 Wigner-D 矩阵
\otimes张量积(外积)
symtraceless\mathrm{symtraceless}对称无迹投影(l=2l=2 通道)
AF|A|_FFrobenius 范数
\odot逐元素相乘
EE能量(或其它标量标签)
Fi=E/xiF_i=-\partial E/\partial x_iii 个原子受到的力
μ\mu偶极矩(向量标签)
r^ij\hat r_{ij}单位方向向量 (xixj)/xixj(x_i-x_j)/|x_i-x_j|
χ\chi手性指标(赝标量)
ϕe,ϕh,ϕx\phi_e,\phi_h,\phi_xEGNN 的边/节点/坐标更新网络(小 MLP)
σE,σF\sigma_E,\sigma_F能量与力的归一化尺度
wFw_F力损失权重

附录 B:参考文献

以下文献按主题分为四组,涵盖了等变深度学习的理论基础、化学与材料领域的应用实践、分子表示与公开数据集,以及对称性数学与开源工具生态。

等变与几何深度学习的理论基础

  1. T. Cohen, M. Welling. Group Equivariant Convolutional Networks. ICML 2016.
  2. N. Thomas, T. Smidt, S. Kearnes, et al. Tensor Field Networks: Rotation- and Translation-Equivariant Neural Networks for 3D Point Clouds. arXiv:1802.08219, 2018.
  3. F. Fuchs, D. Worrall, V. Fischer, M. Welling. SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks. NeurIPS 2020.
  4. M. Weiler, M. Geiger, M. Welling, et al. 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data. NeurIPS 2018.
  5. T. Cohen, M. Weiler, B. Kicanaoglu, M. Welling. Gauge Equivariant Convolutional Networks and the Icosahedral CNN. ICML 2019.
  6. V. G. Satorras, E. Hoogeboom, M. Welling. E(n) Equivariant Graph Neural Networks. ICML 2021.
  7. J. Klicpera, J. Groß, S. Günnemann. Directional Message Passing for Molecular Graphs (DimeNet). ICLR 2020.
  8. K. Schütt, P.-J. Kindermans, H. E. Sauceda, et al. SchNet: A Continuous-Filter Convolutional Neural Network for Modeling Quantum Interactions. NeurIPS 2017.
  9. K. Schütt, F. Arbabzadah, S. Chmiela, et al. Quantum-Chemical Insights from Deep Tensor Neural Networks. Nature Communications 8:13890, 2017.
  10. B. Anderson, T. S. Hy, R. Kondor. Cormorant: Covariant Molecular Neural Networks. NeurIPS 2019.
  11. S. Batzner, A. Musaelian, L. Sun, et al. E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials (NequIP). Nature Communications 13:2453, 2022.
  12. A. Musaelian, S. Batzner, et al. Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics (Allegro). Nature Communications 14:579, 2023.
  13. I. Batatia, D. P. Kovács, et al. MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields. NeurIPS 2022.
  14. M. Geiger, T. Smidt. e3nn: Euclidean Neural Networks. arXiv:2207.09453, 2022.

化学与材料领域的应用实践

  1. M. Rupp, A. Tkatchenko, K.-R. Müller, O. A. von Lilienfeld. Fast and Accurate Modeling of Molecular Atomization Energies with Machine Learning. PRL 108:058301, 2012.
  2. A. P. Bartók, R. Kondor, G. Csányi. On Representing Chemical Environments. PRB 87:184115, 2013.
  3. J. Behler, M. Parrinello. Generalized Neural-Network Representation of High-Dimensional Potential-Energy Surfaces. PRL 98:146401, 2007.
  4. S. Chmiela, A. Tkatchenko, H. E. Sauceda, et al. Machine Learning of Accurate Energy-Conserving Molecular Force Fields. Science Advances 3:e1603015, 2017.
  5. B. Huang, O. A. von Lilienfeld. Ab Initio Machine Learning in Chemical Compound Space. Chemical Reviews 121:10001, 2021.
  6. K. Gubaev, E. V. Podryabinkin, A. V. Shapeev. Machine Learning of Molecular Properties: Locality and Active Learning. JCP 148:241727, 2018.
  7. A. Musaelian, A. Kozinsky, et al. Scaling the Leading Accuracy of Deep Equivariant Models to Biomolecular Simulations. SC 2023.
  8. O. T. Unke, S. Chmiela, H. E. Sauceda, et al. Machine Learning Force Fields. Chemical Reviews 121:10142, 2021.

分子表示与公开数据集

  1. D. Weininger. SMILES, a Chemical Language and Information System. J. Chem. Inf. Comput. Sci. 28:31, 1988.
  2. G. Landrum. RDKit: Open-Source Cheminformatics Software. 2006–.
  3. B. Ramsundar, P. Eastman, P. Walters, V. Pande. Deep Learning for the Life Sciences. O'Reilly, 2019.
  4. A. Gaulton, L. J. Bellis, A. P. Bento, et al. ChEMBL: A Large-Scale Bioactivity Database. Nucleic Acids Research 40:D1100, 2012.
  5. Z. Wu, B. Ramsundar, E. N. Feinberg, et al. MoleculeNet: A Benchmark for Molecular Machine Learning. Chemical Science 9:513, 2018.
  6. R. Ramakrishnan, P. O. Dral, M. Rupp, O. A. von Lilienfeld. Quantum Chemistry Structures and Properties of 134 Kilo Molecules (QM9). Scientific Data 1:140022, 2014.

对称性理论、手性表示与开源工具生态

  1. J.-P. Serre. Linear Representations of Finite Groups. Springer, 1977.(Schur 引理的标准参考)
  2. A. Zee. Group Theory in a Nutshell for Physicists. Princeton University Press, 2016.(SO(3)SO(3) 与 Wigner-D 的物理讲法)
  3. R. Kondor. A Novel Set of Rotationally and Translationally Invariant Features for Molecules. arXiv:0705.1792, 2007.
  4. K. T. Schütt, O. T. Unke, M. Gastegger. Equivariant Message Passing for the Prediction of Tensorial Properties and Molecular Spectra (PaiNN). ICML 2021.
  5. J. Brandstetter, R. Hesselink, E. van der Pol, et al. Geometric and Physical Quantities Improve E(3) Equivariant Message Passing. ICLR 2022.
  6. M. Fey, J. E. Lenssen. Fast Graph Representation Learning with PyTorch Geometric. ICLR Workshop 2019.
  7. DeepChem 开发团队. Democratizing Deep-Learning for Drug Discovery, Quantum Chemistry, Materials Science and Biology. 2016–.
  8. I. Loshchilov, F. Hutter. Decoupled Weight Decay Regularization (AdamW). ICLR 2019.
  9. D. P. Kingma, J. Ba. Adam: A Method for Stochastic Optimization. ICLR 2015.

本教程的全部代码、数据与插图均可完整复现。code/tests_enn.py 提供了涵盖 107 项断言与逐参数有限差分梯度检查的自动化测试套件; code/run_all.sh 可一键串行执行八个实验并输出结构化 JSON 结果; code/make_figures.py 从 JSON 数据自动重绘全部数据图表。 正文中的手绘风格示意插图由图像生成模型按 handraw-style 风格 #097 生成, 对应的提示词保存在 images/prompts/ 目录下,仅供教学示意之用。

《原子智能》· 纸质出版预备版 · PolyAI Team 著