CVPR2022 oral | MetaFormer才是探索Transformer的源泉,衍生PoolFormer速度喜人（一）

1说在前面

为什么学MetaFomer这个模型？

小编其实一直在寻找一个精度与速度都很nice的变形金刚，当我在看李沐大神的Transformer课程的时候，听到介绍MetaFormer，我们都知道范式革命者和提出者都是伟大的，MetaFormer毫无意外就是范式的提出者，其使用pool操作代替了非常耗费算力的MSA，你敢想象吗？作者自己都称之为'Embarrassingly Simple'，这多么骚气，不用想我们都知道pool没有参数学习，计算量是真的小，还说什么呢！肝吧！

下面我们就开始吧，不要错过这个文章，MetaFormer也是CVPR2022的Oral作品，值得反复学习的。文末小编也针对ResNet与MetaFormer的推理速度进行了对比。

2简介

Transformer已经在计算机视觉领域的中获得了极大的影响和成功。由于ViT的开创性工作，使Transformer适应图像分类任务，同时许多后续模型也被开发出来，并在各种计算机视觉任务中取得良好的性能。

图1(a)

如图1(a)所示，通过不指定token mixer，将MetaFormer作为从Transformer中抽象出来的通用架构。当使用注意力/空间MLP作为token mixer时，MetaFormer被实例化为Transformer/MLP-like的模型。

作者认为，Transformer/MLP-like模型的性能优势主要来自通用架构MetaFormer，而不是配备的特定的token mixer。

如图1(a)所示，Transformer编码器由2个组件组成：

• 一个组件是在token之间混合信息的注意力模块，将其命名为token mixer。
• 另一个组件包含残差的模块，如通道mlp和残差连接。

通过将注意力模块视为一个特定的token mixer，作者进一步将整个Transformer抽象为一个通用架构MetaFormer，其中没有指定token mixer。

Transformer的演变长期以来一直归功于基于注意力的token mixer。基于这一普遍的观念，许多注意力模块的变体已经被开发出来，以改善Vision Transformer。然而，最近的一项工作使用Spatial MLPs完全取代了基于注意力模块的作为token mixer，并发现衍生的MLP-like模型可以很容易地在图像分类基准上获得具有竞争力的性能。后续的工作通过数据高效的训练和特定的MLP模块设计，进一步改进了MLP-like模型，逐渐缩小了与ViT的性能差距，并挑战了基于注意力的token mixer的主导地位。

最近的一些方法探索了MetaFormer体系结构中其他类型的token mixer，并表现出了令人鼓舞的性能。例如，FNet用傅里叶变换代替了注意力，并且仍然达到了标准Transformer的97%左右的精度。综合所有这些结果，似乎只要一个模型采用MetaFormer作为一般的体系结构，就可以获得很好的结果。因此，作者假设，与特定的token mixer相比，MetaFormer对于模型实现竞争性能更为重要。

为了验证这一假设，作者应用了一个非常简单的非参数算符，池化，作为token mixer来进行最基本的token混合。如图1(b)，PoolFormer实现了具有竞争力的性能，甚至性能始终优于Transformer和MLP-Like模型，包括DeiT和ResMLP。

图1(b)

所示更具体地说，Poole-m36在ImageNet-1k分类基准上达到了82.1%的Top-1精度，分别超过DeiT/MLP-B24B 0.3%/1.1%，同时参数减少了35%/52%，mac减少了48%/60%。

这些结果表明，MetaFormer即使使用一个简单的token mixer，仍然可以提供良好的性能。因此，作者认为，token mixer值得未来进行更多的研究。

本文的贡献主要有2个方面：

• 首先，将Transformer抽象为一个通用的MetaFormer，并通过经验证明了Transformer/MLP-Like模型的成功很大程度上归因于MetaFormer结构。具体地说，通过只使用一个简单的非参数池化算子作为一个极弱的token mixer，建立了一个简单的模型，发现它仍然可以获得具有很高竞争力的性能。作者希望这个发现能激发更多的研究，致力于改进MetaFormer，而不是专注于token mixer模块。
  
• 其次，对图像分类、目标检测、实例分割和语义分割等多个视觉任务上的PoolFormer进行了评估，发现其与精心设计token mixer的SOTA模型相比具有良好的性能。PoolFormer可以很容易地作为未来MetaFormer架构设计的一个良好的起始Baseline。
• 
  

3相关工作

Transformer最初是在Attention is all you need这篇文章中提出的，主要用于翻译任务，然后迅速在各种NLP任务中流行起来。在语言预训练任务中，Transformer在大规模的无标记文本语料库上进行训练，实现了惊人的性能。受Transformer在NLP中的成功的启发，许多研究者将注意力机制和Transformer应用到视觉任务中。值得注意的是，Chen等人引入了iGPT，其中训练Transformer来自动回归预测图像上的像素，以进行自监督学习。

