3青出于蓝
3.1 Weight Multiplexing
权重共享
的潜力已经在自然语言处理中得到了证实;然而,其在Vision Transformer
中的作用尚不清楚。为了检验这一点,作者直接将跨层权重共享
应用于DeiT-S
和Swin-B
模型,并观察2个问题:
- 训练不稳定
- 性能下降
图4
根据作者的实验分析,不同层之间权值的严格一致性是问题的主要原因。其中,权重共享
后的梯度ℓ2-范数
较大,在不同的Transformer Block
内出现波动,如图4所示。
图5
如图5所示,从CKA值可以看出,权重共享
后模型生成的Feature map与原模型的相关性较小。为了解决这些问题,受到电信领域多路复用技术
的启发,提出了一种新的Transformer Compression
技术,Weight Multiplexing
。它将多层权重组合为共享部分的单个权重,同时涉及转换
和蒸馏
,以增加权重的多样性。
更具体地说,如图2(b)所示,本文提出的权重复用方法包括:
- 在多个
Transformer Block
之间共享权重,可以认为在复用时是一个组合过程; - 在每一层引入转换来模拟解复用;
- 运用知识蒸馏增加压缩前后模型之间特征表示的相似性。
根据Eq.(4),可以将权重复用重新表述如下:
其中,表示第个Transformer Layer
中Transformer Block
的权值。请注意,中的参数数量远少于。
1、Weight Transformation
Weight Transformation
是强加于注意力矩阵和前馈网络。这种转换允许每一层都是不同的,从而提高了参数多样性和模型表示能力。
如图3(右)所示,各层之间不共享Transformation Kernels
的参数,而原始Transformer
中除了LayerNorm
以外的所有其他层都是共享的。由于共享层占据了模型参数的绝大部分,权重复用
后模型大小仅略有增加。
Transformation for MSA
为了提高参数的多样性,分别在Softmax
模块前后插入了2个线性变换。
其中,分别为Softmax
前后的线性变换。这种线性变换可以使每个注意力矩阵不同,同时结合注意力Head之间的信息来增加参数方差。
Transformation for MLP
另一方面,进一步对MLP进行了轻量级转换,以提高参数的多样性。特别地,设输入为,其中表示所有Token
的嵌入向量的第1个位置。然后引入线性变换将Y转换为,其中是线性层的独立权重矩阵。然后是等式3被重新表述为:
为了减少参数的数量并在变换中引入局域性,本文采用Depth-wise Convolutional
来稀疏和共享每个权重矩阵中的权值,并且参数量只有比(K<<N)小得多,其中K是卷积的kernel size。经过转换后,MLP的输出变得更加多样化,提高了参数的性能。
理论上,通过这些转换,权重共享
层可以恢复预训练模型的行为,类似于解复用过程
。这样可以缓解训练不稳定性和性能下降问题,因为这些问题在原始模型中没有观察到。类似的转换也被应用于提高没有权重共享
的Transformer
的性能,如Talking-heads Attention
和CeiT
。
2、Weight Distillation
为了压缩大型预训练模型并解决权重共享
导致的性能下降问题,作者进一步采用权重蒸馏法
将知识从大型模型转移到小型且紧凑的模型。考虑了Transformer Block
的3种蒸馏方法:
- Prediction-Logit Distillation
- Self-Attention Distillation
- Hidden-State Distillation
Prediction-Logit Distillation
Hinton等人证明了深度学习模型可以通过模仿训练过程中表现良好的教师模型的输出行为来获得更好的表现。作者利用这个想法来引入一个预测损失,如下所示:
其中,和分别是学生模型和教师模型预测的对数,T是控制对数平滑度的温度值。在实验中,设置T=1。CE表示交叉熵损失。
Self-Attention Distillation
最近有研究表明,利用Transformer Layer
中的注意力图来指导学生模型的训练是有效。为了解决学生模型和教师模型之间由于Head num
不同所导致的维度不一致问题,受Minilmv2
启发,在MSA
中对Query
、Key
和Value
之间的关系应用了交叉熵损失。
首先在所有的Head
上附加矩阵。例如,定义,以同样的方式定义。为了简化符号,分别用、和来分别表示Q
、K
和V
。然后,可以生成由定义的9个不同的关系矩阵。注意,是注意力矩阵a
,Self-Attention Distillation
损失可以表示为:
其中,表示的第n行。
Hidden-State Distillation
类似地,可以生成隐藏状态的关系矩阵,即由MLP输出的特征。用表示Transformer Layer
的隐藏状态,将基于关系矩阵的隐藏状态蒸馏损失定义为:
其中表示第n行,计算公式为。
总的蒸馏损失
根据作者的观察,仅使用prediction soft labels
比同时使用prediction soft labels
+ ground truth labels
能够产生更好的性能,因此最终蒸馏目标函数表示为:
其中,β和γ分别为超参数,默认值分别为1和0.1。
3.2 Compression Pipeline
阶段1:通过权重转换生成紧凑的架构
给定一个预训练的Vision Transformer
模型,首先在共享每个相邻Transformer Layer
中除了LayerNorm
以外的参数。然后通过在Softmax
层之前和之后插入一个线性层对每一层进行权值变换。此外,还为MLP
引入了一个Depth-wise Convolutional
。这些线性层和转换块的参数不共享。
阶段2:用Weight Distillation
训练压缩后的模型
在这一步中使用所提出的Weight Distillation
方法,将知识从大的预训练模型转移到小的模型。Transformer Block
内部的这种精馏使得学生网络能够再现教师网络的行为,从而从大规模的预训练模型中提取出更多有用的知识。
请注意,只有当教师模型和学生模型都是Transformer
架构时,才会执行此操作。在其他情况下,学生和教师的结构是异质的,只保留Prediction-Logit Distillation
。
4实验
4.1 消融实验
1、Weight Sharing
图4
如图4所示,DeiT-S
和Swin-B
权重共享
后梯度的ℓ2-范数较大,说明权值大小变化较快。此外,权重共享
还导致了不同层间的梯度范数的波动。这可能会导致不同的层的优化空间。特别是,一些层更新得很快,而其他部分几乎没有优化,这使得模型很可能收敛到一个糟糕的局部最优,甚至在训练中出现分歧。
因此,不同层间共享的严格相同的权重会导致训练的不稳定性。然而,权值复用
方法可以通过引入变换来提高参数多样性,从而降低梯度范数和提高层间的平滑度,促进更稳定的训练过程。
图5
在性能分析方面,在图5中对权值共享
和权重复用
之间的特征相似性与CKA进行了比较。CKA值越高,表示两个模型之间的特征表示越相似,从而获得相似的性能。可以观察到,DeiT
和Swin
在应用权重共享
后都存在很大的特征表示偏差,特别是在最后几层,这可能是权重共享
导致性能下降的原因之一。然而提出的权值复用
方法可以提高相似性。
2、Component-wise
采用权值共享
方法后,Swin-T
和DeiT-S
模型的参数数量减少了一半,精度也下降了2%。性能可以通过应用权重精馏
或权重转换
来提高,这表明了它们各自的有效性。值得注意的是权值变换
只引入了少量的参数。
此外,当结合这3个组件时,MiniViT
可以达到最好的精度,甚至超过原始的Swin-T
和DeiT-S
。
3、Number of Sharing Blocks
在表2中,在Swin-T
或DeiT-B
的每个阶段共享2个Block可以显著将参数量从28M减少到16M,86M减少到44M,而Top-1的精度提高了1%。在极端情况下,每个阶段的所有块都共享,Mini-Swin-T
仍然可以以43%的参数优于原始模型。Mini-DeiT-B
可以实现90%的参数降低,性能仅下降2%。
结果表明,Vision Transformer
存在参数冗余,所提出的MiniViT
可以有效地提高参数效率。此外,MiniViT
是可配置的,以满足模型大小和性能的各种要求。
4、Distillation Losses
如表3所示,与仅使用预测损失相比,额外的GT标签导致Swin的性能下降了0.3%,这是由于权值共享
导致学习能力下降。此外,还观察到,在应用Self-Attention Distillation
和Hidden-State Distillation
后,精确度提高了约0.2%,表明提出的Distillation
方法是有效性的。
4.2 SoTA分类
4.3 迁移学习
5参考
[1].MiniViT: Compressing Vision Transformers with Weight Multiplexing