• 问答
  • 技术
  • 实践
  • 资源
搞懂 Vision Transformer 原理和代码,看这篇技术综述就够了(十)
技术讨论

作者丨happy
来源丨极市平台
编辑丨极市平台

本文目录

21 LV-ViT: 56M参数训练视觉Transformer
(来自 新加坡国立大学冯佳时组 ,字节跳动)
21.1 LV-ViT原理分析

22 通过抑制过度平滑来改进视觉Transformer训练
(来自 Facebook)
22.1 原理分析

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型,现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制,不采用 RNN 的顺序结构,使得模型可以并行化训练,而且能够拥有全局信息。

本文介绍的两个工作都不是提出了新的Vision Transformer结构,而都是为了解决:"如何更好地训练视觉Transformer?" 这个问题。

21 LV-ViT: 56M参数训练视觉Transformer

论文名称:Token Labeling: Training a 85.5\% Top-1 Accuracy Vision Transformer
with 56M Parameters on ImageNet

论文地址:

https://arxiv.org/pdf/2104.10858.pdf​arxiv.org

  • 21.1 LV-ViT原理分析:

本文提出了一个叫做 Token Labeling 的方式解决这个问题,Token Labeling是一个新的训练目标,借助它可以使得参数量只有26M/56M/150M的视觉Transformer模型在ImageNet上的精度达到84.4\%/85.4\%/86.2\%,Token Labeling 是本文的唯一创新点和贡献。

作者首先总结和列举了之前有过的几种不同的Training Techniques,看看它们对视觉Transformer模型的效果到底如何。

  • 1 增加模型的深度: 原始的视觉Transformer是12个block,简单地多加几个block对模型性能没什么改变,这一点已经在DeepViT和CaiT的论文中得到了证实,可参考DeepViT和CaiT这2个模型:

科技猛兽:Vision Transformer 超详细解读 (原理分析+代码解读) (八)​zhuanlan.zhihu.com图标

  • 2 在模型中引入卷积使之具有inductive bias的能力: 前作已经指出,在视觉Transformer中加入一些卷积操作,能够更容易地提取到low-level的信息 (即edges 和 lines),这是因为卷积操作具有inductive bias的能力,使得模型在浅层能很好地学习到low-level的信息,这些信息使得输入给 Transformer Block 的patch embedding 更有意义。本文沿用了这一操作,在模型中加入了卷积。同时为了使得卷积的感受野更大,本文使得卷积操作的stride更小,这样一来就能够提供更多的overlapped的信息。在模型中引入卷积这一点已经在CvT,CeiT的论文中得到了证实,具体可以参考下面的链接:

科技猛兽:Vision Transformer 超详细解读 (原理分析+代码解读) (六)​zhuanlan.zhihu.com图标

  • 3 在残差连接中加入scale factor: 原始Transformer模型的前向过程是这样的:

$$
\begin{equation} \begin{split} X \longleftarrow X + \text{SA}(\text{LN}(X)) \ X \longleftarrow X + \text{FF}(\text{LN}(X)) \end{split} \end{equation} \tag{21.1}
$$

这样的过程可以视为:从输入信号 $X$ 开始,不断地在它上面叠加一些其他的信息 $\text{SA}(\text{LN}(X))$ 或者 $\text{FF}(\text{LN}(X))$ ,这里的 $\text{SA}$ 表示multi-head self-attention,而 $\text{FF}$ 表示Feed Forward Network。之前的CaiT的论文提到:那些其他的信息,在叠加在输入信号上面之前,要先乘上一个系数,而且不同的dimension,要乘的系数scale也不同。如下图1所示:

图1:CaiT的结构,右侧分支在叠加在输入信号上面之前,要先乘上一个系数lambda

