NeRF-MAE: Masked AutoEncoders for Self-Supervised 3D Representation Learning for Neural Radiance Fields

TL;DR

NeRF-MAE targets scalable, self-supervised 3D representation learning for Neural Radiance Fields by operating directly on a dense radiance-density grid extracted from NeRFs. The method masks random 3D patches in a canonical 4D grid and trains a standard 3D Swin Transformer encoder with lightweight decoders to reconstruct the masked regions, optimizing both radiance and opacity terms. Trained on over 1.8 million posed RGB images from Front3D, HM3D, Hypersim, and ARKitScenes, NeRF-MAE achieves strong transfer performance on downstream 3D tasks such as 3D object detection, voxel labeling, and voxel super-resolution, outperforming self-supervised 3D pretraining and NeRF-scene baselines by substantial margins. This approach demonstrates that dense NeRF grids enable effective Transformer-based pretraining, offering data-efficient 3D learning from unlabeled 2D data and opening avenues for future improvements in efficiency and cross-view neural rendering collaboration.

Abstract

Neural fields excel in computer vision and robotics due to their ability to understand the 3D visual world such as inferring semantics, geometry, and dynamics. Given the capabilities of neural fields in densely representing a 3D scene from 2D images, we ask the question: Can we scale their self-supervised pretraining, specifically using masked autoencoders, to generate effective 3D representations from posed RGB images. Owing to the astounding success of extending transformers to novel data modalities, we employ standard 3D Vision Transformers to suit the unique formulation of NeRFs. We leverage NeRF's volumetric grid as a dense input to the transformer, contrasting it with other 3D representations such as pointclouds where the information density can be uneven, and the representation is irregular. Due to the difficulty of applying masked autoencoders to an implicit representation, such as NeRF, we opt for extracting an explicit representation that canonicalizes scenes across domains by employing the camera trajectory for sampling. Our goal is made possible by masking random patches from NeRF's radiance and density grid and employing a standard 3D Swin Transformer to reconstruct the masked patches. In doing so, the model can learn the semantic and spatial structure of complete scenes. We pretrain this representation at scale on our proposed curated posed-RGB data, totaling over 1.8 million images. Once pretrained, the encoder is used for effective 3D transfer learning. Our novel self-supervised pretraining for NeRFs, NeRF-MAE, scales remarkably well and improves performance on various challenging 3D tasks. Utilizing unlabeled posed 2D data for pretraining, NeRF-MAE significantly outperforms self-supervised 3D pretraining and NeRF scene understanding baselines on Front3D and ScanNet datasets with an absolute performance improvement of over 20% AP50 and 8% AP25 for 3D object detection.

Figures (13)

• Figure 1: NeRF-MAE Overview: The first large-scale self-supervised pretraining utilizing NeRF's radiance and density grid as an input modality. Our approach uses standard Transformers to learn powerful 3D representations in (a) an opacity-aware dense volumetric masked self-supervised learning objective. (b) when fine-tuned on a small subset of data, our representation improves many 3D downstream tasks such as 3D object detection, super-resolution, and voxel-labeling.
• Figure 2: Quantitative results showing our representation improves with more unsupervised data and as well as better reconstruction quality of NeRFs.
• Figure 3: NeRF-MAE dataset mix:a) Multi-view dataset with camera distribution, b) diverse scenes from different sources i.e. Front3D fu20213d, Hypersim roberts2021hypersim, HM3D ramakrishnan2021hm3d and ScanNet dai2017scannet totaling over 3600 scenes and 1.8M images used for pertaining and evaluating NeRF-MAE.
• Figure 4: NeRF-MAE Comparison: In § \ref{['sec:pretraining_3D_baelines']}, we compare our method to point-based pretraining approaches, which are limited by their modeling of surface-level sparse points. In contrast, NeRF-MAE resembles images in terms of high information density and structural regularity, hence positioning it as a direct extension of image MAE to 3D. Leveraging NeRF's dense volumetric information and opacity-aware reconstruction loss, we achieve superior representation learning.
• Figure 5: NeRF-MAE data processing flow for pretraining: showing a) multi-view training data, b) trained NeRF representation, c) extracted radiance and density grid from NeRF and d) masked pretraining of the radiance and density voxel-grid neural radiance field.