【图像分割轻量化】 A DECODING SCHEME WITH SUCCESSIVE AGGREGATION OF MULTI-LEVEL FEATURES (SASFormer)

A DECODING SCHEME WITH SUCCESSIVE AGGREGATION OF MULTI-LEVEL FEATURES FOR LIGHT-WEIGHT SEMANTIC SEGMENTATION (SASFormer)

一种基于多级特征逐步聚合的轻量级语义分割解码方案——层级视觉应用于轻量级语义分割

论文链接：http://arxiv.org/abs/2402.11201

代码链接：[暂无]

1、摘要

  文中提出了一种针对多尺度架构编码器的新型语义分割解码方案。该解码方案基于多级视觉Transformer，旨在通过引入多级特征的逐步跨注意力聚合，实现计算成本的降低和分割准确性的提高。此外，文中还提出了一种增强多级特征的方法，通过聚合语义进行提升。文章重点在于从注意力分配的角度保持上下文一致性，从而在显著降低计算成本的同时带来性能的提升。在多个流行数据集上的实验表明，在计算成本不变的情况下，所提出的方案优于最先进的语义分割模型，并且广泛的消融研究证明了所提思路的有效性。

2、创新点

  如图2所示，所提出的方案是SASFormer——基于层次视觉Transformer的多级特征连续聚合语义分割。图中的聚合语义提取器（ASE）在SASFormer中起关键作用：

• （1）通过连续的跨注意力操作提取保持注意力分配上下文一致性的聚合语义，以提高分割精度；

• （2）通过主动利用不同层级的下采样特征图，显著降低计算成本。同时，语义组合模块（SCM）被设计用来利用聚合语义作为增强多尺度特征的权重，从而提升精度。

  文中方案的主要贡献总结如下：

• 基于层次视觉Transformer的语义分割新型解码方案——SASFormer。

• 通过聚合语义提取器（ASE）和语义组合模块（SCM），实现多尺度特征的连续聚合，以提高分割精度和降低计算成本。

• 通过在ADE20K和Cityscapes数据集上的对比实验，证明了与最先进的模型（SOTA）在精度和计算成本之间取得良好的平衡，以及引入概念的有效性，通过广泛的消融研究得到验证。

3、背景

  直接将Transformer应用于对计算需求较高的语义分割任务并不合适，因为它生成的是固定分辨率的特征图。为解决这些问题，基于HVT（Hierarchical Vision Transformer）的模型应运而生。这些模型不仅生成多尺度特征，还提出有效的注意力机制，以降低计算成本，如图1。但HVT模型的一个显著局限性在于，它们未能充分利用多级特征的上下文信息，因为它们仅在每个阶段对输入特征进行自注意力操作。

4、原理

Overall Architecture

  所提出的轻量级语义分割模型如图2所示。该模型的编码器处理$H\times W\times 3$尺寸的输入图像。四阶段编码器生成多尺度特征$F_i$的尺寸为$\frac{H}{2^{(i+1)}}\times \frac{W}{2^{(i+1)}}\times C_i$，其中$i\in 1,...,4$表示阶段索引。在SASFormer中，每个多尺度特征在传递给ASE之前都会被调整为固定尺寸。ASE从调整后的特征中提取聚合语义。随后，通过SCM（尺度组合模块），聚合的语义与相应的多尺度特征相结合，然后通过分割头生成最终的分割map。

Accumulated Semantics Extractor (ASE)

  累积语义提取器（ASE）的设计旨在优化计算负担与多级特征高效利用之间的权衡。直接使用多尺度特征可以提高分割精度，但会增加计算复杂度，因为特征维度较大。为了减轻高计算负担，文中从四阶段解码器中获取的多尺度特征$F_i(i\in 1,...,4)$被降采样为单尺度特征$R_i$，其大小固定为$\frac{H}{64}\times \frac{W}{64}\times C_i$，其中$C_i$表示$i$级的通道数。文中设计了连续交叉注意力（SCA），不仅补偿降采样过程中信息的丢失，还能保持不同层级特征之间的上下文一致性。

  ASE由四个Transformer块组成，每个块包含基于MHA的连续交叉注意力（SCA）和前馈网络（FFN），如图2所示。MHA层使用线性投影生成的Q（查询）、K（键）、V（值）进行处理，其计算过程如下：
