This video presents an overview of an ICLR 2024 best paper on "Vision Transformer Models with Registers," addressing a significant issue in traditional Vision Transformer (ViT) architectures. The problem identified is the presence of "artifacts" in attention heat maps, characterized by outlier tokens with disproportionately high L2 norms, frequently found in background regions rather than salient objects. These artifacts, which emerge in larger, well-trained models, are shown to convey minimal local information but paradoxically store global image class data. The proposed hypothesis is that ViTs implicitly "recycle" low-information background tokens to aggregate global context. To combat this, the paper introduces a simple architectural modification: appending learnable "register tokens" to the input sequence. This minor tweak effectively eliminates the artifacts, producing significantly smoother attention maps, and dramatically improves performance in dense prediction tasks like object detection (nearly doubling accuracy for some models like DeiT and DinoV2), while maintaining or slightly enhancing image classification capabilities. Further analysis reveals that these registers learn to attend to diverse and specific elements within an image, offering a dedicated mechanism for global information processing.
Vision Transformer Models with Registers: Addressing Attention Artifacts
0:00:07 Introduction to ViT Models with Registers: The video introduces a topic based on an ICLR 2024 best paper, focusing on enhancing Vision Transformer (ViT) models.
0:00:18 Problem with Traditional ViTs: Attention Artifacts: Standard ViT models (e.g., DeiT, OpenCLIP, DinoV2) often exhibit "peaky outlier values" in their attention heat maps. These artifacts appear in background areas and do not contribute to focusing on the main objects in an image.
These artifacts are tokens (image patches) with output L2 norms approximately ten times higher than regular patches.
They lead to a bimodal distribution of norms, with outlier patches exceeding a norm of 150 (observed in 2.37% of DinoV2 patches).
Artifacts typically emerge in the middle layers of the model (around layer 15 and beyond), after substantial training, and exclusively in larger ViT models (Large, Huge sizes, not Tiny, Small, or Base).
They are found in patches visually similar to their neighbors, suggesting they carry little distinctive local information.
0:06:37 Local vs. Global Information in Artifacts:
Artifact patches perform poorly in tasks requiring local information, such as position prediction and pixel reconstruction (showing high L2 error, indicating local data loss).
Conversely, they demonstrate high accuracy in global image classification, implying they store global contextual information.
0:07:30 Hypothesis on Artifact Formation: The underlying hypothesis is that ViT models learn to identify patches with low local informational content (background) and "recycle" these tokens to accumulate global information, thereby discarding their original local spatial data.
0:08:39 Proposed Solution: Vision Transformer Models with Registers: A simple architectural modification is introduced where learnable "register tokens" (randomly initialized vectors, trained similarly to the CLS token) are appended to the input sequence alongside the image patch embeddings. The number of registers can vary (e.g., 1 to 32).
0:09:32 Key Result 1: Elimination of Artifacts: The integration of registers successfully removes the peaky outlier artifacts from attention heat maps, resulting in significantly smoother and more semantically meaningful attention distributions. The high-norm outlier patches also vanish from the norm distribution.
Object detection accuracy sees a dramatic improvement, in some cases nearly doubling for models like DeiT and DinoV2.
0:11:40 Interpretation of Registers: Individual register tokens are observed to learn distinct and diverse attention patterns, often focusing on specific objects or parts within an image (e.g., a cup, a sweet, a spoon), suggesting they specialize in capturing different features, while the CLS token maintains a broader focus on overall objects.
0:12:36 Conclusion: Register-augmented ViT models provide a dedicated mechanism for global information aggregation, effectively removing attention artifacts, leading to smoother attention maps, and substantially improving performance in dense prediction and object discovery tasks.