颜水成发了个「简单到令人尴尬」的视觉模型,证明Transformer威力源自其整体架构

Transformer 做视觉取得巨大成功,各大变体频频刷榜,其中谁是最强?

早期人们认为是其中的注意力机制贡献最大,对注意力模块做了很多改进。

后续研究又发现不用注意力换成Spatial MLP效果也很好,甚至使用傅立叶变换模块也能保留 97% 的性能。

争议之下,颜水成团队的最新论文给出一个不同观点

其实这些具体模块并不重要,Transformer 的成功来自其整体架构

他们把 Transformer 中的注意力模块替换成了简单的空间池化算子,新模型命名为PoolFormer

这里原文的说法很有意思,"简单到让人尴尬" ……

测试结果上,PoolFormer 在 ImageNet-1K 上获得了82.1%的 top-1 精度。

(PyTorch 版代码已随论文一起发布在 GitHub 上,地址可在这篇推文末尾处获取。)

同等参数规模下,简单池化模型超过了一些经过调优的使用注意力(如 DeiT)或 MLP 模块(如 ResMLP)的模型。

这个结果让一些围观的 CVer 直接惊掉下巴:

太好奇了,模型简单到什么样才能令人尴尬?

PoolFormer

整体结构与其他模型类似,PoolFormer 只是把 token mixer 部分换了一下。

因为主要验证视觉任务,所以假设输入数据的格式为通道优先,池化算子描述如下:

PyTorch 风格的伪代码大概是这样:

池化算子的复杂度比自注意力和 Spatial MLP 要小,与要处理的序列长度呈线性关系。

其中也没有可学习的参数,所以可以采用类似传统 CNN 的分阶段方法来充分发挥性能,这次的模型分了4 个阶段

假设总共有 L 个 PoolFormer 块,那么 4 个阶段分配成 L/6、L/6、L/2、L/6 个。

每个阶段的具体参数如下:

PoolFormer 基本情况介绍完毕,下面开始与其他模型做性能对比。

首先是图像分类任务,对比模型分为三类:

CNN 模型 ResNet 和 RegNetY

使用注意力模块的 ViT、DeiT 和 PVT

使用 Spatial MLP 的 MLP-Mixer、ResMLP、Swin-Mixer 和 gMLP

在 ImageNet-1K 上,无论是按累计乘加操作数(MACs)还是按参数规模为标准,PoolFormer 性能都超过了同等规模的其他模型。

目标检测实例分割任务上用了 COCO 数据集,两项任务中 PoolFormer 都以更少的参数取得比 ResNet 更高的性能。

目标检测

实例分割

最后是 ADE20K语义分割任务,PoolFormer 的表现也超过了 ResNet、ResNeXt 和 PVT。

消融实验

上面可以看出,几大视觉任务上 PoolFormer 都取得了有竞争力的成绩。

不过这还不足以支撑这篇论文开头提出的那个观点。

到底是整体架构重要?还是说 PoolFormer 中的池化模块刚好是一种简单却有效的 Token Mixer?

团队的验证方法是把池化模块直接替换成恒等映射(Identity Mapping)。

结果令人惊讶,替换后在 ImageNet-1K 上也保留了 74.3% 的 Top-1 精度。

在此基础上无论是改变池化核尺寸归一化方法激活函数影响都不大。

最重要的是,在 4 个阶段中把注意力和空间全连接层等机制混合起来用性能也不会下降。

其中特别观察到,前两阶段用池化后两阶段用注意力这种组合表现突出。

这样的配置下稍微增加一下规模精度就可达到 81%,作为对比的 ResMLP-B24 模型达到相同性能需要 7 倍的参数规模和 8.5 倍的累计乘加操作。

最终,消融实验结果说明 Transformer 中具体到 token mixer 这个部分,具体用了哪种方法并不关键。

不增加模型规模的情况下,网络的整体结构才是对性能提升最重要的。

这样的整体结构被团队提炼出来,命名为MetaFormer

NLP 上还会成立吗?

这项研究由颜水成领导的 Sea AI Lab 和来自新加坡国立大学的成员共同完成。

颜水成

论文的最后,团队表示下一步研究方向是在更多场景下继续验证,如自监督学习和迁移学习。

除了视觉任务,也要看看在 NLP 任务上结论是否也成立。

另外发这篇论文还有一个目的:

呼吁大家把研究的重点放在优化模型的基础结构,而不是在具体模块上花太多精力去打磨。

论文地址:

https://arxiv.org/abs/2111.11418

GitHub 仓库:

https://github.com/sail-sg/poolformer