AAAI 2024 | 高效端到端的3D指向性分割，1080Ti即可训练！代码已开源

本篇分享论文 3D-STMN: Dependency-Driven Superpoint-Text Matching Network for End-to-End 3D Referring Expression Segmentation，通过超点-文本匹配和依存驱动，实现端到端的高效 3D 指向性分割。

论文和代码地址

论文题目：

3D-STMN: Dependency-Driven Superpoint-Text Matching Network for End-to-End 3D Referring Expression Segmentation

论文地址：

https://arxiv.org/abs/2308.16632

代码地址：

https://github.com/sosppxo/3D-STMN

动机

3D Visual Grouding 的目标是根据给定的自然语言描述在 3D 场景中定位实例 [1, 2, 3]。近年来，由于其广泛的应用场景，包括自主机器人、人机交互和自动驾驶系统等，已成为学术研究中的热门话题。该领域中，3D 引用表达式分割（3D-RES）任务是一个艰巨的挑战。与仅使用边界框定位目标对象的三维视觉检测任务相比，3D-RES 需要更复杂的理解。它不仅需要在稀疏点云中识别目标实例，而且还需要提供与每个识别的目标实例相对应的精确 3D 掩码。

现有的二阶段方法 [4] 如图 1-（a）所示，在初始阶段，训练一个独立的文本无关的分割模型来生成实例建议，在第二阶段，使用图神经网络构建生成的提议和文本描述之间的联系。这种两阶段范式存在三个主要问题：

1. 分割与匹配解耦会产生过度依赖初步的与文本无关的分割结果。第一阶段的任何不准确或遗漏都可能损害后续匹配阶段的准确性。
2. 该模型忽略了参考句子中固有的层次和依存结构。它的线性语言建模策略在捕获复杂的语义细微差别方面表现不佳，导致定位和分割中的错误步骤。
3. 为了放大第二阶段的召回率效率，第一阶段通过低效的迭代聚类提取密集的候选掩码。这个迭代过程大大减慢了模型的推理速度。

一种自然的方法是采用端到端方法，将文本特征与三维点云中的点直接匹配，如图 1-（b）所示。这种方法已被证明在 2D-RES 任务中非常有效 [5]。然而，它的低召回率导致较难直接迁移到稀疏、不规则的 3D 点云数据。

作为一种解决方案，3D VG（REC）中的 3D-SPS 提出了一种方法，该方法逐步选择语言引导的关键点，并使用该关键点信息回归框。然而，这种方法破坏了 3D 掩码在 3D-RES 任务中的连续性，从而降低了分割质量。

▲ 图1（a）传统两阶段范式，（b）点级端到端范式，（c）超点级端到端范式。

为了解决上述挑战，该工作提出了一种用于端到端 3D-RES 的依存驱动超点文本匹配网络（3D-STMN）。该工作的方法的想法是将指向性表达式与过分割的超点 [6] 进行匹配。如图 1-（c）所示，这些超点最初通过聚类算法聚合，从而获得细粒度的语义单元。与三维点云中的无序点相比，这些超点体现了语义，并且明显更少，在匹配过程中提供了性能和速度的优势。

与 TGNN 中的提议相比，超点是从过度分割派生的细粒度单元，能够覆盖整个场景，从而避免了分割不准确或缺失实例的问题。鉴于此，该工作引入了一种新的 3D-RES 超点文本匹配（STM）机制，利用超点文本特征的聚合来获取目标实例的掩码。

为了从文本角度支持语义解析，该工作设计了一个依存驱动交互（DDI）模块，实现了令牌级交互。该模块利用来自依存语法树的先验信息来引导文本信息的流动。这种结构通过网络架构进一步增强了对不同实例之间关系的推理，从而显着提高了模型的分割能力。

3D-STMN 对 3D-RES 领域做出了如下贡献：

• 论文提出了一种新的基于超点-文本匹配（STM）机制的高效端到端框架 3D-STMN，用于将超点与文本模态对齐，使超点成为多模态表示中极具竞争力的选项。
• 该工作设计了一个依存驱动交互（DDI）模块来利用来自依存语法树的先验信息来引导文本信息的流动，显着提高了模型的分割能力。
• 大量实验表明，该工作的方法在 ScanRefer 基准测试中显着优于之前的两阶段基线， mIoU 提高了 11.7 个百分点，同时也实现了推理速度的惊人增强。

方法

