基于 4D 卷积 Swin Transformer 的小样本分割代价聚合

发表于 更新于 分类于 阅读次数： 本文字数： 6k 阅读时长 ≈ 20 分钟

原文：《Cost Aggregation with 4D Convolutional Swin Transformer for Few-Shot Segmentation》

摘要

本文提出了一种新颖的代价聚合网络，称为体积聚合 Transformer（VAT），用于小样本分割。Transformer 的使用可以通过对全局感受野的自注意力来促进相关图的聚合。然而，用于 Transformer 处理的相关图的标记化可能是有害的，因为标记边界处的不连续性减少了标记边缘附近可用的局部上下文并减少了归纳偏差。为了解决这个问题，我们提出了一种 4D 卷积 Swin Transformer，其中高维 Swin Transformer 之前是一系列小核卷积，这些小核卷积将局部上下文传递给所有像素并引入卷积归纳偏差。我们还通过在金字塔结构中应用 Transformer 来提高聚合性能，其中较粗级别的聚合引导更精细级别的聚合。然后，利用查询外观嵌入，在后续的解码器中过滤变压器输出中的噪声。该模型在小样本分割的所有标准基准测试中实现了新的最先进水平。研究表明，VAT 在语义对应方面取得了最先进的性能，其中代价聚合也起着核心作用。代码和训练模型可在 https://seokju-cho.github.io/VAT/ 获取。

简介

语义分割是一项基本的计算机视觉任务，旨在用其相应的类别来标记图像中的每个像素。在深度神经网络和包含真实地物分割标注的大规模数据集的帮助下，这个方向已经取得了实质性进展 [37,46,3,4,61]。然而，手动标记按像素级划分的图需要大量的工作，因此很难添加新的类别。为了减少对标记数据的依赖，人们越来越关注小样本分割 [49, 55]，其中仅使用少量支持图像及其相关的掩膜来预测查询图像的分割。 小样本分割的关键是有效利用少量支持样本。许多工作尝试通过从样本中提取原型模型并将其用于与查询进行特征比较 [58,10,35,78]。然而，此类方法忽略了支持和查询特征之间的像素级成对关系或特征的空间结构，这可能导致次优结果。 为了解释这种关系，我们观察到小样本分割可以重新表述为语义对应，其目的是在语义相似的图像之间找到像素级对应，这些图像可能包含大量的类内外观和几何变化 [13,14,43]。最近的语义对应模型 [50,25,51,53,42,44,34,65,41] 遵循特征提取、代价聚合和流量估计的经典匹配流程 [54,47]。在代价聚合阶段，对匹配分数进行细化以产生更可靠的对应估计，这一阶段尤为重要，也是许多研究的重点 [53,42,52,22,34,29,41,6]。最近，CATs [6] 提出使用 Vision Transformer [11] 进行代价聚合，但其对输入标记数量的二次复杂度限制了其适用性。它还忽略了匹配成本的空间结构，这可能会损害其性能。 在小样本分割领域，也存在尝试通过交叉注意力 [83] 或图注意力 [81,68,75] 利用成对信息来优化特征的方法。然而，他们仅依赖于原始相关图，而不聚合匹配分数。因此，他们的对应关系可能会因重复模式或背景混乱而产生歧义 [50,25,27,65,17]。为了解决这个问题，HSNet [40] 使用 4D 卷积聚合匹配分数，但其有限的感受野限制了远程上下文的聚合，并且由于使用固定内核，它缺乏适应输入内容的能力。 在本文中，我们介绍了一种新颖的代价聚合网络，体积聚合 Transformer（VAT），它通过提出的 4D 卷积 Swin Transformer 来解决小样本分割任务。具体来说，我们首先扩展 Swin Transformer [36] 及其补丁嵌入模块来处理高维相关图。通过引入 4D 卷积，进一步扩展了补丁嵌入模块，缓解了补丁嵌入带来的问题，即补丁边界附近有限的局部上下文和低归纳偏差。高维补丁嵌入模块被设计为一系列重叠的小核卷积，为每个像素带来局部上下文信息，并赋予卷积归纳偏差。为了进一步提高性能，我们采用金字塔结构来构建我们的架构，该结构将较粗级别的聚合相关图作为更精细级别的附加输入，从而提供分层指导。然后，我们的亲和力感知解码器以利用查询的外观嵌入给出的更高分辨率空间结构的方式对聚合匹配分数进行细化，并最终输出分割掩膜预测。 我们在几个基准测试中证明了我们的方法的有效性 [55,31,30]。我们的工作在小样本分割甚至语义对应的所有基准上都达到了最先进的性能，突出了代价聚合对于这两项任务的重要性，并显示了其在一般匹配方面的潜力。我们还进行了消融研究来证明我们的设计选择的合理性。

