RI-MAE: Rotation-Invariant Masked AutoEncoders for Self-Supervised Point Cloud Representation Learning

📄 arXiv: 2409.00353v2 📥 PDF

作者: Kunming Su, Qiuxia Wu, Panpan Cai, Xiaogang Zhu, Xuequan Lu, Zhiyong Wang, Kun Hu

分类: cs.CV

发布日期: 2024-08-31 (更新: 2024-12-25)

备注: Accepted to AAAI 2025

🔗 代码/项目: GITHUB

💡 一句话要点

提出RI-MAE，解决点云自监督学习中旋转不变性缺失问题。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 点云自监督学习 旋转不变性 掩码自编码器 RI-Transformer 师生学习

📋 核心要点

  1. 现有基于掩码点云建模的自监督学习方法对旋转敏感，性能在旋转变化下急剧下降。
  2. RI-MAE通过RI-Transformer构建旋转不变的潜在空间，并设计双分支师生架构实现旋转不变的自监督重建。
  3. 实验表明，RI-MAE对旋转具有鲁棒性，并在多个下游任务上取得了当前最优的性能。

📝 摘要（中文）

本文提出了一种新的旋转不变掩码自编码器（RI-MAE），用于点云数据的自监督学习，旨在解决现有方法对旋转敏感的问题。RI-MAE主要解决了两个挑战：1）实现旋转不变的潜在表示；2）以旋转不变的方式促进自监督重建。为此，我们引入了RI-Transformer，它具有解耦的几何内容、旋转不变的相对方向和位置嵌入机制，用于构建旋转不变的点云潜在空间。此外，还设计了一种新颖的双分支师生架构，通过在学习到的旋转不变潜在空间中重建掩码补丁来实现自监督学习。每个分支都基于RI-Transformer，并通过额外的RI-Transformer预测器连接。教师编码所有点云补丁，而学生仅编码未掩码的补丁。最后，预测器使用来自学生的输出潜在嵌入来预测掩码补丁的潜在特征，并由来自教师的输出进行监督。大量实验表明，我们的方法对旋转具有鲁棒性，并在各种下游任务上实现了最先进的性能。

🔬 方法详解

问题定义：现有基于掩码点云建模的自监督学习方法在处理具有旋转变换的点云数据时，性能会显著下降。这是因为这些方法学习到的特征表示对旋转不具备不变性，导致模型无法泛化到新的旋转角度。因此，如何提升点云自监督学习模型对旋转的鲁棒性是一个关键问题。

核心思路：RI-MAE的核心思路是构建一个旋转不变的潜在空间，并在该空间中进行掩码点云的重建。通过设计旋转不变的Transformer模块（RI-Transformer）和双分支师生架构，模型可以学习到对旋转不敏感的特征表示，从而提升在旋转变换下的性能。

技术框架：RI-MAE的整体架构包含三个主要模块：RI-Transformer编码器（教师分支和学生分支）、RI-Transformer预测器。教师分支编码所有点云块，学生分支仅编码未被掩码的点云块。预测器利用学生分支的输出预测被掩码的点云块的潜在特征，并由教师分支的输出进行监督。整个流程通过最小化预测器输出与教师分支输出之间的差异来训练模型。

关键创新：RI-MAE的关键创新在于RI-Transformer的设计，它通过解耦几何内容、旋转不变的相对方向和位置嵌入机制，实现了旋转不变的特征表示。此外，双分支师生架构也保证了模型能够有效地学习到旋转不变的潜在空间。

关键设计：RI-Transformer使用了旋转不变的相对位置编码，具体实现方式未知（论文未详细描述）。损失函数采用均方误差（MSE）来衡量预测器输出与教师分支输出之间的差异。掩码比例等超参数的选择未知（论文未详细描述）。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

RI-MAE在多个点云分类和分割任务上取得了显著的性能提升。具体数据未知（论文未提供具体数值），但摘要中提到在各种下游任务上实现了最先进的性能，表明该方法在旋转鲁棒性方面具有明显优势。与现有方法相比，RI-MAE能够更好地处理具有旋转变换的点云数据。

🎯 应用场景

RI-MAE在三维物体识别、场景理解、机器人导航等领域具有广泛的应用前景。通过提升点云数据处理的旋转鲁棒性，可以使模型在更复杂的真实环境中稳定工作。例如，在自动驾驶中，车辆可以更准确地识别旋转后的障碍物；在机器人抓取中，机器人可以更好地处理不同角度的物体。

📄 摘要（原文）

Masked point modeling methods have recently achieved great success in self-supervised learning for point cloud data. However, these methods are sensitive to rotations and often exhibit sharp performance drops when encountering rotational variations. In this paper, we propose a novel Rotation-Invariant Masked AutoEncoders (RI-MAE) to address two major challenges: 1) achieving rotation-invariant latent representations, and 2) facilitating self-supervised reconstruction in a rotation-invariant manner. For the first challenge, we introduce RI-Transformer, which features disentangled geometry content, rotation-invariant relative orientation and position embedding mechanisms for constructing rotation-invariant point cloud latent space. For the second challenge, a novel dual-branch student-teacher architecture is devised. It enables the self-supervised learning via the reconstruction of masked patches within the learned rotation-invariant latent space. Each branch is based on an RI-Transformer, and they are connected with an additional RI-Transformer predictor. The teacher encodes all point patches, while the student solely encodes unmasked ones. Finally, the predictor predicts the latent features of the masked patches using the output latent embeddings from the student, supervised by the outputs from the teacher. Extensive experiments demonstrate that our method is robust to rotations, achieving the state-of-the-art performance on various downstream tasks. Our code is available at https://github.com/kunmingsu07/RI-MAE.