AI科学计算系列 | 等变神经网络与e3nn数学库

背景

近年来，人工智能在传统自然科学问题上开始发挥越来越重要的作用，在一些重要的科学问题上取得了瞩目的成果。

在实际的物理、化学、生物等学科中，往往会遇到各种具有几何特征的对象，比如位置矢量、力矢量。由这些几何对象所描述的方程，其本身不依赖于坐标系的选取，也就是物理学定律的相对性。蕴藏在这些几何特征之下的，是系统所具有的对称性信息。

因此如何在神经网络中保持这些几何特征与对称性，人工智能应用于自然科学问题的一个关键问题。等变神经网络通过模型的设计，使数据从输入到输出始终保持其几何特征与等变性，并展示出相较于传统神经网络更小的数据需求与更优秀的表现。

在介绍等变神经网络之前，我们需要了解几何张量与等变性的概念。

1. 几何张量与等变性

首先回顾一下我们熟悉的矢量。

矢量v→是一个抽象的数学对象，包含了大小与方向，可以用一个箭头来表示。当我们选择一个坐标系后，（三维）矢量便可以通过一组数字——坐标来描述，如图1所示。

图1：三维空间的矢量与坐标

旋转变换g可以作用到矢量上，选取坐标系后，其本身也可以用一个旋转矩阵来描述，例如绕x轴旋转30度，可以写为

于是对矢量v→的旋转变换则可以用矩阵乘法来表示R_ij(g)v_j（这里使用了爱因斯坦求和记号）。

另一类空间变换为反演变换，在反演变换下保持不变称之为偶（even），方向反向（加个负号）称之为奇（odd）。例如一个位置矢量，在反演变换下会反向（奇），而一个磁场强度矢量在反演变换下保持不变（偶）。

几何张量可以视为对矢量这一概念的推广，使其在空间变换（旋转与反演）下满足特定的坐标变换关系。例如，零阶几何张量为标量，在空间变换下保持不变；一阶几何张量为矢量，在空间变换下变换如上所描述；高阶几何张量的变换关系则更为复杂。

这里再提一点，这里我们通过一组坐标或矩阵来表示抽象的几何张量与变换，这就是数学里表示的思想，后面我们还会再详细地讨论这部分。

等变性（在物理中被称为协变性）：映射f:X→Y是G等变的，如果对于所有g∈G，f(D_g^X(x))=D_g^Y(f(x)) 。

简单来说一个映射的等变性是指对输入的变换可以相应地反映在输出上。如果输入是几何张量，那么输出也会是几何张量。

2. 等变神经网络

等变神经网络便是在传统神经网络的基础上，加入了等变性的约束要求。其中网络的每一种操作都要求是等变的，因此整个网络是一个等变映射。

在实际科学计算问题中，我们经常会遇到各种三维几何体。

比如在化学领域，一个分子的结构可以用各个原子的原子数[Z0,Z1,…]与对应的原子坐标[[x0,y0,z0],[x1,y1,z1],…]来描述。而我们想要知道的信息可能是每个原子的势能[E0,E1,…]或原子的受力[[Fx0,Fy0,Fz0],[Fx1,Fy1,Fz1],…]。其中原子数和势能为标量，原子坐标和受力为矢量。

如果这个分子旋转一个角度，分子还是同一个分子，自身的各个性质应该保持不变。反应在数据上，原子数和势能不变，但原子坐标和受力的数据要做相应的变换（乘旋转矩阵）。这便对神经网络提出了一定的要求。解决思路大致可以分为三类：

1. 机理驱动：使用传统的计算方法，比如密度泛函理论，通过底层机理以及对称性要求，直接求解结果。这种思路不需要数据，其结果完全确定，没有泛化能力。

2. 数据驱动：使用传统的神经网络，但需要大量数据增强，即把分子每一个旋转后的数据都作为输入，直接通过网络进行训练。本质上相当于让神经网络自己学习分子的等变性。这种思路输入数据多，网络规模大，训练相对困难，但不需要底层机理，结果不确定性高，泛化能力强。

3. 机理数据融合：使用等变神经网络，基于等变性构建网络，不需要数据增强，网络规模较小，结果不确定性较低，泛化能力较强。

等变神经网络在分子预测与生成中，展示出了优异的效果。

在预测分子势能中，可以大幅降低数据量的需求。

例如：Graphormer是一个GNN与Transformer结合的神经网络，在2021公开催化剂挑战赛中获得冠军。Equiformer是一个结合了等变性、GNN与Transformer的神经网络，同时降低数据量和网络模型的大小，在获得与Graphormer相当的精度的情况下， 算力需求大约减少15倍，训练效率显著提升。

在分子设计中，可以筛去无效输出，提升精度与效果。

例如：GeoDiff是一个等变性与Diffusion结合的神经网络，在分子生成的泛化性与精度方面，均优于传统diffusion神经网络。

图2：NeuqIP

NequIP是另一种等变性与GNN结合的预测性神经网络，其中使用到了高阶几何张量。文中比较了最高不同阶的几何张量对结果的影响。随着网络中最高阶的几何张量从0到3，预测结果的精度也会显著提升。

