具有自蒸馏寄存器的视觉 Transformer

视觉 Transformer (ViT) 已成为视觉处理任务的主导架构，展现出随训练数据和模型规模增加而提升的优异可扩展性。然而，近期研究发现了 ViT 中出现的伪影 token，这些 token 与局部语义不一致。这些异常 token 降低了 ViT 在需要细粒度定位或结构连贯性的任务中的性能。缓解该问题的有效方法是向 ViT 添加寄存器 token，以在训练过程中隐式“吸收”伪影项。鉴于现有大规模预训练 ViT 的可用性，在本文中，我们寻求向现有模型添加寄存器 token 而无需从头重新训练，考虑到其规模，这是不可行的。具体而言，我们提出了事后寄存器 (PH-Reg)，这是一种高效的自蒸馏方法，无需额外标注数据和完全重新训练即可将寄存器集成到现有 ViT 中。PH-Reg 从相同的预训练 ViT 初始化教师和学生网络。教师网络保持冻结且未修改，而学生网络通过随机初始化的寄存器 token 进行增强。通过对教师网络的输入应用测试时增强，我们生成无伪影的去噪密集嵌入，随后仅用于优化学生网络中一小部分未锁定权重。我们表明，我们的方法能有效减少伪影 token 数量，提升学生 ViT 在零样本和线性探测下的分割与深度预测性能。