Inputs for torch.nn.MultiheadAttention

ksasi: <p>Dear Community,</p>
<p>I am reviewing the API definition and documentation for “torch.nn.MultiheadAttention”. It appears the forward method needs <em>query</em> , <em>key</em> , <em>value</em>. Does this mean query, key, value needs to be learnt outside “torch.nn.MultiheadAttention”?</p>
<p>Assuming x is the input, does this mean the following is incorrect? (we are directly passing the same input x for query, key and value instead of obtaining these embeddings)</p>
<p>self_attent = torch.nn.MultiheadAttention(embed_dim=256 , num_heads=8)<br>
attn_output, attn_output_weights = self_attent(x, x, x)</p>
<p>Please advise.</p>
<p>Thanks !!!</p>

nairbv: <p>it’ll do the projections internally. see parameters like <code>q_proj_weight </code> in <a href="https://pytorch.org/docs/stable/_modules/torch/nn/modules/activation.html#MultiheadAttention" class="inline-onebox" rel="noopener nofollow ugc">torch.nn.modules.activation — PyTorch 2.1 documentation</a></p>
<p>Code like in your example <code>self_attent(x, x, x)</code> is pretty typical, and the k/q/v get projected from that just fine. Distinct k/q/v can matter in cases like cross attention, because the size of the query doesn’t have to be the same as the size of the k/v. If you’re doing a decoder-only or encoder-only model it shouldn’t matter.</p>
<p>You can find more detail in how specific cases are implemented in attention.h, e.g.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/main/torch/csrc/api/include/torch/nn/functional/activation.h#L685">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/main/torch/csrc/api/include/torch/nn/functional/activation.h#L685" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/main/torch/csrc/api/include/torch/nn/functional/activation.h#L685" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/main/torch/csrc/api/include/torch/nn/functional/activation.h#L685</a></h4>



    <pre class="onebox"><code class="lang-h">
      <ol class="start lines" start="675" style="counter-reset: li-counter 674 ;">
          <li>TORCH_INTERNAL_ASSERT(key.sizes() == value.sizes());</li>
          <li></li>
          <li>const auto head_dim = embed_dim / num_heads;</li>
          <li>TORCH_CHECK(</li>
          <li>    head_dim * num_heads == embed_dim,</li>
          <li>    "embed_dim must be divisible by num_heads");</li>
          <li>const auto scaling = 1 / std::sqrt(head_dim);</li>
          <li></li>
          <li>Tensor q, k, v;</li>
          <li>if (!use_separate_proj_weight) {</li>
          <li class="selected">  if (torch::equal(query, key) &amp;&amp; torch::equal(key, value)) {</li>
          <li>    // self-attention</li>
          <li>    const auto chunks =</li>
          <li>        F::linear(query, in_proj_weight, in_proj_bias).chunk(3, /*dim=*/-1);</li>
          <li>    q = chunks[0];</li>
          <li>    k = chunks[1];</li>
          <li>    v = chunks[2];</li>
          <li>  } else if (torch::equal(key, value)) {</li>
          <li>    // encoder-decoder attention</li>
          <li>    // This is inline in_proj function with in_proj_weight and in_proj_bias</li>
          <li>    auto _b = in_proj_bias;</li>
      </ol>
    </code></pre>



  </article>

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

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


ksasi: <p>Thankyou. This is very helpful.</p>