The attention scores are a bit scattered here, usually the cls token focuses on certain patches and they are consistent. Are you sure taking the mean is a good idea? Also it might be better to rollout the cls token over multiple blocks.

So I tried taking the product of the cls token weights over other blocks like this but it shows the error

RuntimeError                              Traceback (most recent call last)Cell In[18], line 13     11# Forward pass through all blocks     12forblockin model.blocks:---> 13     x, attn_map = block.attn.forward(x)     14     outputs.append(x)     15     attn_maps.append(attn_map)Cell In[8], line 4,in my_forward_wrapper.<locals>.my_forward(x)      2 def my_forward(x):      3     B, N, C = x.shape----> 4     qkv = attn_obj.qkv(x).reshape(B, N, 3, attn_obj.num_heads, C // attn_obj.num_heads).permute(2, 0, 3, 1, 4)      5     q, k, v = qkv.unbind(0)# make torchscript happy (cannot use tensor as tuple)      7     attn = (q @ k.transpose(-2, -1))* attn_obj.scaleFile /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1518,in Module._wrapped_call_impl(self,*args,**kwargs)   1516return self._compiled_call_impl(*args,**kwargs)# type: ignore[misc]   1517 else:-> 1518return self._call_impl(*args,**kwargs)File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1527,in Module._call_impl(self,*args,**kwargs)   1522# If we don't have any hooks, we want to skip the rest of the logic in   1523# this function, and just call forward.   1524if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks   1525         or _global_backward_pre_hooks or _global_backward_hooks   1526         or _global_forward_hooks or _global_forward_pre_hooks):-> 1527return forward_call(*args,**kwargs)   1529 try:   1530     result = NoneFile /opt/conda/lib/python3.10/site-packages/torch/nn/modules/linear.py:114,in Linear.forward(self, input)    113 def forward(self, input: Tensor) -> Tensor:--> 114return F.linear(input, self.weight, self.bias)RuntimeError: mat1 and mat2 shapes cannot be multiplied (672x224 and 384x1152)

Code:

model=create_model('deit_small_distilled_patch16_224',pretrained=True)# Replace forward function in all blocksforblockinmodel.blocks:block.attn.forward=my_forward_wrapper(block.attn)# Forward pass through the modeloutputs= []attn_maps= []cls_weights= []# Forward pass through all blocksforblockinmodel.blocks:x,attn_map=block.attn.forward(x)outputs.append(x)attn_maps.append(attn_map)cls_weights.append(block.attn.cls_attn_map.min(dim=1).values.view(14,14).detach())# Combine class scores of all blockscls_weight_combined=torch.prod(torch.stack(cls_weights),dim=0)# Resize input image and class weightsimg_resized=x.permute(0,2,3,1)*0.5+0.5cls_resized=F.interpolate(cls_weight_combined.view(1,1,14,14), (224,224),mode='bilinear').view(224,224,1)# Visualizeshow_img(image)show_img(attn_maps[-1])# Attention map from the last blockshow_img(cls_weight_combined)# Combined class weightsshow_img(img_resized.squeeze())# Squeeze the batch dimensionshow_img2(img_resized.squeeze(),cls_resized,alpha=0.8)# Squeeze the batch dimension

Hi @arnavs04

That's a good question. You're right. There may be a better way to merge different attention maps. The code I give contains very basic ways to merge different attention maps.

The error you posted looks like a dimension mismatch error:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (672x224 and 384x1152)

I suggest applying a patch embedding layer to the input image before passing it into the attention block. You should also ensure that other layers (e.g., normalization, mlp) are applied appropriately.

Thank you.

Hankyul.

Ohh, thank you for the help!!

Hi@mae338

I am happy to help you.

• Iny=model(x.unsqueeze(0)),x.unsqueeze(0) is a transformed image that has shapes as(1, 3, 224, 224).
• Inmodel.blocks[-1].attn.forward=my_forward_wrapper(model.blocks[-1].attn), we replace the originalforward to ourmy_forward_wrapper to save attention map as an instance variable ofmodel.blocks[-1].attn.
• Inmodel.blocks[-1].attn.attn_map.mean(dim=1).squeeze(0).detach(), we average attention map in head-dimension.mean(dim=1) so that we can display whole attention.
• To understand better, we recommend you print out every shape of the tensor, which helps you to see the overall workflow of ViT.

Thank you.

Hankyul

