中山大学陈川课题组（GTML）近日发布了等变图神经网络包Equitorch。Equitorch包中以模块化的形式集成了大多等变算子，并基于pyg构建图神经网络，同时提供了详细的文档、示例与教程。
背景
近两年来，随着AI for Science任务的兴起，3D等变图神经网络（3D-Equivariant GNN）获得了广泛的关注，出现了eSCN、NeuqIP、SEGNN、Equiformer(V2)、DPA(-2)等模型，并在分子几何生成、能量预测、粒子运动求解等领域获得了优异表现。
图片说明：3D旋转等变函数：φ(Rx)=Rφ(x)，若将函数的输入进行旋转，则函数的输出也会相应旋转。(图片源自https://arxiv.org/abs/2110.02905)
特点
在Equitorch中，我们以一种模块化的方式集成了现在工作中使用的大多等变算子以及其它相关辅助函数，并统一这些操作的数据格式约定。 对图神经网络相关的部分，如消息传递机制、图数据表示等，我们完全基于Pytorch-Geometric这一广泛使用的图神经网络框架，更方便传统图神经网络的研究者迁移。 我们正逐渐提供详细的文档、示例与教程，尽可能清晰地展示Equitorch中操作的定义与用途。
统一数据格式：
在Equitorch中，我们对数据的维度含义进行了约定：数据的第一个维度总是样本维度，即节点指标或边指标，最后一个维度总为特征维度，所有几何相关的维度均在中间。这里的几何维度可以是笛卡尔坐标x、y、z，等变特征的度(l)与阶(m)，或是球面角(θ与φ)的网格坐标。
如下图即展示了旋转矩阵（左）与等变特征（右）的数据排布。
（图片说明：左边存储旋转矩阵的张量维度为4×3×3×4，右边存储等变特征的张量维度为4×9×4。）
当前实现的操作：
在Equitorch中，我们现在主要实现的操作有：
神经网络模块equitorch.nn（“equitorch.nn”部分如果可以希望用行内代码块，下同）
等变线性变换 （加权）等变张量积 等变激活函数 等变层归一化 基础注意力操作 （径向、角度与球谐）基展开 软截断操作
数学函数equitorch.math
球谐张量、球谐函数、Wigner D矩阵的相关操作 球谐变换与逆球谐变换
数据变换equitorch.transform（基于Pytorch-Geometric的Data数据类型）
通过点坐标构建几何图、获取边的方向向量、长度嵌入、球谐嵌入等
其他通用功能性操作equitorch.utils
我们实现了当下等变神经网络中大部分基础操作，基于这些操作，可以十分灵活地搭建各种等变图神经网络架构。
文档与教程：
在实现操作之外，为明确Equitorch中使用的数学约定与操作定义，我们同时也编写了比较详细的文档，希望可以方便使用者对操作的理解。
（图片说明：TensorProduct张量积模块的文档）
另外，我们也提供了面向操作的教程，尽可能面向操作规则进行介绍，减少复杂的数学背景引入，以方便使用者更快地上手等变神经网络。目前教程包括了对球张量、张量积、SO(3)等变线性操作、以及搭建SO(3)等变图神经网络的部分。
（图片说明：等变性介绍）
（图片说明：张量积的介绍）
（图片说明：张量积规则说明）
（图片说明：SO(3)线性层的等变性说明）
（分割线）
目前，我们只实现了比较基础的操作，在将来我们还会基于这些操作逐渐实现更多的模型。我们已经将Equitorch上传到pip，欢迎研究或希望研究等变神经网络各位老师与同学尝试我们的工作，并对我们的工作提出宝贵的意见和建议。
项目地址: https://github.com/GTML-LAB/Equitorch
文档地址https://equitorch.readthedocs.io/en/latest/
(点击“阅读原文”跳转）
旋转等变性：如果输入的图形旋转，模型的输出也相应地旋转。 平移等变性：如果输入的图形平移，模型的输出也相应地平移。 尺度等变性：输入图形的缩放导致输出的相应缩放。
等变卷积：等变卷积是保证网络对旋转和平移等变性的关键。传统的卷积操作不具备这种对称性，因此等变GNN通常会基于群卷积（Group Convolution）等变换不变性来设计卷积操作。 消息传递机制：消息传递是图神经网络中的核心操作。等变图神经网络的消息传递机制需要确保信息在节点之间传递时保持等变性，因此通常会使用基于几何关系的权重更新策略。 坐标嵌入：图节点的特征通常包括坐标或位置信息。等变图神经网络在处理节点位置时，通过使用旋转或尺度等变的坐标嵌入，将这些位置编码到网络中。
分子建模：在分子建模中，分子结构通常具有旋转和平移对称性，因此等变图神经网络被广泛应用于分子特性预测、药物发现等领域。 3D点云数据：等变GNN可以对3D点云数据进行鲁棒的建模，尤其适用于物体识别、姿态估计等任务。 物理模拟：等变GNN在处理物理系统时（如流体动力学、粒子系统）可以准确捕捉到物体的对称性和运动规律。
SE(3)-Transformer：一种支持3D空间旋转和平移等变性的网络，适合用于分子建模等任务。 EGNN：等变图神经网络（E(n)-Equivariant GNN），主要用于处理具有几何对称性的图数据。 Tensor Field Network：一种利用张量场的卷积操作来实现等变性的网络结构，在分子建模和物理模拟中具有显著效果。
能够处理具有对称性的数据，提供更鲁棒的表现。 能够减少数据增强的需求（如不需要通过旋转、平移扩增数据），因为等变性已经嵌入到模型中。 更加适合处理高维空间的复杂结构数据。
计算复杂度高，实现难度较大。 需要更深入的数学基础（如群论）来理解和设计等变操作。 实际应用中可能需要对不同变换类型进行仔细选择和调试。