$$MHA(Q,K,V)=Softmax(\frac{Q{K}^T}{{\sqrt{d}}_k})V,(1)$$
其中$d_k$表示键K的嵌入维度。MHA层和FFN按照图2所示连接，操作描述为（2）和（3）：
$$A_j^l=MHA(LN(K{V}_j^l),LN(Q_j^l))+Q_j^l,(2)$$
$$S_j^l=FFN(LN(A_j^l))+A_j^l,(3)$$
其中$l\in 1,...,4$表示Transformer块的索引，$j\in 1,2,3$表示作为查询的特征的索引，LN表示层归一化。对于第一个Transformer块（$l=1$），$K{V}_j^l$和$Q_j^l$ 表示为（4）：
$$Q_j=R_{j+1},K{V}_j=\{\begin{array}{ll}{R}_j,& \text{if j = 1}\\ S_j,& \text{otherwise}\end{array},(4)$$

对于剩余的Transformer块，$K{V}_j^l$ 和$Q_j^l$表示为（5）：
$$Q_j^l=S_{j+1}^l,K{V}_j^l=\{\begin{array}{ll}{R}_j,& \text{if j = 1}\\ S_j^l,& \text{otherwise}\end{array}.(5)$$

  在FFN中，文中通过在两个$1\times 1$卷积层之间集成深度卷积层DWconv来增强局部连接，扩张通道的比例设置为4。需要注意的是，由（2）至（5）表示的过程是连续交叉注意力（SCA），与自注意力和常规交叉注意力机制在查询、键和值设置方式上有所不同。因此，文中称$S_2^4,S_3^4,S_4^4$为聚合语义，它们是通过四个Transformer块提取的。在不同层级上的连续过程有效地聚合了富含上下文信息的语义。

Semantic Combining Module (SCM)

  SCM（Semantic Combining Module）用于细化多尺度特征的上下文信息。如图2右侧所示，对于$j\in 2,3,4$，SCM将从ASE（Aggregated Semantic Embedding）中提取的聚合语义$S_j^4$与相应的多尺度特征$F_j$相结合，以提升表示能力。在这个过程中，聚合语义作为权重，与除最低级特征$F_1$以外的对应多尺度特征进行乘法运算。

为了将聚合语义作为权重处理，它们需要通过与上一节中描述的下采样比例进行上采样，以匹配对应多尺度特征的维度。同时，$F_j$和$S_j^4$通过$1\times 1$卷积层，随后进行批量归一化，如图所示。获取增强特征$O_j$的过程可以描述为（6）：

$$Q_j=F_j\times S_j+F_j，wherej=2,3,4.（6）$$

  在（6）中，将$F_j$添加到乘法结果中是为了弥补在将$F_j$下采样到$R_j$过程中产生的信息损失。增强特征与$F_1$一起通过分割头，用于预测分割图。

5、实验

Datasets

  在广泛的实验和消融研究中，文中选择了两个流行的数据集：ADE20K [19] 和 Cityscapes [20]。ADE20K 数据集包含150个类别，总共有25,000张图像，分为20,000张用于训练，2,000张用于验证，以及3,000张用于测试。Cityscapes 数据集包含19个类别，总共5,000张图像，分为2,975张用于训练，500张用于验证，以及1,525张用于测试。

Implementation Details

  在MMsegmentation[21]框架下，文中基于两块RTX 3090 GPU进行了性能评估和消融研究。选择预训练在ImageNet-1K数据集上的MiT[8]作为编码器。所有数据集的训练都在160,000个迭代周期内完成，使用AdamW优化器。文中没有采用BatchNorm层，而是采用了同步批归一化，这样在训练过程中可以跨多个GPU聚合BatchNorm的均值和标准差。ADE20K数据集的批次大小设置为16，Cityscapes数据集的批次大小为8。初始学习率设为0.0001，采用poly学习率策略，衰减因子为1.0。采用了SegFormer中使用的数据增强方法。报告的性能指标包括单尺度下的mean IoU（mIoU）和计算量（FLOPs）。所有这些实验和性能分析都是在单尺度推理情况下进行的。

Comparisons with the State-of-the-Art

  所提方法在ADE20K和Cityscapes数据集上的性能与当前最先进的模型进行比较。如表1所示，在轻量级配置下，所提模型在ADE20K上实现了40.1%的mIoU，4.8M参数和5.0GFLOPs的计算量，而在Cityscapes上则达到了78.3%的mIoU，99.3GFLOPs的计算量。与SegFormer-B0相比，SASFormer-B0在ADE20K上的mIoU提高了2.7%，计算量减少了40.4%，而在Cityscapes上，mIoU提高了1.9%，计算量减少了20.9%。此外，与Lawin Transformer-B0相比，SASFormer-B0在ADE20K上的mIoU提升了1.2%，但计算量减少了5.6%。在Cityscapes上，SASFormer-B0的mIoU提高了1.2%，而计算成本基本保持不变。

  在中等配置下，SASFormer-B2在ADE20K上达到了48.0%的mIoU，30.7M参数和36.6GFLOPs的计算量，而在Cityscapes上则为81.7%的mIoU和512.7GFLOPs。与SegFormer-B2相比，SASFormer-B2在ADE20K上的计算量减少了41.3%，mIoU提高了1.5%，而在Cityscapes上，计算量减少了28.5%，mIoU提高了0.9%。与Lawin-Transformer-B2相比，SASFormer-B2在ADE20K上的计算量减少了18.7%，mIoU提高了0.2%。在Cityscapes上，SASFormer-B2的mIoU与Lawin Transformer-B2相同，但计算量减少了8.9%。总的来说，精心设计的解码器使得在大致相同的参数数量下，所提模型在mIoU上表现更高，计算量更少。

Ablation study

  在消融研究部分，以MiT-B0作为编码器，并使用ADE20K的数据集进行训练和评估。评估了应用于解码器的proposed方案，所有评估都在与基线模型相同的条件下进行。文中验证了连续交叉注意力的有效性，通过对比两种不同的注意力方法：一种是图3(a)中使用了对齐后的R i 特征的自注意力，然后将自注意力提取的语义分割并仅保留最低级的输入给SCM；另一种是图3(b)中的交叉注意力，没有连续连接。如表2所示，连续交叉注意力（SCA）的表现优于两者，证明了通过连续配置的交叉注意力从注意力分配的角度保持上下文一致性，有效地聚合了上下文信息。

  SCM中的融合方式：文中通过实验（6）中的融合方式（7）和（8）来评估融合的有效性：

• $O_j=F_j\times S_j，wherej=2,3,4.$
• $O_j=F_j\times S_j+S_j,wherej=2,3,4.$

  如表3所示，（6）中的融合策略在保持FLOPs和参数数量不变的情况下，提高了mIoU。这证明了（6）中多尺度特征的残差连接在减少下采样过程中损失的负面影响方面是有效的。

  Transformer块的数量。表4中的性能指标来自于不同数量Transformer块的实现。考虑到参数增加相对较小，我们设置4个块作为FLOPs和mIoU之间最优平衡的选项。

  SASFormer作为解码器的适用性。在这类工作中，一个想法的适用性是关键。为了证明提出的解码器在语义分割中的适用性，文中进行了一系列实验，将SASFormer替换现有的基于HVT的语义分割模型的解码器。如表5所示，该解码方案不仅降低了计算成本，还提高了精度，证明了其在处理多级特征方面的有效性。值得注意的是，当SASFormer与HVT编码器结合时，可以有效地处理这类模型固有的问题。图4展示了三个基于HVT的模型与解码器替换为SASFormer的模型的定性结果对比，图中的框标记了多尺度对象的位置，使用SASFormer的模型在复杂区域显示出精确分割的能力。

6、总结

  文中提出了一种简单而强大的语义分割解码器，称为SASFormer。这种基于轻量级Transformer的解码器旨在通过从注意力分配的角度考虑如何保持上下文一致性，以及逐级聚合基于HVT（层次视觉Transformer）编码器的多尺度特征，优化计算成本与分割精度之间的权衡。通过一系列实验和详尽的消融研究，证明了所提方案的有效性。作者认为，该方案适用于需要在计算成本与性能之间取得平衡的其他计算机视觉任务。