Access self-attention matrix of vision transformer

[sophmrtn]

Is your feature request related to a problem? Please describe. I am training a vision transformer using monai and would like to carry out some interpretability analysis. However, the current model code does not save the self-attention matrix during training, and it is not straightforward to pass it from the self-attention block to the model output.

https://github.com/Project-MONAI/MONAI/blob/88fb0b10d85f59e9ee04cbeeddb28b91f2ea209b/monai/networks/blocks/selfattention.py#L56-L65

Describe the solution you'd like An option to output the attn_mat from the self-attention block in the model forward pass (before matrix multiplication with the input) or access it after training as a class attribute.

[a-parida12]

@wyli I would like to help with this issue.

I was thinking, the easiest way to achieve this without changing the API could be by making att_mat before dropout as a class attribute so could be accessed by SABlock().attn_mat. also, a parameter like store_attn:bool in the __init__ to prevent memory overhead when the feature is not required.

Let me know your thoughts.

[wyli]

sounds good, please make sure it's backward compatible, e.g. previously saved checkpoints can still be loaded by default.

[a-parida12]

@wyli so #6271 is merged? So do i put a new merge request for a potential solution?

[a-parida12]

@wyli I just realized that ViT backbone is used by by UNetr and ViTAutoEnc ideally they should have the option to allow access to the the attn_mat else it will always be set as the default value and no way it can be changed by the user of UNetr and ViTAutoEnc. What do you think?

[wyli]

Sure, please help create another feature request to follow up..