ICCV2021 | Tokens-to-Token ViT:在ImageNet上从零训练Vision Transformer

前言

本文介绍一种新的tokens-to-token Vision Transformer(T2T-ViT)，T2T-ViT将原始ViT的参数数量和MAC减少了一半，同时在ImageNet上从头开始训练时实现了3.0%以上的改进。通过直接在ImageNet上进行训练，它的性能也优于ResNet，达到了与MobileNet相当的性能。

本文来自公众号CV技术指南的论文分享系列

关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读。

论文：Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet

代码：https://github.com/yitu-opensource/T2T-ViT

Background

Vision Transformer(ViT)是第一个可以直接应用于图像分类的全Transformer模型。具体地说，ViT将每个图像分割成固定长度的14×14或16×16块(也称为tokens)；然后ViT应用Transformer层对这些tokens之间的全局关系进行建模以进行分类。

尽管ViT证明了全Transformer架构在视觉任务中很有前途，但在中型数据集(例如ImageNet)上从头开始训练时，其性能仍逊于类似大小的CNN对等架构(例如ResNets)。

论文假设，这种性能差距源于ViT的两个主要局限性：

1)通过硬分裂对输入图像进行简单的tokens化，使得ViT无法对图像的边缘和线条等局部结构进行建模，因此它需要比CNN多得多的训练样本(如JFT-300M用于预训练)才能获得类似的性能；

2)ViT的注意力骨干没有很好地像用于视觉任务的CNN那样的设计，如ViT具有冗余性和特征丰富度有限的缺点，导致模型训练困难。

为了验证论文的假设，论文进行了一项初步研究，通过图2中的可视化来调查ViTL/16和ResNet5的获知特征的差异。论文观察ResNet的功能，捕捉所需的局部结构(边、线、纹理等)。从底层(Cv1)逐渐向中间层(Cv25)递增。

然而，ViT的特点却截然不同：结构信息建模较差，而全局关系(如整条狗)被所有的注意块捕获。这些观察结果表明，当直接将图像分割成固定长度的tokens时，原始 ViT忽略了局部结构。此外，论文发现ViT中的许多通道都是零值(在图2中以红色突出显示)，这意味着ViT的主干不如ResNet高效，并且在训练样本不足的情况下提供有限的特征丰富度。

图2.在ImageNet上训练的ResNet50、ViT-L/16和论文提出的T2T-VIT-24的功能可视化。绿色框突出显示学习的低级结构特征，如边和线；红色框突出显示值为零或过大的无效要素地图。注意：这里为ViT和T2T-ViT可视化的特征图不是attention图，而是从tokens重塑的图像特征。

创新思路

论文决意设计一种新的full-Transformer视觉模型来克服上述限制。

1)与ViT中使用的朴素tokens化不同，论文提出了一种渐进式tokens化模块，将相邻tokens聚合为一个tokens(称为tokens-to-token模块)，该模块可以对周围tokens的局部结构信息进行建模，并迭代地减少tokens的长度。具体地说，在每个tokens-to-token(T2T)步骤中，transformer层输出的tokens被重构为图像(restructurization)，然后图像被分割成重叠(soft split)的tokens，最后周围的tokens通过flatten分割的patches被聚集在一起。因此，来自周围patches的局部结构被嵌入要输入到下一transformer层的tokens中。通过迭代进行T2T，将局部结构聚合成tokens，并通过聚合过程减少tokens的长度。

2)为了寻找高效的Vision Transformer主干，论文借鉴了CNN的一些架构设计来构建Transformer层，以提高功能的丰富性，论文发现ViT中通道较少但层数较多的“深度窄”架构设计在同等型号和MAC(Multi-Adds)的情况下性能要好得多。具体地说，论文研究了宽ResNet(浅宽VS深窄结构)、DenseNet(密集连接)、ResneXt结构、Ghost操作和通道注意。论文发现其中，深窄结构对于ViT是最有效和最有效的，在几乎不降低性能的情况下显著地减少了参数数目和MACs。这也表明CNNs的体系结构工程可以为Vision Transformer的骨干设计提供帮助。

基于T2T模块和深度窄骨干网架构，论文开发了tokens-to-token Vision Transformer(T2T-ViT)，它在ImageNet上从头开始训练时显著提高了性能，而且比普通ViT更轻便。

Methods

T2T-ViT由两个主要部分组成(图4)：