本文也采用了这个做法,这一点的作用已经在CaiT,ReZero/Skipinit and Fixup等多篇论文中得到了证实。

  • 4 Re-labeling技术: 在训练阶段总是会采用一些数据增强的方法,比较常用的是random crop方法,但是在random crop以后,标签并不总是准确的,因为ImageNet中的许多图像包括多个object,并且具有Ground Truth标签的对象可能不会保留在random crop之后的图像中。这个问题,当random crop的图像越小时就会越严重,因为crop出来的图像越小,可能越不会包含Ground Truth的label信息,如下图3所示。在使用单标签标注时,图像随机剪裁可能包含与真值完全不同的对象,为训练带来噪声甚至不准确的监督信号。所以作者遵循了之前一篇论文:Re-labeling ImageNet: from Single to Multi-Labels, from Global to Localized Labels的解决方法,论文链接如下:

https://arxiv.org/pdf/2101.05022.pdf​arxiv.org

Re-labeling技术的具体做法是使用多标签对 ImageNet 训练集进行重新标注,具体是怎么做的呢?

Re-labeling论文的作者首先在一个巨大的数据集上训练了一个强大的图像分类器模型,这个数据集是JFT-300M或者InstagramNet-1B,并把这个模型在ImageNet上进行fine-tune。 作者认为这个图像分类器模型非常接近人类的预测了。那么因为这个数据集是私有的,并且训练这个模型非常耗费时间和算力 (上百个GPU days),所以为了后人研究的方便,作者开源了这个模型。

有了这个模型以后,接下来怎么办呢?

我们把一张图片输入这个强分类模型,取它的Global Average Pooling之前的输出特征,我们假设这个特征是 $L\in R^{H\times W\times C}$ ,我们把它叫做Label map,如下图2所示。原来的分类器得到 $1\times 1\times C$ 维的输出结果,而我们只需要 $H\times W\times C$ 维的Label map。接下来假设我们random crop的区域是 $[c_x, c_y, c_w, c_h]$ ,那么在训练时首先把这个区域的图片给resize成224×224的大小作为输入图片,再把这个特征Label map的位于 $[c_x, c_y, c_w, c_h]$ 的区域进行ROIAlign操作,得到 $1\times 1\times C$ 维的输出结果。最后对这个结果进行softmax操作得到最终的 $1\times 1\times C$ 维的输出结果,作为这个区域图像的label进行训练。

图2:得到Label map

这一过程的伪代码如下:

Algorithm A1 ReLabel Pseudo-code
for each training iteration do
   # Load image data and label maps (assume the minibatch size is 1 for simplicity)
   input, label map = get minibatch(dataset)
   # Random crop augmentation
   [cx,cy,cw,ch] = get crop region(size(input))
   input = random crop(input, [cx,cy,cw,ch])
   input = resize(input, [224; 224])
   # LabelPooling process
   target = RoIAlign(label map, coords=[cx,cy,cw,ch], output size=(1; 1))
   target = softmax(target)
   # Update model
   output = model forward(input)
   loss = cross entropy loss(output, target)
   model update(loss)
end for

也就是说,对于Random crop的每一个crop的区域,在参与训练时,其使用的标签都不再是原图的标签了,而是由这个强分类器得到的软标签。而且,使用强分类器生成软标签,解决了标注成本问题;在最终池化层之前使用像素级多标签预测,以充分利用额外的位置特定监督信号。

图3:在random crop以后,标签并不总是准确的

上图3所示就是再提一遍具体的做法:比如输入一张左上角的图片,这个图片里面有房子,有牛,有天空等等。我们在做Random crop的时候得到了4个crop的图片,就是第2行所列的4个图,并且都把它们resize到224×224的大小。这时,我们使用这4张resize到224×224的图进行训练,我们给每张图片分配一个 $K$ 维的score map,这里 $K$ 是数据集的类别数,比如对于ImageNet-1k就是 $K=1000$ 。这个score map就是之前提到的在JFT-300这个超大的私有数据集上训练的强分类器得到的。对于Random crop的每一个crop的区域,在参与训练时,其使用的标签都不再是原图的标签了,而是由这个强分类器得到的软标签。

  • 5 Token labeling:

Token labeling可以用下图4来表述。作者认为模型的输出,即 $[X^{cls},X^1,...,X^N]$ 现在只利用了 $X^{cls}$ ,而剩下的 $N$ 个token其实是可以结合上面的Re-labeling技术加以利用。

一张图片分了patch以后,每个patch都会转化成token,那这些token其实也是有软标签的。我们只需要给刚才那个强分类模型输入一下这个图片,就能够使用相同的技术得到所有这些token的软标签 $[y^1,...,y^N]$ ,这里的每一个 $y$ 都是一个 $K$ 维向量,它也叫作token label。

图4:Token labeling

所以,作者定义了一个 token labeling loss:

$$
\begin{equation} L{aux}=\frac{1}{N}\sum{i=1}^N H(X^i,y^i) \end{equation} \tag{21.2}
$$

所以总的loss函数变成了:

$$
\begin{equation} L{total}=H(X^{cls},y^{cls}) + \beta L{aux}, \end{equation} \tag{21.3}
$$

式中 $\beta$ 是个超参数,设为了0.5。

同时在训练模型时总是使用了CutMix技术,它能提高模型的性能和鲁棒性。但如果我们直接在原始图像上应用CutMix,产生的一些补丁可能包含来自两个图像的内容,导致一个小patch内的混合区域,如图5左图所示。这样的话,很难为每个输出token分配一个干净、正确的软标签。考虑到这种情况,作者重新考虑了CutMix的增强方法,并给出了MixToken,它可以被看作是在Patch embedding后对令牌进行操作的CutMix的修改版本,如图5右侧所示。

图5:CutMix\(左\) 和 MixToken\(右\)

具体来说,对于两个图像 $I_1,I_2$ 及其对应的token label $Y_1=[y_1^1,...,y^N_1]$ 以及 $Y_2=[y^1_2,...,y^N_2]$ ,我们首先把这2个图像变为a sequence of tokens:

$T_1=[t^1_1,...,t^N_1], T_2=[t^1_2,...,t^N_2]$ 。接下来通过一个mask $M$ 得到一个新的sequence of tokens:

$$
\begin{equation} \hat{T}=T_1\odot M + T_2 \odot (1-M), \end{equation} \tag{21.4}
$$

同时也把这2张图片对应的软标签使用这个mask 进行结合:

$$
\begin{equation} \hat{Y}=Y_1\odot M + Y_2 \odot (1-M). \end{equation} \tag{21.5}
$$

class token的标签可以写为:

$$
\begin{equation} \hat{y^{cls}}= \bar{M} y_1^{cls} + (1 - \bar{M}) y_2^{cls}, \end{equation} \tag{21.6}
$$

式中 $\bar{M}$ 时Mask的均值。

这个过程在图4中有所体现,那么结合 token labeling loss,可以利用软标签以更好地训练模型。

Experiments:

超参数
数据增强 CutOut,RandAug,CutMix,MixUp(没有作用,所以去掉)
优化器 AdamW
学习率 0.001*batch size/1024
weight decay rate 0.05
warm up epochs 5

超参数的设置如上表和下图6所示:

图6:具体的超参数设置

不同尺寸大小的模型如下图7所示:

图7:不同尺寸大小的模型

从DeiT-Small模型渐变到LV-ViT-S模型:每一步的涨点,如下图8所示:

从结果我们可以看到,主要的涨点行为来自本文提出的Token-labeling。

图8:从DeiT-Small模型渐变到LV-ViT-S模型的每一步涨点

对比实验1:卷积核大小的影响

如图9所示为第1个ablation study,即探究卷积核设置对性能的影响。这个表的基本设置是图7的第3行,即baseline是16层的ViT模型和MLP ratio=3。图8表明4层卷积已经足够,第1个卷积层最好设置成kernel size=7,stride=2;中间2个卷积层设置成kernel size=3,stride=1;最后个卷积层设置成kernel size=8,stride=4。

