[Paper Review] RangeViT: Towards Vision Transformers for 3D Semantic Segmentation in Autonomous Driving

Angelika Ando, Spyros Gidaris+4 more|arXiv (Cornell University)|Jan 24, 2023

Advanced Neural Network Applications4 citations

TL;DR

RangeViT proposes a vision transformer (ViT)-based method for 3D semantic segmentation of LiDAR point clouds using range projection, leveraging pre-trained ViTs on ImageNet and enhancing performance with a convolutional stem and a lightweight convolutional decoder. It achieves state-of-the-art results among 2D projection-based methods on nuScenes and SemanticKITTI, demonstrating that ViTs pre-trained on natural images can be effectively transferred to sparse, noisy LiDAR data with minimal architectural modifications.

ABSTRACT

Casting semantic segmentation of outdoor LiDAR point clouds as a 2D problem, e.g., via range projection, is an effective and popular approach. These projection-based methods usually benefit from fast computations and, when combined with techniques which use other point cloud representations, achieve state-of-the-art results. Today, projection-based methods leverage 2D CNNs but recent advances in computer vision show that vision transformers (ViTs) have achieved state-of-the-art results in many image-based benchmarks. In this work, we question if projection-based methods for 3D semantic segmentation can benefit from these latest improvements on ViTs. We answer positively but only after combining them with three key ingredients: (a) ViTs are notoriously hard to train and require a lot of training data to learn powerful representations. By preserving the same backbone architecture as for RGB images, we can exploit the knowledge from long training on large image collections that are much cheaper to acquire and annotate than point clouds. We reach our best results with pre-trained ViTs on large image datasets. (b) We compensate ViTs' lack of inductive bias by substituting a tailored convolutional stem for the classical linear embedding layer. (c) We refine pixel-wise predictions with a convolutional decoder and a skip connection from the convolutional stem to combine low-level but fine-grained features of the the convolutional stem with the high-level but coarse predictions of the ViT encoder. With these ingredients, we show that our method, called RangeViT, outperforms existing projection-based methods on nuScenes and SemanticKITTI. The code is available at https://github.com/valeoai/rangevit.

Motivation & Objective

• To investigate whether vision transformers (ViTs) can improve 2D projection-based 3D semantic segmentation of LiDAR point clouds.
• To address ViTs' lack of inductive bias and poor performance on small datasets by adapting the input processing and introducing a hybrid encoder-decoder architecture.
• To evaluate the transferability of ImageNet-pretrained ViTs to LiDAR data despite the domain gap between natural images and range projections.
• To unify network architectures across 2D image and 3D point cloud tasks by reusing ViT backbones across modalities.
• To establish a simple yet effective baseline for future work on ViT-based 3D segmentation using sparse LiDAR data.

Proposed method

• The method performs range projection of 3D LiDAR point clouds into a 2D bird's-eye view feature map.
• Visual tokens are extracted from the projected 2D map using a patch-based tokenization strategy, similar to standard ViT.
• A pre-trained Vision Transformer encoder (e.g., ViT-S) is used to extract high-level representations from the tokens, with weights initialized from ImageNet-21k pre-training.
• A multi-layer convolutional stem replaces the standard linear embedding layer to inject inductive bias and improve feature learning on sparse range projections.
• A lightweight convolutional decoder with skip connections from the convolutional stem refines coarse ViT predictions by fusing low-level spatial details.
• Pixel-wise predictions are upsampled and projected back to 3D space for final semantic labeling.

Experimental results

Research questions

• RQ1 Can vision transformers pre-trained on large-scale natural image datasets be effectively fine-tuned for 3D LiDAR semantic segmentation despite the domain gap?
• RQ2 Does replacing the standard linear embedding layer in ViTs with a convolutional stem improve performance on sparse, noisy LiDAR range projections?
• RQ3 Can a lightweight convolutional decoder with skip connections from the stem enhance the quality of ViT-based predictions in 3D segmentation?
• RQ4 How does the performance of ViT-based models compare to CNN-based projection methods on standard autonomous driving benchmarks?
• RQ5 To what extent do pre-trained ViTs outperform randomly initialized or non-pretrained ViTs in this low-data regime?

Key findings

• RangeViT achieves a mean Intersection over Union (mIoU) of 74.77% on the nuScenes validation set, outperforming all prior 2D projection-based methods.
• On the SemanticKITTI test set, RangeViT achieves an mIoU of 70.6%, ranking first among 2D projection-based methods and closing the gap with voxel-based models like Cylinder3D.
• Ablation studies show that using a ViT-S backbone pre-trained on ImageNet-21k yields 14.3% higher mIoU than a ResNet-50 backbone under the same fine-tuning protocol.
• The ablation with an identity backbone (no feature extraction) yields only 53.73% mIoU, confirming that ViT features are more suitable for transfer learning than simple linear projections.
• The combination of pre-trained ViTs, a convolutional stem, and a decoder with skip connections leads to significant performance gains, especially on rare and low-density classes like pedestrians and bicycles.
• Qualitative results show that RangeViT reduces border artifacts and improves segmentation of complex structures, though some confusion remains between driveable surfaces and sidewalks.