【ViT 论文笔记】Vision Transformer for Small-Size Datasets

论文地址：https://arxiv.org/abs/2112.13492
项目地址：

将 Transformer 结构应用于图像分类任务的 ViT 的性能优于卷积神经网络。 然而，ViT 的高性能源于使用大型数据集（如 JFT-300M）进行预训练，其对大型数据集的依赖被认为是源于其低局部性归纳偏差。
本文提出了 Shifted Patch Tokeniza-tion (SPT) 和 Locality Self-Attention (LSA)，来解决了缺乏局部性归纳偏差的问题，即使在小型数据集上也能从头开始学习。 此外，SPT 和 LSA 是通用且有效的附加模块，可轻松应用于各种 ViT。
实验结果表明，当SPT和LSA同时应用于ViTs时，在具有代表性的小尺寸数据集Tiny-ImageNet中，性能平均提升了2.96%。 尤其是 Swin Transformer 取得了压倒性的 4.08% 的性能提升。

作者归纳了ViT目前存在的两个问题：

1. 糟糕的标记化。 ViT 将给定​​图像划分为大小相等的不重叠块，并将每个块使用相同的线性投影到视觉标记。因此，ViT 的标记化具有排列不变性，这使得能够很好地嵌入补丁之间的关系 。但是非重叠块使得视觉标记具有相对较小的感受野，而较小的感受野会导致 ViT只用很少的像素进行标记，因此，与相邻像素的空间关系没有充分嵌入到每个视觉标记中。
2. 注意力机制存在问题。图像数据的特征维度远大于自然语言和音频信号，因此嵌入的tokens的数量必然很大。因此，token的注意力分布变得平滑。换句话说，ViT 无法在局部关注重要的视觉令牌。上述两个主要问题导致高度冗余的注意力无法集中在目标类上。这种多余的注意力使 ViT很容易将注意力集中在背景上，而不能很好地捕捉目标类的形状。

本文提出了两种解决方案以有效改善 ViT 从零开始学习小型数据集的局部归纳偏差：

1. 移位补丁标记化（SPT），以在标记化过程中进一步利用相邻像素之间的空间关系。 SPT 的思想来源于时间移位模块（TSM）。TSM是一种有效的时间建模，它可以转换一些特征的时间通道。受此启发，本文提出了有效的空间建模，将空间移位的图像与输入图像一起标记。 SPT可以为 ViT 提供比标准标记化更广泛的感受野，具有通过在每个视觉标记中嵌入更多空间信息来增加局部归纳偏差的效果。
2. Locality Self-Attention (LSA)，允许 ViT 在局部实施注意力。 LSA通过排除自我标记和将可学习温度应用于 softmax 函数来减轻注意力分数分布的平滑现象。LSA通过强制每个标记更多地关注与自身有较大关系的标记来诱导注意力局部工作。

所提出的 SPT 和 LSA 可以很容易地以附加模块的形式应用于各种 ViT，而无需更改结构，并且可以有效地提高性能。

具体而言，

对于SPT（如上图左侧图例），本文作者在左上、左下、右上、右下4个方向上对输入图像进行空间移位半个patch size，并与输入图像连接。接着像标准 ViT 一样应用补丁分区，然后依次执行patch flattening、层归一化 [2] 和线性投影三个过程。

LSA （如上图右侧图例）主要通过学习 softmax 函数的温度参数来锐化注意力分数的分布（一般来说，softmax 函数可以通过温度缩放来控制输出分布的平滑度 [11]）。此外，通过应用**对角屏蔽（Diagonal Masking）**来去除自标记（self-token）关系，可以抑制由Query和Key计算的相似度矩阵的对角分量，相对增加了不同token之间的注意力分数，使得注意力分数的分布更加清晰，具体来说，对角掩蔽就是将自身token赋值-∞ ，而其他token不变。

• 个人总结

SPT实际上是把四个对角方向的位移patch加入到编码中来获得更大的局部信息，原理上十分简单粗暴，个人感觉和重叠patch的想法是差不多的，而且是用计算量来换取的大的感受野，但是由于取的5个patch（4个方向+原patch）都包含部分重叠的部分，所以实际上比重叠patch可以更好的学到patch内部像素间关系。

而softmax 函数的温度参数在可学习而非固定时，会低于标准ViT的温度，并且随着深度的增加会有上升的趋势。而温度会影响Token分布的平滑程度（这方面不是特别了解），个人感觉是过于平滑不同的patch的差异性就没有那么明显，因此尖锐一点的分布有利于区分不同的token从而实现对重要token的重点关注。