Vision Transformers met Zelf-Gedistilleerde Registers

Vision Transformers with Self-Distilled Registers

May 27, 2025

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

cs.AI

Samenvatting

Vision Transformers (ViTs) zijn naar voren gekomen als de dominante architectuur voor visuele verwerkingstaken, waarbij ze uitstekende schaalbaarheid demonstreren met toenemende trainingsdata en modelgrootte. Recent onderzoek heeft echter het ontstaan van artefacttokens in ViTs geïdentificeerd die niet overeenkomen met de lokale semantiek. Deze afwijkende tokens verminderen de prestaties van ViTs bij taken die fijnmazige lokalisatie of structurele samenhang vereisen. Een effectieve oplossing voor dit probleem is het toevoegen van registertokens aan ViTs, die tijdens de training impliciet de artefacttermen "absorberen". Gezien de beschikbaarheid van verschillende grootschalige vooraf getrainde ViTs, is het doel van dit artikel om deze uit te rusten met dergelijke registertokens zonder ze vanaf nul opnieuw te moeten trainen, wat gezien hun omvang onhaalbaar is. Specifiek stellen we Post Hoc Registers (PH-Reg) voor, een efficiënte zelfdistillatiemethode die registers integreert in een bestaande ViT zonder extra gelabelde data en volledige hertraining nodig te hebben. PH-Reg initialiseert zowel het leraren- als het studentennetwerk vanuit dezelfde vooraf getrainde ViT. De leraar blijft bevroren en ongemodificeerd, terwijl de student wordt uitgebreid met willekeurig geïnitialiseerde registertokens. Door test-time augmentatie toe te passen op de inputs van de leraar, genereren we gedenoiseerde dichte embeddings zonder artefacten, die vervolgens worden gebruikt om alleen een kleine subset van ontgrendelde studentgewichten te optimaliseren. We laten zien dat onze aanpak effectief het aantal artefacttokens kan verminderen, waardoor de segmentatie en dieptevoorspelling van de student-ViT verbetert onder zero-shot en lineaire 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.