Issue 1:
🚀 Feature

Motivation

Currently when using the nn.MultiHeadAttention layer, the attn_output_weights consists of an average of the attention weights of each head, therefore the original weights are inaccessible. That makes analysis like the one made in this paper very difficult.

Pitch

When the nn.MultiHeadAttention forward is called with need_weights=True (and maybe a second parameter like nead_attn_heads=True), attn_output_weights should be a tensor of size [N,num_heads,L,S],with the weights of each head, instead of the average of size [N,L,S] (following the notation in the docs)

Alternatives

Additional context

A small discussion about this subject with a potential solution was made here

If you guys agree, I'll gladly make a PR.

Issue 2:
🚀 Feature: Return the attention scores from all heads

Return the attention scores from all heads in the MultiHeadAttention Module and multi_head_attention_forward. The current implementation only allows returning the averaged attention score over all heads

In short: add option to change

pytorch/torch/nn/functional.py

Line 4309
in
16c72a5

return attn_output, attn_output_weights.sum(dim=1) / num_heads

into

return attn_output, attn_output_weights

Motivation

The motivating idea behind Multi-Head attention is to perform the attention mechanism in parallel and allow the model to attend to different sequence elements with each head separately. The current torch code (MultiHeadAttention Module and multi_head_attention_forward) allows to return the attn_output_weights along with the attn_output.

However, the attn_output_weights are averaged over all heads. The attention scores for each separate head are not accessible to the outside.

Why should it be possible to get attention scores for each head?

Each head attends to different sequence elements. The information to which sequence elements each head attends to is lost when attention scores are averaged over all heads. The attention scores for each head are important to understand the features that are extracted in each layer (touches: feature importance, explainability).

This averaging behavior in attn_output_weights is inconsistent with attn_output where the embedding dimensions are grouped into head_dim blocks from the same attention head.

pytorch/torch/nn/functional.py

Line 4124
in
16c72a5

head_dim = embed_dim // num_heads

Pitch

I would propose

modifying the need_weights=True option in multi_head_attention_forward to a choice [all, average, none] to control the return behavior of multi_head_attention_forward. The option need_weights=avg would be equivalent to current need_weights=True and need_weights=all would return attn_output_weights without summation and average with num_heads here. Alternatives like changing the current behavior or adding a second boolean need_all_weights flag are possible too.

carrying the need_weights option to the MultiHeadAttention Module to make it easily accessible to the outside when using the MultiHeadAttention module.

I am happy to write a pull request with the modifications in multi_head_attention_forward and MultiHeadAttention Module according to further discussion or to leave the concrete implementation choices to the maintainers.

Overall, it would be great to have this option in the PyTorch code. This would make the code mode cleaner and would allow understanding the different head behaviors without having to locally modify the pytorch functions.

Alternatives

Current alternatives are:

creating multiple MultiHeadAttention modules hardcoded with a single head to retrieve the attention scores of this head (probably less efficient)

copying and modifying the multi_head_attention_forward function and MultiHeadAttention Module locally.

cc @albanD @mruberry @zhangguanheng66