图9:卷积核大小的影响

对比实验2:改进残差连接

这里对比的是上文的第3点改进: 在残差连接中加入不同的scale factor值,如下图10所示。结果发现小的scale factor值效果更好一点,且收敛更快。因为假设残差值不变的话,小的scale factor值会使得 $F(X)$ 这部分的值更大些。比如这一层要实现的功能是:

$$
Y=2X+0.3=X+F_1(X) \
$$

所以针对这一层有 $F_1(X)=X+0.3$ 。假设加入了scale factor = 0.5:

$$
Y=2X+0.3=X+F_2(X)/2 \
$$

所以针对这一层有 $F_2(X)=2X+0.6$ 。

而注意 $F(X)$ 是个神经网络,所以 $F_2(X)$ 比 $F_1(X)$ 具有更大的weight和gradient,保留了更多信息。

图10:改进残差连接实验结果

对比实验3:Re-labeling

前文已经讲到使用Re-labeling这个技术需要一个强分类模型,这里作者使用的这个强分类模型是NFNet-F6,这个模型可以在ImageNet数据集上获得86.5\%的准确率。作者使用它对于每一张图片得到一个1000维的score map。根据图7可知,Re-labeling的技术使得模型性能从82.4\%涨到了82.8\%。

与SOTA模型的对比:

结果如下图11所示:

图11:与SOTA模型的对比

作者将本文的LV-ViT与之前的经典模型ViT,DeiT,T2T-ViT,CaiT,Swin-Transformer等进行了对比。对于小模型,LV-ViT-S和DeiT-S都是26M参数量,但是后者性能由79.9涨到了83.3。中等模型LV-ViT-M↑384与CaiT-36S↑384对比,它使用知识蒸馏来改进性能,但是这将引入更多的计算。中等模型LV-ViT-L↑448与CaiT-36M↑448对比,节约了近一半的参数量。

小结:

本文是对以往前人提出的方法的一个汇总,包含:增加模型的深度,在模型中引入卷积,在残差连接中加入scale factor,Re-labeling 技术和 Token labeling。其中最后一点是本文作者把Re-labeling技术用在了视觉Transformer的具体实现,总体而言创新性略微偏弱,但是故事的切入点找得很好:作者称他们提出了视觉Transformer的训练策略,但本质上而言以上技术大多也可以算是已有的在模型结构上的改进 (如增加模型的深度,在模型中引入卷积,在残差连接中加入scale factor等等)。个人认为本文对业界的贡献之一是让人们开始关注之前的 Re-labeling ImageNet: from Single to Multi-Labels, from Global to Localized Labels 这篇论文,Re-labeling ImageNet这篇论文是对 ImageNet 的标签进行改进的非常有意义的工作,而本文可以看做是把这篇论文的思想应用在视觉Transformer模型的一次有价值的尝试。

22 通过抑制过度平滑来改进视觉Transformer训练

论文名称:Improve Vision Transformers Training by Suppressing
Over-smoothing

论文地址:

https://arxiv.org/pdf/2104.12753.pdf​arxiv.org

  • 22.1 原理分析:

直接在视觉任务上训练视觉Transformer模型容易产生不稳定和次优的结果。之前有的工作通过加入卷积以更有效地提取low-level的特征,从而加速和稳定视觉Transformer模型的训练,这相当于改变了Transformer的结构。而本文所做的工作也是在不改变原有Transformer模型任何结构的前提下通过改变训练方式来稳定训练。

本文作者把视觉Transformer模型训练不稳定的原因归结为:过平滑问题 (Over-smoothing problem),即:self-attention层趋向于把输入图片的不同patch映射成差不多的latent representation

这句话是啥意思呢?

图片划分成patch之后,每个patch会经过一个Transformer Block层得到token,不同的pacth对应着不同的token。但是,当模型越来越深时,不同的token之间的相似性越来越大,大到甚至无法有效地进行区分了。作者称这种现象为过平滑问题 (Over-smoothing problem)。

