Vision Transformers with Self-Distilled Registers

Bo Dai (Google DeepMind & Georgia Tech) · Andrew Luo (University of Hong Kong) · Zipeng Yan (University of Hong Kong) · Yinjie Chen (Zhejiang University) · Chong Zhou (Nanyang Technological University)
artifact tokensdenoised dense embeddingsdepth predictionfine-grained localizationlinear probinglocal semanticsmodel scalabilitypost hoc registersregister tokenssegmentationself-distillationstructural coherenceteacher-student networkstest-time augmentationvision transformerszero-shot learning

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 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 the addition of register tokens to ViTs, which implicitly ''absorb'' the artifact term during training. Given the availability of existing large-scale pre-trained ViTs, in this paper we seek add register tokens to existing models without needing to re-train 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.