Thank you so much, Hankyul. I appreciate your effort to help me. I'd be pleased if you could help me understand this as well:
What about def my_forward(x):
B, N, C = x.shape
qkv = attn_obj.qkv(x)?
more especifically, the x parameter? what does it refer to?
and why did you create the function my_forward?
I also found out that the function my_forward must return x or else it doesn't work and the whole program won't work. I wonder why?

Hi@mae338

I hope this could help you.

• We replace the original Attentionforward() method tomy_forward() method that only inserts two extra codes for saving attention map (attn_obj.attn_map = attn,attn_obj.cls_attn_map = attn[:, :, 0, 2:]). Thus, everything (function, signature, and return value) should be the same as the original function except for additional code.
• x is image tokens (Batch x N x D) input to attention. Image tokens are tokenized by a linear layer, and their shape depends on the model's size. In our case, DeiT-small/16 splits the whole image into 196 (N=198, with 2 extra class, distill tokens) patches and has 384 (D=384) channel dimensions.
• Since the attention block passes the output to the next block such as the MLP block,my_forward() should also pass the return value. If you skip passing the return value to the next block, the next blocks get aNone value, which is an unexpected situation for them, thereby generating an error.

Thank you.

Hankyul

Hi@hankyul2,

I’d like to view the outputs of a pretrained vit model (vit-base-batch-16 224) , especially the input to the mlp head. I tried the same function but it didn’t work. Any suggestions?
Thank you
mae

Hi@mae338

I hope this can help you.

You can extract the ultimate features of a pre-trained ViT byy = model.forward_head(x, pre_logits=True) and visualize them for your purpose. Since it was a long time ago when I upload initial code in first comment, I copied and modified them as:

# dependency!wgethttps://user-images.githubusercontent.com/31476895/167238573-b0cc3a6d-d3ee-462b-8630-a8f253e69bb2.png!pipinstall-Uqfastaitimm==0.6.13huggingface_hub########################################### codeimportnumpyasnpfromPILimportImagefromtimm.modelsimportcreate_modelfromtorchimportnnfromtorchvision.transformsimportCompose,Resize,CenterCrop,Normalize,ToTensordefto_tensor(img):transform_fn=Compose([Resize(249),CenterCrop(224),ToTensor(),Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))])returntransform_fn(img)img=Image.open('167238573-b0cc3a6d-d3ee-462b-8630-a8f253e69bb2.png').convert('RGB')x=to_tensor(img)model=create_model('deit_small_distilled_patch16_224',pretrained=True)x=model.forward_features(x.unsqueeze(0))y=model.forward_head(x,pre_logits=True)print(y.shape)########################################### outputtorch.Size([1,384])##########################################

Thank you.

Hankyul

Hi@tomos7231

1. 384 means input image resolution to ViT model. If you want to extract an attention map using the code above, you should change the input resolution (224 ->384) and the dimension of patches (14x14 ->24x24) like below code.

attn_map=model.blocks[-1].attn.attn_map.mean(dim=1).squeeze(0).detach()cls_weight=model.blocks[-1].attn.cls_attn_map.mean(dim=1).view(24,24).detach()# (14->24)img_resized=x.permute(1,2,0)*0.5+0.5cls_resized=F.interpolate(cls_weight.view(1,1,24,24), (384,384),mode='bilinear').view(384,384,1)# (14->24), (224->384)

2. 14 means the number of patches in each spatial dimension, e.g., H, W of the feature map. This value (14) is determined by patch size (16) because each image is divided by 196 patches, each size is16x16.
3. If you just want to extract an attention map, I would recommend@rwightman's solution, which is more convenient.

Thank you.

Hankyul

Hi@hankyul2

Well understood. Thank you for reply!

@hankyul2
Thank you for such an insightful explanation above.

1. I also tried to apply this code to the 'vit_small_patch16_384' model as shown above, but I encountered the following error.

RuntimeError                              Traceback (most recent call last)Cell In[26], line 24     20 # y = model(image)     23 attn_map = model.blocks[-1].attn.attn_map.mean(dim=1).squeeze(0).detach()---> 24 cls_weight = model.blocks[-1].attn.cls_attn_map.mean(dim=1).view(24, 24).detach()     26 img_resized = image2.permute(0, 2, 3, 1) * 0.5 + 0.5     27 cls_resized = F.interpolate(cls_weight.view(1, 1, 24, 24), (384, 384), mode='bilinear').view(384, 384, 1)RuntimeError: shape '[24, 24]' is invalid for input of size 575

