Paper: https://arxiv.org/pdf/2301.04648v1.pdf
Code: https://github.com/dongbo811/AFFormer
导读
本文提出了一种名为Adaptive Frequency Transformer（AFFormer）的语义分割架构。AFFormer采用并行架构来利用原型表示（prototype representations）作为特定可学习的局部描述，其取代了解码器并在高分辨率特征上保留丰富的图像语义。 虽然删除了解码器能够压缩大部分的推理计算，但并行架构的精度仍受到低计算资源的限制。因此，我们采用异构运算符（CNN和Vision Transformer）进行像素嵌入（pixel embedding）和原型表示，以进一步节省计算成本。此外，从空间域的角度线性化Vision Transformer的复杂度非常困难。由于语义分割对频率信息非常敏感，我们构建了一个轻量级原型学习块，其具有复杂度O（n）的自适应频率滤波器，以替换标准自注意力的O（n^2）复杂度。 在广泛使用的数据集上进行的大量实验表明，AFFormer在保持3M参数的同时实现了优异的精度。在ADE20K数据集上，AFFormer达到41.8 mIoU和4.6 GFLOPs，比Segformer高4.4 mIoU，GFLOPs减少45％。在Cityscapes数据集上，AFFormer达到78.7 mIoU和34.4 GFLOPs，比Segformer高2.5 mIoU，GFLOPs减少72.5％。
引言
AFFormer achieves better accuracy on ADE20K and Cityscapes datasets with significantly lower FLOPs
语义分割是将图像划分为子区域（像素集合）的任务，并被定义为像素级分类任务。与图像分类相比，它有两个独特的特点：像素级稠密预测和多类表示，这通常建立在高分辨率特征上，并且需要图像语义的全局归纳能力。以前的语义分割方法主要关注使用分类网络作为骨干来提取多尺度特征，并设计复杂的解码器头来建立多尺度特征之间的关系。然而，这些改进是以巨大参数量和高计算成本为代价的。 例如，使用轻量级MobilenetV2作为骨干的著名PSPNet包含13.7M个参数和52.2 GFLOPs，输入尺寸为512×512。使用相同骨干的DeepLabV3+需要15.4M个参数和25.8 GFLOPs。这种固有的设计方式限制了这一领域的发展，并阻碍了许多现实世界的应用。因此，我们提出了以下问题：语义分割是否像图像分类一样简单？
最近视觉Transformer（ViTs）在语义分割中显示出巨大的潜力，但在超低计算能力设备上部署时，它们面临着平衡性能和内存使用的挑战。标准Transformer在空间域中具有O(n^2)的计算复杂度，其中n是输入分辨率。现有方法通过减少token数量或滑动窗口来缓解这种情况，但它们对计算复杂度的减少是有限的，甚至会妥协分割任务的全局或局部语义。 同时，语义分割作为一个基础研究领域，具有广泛的应用场景，需要处理各种分辨率的图像。 如上图所示，尽管高效Segformer相比PSPNet和DeepLabV3+取得了巨大的突破，但它仍面临着更高分辨率的巨大计算负担。在512×512的尺度下，尽管Segformer与PSPNet和DeepLabV3+相比非常轻巧，但它的花费几乎是我们的两倍（8.4 GFLOPs vs 4.6 GFLOPs）;在2048×2048的尺度下，甚至需要5倍的GFLOPs（384.3 GFLOPs vs 73.2 GFLOPs）。因此，我们提出了另一个问题：我们能否为超低计算场景设计高效灵活的Transformer网络用于语义分割？
这篇论文提出了一种名为Adaptive Frequency Transformer(AFFormer)的无头（Head-Free）轻量级语义分割架构。AFFormer采用了一种并行架构，利用原型表示作为特定的可学习局部描述来取代解码器，并在高分辨率特征上保留丰富的图像语义。此外，我们采用异构运算符来处理像素嵌入特征和局部描述特征，以节省更多的计算成本。基于Transformer的模块称为原型学习（PL）用于学习原型表示，而基于卷积的模块称为像素描述符（PD），将像素嵌入特征和学习的原型表示作为输入，将它们转换回全像素嵌入空间以保留高分辨率语义。
然而，从空间域的角度来线性化ViT的复杂度仍然非常困难。受到频率对分类任务的影响启发，我们发现语义分割对频率信息也非常敏感。因此，我们构建了一个复杂度为O(n)的轻量级自适应频率滤波器，其作为原型学习来代替标准自注意力O(n^2)。该模块的核心由频率相似核，动态低通和高通滤波器组成，分别从强调重要频率分量和动态过滤频率的角度捕获对语义分割有益的频率信息。最后，通过在高频和低频提取和增强模块中共享权重来进一步减少计算成本。我们还在前馈网络（FFN）层中嵌入了一个简化的深度卷积层，以增强融合效果，减小两个矩阵变换的大小。
通过并行异构架构和自适应频率滤波器的帮助，我们仅使用一个卷积层作为分类层（CLS）用于单尺度特征，实现最佳性能，使语义分割和图像分类一样简单。 我们在三个广泛使用的数据集上证明了AFFormer的优势：ADE20K，Cityscapes和COCO-stuff。仅使用3M个参数，AFFormer显着优于最先进的轻量级方法。在ADE20K上，AFFormer以4.6 GFLOPs的成本达到41.8 mIoU，比Segformer高出4.4 mIoU，同时将GFLOPs减少了45％。在Cityscapes上，AFFormer实现了78.7 mIoU和34.4 GFLOPs，比Segformer高出2.5 mIoU，GFLOPs减少了72.5％。大量实验结果表明，我们的模型可以在计算受限的场景中应用，同时在不同数据集上保持高性能和鲁棒性。
方法
Parallel Heterogeneous Architecture
语义解码器将编码器获得的图像语义传播到每个像素，并恢复下采样中丢失的细节。一种直接的替代方法是在高分辨率特征中提取图像语义，但这会引入大量计算，特别是对于视觉Transformer。相反，我们提出了一种新策略，用原型语义描述像素语义信息。对于每个阶段，给定一个特征F ∈ R^ (H×W×C)，我们首先初始化一个网格G ∈ R^(h×w×C)作为图像的原型，其中G中的每个点都作为局部簇中心，初始状态仅包含周围区域的信息。这里我们使用1 × C向量来表示每个点的局部语义信息。对于每个具体的像素，由于周围像素的语义不一致，每个簇中心之间存在重叠语义。簇中心在其相应区域α^2中被加权初始化，并且每个簇中心的初始化表示为：
其中n=α×α，wi表示xi的权重，α设置为3。我们的目的是更新网格G中的每个簇中心s，而不是直接更新特征F。由于h×w <H×W，这大大简化了计算。
在这里，我们使用基于Transformer的模块作为原型学习来更新每个聚类中心，其中总共包含L层，并更新的中心被称为G'(s)。对于每个更新的聚类中心，我们通过像素描述符恢复它。设F'i表示恢复的特征，它不仅包含来自F的丰富像素语义，而且还包含由聚类中心G'(s)收集的原型语义。 由于聚类中心聚合了周围像素的语义，导致局部细节的损失，PD首先用像素语义对F中的局部细节进行建模。具体来说，F被投影到低维空间，建立像素之间的局部关系，使每个局部块保持明显的边界。然后将G'(s)嵌入到F中，通过双线性插值恢复到原始空间特征F'。最后，它们通过线性投影层进行整合。
Prototype Learning by Adaptive Frequency Filter Motivation
The effect of different frequency components on semantic segmentation.
语义分割是一项极其复杂的像素级分类任务，容易出现类别混淆。频率表示可以作为学习类别差异的新范式，可以挖掘人类视觉忽略的信息。如上图所示，除非大部分频率分量被过滤，人类对频率信息移除具有鲁棒性。然而，模型对频率信息移除非常敏感，甚至移除少量信息也会导致显著的性能下降。这表明，对于模型来说，挖掘更多的频率信息可以增强类别之间的差异，使每个类别之间的边界更加清晰，从而提高语义分割的效果。
Structure of the adaptive frequency filter in prototype learning.
由于特征 F 包含丰富的频率特征，因此网格 G 中的每个聚类中心也收集了这些频率信息。为了提取不同的频率特征，之前的工作提出了一种基于 Fourier 变换和反 Fourier 变换的方法。然而，这种方法带来了额外的计算开销，并且不能在许多硬件上使用。因此，论文提出了一种基于视觉 Transformer 的自适应频率滤波块，直接在空间域中捕获重要的高频和低频特征。 其核心组件如上图所示，公式定义如下：
上面的公式定义了一种自适应频率滤波块的操作。其中，D^(fc)h、D^(lf)m(X)和D^(hf)n(X)分别表示具有H组的频率相似性内核、具有M组的动态低通滤波器和具有N组的动态高通滤波器。||·||表示Concatenation。值得注意的是，这些操作采用并行结构来进一步减少计算成本，这是通过共享权重来实现的。
Frequency Similarity Kernel (FSK)
不同的频率分量分布在G中，我们的目的是选择和增强有助于语义解析的重要分量。为此，我们设计了一个频率相似性内核模块。一般来说，这个模块是由视觉Transformer实现的。给定一个X∈R^((hw)×C)的特征，通过卷积层在G上进行相对位置编码。我们首先使用固定大小的相似性内核A∈R^(C/H×C/H)来表示不同频率分量之间的对应关系，并通过查询相似性内核来选择重要的频率分量。我们将其视为一个函数转移，通过线性层计算频率分量的键K和值V，并通过Softmax操作在频率分量之间归一化键。每个分量都集成了一个相似性内核Ai,j，计算如下:
其中ki表示K中第i个频率分量，vj表示V中第j个频率分量。我们还通过线性层将输入X转换为查询Q，并通过对固定大小相似性内核的交互获徖分量增强输出。
Dynamic Low-Pass Filters (DLF)
低频分量占据绝对图像的大部分能量，并表示大部分语义信息。低通滤波器允许低于截止频率的信号通过，而高于截止频率的信号被阻挡。因此，我们采用典型的平均池化作为低通滤波器。然而，不同图像的截止频率是不同的。为此，我们在多组中控制不同的内核和步幅来生成动态低通滤波器。对于第m组，我们有:
其中Λk×k表示具有核大小为k×k的深度卷积层。此外，我们使用查询和高频特征的Hadamard积来抑制物体内部的高频，这些高频是分割的噪声。FFN有助于融合捕获的频率信息，但具有大量计算，在轻量级设计中通常被忽略。在这里，我们通过引入卷积层来减少隐藏层的维数，以弥补由于维度压缩而缺失的能力。
讨论
对于频率相似性内核，计算复杂度为O(hwC^2)。每个动态高通滤波器的计算复杂度为O(hwCk^2)，这比频率相似性内核的复杂度小得多。由于动态低通滤波器是通过每组的自适应平均池实现的，其计算复杂度约为O(hwC)。因此，模块的计算复杂度与分辨率呈线性关系，这对于语义分割的高分辨率是有利的。
实验
Frequency analysis of stage-2 (left) and stage-3 (right).
Frequency analysis of the core components in PL module.
Comparison to state of the art methods on ADE20K with resolution at 512 × 512.
Comparison to state of the art methods on Cityscapes val set.
Speed-accuracy tradeoffs at different scales on Cityscapes.
Comparison to state of the art methods on COCO-stuff.
Ablation studies on frequency aware statistics.
总结
本文提出了一种无头轻量级语义分割架构 AFFormer。核心思想是从频率角度学习聚类原型的局部描述表示，而不是直接学习所有像素嵌入特征。它删除了复杂的解码器，同时具有线性复杂度Transformer，并将语义分割简化为常规分类。各种实验表明，AFFormer具有强大的精度，在低计算成本下具有优越的稳定性和鲁棒性。