相关工作

小样本分割

受到小样本学习范式 [49, 58] 的启发，小样本学习范式学习仅使用有限数量的样本来学习新任务的模型，小镜头分割受到了相当多的关注。继 [55] 的成功之后，原型网络 [58] 和许多其他工作 [10,45,56,69,35,78,33,76,79,60,84,28] 提出从支持样本中提取原型 ，用于识别查询中的前景特征。此外，受 [82] 的启发，[82] 观察到在特征处理中简单地添加高级特征会导致性能下降，[63] 建议改为利用高级特征来计算先验图像，有助于识别查询图像中的目标。许多变体 [60, 80] 扩展了利用先验图作为聚合特征图的附加信息的想法。 然而，由于基于原型或先验图的方法具有明显的局限性，例如，忽略支持和查询特征之间的成对关系或特征图的空间结构，许多最近的工作[81,68,40,75,32]利用相关图来利用源特征和查询特征之间的成对关系。具体来说，[81,68,75] 使用图注意力，HSNet [40] 提出 4D 卷积来利用多级特征，[32] 将任务表述为最优传输问题。然而，这些方法没有提供聚合匹配分数的方法，仅利用卷积进行代价聚合，或者使用既不可学习也不对严重变形有鲁棒性的手工方法。 最近，[83] 利用 Transformer 并提出在不考虑匹配分数聚合的情况下，使用循环一致的注意力机制来细化特征图，使其变得更具辨别力。[60] 提出了一个全局和局部增强模块，分别使用 Transformer 和卷积来细化特征。[39] 通过冻结编码器和解码器，只关注基于 Transformer 的分类器。与这些工作不同的是，我们提出了 4D 卷积 Swin Transformer，以实现增强且高效的代价聚合。

语义对应

语义对应的目标是找到语义相似图像之间的对应关系，同时应对大的类内外观和几何变化带来的额外挑战 [34,6,41]。这与小样本分割设置高度相似，因为小样本分割也旨在标记具有较大类内变化的同一类的对象，因此最近在这两个任务上的工作都采取了类似的方法。语义对应中的最新方法 [53,42,52,22,34,29,41,6 ] 重点关注代价聚合阶段以找到可靠的对应关系，并证明了其重要性。其中，[41] 提出使用 4D 卷积进行代价聚合，但由于卷积的感受野有限且缺乏适应性而表现出明显的局限性。CAT [6] 解决了这个问题，并通过利用 Transformer [66] 来聚合代价量，从而建立了新的最先进模型。然而，它忽略了相关图的空间结构，并产生较少的归纳偏差，即平移等变性，这限制了其泛化能力 [36,7,8]。此外，当它用于单独聚合相关图时，其二次复杂度可能会限制其适用性。 在本文中，我们提出的方法解决上述问题。

Vision Transformer

最近，自然语言处理（NLP）中的标准架构 Transformer [66] 已在计算机视觉领域得到广泛采用。自从 ViT [11] 的开创性工作以来，许多工作 [39,83,60,23,73,6,36] 都采用 Transformer 来代替 CNN 或以混合方式与 CNN 一起使用。然而，由于 Transformer 序列长度的二次复杂度，常常面临巨大的计算负担。高效 Transformer [70,24,77,72] 旨在通过近似或简化的自注意力来减少计算负载。Swin Transformer [36] 是我们扩展的网络，它通过在预定义的本地窗口内执行自注意力来减少计算。然而，这些工作继承了补丁嵌入引起的问题，我们通过合并 4D 卷积来缓解这些问题。

