MLP-Like Backbone | Strip-MLP跨行Token交互比SWin Transformer更轻更强的性能

集智书童公众号

发布于 2023-09-04 11:33:15

4920

发布于 2023-09-04 11:33:15

本文首发于【集智书童】，白名单账号转载请自觉植入本公众号名片并注明来源，非白名单账号请先申请权限，违者必究。

Token交互操作是MLP模型中的核心模块之一，用于在不同的空间位置之间交换和聚合信息。然而，Token交互在空间维度上的能力很大程度上依赖于特征图的空间分辨率，这限制了模型的表达能力，特别是在深层中，特征被下采样到很小的空间尺寸。为了解决这个问题，本文作者提出了一种名为Strip-MLP的新方法，以3种方式增强Token交互的能力：

首先，本文作者引入了一种新的MLP范例，称为Strip MLP层，它允许Token以跨行（或跨列）的方式与其他Token交互，使得一行（或一列）中的Token能够对相邻但不同行（或列）的条带中的信息聚合做出贡献；
其次，本文作者提出了Cascade Group Strip Mixing Module（CGSMM），以克服由小空间特征尺寸引起的性能下降。该模块使Token能够以在Patch内部和跨Patch的方式更有效地交互，这与特征的空间尺寸无关；
最后，基于Strip MLP层，本文作者提出了一种新的 Local Strip Mixing Module（LSMM），以增强Token在局部区域的交互能力。

大量实验证明，Strip-MLP显著提高了MLP模型在小型数据集上的性能，并在ImageNet数据集上获得了可比甚至更好的结果。尤其是，在Caltech-101数据集上，Strip-MLP模型的平均Top-1准确率比现有的MLP模型提高了2.44%，在CIFAR-100数据集上提高了2.16%。代码：https://github.com/MedProcess/Strip_MLP

1、简介

在计算机视觉中，卷积神经网络（CNN）是最受欢迎的网络Backbone之一，已经取得了一系列的突破性成果。受到自注意力机制在自然语言处理（NLP）领域的巨大成功启发，Transformer模型被引入到计算机视觉领域，并且取得了与最先进的CNN相媲美的结果。虽然ViT及其变种优于传统CNN，但这些模型引入了高计算复杂度来构建注意力图。最近，视觉领域的一些研究表明注意力机制并非必要，提出了更简单的模型架构。

MLP（多层感知器）模型，如MLP-Mixer、gMLP和ViP，使用多层感知器处理数据，显示了改进视觉模型性能的巨大潜力。作为第一个视觉深度MLP网络，MLP-Mixer引入了两种类型的MLP层：

通道混合MLP（CMM）
Token 混合MLP（TMM）

对于CMM，该模块主要在每个 Token 的不同通道之间混合信息。对于TMM，它允许每个空间 Token 在单个MLP层中与所有其他 Token （整个图像）进行交互。然而，这种设计也引入了更多的参数和更高的计算复杂度，容易过拟合。为了解决这个问题，Sparse MLP（SMLP）和Vision Permutator（ViP）提出了一种类似的层，并沿轴向应用一维MLP，在行或列之间共享参数。因此，它减少了模型参数和计算复杂度，避免了常见的过拟合问题。

尽管SMLP和ViP缓解了MLP-Mixer的一些缺点，但这两种方法带来了一个挑战，即当在空间行（或列）上相互作用 Token 时，Token的交互能力非常依赖于特征的空间大小。如图1所示，空间特征分辨率被下采样到一个小尺寸，但通道数更多，这意味着每个 Token 的特征模式主要集中在通道维度而非空间维度上。通过在所有通道之间共享权重来沿空间维度相互作用 Token 会严重忽略不同通道之间的特征模式差异，这可能降低Token的交互能力，特别是在深层中，空间特征分辨率较小。

本文作者将这个问题称为Token的交互困境。以SMLP为例，本文作者详细分析了不同阶段模型的特征分辨率和复杂度。本文作者发现，随着空间特征尺寸逐渐减小， Token 交互层也变得越来越小，使得 Token 交互能力迅速降低。

为了解决上述挑战，本文提出了一种新的高效Strip MLP模型，称为StripMLP，以三种方式丰富 Token 交互层的能力。

