论文阅读笔记：AllSpark: Reborn Labeled Features from Unlabeled in Transformer for Semi-Supervised Semantic

论文：https://arxiv.org/pdf/2403.01818v3.pdf

代码：https://github.com/xmed-lab/AllSpark

1 背景

目前主流的半监督语义分割主要分为两类：

（1）连续的。最具代表性的是基于教师-学生的框架，首先用带标签的数据训练学生模型，然后用学生模型的EMA更新教师模型，在用教师模型生成伪标签来监督学生模型。

（2）并行的。以交叉伪监督（CPS）为代表，在一个小批量内将有标签数据与无标签数据串联，模型通过一致性正则化同时优化两个数据流。

这些先进方法的一个关键特征是对有标签和无标签数据的训练流进行分离，如图1（a）所示。训练流的分离运行有标签的数据占训练过程中占主导地位，因为他们有标签，如图1（c）。然而这些方法都有一个共同的显著缺点，即用于伪无标签的图像生成伪标签的模型受到有限的有标签样本集的严重影响，从而导致低质量的伪标签，进而导致次优的结果。

为此作者提出了一种新的方法：与直接在标注样本上训练不同，作者利用未标注样本生成标注特征，如图1（b）所示。其潜在的原理是标注数据远远少于未标注数据，通过从未标注数据中重生标注特征，作者将多样性引入到标注流中，创造了更具挑战性的学习环境（等于就是根据语义相似度，用未标注数据的特征来代替标注数据的特征，并传入解码器中，用标注的标签来监督重生的特征）。

需要确定一种有效的方法对标注数据进行重生。在Transformer decoder中使用的交叉注意力机制提供了一种合适的方法，因为它被设计用于使用源序列重构目标序列。因此作者基于交叉注意力的概念开发了本文的方法，称为AllSpark，如图2。使其能够以通道集的方式操作操作，以提取通道级的语义信息。使用已标注的数据特征作为query，未标注的数据特征作为key和value，如图2所示。具体来说，AllSpark计算每个通道的标注特征与未标注特征查询之间的相似度矩阵。然后强调相似最高的未标注通道，以重构标注特征。然而实现这一想法的一个关键挑战是小批量内有限的未标注特征，不足以重建正确的标注特征。为了解决这个问题，作者又引入了一个语义内存S-Mem，它从先前未标注的特征中收集通道特征。S-Mem使用一种提出的基于通道的语义分组策略进行迭代更新，该策略首先将代表相同语义的通道分组，然后将它们排入S-Mem的相应类别槽中。

由于卷积操作的接受域较小，不适合将AllSpark插入到基于CNN的主干中。

2 创新点

本文提出了AllSpark模块来解决标记数据占优问题，并构建了一个基于SOTA纯Transformer的SSSS框架，例如SAM，这可以进一步有利于使用较少的标记数据来有效地训练基础模型。

3 方法

本文提出的方法是一个可以插在通用分割模型灵活组件。在训练阶段，将未标注的特征替换为语义内存中的特征，概率图被用于通道语义分组策略。在推理阶段，将输入作为测试图像的隐藏特征，将交叉注意力退化为自注意力机制。

4 模块

4.1 基线与概况

一般的半监督语义分割数据集包含有限的有标签数据和大量的无标签数据，整个数据集由 $N_L$ 个标注样本 { x i l , y i } i = 1 N L \{x_i^l,y_i\}_{i=1}^{N_L} {xil​,yi​}i=1NL​​ 和 $N_U$ 个未标注样本 { x i u } i = 1 N U \{x_i^u\}_{i=1}^{N_U} {xiu​}i=1NU​​ 组成，其中 $x_i\in R^{3\times H\times W}$为输入，$y_i\in R^{K\times H\times W}$为K类的真实标注。基于伪标签的方法首先通过真实数据作为监督来训练模型，然后用预训练的模型作为监督来生成伪标签 $\hat{y}$。最基本的目标函数为：

其中 $L_{CE}$ 为交叉熵损失。

本文从伪标签的基本公式出发，提出AllSpark模块，该模块集成在基于ViT的分割模型的编码器和解码器之间，如图3。AllSpark模块包括一个通道式交叉注意力模块，一个用于存储未标记特征以更好地重建已标记特征的语义内存。以及用于更好的更新语义内存的通道式语义分组策略。AllSpark伪代码如下：

4.2 通道交叉注意力

