点击下方卡片,关注“CVer”公众号
AI/CV重磅干货,第一时间送达
本文转载自:AIWalker
MLP-Mixer: An all-MLP Architecture for Vision
paper: https://arxiv.org/abs/2105.01601
code: https://github.com/google-research/vision_transformer(暂未开源)
本文是谷歌大脑的研究员(原ViT团队)在网络架构设计方面挖的新坑:MLP-Mixer。无需卷积、注意力机制,MLP-Mixer仅需MLP即可达到与CNN、Transformer相媲美的性能。比如,在JFT-300M数据集预训练+ImageNet微调后,所提Mixer-H/14取得87.94%的top1精度。尽管所提方法性能并未达到最优,但本文的目的并不在于达成SOTA结果,而在于表明:简简单单的MLP模型即可取得与当前最佳CNN、注意力模型相当的性能。
Abstract
CNN已成为计算机视觉的首选方案,近来兴起的基于注意力的网络(比如Vision Transformer)也逐渐变的流行。这篇文章表明:尽管卷积与注意力有助于取得比较的好的性能,但其并非必要。
本文提出了一种新颖的架构:MLP-Mixer
,它仅仅依赖于多层感知器(multi-layer perceptrons, MLP)。MLP-Mixer包含两种类型的层:(1) 独立作用于图像块的MLP,即用于混合每个位置的特性;(2) 跨图像块的MLP,即用于混合空域信息。当在大数据集上训练,或者采用先进正则技术训练后,MLP-Mixer在图像分类基准数据集上取得了极具竞争力的性能。我们希望该工作能激发出超越CNN与Transformer之外的进一步研究。
Introduction
尽管CNN已成为计算机视觉领域的首选方案,近来兴起的Vision Transformer(ViT)打开了另一扇门并取得了优异的性能。本文提出了一种区别于CNN和Transformer的新架构MLP-Mixer,一种极具竞争力且概念与技术简单的新方案,它无需卷积与自注意力。相反,MLP-Mixer仅仅依赖于在空域或者特征通道上重复实施的多层感知器;Mixer仅依赖于基础矩阵乘操作、数据排布变换(比如reshape、transposition)以及非线性层。
Figure1给出了Mixer的宏观建构示意图,它以一系列图像块的线性投影(其形状为patches x channels
)作为输入。Mixer采用了两种类型的MLP层(注:这两种类型的层交替执行以促进两个维度间的信息交互):
channel-mixing
MLP:用于不同通道前通讯,每个token独立处理,即采用每一行作为输入;token-mixing
MLP:用于不同空域位置通讯,每个通道图例处理,即采用每一列作为输入。
![](https://i-blog.csdnimg.cn/blog_migrate/f2c31e306b7da8061fb6319abbae7495.jpeg)
在极端情况下,本文所提架构可视作一种特殊CNN,它采用
卷积进行channel mixing
,全感受野、参数共享的的单通道深度卷积进行token mixing
。
尽管如此简单,Mixer取得极具竞争力的结果,当在大数据(如100M数据量)上预训练时,它可以取得接近SOTA的性能:87.94%top1精度。当在适当尺度数据+先进正则技术进行训练时,所提方法同样取得极强性能。
Mixer Architecture
当前深度视觉架构采用以下三种方式进行特征混合:(i) 在给定空域位置,即channel-mixing;(ii) 不同空域位置,即token-mixing;或者上述两种组合。在CNN中,ii采用 卷积或者池化进行,更深的层具有更大的感受野;与此同时,i则采用 卷积进行;更大的核则同时进进行i与ii。在Vision Transformer以及其他注意力架构中,自注意力层同时执行i与ii;而MLP则孩子能够i。Mixer背后的思想在于:将上述两种特征混合进行显示分离且均通过MLP进行实现。
前面的Figure1给出了Mixer的架构示意图,Mixer以序列长度为S的非重叠图像块作为输入,每个图像块将投影到期望的隐层维度C。这将得到一个二维实值输入表 。如果原始输入图像分辨率为 ,每个块的分辨率为 ,就那么序列长度 ,所有块采用相同的投影矩阵进行线性投影。Mixer包含多个等尺寸的层,每个层包含两种MLP模块:
token-mixing MLP:它作用于X的列,从 映射到 ;
Channel-mixing MLP:它作用于X的行,从 映射到 。
每个MLP模块包含两个全连接层与一个非线性层,它们可以定义如下:
其中 表示GELU激活函数, 分别表示token-mixing与channel-mixing MLP中隐层宽度。需要注意的是, 的选择独立于输入图像块的数量。因此,网络的计算复杂度与输入块的数量成线性关系,这不同于ViT的二次关系。由于, 独立于块尺寸,因此,整体计算量与图像的像素数成线性关系,类似于CNN。
正如上述所提到的,相同的channel(token)-mixing MLP作用于X的每一行(列)。对MLP的参数进行绑定就是一种很自然的选择,它可以提供类似卷积特征的位置不变形。然而,在CNN中进行跨通道参数绑定并不常见。比如,CNN的分离卷积对每个通道采用不同的卷积核,这与本文MLP中的处理机制(所有通道采用相同的核)不相同。这种参数绑定可以避免架构随隐层维度C、序列长度S提升而增长过快,进而导致了显著的显存节省。令人惊讶的是:这种参数绑定机制并不会影响性能。
Mixer中的每一层(除了初始块投影层)采用相同尺寸的输入,这种“各向同性”设计类似于Transformer和RNN中定宽;这与CNN中金字塔结构(越深的层具有更低的分辨率、更多的通道数)不同。
除了MLP外,Mixer还采用其他标准架构成分:跳过连接、LayerNorm。此外,不同于ViT,Mixer并没有采用position embedding
,这是因为token-mixing
MLP对于输入的顺序极为敏感。最后,Mixer采用了标注分类头,即全局均值池化+线性分类器。基于Flax与jax的参考实现code如下:
![](https://i-blog.csdnimg.cn/blog_migrate/e0d30fcf6f0a4d7e0cf057fdc48f4d08.jpeg)
Experiments
基于中等与大尺度数据的预训练,我们在不同下游分类任务上对所提MLP-Mixer的性能进行了评估。我们主要对以下三个问题比较感兴趣并进行重点分析。
在下游任务上的精度
总计预训练计算量,这对于在上游数据上从头开始训练模型非常重要;
推理耗时,这对于实际应用非常重要。
注:本文的目的不在于达成SOTA结果,而在于表明:简简单单的MLP模型即可取得与当前最佳CNN、注意力模型相当的性能。
Downstream Taks 我们采用了多个主流下游任务,比如ILSVRC2012、清洗后ReaLLabels、CIFAR10/100、Oxford-IIIT Pets、Oxford Flowers-102等。我们同事还在VTAB-1k数据集上进行了评估。
Pre-training data 我们参考标准的迁移学习步骤:预训练+微调。我们在两个公开数据(ILSVRC2012与ImageNet-21k)上进行预训练。为评估在更大尺度上的性能,我们还在JFT-300M数据集上进行了训练。
Pre-training details 训练超参:优化器为Adam,batch=4096,weight decay以及梯度裁剪、线性学习率warmup,输入分辨率均为224。数据增广方面包含RandAugment、mixup、dropout、stochastic depth等等。
Fine-tuning details 微调超参:优化器为momentum-SGD,batch=512,梯度裁剪,cosine学习率机制+线性warmup,并未采用weight decay。微调阶段采用了更高的分辨率。
Metrics 我们从模型性能、计算复杂度之间的均衡性能。对于计算复杂度,我们计算了两个测度:(1) TPU-v3上的合计预训练耗时,它包含理论FLOPs、硬件上的计算效率以及数据加载效率;(2) TPU-v3上的吞吐量。
![](https://i-blog.csdnimg.cn/blog_migrate/923fd6bf9227f7aac608abd21ffe7277.jpeg)
Models 上表列出了Mixer模型的不同配置以对标CNN与注意力模型。此外,在对标模型方面,ViTs模型包含ViT、HaloNet、BiT等等;CNN模型包含NFNet、MPL、ALIGN等。
Main results
![](https://i-blog.csdnimg.cn/blog_migrate/3381a1004de70bdbb2fbb7866ea381cb.jpeg)
上表给出了最大Mixer模型与SOTA模型在ImageNet、ReaL数据集上的性能对比,从中可以看到:
当在ImageNet-21k+额外正则技术预训练后,Mixer在ImageNet数据集取得非常强的性能:84.15%top1,比其他模型稍弱。此时,类似于ViT,正则技术很有必要,有助于Mixer避免过拟合。
当提升上游数据集尺寸后,Mixer的性能可以进一步提升。具体来说,Mixer-H/14取得了87.94%的top1精度,比BiT-ResNet152x4高0.5%,比ViT-H/14低0.5%。值得一提的是,Mixer-H/14的推理速度要比ViT-H/14快2.5倍,比BiT快2倍。
The role of the model scale
![](https://i-blog.csdnimg.cn/blog_migrate/d0c878a499debb90d0de1eee27877dea.jpeg)
上表对比了不同模型大小、不同分辨率输入时模型的性能对比,从中可以看到:
当在ImageNet上从头开始训练时,Mixer-B/16取得了一个比较合理的精度:76.44%,。这要比ViT-B/16低3%。也就是说,在ImageNet上从头开始训练时,Mixer-B/16出现了过拟合问题。
随预训练数据集变大,Mixer的性能稳步提升。值得一提的是,JFT-300M数据上预训练、224尺寸微调的Mixer-H/14取得了86.32%的精度,仅以0.3%弱于ViT-H/14,但推理速度快2.2倍。
全文到此结束,更多消融实验与分析建议查看原文。
论文PDF和代码下载
后台回复:MLP,即可下载上述论文PDF
后台回复:CVPR2021,即可下载CVPR 2021论文和代码开源的论文合集
后台回复:Transformer综述,即可下载最新的两篇Transformer综述PDF
重磅!图像分割交流群成立
扫码添加CVer助手,可申请加入CVer-图像分割 微信交流群,方向已涵盖:目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch和TensorFlow等群。
一定要备注:研究方向+地点+学校/公司+昵称(如图像分割+上海+上交+卡卡),根据格式备注,可更快被通过且邀请进群
▲长按加小助手微信,进交流群
▲点击上方卡片,关注CVer公众号
整理不易,请给CVer点赞和在看