(CCS-MLP)Rethinking Token-Mixing MLP for MLP-based Vision Backbone

最新推荐文章于 2024-06-22 09:35:30 发布

HHzdh

最新推荐文章于 2024-06-22 09:35:30 发布

阅读量808

点赞数

分类专栏： PyTorch-MLP 文章标签： python pytorch mlp

本文链接：https://blog.csdn.net/weixin_44855366/article/details/120914024

版权

PyTorch-MLP 专栏收录该内容

15 篇文章 5 订阅

订阅专栏

CCS基于循环行列式的MLP视觉主干网络

本文改进MLP-Mixer的TM，提出一种循环特定通道模块，其分为循环和通道特定两部分。循环指的是将权重矩阵用向量+循环平移的方法实现，好处是提升模型的空间平移性能；通道特定指的是分组做CM，原来的CM是权值共享的，具有通道不变性，改成分组后就引入了些通道特殊性。

Abstract

在过去的十年里，我们见证了机器视觉框架的快速发展。卷积神经网络(CNN)通过引入图像处理中的归纳偏差，在众多的计算机视觉任务中取得了优异的性能，并已成为主流框架。近年来，受Transformer在NLP任务中取得的巨大成功的启发，vision Transformer模型应运而生。使用更少的归纳偏差，与CNN的对手相比他们在计算机视觉任务中取得了很好的表现。最近，研究人员研究使用pure-MLP架构来构建视觉结构，进一步减少归纳偏差，获得良好的性能。纯MLP网络建立在channel-mixing MLP和token-mixing MLP，实现patch之间的通信。本文对token-mixing MLP的设计进行了重新思考。我们发现，在现有的基于MLP的结构中，token-mixing MLP是空间特异性的，因此它对空间转换非常敏感。同时，现有的token-mixing MLP的信道不可知特性限制了它们token mixing的能力。为了克服这些局限性，我们提出了一种改进的结构，称为循环特定信道(CCS)token-mixing MLP，它是空间不变的和信道特定的。它需要的参数更少，但在ImageNet1K基准上获得了更高的分类精度。

1 Introduction

卷积神经网络(CNN)[20,13]在过去的几年里已经成为计算机视觉的重要框架。最近，受Transformer[34]在NLP中取得的成就的启发，出现了几种可视化Transformer方法[9,32]。与CNN相比，vision Transformer方法不需要手工制作卷积核，只是简单地堆叠几个Transformer块。尽管vision Transformer方法的归纳偏差要小得多，但它们取得了与CNN模型相当甚至更好的识别精度。最近，研究人员又向前迈进了一步。他们提出了基于MLP的模型[30,31,25]，其中只有MLP层采取更少的归纳偏差。

MLP-Mixer[30]是pure-MLP视觉框架的先驱之作。它建立在两种类型的MLP层上:channel-mixing MLP layer 和 token-mixing MLP layer 。channel-mixing MLP layer用于信道之间的通信。同时，token-mixing MLP layer 在patch之间进行通信。与vision Transformer基于patch间关系自适应分配注意力相比，token-mixing MLP基于patch的空间位置分配固定注意力。实际上，channel-mixing MLP与深度卷积密切相关[4,17,19]。token-mixing MLP和深度卷积之间的差异有三个方面:

首先，token-mixing MLP具有全局接收域，而深度卷积只有局部接收域。全局接收字段使token-mixing MLP能够访问图像中的整个视觉内容。
其次，深度卷积是空间不变的，而channel-mixing MLP不再具有空间不变的特性，因此其输出对空间平移敏感。
最后，对于特定的位置，token-mixing MLP给不同通道中的元素分配相同的权重。而深度卷积则在不同的通道上采用不同的卷积核，以编码更丰富的视觉模式。

通过观察现有的基于MLP的架构中token-mixing MLP的优缺点，我们提出了一种改进的结构——循环通道专用(CCS)token-mixing MLP，它保留了现有的token-mixng MLP的优点，克服了其缺点。与普通的token-mixing MLP类似，我们的CCS token-mixing MLP有一个全局接收字段。同时，我们在权值矩阵中采用循环结构来实现空间不变性。此外，我们采用特定于通道的设计来利用更丰富的方式混合token。在表1中，我们比较了我们的CCS与深度卷积和vanilla token-mixing MLP的特性。此外，得益于循环结构，我们的CCS需要的参数比token-mixing MLP少得多。其中，CCS只需要 $GN$ 个参数，其中 $N$ 为patch的数量， $G$ 为groups的数量。相比之下，token-mixing MLP通常采用 $N^2$ 参数， $N$ 明显大于 $G$ 。通过在两个现有的纯MLP视觉框架中使用提出的CCS替代token-mixing MLP，我们在ImageNet1K上以更少的参数实现了更高的识别精度。

