title: Why embed dimemsion must be divisible by num of heads in MultiheadAttention?
body: <p>I am learning the Transformer. Here is the pytorch document for <a href="https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html#torch.nn.MultiheadAttention" rel="noreferrer">MultiheadAttention</a>. In their <a href="https://github.com/pytorch/pytorch/blob/7a178a8a523d4653a3a2fa10c573b71e7fab1b9a/torch/nn/modules/activation.py#L874" rel="noreferrer">implementation</a>, I saw there is a constraint:</p>
  <pre><code> assert self.head_dim * num_heads == self.embed_dim, &quot;embed_dim must be divisible by num_heads&quot;
  </code></pre>
  <p>Why require the constraint: <code>embed_dim must be divisible by num_heads?</code>  If we go back to the equation</p>
  <p><a href="https://i.stack.imgur.com/T0niJ.png" rel="noreferrer"><img src="https://i.stack.imgur.com/T0niJ.png" alt="MultiHead(Q,K,V)=Concat(head1​,…,headh​)WOwhereheadi​=Attention(QWiQ​,KWiK​,VWiV​)" /></a></p>
  <p>Assume:
  <code>Q</code>, <code>K</code>,<code>V</code> are <code>n x emded_dim</code> matrices; all the weight matrices <code>W</code> is <code>emded_dim x head_dim</code>,</p>
  <p>Then, the concat <code>[head_i, ..., head_h]</code> will be a <code>n x (num_heads*head_dim)</code> matrix;</p>
  <p><code>W^O</code> with size <code>(num_heads*head_dim) x embed_dim</code></p>
  <p><code>[head_i, ..., head_h] * W^O</code> will become a <code>n x embed_dim</code> output</p>
  <p>I don't know why we require <code>embed_dim must be divisible by num_heads</code>.</p>
  <p>Let say we have <code>num_heads=10000</code>, the resuts are the same, since the matrix-matrix product will absort this information.</p>