Dosovitskiy等人提出了以Hard Patches Embedding作为输入的Vision Transformer(ViT)。ViT的实验结果表明，在有监督图像分类任务中，在一个大型数据集(JFT数据集)上预训练的ViT可以获得优异的性能。

DeiT和T2T-ViT进一步证明，仅在ImageNet-1K（∼130万张图像）上进行预训练的ViT可以也获得良好的性能。

现在有许多工作都集中在通过移动窗口、相对位置编码、细化注意力图或合并卷积等来改进Transformer的token mixer方法上。除了类似注意力的token mixer外，仅仅采用mlp作为token mixer仍然可以实现具有竞争力的性能。这一发现挑战了基于注意力的token mixer的主导地位，并在研究界引发了关于哪种token mixer是更好的的激烈讨论。

然而，这项工作的目标既不是参与这场辩论，也不是设计新的复杂的token mixer来实现新的技术状态。相反，本文研究了一个基本的问题：什么才是真正给Transformer及其变体带来成功的原因？作者的答案是一般的架构MetaFormer。作者简单地利用池化操作作为基本的token mixer来探测MetaFormer的能力。

与此同时，一些工作也有助于回答同样的问题。Dong等人证明了在没有残差连接或MLPs的情况下，输出以双指数收敛于秩1矩阵。Raghu等人比较了ViT和CNN之间的特征差异，发现自注意力使全局信息很快聚集，而残差连接可以很快地从lower到higher层继续宁特征传播。不幸的是，它们既没有将Transformer抽象为一般架构，也没有明确解释Transformer性能优秀的原因。

4本文方法

4.1 MetaFormer

首先提出这项工作的核心概念MetaFormer。如图1所示，从Transformer中抽象出来，MetaFormer是一种通用架构，其中没有指定token mixer，而其他组件与Transformer保持相同。输入I首先通过输入Patch Embedding进行处理，比如ViTs的Patch Embedding:

其中，表示序列长度为N、Embedding维数为C的Embedding token。

然后，将Embedding token输入重复的MetaFormer blocks，每个Block包含2个residual sub-blocks。

第1个sub-block

具体地说，第1个sub-block主要包含一个在token之间通信信息的token mixer，该sub-block可以表示为：

其中，Norm(·)表示Normalization，如层Layer Normalization或Batch Normalization；

TokenMixer(·)是指一个主要用于token mixer信息的模块。它是通过最近的Vision Transformer或spatial MLP模型中的各种注意力机制实现的。

请注意，token mixer的主要功能是传播token信息，尽管一些token mixer也可以混合通道，比如注意力。

第2个sub-block

第2个sub-block主要由一个具有非线性激活的双层MLP组成：

其中和是MLP展开比的可学习参数；是一个非线性激活函数，如GELU或ReLU。

MetaFormer的实例化

MetaFormer描述了一种通用的体系结构，通过指定token mixer的具体设计，可以立即获得不同的模型。如图1(a)所示，如果token mixer被指定为注意力或Spatial MLP，那么MetaFormer分别成为Transformer或MLP-Like模型。

4.2 PoolFormer

从Transformer的引入开始，许多工作都非常重视注意力，并专注于各种基于注意力的token mixer组件的设计。相比之下，这些作品很少关注一般的架构，即MetaFormer。

在这项工作中，作者认为这种MetaFormer的通用架构主要有助于最近的Transformer和MLP-Like的模型的成功。为了演示它，作者故意使用了一个简单的池化操作作为token mixer。这个操作符没有可学习的参数，它只是使每个token平均聚合其附近的token特征。

由于这项工作是针对视觉任务的，假设输入是通道优先的数据格式，即。池化运算符可以表示为：

其中，K是pooling size。由于MetaFormer blocks已经有了一个残差连接，因此在上式中删除了输入本身相加的做法。

算法1显示了池化的Pytorch代码：

对应的Pytorch实现如下：

class Pooling(nn.Module):
    def __init__(self, pool_size=3):
        super().__init__()
        self.pool = nn.AvgPool2d(pool_size, stride=1, padding=pool_size // 2, count_include_pad=False)
    def forward(self, x):
        return self.pool(x) - x

众所周知，自注意力和Spatial MLP的计算复杂度与要混合的token的数量呈二次关系。更糟糕的是，Spatial MLP在处理更长的序列时带来了更多的参数。因此，自注意力和Spatial MLP通常只能处理数百个token。相比之下，pool仅仅需要线性的计算复杂度同时还没有任何可学习的参数。

图2

因此，通过采用类似于传统CNNs和最近的分层Transformer的分层结构来利用pool。图2显示了PoolFormer的整体框架。

具体来说，前者有4个阶段，分别有、、和 的token，其中H和W代表输入图像的宽度和高度。

Embedding size分为2组：

1. Small-size model 4个阶段的Embedding size分别为64、128、320、512；
2. Medium-sized model 4个阶段的Embedding size为96、192、384、768。

假设总共有L个PoolFormer blocks，则Stage-1、2、3、4将分别包含L/6、L/6、L/2和L/6个PoolFormer blocks。MLP膨胀比设为4。

PoolFormer blocks的Pytorch实现如下：

class PoolFormerBlock(nn.Module):
    def __init__(self, dim, pool_size=3, mlp_ratio=4., act_layer=nn.GELU, norm_layer=GroupNorm, drop=0., drop_path=0.,
                 use_layer_scale=True, layer_scale_init_value=1e-5):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.token_mixer = Pooling(pool_size=pool_size)
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.use_layer_scale = use_layer_scale
        if use_layer_scale:
            self.layer_scale_1 = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
            self.layer_scale_2 = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
    def forward(self, x):
        if self.use_layer_scale:
            x = x + self.drop_path(self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) * self.token_mixer(self.norm1(x)))
            x = x + self.drop_path(self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * self.mlp(self.norm2(x)))
        else:
            x = x + self.drop_path(self.token_mixer(self.norm1(x)))
            x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x
