[Paper Review] Masked Autoencoders in 3D Point Cloud Representation Learning

Jincen Jiang, Xuequan Lu+3 more|arXiv (Cornell University)|Jul 4, 2022

3D Shape Modeling and Analysis4 citations

TL;DR

This paper proposes MAE3D, a novel masked autoencoder framework for self-supervised 3D point cloud representation learning that splits point clouds into patches, masks a portion, and uses a patch-wise Transformer to reconstruct missing parts via a multi-task loss. It achieves state-of-the-art performance with 93.4% accuracy on ModelNet40 and 86.2% on ScanObjectNN (PB_T50_RS), outperforming prior methods while using a lightweight backbone with only 1.8M parameters.

ABSTRACT

Transformer-based Self-supervised Representation Learning methods learn generic features from unlabeled datasets for providing useful network initialization parameters for downstream tasks. Recently, self-supervised learning based upon masking local surface patches for 3D point cloud data has been under-explored. In this paper, we propose masked Autoencoders in 3D point cloud representation learning (abbreviated as MAE3D), a novel autoencoding paradigm for self-supervised learning. We first split the input point cloud into patches and mask a portion of them, then use our Patch Embedding Module to extract the features of unmasked patches. Secondly, we employ patch-wise MAE3D Transformers to learn both local features of point cloud patches and high-level contextual relationships between patches and complete the latent representations of masked patches. We use our Point Cloud Reconstruction Module with multi-task loss to complete the incomplete point cloud as a result. We conduct self-supervised pre-training on ShapeNet55 with the point cloud completion pre-text task and fine-tune the pre-trained model on ModelNet40 and ScanObjectNN (PB\_T50\_RS, the hardest variant). Comprehensive experiments demonstrate that the local features extracted by our MAE3D from point cloud patches are beneficial for downstream classification tasks, soundly outperforming state-of-the-art methods ($93.4\%$ and $86.2\%$ classification accuracy, respectively).

Motivation & Objective

• To address the challenge of applying masked autoencoder pre-training to irregular 3D point clouds, which lack the regular structure of images or text.
• To develop a general and lightweight self-supervised framework that avoids complex tokenizers or heavy encoders.
• To improve point cloud representation learning by jointly capturing local patch features and global contextual relationships through a patch-wise Transformer.
• To enable strong downstream performance using only the backbone network, without requiring a separate, large Transformer encoder during fine-tuning.
• To demonstrate that pre-training with masked reconstruction enhances feature discriminability and generalization on 3D classification tasks.

Proposed method

• The input point cloud is divided into spatial patches, and a portion of these patches are randomly masked, similar to masked language modeling in BERT.
• A Patch Embedding Module extracts features from unmasked patches using a lightweight backbone (e.g., PointNet or DGCNN).
• A patch-wise Transformer encoder processes the visible patches to learn both local geometric features and high-level contextual relationships across patches.
• A Point Cloud Reconstruction Module uses a multi-task loss to predict and reconstruct the coordinates of the masked patches from the latent representations.
• The model is pre-trained on ShapeNet55 using a point cloud completion pre-text task, followed by fine-tuning on downstream classification benchmarks.
• During inference, only the pre-trained backbone is used for feature extraction, eliminating the need for a heavy decoder or auxiliary Transformer encoder.

Experimental results

Research questions

• RQ1 Can a masked autoencoder framework be effectively adapted to irregular 3D point cloud data, which lacks the regular grid structure of images?
• RQ2 Does learning to reconstruct masked point cloud patches improve the quality of learned representations for downstream 3D classification tasks?
• RQ3 Can a lightweight backbone with minimal parameters (1.8M) outperform heavier models that rely on additional Transformer encoders?
• RQ4 Does the absence of a complex encoder during fine-tuning lead to more robust and generalizable features?
• RQ5 How does the inclusion of a dedicated MAE3D Transformer for latent representation learning affect final classification accuracy?

Key findings

• MAE3D achieves 93.4% classification accuracy on ModelNet40, surpassing state-of-the-art methods.
• On the challenging ScanObjectNN (PB_T50_RS) benchmark, MAE3D achieves 86.2% accuracy, demonstrating strong robustness to noise and object variations.
• The inclusion of MAE3D Transformers during pre-training improves accuracy by 1% compared to a version without them, confirming their contribution to feature learning.
• Despite using only 1.8M parameters, MAE3D outperforms Point-MAE, which requires a 22.0M parameter Transformer encoder, when both are fine-tuned without the encoder.
• When using DGCNN as the backbone, MAE3D achieves 93.4% accuracy on ModelNet40, outperforming Point-MAE (93.1%) even though Point-MAE uses an additional Transformer encoder.
• t-SNE visualization confirms that MAE3D learns more separable and discriminative features than training from scratch, even during pre-training, and further improves after fine-tuning.