表1 深度卷积，MLPmixer[30]中的token-mixing MLP和我们的循环特定信道(CCS)token-mixing MLP

2 Related Work

2.1 Vision Transformer

Vision Transformer (ViT)[9]是在计算机视觉任务中单独采用Transformer层架构的开创性工作。它将图像裁剪成不重叠的patch，并通过一堆Transformer层为这些patch提供信息，以便在patch之间实现通信。使用较少的手工设计，ViT与CNN一样实现了具有竞争力的识别精度。然而，它需要大量的图像进行预训练。DeiT[32]采用了更先进的数据增强方法和更强大的优化器，在中等规模的数据集上进行训练，取得了优异的性能。PVT[35]在ViT中引入了渐进收缩金字塔，提高了识别精度。PiT[15]集成了Transformer块之间的深度卷积，并设计了一个缩小的金字塔结构。同时，token-to-token ViT[39]通过递归聚合相邻的token有效地模拟了局部结构。它以更少的浮点数实现了更高的识别精度。Transformer-iN-Transformer (TNT)[12]也专注于局部结构的建模。它设计了一个额外的Transformer来对每个patch中的内在结构信息进行建模。最近，vision Transformer的重点是利用局域性和稀疏性来提高效率。例如，Swin[23]开发了一个带有可变窗口大小的本地窗口Transformer层。利用局部窗口大大提高了效率，可变窗口范围实现了全局接收域。Twins[5]交替堆叠一个局部密集的Transformer 和一个全局稀疏的Transformer ，也以一种有效的方式实现全局接收域。Shuffle Transformer[18]也采用本地窗口Transformer，通过空间变换实现跨窗口连接。S2ViTE[3]研究了在vision Transformer中集成稀疏性，从模型和数据两方面提高了效率。最近的工作[10,37]研究了结合卷积和Transformer来构建混合视觉框架。

2.2 MLP-based Backbone

最近，MLP-Mixer[30]提出了一种比vision Transformer和CNN归纳偏差更小的纯MLP框架，在图像识别方面取得了优异的性能。它建立在两种类型的MLP层之上：channel-mixing MLP和token-mixing MLP。channel-mixing MLP等效于1×1卷积层。它实现了信道之间的通信。token-mixing MLP实现跨patch通信。它类似于Transformer中的自注意块。但是Transformer中的注意依赖于输入patch，而token-mixing MLP中的注意则与输入不相关。前馈(FF)[25]采用与MLP-Mixer类似的结构，也取得了优异的性能。ResMLP[31]简化了MLP-Mixer中的token-mixing MLP，采用了深度架构，大量叠加层。同时，为了稳定训练，ResMLP提出了一个仿射变换层来代替MLP-Mixer中的层归一化。得益于利用了更深层次的体系结构，ResMLP比MLP-MIxer获得了更好的性能。同时，通过巧妙地权衡隐藏的大小和深度，ResMLP比MLP-Mixer使用更少的FLOPs和更少的参数。最近，gMLP[22]利用一个门控操作来提高token-mixing MLP的有效性，获得比MLP混合更高的识别精度。与此同时，External Attention[11]用全连接层实现的对外部记忆的注意取代了自我注意操作，取得了与vision transformer相当的性能。最近，spatial-shift MLP (S2MLP)[38]采用了一种spatial-shift操作[36,2]来实现patch之间的通信，取得了很好的性能。

3 Preliminary

在本节中，我们简要回顾现有的主流MLP框架，MLPMixer[30]。它由三个模块组成，包括每个patch的全连接层、一个N层的mixer层堆栈和一个用于分类的全连接层。

MLP-Mixer: An all-MLP Architecture for Vision_HHzdh的博客-CSDN博客

Input 对于大小为 $W\times H\times 3$ 的输入图像，我们将其裁剪为大小为 $p\times p\times 3$ 大小的非重叠patcj。对于每个patch，它展开成向量 $p\in \mathbb{R}^{3p^2}$ 。综上，我们得到一组patch特征向量 $P=\left \{ p_0,...,p_{N-1} \right \}$ ，这就是MLP-Mixer的输入。

