Transformers de Visão com Registros de Auto-Distilação

Vision Transformers with Self-Distilled Registers

May 27, 2025

Autores: Yinjie Chen, Zipeng Yan, Chong Zhou, Bo Dai, Andrew F. Luo

cs.AI

Resumo

Os Vision Transformers (ViTs) emergiram como a arquitetura dominante para tarefas de processamento visual, demonstrando excelente escalabilidade com o aumento de dados de treinamento e tamanho do modelo. No entanto, trabalhos recentes identificaram o surgimento de tokens de artefato em ViTs que são incongruentes com a semântica local. Esses tokens anômalos degradam o desempenho dos ViTs em tarefas que exigem localização refinada ou coerência estrutural. Uma mitigação eficaz desse problema é a adição de tokens de registro aos ViTs, que implicitamente "absorvem" o termo de artefato durante o treinamento. Dada a disponibilidade de diversos ViTs pré-treinados em grande escala, neste artigo buscamos equipá-los com tais tokens de registro sem a necessidade de retreiná-los do zero, o que é inviável considerando seu tamanho. Especificamente, propomos o Post Hoc Registers (PH-Reg), um método eficiente de auto-distilação que integra registros a um ViT existente sem exigir dados rotulados adicionais e retreinamento completo. O PH-Reg inicializa tanto a rede professora quanto a estudante a partir do mesmo ViT pré-treinado. A professora permanece congelada e inalterada, enquanto a estudante é aumentada com tokens de registro inicializados aleatoriamente. Ao aplicar aumento de dados em tempo de teste às entradas da professora, geramos embeddings densos denoizados livres de artefatos, que são então usados para otimizar apenas um pequeno subconjunto de pesos desbloqueados da estudante. Mostramos que nossa abordagem pode reduzir efetivamente o número de tokens de artefato, melhorando a segmentação e a previsão de profundidade do ViT estudante em cenários de zero-shot e sondagem linear.

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.