Source-linked AI summary
Stratified Transformer for 3D Point Cloud Segmentation
Xin Lai, Jianhui Liu, Li Jiang, Liwei Wang, Hengshuang Zhao, Shu Liu, Xiaojuan Qi, Jiaya Jia
TL;DR
3D point-cloud segmentation methods often emphasize local aggregation without directly modeling long-range dependencies. Stratified Transformer samples nearby keys densely and distant keys sparsely, adds point embedding and contextual position encoding, and reports state-of-the-art results across major segmentation datasets.
Problem
Existing point-cloud segmentation methods mainly aggregate local features and fail to explicitly model long-range dependencies, while direct global attention has unacceptable O(N^2) memory consumption.
Method
Stratified Transformer uses standard multi-head self-attention with dense nearby-key and sparse distant-key sampling, first-layer point embedding, contextual relative position encoding, and memory-efficient window processing.
Results
The method achieves state-of-the-art performance on S3DIS, ScanNetv2, and ShapeNetPart, including 1.6% mIoU over Point Transformer on S3DIS and 2.1% mIoU over competing methods on ScanNetv2 validation.
Takeaways & Limitations
The results support building direct long-range dependencies at low computational cost and applying standard Transformers to 3D point clouds with strong generalization ability.
Abstract
from arXiv · showhide
3D point cloud segmentation has made tremendous progress in recent years. Most current methods focus on aggregating local features, but fail to directly model long-range dependencies. In this paper, we propose Stratified Transformer that is able to capture long-range contexts and demonstrates strong generalization ability and high performance. Specifically, we first put forward a novel key sampling strategy. For each query point, we sample nearby points densely and distant points sparsely as its keys in a stratified way, which enables the model to enlarge the effective receptive field and enjoy long-range contexts at a low computational cost. Also, to combat the challenges posed by irregular point arrangements, we propose first-layer point embedding to aggregate local information, which facilitates convergence and boosts performance. Besides, we adopt contextual relative position encoding to adaptively capture position information. Finally, a memory-efficient implementation is introduced to overcome the issue of varying point numbers in each window. Extensive experiments demonstrate the effectiveness and superiority of our method on S3DIS, ScanNetv2 and ShapeNetPart datasets. Code is available at https://github.com/dvlab-research/Stratified-Transformer.
1. Introduction
Stratified Transformer addresses the difficulty of modeling long-range dependencies in irregular 3D point clouds while retaining efficient attention. It combines stratified key sampling with point embedding, contextual position encoding, and memory-efficient implementation, achieving state-of-the-art segmentation results.
- Motivation: Most existing point-cloud segmentation methods aggregate local features but do not explicitly model long-range dependencies.Transformer self-attention can capture long-range information, but applying it directly to 3D point clouds remains difficult.
- Stratified Transformer: Stratified key sampling densely selects nearby points and sparsely samples distant points for each query, enlarging the effective receptive field at negligible extra computational cost.The strategy enables direct long-range dependencies while preserving efficient window-based processing.
- Stratified Transformer: First-layer point embedding aggregates local information to address irregular point arrangements and support faster convergence and stronger performance.The method also uses contextual relative position encoding to generate positional bias dynamically from semantic features.
- Stratified Transformer: A memory-efficient implementation reduces unnecessary memory consumption caused by varying point numbers across 3D windows.This design targets windows containing relatively few points.
- Experiments: Experiments report state-of-the-art results on S3DIS, ScanNetv2, and ShapeNetPart, with ablations verifying the benefit of each component.The method is presented as a strong Transformer tailored for 3D point-cloud segmentation.
2. Related Work
Related work includes Transformer methods for images and voxel- or point-based methods for 3D segmentation. Stratified Transformer differs from prior point-based approaches by targeting long-range contextual dependencies rather than only local aggregation.
- Vision Transformer: Vision Transformers represent image patches as tokens and use hierarchical, window-based, or spatial-reduction designs for efficient image understanding.These methods include ViT, PVT, and Swin Transformer, alongside other approaches for incorporating long-range or global dependencies.
- Point Cloud Segmentation: Voxel-based segmentation divides 3D space into regular voxels and applies sparse convolutions, but voxelization can lose accurate position information.Point-based methods instead directly adopt point features for segmentation.
- Point Cloud Segmentation: Stratified Transformer is point-based but differs from Point Transformer by addressing limited effective receptive fields and modeling long-range contextual dependencies.Its focus is broader contextual interaction rather than merely local feature aggregation.
3. Our Method
Stratified Transformer uses hierarchical point-cloud encoding with local embedding, windowed self-attention, stratified key sampling, contextual relative position encoding, and memory-efficient computation for segmentation.
- Overview: The encoder-decoder architecture uses point embedding, downsampling, Transformer blocks, and progressive upsampling for segmentation.Both xyz coordinates and rgb colors are used as inputs, with hierarchical stages producing multilevel features.
- Transformer Block: Window-based self-attention reduces global attention’s O(N^2) memory cost by restricting each query to points within its window.Different windows may contain varying numbers of points, and the resulting complexity is O(N × k), where k is the average points per window.
- Transformer Block: Stratified key sampling combines dense nearby keys with sparse distant keys, enlarging the effective receptive field at negligible extra computation.Sparse distant keys account for about 10% of the final key set, while shifted windows further support cross-window communication.
- First-layer Point Embedding: First-layer point embedding aggregates local neighbors to address irregular point arrangements and improve convergence and performance.KPConv performs best among tested local aggregation methods and adds only 2% FLOPs compared with the whole network.
- Contextual Relative Position Encoding: Contextual relative position encoding generates positional bias from relative coordinates and query-key features, producing more varied key-specific biases than MLP encoding.Relative coordinates are quantized into discrete parts and mapped through learnable lookup tables before contextual interaction with features.
4. Memory-efficient Implementation
The implementation avoids padding-based waste from variable point counts across 3D windows by computing attention over indexed query-key pairs and fused scatter operations. This reduces memory use substantially compared with the vanilla implementation.
- Variable token counts make padding each window to kmax waste memory and computations.Masked self-attention with dummy tokens is identified as a simple but inefficient solution.
- Indexed query-key pairs enable attention computation without padding windows to a common size.The method indexes q and k, computes attention, applies scatter softmax by query, and aggregates indexed values.
- 57% memory savings are achieved compared with the vanilla implementation.The fused CUDA-kernel implementation reaches O(M·Nh) memory complexity.
5. Experiments
Stratified Transformer achieves state-of-the-art segmentation performance across S3DIS, ScanNetv2, and ShapeNetPart, with consistent gains over recent methods. Ablations and robustness tests support the contributions of stratified sampling, point embedding, contextual position encoding, and shifted windows.
- 5.2. Results: Stratified Transformer outperforms recent methods on S3DIS and ScanNetv2, including Point Transformer by 1.6% mIoU on S3DIS and by 2.1% mIoU on ScanNetv2 validation.On S3DIS, it also exceeds MinkowskiNet by 6.6% mIoU; on ScanNetv2 test, it is slightly higher than MinkowskiNet.
- 5.2. Results: On ShapeNetPart, the method achieves new state-of-the-art results for both category mIoU and instance mIoU, with a large category-mIoU margin over Point Transformer.Its instance mIoU is comparable to Point Transformer.
- 5.3. Ablation Study: First-layer point embedding yields gains of 2.0% mIoU on S3DIS and 4.0% mIoU on ScanNetv2.Local aggregation methods outperform linear projection for first-layer point embedding.
- 5.3. Ablation Study: Contextual relative position encoding improves performance by 2.9% mIoU on S3DIS and 1.9% mIoU on ScanNetv2, with the best result when applied to query, key, and value features.Its benefit remains when combined with the stratified Transformer, producing gains of 4.0% and 2.3% on S3DIS and ScanNetv2.
- 5.3. Ablation Study: Shifted windows improve the framework, while the model remains effective without them and reaches 70.1% mIoU in that setting.The method is also robust to permutation, rotation, shift, scale, and jitter perturbations, improving by 0.63% mIoU under 90° z-axis rotation.
6. Conclusion
The paper concludes that Stratified Transformer achieves state-of-the-art results while enlarging the effective receptive field through stratified sampling. It also argues that standard Transformers can provide direct long-range dependencies, strong generalization, and high performance for 3D point cloud segmentation.
- 6. Conclusion: Stratified Transformer enlarges the effective receptive field and combines long-range dependencies with first-layer point embedding and contextual relative position encoding.The conclusion identifies low-cost long-range modeling and strong Transformer performance on 3D point clouds as central outcomes.