Press "Enter" to skip to content

大道至简——多层感知机向Transformer极速靠拢

卷积神经网络一直是计算机视觉领域的核心力量,最近随着Transformer的加入,这一来自自然语言处理领域的注意力机制迎头赶上,大有卷积模型在当年飞速发展之势。

 

正当人们惊呼Transformer将要迅速抢占视觉阵地的时候, 谷歌 的研究人员却从另外的角度来审视模型的发展。 本月初发表的工作指出不需要如此复杂精巧的结构, 仅仅使用简单的多层感知和简单的混合机制,就能实现与先进模型极为接近的性能 ,为视觉识别领域的研究打开的新的视野。

 

 

论文链接:

 

arxiv.org/abs/2105.01601

 

项目网站:

 

https://github.com/google-research/vision_transformer

 

一、从卷积到Transformer

 

再到多层感知机的轮回

 

近十年来,卷积神经网络在与大规模数据和不断增长的算力协同配合下将计算机视觉推升到了前所未有的高度。 而近几个月来来自Transformer的自注意力机制为计算机视觉开拓了全新的道路,克服现有方法弱点、提升模型性能,在深度学习的道路上持续塑造视觉领域的新方向。 人们在惊呼Transformer来了的同时,事情的发展开始有了不同的剧本。

 

在视觉Transformer(ViT)的启发下,研究人员就开始思考复杂的Transforme注意力r机制是不是真的有必要呢?为此,我们需要先来简单回顾一下ViT的架构。这一视觉领域的首个Transformer架构将图像先分成了32×32大小的片元,而后借助线性映射将图片元提取特征与位置编码相加构成Token,最后送入标准的Transformer编码器中进行处理。 T ransformer模 块中不仅需要计算输入序列的QKV,还需要进行多头自注意力计算,天生就需要非常大的计算量和海量的训练数据。

 

 

在这一架构的启发性,研究人员们开始进行了大胆的尝试。原先需要通过Transformer的层层复杂计算学习序列的注意力获取权重,那幺是不是可以将Transfomer模块替换为最简单的多层感知机(Multi-Layer Perceptrons, MLP)模块呢?

 

说干就干,于是研究人员们在充分考察CNN和Transformer的特点和优势后,通过简单的MLP组合和变换,研究人员真的构建出了一个性能强大的模型Mixer。这一模型仅仅在空间位置和特征通道上交替使用MLP构成的基本模块构建,其中仅仅包含了简单的乘法、转置等操作。它综合了极限情况下的卷积特征(1×1卷积),同时也从Transformer中借鉴了将图像片元线性映射为序列,逐层进行通道和Token特征交换的结构。

 

先来看看整体的架构模型:

 

MLP-Mixer由线性映射层、Mixer层和分类头组成。其中Mixer层中有两个MLP层,一个是用于混合不同位置token的,另一个是用于混合不同通道的MLP,每个MLP包含了两个全链接层和GELU非线性。此外模块中还包含了跳接、Dropout和层归一化等部件。

 

上图中展示了模型的整体架构,其输入是一系列图片元经过线性映射后获取的序列向量,维度为SxC(s个patch片元,C个通道),在整个模型中特征的维度一直维持不变。而后的Mixer层则由两个MLP组成,其中Channel-mixing MLP用于在每个片元对于的Token上对不同channel进行操作和沟通;而Token-mixing MLP则对不同空间位置Tokens的相同通道上对特征进行独立地混合,随着网络的深入两种MLP对于输入处理交替进行。

 

为了便于结合卷积模型理解,这一模型可以被视为一种特殊的卷积。其中Channel-mixing部分可以视为1×1卷积,Token-mixing部分则可视为单通道的具有全局感受野的逐深度卷积。但卷积的结构却比仅仅依赖矩阵乘法的MLP的要复杂。

 

下面我们从模型原理上来进行更深入的理解。 现代卷积神经网络模型可以视为一种特征的混合机制,在给定空间位置上或是在不同的空间位置上(也可以同时)进行特征混合。 例如典型的CNN就通过卷积和池化对不同空间位置的特征混合进行特征学习,越深的神经元具有越大的感受野。而1×1的卷积则可视为在同一位置对不同通道的特征进行混合;更大的三维卷积核则可以同时对空间和通道进行操作。另一方面,Transformer这类具有自注意力机制的模型可同时实现全局和通道的特征混合,而MLP模型则具有对给定空间位置处通道的混合能力。 在这样背景的启发性,Mixer的主要观点在于区分清楚了逐个位置的通道混合操作(Channel-Mixing)和不同位置处的token混合操作(Token-Mixing),都使用MLP来实现这两种操作方式。

 