本文方法

问题阐述

小样本分割的目标是在查询图像中，仅使用少量带有注释的示例，对来自未见类别的对象进行分割[67]。为了减轻由训练数据不足引起的过度拟合，我们采用了 training [67]。让我们分别将训练集和测试集表示为 和 ，其中两个集的对象类别不重叠。在 K-shot 设置下，由两个集合组成多个 ，每个 由一个支持集 组成，其中 是第 个支持图像及其对应的掩膜对，查询样本 ，其中 是查询图像， 是其配对的掩膜。在训练期间，我们的模型从 中获取采样的 ，并学习从 和 到预测 的映射。推断时，我们的模型根据 中随机采样的 和 来预测 。

动机和概述

小样本分割的关键是如何有效地利用为查询图像提供的支持样本。传统方法 [63,60,83,79,28] 利用从支持特征中提取的全局或部分级原型，而最近的方法 [81,68,40,75,32,83] 则利用查询图像和支持图像之间的成对匹配关系。然而，由于类内变化、背景混乱和重复模式，探索这种关系非常具有挑战性。最先进的方法之一是 HSNet [40]，它通过 4D 卷积聚合匹配分数。然而，由于感受野有限或卷积核缺乏适应性，单独使用卷积可能会限制性能。虽然还没有方法可以在小样本分割中使用 Transformer 聚合匹配分数，但 CATs [6] 提出了在语义对应中使用 Transformer 进行代价聚合，证明了 Transformer 作为代价聚合器的有效性。另一方面， Transformer 相对于标记数量的二次复杂度可能会限制其在分割中的实用性。缺乏赋予归纳偏差（例如平移等变性）的操作也可能限制其性能。此外，CATs [6] 定义相关图的标记方式忽略了空间结构，这可能会有害。 提出的代价聚合 Transformer (VAT) 旨在克服这些问题。下面，我们首先描述其特征提取和代价计算。然后，我们提出了 Swin Transformer [36] 用于代价聚合的通用扩展。随后，我们提出了 4D 卷积 Swin Transformer 来解决上述问题。最后，我们介绍了几种额外的技术，包括引导金字塔处理（GPP）和亲和感知 Transformer 解码器（ATD），以进一步提高性能，并将它们结合起来完成设计。

特征提取和代价计算

我们从查询和支持图像中提取特征，并按照传统过程计算它们之间的初始代价 [50,59,53,52,65,17,6]。给定查询和支持图像 和 ，我们使用 CNN [16, 57] 生成 个特征图序列，，其中 和 分别表示第 层的查询和支持特征图。我们得到了一个掩膜支持特征 $ _sl=F_sl^l(m_s)$，其中 表示哈达玛积（元素级别的乘法）， 表示一个函数，该函数调整给定张量的大小，然后在第 层沿着通道维度进行扩展。 给定一对特征图 和 ，我们使用 归一化后特征之间的内积来计算相关图，如下所示： 基于 4D 卷积 Swin Transformer 的小样本分割代价聚合 其中 和 表示特征图的 2D 空间位置。正如 [40] 中所做的那样，我们收集从相同空间大小的所有中间特征计算出的相关图，并将它们堆叠起来以获得堆叠相关图 ，其中 和 分别是查询和支持特征图的高度和宽度， 是金字塔层 中的 CNN 层索引 的子集，包含相同的空间大小的相关图。

金字塔 Transformer 编码器