2.If I want to apply this model to a regression problem and estimate multiple parameters with a single model, is it possible to check the attention map for each parameter? In a multi-class classification problem, we can check the attention map when a certain class is classified to see which areas are being focused on to recognize that class. However, in the case of a regression problem, how should this be done?

Hello@rwightman, thank you for showing us how to extract the attention layers and maintaining your wonderful timm library. I would like to ask, I am using fast.ai alongside with timm models. I've trained a ViT for a classification task in my own dataset. What would be the best way to load the weights of my ViT and visualize the attention activations over my input image using the timm visualization helpers?, thank you

I have a bit of a trivial question here and slightly off-topic to what@rwightman discussed, but related to the attention map that@hankyul2 posted in their initial visualization:

I'm assuming that the ViT in OPs question is being trained for a classification task. Is it necessary that the attention map must contain high activations along the diagonal, similar to the attention maps generated while training seq2seq models? This seems quite unintuitive to me, as I wouldn't expect a classification ViT to attend to self-attend to patches in such a manner. I would instead expect every patch to attend to certain "interesting" parts of the image, almost like having certain columns in the attention map with high activations.

Is my understanding incorrect, or should we expect the attention map to have high activations along the diagonal?

Thank you in advance!

@SarthakJShetty You can see the attention maps here after every block for DeiT, in the initial layers you can see the clear diagonal and how the attention maps then change. Please check out mynotebook. It contains visualizations for attention rollout but also without. I hope it can be of help.

My assumption, or rather intuition, behind these diagonal patterns is centered around the class token, which I believe is the reason we don't see information flow immediately, as the class token takes in information from the tokens without much change. However, this is what I think for shallow layers; I only believe that in the later and deeper layers, the attention starts to play out.

Thank you for the clarification @arnavs04! This helps. Now that I think about it, the visualizations that I'm observing are almost always like Attention Map 5 and beyond.

Quick question: when you say "the attention starts to play out.", you mean when the queries actually start attending to relevant parts (and not just predominately self-attending in a diagonal fashion) of the image correct? i.e when the attention actually starts looking non-diagonal and like the attention maps 5-12?

@SarthakJShetty
One of the major reasons why vision transformers (ViTs) are thought to be “better” than CNNs is their ability to share global contextual information right from the beginning (i.e., in the shallow layers). Unlike CNNs, which progress from local to global information, ViTs can, in theory, access global information from the outset. However, attention maps show that not much global information is actually shared in the early stages. This has led to many works proposing modifications to optimize the self-attention mechanism in vision transformers, effectively “tuning” them similarly to CNNs.

In the diagram, you can see how information moves gradually from local to global, resembling the behaviour of a CNN (hence the diagonal structure slowly transitioning to a more globally attended pattern). This example might not be fully representative, as there is usually some global information shared even in the early stages of a vision transformer.

Here is a paper that maybe clear your doubts:Do Vision Transformers See Like Convolutional Neural Networks. I haven't gone through it completely as I had kept it on my reading list.

I am currently working on Vision Transformers. If you have any questions or would like to discuss them with me, let me know! I'll be glad to help!

This is deit architecture which has both a class token and a distillation token for prediction.
When we do index0 we are looking at the similarity scores of all other tokens wrt to class token. And now when we do2: it means, that we're only looking at the similarity tokens of all the patch tokens wrt class token (excluding the class token and distillation token itself)

Thanks for your response @arnavs04! The similarity scores you mentioned are the dot product of two vectors from the query and key matrices. Is that right?

@zichunxx Yup exactly! The (n+2) x (n+2) attention matrix.

Thanks for your generous help! @arnavs04 I have read your notebook which is very thorough and helpful.

I have noticed that some vision transformers are implemented as an encoder without the cls token. In this situation, how do we plot the overlaid image to illustrate which patch is watched with a higher weight? Thanks!

I apologize for the delay, I didn't see the reply.

The attribution map is resized with bilinear interpolation to fit the H x W resolution as the original map. This heatmap is now taken and with 0.5 x heat_map + 0.5 x original_image we get our saliency map. Obviously you can tweak the values instead of 0.5 and 0.5 respectively. This is the goal of post-hoc model agnostic explainability methods for vision transformers.

Awesome