各类Transformer都得稍逊一筹，LV-ViT：探索多个用于提升ViT性能的高效Trick

【导读】本文探索了用于提升ViT性能的各种训练技巧。通过一系列实验对比、改进与组合，本文所提方案取得了SOTA方案，超越了EfficientNet、T2TViT、DeiT、Swin Transformer、CaiT等方案。

标题&作者团队

paper: https://arxiv.org/abs/2104.10858

Code: https://github.com/zihangJiang/TokenLabeling

本文是新加坡国立大学&字节跳动的研究员在Transformer方面的最新研究成果。本文目标不在于提出一种新颖的Transofrmer架构，而是探索了用于提升ViT性能的各种训练技巧。通过一系列实验对比、改进与组合，本文所提方案取得了SOTA方案，超越了EfficientNet、T2TViT、DeiT、Swin Transformer、CaiT等方案。以26M参数量的ViT为例，所提方案在ImageNet上可以取得84.4%的top1精度；当参数量扩展到56M/150M时，模型性能可以进一步提升到85.4%/86.2%，且无需额外数据。

Abstract

本文为ImageNet分类任务中的Vision Transformer(ViT)提供了一个强基线方案。尽管ViT在ImageNet任务中取得了非常好的结果，但它们的性能距离强效CNN仍存在差距。

本文目标不在于提出一种新颖的Transofrmer架构，而是探索了用于提升ViT性能的各种训练技巧。本文结果表明：通过简单调节ViT的架构并引入token labeling，本文所提方案既可以取得比CNN与其他Transformer更佳的性能(同等参数量与计算量)。

以26M参数量的ViT为例，所提方案在ImageNet上可以取得84.4%的top1精度；当参数量扩展到56M/150M时，模型性能可以进一步提升到85.4%/86.2%，且无需额外数据。

上表对比了本文所提方案与CaiT在训练、模型方便的配置信息，可以看到：采用更少的训练技巧、参数量与计算量，所提LV-ViT即可取得比当前SOTA模型CaiT更佳的性能。

上图对比了所提方案与其他Transformer在参数量与性能方面的对比，可以看到：本文所提LV-ViT在56M参数量下可以取得比其他Transformer更佳的性能。

Bag of Traning Techniques for ViT

我们首先对ViT的结构进行简要介绍，然后描述如何通过训练技巧提升ViT的性能。

Vision Transformer

典型的ViT包含以下五个基本成分：

• Patch Embedding：在ViT中，固定尺寸的输入图像首先分解为一系列小的图像块(一般图像块尺寸预定义为)。以输入尺寸为例，它应当有个图像块，每个图像块(其维度为)通过线性层进行投影为特征向量，或者称之为token。所得特征向量将与class token拼接并送入到骨干网络用进行特征编码。
• Positional Encoding：由于transformer的排列不变特性，一般会添加一个额外的可学习位置编码到输入token以便于将每个块与位置信息相结合。这是非常有必要的，且有助于模型更好的学习视觉结构。从经验来看，固定的sinusoidal位置编码或者可学习编码均可，因为它们具有相似的分类性能。
• Multi-Head Self-Attention：多头自注意力(MHSA)旨在构建全局长期依赖。给定输入张量，多头自注意力在X上实施线性变换并得到。假设有H个自注意力头，这些信息将沿通道维度均匀拆分为，。自注意力模块的公式如下：

所有头的输出再沿着通道维度拼接并采用线性投影得到最后的输出，定义如下：

• Feed-Forward：前馈层包含两个线性层与一个非线性激活，定义如下：

• Score Prediction Layer：这里就是一个简单的分类头，略过。

Training Techniques

接下来，我们将介绍一系列有助于提升ViT性能的技术(可参见上图)，包含以下各个方面：

• Network Depth：众所周知，提升模型的深度有助于改善模型的泛化性能与容量。我们通过逐渐提升基线ViT的深度对其如何影响ViT的性能进行了研究。由于添加更多的模块不可避免会引入更多的模型参数，我们减少前馈层的隐含层维度(我们发现对模型性能几乎没有负面影响)。
• Explicit inductive bias：ViT旨在通过自注意力进行全局信息集成，并不会像CNN一样显式引入归纳偏置。尽管DeepVit一文提到：当无显式归纳偏置时，底层模块仍可以隐式学习之。然而，MHSA计算代价要远高于简单的卷积。正如T2TViT一文所提到：块嵌入中的tokenization难以捕获低级与局部结构信息，导致了训练过程中的较低的采样效率。考虑到这些问题，我们提出直接添加更多的卷积层增强块嵌入模块的容量以显式引入归纳偏置。此外，为从更大感受野提取信息，我们采用了更小stride的卷积以提供近邻token的重叠信息。
• Rethinking residual connection：假设LayerNorm表示为LN，我们重写Transformer模块的计算如下：

对此，我们解释为：通过逐渐添加源自其他token的信息调整输入。原始的ViT采用固定比例1添加信息调整，我们提出采用更小比例重缩放残差特征值：

由于更少的信息经过残差分支，这种方式可以增强残差链接；调整残差分支的缩放因子可以改善模型的泛化性能。