在这一部分，我们提出了 4D 卷积 Swin Transformer 来聚合相关图，然后将其合并到金字塔架构中。 使用 Transformer 进行代价聚合。然而，当使用 Transformer 进行代价聚合时，如何为与图像、文本或特征形状不同的相关性图定义标记是一个非常棘手的问题 [66, 11]。处理相关图的第一个尝试是 CAT [6]，它将 4D 相关图重塑为 2D 图，并在 2D 中执行自注意力。这忽略了相关图的空间结构，即支持和查询的空间结构，这可能会限制其性能。为了解决这个问题，我们可以将所有空间条目（例如 ）视为标记，并将 视为标记的特征维度。然而，这会导致大量的计算负担，并且随着相关图的增大而增加。因此不能使用标准 Transformer [66, 11]，推荐使用 [70,24,77,72,36] 中的高效版本。然而，使用简化（或近似）的自注意力对于性能来说可能产生不利影响，这将在 4.4 节中讨论。此外，正如在光流和语义对应的文献中所证明的那样 [59, 52]，相邻像素往往具有相似的对应关系。为了保留相关图的空间结构，我们选择使用 Swin Transformer [36]，因为它不仅提供有效的自注意力计算，而且还保持相关图的平滑特性的同时仍然提供足够的远程自注意力。 为了使用 Swin Transformer [36] 进行代价聚合，我们需要对其进行扩展以处理更高维度的输入，特别是 4D 相关图。我们首先遵循传统的补丁嵌入过程 [11] 来嵌入相关图，因为由于大量的标记，它们无法被 Transformer 处理。然而，我们将补丁嵌入模块扩展为体积嵌入模块（VEM），它可以处理更高维度的输入，使得 。按照类似于补丁嵌入的过程，我们使用大卷积核（例如 ）将相关图重塑为一系列展平的 4D 窗口。然后，我们通过将 的查询和支持空间维度均匀划分为不重叠的子相关性图 来扩展自注意力计算，如图3所示。

我们在每个划分的子相关性图内计算自注意力。随后，我们将窗口沿着之前划分的窗口移动 像素（ 为取整符号），然后在新创建的窗口内执行自注意力。然后，正如原始 Swin Transformer [36] 中所做的那样，我们只需将相关图回滚到其原始形式。在计算自注意力时，我们使用相对位置偏差，并从扩展的参数化偏差矩阵中获取值，如下 [19,20,36]。我们保留 Swin Transformer 块的其他组件不变，例如层归一化 (LN) [1] 和 MLP 层。我们将此扩展称为体积 Transformer 模块 (VTM)。总结一下，整体流程定义为： 4D 卷积 Swin Transformer 虽然提出的使用 Transformer 进行代价聚合的方法可以解决使用 CNN 和使用标准 Transformer 的高计算负担等问题，但它可能无法避免其他变压器存在的问题 [11, 70, 24, 77, 72]：缺乏平移等变性。这主要是由于在自注意力计算之前利用了非重叠操作造成的。虽然 Swin Transformer 通过使用相对位置偏置在一定程度上缓解了这个问题 [36]，但它提供的近似还不够充分。我们认为体积嵌入模块是需要解决的问题，因为它会导致几个问题。首先，使用大的非重叠卷积核仅提供有限的归纳偏差。与重叠操作相比，非重叠操作只能实现相对较低的平移等变性。这种有限的归纳偏差导致相对较低的泛化能力和性能 [74,8,7,36]。此外，我们认为，对于密集的预测任务，由于非重叠的内核而忽略窗口边界会因不连续性而损害整体性能。 为了解决上述问题，我们将体积嵌入模块（VEM）替换为一个由一系列重叠卷积组成的模块，我们称之为体积卷积模块（Volumetric Convolution Module，VCM)。具体来说，我们通过应用 4D 空间最大池化、重叠的 4D 卷积、ReLU 和群归一化（Group Normalization，GN)依次减少支持和查询的空间维度，在每个 4D 位置投影多级相似性向量，即将大小为 的向量投影到一个任意固定的维度，记为 。将感受野视为 4D 窗口，即 ，我们从 得到张量 ，其中 和 是处理后的尺寸。需要注意的是，可以为支持和查询空间维度选择不同大小的 。VCM 的概述如图 4 所示。 总的来说，我们定义这样一个过程如下： 这样一来，我们的模型从额外的归纳偏差中受益，同时更好地处理了窗口边界。 此外，为了稳定学习过程，我们提出了一种额外的技术来强制网络来估计残差匹配分数作为补充细节。考虑到在初始阶段，当输入 时，由于 Transformer 参数的随机初始化会导致推断出了错误的匹配分数，这可能会使学习过程复杂化，因为网络需要从随机的匹配分数中学习完整的匹配细节，我们添加了残差连接以加快学习过程 [16, 6, 85]。 引导金字塔处理 我们遵循 [40, 60] 的方法，通过金字塔处理采用自上而下的方法，如图2所示。 受到最近众多关于语义匹配和小样本分割的研究 [83, 41, 6, 40] 的启发，这些研究表明利用多级特征可以极大提高性能，我们也采用了金字塔式的架构。 在我们的自上而下方法中，即我们所称的引导式金字塔处理（Guided Pyramidal Processing，GPP），更精细级别的相关性图 的聚合由上一级（更粗糙）级别 的聚合相关性图引导。具体来说，一个聚合的相关性图 被上采样为图 up，该映射被添加到下一个级别的相关性图 中，作为引导。重复此过程，直到计算出最精细级别的聚合图并将其传递给解码器。如表 4 所示，GPP 带来了可观的性能提升。 使用 GPP，金字塔 Transformer 编码器最终定义为： 其中 up表示双线性上采样。