对于单个MLP层，受HOG中的交叉块归一化方案的启发，本文作者设计了一个Strip MLP层，允许 Token 以交叉方式与其他 Token 交互，使得每行或每列的 Token 在对其他行或列的贡献上有所不同。
对于 Token 交互模块，本文作者开发了通道级组混合CGSMM，使得一行（或一列）中的 Token 能够对相邻但不同行（或列）的条带中的信息聚合做出贡献，以解决空间特征大小减小但通道数量增加的问题。
考虑到现有方法主要在行（或列）的远程范围内相互作用 Token ，可能不能很好地在局部区域内聚合 Token ，本文作者提出了LSMM，其中包含一个小的Strip MLP单元，以增强在局部区域内的 Token 交互能力。

提出的Strip-MLP模型显著提升了 Token 交互能力，其主要贡献如下：

针对视觉MLP的新型MLP范式：Strip MLP层，以交叉条带方式聚合相邻的 Token ，使得每行或每列的 Token 能够对其他行或列做出不同贡献，从而更有效地交互 Token 。
设计CGSMM和LSMM，分别有效地提升了模型的 Token 交互能力，并在局部区域增强了 Token 的聚合能力；
大量实验证明Strip-MLP显著改进了MLP-based模型的性能。在Caltech-101上，Strip-MLP的平均Top-1准确率比现有的MLP-based模型提高了+2.44%，在CIFAR-100上提高了+2.16%。此外，与传统的MLP-based模型、其他流行的CNN和基于Transformer的模型相比，本文作者的模型在ImageNet-1K上取得了可比甚至更好的性能。

2、相关工作

视觉识别领域的深度神经网络主要可以分为3类：

卷积神经网络（CNNs）
Vision Transformer（ViTs）
基于多层感知器的模型（MLPs）

2.1、基于CNNs的模型

CNNs是视觉任务的事实上的标准深度学习网络模型，并在视觉社区中得到了广泛研究。AlexNet 是一个具有象征意义的模型，在ILSVRC 2012比赛中获得了比其他模型更高的性能。自那时以来，基于CNNs的模型引起了更多关注，并提出了许多更深、更有效的架构。通过卷积和池化层，CNNs在局部区域聚合特征，但在长期依赖关系方面表现不佳，而这些依赖关系在新的视觉模型如Transformer 和MLP模型中得到了优化。

2.2、基于Transformer的模型

Transformer 最初用于机器翻译，并成为所有自然语言处理（NLP）任务的参考模型。受到Transformer在NLP中的巨大成功的启发，ViTs 首次将标准Transformer应用于图像，与SOTA CNNs模型相比取得了出色的结果。DeiT 引入了几种训练策略和蒸馏方法，使ViTs在较小的ImageNet-1K数据集上更加有效。通过引入具有位移窗口的分层Transformer，Swin Transformer 在ImageNet-1K上实现了SOTA准确率，通过在局部窗口和跨窗口连接中使用自注意力提供更高的效率。对于这些模型，自注意力是核心模块，但在获得注意力图时具有较重的计算负担。

2.3、基于MLP的模型

没有卷积和自注意力机制，MLP-Mixer 构建了只使用MLP层的架构，并在图像分类基准上实现了有竞争力的性能。自那时以来，研究人员开发了许多类似MLP的变种模型。[Are we ready for a new paradigm shift? a survey on visual deep mlp]这项工作对视觉深度MLP模型进行了全面的调查，并详细比较了卷积、自注意力机制和 Token 混合MLP之间的内在联系和差异。

Sparse MLP 引入了一种稀疏操作来沿着轴向分别聚合信息，避免了传统MLP的二次计算复杂度。Hire-MLP 通过分层重新排列 Token 提出了一种新颖的基于MLP的架构，以聚合局部和全局空间信息。Wave-MLP 将每个 Token 表示为具有振幅和相位两部分的波函数，具有对不同输入图像的不同内容进行建模的能力。

尽管如此， Token 交互能力在空间特征分辨率变小的情况下明显降低，这在之前的研究中被忽视了。因此，本文旨在在单个MLP层和 Token 交互模块上丰富 Token 交互能力，以提高MLP-based模型的性能。

