S2-MLP v2: Improved Spatial-Shift MLP Architecture for Vision

  • 引入多分支处理的思想,并应用 Split-Attention 来融合不同分支。
  • 受现有工作的启发,使用更小的 patch 和分层金字塔结构。
  • 扩增偏移的种类

Abstract

        最近,基于MLP的视觉框架出现了。基于MLP的视觉网络结构在图像识别方面具有较低的 inductive bias,与CNN和vision Transformers相比具有较强的性能。其中空间移位MLP (S2-MLP)采用直截了当的空间移位操作,其性能优于MLP-Mixer和ResMLP等开创性结构。最近,通过使用金字塔结构的小块patch,Vision Permutator (ViP)和Global Filter Network (GFNet)获得了比S2-MLP更好的性能。本文对S2-MLP视觉框架进行了改进。我们沿着通道维度展开特征图,并将扩展的特征图分解为几个部分。我们对分组的部分进行不同的空间移位操作。同时,我们利用分散注意力的操作来融合这些分开的部分。此外,与同类算法一样,我们采用了更小尺度的patch和金字塔结构来提高图像识别精度。我们将改进的空间移位MLP视觉网络框架称为S2-MLPv2。使用55M参数,我们的中型模型S2-MLPv2在没有自我注意和外部训练数据的情况下,在ImageNet-1K基准using224×224images上实现了83.6%的top-1准确率。

1 Introduction

        近年来,人们对计算机视觉进行了大量的研究,以期在更小的inductive bias下实现更高的性能。出现了两种类型的架构,包括vision transformer (Dosovitskiy et al.,2021;Touvron et al.,2020)和基于MLP的结构(Tolstikhin et al.,2021;Touvron et al.,2021a)。与使用精心设计的卷积核的人工视觉框架CNN (He et al.,2016)相比,vision transformer和基于MLP的结构在图像识别方面都取得了具有竞争力的性能,而且无需昂贵的手工设计。具体来说,vision Transformer模型堆叠了一系列Transformer块,实现了全局接收领域。

        基于MLP的方法,如MLP-Mixer (Tolstikhin et al.,2021)和ResMLP (Touvron et al.,
2021a),通过沿着MLP实现的不同patch进行投影,实现patch之间的通信。与MLP-Mixer和ResMLP不同,spatial-shift MLP (S2-MLP) (Yu et al.,2021b)对patch之间的通信采用了非常直接的空间移位操作,在ImageNet1K数据集上无需外部训练数据就能实现更高的图像识别精度。同时,Vision Permutator (ViP) (Hou et al.,2021)沿着高度和宽度维度对特征表示进行编码,同时利用两级金字塔结构的更细的patch尺寸,取得了比S2-MLP更好的性能。CCS-MLP (Yu et al.2021a)设计了一个循环的token-mixing MLP来实现平移不变性。Global Filter Networks (GFNet) (Rao et al.,2021b)利用二维傅里叶变换将空间patch特征映射到频域,并在频域进行跨patch通信。Rao等人(2021b)指出,频域的token-mixing操作与循环权值的深度卷积等效。为了达到较高的识别精度,GFNet还采用了金字塔结构的较小尺寸的patch。最近,AS-MLP (Lian et al.,2021)对特征图的通道进行了轴向移动,并设计了一个四层金字塔,取得了出色的性能。与此同时,Cycle-MLP (Chen et al.,2021a)设计了多个用于空间投影的伪核,也取得了优异的性能。值得注意的是,AS-MLP (Lian et al.,2021)和Cycle-MLP (Chen et al.,2021a)都是基于精心设计的四层金字塔。

        在这项工作中,我们重新思考了空间移位MLP (S2-MLP) (Yu等人,2021b)的设计,并提出了一个改进的空间移位MLP (S2-MLPv2)。与原S2-MLP相比,主要在两个方面进行了修改:

  • 如图1(b)所示,我们沿着通道维度展开特征图,并将扩展的特征图分成多个部分。对于不同的部分,我们进行不同的空间移位操作。我们利用分裂注意操作(Zhang et al.,2020)来融合这些分裂的部分。

  • 我们采用更小尺度的patch和分层金字塔结构,就像现有的基于MLP的架构,如ViP (Hou et al.,2021)、GFNet (Rao et al.,2021b)、ASMLP (Lian et al.,2021)和loop - mlp (Chen et al.,2021a)。

