<h2 dir="auto"><g-emoji class="g-emoji" alias="rocket" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f680.png">🚀</g-emoji> Feature</h2>
<h2 dir="auto">Motivation</h2>
<p dir="auto">Currently when using the <code class="notranslate">nn.MultiHeadAttention</code> layer, the <code class="notranslate">attn_output_weights</code> 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 <a href="https://arxiv.org/abs/1906.04341v1" rel="nofollow">paper</a> very difficult.</p>
<h2 dir="auto">Pitch</h2>
<p dir="auto">When the <code class="notranslate">nn.MultiHeadAttention</code> forward is called with <code class="notranslate">need_weights=True</code> (and maybe a second parameter like <code class="notranslate">nead_attn_heads=True</code>),  <code class="notranslate">attn_output_weights</code> should be a tensor of size <code class="notranslate">[N,num_heads,L,S]</code>,with the weights of each head, instead of the average of size <code class="notranslate">[N,L,S]</code> (following the notation in the <a href="https://pytorch.org/docs/stable/nn.html#multiheadattention" rel="nofollow">docs</a>)</p>
<h2 dir="auto">Alternatives</h2>
<h2 dir="auto">Additional context</h2>
<p dir="auto">A small discussion about this subject with a potential solution was made <a href="https://discuss.pytorch.org/t/getting-nn-multiheadattention-attention-weights-for-each-head/72195" rel="nofollow">here</a></p>
<p dir="auto">If you guys agree, I'll gladly make a PR.</p>
<h2 dir="auto"><g-emoji class="g-emoji" alias="rocket" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f680.png">🚀</g-emoji> Feature: Return the attention scores from all heads</h2>
<p dir="auto">Return the attention scores from all heads in the <a href="https://github.com/pytorch/pytorch/blob/b5dd6e3e61c59e88ce6e8aee2fa5d08d6c56fd97/torch/nn/modules/activation.py#L831"><code class="notranslate">MultiHeadAttention</code></a> Module and <a href="https://github.com/pytorch/pytorch/blob/16c72a5a6b7085a9651e7bef7e98c80b2d5547e9/torch/nn/functional.py#L4026"><code class="notranslate">multi_head_attention_forward</code></a>. The current implementation <strong>only</strong> allows returning the averaged attention score over all heads</p>
<p dir="auto">In short: add option to change<br>
</p><div class="Box Box--condensed my-2">
  <div class="Box-header f6">
    <p class="mb-0 text-bold">
      <a href="https://github.com/pytorch/pytorch/blob/16c72a5a6b7085a9651e7bef7e98c80b2d5547e9/torch/nn/functional.py#L4309">pytorch/torch/nn/functional.py</a>
    </p>
    <p class="mb-0 color-fg-muted">
         Line 4309
      in
      <a data-pjax="true" class="commit-tease-sha" href="/pytorch/pytorch/commit/16c72a5a6b7085a9651e7bef7e98c80b2d5547e9">16c72a5</a>
    </p>
  </div>
  <div itemprop="text" class="Box-body p-0 blob-wrapper blob-wrapper-embedded data">
    <table class="highlight tab-size mb-0 js-file-line-container" data-tab-size="8" data-paste-markdown-skip="">

        <tbody><tr class="border-0">
          <td id="L4309" class="blob-num border-0 px-3 py-0 color-bg-default" data-line-number="4309"></td>
          <td id="LC4309" class="blob-code border-0 px-3 py-0 color-bg-default blob-code-inner js-file-line"> <span class="pl-k">return</span> <span class="pl-s1">attn_output</span>, <span class="pl-s1">attn_output_weights</span>.<span class="pl-en">sum</span>(<span class="pl-s1">dim</span><span class="pl-c1">=</span><span class="pl-c1">1</span>) <span class="pl-c1">/</span> <span class="pl-s1">num_heads</span> </td>
        </tr>
    </tbody></table>
  </div>
