PaperlayerView exact revision on arXiv

Source-linked AI summary

Swin3D: A Pretrained Transformer Backbone for 3D Indoor Scene Understanding

Yu-Qi Yang, Yu-Xiao Guo, Jian-Yu Xiong, Yang Liu, Hao Pan, Peng-Shuai Wang, Xin Tong, Baining Guo

arXiv:2304.06906v3cs.CV

TL;DR

Pretrained 3D backbones remain limited in scalability and effectiveness for point-cloud understanding. SWIN3D adapts hierarchical Swin Transformer processing to sparse voxels with linear-memory attention and contextual signal encoding, then pretrains on large synthetic data. The resulting backbone transfers across real indoor segmentation and detection tasks and outperforms existing state-of-the-art methods on the reported benchmarks.

  • Problem

    Generic and scalable pretrained 3D backbones remain underdeveloped and often trail state-of-the-art non-pretrained methods across 3D vision tasks.

  • Method

    SWIN3D uses hierarchical sparse-voxel Swin processing, linear-memory self-attention, contextual relative signal encoding, and synthetic Structured3D pretraining.

  • Results

    SWIN3D pretrained on synthetic data generalizes across real ScanNet and S3DIS segmentation and detection, outperforming existing state-of-the-art methods on the reported tasks.

  • Takeaways & Limitations

    The results support a unified pretrained backbone for multiple 3D indoor scene understanding tasks, with large model capacity and pretraining data identified as critical.

  • Takeaways & Limitations

    Larger models tended to overfit the pretraining dataset, so models beyond SWIN3D-L were not explored further.

Abstract

from arXiv · show

The use of pretrained backbones with fine-tuning has been successful for 2D vision and natural language processing tasks, showing advantages over task-specific networks. In this work, we introduce a pretrained 3D backbone, called {\SST}, for 3D indoor scene understanding. We design a 3D Swin transformer as our backbone network, which enables efficient self-attention on sparse voxels with linear memory complexity, making the backbone scalable to large models and datasets. We also introduce a generalized contextual relative positional embedding scheme to capture various irregularities of point signals for improved network performance. We pretrained a large {\SST} model on a synthetic Structured3D dataset, which is an order of magnitude larger than the ScanNet dataset. Our model pretrained on the synthetic dataset not only generalizes well to downstream segmentation and detection on real 3D point datasets, but also outperforms state-of-the-art methods on downstream tasks with +2.3 mIoU and +2.2 mIoU on S3DIS Area5 and 6-fold semantic segmentation, +1.8 mIoU on ScanNet segmentation (val), +1.9 mAP@0.5 on ScanNet detection, and +8.1 mAP@0.5 on S3DIS detection. A series of extensive ablation studies further validate the scalability, generality, and superior performance enabled by our approach. The code and models are available at https://github.com/microsoft/Swin3D .

1. Introduction

SWIN3D addresses the limited scalability and effectiveness of pretrained 3D backbones by adapting Swin Transformer to sparse, unorganized 3D points. Its design supports large-scale synthetic pretraining and strong transfer to real indoor scene tasks.

  • Pretrained 3D backbones remain early-stage and often trail state-of-the-art non-pretrained methods on 3D vision tasks.
  • SWIN3D adapts Swin Transformer to sparse voxels and targets scalability, generality, and reduced network-design complexity.
  • Linear-memory sparse-voxel self-attention addresses high memory complexity, while contextual relative signal encoding captures irregular point signals.
  • A 60 M-parameter SWIN3D model was pretrained on Structured3D, a synthetic dataset ten times larger than ScanNet, then fine-tuned with task-specific decoders.
  • +8.1 mAP@0.5 on S3DIS detection, +2.2 mIoU on 6-fold S3DIS segmentation, +1.9 mAP@0.5 on ScanNet detection, and +1.8 mIoU on ScanNet segmentation were reported.
  • The authors attribute downstream performance to the large backbone and large pretraining dataset and release code and trained models.

2. Related Work

Prior work applied transformers and pretraining to point clouds, but most 3D transformers were tailored to particular tasks. SWIN3D belongs to the effort to develop more general pretrained 3D backbones.

  • Vision transformers: Vision transformers use attention for 2D image classification, semantic segmentation, and object detection, with global attention offering long-range interactions but high cost.
  • 3D transformers for point cloud understanding: 3D point-cloud transformers explored global and local attention, grouped vector attention, voxel hashing, and partition-based pooling for efficiency and scene-level processing.
  • 3D transformers for point cloud understanding: Stratified Transformer adapted Swin design to 3D, expanded the receptive field through stratified windows, and used contextual relative positional encoding.
  • Pretrained 3D backbones: Self-supervised methods pretrained sparse-convolutional, octree-based, and transformer backbones using point-level, contrastive, spatial, depth-map, or masked-signal objectives.

