How can I retrieve the output from each head of MULTIHEADATTENTION?

abc50111: <p>I learned the Multi-Head Attention mechanism from this <a href="https://magazine.sebastianraschka.com/p/understanding-and-coding-self-attention" rel="noopener nofollow ugc">article</a>.<br>
If I don’t misunderstand this article, according to my understanding of the article, if I increase the <code>num_heads</code>, I should receive more output.</p>
<p>However, I’ve observed that regardless of the <code>num_heads</code> I set, the output shape of <code>MULTIHEADATTENTION</code> remains the same, contradicting what I learned from the article.</p>
<p>My codes snipets:</p>
<pre><code class="lang-auto">attn_2_heads = nn.MultiheadAttention(embed_dim=20, num_heads=2)
attn_5_heads = nn.MultiheadAttention(embed_dim=20, num_heads=5)

L, S, N, E = 2, 3, 4, 20
query = torch.randn(L, N, E)
key = torch.randn(S, N, E)
value = torch.randn(S, N, E)
attn_output_2_heads, _ = attn_2_heads(query, key, value)
attn_output_5_heads, _ = attn_5_heads(query, key, value)
print(attn_output_2_heads.shape, attn_output_5_heads.shape)
</code></pre>

J_Johnson: <p>The default behavior of <a href="https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html" rel="noopener nofollow ugc">nn.MultiheadAttention</a> is to average the weights over the heads.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd22a72b32b7fde39f6f8518d4cea82a1244d428.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/bd22a72b32b7fde39f6f8518d4cea82a1244d428" title="IMG_20240117_220838"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd22a72b32b7fde39f6f8518d4cea82a1244d428_2_426x499.jpeg" alt="IMG_20240117_220838" data-base62-sha1="qZau6LvsrBkeyoIUHBnvWpzscEo" width="426" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd22a72b32b7fde39f6f8518d4cea82a1244d428_2_426x499.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd22a72b32b7fde39f6f8518d4cea82a1244d428_2_639x748.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd22a72b32b7fde39f6f8518d4cea82a1244d428.jpeg 2x" data-dominant-color="F1F1F1"><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">IMG_20240117_220838</span><span class="informations">725×849 130 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>If you need the weights per head returned, set <code>average_attn_weights = False</code> during the forward method.</p>

abc50111: <p>Oh, I thought <code>attn_output_weights</code> is attention score.<br>
So should I use <code>attn_output</code> and <code>attn_output_weights</code> to retrieve the output of each head?<br>
How to do that?</p>

J_Johnson: <p>The actual function can be found here:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/02c96f69493ebd06d940a778a2d3d89c660a7ce4/torch/nn/functional.py#L5155">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/02c96f69493ebd06d940a778a2d3d89c660a7ce4/torch/nn/functional.py#L5155" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/02c96f69493ebd06d940a778a2d3d89c660a7ce4/torch/nn/functional.py#L5155" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/02c96f69493ebd06d940a778a2d3d89c660a7ce4/torch/nn/functional.py#L5155</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="5145" style="counter-reset: li-counter 5144 ;">
          <li>            )</li>
          <li>    return mask</li>
          <li></li>
          <li>def _none_or_dtype(input: Optional[Tensor]) -&gt; Optional[DType]:</li>
          <li>    if input is None:</li>
          <li>        return None</li>
          <li>    elif isinstance(input, torch.Tensor):</li>
          <li>        return input.dtype</li>
          <li>    raise RuntimeError("input to _none_or_dtype() must be None or torch.Tensor")</li>
          <li></li>
          <li class="selected">def multi_head_attention_forward(</li>
          <li>    query: Tensor,</li>
          <li>    key: Tensor,</li>
          <li>    value: Tensor,</li>
          <li>    embed_dim_to_check: int,</li>
          <li>    num_heads: int,</li>
          <li>    in_proj_weight: Optional[Tensor],</li>
          <li>    in_proj_bias: Optional[Tensor],</li>
          <li>    bias_k: Optional[Tensor],</li>
          <li>    bias_v: Optional[Tensor],</li>
          <li>    add_zero_attn: bool,</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>It is conditional on the choice of arguments given. But you could start there to determine how to apply the attn_output_weights to get the attn_output shape you want in your specific case.</p>