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

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.

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

The paper "NeRF-MAE: Masked AutoEncoders for Self Supervised 3D Representation Learning for Neural Radiance Fields" introduces an innovative approach for pretraining neural radiance fields (NeRF) using masked autoencoders (MAE). This technique leverages the volumetric data structure of NeRFs for learning effective 3D representations from 2D posed images in a self-supervised manner. The authors propose NeRF-MAE, which entails employing a Swin Transformer-based architecture, optimized for processing NeRF-generated 3D radiance and density grids, alongside an opacity-aware masked reconstruction objective.

Contributions

1. Self-Supervised Learning Framework: The paper presents a method for large-scale self-supervised learning by effectively utilizing NeRF's volumetric nature. The approach contrasts with prior 3D representation techniques that often rely on irregular data structures such as point clouds or depth maps.
2. Transformer Architecture: NeRF-MAE employs a standard 3D Swin Transformer as the backbone encoders, demonstrating the transformers' adaptability in capturing both local and global scene attributes from the dense volumetric data provided by NeRFs.
3. Scalability and Transfer Learning: The paper highlights the scalability of NeRF-MAE to pretrain using over 1.6 million images across diverse datasets like Front3D, HM3D, and Hypersim. This pretraining significantly enhances the performance of various downstream 3D tasks, demonstrating substantial improvements in 3D object detection, semantic voxel labeling, and voxel super-resolution.
4. Improvements Over Baselines: Comparatively, NeRF-MAE surpasses state-of-the-art baselines in challenging tasks. For instance, it outperforms competitors in 3D object bounding box prediction with absolute performance gains of over 20% AP50 on datasets like Front3D and ScanNet.

Technical Insights

• The 3D volumetric grid extracted from NeRFs is used to construct semantically rich and spatially aware 3D representations, aligning with the characteristics of structured data seen in images.
• Masked autoencoders are applied in a novel way to 3D grids, reconstructing randomly masked spatial patches via a reconstruction loss that considers both radiance and opacity, enhancing the learning of scene semantics and structure.
• The pretraining methodology successfully bridges modalities between 2D and 3D data, employing these representations for effective transfer learning across scene understanding tasks.

Implications and Future Directions

The implications of NeRF-MAE are significant for the fields of computer vision and robotics, where robust 3D representations are crucial. By learning representations directly from dense 3D data, this method enhances tasks such as 3D object detection and scene segmentation without requiring detailed annotations, which are often expensive to obtain.

Future directions could involve the integration of more efficient transformer networks, potentially utilizing linear attention mechanisms to reduce computational complexity and speed training. Additionally, advancing the communication between neural rendering and masked reconstruction could further optimize the learning process. Cross-disciplinary applications, such as improving real-time SLAM systems with better 3D scene understanding, demonstrate promising potential for extended research based on the foundation laid by NeRF-MAE.

In summary, NeRF-MAE demonstrates a significant advancement in utilizing self-supervised learning techniques for 3D representation learning, highlighting the potential for broader applications in 3D computer vision tasks and beyond.