等变神经网络:原理与脉络
引言:从不变到等变
我们在 机器学习分子力场:综述 和 不变神经网络:从SchNet到GemNet 的结尾都提到了同一个矛盾:能量是旋转不变的标量,但力是旋转等变的矢量。不变网络把中间表示压成标量,等于强迫所有张量信息走”标量化→存储→求导还原”的迂回路线。等变神经网络的思路则截然相反:既然物理定律在$E(3)$群作用下具有协变性,不如让网络的每一层都保持这种协变性。
形式化地说,对群$G$和其表示$\rho$,若映射$f$满足
则称$f$关于$G$等变(equivariant);当$\rho_{\mathrm{out}}$取平凡表示时,它退化为不变(invariant)。力场问题恰好同时需要两者:能量头输出不变标量,力头输出等变矢量,而后者可以由前者对坐标求导自动获得——前提是整个网络严格等变。
数学准备
群与表示
分子模拟关心的对称群是欧几里得群$E(3)$:平移、旋转、反射的全体。平移通常通过以相对坐标$\vec{r}_{ij}$为输入来平凡地处理;反射的处理则见仁见智(手性问题需要区分对映异构体时必须破坏镜面对称),多数力场模型选择只对$SO(3)$(真旋转)等变、对宇称做适当标注。所以核心困难集中在$SO(3)$上。
群的表示是把群元映为矩阵的同态:$\rho:G\to GL(V)$。表示可以被分解为不可约表示(irreducible representation, irrep)的直和——它们是不能继续分解的最小构件。
不可约表示与球谐函数
$SO(3)$的不可约表示由角动量量子数$l=0,1,2,\dots$标记,维度为$2l+1$:
| $l$ | 维度 | 物理对应 | 宇称 |
|---|---|---|---|
| 0 | 1 | 标量(能量、电荷) | 偶/奇 |
| 1 | 3 | 矢量(力、偶极矩) | 偶/奇 |
| 2 | 5 | 二阶张量(极化率) | 偶/奇 |
$l>0$的irrep不是实数数组意义上的向量,而是球谐函数展开系数:任何定义在球面上的函数都可以用球谐函数$Y_l^m(\hat{r})$展开,
旋转作用在系数上的矩阵就是Wigner D矩阵:$\rho_l(R)_{mm’}=D^l_{mm’}(R)$。一个等变特征向量的通用形态是若干irrep的直和:
例如NequIP默认使用到$l_{\max}=1$或$2$的特征,MACE常用到$l_{\max}=2$或$3$,EquiformerV2甚至用到$l_{\max}=8$——阶数越高,角度分辨能力越强,计算量也越大。
张量积与耦合
两个irrep如何组合出第三个?答案是Clebsch-Gordan级数(耦合张量积):
其中$C$是Clebsch-Gordan系数——本质上就是量子力学里角动量耦合的那套系数。这个运算有一个决定性的优点:它是天然等变的,两个等变张量的张量积仍然等变。等变网络的全部积木——卷积、非线性、注意力——最终都要靠它搭出来。
最简单也最重要的例子:把邻居方向$\hat{r}_{ij}$用球谐函数展开得到等变滤波器$Y_l(\hat{r}_{ij})$,与邻居特征$h_j$做张量积再径向加权求和,就得到一个等变版本的图卷积——这正是Tensor Field Network和NequIP的核心算子。
还有一个工程上极其重要的恒等式。当两个irrep相同时,耦合张量积可以写成显式的对称形式:
eSCN正是利用这类结构把稠密的CG张量积稀疏化,才让高阶$L$的计算变得可行(详见 eSCN:把SO(3)卷积拆成稀疏矩阵 )。
等变网络的基本构造
有了上面的语言,一个典型的等变消息传递层可以这样组装:
- 嵌入层:原子种类和距离编码为标量特征($l=0$),坐标$\hat{r}_{ij}$展开为球谐矢量$Y_l(\hat{r}_{ij})$;
- 等变卷积:$h_i^{(l)}\leftarrow \sum_j W(r_{ij})\left(Y^{(l)}(\hat{r}_{ij})\otimes h_j\right)^{(l)}$;
- 等变非线性:难点所在——普通激活函数作用在irrep分量上会破坏等变性。标准方案是门控(gating):用$l=0$的标量特征经sigmoid生成门控信号去缩放高阶分量,或者使用S² Activation(在球面上做傅里叶变换、逐点激活、再变换回来);
- 输出头:能量取$l=0$分量的聚合;力由$\vec{F}_i=-\partial E/\partial \vec{r}_i$给出,由于全网络等变,求导后的力自动等变,无需任何后处理。
实践中的实现几乎都依赖e3nn库——它提供了irrep张量的数据结构、CG张量积和Wigner D矩阵的高效计算,是整个等变生态的地基。
发展简史
等变网络的思想并非力场领域的发明,这里按时间顺序梳理关键节点:
- 2018年,Tensor Field Networks(Thomas et al.):奠基之作,首次提出基于球谐函数和张量积的等变点云网络,并在简单的分子能量回归上验证了样本效率优势。
- 2019年,SE(3)等变的3D CNN(Weiler et al.)与Cormorant:将等变卷积推广到体积数据;Cormorant首次用于分子势能面并引入软等变思想。
- 2020年,SE(3)-Transformer(Fuchs et al., NeurIPS):把等变性与注意力机制结合,注意力权重本身设计为不变标量,从而不破坏等变性。它证明了”等变+Transformer”的兼容性,为后来的Equiformer铺路。
- 2021年,NequIP(Batzner et al., Nature Communications):第一个在真实力场基准上展示数量级样本效率提升的等变消息传递模型,一锤定音地宣告等变架构的价值,MLIP领域从此转向。详见 NequIP:等变消息传递的开山之作 。
- 2022年,Allegro(Musaelian et al.):strictly-local的等变架构,绕开全局消息传递,为超大规模并行模拟而生。详见 Allegro:为十亿原子而生的严格局域等变势 。
- 2022年,MACE(Batatia et al., Commun. Physics):把多体关联(ACE思想)注入等变框架,用更少的层数达到更高的表达力,后来成为通用大模型时代的主力架构之一。详见 MACE:高阶等变消息传递与原子簇展开 。
- 2022–2024年,Equiformer系列(Liao et al.):等变Transformer的系统化,v1验证可行性,v2借eSCN卷积把可扩展性推到大模型级别,v3完成工程优化与精度收官。详见 Equiformer系列:等变Transformer的三级跳 。
- 2023–2024年,eSCN与eSEN(Passaro & Zitnick;Fu et al.):分别解决了高阶张量积的计算瓶颈和力的质量问题,是OC20榜单上两代标杆。详见 eSCN:把SO(3)卷积拆成稀疏矩阵 与 eSEN:平滑能量网络与力的质量 。
代价与收益
等变不是免费的午餐。它的代价包括:CG张量积的计算量随$l_{\max}$急剧增长;实现复杂度高、调试困难;对小体系而言可能过度参数化。
但收益同样是结构性的:样本效率的数量级提升(NequIP在MD17上只需千分之一的数据即可匹配不变网络)、物理量预测的正确归纳偏置(矢量性质直接输出而非靠求导拼凑)、以及跨元素、跨化学空间的外推能力——最后这一点正是通用大模型时代得以开启的技术前提。
下一篇文章我们从NequIP开始,逐一解析这些模型。