前面的图示中详细描述了Mixer的流程,图像首先被分解为S个PxP大小的片元,而后通过相同的投影矩阵被线性映射为C维的输入模型中,模型输入的为SxC大小的矩阵。每个Mixer层包含了两个MLP模块,其中第一个Token-mixing模块首先对输入X的列方向进行作用,而第二个Channel-mixing则对X的行方向的通道进行作用。 每个MLP模块包含了两个全连接层和一个GELU非线性变换,其数学形式可以表达如下:

 

 

由于每个Mixer层中的MLP模块权重共享且与输入片元数和片元大小无关,所以模型可以与CNN模型一样实现相对于输入像素的线性复杂度。Mixer结构有很多独特的架构方式:Mixer中的每一层的输入输出尺寸相同,这种设计借鉴了Transformer和深度RNN的架构;同时固定的宽度也与CNNs中不断缩小的特征图结构大相径庭,所有层的输出宽度都固定了下来;值得一提的是,由于MLP的全连接结构对于输入Token对于的Patch顺序信息非常敏感,可以学习到位置信息,所以在其中省略位置编码。

 

二、实验验证

 

为了验证这一新想法的性能,研究人员按照ViT相似的方式进行了训练和测试:在中等/大规模数据集上预训练,在中小规模数据集上进行测试。主要关注三个方面的性能:下游任务的精度;预训练的整体计算开销;推理时的吞吐量。需要指出的是,本文的目标不是超过目前的SOTA模型,而是展示简单的MLP模块构建的模型也可以实现与目前最好的CNN和Transformer模型相媲美的结果。

 

在实验中研究人员一共构造了7中不同规模的模型进行比较,下表展示了不同模型的具体构造参数:

 

 

其中S,B,L,H分别代表了small,base,large,huge模型,后面的数字是输入patch片元的大小。

 

模型的训练和测试结果如下表所示,其中分别比较了在ImageNet21K(中等规模数据)上和JFT-300M(大规模数据)上进行训练的结果。

 

 

可以看到Mixer在ImageNet21k上取得了不错的结果。但与最先进的模型相比,由于没有Mixer中正则项所以会产生一定的过拟合,这与ViT的中观察到的现象一致。但随着Mixer可用的训练数据集增加,其性能得到了大幅改善。在JFT-300M数据集上模型的top-1达到了87.94%,甚至超过了BiT-ResNet153x4性能0.5%,而且仅仅比ViT-H/14低0.5%.而最值得一提的是,Mixer-H/14比ViT-H/14快2.5倍,比BiT快近2倍。整体计算效率如下图所示, 其精度与效率的平衡做的比大多数卷积模型要好。 下图右侧还展示了在不同架构下Mixer预训练数据对于下游任务的精度影响, 数据集的大小起到了重要的作用。

 

 

与SOTA模型在精度/训练开销上进行比较。模型在ImageNet21k,JFT或者带有噪声的网络数据集上进行训练。实验表面Mixer具有和ViT等先进模型一样优异的性能,同时随着数据量的增加,其精度提升幅度较大与先进模型间的差距也在大幅缩短。

 

此外研究人员还考虑了模型规模的影响。 增加模型大小(影响训练开销和推理速度)或者增加调优阶段输入图像的分辨率(影响推理速度),来观察训练的开销和推理速度。下图和下表可以发现,从零开始训练模型的时Mixer-B/16模型可以实现76.44%的top-1精度,比ViT-B/16低3%。在相同训练损失的情况下意味着Mixer-B/16存在一定程度的过拟合,这种情况在Mixer-L/16与ViT-L/16相比时更为明显。

 

 

模型精度与训练开销的关系;模型精度与推理速度的关系。

 

此外, 随着预训练数据集的增加Mixer的性能也得到了稳步提升。 上图中可以明显地看到模型的性能在持续提升,基本接近最先进算法的性能。随着训练数据集的增加,Mixer的过拟合得到了有效的缓解,其性能提升速度快于BiT,同时在BiT进入瓶颈期后Mixer还有继续增长的潜力。

 

 

在不 同模型规模和预训练数据集上的比较。

 

最后研究人员们还对模型中的权重进行了可视化来观察是否具有和CNN类似的分布。下图展示了在JFT-300M预训练和前面几层Mixer层中Token-mixing MLP的权重,因为Token-mixing可以在整图的空间位置上进行交互,所以一些特征会对整张图进行处理,而一些特征则集中在小区域中。 第一个Token-mixing中含有更多的局域交互,而第二第三层中则含有更大范围的交互。

 

 

Mixer-B/16中第一二三层中的Token-mixingMPL权重,其196个权重Reshpe为了14×14的特征图。

 

如果想要了解更多细节和实验原理,请参考论文和项目网站。

 

编译: T.R From: Google

 

Illustrastion  by Oleg Shcherba  from Ico ns8

 

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注