Embed_dim must be divisible by num_heads

slavavs: <p>I get an error: AssertionError: embed_dim must be divisible by num_heads<br>
(assert self.head_dim * num_heads == self.embed_dim, “embed_dim must be divisible by num_heads”)<br>
I can’t understand how to fix it.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class Trans(nn.Module):
    def __init__(self, d_model=10):
        super().__init__()
        self.tr = nn.Transformer(d_model)
        
    def forward(self, src, tgt):
        out = self.tr(src, tgt)
        
        return out #output: (T, N, E)
    
net = Trans() 

sr = torch.randn(5, 1, 10) #src: (S, N, E)
tg = torch.randn(1, 1, 10) #tgt: (T, N, E)
outputs = net(sr, tg) 

print(outputs)
</code></pre>

andreaskoepf: <p>The <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.Transformer" rel="nofollow noopener">nn.Transformer</a> module by default uses 8 attention heads. Since the MultiHeadedAttention impl slices the model up into the number of head blocks (simply by a view operation) the model dimension must be a divisible by the number of heads. Please see also the documentation of <a href="https://pytorch.org/docs/stable/nn.html?highlight=multiheadattention#torch.nn.MultiheadAttention" rel="nofollow noopener">nn.MultiheadAttention</a>.</p>

cgarciae: <p>I don’t know why the implemented it like this, its a bit annoying. If instead of <code>embed_dim</code> being an input they asked you for <code>head_dim</code> and they calculated <code>embed_dim</code> as:</p>
<pre><code class="lang-python">self.embed_dim = self.head_dim * num_heads
</code></pre>
<p>It would be much easier to understand because you can do a per-head reasoning when defining the shapes, and it would also guarantee you never get an error. The easiest way to deal with this is just to do the math on your own outside:</p>
<pre><code class="lang-python">nhead = 5
head_dim = 32
dmodel = nhead * head_dim

net = nn.Transformer(d_model=d_model, nhead=nhead)
</code></pre>

Nat: <p>I am also confused with the PyTorch implementation.</p>
<p>In the <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L4837" rel="noopener nofollow ugc">multi_head_attention_forward</a> method Q, K and V matrices are first calculated with respect to the shapes of the source and target tensors (S,N,E) and (T,N,E), E being the size of each embedding.<br>
But <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L5024" rel="noopener nofollow ugc">then</a> they are reshaped splitting the embedding dimension as <a class="mention" href="/u/cgarciae">@cgarciae</a> just said<br>
So why ? As the goal of MHA is to project these matrices on each head parameters and then concat them and project again on W_0 ?<br>
The formula is even in the <a href="https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html" rel="noopener nofollow ugc">documentation</a> but is not applied properly ?<br>
This is very unclear, is there something I totally missed ?<br>
EDIT: I found what I missed, I got confused by the formula itself which doesn’t distingue the Q/K/V slices for each head</p>

Nick_Nguyen: <p>Hi <a class="mention" href="/u/nat">@Nat</a><br>
From my understanding, Pytorch forces the embedding size to be consistent all over the computation. Hence, the embed_dim must be divisible by num_heads so later on when you “concatenate” all heads, the matrix size will be embed_dim.</p>
<p>The use of W0 in the documentation you showed above is not for reshaping the concatenate of heads back to embed_dim.<br>
Here is the proof.<br>
You can notice this <a href="https://github.com/pytorch/pytorch/blob/bacf8ecbd16567de7f7bcf126ab22d62cf01f8a3/torch/nn/functional.py#L5084" rel="noopener nofollow ugc">code line</a><br>
attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)<br>
attn_output = linear(attn_output, out_proj_weight, out_proj_bias)</p>
<p>attn_output is reshaped back to embed_dim already (no need W0, out_proj_weight, to help reshape)<br>
In the 2nd line, out_proj_weight (W0) is initiated with <a href="https://github.com/pytorch/pytorch/blob/bacf8ecbd16567de7f7bcf126ab22d62cf01f8a3/torch/nn/modules/activation.py#L939" rel="noopener nofollow ugc">dimension (embed_dim, embed_dim)</a><br>
So you can see that the use of W0 is not for reshaping concatenate heads matrix.</p>
<p>To summary, with this implementation, Pytorch forces the matrix concatenate attention heads to have dimension is embed_dim. Hence, follow Attention is all you need paper, the purpose of using W0 to reshape concatenate heads to embed_dim is no need. Adding more evidence this statement, as shown above, out_proj_weight (W0) is initiated with size (embed_dim, embed_dim). So, W0 is not used for reshaping the “concatenate” attention heads matrix.</p>

ecolss: <p>Yes, this is quite confusing, why it’s required to have <code>embed_dim = head_dim * num_heads</code>?</p>

jasperhyp: <p>I know it’s pretty late but I also came across the same issue.  After some search my guess is that PyTorch just didn’t implement the original version, which would have linearly transformed the input to multiple heads instead of reshape-and-split. A nice snippet showing what the original version should look like can be found <a href="https://nn.labml.ai/transformers/mha.html" rel="noopener nofollow ugc">here</a>.  I’m not sure if this would negatively impact the performance, and whether or not one need to manually append linear layers before and after though.</p>

Goldname: <p>I do not understand why they enforce embedding dimension of input = embedding dimension post linear transform.</p>

fiendfish: <p>This really has nothing to do with torch, but is an implementation detail / design  choice of how multi-head attention usually is implemented.</p>
<p>If you want to have multi-head attention with d_head == d_model, each head would require its own QKV matrix.</p>
<p>With the Torch implementation, you instead have one QKV matrix that projects into a space of size d_model.<br>
This just gets cut up into head_count slices of size d_head = d_model / head_count.</p>
<p>This lets each head do a normal attention on these reduced QKV vectors.</p>
<p>After that, you just concat and get back to the original d_model, and do the final matmul and non-linearity.</p>
<p>This has the upsides:</p>
<p>head_count doesn’t increase the QKV matmul cost</p>
<p>heads can share a single KV cache</p>
<p>heads do attention in a much lower-dimensional space</p>
<p>In practice, you want lots of heads because each head can focus on a specific thing.<br>
This works fine even in a lower-dimensional space.</p>
<p>If compute cost scaled linearly with head count, big head counts simply wouldn’t be feasible.</p>