1. 引言

在AI4Science领域，深度学习已被广泛应用到多个领域中，包括药物发现、蛋白质结构预测、动力学模拟等。在科学研究中，高质量数据的数据往往昂贵而稀少，尤其是以三维结构形式呈现的数据。这些数据通常仅代表空间中特定位置的状态，但它们的物理属性（例如能量分布、电势场、力学性质等）在经历空间的旋转和平移变换时，能够展现出一定不变性或等变性。即如果数据在空间位置上发生变化，其核心物理特性会保持稳定或以可预测的方式变化。而几何深度学习则将几何先验知识融入模型设计之中，以提高模型对科学数据的使用效率和训练质量。本文旨在深入探讨当前几何深度学习领域基于GNN的三种主流框架及其代表性工作。

2. 群不变GNN

第一种框架是维持群不变性（G-invariant）的GNN。许多情况下物体的特性与其空间位置无关，因此在某些问题中，我们希望网络在输入发生旋转时仍能产生稳定的输出。群不变的GNN通过将局部邻域内的几何信息标量化来实现这一点，比如两点间的距离，这样的测量不会因为物体的旋转而受影响。然后，网络聚合来自局部邻域的标量，以生成其表示。

一个较新的代表性工作是发表在NeurIPS 2021的GemNet（GemNet: Universal Directional Graph Neural Networks for Molecules）。GemNet讨论了不变性网络可能具有的“毕加索问题”，即如果一张人脸图片只有眼睛部分经过旋转，那么一个“旋转不变”的CNN可能无法识别出这种错误。GemNet希望在保持相对旋转信息的同时，对全局旋转保持不变。

GemNet的核心创新在于使用有向边嵌入和两跳信息传递。模型首先利用球形傅里叶-贝塞尔基底结合多项式径向包络来表示原子间的相对方向信息，并将基底分为CBF、RBF、SBF三个部分，以纳入更多可用的几何信息。其次，模型通过选择基于邻近原子方向的球面上的点来实现球形表示的离散化。例如，原子a的嵌入方向可以由原子c定义，形成边嵌入ca。然后基于原子b更新原子a的球形表示，通过边ca和db进行两跳信息传递，其中原子c和d定义嵌入方向，如下图。此外，模型对于一个点对设计了一对有向嵌入，并使用独立的参数保证它们互不相同。

尽管GemNet的架构在性能上有所提升，但其两跳消息传递机制也带来了额外的计算负担。为了缓解这一问题，作者设计了高效的双线性层，并采用了下采样技术。在COLL数据集上使用单跳消息传递的消融模型表现一般，但在MD17上表现良好，这说明在某些情况下单跳消息传递已足够，但在复杂任务中，两跳消息传递能够展现更强的学习能力。

3. 使用笛卡尔向量的群等变GNN

在一些特定应用场景中，我们希望模型能够根据输入的旋转变化产生相应的等变输出。例如，在蛋白质复合物预测问题中，两个蛋白质无论处于何种初始位置，都不会影响它们形成复合物的反应。这种类型的网络即为群等变图神经网络（G-equivariant GNN），其中一种有代表性的模型是使用笛卡尔向量的等变网络EGNN（E(n) Equivariant Graph Neural Networks，ICML 2021）。

EGNN在消息传递过程中将等变特征和不变特征分开处理，分别表示为x和h。模型传递的消息m由不变特征、点对之间的距离以及边的初始特征这两个标量组成，保持了不变性。不变特征h是通过直接聚合消息得到的，而等变特征x则是通过将消息与两点之间的向量相乘得到的。

相比于其他方法，EGNN在中间层不需要进行计算成本高昂的高阶表示处理，例如，它避免了一些其他方法所需进行的球面谐波等昂贵计算，同时仍能达到或超越竞争对手的性能。此外，该模型还可轻松扩展至更高维空间，而不会带来显著的计算量增加。

4. 使用球形张量的群等变GNN

群等变图神经网络的主流是使用球形张量框架，一些工作称该框架为e3nn。这个框架可以用来实例化多种模型，如TFN、Cormorant、SE(3)-Transformers以及MACE。

这类模型使用高阶球形张量作为节点特征，其球面谐波基的数量从最低阶数l=0开始，可扩展到任意高的阶数l=L。前两个阶数（即l=0,l=1）分别对应标量特征和向量特征，模型通过张量积聚合邻域特征，以及向量的高阶球面谐波表示，来获得更高阶的张量特征：

其中，张量积的权重是通过可学习的相对距离径向基函数计算得出的，要得到更新后的高阶张量的阶数，要将上式扩展为：

其中C为Clebsch-Gordan系数，确保了更新特征的群等变性。在这个模型里，当表示仅限于标量（即l=0）时，获得的更新形式类似于群不变的GNN。同样，当使用标量和向量（即l=1）表示时，得到的更新形式类似于使用笛卡尔向量的群等变GNN。

本文介绍的MACE发布于NIPS2022（MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields）。它在e3nn框架中提供了一种计算高k体阶特征的高效方法，即通过原子团簇展开。这一方法首先聚集邻域特征，然后对这些邻域特征进行k-1次重复的自张量积。类似于：

这种方法避免了在更标准的多体展开中所需的对所有k元组进行对称化或生成的复杂工作。等式实际上相当于计算向量的幂，隐含地包括了高阶特征。MACE通过减少所需的消息传递迭代次数，显著提高了计算效率。在实际应用中，MACE显示出在多种基准任务上超越现有方法的潜力，在保持准确性的同时提高了计算效率。

5. 总结

本文探讨了AI4Science领域中几何图神经网络的三种主流框架及其代表性工作，即群不变GNN、使用笛卡尔向量的群等变GNN，以及使用球形张量的群等变GNN。这些工作具有良好的性能，在效率上具有优势，有充足的理论支持，是很有潜力的研究方向。