图1:S2-MLP (Yu et al.,2021b)和拟议的S2-MLPv2的空间移位操作之间的比较。在S2-MLP中,通道平均分为四个部分,每个部分沿不同方向移动。MLP是在移位的信道上进行的。相比之下,在S2-MLPv2中,c通道特征图被扩展为3c通道特征图。然后将扩展后的映射沿通道维度平均分成三部分。对于每个部分,我们进行不同的空间移位操作。然后将移位的部分合并。

        我们将改进的空间移位MLP架构称为S2-MLPv2。在图1中,我们可视化了原始空间移位MLP (S2-MLP)和改进版本S2-MLPv2之间的区别。我们在公众基准ImageNet-1K上进行的实验表明,所提出的S2-MLPv2的图像识别精度达到了最高水平。具体来说,使用55M参数,采用224×224的图片,我们的中等规模模型S2-MLPv2在没有自我注意和外部训练数据的情况下实现了83.6%的top-1准确率。

2 Related Work

        vision Transformer. vision Transformer (ViT) (Dosovitskiy et al.,2021)将图像切割到16×16的patches中,并将每个patch作为Transformer输入中的一个token。这些patch/token由一堆Transformer层处理,以便彼此通信。利用大规模的预训练数据集,其图像识别精度与CNN相比具有一定的竞争力。DeiT (Touvron et al.,2020)采用了更先进的优化器和数据增强方法,使用中等规模的预训练数据集取得了很好的结果。金字塔vision Transformer (PvT) (Wang et al.,2021b)和PiT (Heo et al.,2021)利用金字塔结构,逐渐缩小空间维度,扩大隐藏尺寸,获得更好的性能。Tokensto-Token (T2T) (Y yuan et al.,2021)和Transformer-in-Transformer (TNT) (Han et al.,2021)提高了建模每个patch/token局部结构的有效性。为了克服全局自我注意的低效性,Swin (Liu et al.,2021b)在局部窗口内进行自我注意,但通过改变窗口设置实现全局接收场。Shuffle Transformer (Huang et al.,2021)也利用局部自注意窗口,通过切换空间维度和特征维度实现跨窗口通信。Twins (Chu et al.,2021a)通过全局下采样注意增强了局部窗口内的自我注意。DynamicViT (Rao et al.,2021a)和SViTE (Chen et al.,2021b)利用稀疏性实现高效率。CaiT (Touvron et al.,2021b)通过叠加数十层来探索极深的结构。最近,PVTv2 (Wang et al.,2021a)利用重叠patch嵌入、卷积前馈网络和线性复杂度注意层改进了PvT。CSwin Transformer (Dong et al.,2021)通过交叉形状窗口并行计算水平和垂直条纹中的自我注意来改进Swin。Focal Transformer (Y ang et al., 2021)还开发了更先进的局部窗口,仅在局部参与细粒token,但在全球进行了总结。

        MLP-based架构。MLP-Mixer(Tolstikhin et al.,2021)是基于MLP的视觉架构的开创性工作。它提出了一个token-mixing MLP,由两个完全连接的层组成,用于补丁之间的通信。Res-MLP (Touvron等人,2021a)简化了token-mixing到单全连接层,并使用更多的中继层探索更深的架构。Spatialshift MLP backbone (S2-MLP) (Y u et al.,2021b)采用跨patch通信的空间移位操作。Vision Permutator (ViP) (Hou et al.,2021)分别沿着高度维度和宽度维度混合token。同时,ViP采用金字塔结构作为PvT (Wang et al.,2021b),其性能明显优于MLP-mixer, Res-MLP和S2MLP。CCS-MLP (Y u et al.,2021a)重新思考了MLP-Mixer和Res-MLP中token-mixing MLP的设计,并提出了一种循环信道专用MLP。具体地说,他们将token-mixing MLP的权矩阵设计为一个循环矩阵,参数较少。同时,利用快速傅里叶变换(FFT)可以有效地计算出向量与循环矩阵的乘法。Global Filter Network (GFNet) (Rao et al.,2021b)通过2D FFT将patch特征映射到频域,并在频域混合令牌。Rao等人(2021b)证明,GFNet中的全局滤波器与滤波器大小H×W等价于一个深度全局圆卷积。同时,GFNet还采用了金字塔结构来提高识别精度。在本工作中,我们重新思考了S2-MLP的设计,大大提高了它的图像识别性能。w\times h

