【CV Transformer 论文笔记】PS-ViT: Vision Transformer with Progressive Sampling

论文地址：https://openaccess.thecvf.com/content/ICCV2021/papers/Yue_Vision_Transformer_With_Progressive_Sampling_ICCV_2021_paper.pdf
项目地址：https://github.com/yuexy/PS-ViT

ViT直接将纯Transformer架构应用于图像分类，通过简单地将图像分成固定长度的Tokens，并使用Transformer来学习这些tokens之间的关系。然而，这种朴素的token化可能会破坏对象结构，将grid分配给背景等不感兴趣的区域，并引入干扰信号。（下图a）
为了缓解上述问题，本文提出了一种迭代渐进采样策略来定位判别区域。不从固定位置采样，而是在每次迭代中，当前采样步骤的嵌入被馈送到Transformer编码器层，并且预测一组采样偏移以更新下一步的采样位置。（下图b）

渐进采样（progressive sampling）是可微的，与 ViT 结合后可以自适应地学习关注图像的哪些部分。
由于训练是任务驱动的，PS-ViT 倾向于对与语义结构相关的对象区域进行采样，即更关注前景对象而较少关注模糊的背景。
PS-ViT 减轻了 ViT 中的结构破坏问题，并自适应地关注对象的有兴趣的区域。
在 ImageNet 上从头开始训练时，PS-ViT 在 top-1 精度方面比普通 ViT 高 3.8%，参数减少了约 4 倍，FLOP 减少了 10 倍。

Progressive Sampling

渐进采样模块是一个迭代框架，如下图所示。

在每次迭代中，通过将采样位置与上次迭代的偏移向量相加来更新采样位置：

每次迭代的初始标记在采样位置的输入特征图上采样，如下所示：

其中 T′t表示迭代 t 的初始采样标记。 由于 pt 的元素是小数，采样是通过双线性插值操作实现的，该操作是可微的。
将初始采样标记、上一次迭代的输出标记和当前采样位置的位置编码按元素相加，以获得当前迭代的输出标记，然后被馈送到Transformer编码器，如下公式所示：

其中Wt是将采样位置pt投影到位置编码矩阵Pt的线性变换，所有迭代共享相同的Wt。
ViT 使用补丁索引的二维正弦位置嵌入。然而，在PS-ViT中由于采样位置是非等距的，因此将采样位置的归一化绝对坐标投影到一个嵌入空间作为位置嵌入。采样位置偏移量预测如下：

Mt是可学习线性变换矩阵。

整体架构

如图所示，PS-ViT 的架构由4个主要组件组成：

1. 特征提取器模块来预测密集Tokens；
2. 渐进式采样模块，用于对判别位置进行采样；
3. vision transformer模块，遵循ViT [12]和DeiT [38]的类似配置；
4. 分类模块。

因为卷积算子在空间局部上下文建模方面特别有效，作者使用了ResNet50第一阶段的卷积骨干和前两个残差块作为特征提取器模块。

vision transformer模块遵循 ViT [12] 和 DeiT [38] 中采用的架构。
同时添加了一个额外的分类标记Tcls到渐进采样模块中最后一次迭代的输出标记TN中，并将它们馈送到vision transformer模块。形式上，

由于位置信息已经在渐进采样模块中融合到TN中，所以不需要再添加位置嵌入。通过ViT模块细化的分类标记最终用于预测图像类别。
端到端训练使用交叉熵损失。