3. Architecture overview

Naively extending window attention to 3D is hindered by high memory cost and irregular point signals. SWIN3D combines sparse-voxel window attention, contextual relative signal encoding, and hierarchical multiscale processing to address these issues.

  • Naive 3D extension of Swin Transformer: Naive 3D Swin attention partitions point clouds into regular and shifted windows, but this extension does not produce superior performance.
  • Memory complexity: 3D self-attention can require about O(NwM^4) memory because each M × M × M window averages O(M^2) non-empty voxels.
  • Signal irregularity: 3D points may occur anywhere within occupied voxels and carry varying signals such as RGB, creating irregular positional and pointwise-signal relationships.
  • SWIN3D architecture: SWIN3D uses sparse voxel features within regular and shifted windows, memory-efficient self-attention, and contextual relative signal encoding.
  • SWIN3D architecture: Its five-stage hierarchical blocks operate at different voxel resolutions and provide a multiscale encoder compatible with task-specific decoders.

4.1. Voxelization

Voxelization converts unstructured point clouds into a five-level hierarchical sparse voxel grid. Representative points provide raw features and contextual information for the backbone.

  • Point cloud input: Point-cloud inputs typically contain point position, color, and normal signals; a common setting concatenates 3D coordinates with RGB color.
  • Voxelization: The method constructs a five-level hierarchical sparse voxel grid, starting with a 2 cm finest-level voxel size for indoor scenes and doubling it per level.
  • Representative points: At the finest level, one point is randomly selected from each voxel as its representative point.
  • Representative points: At coarser levels, the representative point closest to the center of child representatives becomes the parent voxel’s representative point.
  • Voxelization: Voxelization assigns unstructured points to sparse grids, while representative-point signals feed initial embedding and contextual self-attention.

4.2. Initial feature embedding

The initial feature embedding lifts raw voxel features into a higher-dimensional representation using sparse convolution, batch normalization, and ReLU. It uses positional offsets and point signals while avoiding absolute positions.

  • Initial feature embedding: Each voxel feature concatenates its positional offset rv − cv with other point signals stored at rv.Absolute point position is omitted so convolution can learn local priors.
  • Initial feature embedding: Compared to KPConv utilized by [28], the initial feature embedding is much lighter and five times faster.The comparison concerns the embedding design used before transformer processing.

4.3. SWIN3D block

The SWIN3D block adapts hierarchical regular and shifted-window Swin attention to sparse 3D voxels, addressing memory growth and irregular point signals. It combines memory-efficient self-attention with contextual relative signal encoding and standard transformer components.

  • SWIN3D block: SWIN3D operates on regular and shifted 3D windows, modifying multi-head self-attention for memory efficiency and signal irregularity.The block includes LayerNorm and an MLP layer alongside the revised attention module.
  • Memory-efficient self-attention: O(NwM^4) memory cost arises in vanilla 3D window attention because each window contains approximately O(M^2) non-empty voxels.The large number of 3D windows makes memory usage substantially higher than in typical fixed-size 2D images.
  • Memory-efficient self-attention: The revised attention postpones SoftMax normalization, computes numerator and denominator simultaneously, and avoids storing all attention coefficients.Exponentiated attention scores are computed on the fly; during training they are recomputed twice, with negligible additional computation cost reported.
  • Contextual relative signal encoding: Contextual relative signal encoding extends positional encoding to differences in point positions, RGB colors, and normal components.The trainable functions are implemented with learnable lookup tables, using signal-specific quantization ranges and table lengths.
  • Self-attention formulation: The attention score uses scaled query-key dot products plus contextual signal encoding, while the output remains a weighted sum of projected voxel features.Query, Key, and Value are produced by linear projection matrices, and the contextual functions map signal differences to the head channel dimension.

4.4. Downsample

SWIN3D downsamples sparse voxel features between hierarchy levels using KNNPooling. Features from 16 nearest voxels at the finer level are maxpooled and assigned to each coarser-level voxel.

  • Downsample: KNNPooling downsamples level-l sparse voxel features by lifting them through LayerNorm and a fully connected layer, then maxpooling k-nearest features.The default neighborhood size is k = 16.

5. SWIN3D Backbone pretraining