其实这个现象在之前的DeepViT这篇论文中已有所提及,详见下面的链接。 只是DeepViT称这种现象为注意力坍塌 (attention collapse):随着ViT的深入,self-attention 在生成不同注意力以捕获rich representation方面变得不那么有效了。

科技猛兽:Vision Transformer 超详细解读 (原理分析+代码解读) (八)​zhuanlan.zhihu.com图标

但不同的是,DeepViT描述的是不同的Block之间的相似性随着模型的加深而增加;本文所描述的是不同token之间的相似性随着随着模型的加深而增加。

作者首先定义了一种不同patch representation之间的Layer-wise cosine similarity:假设输入图片 $x$ 及其对应的patch representation为 $\bm{h} = (h^{\textit{cls}}, h_1, \cdots, h_n),h_j \in R^d$ ,定义smoothness为它的patch representation中所有token之间的cosine similarity:

$$
\mathrm{CosSim}(\bm{h}) = \frac{1}{n(n-1)} \sum_{i\neq j} \frac{h_i^\top h_j}{\parallel h_i\parallel \parallel h_j\parallel}, \tag{22.1}
$$

式中 $\parallel \cdot \parallel$ 为Euclidean norm。

如下图12所示为所有层的cosine similarity值,可以观察到随着层数的加深,DeiT模型的cosine similarity越来越大。一个24层的DeiT-Base模型的最后一层的不同token之间的pairwise相似度达到了0.9,表明所学到的patch representation之间的高度相关性和重复性。

图12:随着层数的加深,cosine similarity越来越大

此外,作者还通过另外一个指标来描述这种相似性:Layer-wise standard deviation

我们希望每一层的某一个patch,它attend to其他所有patch的程度是不一样的,这样才能捕获到图片中有意义的区域的信息。Layer-wise standard deviation就是为了衡量这种能力的大小。假设某个patch $h_i$ ,其softmax attention score为 $S(h_i)\in R^n$ ,这是一个 $n$ 维的向量,作者将其求标准差来表示相似度的大小,相似度越小,标准差就越大。所以说每一个layer我们能求出来 $n$ 个标准差。将它们取平均值来代表这一层的不同token之间的相似度的大小,得到的结果如下图13所示:

图13:一层的不同token之间的相似度的大小,standard deviation越小,相似度越大

从结果可以看出,传统的Transformer模型每一层的不同token之间的standard deviation都很小,表示相似度都很大,也就造成了over-smoothing的问题。但是在NLP模型中,每个patch都有其对应的label,所以不同的token之间的相似性不易变得很大,也就不容易出现over-smoothing的问题。

为了解决这个问题,作者提出了一系列方法:

1 添加相似度罚项:

最直接的避免每一层的不同token之间的相似度过大的办法是添加罚项,基于此作者使用了一个新的loss函数,对于最后一层的patch representation ${h} = (h^{\textit{cls}}, h_1, \cdots, h_n)$ 来讲:

$$
\ell{\textit{cos}} = \frac{1}{n(n-1)} \sum{i\neq j} \frac{h_i^\top h_j}{\parallel h_i\parallel \parallel h_j \parallel}. \tag{22.2}
$$

来最小化相似性。这种做法已经有相关的工作在GNN和其他NLP任务上实现了,可以看做是增加最后一层的patch representation的表达能力。

2 深层特征与浅层对应特征接近 (+ Patch Contrastive Loss):

作者认为浅层的patch应该与深层对应的那个的patch的值比较接近,这样一来patch之间的区分度就能够在整个网络中维持住。对于一个给定的图片,假设 $e$ 是第1层的patch representations,而 $h$ 是最后1层的patch representations,作者使用了对比学习的损失:

$$
\ell{\textit{cons}} = - \frac{1}{n} \sum{i=1}^n \log \frac{\exp(e_i^\top h_i)}{\exp(e_i^\top h_i) + \exp(ei^\top (\sum{j=1}^n h_i / n))}, \tag{22.3}
$$

