[Paper Review] Mask-Attention-Free Transformer for 3D Instance Segmentation

Xin Lai, Yuhui Yuan+4 more|arXiv (Cornell University)|Sep 4, 2023

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes a mask-attention-free Transformer for 3D instance segmentation that replaces mask attention with an auxiliary center regression task to overcome slow convergence caused by low-recall initial masks. By using dense position queries, relative position encoding, and iterative refinement, the method achieves 4× faster convergence and sets a new state of the art on ScanNetv2, outperforming prior methods with fewer training epochs.

ABSTRACT

Recently, transformer-based methods have dominated 3D instance segmentation, where mask attention is commonly involved. Specifically, object queries are guided by the initial instance masks in the first cross-attention, and then iteratively refine themselves in a similar manner. However, we observe that the mask-attention pipeline usually leads to slow convergence due to low-recall initial instance masks. Therefore, we abandon the mask attention design and resort to an auxiliary center regression task instead. Through center regression, we effectively overcome the low-recall issue and perform cross-attention by imposing positional prior. To reach this goal, we develop a series of position-aware designs. First, we learn a spatial distribution of 3D locations as the initial position queries. They spread over the 3D space densely, and thus can easily capture the objects in a scene with a high recall. Moreover, we present relative position encoding for the cross-attention and iterative refinement for more accurate position queries. Experiments show that our approach converges 4x faster than existing work, sets a new state of the art on ScanNetv2 3D instance segmentation benchmark, and also demonstrates superior performance across various datasets. Code and models are available at https://github.com/dvlab-research/Mask-Attention-Free-Transformer.

Motivation & Objective

• To address the slow convergence issue in existing 3D instance segmentation methods based on Transformers.
• To identify the root cause of training instability as low recall in initial instance masks due to mask attention.
• To replace mask attention with a center regression-based guidance mechanism to improve training efficiency and accuracy.
• To design position-aware components such as dense position queries and relative position encoding to support the new training paradigm.
• To achieve state-of-the-art performance with significantly reduced training time.

Proposed method

• Introduces learnable, densely distributed position queries that span the 3D space to ensure high recall of object instances from the start of training.
• Replaces mask attention with a center regression task that uses predicted center coordinates to guide cross-attention, avoiding reliance on low-quality initial masks.
• Employs relative position encoding (RPE) in cross-attention layers to model spatial relationships between queries and points, replacing hard masking.
• Applies iterative refinement of position queries across decoder layers to progressively improve localization accuracy.
• Introduces center-based matching and a center loss during training to align predictions with ground-truth instance centers.
• Uses a two-stream query mechanism: content queries for feature aggregation and position queries for spatial guidance, both refined iteratively.

Experimental results

Research questions

• RQ1 Why do existing Transformer-based 3D instance segmentation methods suffer from slow convergence?
• RQ2 Can the low recall of initial instance masks in mask attention mechanisms be a primary cause of training instability?
• RQ3 Can replacing mask attention with a center regression task improve convergence speed and segmentation accuracy?
• RQ4 How can positional priors be effectively leveraged in a mask-attention-free Transformer for 3D instance segmentation?
• RQ5 What design choices—such as dense position queries, relative position encoding, and iterative refinement—contribute most to performance gains?

Key findings

• The proposed method achieves a new state of the art on the ScanNetv2 3D instance segmentation benchmark with a 63.9% mAP@50 and 73.5% mAP@25.
• The model converges 4× faster than the baseline, outperforming a 512-epoch baseline with only 128 epochs of training.
• Ablation studies confirm that center matching and center loss each contribute over 1.6% mAP improvement, with combined removal causing a 2.0% drop.
• Relative position encoding outperforms absolute and content-conditioned position encodings, proving the importance of modeling spatial relationships.
• Iterative refinement of position queries improves performance by 0.9% mAP, demonstrating its value in refining spatial localization.
• The method generalizes well, showing superior performance across multiple datasets including ScanNetv2, ScanNet200, and S3DIS.