Vision Transformers avec Registres d'Auto-Distillation

Vision Transformers with Self-Distilled Registers

May 27, 2025

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

cs.AI

Résumé

Les Vision Transformers (ViTs) sont devenus l'architecture dominante pour les tâches de traitement visuel, démontrant une excellente scalabilité avec l'augmentation des données d'entraînement et de la taille des modèles. Cependant, des travaux récents ont identifié l'émergence de tokens d'artefacts dans les ViTs qui sont incohérents avec la sémantique locale. Ces tokens anormaux dégradent les performances des ViTs dans les tâches nécessitant une localisation fine ou une cohérence structurelle. Une atténuation efficace de ce problème consiste à ajouter des tokens de registre aux ViTs, qui absorbent implicitement les termes d'artefacts pendant l'entraînement. Étant donné la disponibilité de divers ViTs pré-entraînés à grande échelle, cet article vise à les équiper de tels tokens de registre sans nécessiter de ré-entraînement complet, ce qui est irréalisable compte tenu de leur taille. Plus précisément, nous proposons Post Hoc Registers (PH-Reg), une méthode d'auto-distillation efficace qui intègre des registres dans un ViT existant sans nécessiter de données étiquetées supplémentaires ni de ré-entraînement complet. PH-Reg initialise les réseaux enseignant et étudiant à partir du même ViT pré-entraîné. L'enseignant reste figé et inchangé, tandis que l'étudiant est augmenté avec des tokens de registre initialisés aléatoirement. En appliquant une augmentation au moment du test aux entrées de l'enseignant, nous générons des embeddings denses débruités exempts d'artefacts, qui sont ensuite utilisés pour optimiser uniquement un petit sous-ensemble de poids déverrouillés de l'étudiant. Nous montrons que notre approche peut efficacement réduire le nombre de tokens d'artefacts, améliorant la segmentation et la prédiction de profondeur du ViT étudiant en mode zero-shot et par sondage linéaire.

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.