This section evaluates SWIN3D’s efficiency, scalability, and pretraining setup, including memory-efficient self-attention, Structured3D data preparation, and downstream fine-tuning.

  • Model efficiency: Memory-efficient self-attention saves around 25% GPU memory versus vanilla self-attention, with quadratic savings as head number or window size increases.Memory was measured during forward-and-backward iterations on ScanNet data.
  • Support to large models: SWIN3D’s memory reduction makes the backbone suitable for designing large models, although larger models than SWIN3D-L overfit the pretraining dataset.The authors therefore did not explore those larger models further.
  • Pretraining data preparation: Structured3D provides 21835 rooms in 3500 synthetic scenes, making it one order larger than ScanNet and other real indoor datasets described.RGBD and panoramic images are projected into 3D points using camera parameters, then reduced by retaining one point per 1 cm^3 cube.
  • Backbone pretraining: SWIN3D is pretrained for 3D semantic segmentation with 25 labels, using an encoder and simple UNet-like decoder with interpolation and skip connections.Four rare labels—counter, box, toilet, and bathtub—are excluded from pretraining.
  • Fine-tuning for downstream tasks: The pretrained backbone is loaded as a multi-resolution feature encoder, paired with task-specific decoders whose weights are randomly initialized for downstream fine-tuning.The pretrained network weights and look-up tables initialize the encoder.

6. Experimental Analysis

Experiments evaluate SWIN3D across segmentation, detection, pretraining, and scalability settings on ScanNet, S3DIS, and Structured3D. Pretraining, larger models, and more synthetic data consistently improve downstream performance, while SWIN3D surpasses compared methods across tasks.

  • Semantic segmentation: SWIN3D-L exceeds supervised methods by 1.0 mIoU on ScanNet val, +2.3 mIoU on S3DIS Area5, and +2.2 mIoU on S3DIS 6-fold.Using point normals further improves SWIN3Dn-L by +0.8 mIoU on ScanNet val.
  • Semantic segmentation: Unsupervised pretrained methods remain below supervised methods and SWIN3D, with restricted benefits attributed to convolutional encoders’ weaker extraction of data priors.The paper contrasts these methods with transformer-based SWIN3D.
  • Pretraining ablation: Pretrained SWIN3D variants consistently outperform their from-scratch counterparts, demonstrating the effectiveness of pretraining.Larger SWIN3D models trained from scratch overfit because ScanNet and S3DIS provide limited training data.
  • 3D detection: SWIN3D-L+FCAF3D improves S3DIS detection by 8.1 points in mAP@0.5, while SWIN3D variants improve ScanNet detection over their task-specific baselines.On ScanNet, SWIN3D-S+FCAF3D and SWIN3D-L+FCAF3D each gain 2.7 points in mAP@0.25; SWIN3D-L+CAGroup3D achieves a new mAP@0.50 record.
  • Model scalability: Both SWIN3D-S and SWIN3D-L improve substantially as pretraining data increases, with SWIN3D-L benefiting more from large datasets.The study uses 10%, 33%, and 100% of Structured3D for pretraining.
  • Ablation study: Pretraining on extensive synthetic Structured3D data benefits the backbones more than pretraining on limited real ScanNet data, and cRSE scales better than cRPE.The ablations identify training-data size as the most critical factor and show that capturing irregular signals is important for backbone learning.

7. Conclusion

The paper concludes that SWIN3D is a scalable and general pretrained 3D backbone for indoor scene understanding. Future work will extend pretraining with additional data and combine image and 3D backbones.

  • Conclusion: SWIN3D demonstrates scalability, generality, and superior performance for indoor scene understanding through extensive experiments.The conclusion identifies SWIN3D as a novel 3D backbone.
  • Future work: Future work includes self-supervised pretraining with more real and synthetic data, including outdoor 3D data.The proposed direction aims to enrich SWIN3D’s representation.
  • Future work: Future work also proposes integrating pretrained image and 3D backbones, leveraging high-resolution multiview images accompanying many 3D point clouds.The paper frames this as a way to enhance 3D learning.

Appendix: Efficient self-attention implementation

The implementation improves efficient self-attention through GPU-oriented kernel and memory optimizations. These changes reduce memory access and atomic-operation overhead while supporting half-precision execution.

  • Kernel scheduling: Kernel scheduling addresses variable point counts across windows, which otherwise makes CUDA execution difficult to optimize.The design targets better use of GPU memory bandwidth and computational capability.
  • Half-precision support: Half-precision support is enabled for implemented CUDA kernels, with channel-last lookup-table layouts allowing paired 16-bit elements in one 32-bit access.The reorganization applies to query, key, and lookup-table memory accesses.
  • Computation speedup: Combining coefficient-weight and cRPE/cRSE calculations in one kernel reduces query-key memory access and accelerates GPU computation.The optimization fuses related self-attention operations.
  • Atomic operations: Shared memory combines two consecutive 16-bit atomic operations into one 32-bit operation, reducing atomic-operation overhead.This addresses inefficiency from half-precision atomic conflicts involving double writing and reading.
Loading 2304.06906v3…