DeepViT: Towards Deeper Vision Transformer
发现问题
通过余弦相似度来衡量两个注意力图的相似度。发现ViT 中:
- 不同block中的注意力图相似,层数越深越相似,甚至一样。
- block内部,不同head的注意力图是多样性的。
进一步实验:通过将最后几个block的Q和K设置为一样的值,以使注意力图一样。测试发现掉点不严重,这更加说明了,最后深层的注意力图很相似。表明,注意力坍塌 (attention collapse)确实是存在的。
解决方法(1,2,3,4)
方法1:增加嵌入维度embedding dimension
- 原理:增加每个token的表达能力,使得生成的注意力图可以更加多样化,每个块的注意力图之间的相似度可以减少。
- 实验得到:随着embedding dimension的增长,相似的block的数量在下降,同时模型的Acc在上升。
- 该方法不足:增加embedding dimension也会显著增加计算成本,带来的性能改进往往会减少,且需要更大的数据量来训练,增加了过拟合的风险。
方法2(本文核心):Re-attention 模块:Re-attention以一种可学习的方式交换来自不同attention heads的信息来再生注意力图
-
普通的self-attention的attention map的特点是:只学习到一些局部的patch之间的relationship,而且在网络的深层时attention map很接近。
-
Re-attention的attention map的特点是:学习到更大范围内的patch之间的relationship,而且在网络的深层时attention map有差别。
实现代码:只是在ViT的基础上,多了红色框部分而已。
方法3:给self-attention加一个temperature λ
-
原理: λ的减小可以让attention map的分布变得更加的sharp,使得attention map的数值拉开,在实际实现的时候作者使得 λ的值随着层数的加深而逐渐衰减,或者干脆把 λ设置为可学习的参数。
-
不足:只起到非常小的作用。
方法4:随机drop掉attention map中的一些值。
-
原理:因为对于不同的block来说,会随机drop掉不同的值,故attention maps之间的相似性可以得到减少。
-
不足:只起到非常小的作用。
总结
DeepViT能够有效缓解注意力坍塌,使得ViT模型的性能可以随着层数的加深而增加。