MetaFormer/PoolFormer学习笔记及代码

MetaFormer/PoolFormer学习笔记及代码

MetaFormer Is Actually What Y ou Need for Vision
code：https://github.com/sail-sg/poolformer

Abstract

变形金刚在计算机视觉任务中显示出巨大的潜力。人们普遍认为，他们基于注意力的token mixer模块对他们的能力贡献最大。然而，最近的研究表明，Transformers中基于注意力的模块可以被空间MLP所取代，得到的模型仍然表现良好。基于这一观察，我们假设变压器的一般架构，而不是特定的令牌混频器模块，对模型的性能更为重要。为了验证这一点，我们故意用令人尴尬的简单空间池算子替换Transformers中的注意力模块，以仅进行基本令牌混合。令人惊讶的是，我们观察到衍生模型称为PoolFormer，在多个计算机视觉任务上实现有竞争力的性能。例如，在ImageNet-1K上，PoolF-former达到82.1%的top-1精度，在参数减少35%/52%和MAC减少50%/62%的情况下，比调整良好的视觉变换器/MLP样基线DeiT-B/ResMLP-B24的精度高出0.3%/1.1%。PoolFormer的有效性验证了我们的假设，并促使我们提出“MetaFormer”的概念，这是一种从Transformers中抽象出来的通用架构，无需指定令牌混合器。基于广泛的实验，我们认为MetaFormr是在视觉任务上实现最近Transformer和MLP类模型的优异结果的关键。这项工作需要更多致力于改进MetaFormer的未来研究，而不是专注于令牌混频器模块。此外，我们提出的PoolF-ormer可以作为未来MetaFormer架构设计的起始基线。

1. Introduction

Transformers在计算机视觉领域获得了很大的兴趣和成功[3、8、44、55]。由于视觉变换器（ViT）[17]的开创性工作使纯变换器适应图像分类任务，因此开发了许多后续模型，以进一步改进并在各种计算机视觉任务中实现良好的性能[36、53、63]。

如图1（a）所示，变压器编码器由两个组件组成。一个是注意力模块，用于在令牌之间混合信息，我们将其称为令牌混合器。另一个组件包含其余模块，例如信道MLP和剩余连接。通过将注意力模块视为一个特定的令牌混合器，我们进一步将整个转换器抽象为一个通用架构元，其中未指定令牌混合器。如图1（a）所示。

Transformers的成功长期以来都归功于基于注意力的令牌混合器[56]。基于这一共同信念，开发了注意力模块的许多变体[13、22、57、68]，以改进视觉转换器。然而，最近的一项工作[51]将注意力模块完全替换为空间MLP作为令牌混合器，并发现衍生的类似MLP的模型可以在图像分类基准上获得具有竞争力的性能。后续工作[26、35、52]通过高效的数据训练和特定的MLP模块设计，进一步改进了类似MLP的模型，逐渐缩小了与ViT的性能差距，并挑战了注意力作为令牌混合器的主导地位。

最近的一些方法[32、39、40、45]探索了元前体架构中的其他类型的令牌混合器，并证明了令人鼓舞的性能。例如，[32]用傅立叶变换代替注意力，仍然达到普通变压器精度的97%左右。综合所有这些结果，似乎只要模型采用元模型作为通用架构，就可以获得有希望的结果。因此，我们假设与特定令牌混频器相比，MetaFormer对于模型实现竞争性能更为重要。

为了验证这一假设，我们应用了一个非常简单的非参数算子pooling作为令牌混合器，只进行基本的令牌混合。令人惊讶的是，这种被称为PoolFormer的衍生模型实现了具有竞争力的性能，甚至始终优于经过良好调优的Transformer和类似MLP的模型，包括DeiT[53]和ResMLP[52]，如图1（b）所示。更具体地说，PoolFormer-M36在ImageNet-1K分类基准上实现了82.1%的top-1精度，比调整良好的vision Transformer/MLP-like基线DetB/ResMLP-B24高出0.3%/1.1%，参数减少35%/52%，MAC减少50%/62%。这些结果表明，即使使用简单的令牌混合器，MetaFormer仍然可以提供有希望的性能。因此，我们认为，MetaFormer是我们对视觉模型的实际需求，而视觉模型对于实现具有竞争力的性能更为重要，而不是特定的令牌混合器。请注意，这并不意味着令牌混合器无关紧要。MetaFormer仍然有这个抽象组件。这意味着令牌混合器不限于特定类型，例如注意力