但是当涉及到高阶几何张量时，其运算规则相较于标量和矢量会变得复杂很多，其中关键的运算包括球谐函数和张量积。简单理解，球谐函数可以将球坐标的矢量映射成高阶几何张量，张量积则是将两个几何张量相乘，获得更高阶的几何张量。

e3nn数学库是一个方便进行这类高阶几何张量运算的数学库。为了深入理解这些几何张量运算是如何进行的，以及如何使用e3nn来构建等变神经网络，我们需要介绍一些基础的群表示论的背景知识。

3.变换群与表示

群是一种代数结构，在一个集合上定义一种二元运算，这种运算满足以下4点：

1. 封闭性：两个群元素的运算仍是群元素。

2. 存在单位元。

3. 存在逆元素，使得两者的运算为单位元。

4. 满足运算的结合律。

e3nn中的nn指神经网络（neural network），而e3则是指E(3)，即三维欧几里得空间的空间变换群，可分解为平移、旋转（SO(3) 特殊正交群）与反演（Z2二阶循环群）。

其中平移的等变性在卷积中已经满足，因此我们重点关注旋转与反演（SO(3)×Z2=O(3)）。

但是这里群元素作为一个变换操作，是抽象的对象。为了更好地研究其性质与应用，数学上会将其同态映射到我们熟悉的对象（线性变换）上，并保持数学结构相似，这便是群的表示。

群表示：线性空间V上的所有可逆线性变换构成一个群GL(V)，同态映射R:G→GL(V)称为群G在线性空间V上的一个表示。

举个例子，在最开始我们为了描述抽象的矢量，通过选取坐标系将其表示为三维线性空间的一组坐标，而旋转变换则表示为这个线性空间上的一个旋转矩阵。于是旋转变换对矢量的作用就可以通过矩阵乘法计算得到。

一个矩阵的表示可以分解为多个不可约子矩阵的直和（分块对角），将这个概念抽象化到表示的概念：可约表示可以分解为多个不可约表示的直和。当我们知道了所有不可约表示的性质，那么就可以获得整个表示的信息。

O(3) 群的不可约表示分解相对简单，其每种不可约表示可以用 (l,p) 标记，l=0,1,2,... 为阶数，p=e,o为奇偶性，l阶不可约表示的维数为 2l+1 。例如一个矢量，阶数为1，奇偶性为奇，因此可以简记为1o。

在神经网络中，不同阶的不可约表示的地位类似于CNN中不同的通道。

举个例子：假如我们知道一个长度为13的数据的表示为'1x2e+2x1o+2x0e'，意味着这组数据可以分解为一个（无迹对称）二阶张量，两个矢量，两个标量。而对应的13×13空间变换矩阵，则可以分解为一个5×5，两个3×3，两个1×1的子变换矩阵的直和。

4.张量积

表示的张量积可以理解为其对应的表示空间的张量积。

一般而言，两个表示的张量积是可约的，维数则为相乘。

以两个矢量的张量积为例，矢量的维数为3，那么张量积的结果的维数为9，可以分解为一个标量（维数为1），一个赝矢量（维数为3）和一个无迹对称二阶张量（维数为5）。

1o⊗1o=0e⊕1e⊕2e

两个不可约表示的张量积的每一个不可约表示分量可以通过以下公示来计算。

其中 C_{(L1,m1),(L2,m2)}^(L3,m3) 为Clebsch-Gordan系数。

上面这个例子中，如果只取张量积的标量分量，那么这就是我们熟悉的矢量的点乘；如果只取张量积的矢量分量，这就是我们熟悉的矢量的叉乘。

对于一般的可约表示的张量积，我们需要先将其分解成不可约表示，然后每一组不可约表示两两相乘，得到张量积的不可约表示。

用图形表示，将两个张量积的输入写在左上和右上方，输出写在下方。上面两个矢量的张量积便可以表示为图3(左)的形式，而更一般的可约表示，其张量积会更加复杂，如图3(右)所示。

图3：张量积的图形表示

e3nn数据库提供了O(3)群表示的核心计算模块，包括不可约表示、旋转变换、球谐函数、Wigner函数（CG系数）和张量积。此外还提供了在此之上构建的等变网络层，例如等变激活层、等变线性层和等变卷积层等。通过类似于e3nn的等变数学库，可以更方便地构建等变神经网络。

参考文献

[1] Ying, Chengxuan, et al. "Do transformers really perform badly for graph representation?."Advances in Neural Information Processing Systems34 (2021): 28877-28888.

[2] Liao,Yi-Lun,and,Tess Smidt. "Equiformer: Equivariant graph attention transformer for 3d atomistic graphs."arXiv preprint arXiv:2206.11990(2022).

[3] Xu, Minkai, et al. "Geodiff: A geometric diffusion model for molecular conformation generation."arXiv preprint arXiv:2203.02923(2022).

[4] Batzner, Simon, et al. "E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials."Nature communications13.1 (2022): 1-11.

[5]Geiger, Mario, and Tess Smidt. "e3nn: Euclidean neural networks."arXiv preprint arXiv:2207.09453(2022).
————————————————
版权声明：本文为CSDN博主「小乐快乐」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/weixin_45666880/article/details/129061365