3D-STMN 包含超点-文本匹配（STM）机制和依存驱动交互（DDI）模块。STM 利用密集注释的超点-文本对，而不是更稀疏的实例-文本对，直接将语言指示与它们对应的超点进行关联，从而有效地利用跨模态语义关系。DDI 模块利用依存语法树引导文本信息流，辨别表达式中的重要单词及其相关描述符之间错综复杂的关系，以提高模型的分割能力。图 2 展示了 3D-STMN 的详细构成。

▲ 图2 3D-STMN框架图

3.1 超点-文本匹配（STM）

视觉端，输入的点云 首先通过一个 Sparse 3D U-Net 获取点级特征 ，与此同时利用 [6] 的无参数方法从 中提取 个超点 [6]，接着通过超点池化层提取超点级特征：

其中， 表示第 个超点的特征， 表示第 个超点所包含的点的集合， 表示超点平均池化操作。

文本端，该工作将 BERT 提取的单词级特征送入 DDI 模块，以构建依存树并输出依存驱动的特征：

为了进行超点-文本匹配，该工作首先将超点特征 投影到与依存驱动特征 对应的 维子空间中得到 ，然后对 采用一种过滤方法从而提高后续处理的效率。具体来说就是基于超点与其对应描述之间的相关性得分 ，获取了 个超点。线性投影与过滤过程可以表示为：

其中， 是可学习的参数， 起到全局特征的作用， 表示连接操作， 表示与描述相关的超点特征。

为了提高文本分割核的判别能力，该工作引入了一个超点-单词聚合模块（Superpoint-Word Aggregation，SWA），该模块旨在细化超点和文本描述之间的多轮模态交互。在第 层，SWA 自适应地聚合超点特征，使每个单词能够吸收相关超点特征的视觉信息。自适应的超点-单词交叉注意力利用依存驱动特征 通过融合相关超点的信息来精炼词特征：

其中， 是超点-单词交叉注意力的输出， 为可学习的参数， 是超点注意力掩码。给定来自预测头的预测超点掩码 ，超点注意力掩码通过阈值 过滤超点，如下所示：

其中， 表示第 i 个单词对第 j 个超点的注意力。根据经验，该工作将 设置为 0.5。通过堆叠 SWA 层，超点注意力掩码 自适应地限制了目标实例之间的交叉注意力。

基于 SWA 模块的超点和文本描述之间的多轮模态交互过程，可以描述如下：

其中，， 是多轮交互的轮次。

接着该工作对 和 进行矩阵乘法，以获得捕捉所有超点和单词之间关系的响应图。这个计算过程可以描述如下：

其中， 表示 的转置， 是响应图， 表示 sigmoid 函数。特别地， 是第 个单词的响应图，基于该响应图该工作可以生成对应于第 个单词的分割结果和注意力掩码 。

为了获得最终的掩码，该工作选择与具有最高相关得分的描述相关超点对应的单词的响应图 ：

其中， 返回对应于最大值的索引。 是可学习的参数， 表示第 个单词和第 个与描述相关的超点之间的注意力得分， 表示第 个词的视觉相关性得分。

3.2 依存驱动的交互（DDI）

对于一个由 个句子和总共 个单词组成的目标对象描述文本，该工作首先使用 Stanford CoreNLP 工具包 [8] 获取 个依存树。然后，该工作通过合并这些 个依存树的 ROOT 节点，将它们合并成一个图。因此，对于每个描述，依存图具有 个节点 和 条边 。每个节点代表一个单词，包括特殊标记 “ROOT”，而每条边代表一种依存关系类型。

然后，该工作采用具有边特征的图 Transformer 层，以更有效地利用描述依存图中丰富的特征信息，这些信息以包括依存关系在内的边属性的形式存储。针对单词级特征 ，根据如下式子可以得到文本特征 ：

其中， 是一个可学习的参数， 表示随机初始化的 ROOT 节点特征，表示连接操作。根据该文本特征中的元素对应的索引直接推导出节点特征 。对于边特征 ，该工作为每个依存关系分配一个唯一的 ID，通过线性投影传递，并获得 维的隐藏特征 。

其中， 和 是线性投影层的参数。

接着针对与计算的 维节点位置编码提取 维特征，然后加到节点特征 中。

其中 ，。注意，拉普拉斯位置编码仅在输入层将节点特征中添加，而不是在中间的图 Transformer 层中添加。

接下来，该工作将继续定义第 层的更新方程：