来迫使浅层和深层相对应的patch representations接近,而与其他的patch representations疏远,如下图14所示。

图14:浅层和深层相对应的patch representations接近,而与其他的patch representations疏远

3 每个Patch都有监督信息 (+ Patch Mixing Loss):

与上篇文章LV-ViT的思路相似,作者希望每个patch都能有监督信息,这样一来就可以使用CutMix这个数据增强策略。具体而言是给每个patch的向量通过一个共享的分类器得到分类的结果。比如说图15中输入一个image,并经过了CutMix,这个image就包含了Dog和Bird,其中Dog对应的Patch是蓝色的部分,Bird对应的Patch是草绿色的部分。这些patch都会有对应的监督信息,所以模型能够学习到更informative和更useful的patch representation。

Patch Mixing Loss定义为:

$$
\ell\textit{token} = \frac{1}{n} \sum{i=1}^n \ell_\textit{ce}(g(h_i), y_i), \tag{22.4}
$$

图15:每个patch都有监督信息

作者结合以上的3个损失函数训练视觉Transformer模型,而不对模型结构有任何改进。

Experiments:

再次强调本文所使用的模型没有进行任何结构上的变化,即:没有添加卷积等操作以方便地提取low-level的特征。

与SOTA模型的对比:

作者对比了很多最近出的视觉Transformer模型,包含一些与CNN结合的模型,如下图16所示。

本文方法结合DeiT模型可以在ImageNet上达到85.0\% 的 top-1 accuracy,略低于与CNN结合的模型。

图16:与SOTA模型的对比

对比实验1:各个模块的性能

如下图17所示为对比各个模块的性能,发现Patch Contrastive Loss 和 Patch Mixing Loss都能够提升性能。把这2个loss都加上去的话就可以涨到82.6\%的top 1 accuracy。

图17:各个模块的性能

对比实验2:训练更稳定

标准Transformer的训练并不稳定。删除一些正则化或增加模型深度会使训练变得困难,并产生更差的性能,如下图18所示。将本文提出的几个损失引入训练过程,可以使得训练Transformer结构更稳定。删除一些正则化或使用更大的模型不会导致更差的结果。

图18:标准Transformer的训练并不稳定,使用本文方法可使训练更加稳定

对比实验3:更大的drop path rate

Drop path rate是指Transformer的一个超参数,代表Dropout的比例。如下图19所示,添加了Patch Contrastive Loss 和 Patch Mixing Loss以后可以使用更大的drop path rate。

图19:添加了Patch Contrastive Loss 和 Patch Mixing Loss以后可以使用更大的drop path rate

小结:

本文提出了3个对视觉Transformer模型的训练非常有意义的损失函数,可以在不改变Transformer任何架构的前提之下使它的训练更加稳定,且能获得性能上的提升。其出发点也是从降低各个token之间的相似性入手,Patch Contrastive Loss使得深层特征与浅层对应特征更加接近,Patch Mixing Loss给每个Patch提供一个监督信息。

总结:

本文所介绍的2个工作都不是提出了新的Vision Transformer结构,而都是为了解决:"如何更好地训练视觉Transformer?" 这个问题。第1篇是对以往前人提出的方法的一个汇总,包含:增加模型的深度,在模型中引入卷积,在残差连接中加入scale factor,Re-labeling 技术和 Token labeling。其中最后一点是本文作者把Re-labeling技术用在了视觉Transformer的具体实现,它告诉我们通过使用Re-labeling 的技术也可以改进视觉Transformer的性能。第2篇提出了3个对视觉Transformer模型的训练非常有意义的损失函数:相似度罚项,Patch Contrastive Loss 和 Patch Mixing Loss, 可以在不改变Transformer任何架构的前提之下使它的训练更加稳定,且能获得性能上的提升。

  • 0
  • 0
  • 1074
收藏
暂无评论