• Re-labeling：RandomCrop是训练阶段最常用的一种数据增光方法。ImageNet一种单标签基准数据集，然而其标签并不总是精确，甚至裁剪后的图像可能并不包含标签对应的目标。为解决该问题，我们在训练数据及中添加了re-labeling策略，为每个图像重设一个K=1000维得分图以得到更精确的标签。不同于知识蒸馏需要老师模型在线生成监督标签，re-labeling是一种廉价操作且可视作预处理操作。在训练过程中，我们仅需对每个裁剪图像计算K维稠密得分。因此我们能够以近乎可忽略的额外计算改善标签质量。
• Token Labeling：由于可以从所有token中收集全局信息，为ViT添加特定的token(比如cls)是一种常见操作。已有研究表明：在所有token中定义训练目标有助于改善训练时的采样效率。更进一步，基于re-labeling提供的稠密得分图，我们可以采用它为每个图像块与其对应token提供一个独特标签，我们称之为token labeling。具体来说，我们添加了一个token labeling损失函数，它利用了每个训练图像的稠密得分图并计算每个token与对应标签之间的交叉熵损失作为辅助损失。假设ViT的输出为，K维得分图为，整个图像的标签为。辅助损失定义如下：

注：H表示交叉熵损失。因此，总体损失定义如下：

• MixToken：已有研究表明：MixUp、CutMix等数据增广方法有助于提升ViT的性能与鲁棒性。然而，ViT依赖于块标记化将每个输入图像转换成token序列，本文所提token labeling策略同样是在块层面进行处理。如果我们直接将CutMix用于原始图像，某些图像块可能包含来自两个图像的内容，进而导致混合区域问题，见下图。

当实施token labeling时，我们很难为每个输出token赋予干净而正确的标签。考虑到这个问题，我们重新思考了CutMix并提出了MixToken，它可以视作改进版CutMix，见上图右。具体来说，假设两个图像表示为，他们对应的标签为，我们首先将两个图像送入块嵌入模块将每个图像序列化得到；然后，我们通过MixToken按照如下方式生成新的序列token：

注：M的生成方式与CutMix相同。对于每个标签，我们采用类似方式混合：

cls的标签重写如下：

Experiments

直接看结果咯，先看一下不同训练技巧的影响，再看所提方法与其他方案的性能对比。

我们采用DeiT-Small作为基线模型，上表给出了所提方法性能一步步演进，随着更多训练技巧的添加，模型性能逐渐提升。

Explicit inductive bias 上表给出了关于块嵌入的消融实验分析，可以看到：

• 相比原始单层块嵌入，添加更多的卷积层并收缩块嵌入的核尺寸可以一致性提升模型性能。
• 当卷积层数提升到6后，模型性能基本达到饱和，未发现性能提升；这意味：块嵌入中采用4个卷积层足矣。
• 采用更大的stride可以减少特征尺寸，进而大大减少计算量，但同时也带来了性能损失。最终，我们在块嵌入模块中k7s2+k3s1+k3s1+k8s8的配置。

Enhance residual connection 上图给出了残差缩放因子的性能对比，可以看到：引入残差缩放因子可以带来性能提升。与此同时，我们发现：采用更小的缩放因子可以导致更好地性能、更快的收敛速度。

Re-labeling 我们采用预训练NFNet-F6对ImageNet数据进行重标记，每个图像可以得到1000维得分图。从Table3可以看到：re-labeling可以带来0.4%的性能提升。

MixToken 上表给出了Mix-Token与Token labeling所带来的性能提升，可以看到：

• 均采用Re-labeling时，相比CutMix，MixToken可以带来0.2%性能提升；
• 相比Re-labeling，Token labeling可以带来额外的0.3%性能提升。

上表对比了不同数据增广的性能对比，可以看到：CutOut+RandAug的组合带来的最高性能。

上表对比了不同模型配置下的性能对比，可以看到：提升模型深度与嵌入维度可以进一步提升模型性能。

sota

上表给出了本文所提方案与CNN、其他Transformer在计算量、性能方面的对比，可以看到：

• 在小型模型方面，当测试分辨率为时，所提方法(仅需26M参数)取得了83.3%精度，以3.4%指标高于DeiT-S；
• 在中等模型方面，当测试分辨率为时，所提方法取得了与CaiT-S36相同的精度(85.4%)且具有更少的参数量和计算量；
• 在大型模型方面，当测试分辨率为时，所提方法取得了86.2%精度(与CaiT-M36相当)且仅需一半计算量和参数量；
• 所提方法同样取得了优于CNN模型(比如EfficientNet)与其他Transformer(比如T2TViT、Swin Transformer)的性能。

本文亮点总结

1. 本文目标不在于提出一种新颖的transformer架构，而是探索了用于提升ViT性能的各种训练技巧。本文结果表明：通过简单调节ViT的架构并引入transformer，本文所提方案既可以取得比CNN与其他Transformer更佳的性能(同等参数量与计算量)。

2. 本文通过逐渐提升基线ViT的深度对其如何影响ViT的性能进行了研究。由于添加更多的模块不可避免会引入更多的模型参数，我们减少前馈层的隐含层维度(对模型性能几乎没有负面影响)。