跳转至

SE(3) Transformer

论文标题:SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks
发表期刊:NIPS 2020
代码链接:github

Abstract

1.本文贡献——提出SE(3)-Transformer
  • 一种适用于三维点云和图的自注意力模块变种。
  • 在连续的三维旋转/平移下,具有等变性性质。
2.等变性
  • 等变性:对于确保在输入数据受到扰动变换时,模型性能的稳定性和可预测性至关重要。
  • 等变性的积极推论:模型内部权重绑定的增加。
  • SE(3)-Transformer:利用自注意力的优势,对包含大量点且点数可变的点云和图进行操作,同时保证SE(3)等变性,以增强鲁棒性。
3.模型效果评估
  • 在一个玩具N-body质点模拟数据集上评估模型,展示了模型在输入旋转下的预测鲁棒性。
  • 在两个真实世界数据集ScanObjectNN和QM9上,取得了与非等变注意力baseline,以及无注意力等变模型相比具有竞争力的表现。在所有情况下,SE(3)-Transformer均优于上述两种模型。

1 Introduction

1.自注意力机制
  • 受欢迎程度急剧上升
  • 优点
    • 实现简单
    • 在各类任务上的高效性
      • 语言建模
      • 图像识别
      • 基于图的问题
  • 缺点
    • 自注意力机制广泛的应用性也意味着,对于特定任务,它未能充分利用已知的底层结构信息。
  • 本文工作
    • 提出SE(3)-Transformer,一种专门针对三维点云和图数据的自注意力机制,遵循等变性约束,从而提高了对扰动变换的鲁棒性,以及总体性能。
2.点云
  • 在许多领域中普遍存在、表现为各种形式
    • 三维物体扫描
    • 三维分子结构
    • N-body质点模拟
  • 挑战
    • 在能够在适应于输入数据中变化的点个数的同时,考虑到点位置的不规则采样。
  • 重要属性
    • 应对输入整体姿态的全局变化,输出保持不变。
    • 输入点云的三维平移、旋转,不应影响输出。
  • 本文发现
    • 在自注意力机制上显式施加等变性约束,可以解决上述挑战。
3.SE(3)-Transformer对注意力机制的应用
  • 方法概述
    • 使用自注意力机制,作为一种十分适合稀疏、未体素化的点云数据的数据依赖滤波器,同时考虑并充分利用任务的对称性。
  • 自注意力与点集的关系
    • 自注意力本身:是一种点集之间的伪线性映射。
    • 自注意力机制的两个组成部分
      • 依赖于输入的注意力权重(attention weights)
      • 值嵌入(value embedding):输入的嵌入表示。
    • 举例(图1)
      • 分子图
      • 每个原子都附有一个value embedding向量
      • 注意力权重:用边表示,边的宽度对应于注意力权重的大小。
    • SE(3)-Transformer的全局显式约束
      • attention weight:对全局姿态保持不变(invariant)。
      • value embedding:对全局姿态保持等变(equivariant)。
  • 等变性
    • 等变性:是卷积的平移权重绑定(translational weight-tying)的推广。
    • 确保一层的输入的变换能够表现为对输出的等效变换。
    • SE(3)等变性
      • 将传统卷积中已知的二维平移权重,绑定推广到三维旋转平移。
      • 将可学习函数的空间限制到遵循任务对称性的子空间,从而减少了可学习参数的数量。
      • 输入中特征之间的相对位置信息得以保留,提供了更丰富的不变性形式。
4.相关工作
  • 张量场网络(TFN, Tensor Filed Network)及其体素化等价物——三维导向CNN
    • 提供框架,用于构建在点云上运行的SE(3)等变卷积网络。
  • 采用自注意力而非卷积的优势
    • 允许自然地处理边缘特征,将TFN扩展到图。
    • 非线性等变层的首批例子之一
      • 在Section 3.2中,展示了:SE(3) Transformer相对于TFN,松弛了滤波器上的强角度约束(strong angular constraints),从而提升表征能力。
    • 强角度约束:在等变性文献中,被指出严重限制性能。
  • 提供了一个更高效的实现
    • 更高效的主要原因:球谐函数(spherical harmonics)的GPU加速版本。
    • 在实验中的TFN baseline利用了这一点,并采用了显著放大的架构。
5.贡献总结
  • 引入一种新颖的自注意力机制
    • 对其输入的全局旋转和平移,能确保具有不变性。
    • 对输入点标签的排列也具有等变性。
  • 解决等变神经网络的角度限制问题
    • 表明SE(3)-Transformer解决了当前SE(3)等变神经网络受到角度受限滤波器影响的问题。
  • 引入球谐函数的Pytorch实现
    • 在CPU上的速度比Scipy快10倍
    • 在GPU上快100至1000倍
    • 直接解决了TFNs[28]
    • 对于ScanObjectNN模型,与使用lielearn库中SH构建的网络相比,实现了约22倍的前向传播速度提升(见附录C)。
  • 背景介绍
    • 自注意力
    • 图神经网络
    • 等变性
  • 任务
    • 基于点云的机器学习任务
      • 对象分类
      • 分割
    • 输入:点云
      • 由\(n\)个坐标向量\(x_i \in \mathbb{R}^3\)组成的集合。
      • 每个点可能带有的附加特征:\(f_i \in \mathbb{R}^d\)。
2.1 The Attention Mechanism