自蒸馏寄存器视觉Transformer
Vision Transformers with Self-Distilled Registers
May 27, 2025
作者: Yinjie Chen, Zipeng Yan, Chong Zhou, Bo Dai, Andrew F. Luo
cs.AI
摘要
视觉Transformer(ViTs)已成为视觉处理任务的主导架构,展现出随着训练数据和模型规模增加而优异的扩展性。然而,近期研究揭示了ViTs中出现的与局部语义不符的伪影标记,这些异常标记在需要精细定位或结构一致性的任务中降低了ViT的性能。一种有效的缓解方法是在ViTs中加入寄存器标记,这些标记在训练过程中隐式地“吸收”伪影。鉴于多种大规模预训练ViTs的可用性,本文旨在无需从头重新训练(考虑到其规模,这并不可行)的情况下,为它们配备此类寄存器标记。具体而言,我们提出了事后寄存器(PH-Reg),一种高效的自蒸馏方法,它能在无需额外标注数据和完全重新训练的情况下,将寄存器整合到现有ViT中。PH-Reg从同一预训练ViT初始化教师和学生网络,教师网络保持冻结且未修改,而学生网络则通过随机初始化的寄存器标记进行增强。通过对教师网络的输入应用测试时增强,我们生成无伪影的降噪密集嵌入,随后仅用于优化学生网络中一小部分解锁的权重。我们证明,该方法能有效减少伪影标记的数量,在零样本和线性探测下提升学生ViT的分割和深度预测性能。
English
Vision Transformers (ViTs) have emerged as the dominant architecture for
visual processing tasks, demonstrating excellent scalability with increased
training data and model size. However, recent work has identified the emergence
of artifact tokens in ViTs that are incongruous with the local semantics. These
anomalous tokens degrade ViT performance in tasks that require fine-grained
localization or structural coherence. An effective mitigation of this issue is
to the addition of register tokens to ViTs, which implicitly "absorb" the
artifact term during training. Given the availability of various large-scale
pre-trained ViTs, in this paper we aim at equipping them with such register
tokens without the need of re-training them from scratch, which is infeasible
considering their size. Specifically, we propose Post Hoc Registers (PH-Reg),
an efficient self-distillation method that integrates registers into an
existing ViT without requiring additional labeled data and full retraining.
PH-Reg initializes both teacher and student networks from the same pre-trained
ViT. The teacher remains frozen and unmodified, while the student is augmented
with randomly initialized register tokens. By applying test-time augmentation
to the teacher's inputs, we generate denoised dense embeddings free of
artifacts, which are then used to optimize only a small subset of unlocked
student weights. We show that our approach can effectively reduce the number of
artifact tokens, improving the segmentation and depth prediction of the student
ViT under zero-shot and linear probing.Summary
AI-Generated Summary