Per-patch fully-connected layer 对于每个patch特征向量 $p_i\in P$ , per-patch全连接层将其映射到一个C维向量：

$x_i\leftarrow W_0 {p_i}+b_0$ ……(1)

其中 $W_0\in \mathbb{R}^{3p^2\times C}$ ， $b_0\in \mathbb{R}^{C}$ 为per-patch fully-connected层的权值。

Mixer layers MLP-Mixer堆叠L个相同大小的混合层，每一层包含两个MLP块：token-mixing MLP块和channel-mixing MLP块。设每个Mixer层输入中的patch特征为 $X=[x_0,...,x_{N-1}]\in \mathbb{R}^{C\times N}$ ，其中C为通道数量，N是patch的数量。channel-mixing MLP块通过下式在通道维度上投影patch特征：

$U=U+W_2\sigma [W_1LayerNorm(X)]$ ……（2）

$W_1\in \mathbb{R}^{rC\times C}$ 表示全连接层的权重，将特征从C维增加到rC维（r>1是膨胀比)。 $W_2\in \mathbb{R}^{C\times rC}$
表示全连接层的权值，其特征维数从rC维降至c维。LayerNorm(·)表示Transformer-based models中广泛使用的层归一化[1]，而σ(·)表示GELU[14]实现的激活函数。

然后channel-mixing块的输出U输入到token-mixing块中进行patch之间的通信：

$Y=U+\sigma[LayerNorm(U)W_3]W_4$ ……（3）

$W_3\in \mathbb{R}^{N\times M} ,W_4\in \mathbb{R}^{M\times N}$ 表示沿标记维度投射patch特征的全连接层的权重。最近ResMLP[31]采用了简化的token-mixing块：

$Y=U+LayerNorm(U)W_3$ ……（4）

$W_3\in \mathbb{R}^{N\times N}$ 为方阵，在ResMLP[31]中的实验表明，简化的token-mixing块的性能与公式(3)中定义的原token-mixing块相当。除非另有说明，下面提到的token-mixing是公式(4)中定义的简化形式。

Classification head 在L层Mixer之后，N个patch的特征就生成了。通过平均池化，它们进一步聚合成一个全局矢量。然后将全局向量送入全连接层进行分类。

4 Circulant Channel-Specific Token-mixing MLP

在本节中，我们将介绍我们的循环信道特定(CCS)token-mixing MLP。正如我们在介绍中提到的，MLP-Mixer和ResMLP中的普通token-mixing MLP是特定于空间和与通道无关的。空间特性使得token-mixing MLP对空间平移敏感。同时，通道无关配置限制了其mixing token的能力。设计CCS token-mixing MLP的动机是为了实现空间不可知的特性，同时以更有效的方式获得特定于信道的token-mixing。为此，我们对普通token-mixing MLP做了两处修改，如图1、算法1所示，详细信息如下所示。

4.1 Circulant structure

在循环结构中，我们在等式（4）中设计了权重矩阵 $W_3$ 。具体来说，我们将权重矩阵 $W_3$ 设置为如下形式:

图1所示比较现有基于MLP的框架中Vanilla token-mixing MLP和提出的循环通道特定(CCS)令牌混合MLP。该算法对token-mixing MLP的权值施加循环结构约束。同时，我们的CCS沿着通道维度将input X分成多个组。它在不同的组上使用不同的权重矩阵，以更灵活的方式mixing token。

它由一个向量 $w=[w_0,...,w_{N-1}]$ 完全指定，它出现在 $W_3$ 的第一列。对于mixing token来说，循环结构自然是空间不可知的。下面我们将详细介绍。设等式(4)中的LayerNorm(U)表示为 $\hat{U}=[\hat{u}_0,...,\hat{u}_{N-1}]$ ，其中 $\hat{u}_i(i\in [0,N-1])$ 为层归一化后的patch feature。我们把等式(4)写成：

$Y=U+\hat{U}W_3$ ……（6）

让我们用 $\hat{u}_i$ 表示的 $\hat{U}$ 的第i列和用 $\bar{u}_i$ 表示 $\hat{U}W_3$ 的第i列。那么很容易得到它：