3、本文方法

在本节中，本文作者首先介绍Strip-MLP的总体架构。然后，本文作者详细展示模型的关键组成部分：

Strip MLP Layer
Cascade Group Strip Mixing Module (CGSMM) 和
Local Strip Mixing Module (LSMM)

并对Strip MLP和传统MLP模型在参数和复杂度方面进行比较分析。

最后，本文作者定义了4种不同尺寸的架构变体，以比较模型的性能。

3.1、整体架构

1、Patch Embedding

Strip-MLP模型的概述如图2(a)所示。本文作者将Strip-MLP设计为多阶段模型的层次结构。给定输入图像

，StripMLP模型首先将其分成一系列图像块（也称为Token），其中

Ip ∈ R^{hw×c}

（

：通道数量）是图像块的大小为

p×p

，图像块的数量为

hw(h = H/p,w = W/p)

。然后，所有图像块都被线性投影到所需的更高维度

的特征空间

(X ∈ R^{hw×C})

。

2、Mixing Block

混合块的目的是增强不同空间位置和通道的特征之间的交互。为了有效地聚合空间和通道信息，本文作者设计了Strip混合块和通道混合块两个子块。Strip混合块由CGSMM和LSMM组成，分别在全局和局部层次上更高效地混合和聚合空间信息。

Patch Merging模块旨在将特征合并，其中空间维度将缩减为2×2，通道维度增加2倍，从而使模型将特征从

H/p×W/p

逐步下采样为

H/8p×W/8p

。

为了获得多尺度特征，本文作者对第1和第2阶段的输出特征应用了一个单一的卷积层。然后，本文作者将结果特征添加到第3和第4阶段的输入特征中。

3、Head Layer

头部层将由多个块提取的特征输入全局平均池化（GAP）层以减少特征维度。最后，特征将被输入到完全连接的头部层进行分类。

3.2、Strip Mixing Block和Channel Mixing Block

1、Strip Mixing Block

为了提高 Token 的交互能力，本文作者设计了该块以并行的方式聚合长距离和短距离的交互。如图2(b)所示，本文作者在通道维度上分割特征，将一半的通道特征输入CGSMM来建模长距离交互，将剩余的特征输入LSMM来聚合局部交互。给定输入特征

X ∈ R^{H×W×C}

，该块可以表示为：

其中Xm和Y分别是该块的中间特征和输出特征。

X_m

被分割为

X_{m_1}

和

X_{m_2}

，每部分有一半的通道。DWSC代表深度可分离卷积，卷积核大小为3×3。MLP是全连接(FC)、批归一化和GELU激活层的串联连接。

Cat(*)

表示连接操作。

2、Channel Mixing Block

该块旨在聚合 Token 的通道间信息，基本结构如图2(c)所示。本文作者遵循[ConvNeXt]中的 Inverted Bottleneck结构和[ConvNeXt v2]中的全局响应归一化（GRN），设计了通道混合块来增加通道的对比度和选择性。

3.3、CGSMM和LSMM

首先，本文作者介绍Strip MLP层。基于这一层，本文作者设计了Cascade Group Strip Mixing Module (CGSMM)，采用简单而有效的策略，将特征沿通道维度分割成Patches，并以within-patch和cross-patch的方式更有效地交互 Token ，这与特征的空间尺寸无关。此外，现有方法的设计只允许 Token 在行（或列）的长距离范围内进行交互并共享权重；因此，由此得到的模型可能难以同时有效地聚合全局和局部信息。

因此，本文作者设计了Local Strip Mixing Module (LSMM)来更高效地捕获局部交互。

1、Strip MLP层

在MLP-based模型中，大多数MLP层独立地处理数据的每一行和每一列，如公式(3)所示，这可能导致 Token 交互的低效性。受HOG（方向梯度直方图）交叉块规范化方案的启发，该方案将块重叠，使得每个标量单元响应都能贡献给不同的块，本文作者提出了Strip MLP层。所提出的层在空间方向上对相邻行或列的“strip”数据应用MLP，以以cross-strip方式聚合特征。给定输入

，本文作者在公式(4)和(5)中表示Strip MLP层（以strip宽度为3为例）：

