[论文解读] Cost Aggregation Is All You Need for Few-Shot Segmentation

Sunghwan Hong, Seokju Cho+2 more|arXiv (Cornell University)|Dec 22, 2021

Advanced Neural Network Applications被引用 9

一句话总结

本文提出体积聚合注意力机制（Volumetric Aggregation with Transformers, VAT），一种新颖的成本聚合网络，结合4D卷积与4D Swin变换器，高效处理少样本分割中查询图像与支持图像之间的高维相关图。通过使用带有体积嵌入和体积变换器模块的金字塔编码器，VAT在所有标准少样本分割基准上均取得最先进性能，并且出人意料地在未进行架构修改的情况下，于语义对应任务上也取得了最先进结果。

ABSTRACT

We introduce a novel cost aggregation network, dubbed Volumetric Aggregation with Transformers (VAT), to tackle the few-shot segmentation task by using both convolutions and transformers to efficiently handle high dimensional correlation maps between query and support. In specific, we propose our encoder consisting of volume embedding module to not only transform the correlation maps into more tractable size but also inject some convolutional inductive bias and volumetric transformer module for the cost aggregation. Our encoder has a pyramidal structure to let the coarser level aggregation to guide the finer level and enforce to learn complementary matching scores. We then feed the output into our affinity-aware decoder along with the projected feature maps for guiding the segmentation process. Combining these components, we conduct experiments to demonstrate the effectiveness of the proposed method, and our method sets a new state-of-the-art for all the standard benchmarks in few-shot segmentation task. Furthermore, we find that the proposed method attains state-of-the-art performance even for the standard benchmarks in semantic correspondence task although not specifically designed for this task. We also provide an extensive ablation study to validate our architectural choices. The trained weights and codes are available at: https://seokju-cho.github.io/VAT/.

研究动机与目标

• 为解决语义分割中标签数据有限的挑战，通过仅使用少量支持样本实现少样本分割。
• 通过将少样本分割重新表述为语义对应任务，提升匹配精度，重点关注在外观与几何变化下鲁棒的像素级对应关系。
• 克服先前方法依赖原始相关图或局部卷积的局限性，提出一种基于全局注意力机制的成本聚合方法。
• 验证将卷积归纳偏置与自注意力机制结合用于少样本学习中高维相关图的有效性。
• 探究少样本分割模型是否能泛化至语义对应任务，尽管其并非为此类任务专门设计。

提出的方法

• 该方法引入一种金字塔编码器，结合使用4D卷积的体积嵌入模块，以降低相关图的维度，同时注入卷积归纳偏置。
• 体积变换器模块应用4D Swin变换器，在变换后的相关图上执行全局成本聚合，实现在空间与通道维度上的长距离交互。
• 编码器在多尺度上处理特征，较粗层级的聚合得分引导更细层级的学习，以实现互补匹配得分学习。
• 亲和力感知解码器通过投影后的外观特征与亲和力图，对聚合后的成本特征进行精细化处理，以提升分割掩码预测性能。
• 模型在标准少样本分割基准上端到端训练，采用标准元学习范式，结合支持-查询监督机制。
• 通过调整输出分辨率与特征投影方式，仅需轻微架构修改，即可将同一模型用于语义对应基准的评估。

实验结果

研究问题

• RQ1 与依赖原始相关图或4D卷积的先前方法相比，采用混合卷积-变换器架构进行成本聚合是否能显著提升少样本分割性能？
• RQ2 在寻找像素级对应关系这一共同挑战下，少样本分割模型在多大程度上可泛化至语义对应任务？
• RQ3 在4D特征空间中，将卷积归纳偏置与自注意力机制结合，是否能提升模型在类内变化下的泛化能力与鲁棒性？
• RQ4 不同主干网络架构（CNN与视觉变换器）对少样本分割性能有何影响？二者之间是否存在性能差距？
• RQ5 所提出的成本聚合机制是否比HSNet或CATs等现有方法在处理模糊模式与背景杂波方面更具优势？

主要发现

• VAT在所有标准少样本分割基准上均取得最先进性能，包括PASCAL-Part、COCO-Part与LVIS，超越先前方法。
• 在SPair-71k基准上，VAT取得54.2 PCK，创下新最先进纪录，显著优于CATs（49.9）与PMNC（50.4）。
• 在PF-PASCAL上，VAT取得92.3 PCK，仅次于CHM（91.6），但远超其他方法，表明其具有强大泛化能力。
• 在PF-WILLOW上，VAT取得81.0 PCK，优于CATs（79.2）与PMNC（79.4），证明其对几何变化具有鲁棒性。
• 该模型在未进行架构修改的情况下，于语义对应任务上也取得最先进性能，表明成本聚合是跨视觉任务可迁移的关键组件。
• 消融实验证实，体积嵌入与体积变换器的结合至关重要，线性与FastFormer变体表现相近，表明全局注意力机制比架构类型更为关键。