采用超令牌采样的视觉 Transformer

视觉 transformer 已在许多视觉任务上取得令人印象深刻的性能。然而，其在浅层捕获局部特征时可能遭受高度冗余。因此人们利用局部自注意力或早期卷积，但这牺牲了捕获长程依赖的能力。由此产生一个挑战：我们能否在神经网络的早期阶段实现高效且有效的全局上下文建模？为解决此问题，我们从超像素的设计中汲取灵感（超像素减少了后续处理中的图像基元数量），并将超令牌引入视觉 transformer。超令牌试图提供语义上有意义的视觉内容镶嵌，从而减少自注意力中的令牌数量并保留全局建模。具体而言，我们提出一种简单而强大的超令牌注意力（STA）机制，包含三个步骤：首先通过稀疏关联学习从视觉令牌中采样超令牌，其次在超令牌上执行自注意力，最后将其映射回原始令牌空间。STA 将普通全局注意力分解为稀疏关联图与低维注意力的乘积，从而在捕获全局依赖时实现高效率。基于 STA，我们开发了一种分层视觉 transformer。大量实验证明了其在各种视觉任务上的强大性能。特别是在无任何额外训练数据或标签的情况下，它以少于 100M 参数在 ImageNet-1K 上达到 86.4% 的 top-1 准确率。它还在 COCO 检测任务上达到 53.9 box AP 和 46.8 mask AP，在 ADE20K 语义分割任务上达到 51.9 mIOU。代码发布于 https://github.com/hhb072/STViT。