$\bar{u}_i=\sum_{j=0}^{N-1}w_j\hat{u}_{(i+j)%N},\forall i\in [0,N-1]$ ……（7）

其中%表示拼接循环的意思。如等式(7)所示，获取每个patch 特征 $\bar{u}_i$ 的混合操作对位置不变。因此，带有循环结构的token-mixing MLP的 $W_3$ 是空间不可知的，它对空间转换不敏感。同时，循环结构从减少了参数数量，从 $N^2$ 减少为 $N$ 。同时，利用快速傅里叶变换(FFT)计算N维向量与N×N循环矩阵的乘法运算复杂度仅为 $O(NlogN)$ ，比普通的 $O(N^2)$ )少。效率更高。具体来说，给定一个向量 $x\in \mathbb{R}^N$ ，通过这种方法可以有效地计算矩阵向量乘法Wx：

$xW=FFT[IFFT(x)\odot FFT(w)]$ ……（8）

其中FFT表示快速傅里叶变换，IFFT表示快速反傅里叶变换， $\odot$ 表示j矩阵元素相乘。由于FFT和IFFT都只需要 $O(NlogN)$ 的计算复杂度，所以总复杂度仅为 $O(NlogN)$ 。相比之下，现有方法中的普通token-mixing MLP需要 $O(N^2)$ 的计算复杂度。但当patch数量N较小时，复杂度为 $O(NlogN)$ 的FFT与复杂度为 $O(N^2)$ 的普通token-mixing MLP相比，无法表现出其效率优势。

4.2 Channel-Specific settings

MLP-Mixer[30]中的token-mixing MLP在不同的通道中具有相同的权重。即在 $\hat{U}\in \mathbb{R}^{C\times N}$ 中，对C个通道中每个通道应用相同的token-mixing MLP。一个简单的扩展是设计出可分离的MLP，即 $\left \{ MLP_c \right \}_{c=1}^{C}$ 。对于 $\hat{U}$ 沿着通道的每一次滑动， $U[c,:]\in \mathbb{R}^{N}$ ，都由特定的MLPc进行处理。这个扩展直接将MLP-Mixer中的的vanilla token-mixing MLP中的参数数量从 $N^2$ 增加到 $N^2C$ 。但是参数 $N^2C$ 的数量非常大，这使得网络容易过度拟合。因此，它并没有像MLP-Mixer[30]补充材料中所说的那样在识别精度方面带来明显的提高。

相比之下，在我们的循环结构设置中，即使我们采用C个分开的MLP来处理 $\hat{U}$ 的不同通道，参数的数量也只从 $N$ 增加到 $NC$ 。参数的数量仍然处于中等规模。因此，我们的实验表明，使用可分离循环结构的MLP比使用单一循环结构的MLP具有更好的性能。为了进一步减少参数的数量，我们将C个通道分成G个组，每个组包含 $\frac{C}{G}$ 个通道。

对于每个组，我们使用一个特定的MLP。因此，在这个配置中，我们的循环结构的token-mixing MLP的参数数量增加到 $GN$ 。值得注意的是，我们的CCS token-mixing MLP的计算复杂度与组数G无关。因此，在token-mixing中使用G个MLP并不比使用单一MLP带来更多的计算成本。我们在图1中可视化了普通token-mixing MLP和我们的CCS token-mixing MLP的体系结构。我们将我们的CCS与表2中token-mixing 的MLP进行了比较。如表所示，我们的CCS token-mixing MLP比普通的token-mixing 使用更少的参数。在算法1中给出了提出的CCS token-mixing MLP的伪代码。

表2 CCS和普通的token-mixing MLP在参数数量和计算复杂度上的比较。

5 Experiments

6 Conclusion

现有基于MLP的视觉网络中的toke-mixing MLP是特定于空间和信道的。空间特定的结构使它对空间平移敏感。同时，通道不可知属性限制了其token-mixing的能力。为了克服这些限制，我们提出了一个循环特定信道(CCS)token-mixing MLP，它是空间不可知的和信道特定的。空间不可知的属性使我们的CCS token-mixing MLP对空间转换更稳定，而通道特定的属性使它能够编码更丰富的视觉模式。同时，利用循环结构，显著减少了token-mixng MLP层的参数数量。在ImageNet1K数据集上的实验表明，将现有的token-mixing MLP替换为我们的CCS token-mixing MLP，在参数较少的情况下获得了较高的识别精度。