Register and CLS tokens yield a decoupling of local and global features in large ViTs

Abstract: Recent work has shown that the attention maps of the widely popular DINOv2 model exhibit artifacts, which hurt both model interpretability and performance on dense image tasks. These artifacts emerge due to the model repurposing patch tokens with redundant local information for the storage of global image information. To address this problem, additional register tokens have been incorporated in which the model can store such information instead. We carefully examine the influence of these register tokens on the relationship between global and local image features, showing that while register tokens yield cleaner attention maps, these maps do not accurately reflect the integration of local image information in large models. Instead, global information is dominated by information extracted from register tokens, leading to a disconnect between local and global features. Inspired by these findings, we show that the CLS token itself, which can be interpreted as a register, leads to a very similar phenomenon in models without explicit register tokens. Our work shows that care must be taken when interpreting attention maps of large ViTs. Further, by clearly attributing the faulty behaviour to register and CLS tokens, we show a path towards more interpretable vision models.

• The paper finds that register and [CLS] tokens in large Vision Transformers decouple local image patches from global representations, leading to cleaner but less representative attention maps.
• The [CLS] token behaves similarly to explicit register tokens in larger models, storing global information separately, a phenomenon correlated with increased model size due to overparameterization.
• This feature decoupling impacts model interpretability for tasks requiring fine-grained local understanding and suggests future ViT architectures should aim for better integration of local and global features.

Decoupling of Local and Global Features in Large Vision Transformers: An Analysis of Register and [CLS] Tokens

The paper "Register and [CLS] tokens yield a decoupling of local and global features in large ViTs" presents a detailed examination of the attention mechanisms and feature hierarchies in Vision Transformers (ViTs), particularly focusing on the implications of using register tokens alongside the [CLS] token. The authors aim to elucidate the phenomenon through which large ViTs exhibit a disconnect between local image patches and global feature representations. This investigation specifically targets the DINOv2 model and highlights how it manipulates attention maps, impacting both model interpretability and performance on dense image tasks.

Key Findings

1. Register Tokens and Attention Maps: The introduction of register tokens was originally intended to store global image features, thus removing the need for high-norm patch tokens which muddy the interpretation of attention maps. Register tokens indeed lead to cleaner attention maps, but in large models, they do not faithfully represent the mechanism through which the final global output is constructed. The global representation becomes dominated by information from the register tokens rather than from the local patches.
2. [CLS] Token as an Implicit Register: The study also shows that the [CLS] token in models without explicit register tokens behaves similarly, acting as a storage for global information through its residual connections. This demonstrates that even in the absence of explicit register tokens, large models can abstract global information away from the patch tokens. In smaller models where the patch integration assumption holds, the patch tokens align more closely with global representations.
3. Model Size and Feature Disconnection: The results intriguingly demonstrate that the size of the model correlates with the degree of feature disconnection. Larger models tend to rely more heavily on non-local tokens like register and [CLS] tokens. This is likely due to the overparameterized architecture which facilitates premature convergence of global representations within the intermediate layers, leading to simplistic last-layer representations.
4. Neural Collapse and Overparameterization: The observed behavior aligns with the neural collapse theory, which predicts a convergence towards low-dimensional representations, especially in overparameterized models. This tendency towards simplistic, high-level abstractions could inherently decouple the global representation from local features, thus impacting the model's interpretability in tasks requiring an understanding of local image regions.

Theoretical and Practical Implications

The study provides a crucial insight into how modifications to transformer architectures—such as the inclusion of register tokens or reliance on the [CLS] token—can have far-reaching consequences on model performance and interpretability. For practitioners, it suggests caution in using large ViTs for tasks requiring fine-grained local feature integration. The paper implies that future iterations of ViTs, particularly for applications necessitating interpretable attention maps, should consider the elimination of register and [CLS] tokens or at least adjust the architecture to enforce a more meaningful integration of local and global features.

Future Directions

Moving forward, there are several promising avenues for research and development:

• Model Architecture: Designing alternative architectures that can maintain the patch integration assumption without sacrificing the simplicity or flexibility offered by the transformer model.
• Regularization Techniques: Employing regularization strategies that can prevent the dominance of non-local tokens, ensuring a distributed and meaningful integration of local features into global representations.
• Empirical Validation: Extending comprehensive evaluations to a variety of datasets and downstream tasks, especially those requiring detailed understanding of image spatial hierarchy, to confirm the general applicability of the findings.

This investigation underscores the nuanced challenges and considerations in designing interpretability into large, complex models like Vision Transformers. The sophisticated interaction between model components as revealed through this study could guide the next generation of interpretable and efficient ViT architectures.