CAT-Seg: Cost Aggregation for Open-Vocabulary Semantic Segmentation

项目主页：https://ku-cvlab.github.io/CAT-Seg/
代码：https://github.com/KU-CVLAB/CAT-Seg
论文：https://arxiv.org/abs/2303.11797

这份工作是针对于open-vocabulary semantic segmentation的，整体的结构非常的简单和直接，但是获得了非常好的效果。

开放词汇语义分割旨在使用根据文本描述给出的类别来标记图像中的每个像素，其中可能包括在训练期间没有见过的类。
本文的训练与测试设定与OVSeg一致，是基于跨数据集的评估方式，并使用mIoU作为评估指标：

训练中，在COCO-Stuff-171上训练。

测试中，在ADE20K-150/847、PASCAL-Context-59/459、PASCAL VOC-20/20b（定义了背景类）。

为了解决这个任务，最近的工作通常尝试利用大规模的预训练视觉语言模型，例如 CLIP 和 ALIGN。因为它们具有显着的开放性词汇识别能力。这是通过广泛的训练过程实现的，涉及从数百万到数十亿的大型图像描述对。
然而，由于图像和像素之间的粒度差异，将这些经过图像级监督训练的模型的知识转移到像素级预测任务，并在推理中处理任意训练时不可见的类别，非常具有挑战性。
为了应对这一挑战，最近的方法主要有两种范式：

基于生成的类别无关的区域提案，将这些区域用作 CLIP编码器的输入。尽管获得了不错的性能，但这些提案不仅表现出对训练数据集的偏见，还由 CLIP 独立处理并忽略了图像中的全局上下文。典型工作有：
- OVSeg: Open-Vocabulary Semantic Segmentation with Mask-adapted CLIP
- ZegFormer: Decoupling Zero-Shot Semantic Segmentation
- ZSSeg: A Simple Baseline for Open-Vocabulary Semantic Segmentation with Pre-trained Vision-language Model
MaskCLIP: Extract Free Dense Labels from CLIP 中直接修改 CLIP 编码器中的最后一个池化层以获得密集图像嵌入。这种方法会引入大量无法细化的噪声，而可能导致错误预测。通常用的微调编码器来获得任务特定特征的策略可以缓解此问题，同时也已表明使用传统的微调方法会损害 CLIP 的能力。

此外，现有的方法往往没有考虑到开放词汇设置带来的额外挑战。与传统的分割设置不同，其中预定义类别之间的关系信息可以用于推理期间更好地理解场景，开放词汇设置则在推理时假设有任意类，防止模型在没有明确考虑的情况下学习这种关系信息。
为此，鉴于CLIP的编码器已经在大规模数据上进行了训练，从而具有显著的开放词汇能力，本文的目标是：

有效地转移其知识用于像素级预测任务，并通过联合聚合两种模态，将图像中的对象与以文本描述方式给定的类别联系起来。
对CLIP图像编码器进行微调，使模型能够适应像素级预测任务。然而，由于前面提到的直接优化会降低开放词汇能力的问题，探索一种在不损害其预训练知识的情况下优化CLIP图像编码器的替代方案是很重要的。

在回顾 CLIP 训练过程时发现，在图像嵌入和文本嵌入之间构造的cost volume通常被定义为余弦相似图，用于损失计算。由此，基于保留cost volume的使用可作为替代方法的洞察力，本文提出一个图像文本匹配成本来明确表示它们的关系，并引入一个基于成本聚合的框架来处理cost volume。通过聚合，利用类别和视觉语义之间的关系，将图像中的对象与给定类别关联起来，同时使得 CLIP 的表征可以适应于像素级任务。
作者们也观察到直接优化 CLIP 嵌入会损害其开放词汇能力，为此提出了cost aggregation-based method，从而直接优化图像文本相似性图。在图二中证明了这种方法可以有效地转移 CLIP 的知识，而不会削弱其开放词汇能力。

模型细节

整个模型包括三部分：

使用CLIP的B,D,H,W图像嵌入和B,T,P,D文本嵌入计算成对相似性（消除D维度计算余弦相似度），从而得到cost colume，即B,P,T,H,W大小的Tensor，并利用特定的卷积层将这里的P变换为C，从而得到B,C,T,H,W。
1. 这里值得说明的是文本嵌入的形状。本文中，T个类别文本中的每一个都使用了P=80个特定的prompt模板。而B维度则直接是对原始T,P,C大小的Tensor的重复，因为对于每个样本来说，用于对比的文本信息实际上是一样的。
  1. 在cost computation中如上所述，通过卷积进行了融合变换
  2. 而对于上图中用于后续引导的EL，则直接平均归一化后消除这一维度。
2. 在整个后续的处理过程中，维度T（对应类别数目）直接与B合并，前面的维度D被当做特征的通道数，始终保持并直到最后缩小为1。
采用基于 Transformer 的模块，将聚合过程分解为空间聚合和类别聚合两个阶段交替执行数次。并借助额外的称为嵌入引导技术，利用额外的图像编码器的最深层嵌入和原始的文本嵌入作为辅助引导信息，有效地聚合cost volume。
1. 空间聚合使用Swin Transformer，来基于图像和文本嵌入计算的像素级相似性推理空间关系。**这里的K和Q是对cost表征与来自独立图像编码器的深层特征沿通道的拼接，而V仍然是前者。**由于cost表征形状调整为B,T,H,W,C，所以图像特征对应重复T次后得到相同形状。并且无偏移版本后会接一个有偏移的版本。
2. 类别聚合使用线性注意力来聚合cost表征和来自CLIP文本编码器的类别文本表征。基于不使用位置嵌入的注意力的形式既放松了对类别数量的约束，也保证了类别之间的无序性。这里使用的注意力形式与"Transformers are RNNs"中的形式一一致，即将QK的计算拆开了。这里的Q和K同样是两个特征的拼接，V仅仅是前者。
  1. 这里为了拼接二者，对形状为B,T,C的文本特征通过重复得到了形状为B,H,W,T,C，从而与cost特征对应，即每个空间位置附加相同的文本特征。之后Attention计算的是两个**(BHW),T,C**的特征的维度T之间的成对关系。
  2. 同时为了进一步降低复杂度，这里的cost和文本特征在运算之前通过池化缩小了，计算过Transformer后，再对HW进行线性插值放大。
使用基于转置卷积的简单的卷积解码器来处理引导聚合后的cost表征，同时结合额外的图像编码器的浅层图像特征，补充细节信息。最终得到特征为B,T,H’,W’。这一过程沿用了Cost Aggregation with 4D Convolutional Swin Transformer for Few-Shot Segmentation中的设计。

这里要注意的是：

CLIP的文本编码器是固定的，图像编码器仅训练注意力模块和位置嵌入参数。 作者认为微调这些层就足以迁移图像表征到像素级别了。而实际的性能表现也展现出了这一策略的有效性。
作为对比的基准模型，Feature Aggregation与锁体模型的差异主要在于深层嵌入的聚合方式，Feature Aggregation直接拼接，而非通过cost computation。
训练时，对数据采用crop策略，而推理时，使用了基于Patch的重叠划窗进行推理以适应高分辨率输入图像，而无需使用放缩模型中的位置嵌入等策略。具体使用了128x128的步长和大小384x384的窗口。结果可见表9。