[Paper Review] Vision Transformers Need Registers

Timothée Darcet, Maxime Oquab+2 more|arXiv (Cornell University)|Sep 28, 2023

Domain Adaptation and Few-Shot LearningComputer Science32 references48 citations

TL;DR

The paper identifies high-norm outlier tokens causing artifacts in Vision Transformer feature maps and shows that adding learnable register tokens to the input sequence removes these artifacts and improves dense prediction and object discovery across supervised, text-supervised, and self-supervised ViTs.

ABSTRACT

Transformers have recently emerged as a powerful tool for learning visual representations. In this paper, we identify and characterize artifacts in feature maps of both supervised and self-supervised ViT networks. The artifacts correspond to high-norm tokens appearing during inference primarily in low-informative background areas of images, that are repurposed for internal computations. We propose a simple yet effective solution based on providing additional tokens to the input sequence of the Vision Transformer to fill that role. We show that this solution fixes that problem entirely for both supervised and self-supervised models, sets a new state of the art for self-supervised visual models on dense visual prediction tasks, enables object discovery methods with larger models, and most importantly leads to smoother feature maps and attention maps for downstream visual processing.

Motivation & Objective

• Identify artifacts in ViT feature maps across supervised and self-supervised models.
• Characterize the properties and loci of high-norm outlier tokens.
• Propose a simple architectural fix using register tokens to mitigate artifacts.
• Demonstrate improved dense prediction performance and smoother feature/attention maps.

Proposed method

• Analyze attention map artifacts and define high-norm tokens as outliers (norm > 150).
• Examine when and where outliers appear during training and across model sizes.
• Probe information content of high-norm tokens via linear models on patch embeddings.
• Introduce N learnable input tokens (registers) added after patch embedding to absorb outliers.
• Evaluate impact on downstream tasks via linear probing, zero-shot classification, and object discovery with and without registers.
• Ablate the number of registers to study performance and artifact reduction.

Experimental results

Research questions

• RQ1 What causes artifacts in ViT attention maps across supervision types?
• RQ2 Do high-norm outlier tokens carry local or global information, and how do they affect dense prediction tasks?
• RQ3 Can adding register tokens remove artifacts without harming downstream performance across training paradigms?
• RQ4 How does the number of registers influence artifact suppression and downstream task performance?

Key findings

• Artifacts correspond to a small fraction (~2%) of tokens with much higher norms (≈10x) and appear in middle layers after substantial training of large ViTs.
• High-norm tokens contain less local information about patch position and pixels, but store more global information about the image.
• Appending learnable register tokens after patch embedding eliminates high-norm outliers and smooths feature/attention maps.
• Models trained with registers show no performance regression on ImageNet linear probing, ADE20k segmentation, or NYUd depth estimation; some cases improve slightly.
• Register-based models enable improved unsupervised object discovery (e.g., LOST) across DeiT-III, OpenCLIP, and DINOv2 backbones; one register suffices to remove artifacts, with more registers potentially aiding dense tasks.