亲和感知 Transformer 解码器

给定金字塔 Transformer 编码器生成的聚合相关图，基于 Transformer 的解码器生成最终的分割掩膜。为了提高性能，我们建议借助从查询特征图中获得的外观嵌入在解码器内进行进一步聚合。查询的外观嵌入可以在两个方面提供帮助。首先，外观亲和力信息是过滤匹配分数中的噪声的有效指导，正如立体匹配文献中所证明的那样，例如代价体积过滤（Cost Volume filtering，CVF）[18, 59]。此外，可以利用外观嵌入提供的更高分辨率的空间结构来提高上采样质量，从而产生一个保留了细节高度精确的预测掩模 。 对于我们的亲和感知 Transformer 解码器（ATD）的设计，我们对 在支持图像维度上取平均，将其与来自查询特征图的外观嵌入连接起来，然后通过Transformer [66, 70, 72, 36] 进行聚合，随后进行双线性插值。该过程定义如下： 其中 是通过对 在支持图像的空间维度上进行平均池化得到的， 是一个线性投影，， 表示连接操作。我们在双线性上采样后立即对输出进行逐步的细化，以重新捕获细节并整合外观信息。

扩展至 K-Shot 设置

给定 对支持图像和掩膜 和查询图像 ，我们的模型前向传递 次以获得 个不同的查询掩码 。我们将每个空间位置处的所有 个预测相加，如果总和除以 超过阈值 ，则该位置被预测为前景，否则为背景。

实验

实施细节

我们使用在 ImageNet [9] 上预训练的 ResNet50 和 ResNet101 [16]，在训练期间保持权重不变，这是按照 [40, 82] 的方法进行的。我们没有使用数据增强进行训练，详细内容可以参考补充材料。我们将输入图像的尺寸设置为 417 或 473，参考了 [28, 2] 的设置。Swin Transformer 的窗口大小设置为 4。我们使用 AdamW [38] 优化器，学习率设置为 5e−4。用于代价计算的特征图设置 conv3_x ( = 3 )，conv4_x ( = 4 ) 和 conv5_x ( = 5 )。K-shot 阈值 设置为 0.5，嵌入维度 设置为 128。对于外观亲和力，当在 FSS-1000 [30] 训练时，我们选择从 conv2_x，conv3_x 和 conv4_x 获取最后一层，在 PASCAL-5i [55] 和 COCO-20i [31] 训练时，我们排除了 conv4_x。对于 conv2_x，conv3_x 和 conv4_x，我们将 设置为 16、32 和 64。