其中

W_1∼W_4

是MLP层的权重，i和j是行和列中的 Token 索引。

Strip MLP层的优势主要体现在两个方面：

一方面，Strip MLP层使得 Token 能够同时在短距离和长距离上与其他 Token 交互。
另一方面，类似于HOG交叉块规范化过程，每一行（或列）不仅为当前行（或列）的 Token 聚合服务，还为相邻行（或列）的特征聚合做出贡献。

例如，在公式(4)中，

X_{∗,j}

对于聚合

X^h_{∗,j−1}

，

X^h_{∗,j}

和

X^h_{∗,j+1}

的处理特征做出了不同的贡献，从而使得 Token 能够以cross-strip方式更有效地进行交互。

2、Cascade Group Strip Mixing Module

该模块的架构如图3(a)所示。该模块在行和列方向上以级联模式应用Strip MLP层。由于行上的Strip MLP层操作与列上的类似，本文作者以其中一个为例来说明。

数据在通道维度上的patch分割

给定输入特征

X_{m_1} ∈ R^{H×W×C/2}

，该模块首先将特征重新排列为

X^{H×C/2×W}_{m_1}

，并在通道维度上将特征分割成P个Patches，然后在列维度上将它们连接在一起（

X^{H×C/2P×PW}

）。

Group Strip MLP层

为了提高 Token 的交互能力，本文作者提出了Group Strip MLP层（GSML）来在within-patch和cross-patch的方式下与不同块中的 Token 交互。

具体而言，本文作者在不同的块上应用了不共享的Strip MLP层权重，而在同一块中的 Token 之间共享权重以进行交互。然后，本文作者将特征恢复到原始形状，并将其与输入特征连接起来。

为了交互跨块的 Token ，本文作者在它们之间应用了一个通道全连接（Channel FC）层。

3、Local Strip Mixing Module

该模块如图3(b)所示。为了更好地聚合空间维度上的局部交互，本文作者定义了一个小的Strip MLP

单元，其中strip的宽度和长度分别为3和7。给定输入特征

X^{H×W×C/2}_{m_2}

，本文作者同时在行和列方向上聚合局部交互。本文作者使用一个重新加权模块来对所有分支进行求和。

3.4、参数及复杂性分析

在多阶段处理架构中，Token的交互困境变得更加严重，因为在深层中，特征的空间分辨率将被下采样，导致Token的交互能力减弱。

在本节中作者通过参数和复杂度的比较分析来展示所提出的CGSMM的有效性。考虑到大多数基于MLP的模型具有类似的Token交互模型结构，本文作者通过与在各种数据集上表现良好的流行模型Sparse MLP 进行参数和复杂度的分析。

Sparse MLP首先将MLP应用于X的列和行，将

R^{3C×W×H}

映射到

R^{CW×H}

。然后，该模型将两个处理过的特征与输入特征

进行连接，并在通道维度上使用线性层将特征从

R^{HW×3C}

融合到

R^{HW×C}

。第一步的参数数量和FLOPs分别为

W^2+H^2

和

CHW(H + W)

；融合步骤的参数数量和FLOPs分别为

3C^2

和

3HWC^2

。

相比之下，Strip MLP层用于Token交互的参数数量和FLOPs分别为

3P(H^2 + W^2)

和

3CHW(H + W)

。融合步骤有

4C^2

个参数和

4HW C^2

个FLOPs。特别地，对于

H_1 = W_1 = 56

，

H_4 = W_4 = 7

，

C_1 = 112

，

C_4 = 112 × 2^3 = 896

，以及

P = C_4

，本文作者计算了两个层的参数数量和FLOPs。

如表1所示，在Sparse MLP块中，第4阶段的Token交互参数数量（仅0.10k）比第1阶段减少了62.70倍。此外，大部分块的参数数量和FLOPs集中在融合步骤中。例如，在第4阶段，仅有0.01%的参数（0.52%的FLOPs）用于Token交互步骤。基于上述分析，为了提高Token的交互能力，最好重新设计块以平衡参数数量和FLOPs。

