[论文解读] GeoMAE: Masked Geometric Target Prediction for Self-supervised Point Cloud Pre-Training

Xiao Yu Tian, Haoxi Ran+2 more|arXiv (Cornell University)|May 15, 2023

3D Surveying and Cultural HeritageEarth and Planetary Sciences被引用 3

一句话总结

GeoMAE 提出了一种用于点云的自监督预训练框架，通过掩码几何目标预测（具体为质心、法线和曲率预测）以及占据预测，学习鲁棒的三维表征。通过在掩码体素化点云上使用基于 Transformer 的编码器-解码器进行训练，并结合这些几何感知目标，GeoMAE 在 3D 检测、分割和多目标跟踪任务上实现了最先进性能，在 nuScenes 数据集上分别实现了 3.38 mAP、2.1 mIoU 和 1.7 AMOTA 的性能提升。

ABSTRACT

This paper tries to address a fundamental question in point cloud self-supervised learning: what is a good signal we should leverage to learn features from point clouds without annotations? To answer that, we introduce a point cloud representation learning framework, based on geometric feature reconstruction. In contrast to recent papers that directly adopt masked autoencoder (MAE) and only predict original coordinates or occupancy from masked point clouds, our method revisits differences between images and point clouds and identifies three self-supervised learning objectives peculiar to point clouds, namely centroid prediction, normal estimation, and curvature prediction. Combined with occupancy prediction, these four objectives yield an nontrivial self-supervised learning task and mutually facilitate models to better reason fine-grained geometry of point clouds. Our pipeline is conceptually simple and it consists of two major steps: first, it randomly masks out groups of points, followed by a Transformer-based point cloud encoder; second, a lightweight Transformer decoder predicts centroid, normal, and curvature for points in each voxel. We transfer the pre-trained Transformer encoder to a downstream peception model. On the nuScene Datset, our model achieves 3.38 mAP improvment for object detection, 2.1 mIoU gain for segmentation, and 1.7 AMOTA gain for multi-object tracking. We also conduct experiments on the Waymo Open Dataset and achieve significant performance improvements over baselines as well.

研究动机与目标

• 为解决点云表征学习中缺乏有效自监督学习目标的问题，这些目标能利用内在几何特性。
• 识别并利用点云中独特且信息丰富的几何特征（如质心、法线和曲率），这些特征在现有方法中尚未被充分使用。
• 通过在大规模未标注点云上使用几何感知的代理任务进行预训练，提升下游 3D 感知任务（检测、分割、跟踪）的性能。
• 证明几何特征重建作为自监督信号比单独的坐标或占据预测更有效。
• 验证在统一的掩码自编码框架中结合多个几何预测头对点云表征学习的有效性。

提出的方法

• 该方法首先将输入点云体素化为 3D 网格，将每个体素转换为可学习的特征 token。
• 应用随机掩码以丢弃体素 token 的组，模拟类似于图像 MAE 的掩码输入，但针对 3D 几何特性进行了适配。
• 可见的 token 由稀疏 Transformer 编码器编码，以从剩余点中学习全局几何上下文。
• 轻量级 Transformer 解码器并行预测每个掩码体素的四个几何目标：占据、质心、法线和曲率。
• 使用结合 L2 损失的多任务损失对模型进行预训练，分别用于占据、质心、法线和曲率预测。
• 在下游 3D 感知任务上，使用相同的主干架构对预训练编码器进行微调，无需额外的预训练数据。

实验结果

研究问题

• RQ1 质心、法线和曲率等几何特征能否作为点云表征学习的有效自监督信号？
• RQ2 与单目标掩码自编码相比，组合多个几何预测目标（占据、质心、法线、曲率）如何提升特征学习？
• RQ3 所提出的几何感知预训练框架是否在下游 3D 感知任务上优于现有的对比学习和掩码建模基线方法？
• RQ4 几何特征在 3D 检测、分割和多目标跟踪中对泛化能力的贡献程度如何？
• RQ5 预训练数据规模如何影响所提出的 GeoMAE 框架的性能提升？

主要发现

• 在 nuScenes 数据集上，GeoMAE 相较于基线方法在 3D 目标检测任务上实现了 3.38 mAP 的提升，在 3D 分割任务上实现了 2.1 mIoU 的增益。
• 在 nuScenes 上，模型在多目标跟踪任务上实现了 1.7 AMOTA 的性能提升，表明其在序列感知任务中具有强大的泛化能力。
• 在 Waymo Open Dataset 上，GeoMAE 相较于现有自监督方法实现了显著的性能提升，证实了其在多样化户外点云基准上的泛化能力。
• 消融研究显示，每个几何预测头均对整体性能有所贡献，且四种目标的组合实现了最佳结果。
• GeoMAE 的性能随着预训练数据集规模的增大而提升，表明所提出的几何目标具有可扩展性和鲁棒性。
• 该方法在所有评估的下游任务上均优于对比学习基线方法（如 PointContrast、STRL、BYOL 和 SwAV），确立了自监督 3D 点云学习的新 SOTA。