导读
自注意力机制是最近非常火热的内容。今天这篇文章，作者基于自注意力机制提出SE(3)-Transformer，实现了三维数据中旋转平移的等变性。等变性有什么意义？作者取得了什么样的成果呢？
接下来让我们一起走进今天的文章吧！
SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks
作者：Fabian B. Fuchs、Daniel E. Worrall等
期刊：arXiv:Learning
时间：2020
地址：https://arxiv.org/abs/2006.10503
摘要
我们引入了SE(3)-Transformer。一种用于三维点云和图的自注意力模块的变体，在连续的三维旋转平移下是等变的。在存在脏数据输入转换时，等变性对于确保稳定和可预测的性能非常重要。等变性的一个积极的推论是增加了模型内的权重绑定。SE(3)-Transformer利用自注意力的优势来操作大点云和具有不同数量点的图，同时保证SE(3)-等变性的鲁棒性。我们在一个小N体粒子模拟数据集上评估我们的模型，展示了在输入旋转情况下预测的稳健性。我们进一步实现了两个真实世界的数据集，ScanObjectNN和QM9的竞争性性能。在所有情况下，我们的模型都优于强的非等变注意力基准和无注意力的等变模型。
1
引言
自注意力机制[31]近年来受到了越来越多的欢迎。它们相对简单的实现加上对广泛任务(如语言建模[31]、图像识别[18]或基于图的问题[32])的高效率，使它们成为一个有吸引力的组件。然而，应用程序的通用性意味着，对于特定的任务，现有底层结构的知识是没有用的。在本文中，我们提出了如图1所示的SE(3)-Transformer，一种专门针对3D点云和图数据的自注意力机制，它遵循等变性约束，提高了对损失变换的鲁棒性和总体性能。
图1:A) SE(3)-Transformer的每一层从点云映射到点云(或图到图)，同时保证等变性。对于分类，后面是不变池层和MLP。B)在每一层，对每个节点执行注意。在这里，红结点照顾它的邻居。注意权重(由线厚表示)是输入旋转时不变的。
点云数据广泛存在于各个领域，并以不同的形式呈现，如3D物体扫描[29]、3D分子结构[21]或n体粒子模拟[14]。寻找既能适应输入中不同数量的点，又能尊重点位置的不规则采样的神经结构是一项挑战。此外，一个重要的性质是，这些结构应该是不变的整体输入姿态的变化；即输入点云的三维平移和旋转不应影响输出。在本文中，我们发现在自我注意机制上明确施加等变性约束可以解决这些挑战。SE(3)-Transformer使用自注意力机制作为数据依赖过滤器，特别适合稀疏、非体素化的点云数据，同时尊重和利用手头任务的对称性。
自注意力本身是点集之间的伪线性映射。可以看出，它由两个部分组成：依赖于输入的注意力权重和输入的嵌入，称为值嵌入。在图1中，我们展示了一个分子图的例子，每个原子上都有一个值嵌入向量，注意力权重表示为边，宽度与注意力权重大小对应。在SE(3)-Transformer中，我们明确地设计了注意力权重对全局姿态等变。进一步，我们设计了值嵌入与全局位姿等变。等变性一般化了卷积的平移权重绑定。它确保层的输入转换等价于输出转换。特别地，SE(3)-等变性是由传统卷积已知的二维权重绑定到三维旋转平移的推广。这将可学习函数空间限制为一个子空间，该子空间符合任务的对称性，从而减少了可学习参数的数量。同时，它为我们提供了一种更丰富的不变性形式，因为它保留了输入特征之间的相对位置信息。
与我们最接近的作品是张量场网络(TFN)[28]和它们的体素化等价物，3D可操作CNN[37]。这为构建点云上的SE(3)等变卷积网络提供了框架。使用自我注意力而不是卷积有几个优点。(1)允许对边缘特征进行自然处理，将TFN扩展到图设置中。这是非线性等变层的第一个例子。在第3.2节中，我们展示了与TFN相比，我们提出的方法减轻了滤波器的强角约束，因此增加了表征能力。在等变性文献中已经指出了这一约束严重限制了性能[36]。此外，我们提供了一个更有效的实现，主要是由于一个GPU加速版本的球谐函数。我们实验中的TFN基准利用了这一点，并使用了比[28]中使用的架构显著扩大的架构。
我们的贡献如下：
1、我们引入了一种新颖的自注意力机制，保证对其输入的全局旋转和平移不变。它对输入点标签的排列也是等变的。
2、我们证明了SE(3)-Transformer解决了SE(3)-等变神经网络并发的问题，该问题受到角度约束滤波器的影响。
3、我们引入了Pytorch实现的球面谐波，它在CPU上比Scipy快10倍，在GPU上比Scipy快100-1000倍。这直接解决了TFNs[28]的瓶颈。例如，对于ScanObjectNN模型，与lielearn库中使用SH构建的网络相比，我们实现了≈22的前向传播(见附录C)。
4、代码在：
https://github.com/FabianFuchsML/se3-transformer-public
2
背景和相关工作
在本节中，我们将介绍自注意力、图神经网络和等变性的相关背景材料。我们关注的是基于点云的机器学习任务，如目标分类或分割。在这个任务中，我们给出一个点云作为输入，表示为n个坐标向量 的集合，每个点都有可选的特征
1、注意力机制
标准注意机制[31]可以被认为是由三个项组成的：一组查询向量 (i 从 1 到 m)，一组键向量 ( j 从 1 到 n)，和一组值向量 ( j 从 1 到 n)，其中r和p是低维嵌入的维数。我们通常将键 和值 解释为连接到同一个点j。对于给定的查询
我们用softmax作为作用于权重的非线性。一般来说，查询向量的数量不必等于输入点的数量[16]。在自注意力的情况下，查询、键和值向量是输入特征的嵌入，所以（2）：
其中， 在最普遍的情况下是神经网络[30]。对于我们来说，查询 与输入中的点i相关联，点i的几何位置是 。因此，如果有n个点，就有n个可能的查询。对于查询 ，我们说节点i关注所有其他节点
在深度学习的广泛任务(如语言建模[31]、图像识别[18]、基于图的问题[32]和关系推理[30,9])的成功推动下，最近的一系列工作将自注意力算法的形式应用于点云数据[44,42,16]。其中一个例子是Set Transformer[16]。当应用于ModelNet40[41]上的对象分类时，Set Transformer的输入是这些点的笛卡尔坐标。每一层都进一步嵌入这个位置信息，同时动态地从其他点查询信息。最终的每点嵌入被下采样并用于对象分类。
排列等变性：自注意力的一个关键性质是排列等变性。点标签 的排列导致自注意力输出的排列。这保证了注意力输出不会任意依赖于输入点的顺序。Wagstaff et al.[33]最近表明，这种机制在理论上可以近似所有置换等变函数。SE(3)-transformer是这种注意机制的一个特例，继承了置换等变性。但它将可学习函数的空间限制在旋转和平移等变函数上。
2、图神经网络
注意力与点云的大小呈二次关系，因此引入邻域是有用的：不是每个点都关注所有其他点，而是只关注最近的邻域。有邻域的集合自然用图表示。注意以前已经被引入到图中，命名为intra-、self-、vertex-或graph-attention [17,31,32,12,26]。这些方法由Wang等人[34]使用非局部神经网络统一。它的形式很简单（3）：
其中w和h是神经网络， 将和归一化为邻居
3、等变性
对 ，给定一组变换 ，其中G是一个抽象群，一个映射 被称为是等变的，如果对于每个g，都存在一个变换 使得（4）：
指标g可以被认为是描述这个变换的参数。给定一个对 ，我们可以求解满足上面方程的等变映射 的族。此外,如果 是线性的， 也是线性的，因此已经有一个非常丰富和发达的理论来寻找 [6]。在等变性文献中，深度网络是由交错线性映射 和等变非线性建立的。在三维旋转平移的情况下，已经证明了 的合适结构是TFN[28]，解释如下。注意，Romero等人最近为基于像素的图像数据引入了一个2D旋转平移等变注意力模块。
群表示：一般来说，变换 称为群表示。正式地说，一个群表示 是G到 的可逆矩阵GL (N)的集合的映射。更精确地说， 群同态，也就是说，它满足以下性质：对于所有的 ， 。特别对于三维旋转G = SO(3)，我们有几个有趣的性质:1)它的表示是正交矩阵，2)所有的表示都可以分解为（5）
其中 的基的变换矩阵[5]；每个 （ ）是 的矩阵，称为Wigner-D矩阵； 是沿着对角线的矩阵的和或连接。Wigner-D矩阵是SO(3)的不可约表示，认为它们是可能的最小表示。根据 (即设置 )变换的向量称为 向量。 向量在旋转下是不变的， 向量根据三维旋转矩阵旋转。注意， 向量的长度是 。它们可以进行叠加，根据Eq.(5)变换形成特征向量 。
张量场网络：张量场网络(TFN)[28]是一种神经网络，它在SE(3)-等变性约束下将点云映射到点云，即一组三维旋转和平移。对于点云，输入是形式如下（6）的向量场 ： 。
其中 是一个狄拉克δ函数， 是一个3D点坐标， 是点特征，表示原子序数或点身份。为了满足等变性，TFN变换的特征在Eq.(5)下，其中 每个 是不同类型的向量的串联，其中 的子向量被写成 。一个TFN层从 特征到 特征计算一个连续的、可学习的权重核 。位置 的TFN层的 输出为（7）：
我们也可以包含输入通道上的和，但在这里我们省略了它。Weiler et al. [37]， Thomas et al.[28]和Kondor[15]证明了核 位于等变基 的范围内。核是这些基础核的线性组合，其中半径 的第J个系数是一个可学习的函数
每一个基础核 是由大小为 的Clebsch-Gordan矩阵 的线性组合形成的，其中，第J,m个线性组合系数是第J个球谐函数 的第m维数。注意，只有当 且 时， ，这将核简化为标量 乘以恒等式， ，称为自交互[28]。因此，我们可以将TFN层重写为（9）:
等式(7)和等式(9)以消息传递的形式表示卷积，其中消息从所有节点和特征类型聚合。它们也是等式(3)中非局部图运算的一种形式，其中权值是边上的函数，特征
3
方法
在这里，我们提出了SE(3)-Transformer。该层可以分解为如图2所示的步骤过程，我们将在下一节中描述。这些就是从点云构造图，构造图上的等变边函数，如何传播图上的SE(3)-等变消息，以及如何聚合它们。我们还引入了自我交互层的另一种选择，我们称之为注意力自我交互。
图2：用我们的等变注意机制更新节点特征分四个步骤。附录中提供了更详细的描述，特别是对步骤2的描述。步骤3和步骤4将图形网络视角可视化:特征从节点传递到边，以旋转等变的方式计算键、查询和值，这些都依赖于特征和相对位置。
1、邻居
给定点云 ，我们首先引入一个邻居
2、SE(3)-Transformer
SE(3)-Transformer本身由三个组件组成。这些是1)沿边注意权重 ，构造为在每个边ij上的SE(3)-不变量，2)沿边SE(3)-等变值消息，在节点之间传播信息，如等式(7)的TFN卷积所示，3)线性/注意自交互层。每个邻居的注意力如下（10）:
这些组件如图2所示。如果我们去掉注意力权值，我们就有一个张量场卷积，如果我们去掉 对 的依赖，我们就有了一个常规的注意力机制。假设注意力权重 不变，则公式(10)与SE(3)-变换等价。这是因为它只是等价值消息的线性组合。不变的注意力权重可以通过等式(11)所示的点积注意结构来实现。该机制由节点i上的查询向量 和 附近每条边ij上的键向量
是直接和，也就是这个例子中的向量连接。线性嵌入矩阵 和 为TFN型(c.f. 等式(8))。注意权重 是不变的，原因如下。如果输入特征 是SO(3)-等变的，则查询 和键向量 也是SE(3)-等变的，因为线性嵌入矩阵是TFN类型的。在相同表示 下变换的SO(3)-等变向量的内积是不变的，因为如果 和 ，那么
注解：角调整当将等式(10)视为与数据相关的核 在文献中，SO(3)等变核被分解为可学习的径向函数 和不可学习的角核 的乘积之和(c.f. 公式(8))。 的固定角相依性是非交换代数中等变性条件的一个奇怪的伪命题，虽然是保证等变性的必要条件，但它被认为是对核的表达性的过度约束。有趣的是，注意力权重 引入了一种方法来调节 的角轮廓，同时保持等变性。
通道、自交互层和非线性 与传统的神经网络类似，SE(3)-Transformer可以直接扩展到每个表示度 的多个通道，为简便起见，本文省略了。这为自交互层搭建了舞台。注意力层（由图2和公式(10)中的(1)和(2)）聚集节点上的信息和输入表示度k。相比之下，自交互层(c.f. 公式(10)中的(3))只在相同程度的特征之间和一个节点内交换信息——非常类似于CNN中的1x1卷积。自交互是一种优雅的可学习跳跃连接形式，将信息从L层的查询点i传输到L + 1层的查询点i。这是至关重要的，因为，在SE(3)-Transformer中，点不关注自己。在我们的实验中，我们使用了两种不同类型的自交互层：(1)线性的(2)有注意力的两种形式：
线性的：在Schütt等[25]之后，输出通道是输入通道的学习线性组合，每个表示度使用一组权值 正如Thomas等人[28]所提出的，这之后是一个基于范数的非线性。
有注意力的：我们提出线性自交互、注意力自交互的扩展，将自交互与非线性相结合。我们将学习到的标量权值 替换为MLP输出的注意力权值，如等式(13) ( 意味着连接)。由于特征的内积的不变性，这些权值是SE(3)不变的，在相同的表示下变换。
3、节点和边的特征
点云数据通常有附加到点(节点特征)和点之间的连接(边缘特征)的信息，我们都希望将这些信息作为输入传递到网络的第一层。通过等式(6)和(10)中的张量 可以直接包含节点信息。为了整合边缘信息，请注意 是多个邻域的一部分。可以用公式(10)中的 代替 。现在， 可以携带不同的信息，取决于我们目前正在关注的邻居 。换句话说， 既可以携带关于节点j的信息，也可以携带关于边ij的信息。或者，如果边缘信息是标量，则可以将其纳入权矩阵 和 中作为径向网络的输入(见图2中的步骤2)。
4
实验
我们在三个数据集上测试了SE(3)-Transformer的有效性，每个数据集都测试了模型的不同方面。N-Body问题是一个等变任务：输入的旋转应该导致粒子位置和速度的旋转预测。接下来，我们评估一个真实世界的对象分类任务。在这里，网络面对的是大的噪声数据点云，只有在引力轴周围对称。最后，我们在一个分子性质回归任务上测试了SE(3)-Transformer，这揭示了它合并丰富的图结构的能力。我们比较公开可用的，最先进的结果，以及一套我们自己的基准。具体来说，我们比较了Set-Transformer[16](非等变注意模型)和张量场网络[28](类似于SE(3)-Transformer但不利用注意)。与[27,39]类似，我们通过对输入和输出应用均匀采样的SO(3)-变换来度量等方差的准确性。两者之间的距离，在样本上平均，产生等方差误差。注意，与Sosnovik等人的[27]不同，误差不是平方的：
1、N-Body 模拟
在这个实验中，我们使用了来自Kipf等人[14]的数据集。五个粒子各带一个正电荷或一个负电荷，相互之间相互排斥或相互吸引。网络的输入是粒子在特定时间步长的位置、速度和电荷。该算法的任务是预测500时间步后的相对位置和速度。我们故意将其表述为一个回归问题，以避免需要迭代地预测多个时间步骤。尽管将等变注意与LSTM等结合无疑是未来研究的一个有趣方向，但我们的目标是测试我们的核心贡献，并将其与相关模型进行比较。这个任务使自己有别于其他两个实验不是不变但等变化:当输入旋转或翻译,分别输出的变化(见图3)。
图3:基于传统自注意力(左)和旋转等变版本(右)的模型预测了五体问题中未来的位置和速度。左边的图分别显示了时间步长t = 0时的输入位置、t = 500时的地面真值位置和各自的预测。右边的图显示了输入旋转10度的预测位置和速度。虚线显示了完全等变模型的预测位置。
我们训练一个SE(3)-Transformer 具有4个等变层，每层紧随其后的是一个注意力自交互层(细节在附录中提供)。表1显示了定量结果。我们的模型优于基于注意力但非旋转的等变方法(Set Transformer)和不利用注意力的等变方法(张量场)。等方差误差表明，我们的方法确实是完全旋转等方差的计算精度。
2、在ScanObjectNN上的真实世界目标分类
基于点云的目标分类是一项研究较多的任务。有趣的是，目前绝大多数的神经网络方法都是在标量坐标上工作，而没有考虑向量特定的归纳偏差。最近的一些工作探索了旋转不变点云分类[45,47,4,22]。我们的方法通过使用直接作用在点云上的旋转平移等变层来将自身分离开来，而无需事先将其投影到球体上[22,45,7]。这允许在通过网络传输关于局部和全局方向的信息时进行权重绑定并提高样本效率——类似于二维图像上CNN的平移等方差。为了测试我们的方法，我们选择ScanObjectNN，这是一个最近引入的用于真实世界对象分类的数据集。基准测试提供了跨越15个不同类别的2902个对象的点云。我们只使用点的坐标作为输入，对象类别作为训练标签。我们训练了具有线性自交互的4个等变层的SE(3)-Transformer，然后是最大池化层和MLP。有趣的是，这个任务并不是完全旋转不变的，在统计意义上，因为对象是相对于重力轴对齐的。这将导致部署完全SO(3)不变模型时的性能损失(见图4a)。换句话说:当查看一个新对象时，知道up在哪里很有帮助。我们通过将z组件作为一个type-0字段，将x、y位置作为一个额外的type-1字段来创建算法的SO(2)不变版本(见附录)。我们把这种型号称为SE(3)-Transformer+z。这样，模型就可以通过抑制和促进不同的输入来学习要遵循哪些对称(对比图4a和图4b)。在表2中，我们将我们的模型与当前最先进的对象分类4进行了比较。尽管数据集没有发挥我们模型的优势(完全SE(3)-不变性)和更少的输入点，但性能与专门为对象分类设计的模型相比是有竞争力的——例如，PointGLR[23]是在更大的ModelNet40数据集[41]上预先训练的。有关性能与使用的输入点数的讨论，请参见附录D.1.2。
图4:ScanObjectNN: x轴显示了测试集中的数据增强。x值对应于x-y平面上任意轴的最大旋转。如果训练集和测试集都没有旋转(a中的x = 0)，通过提供附加的坐标的z分量来打破SE(3)-Transformer的对称性，标量输入可以显著改善性能。有趣的是，当训练集出现旋转不变问题(b中暗红色圆圈和深紫色三角形的强烈重叠)时，模型学会忽略额外的、破坏对称的输入。
3、QM9
QM9回归数据集[21]是一个公开可用的化学性质预测任务。有134k分子，每个分子最多有29个原子。原子被表示为一个5维单热节点嵌入到一个由4种不同化学键类型连接的分子图中(更多细节见附录)。给出了每个原子的位置。我们在表3中显示了Anderson等人[1]对6个回归任务的测试集的结果，越低越好。该表分为非等变模型(上)和等变模型(下)。我们最接近的模型是Cormorant和TFN(自己的实现)。我们看到，虽然不是最先进的，但我们提供了有竞争力的性能，特别是针对Cormorant和TFN，它们在SE(3)的不可约表示下转换(像我们)，不像LieConv(T3)，使用SE(3)的左正则表示，这可能解释了它的成功。
5
结论
我们提出了一个基于注意力的神经结构，专门为点云数据设计。该体系结构保证了对输入的旋转和平移的鲁棒性，避免了对训练时间数据的增加的需要，并确保了对任意坐标系选择的稳定性。自注意力的使用允许非均匀性、数据自适应过滤器，而邻居关系的使用可以扩展到大型点云。我们还介绍了注意机制作为一个数据依赖的非线性的解释，增加了我们可以在等变网络中使用的等变非线性的列表。此外，我们提供了一个加速球谐计算高达3个数量级的代码。这种加速使我们能够训练更大版本的SE(3)-Transformer和张量场网络[28]，并将这些模型应用于真实的数据集。
我们的实验表明，在旋转-平移-等变模型中增加注意力始终能够提高训练的准确性和稳定性。特别是对于大型社区，注意力被证明对模型收敛至关重要。另一方面，与传统的注意相比，添加等方差约束也提高了我们所有实验的性能，同时提供了关于输入数据旋转的鲁棒性的数学保证。
长按二维码关注