[论文解读] VN-Transformer: Rotation-Equivariant Attention for Vector Neurons

Serge Assaad, Carlton Downey+3 more|arXiv (Cornell University)|Jun 8, 2022

Optical measurement and interference techniques被引用 11

一句话总结

VN-Transformer 引入了一种旋转等变注意力机制，消除了对向量神经元模型中复杂预处理的依赖，支持直接输入点云数据并处理非空间属性。通过 ε-近似等变性，该方法在 3D 形状分类和运动预测任务中实现了最先进性能，推理速度更快、鲁棒性更强，在 Waymo Open Motion Dataset 上将平均距离误差降低了 27%，相比标准 Transformer。

ABSTRACT

Rotation equivariance is a desirable property in many practical applications such as motion forecasting and 3D perception, where it can offer benefits like sample efficiency, better generalization, and robustness to input perturbations. Vector Neurons (VN) is a recently developed framework offering a simple yet effective approach for deriving rotation-equivariant analogs of standard machine learning operations by extending one-dimensional scalar neurons to three-dimensional "vector neurons." We introduce a novel "VN-Transformer" architecture to address several shortcomings of the current VN models. Our contributions are: $(i)$ we derive a rotation-equivariant attention mechanism which eliminates the need for the heavy feature preprocessing required by the original Vector Neurons models; $(ii)$ we extend the VN framework to support non-spatial attributes, expanding the applicability of these models to real-world datasets; $(iii)$ we derive a rotation-equivariant mechanism for multi-scale reduction of point-cloud resolution, greatly speeding up inference and training; $(iv)$ we show that small tradeoffs in equivariance ($ε$-approximate equivariance) can be used to obtain large improvements in numerical stability and training robustness on accelerated hardware, and we bound the propagation of equivariance violations in our models. Finally, we apply our VN-Transformer to 3D shape classification and motion forecasting with compelling results.

研究动机与目标

• 开发一种旋转等变的 Transformer 架构，以消除向量神经元模型中对复杂预处理的依赖。
• 将向量神经元框架扩展至处理真实点云数据中的非空间属性。
• 在保留旋转等变性的同时，实现点云数据的高效多尺度特征压缩。
• 分析近似等变性（ε）与数值稳定性之间的权衡，以提升在加速硬件上的训练性能。
• 在 3D 形状分类和运动预测任务中，实现具有等变性保证的最先进性能。

提出的方法

• 提出一种基于弗罗贝尼乌斯内积的注意力机制，将标量自注意力推广至矩阵值向量神经元，确保旋转等变性。
• 引入可学习的潜在特征机制，将点云分辨率从 1024 个 token 降低至 32 个，训练速度提升约 2 倍，且精度损失极小。
• 设计两种机制，将非空间属性（如强度、语义）整合进等变模型中，同时保持旋转等变性。
• 在 VN-LinearWithBias 层中采用 ε-近似等变性，提升硬件加速器上的数值稳定性和训练鲁棒性。
• 在端到端训练中将 VN-Transformer 应用于 3D 形状分类（ModelNet40）和运动预测（Waymo Open Motion Dataset）。
• 在 Transformer 解码器头中使用残差连接和层归一化，并适配向量神经元形式化表达。

实验结果

研究问题

• RQ1 能否直接从向量神经元推导出旋转等变注意力机制，而无需依赖边卷积或预处理？
• RQ2 如何在不破坏等变性的前提下，将非空间属性整合进旋转等变模型？
• RQ3 ε-近似等变性对硬件加速环境下的数值稳定性和训练效率有何影响？
• RQ4 能否实现多尺度特征压缩的等变性，以提升点云模型的推理速度？
• RQ5 在旋转扰动下，VN-Transformer 是否优于标准 Transformer 在运动预测任务中的表现？

主要发现

• 在 Waymo Open Motion Dataset 上，当 ε = 10⁻⁶ 且加入速度特征时，VN-Transformer 将平均距离误差（ADE）降低至 3.67，相比原始 Transformer 提高了 27%。
• 通过潜在特征机制将 token 数从 1024 减少至 32 后，训练速度提升 2 倍（每秒步数），在 ModelNet40 上仅造成 1.7% 的精度下降。
• 当 ε = 10⁻⁶ 时，VN-Transformer 在 WOMD 上实现 3.95 ADE，优于标准 Transformer 的 5.01 ADE，即使后者使用了 z 轴旋转数据增强。
• 模型在旋转下保持强泛化能力：预测结果在输入旋转下保持等变性，如图 7 所示，而标准 Transformer 则不具备此特性。
• VN-PointNet 和 VN-Transformer 均从点阵状空间特征中显著受益，表明空间与非空间信息实现了有效融合。
• 理论边界分析表明，等变性破坏的误差以受控方式传播，总误差受各层利普希茨常数与 ε 的乘积所限制。