Transformadores de Visión con Registros de Auto-Destilación

Vision Transformers with Self-Distilled Registers

May 27, 2025

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

cs.AI

Resumen

Los Transformers de Visión (ViTs) han surgido como la arquitectura dominante para tareas de procesamiento visual, demostrando una excelente escalabilidad con el aumento de datos de entrenamiento y tamaño del modelo. Sin embargo, trabajos recientes han identificado la aparición de tokens de artefactos en los ViTs que son incongruentes con la semántica local. Estos tokens anómalos degradan el rendimiento de los ViTs en tareas que requieren localización fina o coherencia estructural. Una mitigación efectiva de este problema es la adición de tokens de registro a los ViTs, que implícitamente "absorben" el término de artefacto durante el entrenamiento. Dada la disponibilidad de varios ViTs preentrenados a gran escala, en este artículo buscamos equiparlos con dichos tokens de registro sin la necesidad de reentrenarlos desde cero, lo cual es inviable considerando su tamaño. Específicamente, proponemos Registros Post Hoc (PH-Reg), un método eficiente de auto-distilación que integra registros en un ViT existente sin requerir datos etiquetados adicionales ni un reentrenamiento completo. PH-Reg inicializa tanto la red maestra como la red estudiante a partir del mismo ViT preentrenado. La red maestra permanece congelada y sin modificaciones, mientras que la red estudiante se aumenta con tokens de registro inicializados aleatoriamente. Al aplicar aumentación en tiempo de prueba a las entradas de la red maestra, generamos embeddings densos libres de artefactos, que luego se utilizan para optimizar solo un pequeño subconjunto de pesos desbloqueados de la red estudiante. Demostramos que nuestro enfoque puede reducir efectivamente el número de tokens de artefactos, mejorando la segmentación y predicción de profundidad del ViT estudiante bajo evaluación zero-shot y linear probing.

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.