Towards Compact 3D Representations via Point Feature Enhancement Masked Autoencoders

TL;DR

This work tackles the inefficiency and limited representational richness of existing MAE-based point-cloud pretraining by introducing Point-FEMAE, a dual-branch framework that jointly learns global and local features. A shared Transformer encoder processes unmasked patches, while a Local Enhancement Module in the local branch captures fine-grained local context, guided by a hybrid global-local masking scheme. Reconstruction is performed in both branches with separate decoders, using Chamfer Distance as the loss, and the encoder from the local branch is reused during fine-tuning for downstream tasks. Empirical results on ShapeNet pretraining, plus downstream evaluations on ScanObjectNN and ModelNet40, show substantial improvements over Point-MAE and competitive efficiency versus cross-modal methods, backed by thorough ablations. The findings demonstrate that carefully designed masking and local feature augmentation yield compact, high-quality 3D representations with practical gains for real-world point-cloud understanding.

Abstract

Learning 3D representation plays a critical role in masked autoencoder (MAE) based pre-training methods for point cloud, including single-modal and cross-modal based MAE. Specifically, although cross-modal MAE methods learn strong 3D representations via the auxiliary of other modal knowledge, they often suffer from heavy computational burdens and heavily rely on massive cross-modal data pairs that are often unavailable, which hinders their applications in practice. Instead, single-modal methods with solely point clouds as input are preferred in real applications due to their simplicity and efficiency. However, such methods easily suffer from limited 3D representations with global random mask input. To learn compact 3D representations, we propose a simple yet effective Point Feature Enhancement Masked Autoencoders (Point-FEMAE), which mainly consists of a global branch and a local branch to capture latent semantic features. Specifically, to learn more compact features, a share-parameter Transformer encoder is introduced to extract point features from the global and local unmasked patches obtained by global random and local block mask strategies, followed by a specific decoder to reconstruct. Meanwhile, to further enhance features in the local branch, we propose a Local Enhancement Module with local patch convolution to perceive fine-grained local context at larger scales. Our method significantly improves the pre-training efficiency compared to cross-modal alternatives, and extensive downstream experiments underscore the state-of-the-art effectiveness, particularly outperforming our baseline (Point-MAE) by 5.16%, 5.00%, and 5.04% in three variants of ScanObjectNN, respectively. The code is available at https://github.com/zyh16143998882/AAAI24-PointFEMAE.

Figures (7)

• Figure 1: Differences in data distribution between pre-training and fine-tuning. (a) Global Masked Point Cloud (GMPC) input during pre-training with global random masking. (b) Local Masked Point Cloud (LMPC) input during pre-training with local block masking. (c) Complete Point Cloud (CPC) input during downstream fine-tuning.
• Figure 2: The pipeline of our Point-FEMAE. During the pre-training stage, we perform mask reconstruction in both the global and local branches to learn compact 3D representations. During the fine-tuning stage, we only employ the encoder of the local branch to learn the 3D representation of downstream data.
• Figure 3: The Encoder Layer's structure, where each layer incorporates a globally oriented Transformer Block and a locally oriented LEM. Within the LEM, information from $k$ nearest neighbor patches is fused based on the patches' coordinates, facilitating a broader scope of local perception.
• Figure 4: Structure of the Local Enhancement Module (LEM) with Scalable Parameters.
• Figure 5: Effect of different $K$-values in LEM.