本文的贡献有两个方面。首先，我们将Transformer抽象为一个通用的架构元模型，并实证证明Transformer/MLP类模型的成功在很大程度上归功于元模型架构。具体来说，通过仅使用一个简单的非参数算子pooling作为MetaFormer的极弱令牌混合器，我们构建了一个名为PoolFormer的简单模型，发现它仍然可以实现高度竞争的性能。我们希望我们的发现能激发更多致力于改进元化器而不是专注于令牌混频器模块的未来研究。其次，我们在多个视觉任务（包括图像分类[14]、对象检测[34]、实例分割[34]和语义分割[67]）上评估了所提出的PoolFormer，并发现它与使用令牌混合器的复杂设计的SOTA模型相比具有竞争性能。PoolFormer可以很容易地作为未来MetaFormer架构设计的良好起点。

2. Related work

Transformer首先由[56]提出用于翻译任务，然后迅速在各种自然语言处理任务中流行。在语言预训练任务中，Transformer在大规模未标记文本语料库上进行训练，并取得了惊人的性能[2，15]。受Transformers在自然语言处理中的成功启发，许多研究人员将注意力机制和Transformers应用于视觉任务[3、8、44、55]。值得注意的是，陈等人介绍了iGPT[6]，其中训练变压器以自回归预测图像上的像素，用于自监督学习。Dosovitskiy等人提出了以硬补丁嵌入为输入的视觉变换（ViT）[17]。他们表明，在有监督的图像分类任务中，在大型属性数据集（具有3亿图像的JFT数据集）上预训练的ViT可以实现优异的性能。DeiT[53]和T2T ViT[63]进一步证明ViT仅在ImageNet-1K上预训练(∼ 130万张图像）从头开始可以实现良好的性能。已经完成了很多工作关于哪个令牌混合器更好[7，26]。然而，这项工作的目标既不是参与这场辩论，也不是设计新的复杂令牌混合器以达到新的技术水平。相反，我们研究了一个基本问题：是什么真正为变形金刚及其变体的成功负责？我们的答案是通用架构，即MetaFormer。我们只是利用池作为基本的令牌混合器来探索MetaFormer的功能。

同时，一些作品有助于回答同样的问题。Dong等人证明，在没有剩余连接或MLP的情况下，输出以双指数方式收敛到秩一矩阵[16]。Raghu等人[43]比较了ViT和CNN之间的特征差异，发现自注意力可以早期收集全局信息，而剩余连接可以将特征从较低层传播到较高层。Park等人[42]表明，多头自注意力通过平坦损失景观来提高准确性和泛化。不幸的是，他们没有将转换器抽象为一个通用架构，并从通用框架的角度对其进行研究

3. Method

3.1. MetaFormer

我们首先提出了这项工作的核心概念“MetaFormer”。如图1所示，从Transformers[56]中抽象出来，MetaFormer是一种通用架构，其中未指定令牌混合器，而其他组件与Transformers保持相同。输入$I$首先通过输入嵌入进行处理，例如VIT的面片嵌入[17]，

其中$X\in R^{N\times C}$表示序列长度为N、嵌入维数为C的嵌入令牌。

然后，嵌入令牌被馈送到重复的元前体块，每个元前体包括两个剩余子块。具体来说，第一个子块主要包含令牌混合器，用于在令牌之间通信信息，该子块可以表示为

其中，范数（·）表示归一化，例如层归一化[1]或批量归一化[28]；TokenMixer（·）是指主要用于混合令牌信息的模块。在最近的视觉变换模型[17,63,68]中，它通过各种注意机制实现，或者在类似MLP的模型中通过空间MLP实现[51,52]。注意，令牌混合器的主要功能是传播令牌信息，尽管一些令牌混合器也可以混合通道，例如注意力。

第二个子块主要由具有非线性激活的双层MLP组成，

其中$W_1\in R^{C\times rC}$和$W_2\in R^{rC\times C}$是具有MLP膨胀比r的可学习参数；σ（·）是非线性激活函数，如GELU[25]或ReLU[41]。

Instantiations of MetaFormer

MetaFormer描述了一种通用架构，通过指定令牌混合器的具体设计，可以立即获得不同的模型。如图1（a）所示，如果将令牌混合器指定为注意力或空间MLP，则MetaFormer将分别成为Transformer或类似MLP的模型。

3.2. PoolFormer

从Transformers[56]的介绍开始，许多工作都非常重视注意力，并专注于设计各种基于注意力的令牌混频器组件。相比之下，这些作品很少关注一般的架构，即元模型。

在这项工作中，我们认为，这种元前体通用架构主要有助于最近Transformer和MLP类模型的成功。为了演示它，我们故意使用了一个令人尴尬的简单操作符pooling作为令牌混合器。该操作符没有可学习的参数，它只是使每个令牌平均聚合其附近的令牌特征。

由于这项工作的目标是视觉任务，我们假设输入为通道优先数据格式，即$T\in R^{C\times H\times W}$。池算子可以表示为

其中K是池大小。由于MetaFormer块已经具有残余连接，因此在等式（4）中添加输入本身的减法。池的类似Pytork的代码如算法1所示。

众所周知，自注意力和空间MLP的计算复杂度是要混合的令牌数的二次方。更糟糕的是，在处理更长的序列时，空间MLP会带来更多的参数。因此，自注意力和空间MLP通常只能处理数百个令牌。相反，pooling需要与序列长度成线性关系的计算复杂度，而不需要任何可学习的参数。因此，我们通过采用类似于传统CNN[24,31,49]和最近的层次变换器变体[36,57]的层次结构来利用池。图2显示了PoolFormer的总体框架。具体来说，PoolFormer有4个阶段，分别具有$h/4\times w/4、h/8\times w/8、h/16\times w/16和h/32\times w/32$令牌，其中H和W表示输入图像的宽度和高度。嵌入尺寸分为两组：1）小尺寸模型，嵌入尺寸为64、128、320和512，响应四个阶段；2） 嵌入维数为96、192、384和768的中型模型。假设总共有L个池形成器块，阶段1、2、3和4将分别包含L/6、L/6，L/2和L/6池形成器。MLP展开比设置为4。根据上述简单的模型缩放规则，我们获得了池形成器的5种不同模型尺寸，其超参数如表1所示。

4. Experiments

5. Conclusion and future work

在这项工作中，我们将Transformers中的注意力抽象为令牌混合器，将Transformer整体抽象为一般称为元成型器的架构，其中未指定令牌混合器。我们并没有将重点放在特定的令牌混合器上，而是指出，MetaFormer实际上是我们需要的，以确保实现合理的性能。为了验证这一点，我们特意将令牌混合器指定为元形成器的极其简单的池。研究发现，派生的PoolFormer模型可以在不同的视觉任务上实现有竞争力的性能，这很好地支持“元模型实际上是视觉所需要的”。

未来，我们将在更多不同的学习环境下，例如自监督学习和转移学习，进一步评估PoolFormer。此外，很有意思的是，看看PoolFormer是否仍然适用于自然语言处理任务，以进一步支持自然语言处理领域中“MetaFormer实际上是您所需要的”的说法。我们希望这项工作能够激发更多致力于改进基本架构元模型的未来研究，而不是过多关注令牌混频器模块。

附录

改进的Layer Normalization

PoolFormer block