Questions about `torch.nn.MultiheadAttention`

username: namespace_pt
cooked: <p>I don’t understand how <code>nn.MultiheadAttention</code> module work.</p>
<ol>
<li>
<p>What is the meaning of <code>k_dim</code> and <code>v_dim</code> respectively in <code>__init__</code>? Are they related to <em>key</em> and <em>value</em> parameters in <code>forward()</code>?</p>
</li>
<li>
<p>Why must the <code>embedding_dim</code> be divisible by <code>num_heads</code>? What is the meaning of <code>head_dim</code> derived from such a fomulation?</p>
<p>I thought <code>multi-head self-attention</code> works this way:</p>
</li>
</ol>
<ul>
<li>
<code>embedding_dim</code> are projected into <code>query_dim</code>, <code>key_dim</code>, <code>value_dim</code> respectively</li>
<li>after computing attention weights between query and key, the <code>value_dim</code> are collected according to the attention weights of each query, which forms exactly the output of current attention head.</li>
<li>Finally output of all heads are concatenated into one matrix of <code>~</code>*<code>head_num</code>*<code>value_dim</code> and projected into <code>embedding_dim</code>. Thus, I don’t see where should the module divide <code>embedding_dim</code> by <code>num_heads</code>.</li>
</ul>
<ol start="3">
<li>What is the difference between <code>torchtext.nn.MultiheadAttention</code> and <code>torch.nn.MultiheadAttention</code>?</li>
</ol>

username: steffenN
cooked: <p>UP! I am very confused by this as well! Can anyone clarify?</p>
<p>Also: how can kdim be different from embed_dim (which is hardcoded to be qdim, which also makes no sense / is not very general)? In order to take the dot product between queries and keys, they have to have the same length! The only one between q,k,v that should be able to have a different length is v.</p>
<p>And vdim should also determine the size of the output, but then out_proj (the Matrix to projekt that output to the final output of the layer) is set up with the shape (embed_dim, embed_dim), when it should be (vdim, embed_dim), or better (vdim, out_dim) to be more general.</p>
<p>It seems that this layer can only work when embed_dim=kdim=vdim, which is the case any way when you do not specify kdim and vdim, so whats the point?</p>
<p>I have not yet tried to play around with it to see if it really doesn’t work when kdim or vdim are different from embed_dim, but just from looking at the code and the doc, this seems all a bit messy and not very well documented at all!</p>

username: AbdulsalamBande
cooked: <p><a class="mention" href="/u/namespace_pt">@namespace_pt</a>  <a class="mention" href="/u/steffenn">@steffenN</a></p>
<p>Question1 : <code>k_dim and v_dim</code> are the dimensions of your inputs(same as key and value in self attention)<br>
Question 2 : It must be divisible because. <code>embedding_dim</code> is divided across different heads. So if your <code>embedding_dim = 300</code> and you have <code>num_heads = 2</code>. The first head words on  <code>150</code> part of the embedding and the second head works on the other <code>150</code>, the results of the two heads are later concatenated.</p>
<p>Please watch <a href="https://discuss.pytorch.org/t/masking-in-torch-nn-multiheadattention/132060">torch.nn.MultiheadAttention</a>. It can clear all your doubts</p>

username: Qualia
cooked: <p>Hi <a class="mention" href="/u/abdulsalambande">@AbdulsalamBande</a>,</p>
<p>Thank you for your answer, I am also confused by point 2. I understand your answer, though, it doesn’t make sense to me that it works like that.</p>
<p>It is very weird to me to prevent a head from seeing parts of the features of the input. Going back to your example, you deprive each head from half of the feature of the input.</p>
<p>Extending your example, let’s say head 1 specializes in semantics and head 2 specializes in logic connectors in the sentence. If you give only half of the features to head 1, then it is difficult for the head to do its job because the split it receives may miss some important features from the input that contain a lot of information on semantics.<br>
The same goes for head 2, the split it receives may miss some important features regarding logic.</p>
<p>I don’t understand how this architecture can be efficient, I am lacking the intuition here. What would make sense is that the split would happen AFTER the projection of Q, K and V. But then, we would not need <code>embedding_dim</code> to be divisible by <code>num_heads</code>.</p>
<p>Thanks</p>

username: mrich
cooked: <aside class="quote no-group" data-username="Qualia" data-post="5" data-topic="103752">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/q/ce7236/48.png" class="avatar"> Qualia:</div>
<blockquote>
<p>Thank you for your answer, I am also confused by point 2. I understand your answer, though, it doesn’t make sense to me that it works like that.</p>
<p>It is very weird to me to prevent a head from seeing parts of the features of the input. Going back to your example, you deprive each head from half of the feature of the input.</p>
</blockquote>
</aside>
<p>Your transformer model has a certain dimensionality, like 512. Each “head” gets parts of that vector to hold it’s representation. So if you have 512 dimensionality vector representation, and 8 heads, each head gets 512/8 = 64 numbers to store it’s representation. The dimensionality is not about the <em>features</em> the head gets access to, it’s about the amount of numbers the heads gets to use to store <em>its representation</em>.</p>