I am trying to experiment with Transformers. After training, I would like to visualize the contribution matrix from an input. If I understood correctly, the contribution matrix is a square one, indicating the contribution of each element with other elements from the sequence (diagonal values are high). This matrix is computed as in the linked picture, Q and K being respectively the query and key matrix, and d the number of features of each element.
Is there a way to easily get the result of this operation in order to visualize this contribution matrix?
I think that these contribution matrices can be obtained from the TransformerEncoderLayer. So, in the function at line 350 in this file: pytorch/transformer.py at master · pytorch/pytorch · GitHub, I call the function
self.self_attn with parameter
need_weights=True, average_attn_weights=False, and then save the output_weights as tensors. Is this the right way ?
Thanks for the help