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

发表
Andrew LuoAndrew Luo 提交
作者: Yinjie Chen, Zipeng Yan, Chong Zhou, Bo Dai, Andrew LuoAndrew F. Luo

摘要

Vision Transformers (ViTs) 已成为视觉处理任务的主流架构,随着训练数据和模型尺寸的增加,展现出卓越的可扩展性。然而,最近的研究发现 ViTs 中出现了与局部语义不一致的伪影 tokens。这些异常 tokens 会降低 ViT 在需要细粒度定位或结构连贯性任务中的性能。解决此问题的一个有效方法是在 ViTs 中添加寄存器 tokens,它们在训练期间隐式地"吸收"伪影项。考虑到现有各种大型预训练 ViTs 的可用性,在本文中,我们的目标是在无需从头开始重新训练(考虑到其尺寸是不可行的)的情况下,为其配备此类寄存器 tokens。具体来说,我们提出了 Post Hoc Registers (PH-Reg),一种高效的自蒸馏方法,它将寄存器集成到现有 ViT 中,无需额外的标记数据和完全重新训练。PH-Reg 使用相同的预训练 ViT 初始化教师和学生网络。教师保持冻结且未经修改,而学生则添加了随机初始化的寄存器 tokens。通过对教师的输入应用测试时增强,我们生成去噪的、无伪影的密集嵌入,然后仅使用这些嵌入来优化解锁的学生权重的一小部分。我们展示了我们的方法可以有效减少伪影 tokens 的数量,提高学生 ViT 在零样本和线性探测下的分割和深度预测性能。
查看 arXiv 页面查看 PDF

评论

Andrew LuoAndrew Luo
论文作者
论文提交者

视觉 Transformer (ViT) 已成为视觉处理任务的主导架构,随着训练数据和模型规模的增加,表现出卓越的可扩展性。然而,最近的研究发现 ViT 中出现了与局部语义不一致的伪影标记(artifact tokens)。这些异常标记会降低 ViT 在需要精细定位或结构连贯性的任务中的性能。解决此问题的有效方法是向 ViT 添加寄存器标记(register tokens),这些标记在训练期间会隐式地“吸收”伪影项。考虑到各种大型预训练 ViT 的可用性,在本文中,我们的目标是在不从头开始重新训练的情况下为其配备此类寄存器标记,考虑到其规模,从头训练是不可行的。具体来说,我们提出了 Post Hoc Registers (PH-Reg),这是一种高效的自蒸馏方法,可以在不需要额外标注数据和完全重新训练的情况下,将寄存器集成到现有 ViT 中。PH-Reg 从同一个预训练 ViT 初始化教师网络和学生网络。教师网络保持冻结和不变,而学生网络则添加了随机初始化的寄存器标记。通过对教师网络的输入应用测试时增强,我们生成了无伪影的去噪密集嵌入,然后仅使用这些嵌入来优化少量未锁定的学生网络权重。我们表明,我们的方法可以有效减少伪影标记的数量,提高学生 ViT 在零样本和线性探查下的分割和深度预测性能。