Vision Transformers с самообучаемыми регистрами

Vision Transformers with Self-Distilled Registers

May 27, 2025

Авторы: Yinjie Chen, Zipeng Yan, Chong Zhou, Bo Dai, Andrew F. Luo

cs.AI

Аннотация

Трансформеры для обработки изображений (Vision Transformers, ViTs) стали доминирующей архитектурой для задач визуальной обработки, демонстрируя превосходную масштабируемость с увеличением объема обучающих данных и размера модели. Однако в последних исследованиях было выявлено появление артефактных токенов в ViTs, которые не соответствуют локальной семантике. Эти аномальные токены ухудшают производительность ViTs в задачах, требующих точной локализации или структурной согласованности. Эффективным способом устранения этой проблемы является добавление регистровых токенов в ViTs, которые неявно "поглощают" артефактные термины в процессе обучения. Учитывая доступность различных крупномасштабных предобученных ViTs, в данной работе мы стремимся оснастить их такими регистровыми токенами без необходимости их повторного обучения с нуля, что нецелесообразно из-за их размера. В частности, мы предлагаем метод Post Hoc Registers (PH-Reg), эффективный метод самодистилляции, который интегрирует регистры в существующий ViT без необходимости дополнительных размеченных данных и полного переобучения. PH-Reg инициализирует как учителя, так и ученика из одного и того же предобученного ViT. Учитель остается замороженным и неизменным, в то время как ученик дополняется случайно инициализированными регистровыми токенами. Применяя аугментацию на этапе тестирования к входам учителя, мы генерируем очищенные плотные эмбеддинги, свободные от артефактов, которые затем используются для оптимизации только небольшого подмножества разблокированных весов ученика. Мы показываем, что наш подход может эффективно сократить количество артефактных токенов, улучшая сегментацию и предсказание глубины ученика ViT в условиях zero-shot и линейного зондирования.

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.