0%

论文阅读-image-classification-MLP-Mixer: An all-MLP Architecture for Vision

动机

  • 计算机视觉历史证明,更大数据集的使用和增加的计算能力通常会导致范式的转变。例如:CNN-ViTViT延续了去除模型中人工设计的视觉特征和归纳偏置的长期趋势,进一步依赖于从原始数据中学习。

贡献

  • 证明了卷积和注意力能够取得良好的性能,但都不是必要的;
  • 提出了一种专门基于多层MLP体系的结构:MLP-Mixer,不使用卷积或自注意力,架构完全基于多层MLP,这些MLP在空间位置或特征通道上重复应用,只依赖于1) 基本的矩阵乘法,2) 数据布局更改(reshape和转置),3) scalar非线性。它包含两种类型的层:一种是MLP独立应用于图像patch(即“混合”每个位置的特征),一种是MLP应用于跨patch之间(即“混合”空间信息);
  • MLP-Mixer在大型数据集或现有正则化方案上进行训练时,可以在图像分类基准上获得有竞争力的分数,其预训练和推理成本可与最先进的模型相媲美。

Mixer的结构

Mixer的输入为一系列线性投影图像patch(也称为token),构成形状为 \text{patches}\times\text{channels} tableMixer使用两种类型的MLP层:channel-mixing MLPstoken-mixing MLPschannel-mixing MLP允许不同通道之间的通信,它独立地对每个token进行操作,并将table中的每一行作为输入。token-mixing MLP则允许不同空间位置(token)之间的通信,它独立地对每个通道进行操作,并将table中的每一列作为输入。这两种类型的层是交叉的,以支持两个输入维度的交互。(CNN角度来看:使用 1\times 1 卷积进行channel-mixing,使用全感受野的单通道深度卷积和参数共享进行token-mixing)

神经网络通常包含多个网络层,这些网络层混合1)给定的空间位置的特征 2)不同空间位置之间的特征,或同时混合这两种特征。CNN使用 N\times N(N>1) 卷积和池化层实现2),更深层的神经元具有更大的感受野,且使用 1\times1 卷积实现1),更大的核则同时实现1)和2)。在ViT和其它基于注意力的架构中,自注意力层执行1)和2),MLP网络块执行1)。在Mixer中,channel-mixing(单个位置)和token-mixing(跨位置)分别对应1)和2)。

如上图,Mixer的输入为一个包含 S 个非重叠的图像patchs的序列,每个图像patch被相同的线性层映射到隐藏维度 C 上,由此得到一个输入table,记作 X\in \mathbb{R}^{S\times C} 。假若原始输入图像的大小为 (H,W) ,每个patch的尺寸为 (P,P) ,则patch的数量为 S=\frac{HW}{P^2} Mixer包含多个相同大小的网络层,每层包含两个MLP块:1)token-mixing MLP:作用于 X 的列(即X的转置),映射为 \mathbb{R}^S\rightarrow\mathbb{R}^S ,在所有列之间共享;2)channel-mixing MLP:作用于 X 的行,映射为 \mathbb{R}^C\rightarrow\mathbb{R}^C ,跨所有行共享。每个MLP块包含两个全连接层和一个独立应用于其输入数据张量每一行的非线性层。

token-mixing MLP的维度选择与输入图像patch的数量无关,因此网络计算复杂度与patch的数量线性相关,不同于ViT的二次复杂度。channel-mixing MLP的维度选择与图像patch的大小无关,因此网络复杂度与图像的像素成线性关系,与CNN一致。

部分实验结果

在大中型数据集上预训练,在中小型下游分类任务上微调并测试,主要关注下游任务的准确率,预训练的计算开销,推理时间。

预训练数据集:ILSVRC2021 ImageNetImageNet -21kJFT-300M数据集

下游任务:ILSVRC2021 ImageNet带原始验证集标签和更改后的ReaL标签、CIFAR-10CIFAR-100Oxford- IIIT PetsOxford Flowers-102。此外,Visual Task Adaptation Benchmark (VTAB-1k)包含 19 个数据集。

Avg 5指在五个下游任务上的平均精度,Throughput指每个TPUv3核心每秒处理多少张图像,TPUv3指每个核心运行的时间。

  • 当预训练数据集增大时,Mixer的性能提升明显;
  • 相比于BiT和ViT,推理速度Mixer更快

上图左边展示了精度与预训练成本的关系,大致呈正相关,且Mixer是具有竞争力的;

上图右边展示了精度与训练数据量的关系,当数据量较少时Mixer过拟合严重,随着数据量增大时,Mixer性能快速增加。

两种方式影响模型大小:

(1) 在预训练时增加模型大小(层数、隐藏维数、MLP宽度),同时影响预训练计算量和测试时吞吐量;

(2) 微调时提高输入图像的分辨率,仅影响测试时吞吐量

上图显示:尽管Mixer在模型规模较小时略低于最新模型,但随着模型的增大、其竞争力越来越强。