1)一个层次化的“Tokens-to-Token模块”(T2T模块)，用于对图像的局部结构信息进行建模，并逐步减少tokens的长度；

2)一个有效的“T2T-ViT骨干”，用于从T2T模块中提取对tokens的全局关注关系。

在研究了几种基于CNN的体系结构设计后，对主干采用深窄结构，以减少冗余度，提高特征丰富性。

图4.T2T-ViT的整体网络架构。在T2T模块中，首先将输入图像soft split为patches，然后将其展开为token T0序列。在T2T模块中，token的长度逐渐减小(在这里使用两次迭代和输出Tf)。然后，T2T-VIT主干将固定token作为输入并输出预测。两个T2T块与图3相同，PE为位置嵌入。

Tokens-to-Token

Tokens-to-Token(T2T)模块旨在克服ViT中简单tokens化的限制。它将图像逐步结构化为表征，并对局部结构信息进行建模，这样可以迭代地减少表征的长度。每个T2T流程有两个步骤：重组和Soft Split(SS)(图3)。

图3.T2T流程图解。

经过变换和reshape后，tokens Ti被重构为图像Ii，然后重叠split为tokens Ti+1。具体地说，如粉色面板中所示，输入Ii的四个tokens(1、2、4、5)被串联以形成一个tokens 在Ti+1。T2T transformer可以是普通的transformer 层或有限GPU存储器中的像Performer层这样的其他高效transformer。

在进行soft split时，每个块的大小为k×k，在图像上叠加s个，其中k−类似于卷积运算中的步长。因此，对于重建图像I_∈_rh×w×c，soft split后的输出tokens的长度为

每个分割patches的大小为k×k×c。将空间维度上的所有patches展平，以To表示。在soft split之后，为下一个T2T过程馈送输出tokens。

通过迭代进行上述重构和soft split，T2T模块可以逐步减少tokens的长度，并转换图像的空间结构。T2T模块中的迭代过程可以表示为

对于输入image I0，首先应用soft split将其分割为tokens：T1=SS(I0)。在最终迭代后，T2T模块的输出tokens Tf具有固定的长度，因此T2T-ViT的主干可以对Tf上的全局关系进行建模。

T2T-ViT Backbone

论文探索了不同的VIT体系结构设计，并借鉴了CNN的一些设计，以提高骨干网的效率，增强学习特征的丰富性。由于每个transformer层都有跳跃连接，一个简单的想法是采用如DenseNet的密集连接来增加连通性和特征丰富性，或者采用Wide-ResNets或ResNeXt结构来改变VIT主干中的通道尺寸和头数。

论文探讨了从CNN到VIT的五种架构设计：

密集连接如DenseNet；
深-窄与浅-宽结构如宽ResNet]；
通道注意如挤压-激励(SE)网络；
多头注意层中更多的分头如ResNeXt；
Ghost操作如Ghost Net。

实验发现：1)采用简单降低通道维数的深窄结构来减少通道中的冗余，增加层深来提高VIT中的特征丰富度，模型尺寸和MACs都有所减小，但性能有所提高；2)SE块的通道关注度也提高了VIT，但效果不如深窄结构。

基于这些发现，论文为T2T-VIT骨干网设计了一种深窄结构。具体地说，它具有较小的通道数和隐藏维度d，但具有更多的层b。对于T2T模块最后一层定长的Token，论文在其上拼接一个类Token，然后添加正弦位置嵌入(PE)，与VIT一样进行分类：

T2T-ViT Architecture

T2T-VIT的结构细节。T2T-VIT-14/19/24的型号尺寸与ResNet50/101/152相当。T2T-VIT-7/12的型号大小与MobileNetV1/V2相当。对于T2T transformer 层，在有限的GPU内存下，论文采用了T2T-VITT-14的transformer层和T2T-VIT-14的Performer层。对于VIT，‘S’表示小，‘B’表示基本，‘L’表示大。‘VIT-S/16’是原始VIT-B/16的变体，具有更小的MLP大小和层深。

Conclusion

如图1所示，论文的215M参数和5.2G MACS的T2T-ViT在ImageNet上可以达到81.5%的TOP-1准确率，远远高于ViT的48.6M参数和10.1G MACs的TOP-1准确率(78.1%)。这一结果也高于流行的类似大小的CNN，如具有25.5M参数的ResNet50(76%-79%)。此外，论文还通过简单地采用更少的层来设计T2T-ViT的精简变体，取得了与MobileNets(图1)相当的结果。