def basic_blocks(dim, index, layers, pool_size=3, mlp_ratio=4., act_layer=nn.GELU, norm_layer=GroupNorm, drop_rate=.0,
                 drop_path_rate=0., use_layer_scale=True, layer_scale_init_value=1e-5):
    blocks = []
    for block_idx in range(layers[index]):
        block_dpr = drop_path_rate * (block_idx + sum(layers[:index])) / (sum(layers) - 1)
        blocks.append(PoolFormerBlock(dim, pool_size=pool_size, mlp_ratio=mlp_ratio, act_layer=act_layer,
                                      norm_layer=norm_layer, drop=drop_rate, drop_path=block_dpr,
                                      use_layer_scale=use_layer_scale, layer_scale_init_value=layer_scale_init_value,))
    blocks = nn.Sequential(*blocks)
    return blocks

根据上述简单的模型比例规则，得到了5种不同的PoolFormer的模型尺寸，其超参数如表所示:

不同配置的模型Pytorch实现如下：

def poolformer_s12(pretrained=False, **kwargs):
    """
    PoolFormer-S12 model, Params: 12M
    --layers: [x,x,x,x], numbers of layers for the four stages
    --embed_dims, --mlp_ratios:
        embedding dims and mlp ratios for the four stages
    --downsamples: flags to apply downsampling or not in four blocks
    """
    layers = [2, 2, 6, 2]
    embed_dims = [64, 128, 320, 512]
    mlp_ratios = [4, 4, 4, 4]
    downsamples = [True, True, True, True]
    model = PoolFormer(layers, embed_dims=embed_dims, mlp_ratios=mlp_ratios, downsamples=downsamples, **kwargs)
    model.default_cfg = default_cfgs['poolformer_s']
    return model
def poolformer_s24(pretrained=False, **kwargs):
    """
    PoolFormer-S24 model, Params: 21M
    """
    layers = [4, 4, 12, 4]
    embed_dims = [64, 128, 320, 512]
    mlp_ratios = [4, 4, 4, 4]
    downsamples = [True, True, True, True]
    model = PoolFormer(layers, embed_dims=embed_dims, mlp_ratios=mlp_ratios, downsamples=downsamples, **kwargs)
    model.default_cfg = default_cfgs['poolformer_s']
    return model
def poolformer_s36(pretrained=False, **kwargs):
    """
    PoolFormer-S36 model, Params: 31M
    """
    layers = [6, 6, 18, 6]
    embed_dims = [64, 128, 320, 512]
    mlp_ratios = [4, 4, 4, 4]
    downsamples = [True, True, True, True]
    model = PoolFormer(layers, embed_dims=embed_dims, mlp_ratios=mlp_ratios, downsamples=downsamples, layer_scale_init_value=1e-6, **kwargs)
    model.default_cfg = default_cfgs['poolformer_s']
    return model
def poolformer_m36(pretrained=False, **kwargs):
    """
    PoolFormer-M36 model, Params: 56M
    """
    layers = [6, 6, 18, 6]
    embed_dims = [96, 192, 384, 768]
    mlp_ratios = [4, 4, 4, 4]
    downsamples = [True, True, True, True]
    model = PoolFormer(layers, embed_dims=embed_dims, mlp_ratios=mlp_ratios, downsamples=downsamples, layer_scale_init_value=1e-6, **kwargs)
    model.default_cfg = default_cfgs['poolformer_m']
    return model
def poolformer_m48(pretrained=False, **kwargs):
    """
    PoolFormer-M48 model, Params: 73M
    """
    layers = [8, 8, 24, 8]
    embed_dims = [96, 192, 384, 768]
    mlp_ratios = [4, 4, 4, 4]
    downsamples = [True, True, True, True]
    model = PoolFormer(layers, embed_dims=embed_dims, mlp_ratios=mlp_ratios, downsamples=downsamples, layer_scale_init_value=1e-6, **kwargs)
    model.default_cfg = default_cfgs['poolformer_m']
    return model