其中， 表示可学习的参数。

考虑到依存图结构中缺乏远程连接，该工作引入了自注意机制，并将其与图注意力并行使用。输出 与 的自注意力输出相加，并通过残差连接和归一化层进行处理，得到输出 。 和 分别传递给独立的前馈网络，并通过残差连接和归一化层进行处理，如下所示：

其中，, , , 是可学习的参数，, , 是中间表示， 表示 的第 个自注意力输出。

最后，第 层 DDI 输出的依存驱动特征 通过 的拼接获得。

3.3 训练损失

训练一个超点-引用表达式匹配网络很简单：给定真实二进制掩码 ，该工作首先通过超点池化得到相应的超点掩码 ，然后在最终的响应图 上应用二元交叉熵（BCE）损失函数。该操作可以表示为：

其中， 为叠加点平均池化操作， 为第 个超点 的二值化掩码值。 表示掩码值是否大于 50%。

尽管二元交叉熵（BCE）损失函数独立处理每个超点，但在解决前景-背景样本不平衡的问题上存在不足。为了解决这个问题，该工作使用 Dice 损失函数：

在 STM 模块中，该工作参考 [3] 应用 对描述相关性得分 使用交叉熵损失进行监督。对 的监督是基于某点是否属于描述中提到的对象类别。

此外，该工作参考 [7] 为 proposal 质量预测添加了一个简单的辅助得分损失 ：

其中， 是分割核对 proposal 的质量得分预测。 表示 proposal 掩码预测与分配的真实掩码之间的 IoU 是否高于 50%。

最终的训练损失函数 可以表示为：

其中 ， ， 和 是用于平衡这四个 loss 的超参数。该工作将其设置为 ，，。

实验

该工作是在 3D 参考数据集 ScanRefer 上进行训练并测试的，该数据集包含 51,583 个自然语言表达式，涉及 800 个 ScanNet 场景中的 11046 个对象。所有实验都是用 PyTorch 实现的，并在单个 NVIDIA Tesla A100 GPU 上进行训练。

该工作采用平均交并比（mIoU）和交并比超过的准确率（Acc@ IoU）作为该工作的评估指标。交并比超过 的准确率表示预测掩码与真实标注之间的重叠部分的 IoU 大于 的描述文本的比例，其中 。

▲ 表1 ScanRefer数据集上的3D-RES结果

此前唯一研究 3D-RES 任务的工作是 TGNN。如表 1 所示，该工作提出的 3D-STMN 在 Acc@0.25、Acc@0.5 和 mIoU 方面取得了显著的改进，整体性能提升分别达到了 17.1%、8.4% 和 11.7%。

在推断速度方面，该工作计算了每个描述的平均推断时间。该工作的 3D-STMN 表现出明显优势，比两阶段的 TGNN 快了 95.7 倍。由于该工作模型的推断时间控制在 0.3 秒以内，使得 3D-RES 的实时应用成为可能。无论是使用 BERT 还是 GRU 特征，该工作的 3D-STMN 都明显优于 TGNN，突出了该工作模型的鲁棒性和推断能力。在 “Unique” 设置中，该工作的模型将 Acc@0.25 提升了 30 个点，突显了其对于唯一对象的精确性。

4.1 消融

该工作首先对于 STM 机制进行了消融实验。首先针对是否使用超点级特征，从表 2 中可以看出，在相同的设置下，第二行在所有指标上都显着优于第一行，证明了使用超点作为表示的有效性。

接下来，在第 3-6 行，该工作添加了 DDI 模块。无论 DDI 模块的结构如何，它都会极大地增强分割核的性能，导致所有指标的显着改进，证明了依存驱动特征的细粒度分辨能力。在 STM 框架中，分割内核策略的选择在依存驱动特征如何构建内核进行分割方面起着关键作用。

该工作测试了三种不同的策略：i）Root：它使用根节点的嵌入来制定分割内核；ii）Top1：利用得分最高的词嵌入，这是通过沿超点维度平均词嵌入得出的；iii）Average：利用通过平均所有单词的嵌入计算的嵌入。该工作的发现，如表 2 所示，Top1 策略因其灵活的视觉自适应性而成为最有效的。因此，后续实验选择了此设置。

▲ 表2 STM 机制的消融，其中 “w/o DDI” 表示直接使用 [CLS] Token 去生成分割核，而非采用 DDI 模块。

