Attention weights with multiple heads in nn.MultiheadAttention?

Sam_Lerman: <p>I’m confused how <code>attn_output_weights</code> is specified to have shape <code>(N, L, S)</code> regardless of the number of heads. Wouldn’t there be a unique set of weights for each head?</p>
<p><a href="https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html</a></p>

eqy: <p>This seems to be because the attention weights are averaged across all of the heads:</p><aside class="onebox githubissue">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/34537" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/34537" target="_blank" rel="noopener nofollow ugc">nn.MultiHeadAttention should be able to return attention weights for each head.</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-03-10" data-time="13:18:18" data-timezone="UTC">01:18PM - 10 Mar 20 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/ironcadiz" target="_blank" rel="noopener nofollow ugc">
          <img alt="ironcadiz" src="https://avatars.githubusercontent.com/u/13740457?v=4" class="onebox-avatar-inline" width="20" height="20">
          ironcadiz
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">enhancement</span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">module: nn</span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">oncall: transformer/mha</span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">triaged</span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## 🚀 Feature

## Motivation
Currently when using the `nn.MultiHeadAttention` <span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">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](https://arxiv.org/abs/1906.04341v1) 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](https://pytorch.org/docs/stable/nn.html#multiheadattention))

## Alternatives

## Additional context
A small discussion about this subject with a potential solution was made [here](https://discuss.pytorch.org/t/getting-nn-multiheadattention-attention-weights-for-each-head/72195)

If you guys agree, I'll gladly make a PR.</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>