当前,卷积神经网络(CNN)和基于自注意力的网络(如近来大火的 ViT)是计算机视觉领域的主流选择,但研究人员没有停止探索视觉网络架构的脚步。近日,来自谷歌大脑的研究团队(原 ViT 团队)提出了一种舍弃卷积和自注意力且完全使用多层感知机(MLP)的视觉网络架构,在设计上非常简单,并且在 ImageNet 数据集上实现了媲美 CNN 和 ViT 的性能表现。
计算机视觉的发展史证明,规模更大的数据集加上更强的计算能力往往能够促成范式转变。虽然卷积神经网络已经成为计算机视觉领域的标准,但最近一段时间,基于自注意力层的替代方法 Vision Transformer(ViT)实现新的 SOTA 性能。从技术上讲,ViT 模型延续了长久以来去除模型中手工构建特征和归纳偏置的趋势,并进一步依赖基于原始数据的学习。
近日,原 ViT 团队提出了一种不使用卷积或自注意力的 MLP-Mixer 架构(简称 Mixer),这是一种颇具竞争力并且在概念和技术上都非常简单的替代方案。
Mixer 架构完全基于在空间位置或特征通道重复利用的多层感知机(MLP),并且仅依赖于基础矩阵乘法运算、数据布局变换(如 reshape 和 transposition)和非线性层。
- 论文地址:https://arxiv.org/pdf/2105.01601.pdf
- 项目地址:https://github.com/google-research/vision_transformer/tree/linen
结果表明,虽然 Mixer 架构很简单,但取得了极具竞争力的结果。当在大型数据集(约 1 亿张图像)上进行预训练时,该架构在准确率 / 成本权衡方面能够媲美 CNN 和 ViT,实现了接近 SOTA 的性能,在 ImageNet 数据集上取得了 87.94% 的 top1 准确率。
对于该研究提出的 Mixer 架构,特斯拉 AI 高级总监 Andrej Karpathy 认为:「很好!1×1 卷积通常利用深度卷积实现堆叠或交替,但在这里,通道或空间混合得到简化或者实现完全对称。」
另一用户表示:「CV 领域网络架构的演变从 MLP 到 CNN 到 Transformer 再回到 MLP,真是太有意思了。」
不过,谷歌 DeepMind 首席科学家 Oriol Vinyals 也提出了质疑,他认为:「per-patch 全连接,那不就是卷积吗」
那就先来看一下MLP-Mixer这个新框架吧,它不使用卷积或自注意力机制。相反,Mixer体系架构完全基于在空间位置或特征通道上重复应用的多层感知器(MLP),它只依赖基础的矩阵乘法操作、数据排布变换(比如reshape、transposition)以及非线性层。
下图展示了MLP-Mixer的整体结构:
首先,它的输入是一系列图像块的线性投影(其形状为patches x channels),其次,Mixer使用两种类型的MLP层:
1、通道混合MLP(channel-mixing MLPs ):用于不同通道之间进行通信,允许对每个token独立操作,即采用每一行作为输入。
2、token混合MLP(The token-mixing MLPs ):用于不同空间位置(token)之间的通信;允许在每个通道上独立操作,即采用每一列作为输入。
以上两种类型的MLP层交替执行以实现两个输入维度的交互。
在极端情况下,MLP-Mixer架构可以看作一个特殊的CNN,它使用1×1通道混合的卷积,全感受域的单通道深度卷积以及token混合的参数共享。典型的CNN不是混合器的特例,卷积也比MLPs中的普通矩阵乘法更复杂(它需要额外的成本来减少矩阵乘法或专门实现)不过,尽管它很简单,MLP-Mixer还是取得了很不错的结果。
当对大型数据集进行预训练时(大约100万张图片),它达到了之前CNNs和Transformers在ImageNet上的最佳性能:87.94%的 top-1 验证准确率。当对1-10万张图片大小的数据集进行预训练时,结合现代正则化技术( regularization techniques),Mixer同样取得了强大的性能。
Mixer 混合器架构
一般来讲,当今深度视觉体系结构采用三种方式进行特征混合:
(i)在给定的空间位置;
(ii)不同的空间位置之间;
(iii)将上述两种方式组合。
在CNNs中,(ii)是采用N× N进行卷积和池化,其中N>1;(i)采用1×1卷积;较大的核则同时执行(i)和(ii)。通常更深层次的神经元有更大的感受野。
在Transformer和其他注意力架构中,自注意力层允许同时执行(i)和(ii),而MLP只执行(i)。Mixer架构背后的思想是:通过MLP实现每个通道混合操作(i)和 token混合操作(ii)的显著分离。
在上图体系架构中,Mixer将序列长度为S的非重叠的图像块作为输入,每个图像块都投影到所需的隐层维度C,并产生一个二维实值输入X∈ RS×C。如果原始图像的分辨率为(H x W),每个图像块的分辨率为(P x P),那么图像块的数量则为S=HW/P2。所有的块都采用相同的投影矩阵进行线性投影。
Mixer由等尺寸的多层组成,每层有两个MLP块。第一个是token mixing MLP块:它作用于X的列,从RS映射到R S,可在所有列中共享。第二个是Channel-mixing MLP块:它作用于X的行,从Rc映射到 R C,可在所有行中共享。每个MLP块包含两个全连接层和一个独立于输入的非线性层。其基本方程如下:
图中,Ds Dc分别代表token-mixing与channel-mixing MLP中隐层宽度。由于Ds的选择独立于输入图像块的数量,因此,网络的计算复杂度与输入块的数量成线性关系;此外,Dc独立于块尺寸,整体计算量与图像的像素数成线性关系,这类似于CNN。
如上文所说,相同的通道混合MLP(或令牌混合MLP)应用于X的每一行和列,在每一层内绑定通道混合MLP的参数都是一种自然选择,它提供了位置不变性,这是卷积的一个显著特征。
不过,跨通道绑定参数的情况在CNN中并不常见。例如CNN中可分离卷积,将不同的卷积核独立应用于每个通道。而Mixer中的token 混合MLP可以对所有通道共享相同的核(即获得完全感受野)。通常来讲,当增加隐层维数C或序列长度S时,这种参数绑定可以避免体系架构增长过快,并且节省内存。令人没想到的是,这种绑定机制并没有影响性能。
Mixer中的每个层(除了初始块投影层)接收相同大小的输入。这种“各向同性”设计最类似于使用固定宽度的Transformer和RNN。这与大多数CNN不同,CNN具有金字塔结构:越深的层具有更低的分辨率,更多的通道。需要注意的是,以上是典型的设计,除此之外也存在其他组合,例如各向同性网状结构和金字塔状VIT。除了MLP层之外,Mixer还使用了其他标准的体系结构组件:Skip 连接和层规范化。
此外,与ViTs不同,Mixer不使用位置嵌入,因为token混合mlp对输入token的顺序敏感,因此可以学习表示位置。最后,Mixer使用一个标准的分类head和一个线性分类器。
更多实验结果和代码
该研究用实验对 MLP-Mixer 模型的性能进行了评估。其中,模型在中大规模数据集上进行预训练,采用一系列中小型下游分类任务,并对以下三个问题进行重点研究:
- 在下游任务上的准确率;
- 预训练的总计算成本,这对于在上游数据集上从头开始训练模型非常重要;
- 推断时的吞吐量,这在实际应用中非常重要。
该研究的实验目的不是展示 SOTA 结果,而在于表明:一个简单的基于 MLP 的模型就可以取得与当前最佳的 CNN、基于注意力的模型相媲美的性能。
下表 1 列出了 Mixer 模型的各种配置以对标一些最新的 SOTA CNN 和基于注意力的模型:
下表 2 给出了最大 Mixer 模型与 SOTA 模型的性能对比结果:
当在 ImageNet-21k 上进行带有额外正则化的预训练时,Mixer 实现了非常好的性能(ImageNet 上 84.15% top-1),略低于其他模型。当上游数据集的大小增加时,Mixer 的性能显著提高。具体来说,Mixer-H/14 在 ImageNet 上取得了 87.94% top-1 的准确率,比 BiT-ResNet152x4 高 0.5%,比 ViT-H/14 低 0.5%。值得一提的是,Mixer-H/14 的运行速度要比 ViT-H/14 快 2.5 倍,比 BiT 快 2 倍。
图 2(左)展示了表 2 中 SOTA 模型在 ImageNet 数据集上的准确率、训练成本帕累托前沿(Pareto frontier):
下表展示了在多种模型和预训练是数据集规模上,Mixer 和其他一些模型的性能对比结果。
由上表可得,当在 ImageNet 上从头开始训练时, Mixer-B/16 取得了一个合理的 top-1 准确率 76.44%,这要比 ViT-B/16 低 3%。随着预训练数据集的增大,Mixer 的性能逐步提升。值得一提的是,在 JFT-300M 数据集上预训练、微调到 224 分辨率的 Mixer-H/14 取得了 86.32% 的准确率,比 ViT-H/14 仅低 0.3%,但运行速度是其 2.2 倍。
下图是原论文附带的代码,很简单,只有43行。
更多细节可参考论文原文。
代码语言:javascript复制