Interpreting Attention matrix weights in MultiheadAttention

neurospacer: <p>Hi Pytorch Universe,<br>
I have an sequences of human’s poses from multiple angles and ground true pose class. So my data tensor is in form of <code>[1, 97, 1, 128, 128]</code>, where dims are <code>[batch_size, timepoint, channel, height, width]</code>.<br>
Class is only related to whole sequence, not each timepoint. Therefore, this task could be assumed as weakly-supervised multi-instance learning.</p>
<p>I build a transformer encoder with a CNN backbone, without  class token and positioning encoding. First input tensor is embedded to <code>[1,97,128]</code> (97 tokens of 128 embedding vector) after backbone CNN. Then it rearraranged to <code>[97,1,128]</code> and forwarded to Transformer encoder with pytorch’s <code>nn.MultiheadAttention(128, 8, dropout)</code>.<br>
What I am trying to understand is - I acquire the attention matrix with shape <code>[1,8,97,97]</code>, where 8 is number of head, 97 - is number if tokens (timepoints). How can interpret the attention matrix? Is it possible to create a heatmap to overlay on timepoints of original tensor <code>[1, 97, 1, 128, 128]</code>?<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/9/f9edbbbfbf4a949bfab93b856345e04e759b8343.png" data-download-href="https://discuss.pytorch.org/uploads/default/f9edbbbfbf4a949bfab93b856345e04e759b8343" title="att_matcs"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/9/f9edbbbfbf4a949bfab93b856345e04e759b8343_2_495x500.png" alt="att_matcs" data-base62-sha1="zEYemexbTtPcSN4R8E6wy68egKf" width="495" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/9/f9edbbbfbf4a949bfab93b856345e04e759b8343_2_495x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/9/f9edbbbfbf4a949bfab93b856345e04e759b8343_2_742x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/9/f9edbbbfbf4a949bfab93b856345e04e759b8343.png 2x" data-dominant-color="6F4D86"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">att_matcs</span><span class="informations">802×809 155 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

J_Johnson: <p>Normally, attention matrices show which tokens in the sequence(in your case, time steps) are most relevant to solving the task. And what usually gets highlighted is along the main diagonal, with occasional spots else where, indicating a reordering of tokens may make more sense in decoding(i.e. translation).</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/6/b67491f14df0ed8d305223857e650245cc498b49.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/b67491f14df0ed8d305223857e650245cc498b49" title="bahdanau-fig3"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/6/b67491f14df0ed8d305223857e650245cc498b49_2_558x500.jpeg" alt="bahdanau-fig3" data-base62-sha1="q24Ei9yiDTjcn90lmV57LlR6dXb" width="558" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/6/b67491f14df0ed8d305223857e650245cc498b49_2_558x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/6/b67491f14df0ed8d305223857e650245cc498b49_2_837x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/6/b67491f14df0ed8d305223857e650245cc498b49_2_1116x1000.jpeg 2x" data-dominant-color="6C6C6C"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">bahdanau-fig3</span><span class="informations">1654×1482 106 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>But with your attention matrices, there appear to be primarily vertical lines. It might mean those brighter sections (time steps) are what the model thinks it should focus on to classify the sequence. That’s assuming it’s been sufficiently trained.</p>

neurospacer: <p>Hey <a class="mention" href="/u/j_johnson">@J_Johnson</a>,<br>
thank for your reply. In the case of NLP example you provided, the attention matrix, is a vector dot product of Q and V, and, as a measure of similarity, shows how token representation is similar to it’s target. However, I still struggle to understand, what the attention matrix shows in my case. So the there is two matrices Q and V acquired by linearly mapping input tensor to Q K V matrices, duplicated with weights <code>n_head</code> times. So the Q and V in each head are just two input tensor projections. Then attention matrix is computed with dot product of Q and V. As I understood, the dot product is computed between two linear projections of the input tensor. If I look at the attention matrices above,  Head 1 Layer 1, have high weight values in around 10/97 sequences, it means that we have high similarity between vector in V with Q? Does my intuition right?<br>
Why the pattern is mostly vertical?</p>

J_Johnson: <p>If you’re using the PyTorch Transformer, then it is self-attention. And so it is simply determining which time steps are most relevant.</p>
<p>The vertical lines are a bit odd and I don’t know what to tell you, in that case. How is the performance of the model after training?</p>
<p>If the performance is fitting to the data, it’s possible that the vertical is mapping the time steps while the horizontal is some type of mapping to the output probabilities.</p>

neurospacer: <p>The performance of multi-class classification, in terms of accuracy is more than good: training, test and hold-out validation are &gt; 0.94 ± 0.148.<br>
Thanks for your thoughts</p>

ado_sar: <p>Shouldn’t the weights always be highest in the diagonal? I mean dot product which measures  similarity between x_i and y_j should put more weight to a_ij when i=j.</p>

J_Johnson: <p>Right. That’s what tends to happen.</p>

neurospacer: <p>In a simple self-attention mechanism, this would indeed be the case. However, in transformer, vector columns of X input matrix is weighted with Wq, Wk, Wv matrices, that initilized differently. I assume if you would initilize them with the same weights, then diagonal elements likely to be highest.</p>