MultiheadAttention module - How do I have to set embed_dim, kdim, vdim?

dlsf: <p>I tried to set k_dim differently than embed_dim, but I got an error stating they should be the same!? I am wondering how one needs to set the values in the module, what’s the point if you can’t set them differently?</p>
<p>I looked at the source code, now I am wondering if I misunderstand the implementation or if there is an error in the implementation? The dimension for the q embedding should be the same as kdim in my opinion (in red), as q and v are used for calculating the correlations. Why is it set to embed_dim?</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/7/37d1396a39783d6fed0086b89b189d0b3e627a91.png" data-download-href="https://discuss.pytorch.org/uploads/default/37d1396a39783d6fed0086b89b189d0b3e627a91" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/7/37d1396a39783d6fed0086b89b189d0b3e627a91_2_606x500.png" alt="image" data-base62-sha1="7XMwlSz3wPtXVtaixR0X5XNcOIh" width="606" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/7/37d1396a39783d6fed0086b89b189d0b3e627a91_2_606x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/7/37d1396a39783d6fed0086b89b189d0b3e627a91_2_909x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/7/37d1396a39783d6fed0086b89b189d0b3e627a91_2_1212x1000.png 2x" data-dominant-color="E9E8EB"><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">image</span><span class="informations">1360×1122 164 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>

ptrblck: <p>I don’t get an error, which would state that <code>kdim</code> and <code>vdim</code> should be equal to the <code>embed_dim</code> as seen here:</p>
<pre><code class="lang-python">embed_dim = 10
num_heads = 2

multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)

L, S, N, E = 2, 3, 4, embed_dim
query = torch.randn(L, N, E)
key = torch.randn(S, N, E)
value = torch.randn(S, N, E)
attn_output, attn_output_weights = multihead_attn(query, key, value)


# with different kdim, vdim
kdim = 5
vdim = 6
multihead_attn = nn.MultiheadAttention(embed_dim, num_heads, kdim=kdim, vdim=vdim)

L, S, N, E = 2, 3, 4, embed_dim
query = torch.randn(L, N, E)
key = torch.randn(S, N, kdim)
value = torch.randn(S, N, vdim)
attn_output, attn_output_weights = multihead_attn(query, key, value)
</code></pre>
<p>Could you post a code snippet to reproduce this issue, as I might misunderstand it?</p>

dlsf: <p>Thanks for your answer. I think there is a misunderstanding from my side: I expected the module (nn.MultiheadAttention) to embed q, k, v values inside the module according to kdim, vdim and ideally qdim=kdim. The way the module is implemented, unfortunately it is not possible to use different embedding dimensions for the attention operation.</p>
<p>In my opinion it would be much more useful to be able to perform attention in individual dimensions for kdim and vdim. Is there a function or module that only performs attention without internal embedding stage?</p>

akuz: <p>I can’t believe this is still not fixed after 3 years. In the posted screenshot, obviously the highlighted embed_dim should be replaced with self.kdim (kdim and qdim must be equal to be able to compute scalar products). I guess everyone must be rolling their own attention code who needs kdim different from embed dim. I just found this amazing project, look for “Transformer’s attention needs more attention” on this page: <a href="https://einops.rocks/pytorch-examples.html" class="inline-onebox" rel="noopener nofollow ugc">Writing better code with pytorch+einops</a></p>