CodedVTR: Codebook-based Sparse Voxel Transformer with Geometric Guidance

TL;DR

CodedVTR introduces a codebook-based and geometry-aware attention mechanism to enhance data efficiency and generalization in 3D sparse voxel transformers, improving semantic segmentation performance.

Contribution

It presents a novel codebook-based attention and geometry-aware self-attention for 3D transformers, addressing generalization and data efficiency issues.

Findings

Improves semantic segmentation accuracy in 3D tasks.

Enhances generalization with codebook-based attention.

Achieves consistent performance gains in indoor and outdoor scenarios.

Abstract

Transformers have gained much attention by outperforming convolutional neural networks in many 2D vision tasks. However, they are known to have generalization problems and rely on massive-scale pre-training and sophisticated training techniques. When applying to 3D tasks, the irregular data structure and limited data scale add to the difficulty of transformer's application. We propose CodedVTR (Codebook-based Voxel TRansformer), which improves data efficiency and generalization ability for 3D sparse voxel transformers. On the one hand, we propose the codebook-based attention that projects an attention space into its subspace represented by the combination of "prototypes" in a learnable codebook. It regularizes attention learning and improves generalization. On the other hand, we propose geometry-aware self-attention that utilizes geometric information (geometric pattern, density) to guide attention learning. CodedVTR could be embedded into existing sparse convolution-based methods, and bring consistent performance improvements for indoor and outdoor 3D semantic segmentation tasks