3 Preliminary

3.1 Spatial-Shift MLP (S2-MLP)

       在本节中,我们简要回顾S2-MLP (Y u等人,2021b)架构的结构。该算法由patch embedding层、S2-MLP块堆栈和分类头组成。

        S2-MLP: Spatial-Shift MLP Architecture for Vision_HHzdh的博客-CSDN博客

        Patch embedding layer 。首先将一个大小为W\times H\times 3的输入图像裁剪成w\times h的patches。每个patch为p\times p\times 3,以及p=\frac{W}{w}=\frac{H}{h},然后通过全连通层将每个patch映射到d维向量。

图2:S2-MLP模块的结构

 

        Spatial-shift MLP block。如图2所示,它由用于channel-mixing的MLP层和用于混合patches的空间移位层组成。下面我们只介绍空间移位模块。给定一个输入张量\chi \in \mathbb{R}^{w\times h\times c},它首先沿着通道维将\chi等分为四个部分\left \{ \chi _i \right \}_{i=1}^{4},并沿四个方向移动:

        值得注意的是,S2-MLP (Yu et al.,2021b) 堆叠的N个Spatial-shift MLP块具有相同的设置,并没有利用金字塔结构作为其MLP-backone对等物,如Vision Permutator (Hou et al.,2021)和Global Filter Network (GFNet) (Rao et al.,2021b)。

3.2 Split Attention

        Vision Permutator (Hou et al.,2021)采用ResNeSt (Zhang et al.,2020)中提出的注意分割方法对来自不同操作的多个特征地图进行增强。具体来说,我们将K个大小为n\times c的特征映射表示为\left [ X_1,X_2,...,X_K \right ],其中n为patch的数量,c为通道的数量,split-attention操作首先对它们进行平均,得到:

a=\sum_{k=1}^{K}IX_k……(2)

其中I\in \mathbb{R}^{n}全为1的一维行向量。然后a\in \mathbb{R}^c通过MLP堆栈并生成:

\hat{a}=\sigma(aW_1)W_2……(3)

其中σ为GELU实现的激活函数,W_1\in \mathbb{R}^{c\times \bar{c}},W_2\in \mathbb{R}^{\bar{c}\times Kc}是MLP的权重,输出\hat{a}\in \mathbb{R}^{K\times c}。然后将\hat{a}重塑为矩阵ˆ\hat{A}\in \mathbb{R}^{K\times c},再通过沿第一维的softmax函数进行处理,生成\bar{A}=softmax(\hat{A})\in \mathbb{R}^{K\times c}然后生成参与特征映射\hat{X}

\hat{X}\left [ i,: \right ]=\sum_{k=1}^{K}X_k\left [ i,: \right ]\odot \bar{A[k,:]}……(4)

\odot表示两个向量在元素上的乘法。

4 S2-MLPv2

        在本节中,我们将介绍拟议的S2-MLPv2体系结构。与S2-MLP结构类似,S2-MLPv2结构由patch embedding层、S2-MLPv2块堆栈和分类头组成。由于我们在上一节中已经介绍了patch嵌入层,下面我们只介绍提出的S2-MLPv2块。

图3:S2-MLPv2块的结构

 

4.1 S2-V2 Block

        所述S2-MLPv2模块由两部分组成,即S2-MLPv2组件和通道混合MLP (CM-MLP)组件。给定输入特征\chi \in \mathbb{R}^{w\times h\times c},执行:

                                             \gamma =S^{2}.MLPv2(LN(\chi))+\chi

Z=CM.MLP(LN(\gamma))+\gamma……(5)

        channel-mixing MLP (CM-MLP)采用与MLP-Mixer (Tolstikhin et al.,2021)和ResMLP (Touvron et al.,2021a)相同的结构,因此我们在这里跳过它们的细节。下面我们只详细介绍提出的S2-MLPv2组件。

        给定一个输入特征映射\chi \in \mathbb{R}^{w\times h\times c},S2-MLPv2组件首先通过一个MLP将从c个通道扩展为3c:

\hat{\chi}={MLP}_1(\chi)\in \mathbb{R}^{w\times h\times 3c}……(6)

        然后将扩展后的feature map即\hat{\chi}沿着通道维度平均分割为三部分:

\chi_1=\hat{\chi}[:,:,1:c],\chi_2=\hat{\chi}[:,:,c+1,2c],\chi_3=\hat{\chi}[:,:,2c+1,3c]……(7)

        它通过两个空间移位层SS_1(\cdot )SS_2(\cdot )\chi_1\chi_2进行移位。SS_1(\cdot )进行与等式(1)相同的空间移位操作(即按照之前的移位方法)。相比之下,SS_2(\cdot )SS_1(\cdot )进行了非对称的空间移位操作。具体来说,给定特征图\chi_2,{SS_2} (\chi_2)的行为为:

        请注意,我们有意将SS_1(\cdot )SS_2(\cdot )设计成非对称结构,以便它们相互补充。同时,我们没有移动\chi_3,只是保留它。

        然后将\left \{ \chi_k \right \}_{k=1}^{3}重塑成矩阵\left \{ X_k \right \}_{k=1}^{3},其中X_k\in \mathbb{R}^{wh\times c},将其以式2、式3、式4的形式输入到split-attention(SA)模块中生成:

\hat{X}=SA(\left \{ X_k \right \}_{k=1}^{3})……(9)

        然后,feature map A进一步输入到另一个MLP层生成输出:

\bar{X}=MLP_2(\hat{X})……(10)

         所提出的S2-MLPv2模块的结构如图3所示,详细内容列在算法1中。

算法1  我们的S2-MLPv2模块的伪代码

 

4.2 Pyramid Structure

        继vision Permutator (Hou et al.,2021)之后,我们也利用了两层金字塔结构来提高性能。为了与Vision Permutator (Hou et al.,2021)进行公平比较,我们采用了完全相同的金字塔结构。详细信息见表1。我们注意到类似PVTv2 (Wang et al.,2021a)、as - mlp (Lian et al.,2021)和Cycle-MLP (Chen et al.,2021a)等工作采用了更先进的金字塔结构,早期块中的patch更小。小块可能更擅长捕捉细粒度的视觉细节,从而获得更高的识别精度。然而,由于有限的计算资源,我们不可能重新实现所有这些金字塔结构。此外,我们还注意到Vision Permutator设计了一个具有相当多参数和FLOPs的大模型。然而,由于计算资源有限,大模型对我们来说也不可行。

表1:在我们的S2-MLPv2中使用的两级金字塔结构的配置。我们利用了小型和中型设置,它们与Vision Permutator (Hou et al.,2021)完全相同,以进行公平的比较。AS-MLP (Lian et al.,2021)和Cycle-MLP (Chen et al.,2021a)采用更高级的四层金字塔结构,小尺度的patch。

5 Experiments

6 Conclusion

        本文对空间移位MLP (S2-MLP)模型进行了改进,提出了S2-MLPv2模型。将feature map进行扩展,并将扩展后的feature map分为三部分。它将每个部分单独移动,然后通过split-attention融合分开的特征图。同时,我们利用层次金字塔来提高其建模细粒度细节的能力,以获得更高的识别精度。在没有外部训练数据集的情况下,采用224×224的images,我们的s2-mlv2-medium模型在ImageNet1K数据集上取得了83.6%的top-1准确率,这是目前基于MLP的方法中最先进的性能。同时,与基于transformer的方法相比,我们的S2-MLPv2模型在不需要自我注意的情况下,参数更少,达到了相当的精度。

        与基于MLP的先驱作品如MLP-Mixer、ResMLP以及最近类似MLP的模型如Vision Permutator和GFNet相比,空间移位MLP的另一个重要优势是,空间移位MLP的形状对图像的输入尺度是不变的。因此,经过特定尺度图像预训练的空间移位MLP模型可以很好地应用于具有不同尺寸输入图像的下游任务。未来的工作将致力于不断提高空间移位MLP体系结构的图像识别精度。一个有希望且直接的方向是尝试更小尺寸的patch和更高级的四层金字塔,如CycleMLP和AS-MLP,以进一步减少FLOPs和缩短基于transformer模型之间的识别差距。

  • 1
    点赞
  • 10
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值