在CGSMM中，本文作者采用了一种简单但有效的策略，即沿着通道维度将特征分割成块，并以Patch内和Patch间的方式相互作用。无论空间分辨率如何降低，该模块仍然可以以通道方式在不同块中相互作用。在表1中，本文作者将第4阶段中Token交互层的参数数量占总参数数量的比例从0.01%提高到2.01%。

虽然本文作者的CGSMM在考虑到仅有一个Token交互层时带来了更多的参数和计算，但由于Token交互能力的提高以及本文作者对T1∼T4进行了小型模型配置，模型的总参数数量和整体计算复杂性却有所减少。此外，第4节的实验结果显示，与其他SOTA模型相比，本文作者的模型在更少的参数和FLOPs下取得了更好的性能。

3.5、结构变体

本文作者开发了4个Strip-MLP网络的变体：Strip-MLP-T∗（轻型微小），Strip-MLP-T（微小），Strip-MLP-S（小型），Strip-MLP-B（基础），它们与基于MLP的模型和Swin Transformer相比具有类似或更小的模型大小。

这4个模型的超参数如下：

Strip-MLP-T∗：C = 80，{T1∼T4} = {2, 2, 6, 2}；
Strip-MLP-T：C = 80，{T1∼T4} = {2, 2, 12, 2}；
Strip-MLP-S：C = 96，{T1∼T4} = {2, 2, 18, 2}；
Strip-MLP-B：C = 112，{T1∼T4} = {2, 2, 18, 2}。

其中，C表示第一阶段隐藏层的通道数，T1 ∼ T4表示在每个阶段的Strip Mixing Block和Channel Mixing Block中重复的次数，如图2所示。

4、实验

4.1、Results on Caltech-101

4.2、Results on CIFAR-100

4.3、Transfer Learning on Small Datasets

4.4、Image Classification on ImageNet-1K

4.5、消融实验

为了更好地展示本文作者提出方法的有效性，本文作者对模型设计的关键组件进行了消融实验。由于GPU资源有限，本文作者在Caltech-101和CIFAR-100两个数据集上进行了消融研究。

1、Strip MLP中Strip宽度的影响

Strip宽度影响着 Token 交互的范围，并确定了每行或每列对相邻 Token 的贡献范围。本文作者进一步通过将宽度从1变化到7，并以2为步长进行实验来验证其影响。表6显示了本文作者模型在Caltech-101和CIFAR-100上的性能。当宽度变大时，性能增加并趋于饱和，表明Strip MLP层改善了 Token 交互能力。在前面的所有比较实验中，本文作者根据消融实验的结果一致将Strip宽度设置为3。

2、CGSMM中Patch数量的影响

不同的Patch数量会带来不同程度的 Token 交互能力改进，并影响模型性能。在表7中，本文作者设计了5种不同的Patch数量，从C减少到1，以显示CGSMM的有效性。在没有Patch分割（P = 1）和GSML操作的情况下，模型在Caltech-101上的性能下降了1.16%（从92.09%下降到90.93%），在CIFAR-100上下降了1.47%（从86.50%下降到85.03%），这一致地证明了CGSMM在提高 Token 交互能力方面的有效性。

从实验结果中本文作者可以观察到，最优的Patch数量在数据集之间有所不同，因此最优的Patch数量应通过验证实验来确定。在其他消融研究中，本文作者一致将Patch数量设置为C/4。

3、GSML的级联与并行结构

将GSML应用于级联结构使得 Token 在一个模块中与整个2D空间的其他 Token 交互，而并行结构需要两个模块，可能会降低 Token 交互的效率。在表8中，本文作者测试了CGSMM的级联结构和并行结构之间的效果。本文作者的实验表明，级联结构相较于并行结构在Caltech-101和CIFAR-100数据集上获得了更高的准确率，分别增加了+0.48%和+0.61%。

4、CGSMM和LSMM的作用

为验证两个互补分支的重要性，本文作者进行了实验，其中只保留了CGSMM或LSMM中的一个分支。表9显示了本文作者在Caltech-101和CIFAR-100上的试验结果。值得注意的是，本文作者发现从任何一个分支中移除任何模块都会显著降低模型的性能。这些结果突显了CGSMM和LSMM在丰富 Token 交互能力方面的重要作用。