CaiT通过LayerScale层来保证深度ViT训练的稳定性,加上将特征学习和分类信息提取隔离的class-attention层达到了很不错的性能,值得看看
来源:晓飞的算法工程笔记 公众号
论文: Going deeper with Image Transformers
自ResNet出现以来,残差架构在计算机视觉中非常突出:
其中函数\(g_l\)和\(R_l\)定义了网络如何更新第l
层的输入\(x_l\)。函数\(g_l\)通常是恒等式,而残差分支\(R_l\)则是网络构建的核心模块,许多研究都着力于残差分支\(R_l\)的变体以及如何对\(R_l\)进行初始化。实际上,残差结构突出了训练优化和结构设计之间的相互作用,正如ResNet作者所说的:残差结构没有提供更好的特征表达能力,之所以取得更好的性能,是因为残差结构更容易训练。
目前很火的ViT网络可认为是实现了一种特定形式的残差架构:在将输入图像转换为一组\(x_0\)的向量之后,网络交替进行自注意力层 (SA) 与前馈网络 (FFN) 处理:
其中\(\eta\)是LayerNorm算子。
对于卷积神经网络和应用于NLP或语音任务的Transformer,如何对残差架构的残差分支进行归一化、加权或初始化受到了广泛关注。作者也在ViT上对不同初始化、优化和架构设计之间的相互作用进行了分析,并且提出了LayerScale层。LayerScale层包含一个初始权值接近于零的可学习对角矩阵,加在每个残差模块的输出上,可以有效地改进更深层架构的训练。
此外,作者还提出了class-attention层。类似于编码器/解码器架构,显示地将用于token间特征提取的transformer层与将token整合成单一向量进行分类的class-attention层分开,避免了两种目标不同的处理混合的矛盾现象。
通过实验验证,论文的主要贡献如下:
作者的目标是在提高Transformer架构的深度同时,提升图像分类训练优化的稳定性。在ViT和DeiT两项工作中,都没有研究仅在Imagenet上训练时,更大的深度可以带来任何好处:更深的ViT架构性能反而更低,而DeiT则只考虑了12层的架构。
图1展示了可能有助于优化的主要变体,图a是标准的预归一化结构。图b则是Fixup、ReZero和SkipInit这类引入可学习标量\(a_l\)的结构,该类结构会同时去掉预归一化层和学习率warmup:
ReZero简单地初始化为\(\alpha = 0\),而Fixup则初始化为\(\alpha = 1\)并进行其他修改:采用不同的权值的初始化策略,添加了几个偏置权值。但在作者的实验中,即使对超参数进行了调整,这些方法也难以收敛。
经过观察,移除warmup和层归一化是导致Fixup和T-Fixup训练不稳定的原因。因此作者重新引入这两部分,使Fixup和T-Fixup在DeiT模型上收敛,如图1c所示。当深度增加时,以较小的值初始化的可学习标量\(a_l\)确实有助于收敛。
作者提出的LayerScale对输出进行通道级别的乘法,而不是单个标量,如图1d所示,将权值更新与特定输出通道关联。公式上,可认为LayerScale是对每个残差分支输出的对角矩阵乘积:
其中参数\(\lambda_{l,i}\)和\(\lambda^{'}_{l,i}\)是可学习权值,初始化为一个固定的小值\(\varepsilon\):
该公式类似于其他归一化策略,如ActNorm或LayerNorm,但是在残差分支的输出上执行。此外,实际目的也有很大区别:
添加这些参数不会改变架构的特征表达能力,因为也可以集成到SA和FFN层的矩阵参数中,无需更改网络的实现。
CaiT架构如图2右,设计核心旨在规避ViT架构要求权值训练同时优化两个相互矛盾的目标的问题。两个矛盾的目标分别是:
为此,CaiT的核心就是将上面两个矛盾完全分隔开。
作为对比,在网络中间中插入Class token,这样前面的层可以专注于执行自我注意计算。作为不受矛盾目标影响的baseline,作者还考虑了将输出的平均池化用于分类的做法
CaiT包含两个不同的处理阶段:
class-attention阶段依次交替由多头类注意(CA)和FFN组成的层,在这个阶段只有class token会更新。
CA的作用是从token中提取信息,与SA 类似,但CA只计算class toekn \(x_{class}\)和\(x_{class}\)与冻结的token \(x_{patches}\)的集合之间的注意力。
定义具有h
个head和p
个token的网络,d
为token维度,将多头类注意力参数化为投影矩阵\(W_q、W_k、W_v、W_o \in \mathbb{R}^{d\times d}\)和偏置\(b_q, b_k, b_v, b_o \in R_d\)。基于上述定义,CA参数分支的计算可公式化为:
CA从特征token中提取有用信息整合到class token中。在实验中发现,第一个CA和FFN模块提供了主要的性能提升,叠加第二个模块就足以达到性能提升上限。
CA函数在内存和计算方面也比SA更轻量,因为CA只计算class token和token集合之间的注意力:
这意味着,与token数量成二次方的计算复杂度在CaiT层中变为线性计算负责度。
直接扩展网络深度,对不同训练超参数进行分析。
对比不同随机深度丢弃比例以及不同归一化策略的性能。
不同层的残差分支的权重可视化,使用LayerScale的权重会比较平稳。
对图2的三种架构进行对比分析,最后是对比不同的SA和CA组合比例。
不同大小的CaiT模型性能对比以及对应的超参。
与SOTA模型对比,\(\uparrow\)代表使用高像素finetune,\(\gamma\)代表使用Deit的蒸馏训练。
从ImageNet预训练迁移到其它分类数据集的性能对比。
从DeiT过渡到CaiT的性能对比。
Head数量对性能的影响。
数据增强Crop Ratio对性能的影响。
CaiT通过LayerScale层来保证深度ViT训练的稳定性,加上将特征学习和分类信息提取隔离的class-attention层达到了很不错的性能,值得看看。
如果本文对你有帮助,麻烦点个赞或在看呗~
更多内容请关注 微信公众号【晓飞的算法工程笔记】