[论文解读] Mask-Attention-Free Transformer for 3D Instance Segmentation

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

Advanced Neural Network Applications被引用 4

一句话总结

本文提出一种无需掩码注意力机制的3D实例分割Transformer，通过引入辅助中心回归任务替代掩码注意力，以解决因初始掩码召回率低导致的收敛缓慢问题。通过使用密集位置查询、相对位置编码和迭代优化，该方法实现4倍于基线的收敛速度，并在ScanNetv2上达到新的SOTA性能，且训练轮次更少。

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.

研究动机与目标

• 为解决基于Transformer的现有3D实例分割方法中存在的收敛缓慢问题。
• 识别训练不稳定的根源为掩码注意力机制导致的初始实例掩码召回率低。
• 通过基于中心回归的引导机制替代掩码注意力，以提升训练效率与精度。
• 设计位置感知组件，如密集位置查询与相对位置编码，以支持新型训练范式。
• 在显著减少训练时间的前提下，实现SOTA性能。

提出的方法

• 引入可学习的、在3D空间中密集分布的位置查询，确保从训练初期即实现对物体实例的高召回率。
• 用基于中心回归的任务替代掩码注意力，利用预测的中心坐标引导交叉注意力，避免依赖低质量的初始掩码。
• 在交叉注意力层中引入相对位置编码（RPE），以建模查询与点之间的空间关系，替代硬性掩码。
• 在解码器各层中对位置查询进行迭代优化，逐步提升定位精度。
• 在训练过程中引入基于中心的匹配机制与中心损失，使预测结果与真实实例中心对齐。
• 采用双流查询机制：内容查询用于特征聚合，位置查询用于空间引导，两者均通过迭代方式优化。

实验结果

研究问题

• RQ1 为何基于Transformer的现有3D实例分割方法存在收敛缓慢的问题？
• RQ2 掩码注意力机制中初始实例掩码召回率低，是否是导致训练不稳定的主因？
• RQ3 能否通过用中心回归任务替代掩码注意力，来提升收敛速度与分割精度？
• RQ4 如何在无掩码注意力机制的3D实例分割Transformer中有效利用位置先验？
• RQ5 哪些设计选择——如密集位置查询、相对位置编码与迭代优化——对性能提升贡献最大？

主要发现

• 所提方法在ScanNetv2 3D实例分割基准上达到新的SOTA性能，mAP@50为63.9%，mAP@25为73.5%。
• 模型收敛速度比基线快4倍，仅用128个训练轮次即超越基线在512个轮次下的性能表现。
• 消融实验表明，中心匹配与中心损失各自带来超过1.6%的mAP提升，两者同时移除导致mAP下降2.0%。
• 相对位置编码优于绝对位置编码与内容相关位置编码，证明建模空间关系的重要性。
• 对位置查询进行迭代优化使性能提升0.9% mAP，证明其在精化空间定位中的价值。
• 该方法泛化能力优异，在多个数据集（包括ScanNetv2、ScanNet200与S3DIS）上均表现出优越性能。