TLDR: A new method called “Token Discarding” and metrics (A-TSI, M-TSI) are introduced to detect spurious correlations in Vision Transformers (ViTs). By removing image tokens and measuring prediction changes, the method identifies when ViTs rely on irrelevant features instead of the main object. Experiments on ImageNet show training methods impact spurious correlations, with DINO (self-supervised) being more robust. The method also highlights problematic dataset classes and has been successfully applied to breast cancer detection, improving AI trustworthiness by ensuring models focus on relevant features.
In the rapidly evolving world of artificial intelligence, especially in computer vision, models are becoming incredibly powerful. However, this power comes with a hidden challenge: the ability of these models to learn and exploit “spurious correlations.” These are unintended patterns in the data that, while statistically present, are not truly relevant to the task at hand. Imagine a model learning to identify a specific animal not by its features, but by a common background element in all its training images. This can lead to correct predictions for the wrong reasons, making the models unreliable and less generalizable to new, unseen data.
This issue is particularly pronounced in Vision Transformers (ViTs), a type of neural network that has gained immense popularity. ViTs process images by breaking them down into small “tokens” and using an “attention mechanism” to understand relationships between these tokens. While powerful, this mechanism can inadvertently connect distant and irrelevant parts of an image, exacerbating the problem of spurious correlations.
Traditionally, researchers have approached this problem in two ways: by developing robust training methods to prevent these correlations from being learned, or by detecting them after a model has been trained. The latter often relies on interpretability methods like GradCAM, which highlight parts of an image a model focuses on. However, these methods can sometimes be misleading, raising concerns about their accuracy.
A recent research paper, titled “Detecting Regional Spurious Correlations in Vision Transformers via Token Discarding,” introduces a novel and more robust approach to identify these problematic correlations. Authored by Solha Kang, Esla Timothy Anzaku, Wesley De Neve, Arnout Van Messem, Joris Vankerschaver, Francois Rameau, and Utku Ozbulak, this work leverages the unique properties of Vision Transformers to offer a clearer picture of what a model is truly paying attention to. You can read the full paper here.
The Token Discarding Method
The core of the proposed method lies in “token discarding.” Unlike traditional image masking, which can introduce biases in older convolutional neural networks (CNNs) by creating artificial missing regions, ViTs naturally handle the removal of individual tokens. The researchers systematically remove one token (a small patch of the image) at a time and observe how this affects the model’s confidence in its prediction. A significant drop in confidence indicates that the removed token was highly “influential” to the model’s decision.
To quantify spurious correlations, the paper introduces two new metrics: the Average Token Spuriosity Index (A-TSI) and the Maximum Token Spuriosity Index (M-TSI). These indices compare the influence of tokens located inside the bounding box of the actual object of interest (core features) with those outside it (potentially spurious features). If a TSI score is less than 1, it means the model is primarily relying on relevant features. A score of 1 indicates equal reliance, and a score greater than 1 signals that the model is leaning more on spurious, irrelevant features for its prediction. The higher the TSI, the stronger the spurious correlation.
Key Findings and Insights
The researchers conducted extensive experiments using the widely-used ImageNet dataset, employing both supervised and self-supervised Vision Transformers (DINO and MAE). Their findings revealed several crucial insights:
- Training Methodology Matters: The way a model is trained significantly impacts its susceptibility to spurious correlations. DINO, a self-supervised model, consistently exhibited fewer spurious correlations compared to supervised models and even another self-supervised model, MAE. This suggests that certain self-supervised training routines can lead to more robust models.
- Spurious Correlations Lead to Errors: Images that were incorrectly classified by the models generally had higher TSI scores, indicating that misclassifications often stem from the model relying on irrelevant features.
- Problematic ImageNet Classes: The study identified specific classes within the ImageNet dataset that are particularly prone to spurious correlations across different models. Examples include “space bar,” “ping-pong ball,” and “puck.” The underlying reasons for these issues were diverse, ranging from inconsistent labeling and the presence of strong secondary objects (like humans in sports images) to very small objects of interest or visual similarities between different classes.
- Real-World Application: The method was successfully applied to a medical imaging case study: invasive breast mass classification on MRI images. It effectively identified instances where models focused on irrelevant regions, such as chest fat tissue, instead of the actual breast tissue, highlighting its potential for building more trustworthy AI in critical applications.
- M-TSI for Small Objects, A-TSI for Large: The choice between A-TSI and M-TSI depends on the object’s size. M-TSI is more effective at detecting spurious correlations when the object of interest is small, while A-TSI provides a more reliable measure for larger objects.
Limitations and Future Directions
While powerful, the token discarding method has some limitations. It requires significant computational resources to generate token influence maps and relies on annotations (like bounding boxes) to define the object of interest. The researchers are exploring ways to mitigate these issues, such as using attention maps as a more efficient proxy for token influence, especially in certain self-supervised models like DINO. They also aim to develop a more general framework that eliminates the need for explicit annotations altogether.
Also Read:
- Enhancing Vision Transformers with Challenging Synthetic Negatives
- Uncovering the True Drivers of AI Vision: A Causal Approach to Feature Explanation
Broader Impact
The ability to detect and quantify spurious correlations has profound implications for the trustworthiness and fairness of AI systems. It provides a principled way to debug models, identify biases, and improve the integrity of datasets. This is particularly vital in high-stakes fields like medical diagnosis, where relying on irrelevant features can have serious consequences. The Token Spuriosity Index framework is also adaptable to other transformer-based models beyond computer vision, including those used in natural language processing, paving the way for more reliable and interpretable AI across various domains.