在表 3 中，该工作探索了 DDI 模块结构的四个不同的版本：i）GA（仅图注意力层），ii）SA-GA（自注意力层后串联图注意力层），iii）GA-SA（图注意力层后串联自注意力层），iv）GA SA（图注意力层和自注意力层并联）。

该工作的研究结果表明，与 DDI 模块的缺失相比，GA 配置带来了性能的显着提高。这强调了模型中详细的依存驱动交互的关键作用。在将 SA 的密集连接附加给 GA（SA-GA 和 GA-SA）后，由于添加了远程连接，模型的 “Overall” 性能有所提高，证明了 SA 在增强 GA 结构有效性方面的互补作用。

最后，通过结合并行自注意力（GA SA），在所有设置中都实现了性能的显着提高。这突出了在同时补充远程连接的同时利用并行连接的有效性，保留了依存树的显式建模能力并保持信息的有序交互。

▲ 表3 DDI 模块的消融

4.2 定性分析

该工作对 ScanRefer 验证集进行了定性比较，展示了 3D-STMN 相对于 TGNN 在辨别能力上的显著优势。图 3 直观地展示了无论测试样本的难度级别如何，3D-STMN 在准确定位目标对象的注意力图上都有着优越的性能。3D-STMN 生成的注意力高度聚焦，展现了卓越的精确性。

相反，TGNN 在辨别方面存在困难，在多个语义相似对象上显示出显著高的注意力值，如图 3 的（a）、（b）和（c）所示。值得注意的是，当面临包含多个与目标相似的对象，并伴随着更长、更复杂的文本描述（如图 3 的（a）和（c）所示）的场景时，TGNN 无法区分和准确定位目标，其性能可与随机猜测相似。

相反，3D-STMN 能够对这些具有挑战性的样本进行精确的分割。类似于人类，它对靠近目标的对象有微妙但明显的关注，并将它们与背景区分开来，如图 2 的（c）中所示。

▲ 图3 3D-STMN 和 TGNN 的预测结果和注意图的可视化

总结

这项研究提出了 3D-STMN，这是一种高效且密集对齐的 3D-RES 端到端方法。

通过使用超点文本匹配（STM）机制，该模型成功地摆脱了传统两阶段范式的局限性。使得模型可以利用端到端密集监督，具备精确分割和快速推理速度的优点。特别地，该模型每个场景实现了不到 1 秒令人印象深刻的推理速度，这使得它非常适合实时应用，并且高度适用于时间敏感的场景。此外，所提出的依存驱动交互（DDI）模块显着提高了模型对引用表达式的理解。通过显式建模依赖关系，该模型展示了改进的定位和分割能力，性能显着提高。

得益于精细、高效的设计，该模型甚至可以在 1080Ti 上进行训练！

更多细节请参考原文！

参考文献

[1] CHEN D Z, CHANG A X, NIESSNER M. Scanrefer: 3d object localization in rgb-d scans using natural language[C] // European conference on computer vision. 2020 : 202–221.

[2] ACHLIOPTAS P, ABDELREHEEM A, XIA F, et al. Referit3d: Neural listeners for fine-grained 3d object identification in real-world scenes[C] // Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part I 16. 2020 : 422–440.

[3] LUO J, FU J, KONG X, et al. 3d-sps: Single-stage 3d visual grounding via referred point progressive selection[C] // Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022 : 16454–16463.

[4] HUANG P-H, LEE H-H, CHEN H-T, et al. Text- guided graph neural networks for referring 3d instance segmentation[C] // Proceedings of the AAAI Conference on Artificial Intelligence : Vol 35. 2021 : 1610–1618.

[5] YANG Z, WANG J, TANG Y, et al. Lavt: Language-aware vision transformer for referring image segmentation[C] // Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022 : 18155–18165.

[6] LANDRIEU L, SIMONOVSKY M. Large-scale point cloud semantic segmentation with superpoint graphs[C] // Proceedings of the IEEE conference on computer vision and pattern recognition. 2018 : 4558–4567.

[7] SUN J, QING C, TAN J, et al. Superpoint transformer for 3d scene instance segmentation[C] // Proceedings of the AAAI Conference on Artificial Intelligence : Vol 37. 2023 : 2393–2401.

[8] MANNING C D, SURDEANU M, BAUER J, et al. The Stanford CoreNLP natural language processing toolkit[C] // Proceedings of 52nd annual meeting of the association for computational linguistics: system demonstrations. 2014 : 55–60.