Title: Rethinking Local Perception in Lightweight Vision Transformer
导读
本文主要介绍了一种轻量级Vision Transformer
架构——CloFormer
,用于处理移动端的图像任务。CloFormer
引入了 AttnConv
,这是一种结合了注意力机制和卷积运算的模块,能够捕捉高频的局部信息。相比于传统的卷积操作,AttnConv 使用共享权重和上下文感知权重,能够更好地处理图像中不同位置之间的关系。实验结果表明,CloFormer 在图像分类、目标检测和语义分割任务中具有优越的性能。
动机
众所周知,自从 ViTs 提出后,Transformer 基本横扫各大 CV 主流任务,包括视觉识别、目标检测和语义分割等。然而,一个比较棘手的问题就是这个架构参数量和计算量太大,所以一直被广受诟病。因此,后续有不少工作都是朝着这个方向去改进,例如 Swin-Transformer 在局部非重叠窗口中进行注意力计算,而 PVT 中则是使用平均池化来合并 token 以进一步压缩耗时。然而,这些 ViTs 由于其参数量太大以及高 FLOPs 并不适合部署到移动设备。如果我们直接缩小到适合移动设备的尺寸时,它们的性能又会显著降低。
MobileViT
因此,出现了不少工作聚焦于探索轻量级的视觉变换器,使 ViTs 适用于移动设备,CVHub 此前也介绍过不少的工作,大家有兴趣可以翻阅历史文章读读。例如,苹果团队提出的 MobileViT 研究了如何将 CNN 与 Transformer 相结合,而另一个工作 MobileFormer 则将轻量级的 MobileNet 与 Transformer 进行融合。此外,最新提出的 EdgeViT 提出了一个局部-全局-局部模块来聚合信息。以上工作的目标都是设计具有高性能、较少参数和低 FLOPs 的移动端友好型模型。
EdgeViT
然而,本文作者从频域编码的角度认为,在现有的轻量级模型中,大多数方法只关注设计稀疏注意力,以有效地处理低频全局信息,而使用相对简单的方法处理高频局部信息。具体而言,大多数模型如 EdgeViT 和 MobileViT,只是简单使用原始卷积提取局部表示,这些方法仅使用卷积中的全局共享权重处理高频本地信息。其他方法,如 LVT ,则是首先将标记展开到窗口中,然后使用窗口内的注意力获得高频信息。这些方法仅使用特定于每个 Token 的上下文感知权重进行局部感知。
LVT
虽然上述轻量级模型在多个数据集上取得了良好的结果,但没有一种方法尝试设计更有效的方法,即利用共享和上下文感知权重的优势来处理高频局部信息。基于共享权重的方法,如传统的卷积神经网络,具有平移等变性的特征。与它们不同,基于上下文感知权重的方法,如 LVT 和 NAT,具有可以随输入内容变化的权重。这两种类型的权重在局部感知中都有自己的优势。
NAT
受此启发,本文设计了一种轻量级视觉变换器——CloFormer,其利用了上下文感知的局部增强。特别地,CloFormer 采用了双分支设计结构。
局部分支
在局部分支中,本文引入了一个精心设计的 AttnConv,一种简单而有效的卷积操作符,它采用了注意力机制的风格。所提出的 AttnConv 有效地融合了共享权重和上下文感知权重,以聚合高频的局部信息。具体地,AttnConv 首先使用深度卷积(DWconv)提取局部表示,其中 DWconv 具有共享权重。然后,其使用上下文感知权重来增强局部特征。与 Non-Local 等生成上下文感知权重的方法不同,AttnConv 使用门控机制生成上下文感知权重,引入了比常用的注意力机制更强的非线性。此外,AttnConv 将卷积算子应用于 Query 和 Key 以聚合局部信息,然后计算 Q 和 K 的哈达玛积,并对结果进行一系列线性或非线性变换,生成范围在 [-1,1] 之间的上下文感知权重。值得注意的是,AttnConv 继承了卷积的平移等变性,因为它的所有操作都基于卷积。
全局分支
在全局分支中则使用了传统的注意力机制,但对 K 和 V 进行了下采样以减少计算量,从而捕捉低频全局信息。最后,CloFormer 通过简单的方法将局部分支和全局分支的输出进行融合,从而使模型能够同时捕捉高频和低频信息。总的来说,CloFormer 的设计能够同时发挥共享权重和上下文感知权重的优势,提高其局部感知的能力,使其在图像分类、物体检测和语义分割等多个视觉任务上均取得了优异的性能。
方法
如上图2所示,CloFormer 共包含一个卷积主干和四个 stage,每个 stage you Clo 模块 和 ConvFFN 组合而成的模块堆叠而成 。具体得,首先将输入图像通过卷积主干传递,以获取 token 表示。该主干由四个卷积组成,每个卷积的步长依次为2、2、1和1。接下来,tokens 经历四个 Clo 块和 ConvFFN 阶段,以提取分层特征。最后,再利用全局平均池化和全连接层生成预测结果。可以看出,整个架构非常简洁,支持即插即用!
ConvFFN
为了将局部信息融入 FFN 过程中,本文采用 ConvFFN 替换了常用的 FFN。ConvFFN 和常用的 FFN 的主要区别在于,ConvFFN 在 GELU 激活函数之后使用了深度卷积(DWconv),从而使 ConvFFN 能够聚合局部信息。由于DWconv 的存在,可以直接在 ConvFFN 中进行下采样而不需要引入 PatchMerge 模块。CloFormer 使用了两种ConvFFN。第一种是在阶段内的 ConvFFN,它直接利用跳跃连接。另一种是连接两个阶段的 ConvFFN,主要用于下采样操作。
Clo block
CloFormer 中的 Clo块 是非常关键的组件。每个 Clo 块由一个局部分支和一个全局分支组成。在全局分支中,我们首先下采样 K 和 V,然后对 Q、K 和 V 进行标准的 attention 操作,以提取低频全局信息。
虽然全局分支的设计能够有效减少注意力机制所需的浮点运算次数,并且能够获得全局的感受野。然而,它在处理高频局部信息方面的能力不足。为了解决这个问题,CloFormer 引入了局部分支,并使用 AttnConv 对高频局部信息进行处理。AttnConv 可以融合共享权重和上下文感知权重,能够更好地处理高频局部信息。因此,CloFormer 结合了全局和局部的优势来提高图像分类性能。下面我们重点讲下 AttnConv 。
AttnConv
AttnConv 是一个关键模块,使得所提模型能够获得高性能。它结合了一些标准的 attention 操作。具体而言,在AttnConv 中,我们首先进行线性变换以获得 Q、K和V。在进行线性变换之后,我们再对 V 进行共享权重的局部特征聚合。然后,基于处理后的 V 和 Q ,K 进行上下文感知的局部增强。具体流程可对照流程图理解,我们可以将其拆解成三个步骤。
Local Feature Aggregation
为了简单起见,本文直接使用一个简单的深度卷积(DWconv)来对 V 进行局部信息聚合。
Context-aware Local Enhancement
在将具有共享权重的局部信息整合到 V 中后,我们结合 Q 和 K 生成上下文感知权重。值得注意的是,我们使用了与局部自注意力不同的方法,该方法更加简洁。具体而言,我们首先使用两个 DWconv 对 Q 和 K 分别进行局部信息聚合。然后,我们计算 Q 和 K 的 Hadamard 乘积,并对结果进行一系列转换,以获取在 -1 到 1 之间的上下文感知权重。最后,我们使用生成的权重来增强局部特征。
Fusion with Global Branch
在整个 CloFormer 中,我们使用一种简单直接的方法来将局部分支和全局分支的输出进行融合。具体而言,本文是通过将这两个输出在通道维度上进行直接拼接,然后再通过一个 FC 层聚合这些特征并结合残差输出。
最后,上图展示了三种不同的方法。相比于传统卷积,AttnConv 中的上下文感知权重使得模型能够更好地适应输入内容。相比于局部自注意力机制,引入共享权重使得模型能够更好地处理高频信息,从而提高性能。此外,生成上下文感知权重的方法引入了更强的非线性性,也提高了性能。需要注意的是,AttnConv 中的所有操作都基于卷积,保持了卷积的平移等变性特性。
实验
通过在主流计算机视觉任务上的广泛实验,如 ImageNet 分类、COCO 目标检测/实例分割 和 ADE20K 语义分割,充分证明了 CloFormer 的优越性。如图表所示,CloFormer 在具有相似 FLOPs 的模型中实现了最佳性能。
此外,在消融实验中,作者通过频谱分析表明所提方法通过有机地结合了全局和局部分支,能够有效感知高频和低频信息。
总结
本论文提出了 CloFormer,一种具有上下文感知的局部增强机制的轻量级视觉 Transformer,并且开发了一种新颖的局部感知方法。CloFormer 在类似的 FLOPs 和模型大小的模型中取得了竞争性的性能。特别地,经过精心设计的AttnConv 利用共享权重和上下文感知权重有效地提取高频局部特征表示。此外,采用了双分支结构来混合高频和低频信息。大量实验证明 CloFormer 是一种轻量高效的视觉骨干网络,胜过了许多现有的 SOTA 方法。