[论文解读] Vision Transformers Need Registers

Timothée Darcet, Maxime Oquab+2 more|arXiv (Cornell University)|Sep 28, 2023

Domain Adaptation and Few-Shot Learning参考文献 32被引用 48

一句话总结

该论文识别出高范数异常令牌导致 Vision Transformer 特征图中的伪影，并表明在输入序列中加入可学习的寄存令牌可以消除这些伪影、并在监督、文本监督和自监督 ViTs 的密集预测与对象发现方面提升。

ABSTRACT

Transformers have recently emerged as a powerful tool for learning visual representations. In this paper, we identify and characterize artifacts in feature maps of both supervised and self-supervised ViT networks. The artifacts correspond to high-norm tokens appearing during inference primarily in low-informative background areas of images, that are repurposed for internal computations. We propose a simple yet effective solution based on providing additional tokens to the input sequence of the Vision Transformer to fill that role. We show that this solution fixes that problem entirely for both supervised and self-supervised models, sets a new state of the art for self-supervised visual models on dense visual prediction tasks, enables object discovery methods with larger models, and most importantly leads to smoother feature maps and attention maps for downstream visual processing.

研究动机与目标

• 在有监督和自监督模型中识别 ViT 特征图中的伪影。
• 表征高范数异常令牌的性质及其所在位置。
• 提出使用寄存令牌的简单架构修复以减轻伪影。
• 展示密集预测性能的提升以及更平滑的特征/注意力图。

提出的方法

• 分析注意力图伪影并将高范数令牌定义为异常值（范数 >150）。
• 在训练过程中及不同模型尺寸下，检查异常值出现的时间和位置。
• 通过线性模型对 patch embedding 评估高范数令牌的信息含量。
• 在 patch embedding 之后添加 N 个可学习的输入令牌（寄存令牌）以吸收异常值。
• 通过线性探测、零-shot 分类以及有/无寄存令牌的对象发现评估对下游任务的影响。
• 对寄存令牌数量进行消融研究以探究性能和伪影减少。

实验结果

研究问题

• RQ1 在不同监督类型下，ViT 注意力图中的伪影如何产生？
• RQ2 高范数异常令牌携带局部信息还是全局信息，它们如何影响密集预测任务？
• RQ3 添加寄存令牌是否能在不损害下游性能的前提下消除伪影，适用于各种训练范式？
• RQ4 寄存令牌数量如何影响伪影抑制及下游任务性能？

主要发现

• 伪影对应于一小部分令牌（约 2%）具有远高于其他令牌的范数（约 10 倍），并在对大型 ViTs 进行大量训练后出现在中间层。
• 高范数令牌包含关于 patch 位置信息和像素的局部信息更少，但存储了关于图像的全局信息更多。
• 在 patch embedding 之后追加可学习的寄存令牌可以消除高范数异常值并平滑特征/注意力图。
• 使用寄存令牌训练的模型在 ImageNet 线性探测、ADE20k 分割或 NYUd 深度估计上没有性能回归；某些情况略有提升。
• 基于寄存的模型在 DeiT-III、OpenCLIP 和 DINOv2 主干上实现改进的无监督对象发现（如 LOST）；一个寄存令牌即可消除伪影，更多寄存令牌可能有助于密集任务。