Vision Transformers mit selbst-distillierten Registern

Vision Transformers with Self-Distilled Registers

May 27, 2025

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

cs.AI

Zusammenfassung

Vision Transformers (ViTs) haben sich als dominante Architektur für visuelle Verarbeitungsaufgaben etabliert und zeigen eine hervorragende Skalierbarkeit bei zunehmenden Trainingsdaten und Modellgrößen. Jüngste Arbeiten haben jedoch das Auftreten von Artefakt-Tokens in ViTs identifiziert, die nicht mit der lokalen Semantik übereinstimmen. Diese anomalen Tokens beeinträchtigen die Leistung von ViTs bei Aufgaben, die eine feinkörnige Lokalisierung oder strukturelle Kohärenz erfordern. Eine effektive Lösung für dieses Problem ist die Hinzufügung von Register-Tokens zu ViTs, die während des Trainings implizit die Artefakt-Terme „absorbieren“. Angesichts der Verfügbarkeit verschiedener großskaliger, vortrainierter ViTs zielen wir in dieser Arbeit darauf ab, diese mit solchen Register-Tokens auszustatten, ohne sie von Grund auf neu trainieren zu müssen, was angesichts ihrer Größe nicht praktikabel ist. Konkret schlagen wir Post Hoc Registers (PH-Reg) vor, eine effiziente Selbst-Distillationsmethode, die Register in einen bestehenden ViT integriert, ohne zusätzliche gelabelte Daten und vollständiges Retraining zu erfordern. PH-Reg initialisiert sowohl das Lehrer- als auch das Schüler-Netzwerk aus demselben vortrainierten ViT. Der Lehrer bleibt eingefroren und unverändert, während der Schüler mit zufällig initialisierten Register-Tokens erweitert wird. Durch die Anwendung von Test-Time-Augmentation auf die Eingaben des Lehrers erzeugen wir denoisierte, dichte Embeddings ohne Artefakte, die dann verwendet werden, um nur eine kleine Teilmenge der entsperrten Gewichte des Schülers zu optimieren. Wir zeigen, dass unser Ansatz die Anzahl der Artefakt-Tokens effektiv reduzieren kann und die Segmentierung und Tiefenvorhersage des Schüler-ViTs unter Zero-Shot- und Linear-Probing-Bedingungen verbessert.

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.