This project demonstrates how to visualize the attention flow in a Vision Transformer (ViT) model using the "Attention Rollout" method, as described in the paper "Quantifying Attention Flow in Transformers" (arXiv:2005.00928). By applying this technique, we can identify which parts of an input image are most attended to by the model, particularly in relation to the Class (CLS) token, providing insights into the model's decision-making process.
The core idea is to aggregate attention weights across all layers of the Transformer. For each layer, the attention matrix (averaged across heads) is added to an identity matrix and then normalized. These modified attention matrices are then multiplied sequentially across layers to produce a final "rollout" matrix. The first row of this rollout matrix, excluding the CLS token's self-attention, represents the total attention accumulated from the CLS token to each image patch token. This resulting attention map can then be reshaped and visualized over the original image.
Here are examples of the output generated by the script:
This project implements the "Attention Rollout" method proposed in:
Quantifying Attention Flow in Transformers Samira Abnar, Willem Zuidema arXiv:2005.00928 Link to paper on arXiv