在通道交叉注意力机制中，标记的数据特征作为query，而未标记的特征作为key和value，如图3所示。具体来说，先计算已标注特征的每个通道与未标注特征之间的相似度，相似度最高的未标注通道在重建已标注特征中起最显著的作用，其潜在原理是，尽管未标记的特征相对于已标注的特征可能来自不同的类，但他们仍然具有可以共享的通道通用信息，如纹理。

给定编码阶段后标记和未标记数据的隐藏特征 $[h^l,h^u]$，将其拆分，然后以多头的方式将 $h^l$ 设置为query，$h^u$ 设置为key和value：

其中 $w_q,w_k,w_v\in R^{C\times 2C}$ 是变换权重，$h^l,h^u\in R^{C\times d}$，$d$ 是序列长度（这里应该就是特征图像素个数），$C$ 是通道维度。通道注意力定义为：

其中 $\psi (\cdot )$ 和 $\sigma (\cdot )$ 表示实例归一化和softmax函数。

实例归一化在每个样本的HW维度上进行归一化。

与传统的自注意力机制比，通道注意力可以捕获通道间的长距离依赖。为了进一步细化未标注特征，作者还应用了通道自注意力，获得细化的未标注特征 ${\hat{h}}_i^u$，类似于公式（2）（3）。唯一不同的是将 $q$ 替换成 $q=h^u{w}_q$，重建后的 ${\hat{h}}_i^l$ 和 ${\hat{h}}_i^u$ 随后被送入解码器生成最终的预测。

4.3 类别平衡的语义内存

直接利用单个小批量未标记的特征不足以有效的重建已标注的特征，为了克服这一局限，需要拓展未标注特征的空间。因此，作者引入了一个先进先出队列来存储大量未标记的特征，称为类别级语义内存。S-Mem的形状为 $R^{K\times C\times d}$，其中 $K$ 表示类的个数。对于每一类，S-Mem存储 $C$ 个通道，每个通道有 $d$ 个patch。

直接存储先前未标记的特征并不适用于语义分割任务，因为类别往往是不平衡的，为了保证每个类都有足够的语义信息，需要建立一个类平衡的语义内存。因此，作者引入通道语义分组策略来确定未标记特征中每个通道的语义表示，然后将它们分组并添加到S-Mem的相应类中，如图4。

具体来说，计算未标注特征 $h^u\in R^{C\times d}$ 和概率token $\hat{p}\in R^{K\times d}$ 之间的相似度。概率token是概率图 $p\in R^{K\times H\times W}$ 通过 resize 和 reshape 得到的向量。一旦确定了每个通道的语义表示，就可以根据它们各自的语义类别将它们存储在内存中。相似度矩阵 $Sim\in R^{K\times C}$ 定义为：$Si{m}_{i,j}=\psi ({\hat{p}}_i^T{h}_i^u)$。这样，通过比较第 $i\in [0,C]$ 个通道的语义信息和图像的整体语义上下文，可以确定第 $i$ 个通道所代表的最可能的语义类别，然后将代表相同语义类别的通道组合在一起，将它们分配到语义内存中。

推理阶段，仍然可以利用与S-Mem的交叉注意力机制，但这往往会带来计算负担并只提供了很小的改进，为了提高效率，在推理阶段去除S-Mem和CSG。因此，交叉注意力的两个输入变得相同，可以简化为自注意力。

5 效果

5.1 与SOTA方法对比

5.2 消融实验

各个模块的消融实验如表5。

图5展示了 All Spark模块的可视化验证，证明了重建后的标记特征与表示特征相比，具有更大的多样性， 这有助于环节标记数据的支配性，且与自身特征的自注意力比，更大的特征空间使得特征更接近于目标。

图6验证了通道语义分组策略可以为大多数通道准确的分配准确的语义标签。y轴表示选择并添加到语义内存中的类别的数量。

图7将部分语义内存可视化，总的来说，语义内存存储了与类相关的想当准确的特征。

表6展示了之前具有Transformer骨干的SOTA在切换到Segformer后性能都有下降，因为Transformer的训练需要更多的标注数据以避免过拟合。而AllSpark在Resnet101+DeepLabV3+上也表现出交叉的性能，这归因于CNN有限的感受野。所以本文方法不适合CNN-based模型。

不适合CNN的原因如图8所示，在分析大物体（如汽车）时，注意力图优先考虑的是几个明显的局部部分（如轮子）。因此这些特征经过AllSpark处理后，存在错误的查询对象（例如飞机起落架的轮子）。