NeurIPS 2022 | 沈春华团队：搬来「ATM」打造语义分割的新范式！

论文链接：

https://arxiv.org/abs/2210.05844

本文对朴素ViT在语义分割中的应用进行了探索，提出了一种新的语义分割范式SegViT。不同于已有基于ViT语义分割方案的像素级表达方式，本文采用ViT最基本的成分（注意力机制）直接生成掩码(Mask)。具体来说，本文提出ATM(Attention-to-Mask)模块，它将所学习类别token与特征之间的相似性直接转换为分割Mask。实验结果表明：在ADE20K数据集上，所提SegViT(55.2%mIoU)取得了明显优于其他ViT方案的性能，在COCO-Stuff-10K(50.3%%mIoU)与PASCAL-Context(65.3%mIoU)数据集上取得了新的SOTA指标。

此外，为更进一步降低ViT骨干的计算复杂度，本文还提出了QD(query-based 下采样)与QU(query-based 上采样)用于构建Shrunk结构，搭配上该结构所得SegViT可以在保持相当性能(55.1%mIoU)同时节省40%的计算量(373.5GFLOPs vs 637.9GFLOPs)。

一、本文出发点

注意力机制中的Q与K的点乘操作用于计算两者之间的相似性，而在语义分割中同类区域特征应与其对应类别具有更大的相似性，上图则很好的说明了特征与类别之间的相似性。在Attention的基础上，简单的添加Sigmoid操作即可将相似性图转换为Mask。受此启发，作者设计了一种用于语义分割的ATM模块，并将其与朴素ViT的组合称之为SegViT。

众所周知，ViT骨干具有极高的计算复杂度，为进一步降低计算复杂度，作者提出了QD与QU以构建Shrunk结构，其中QD用于降低特征分辨率，而QU则用于辅助骨干进行分辨率重建。但搭配上Shrunk结构后，所提SegViT能够在保持相当性能的同时节省40%的计算量。

二、本文方案

上图给出了本文所提方案架构示意图，它由基于ViT骨干的Encoder与基于ATM的Decoder构成。接下来，我们对Encoder部分进行简要介绍，重点针对Decoder部分进行介绍。

Encoder

Decoder

Mask为跨注意力的中间输出，ATM模块的最后输出将被用于分类。作者对其添加一个线性变换+Softmax激活以输出得到类别概率预测。在推理阶段，类别概率与Mask的点乘输出作为最终的语义输出。

由于朴素ViT并不具备多阶段多尺度特征，FPN这种多尺度特征聚合方式就不再适用。那么，如何进行底层与高层特征聚合呢？作者基于ATM将ViT骨干不同位置的特征进行聚合并作为SegViT的Decoder模块。基于此，作者进一步引入了QD与QU构建了Shrunk版本的SegViT，在大幅降低计算量的同时保持同等性能。

Shrunk Decoder

众所周知，同等性能前提下，朴素ViT具有更高的计算量复杂。为降低所提方案的计算复杂度，作者提出了一种基于QD与QU的Shrunk结构，可参考上图。由于注意力模块的输出尺寸受Q的尺寸决定，最简单的方式就是对齐进行下采样(如最近邻下采样，见上图b)，然后送入后续Transformer层处理，但这种方式会导致信息损失，进而导致性能的退化。

为补充朴素Shrunk处理导致的信息损失，作者进一步添加了两个并行QU模块进行高分辨率特征重建。通过该方式的加持，SegViT-Shrunk可以大幅降低计算量(约40%)，同时保持相当的性能。

Loss

在损失函数方面，作者对类别预测概率部分采用分类损失进行监督，对Mask部分采用Mask Loss(它有Focal Loss与Dice Loss组成)进行监督，定义如下：

三、本文实验

为验证所提方案的有效性，作者在ADE20K、COCO-Stuff-10K以及PASCAL-Context三个数据集上进行了验证。需要注意的是，ViT预训练模型并非官方预训练，而是参考Segmenter与Structtoken中选用了更强的预训练模型。这是因为：两个版本的预训练模型在最终模型上的性能差距竟然可以达到2.9%mIoU(51.7% vs 54.6%)。