</div>
<br>
into<br>
<code class="notranslate">return attn_output, attn_output_weights</code><p></p>
<h2 dir="auto">Motivation</h2>
<p dir="auto">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 (<a href="https://github.com/pytorch/pytorch/blob/b5dd6e3e61c59e88ce6e8aee2fa5d08d6c56fd97/torch/nn/modules/activation.py#L831"><code class="notranslate">MultiHeadAttention</code></a> Module and <a href="https://github.com/pytorch/pytorch/blob/16c72a5a6b7085a9651e7bef7e98c80b2d5547e9/torch/nn/functional.py#L4026"><code class="notranslate">multi_head_attention_forward</code></a>) allows to return the <code class="notranslate">attn_output_weights</code> along with the <code class="notranslate">attn_output</code>.</p>
<p dir="auto">However, the <code class="notranslate">attn_output_weights</code> are averaged over all heads.  The attention scores for each separate head are not accessible to the outside.</p>
<p dir="auto">Why should it be possible to get attention scores for each head?</p>
<ol dir="auto">
<li>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).</li>
<li>This averaging behavior in <code class="notranslate">attn_output_weights</code> is inconsistent with <code class="notranslate">attn_output</code> where the embedding dimensions are grouped into <code class="notranslate">head_dim</code> blocks from the same attention head.<br>
<div class="Box Box--condensed my-2">
  <div class="Box-header f6">
    <p class="mb-0 text-bold">
      <a href="https://github.com/pytorch/pytorch/blob/16c72a5a6b7085a9651e7bef7e98c80b2d5547e9/torch/nn/functional.py#L4124">pytorch/torch/nn/functional.py</a>
    </p>
    <p class="mb-0 color-fg-muted">
         Line 4124
      in
      <a data-pjax="true" class="commit-tease-sha" href="/pytorch/pytorch/commit/16c72a5a6b7085a9651e7bef7e98c80b2d5547e9">16c72a5</a>
    </p>
  </div>
  <div itemprop="text" class="Box-body p-0 blob-wrapper blob-wrapper-embedded data">
    <table class="highlight tab-size mb-0 js-file-line-container" data-tab-size="8" data-paste-markdown-skip="">

        <tbody><tr class="border-0">
          <td id="L4124" class="blob-num border-0 px-3 py-0 color-bg-default" data-line-number="4124"></td>
          <td id="LC4124" class="blob-code border-0 px-3 py-0 color-bg-default blob-code-inner js-file-line"> <span class="pl-s1">head_dim</span> <span class="pl-c1">=</span> <span class="pl-s1">embed_dim</span> <span class="pl-c1">//</span> <span class="pl-s1">num_heads</span> </td>
        </tr>
    </tbody></table>
  </div>
</div>
</li>
</ol>
<h2 dir="auto">Pitch</h2>
<p dir="auto">I would propose</p>
<ol dir="auto">
<li>modifying the <code class="notranslate">need_weights=True</code> option in <a href="https://github.com/pytorch/pytorch/blob/16c72a5a6b7085a9651e7bef7e98c80b2d5547e9/torch/nn/functional.py#L4026"><code class="notranslate">multi_head_attention_forward</code></a> to a choice [<code class="notranslate">all</code>, <code class="notranslate">average</code>, <code class="notranslate">none</code>] to control the return behavior of <code class="notranslate">multi_head_attention_forward</code>. The option <code class="notranslate">need_weights=avg</code> would be equivalent to current <code class="notranslate">need_weights=True</code> and <code class="notranslate">need_weights=all</code> would return <code class="notranslate">attn_output_weights</code> without summation and average with <code class="notranslate">num_heads</code> <a href="https://github.com/pytorch/pytorch/blob/16c72a5a6b7085a9651e7bef7e98c80b2d5547e9/torch/nn/functional.py#L4309">here</a>. Alternatives like changing the current behavior or adding a second boolean <code class="notranslate">need_all_weights</code> flag are possible too.</li>
<li>carrying the <code class="notranslate">need_weights</code> option to the <a href="https://github.com/pytorch/pytorch/blob/b5dd6e3e61c59e88ce6e8aee2fa5d08d6c56fd97/torch/nn/modules/activation.py#L831"><code class="notranslate">MultiHeadAttention</code></a> Module to make it easily accessible to the outside when using the <code class="notranslate">MultiHeadAttention</code> module.</li>
</ol>
<p dir="auto">I am happy to write a pull request with the modifications in <a href="https://github.com/pytorch/pytorch/blob/16c72a5a6b7085a9651e7bef7e98c80b2d5547e9/torch/nn/functional.py#L4026"><code class="notranslate">multi_head_attention_forward</code></a> and <a href="https://github.com/pytorch/pytorch/blob/b5dd6e3e61c59e88ce6e8aee2fa5d08d6c56fd97/torch/nn/modules/activation.py#L831"><code class="notranslate">MultiHeadAttention</code></a> Module according to further discussion or to leave the concrete implementation choices to the maintainers.</p>
<p dir="auto">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.</p>
<h2 dir="auto">Alternatives</h2>
<p dir="auto">Current alternatives are:</p>
<ol dir="auto">
<li>creating multiple MultiHeadAttention modules hardcoded with a single head to retrieve the attention scores of this head (probably less efficient)</li>
<li>copying and modifying the <a href="https://github.com/pytorch/pytorch/blob/16c72a5a6b7085a9651e7bef7e98c80b2d5547e9/torch/nn/functional.py#L4026"><code class="notranslate">multi_head_attention_forward</code></a> function and <a href="https://github.com/pytorch/pytorch/blob/b5dd6e3e61c59e88ce6e8aee2fa5d08d6c56fd97/torch/nn/modules/activation.py#L831"><code class="notranslate">MultiHeadAttention</code></a> Module locally.</li>
</ol>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/albanD/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/albanD">@albanD</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/mruberry/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/mruberry">@mruberry</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/zhangguanheng66/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/zhangguanheng66">@zhangguanheng66</a></p>