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.

MONAI/monai/networks/blocks/selfattention.py

Lines 56 to 65 in 88fb0b1

	def forward(self, x):
	output = self.input_rearrange(self.qkv(x))
	q, k, v = output[0], output[1], output[2]
	att_mat = (torch.einsum("blxd,blyd->blxy", q, k) * self.scale).softmax(dim=-1)
	att_mat = self.drop_weights(att_mat)
	x = torch.einsum("bhxy,bhyd->bhxd", att_mat, v)
	x = self.out_rearrange(x)
	x = self.out_proj(x)
	x = self.drop_output(x)
	return x

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..