[论文解读] Cost Aggregation with 4D Convolutional Swin Transformer for Few-Shot Segmentation

Sunghwan Hong, Seokju Cho+3 more|arXiv (Cornell University)|Jul 22, 2022

Advanced Neural Network Applications被引用 10

一句话总结

本文提出了一种新型的4D卷积Swin Transformer——体积分量聚合网络（Volumetric Aggregation with Transformers, VAT），用于少样本分割中的代价聚合。通过在Swin Transformer之前引入小感受野卷积以保留局部上下文信息与归纳偏置，并采用外观感知解码的分层金字塔结构，VAT在所有主流少样本分割基准和语义对应任务上均取得了最先进性能。

ABSTRACT

This paper presents a novel cost aggregation network, called Volumetric Aggregation with Transformers (VAT), for few-shot segmentation. The use of transformers can benefit correlation map aggregation through self-attention over a global receptive field. However, the tokenization of a correlation map for transformer processing can be detrimental, because the discontinuity at token boundaries reduces the local context available near the token edges and decreases inductive bias. To address this problem, we propose a 4D Convolutional Swin Transformer, where a high-dimensional Swin Transformer is preceded by a series of small-kernel convolutions that impart local context to all pixels and introduce convolutional inductive bias. We additionally boost aggregation performance by applying transformers within a pyramidal structure, where aggregation at a coarser level guides aggregation at a finer level. Noise in the transformer output is then filtered in the subsequent decoder with the help of the query's appearance embedding. With this model, a new state-of-the-art is set for all the standard benchmarks in few-shot segmentation. It is shown that VAT attains state-of-the-art performance for semantic correspondence as well, where cost aggregation also plays a central role.

研究动机与目标

• 解决标准视觉Transformer在少样本分割代价聚合中因基于图像块的标记化而导致的局部上下文丢失与归纳偏置减弱的问题。
• 通过在Swin Transformer之前引入小感受野4D卷积，以保留空间结构并增强局部特征表示，从而提升代价聚合性能。
• 通过引入分层金字塔架构，使粗粒度层级的聚合结果引导细粒度层级的处理，实现特征精炼。
• 通过引入外观感知解码器，利用查询特定嵌入减少Transformer输出中的噪声，提升分割预测的准确性。
• 通过在少样本分割与语义对应任务中均取得最先进结果，证明所提代价聚合方法具备良好的泛化能力。

提出的方法

• 提出一种4D卷积Swin Transformer，将Swin Transformer的图像块嵌入模块扩展至处理4D相关图（批量、通道、H、W），实现对匹配代价体素的全局自注意力计算。
• 在Transformer之前引入一系列重叠的小感受野4D卷积，以注入局部上下文信息与归纳偏置，缓解图像块边界处的不连续性问题。
• 设计一种分层金字塔架构，将来自粗分辨率层级的聚合代价图作为附加输入送入细分辨率层级，实现多尺度引导。
• 实现一种亲和力感知解码器，利用查询的外观嵌入对最终聚合的代价体素进行精炼，以增强空间细节并减少噪声。
• 采用端到端训练方式，使用标准分割损失（如交叉熵）在支持-查询对上进行优化，特征提取通过主干网络（如ResNet）完成。
• 采用可学习的位置编码，适配4D空间相关图，以在自注意力计算过程中保留空间关系。

实验结果

研究问题

• RQ1 4D卷积Swin Transformer能否在保留局部上下文与归纳偏置的前提下，有效聚合少样本分割中的代价体素？
• RQ2 在多尺度上进行分层聚合是否能提升代价体素精炼在少样本分割中的鲁棒性与准确性？
• RQ3 所提出的代价聚合机制是否能泛化至少样本分割之外的其他匹配任务，如语义对应？
• RQ4 在解码器中引入外观嵌入相较于标准Transformer基线方法，能否显著提升分割质量？
• RQ5 将标准4D卷积替换为混合卷积-Transformer方法，在模糊或杂乱输入下，能否有效降低噪声并提升性能？

主要发现

• VAT在所有标准少样本分割基准上均取得最先进性能，包括PASCAL-5i、COCO-20i与FSS-1000，在COCO-20i（5-shot）上达到54.9的平均mBA，在PASCAL-5i（5-shot）上达到54.8的平均mBA。
• 在语义对应任务的SPair-71k基准上，VAT达到55.5的平均准确率，超越此前最先进方法如CATs（49.9）与PMNC（50.4）。
• 消融实验表明，4D卷积主干显著提升了性能，有效减少了边界伪影并增强了局部上下文，尤其在纹理重复区域表现更优。
• 分层金字塔结构显著促进特征精炼，消融实验显示相比非金字塔变体，mBA提升1.5–2.0%。
• 外观感知解码器有效降低了最终代价体素中的噪声，使分割掩码更清晰、更准确，尤其在背景杂乱的挑战性场景中表现突出。
• 模型在语义对应任务上也展现出良好泛化能力，表明高效的代价聚合是少样本分割与对应任务中的关键瓶颈。