Answers.Body: <p>The queries, keys and values are calculated then chunked so that each chunk depends on (is a linear combination of) all values of the input embedding.</p>
  <p>As for understanding an implementation, I didn't bother with pytorch but instead understood this code <a href="http://einops.rocks/pytorch-examples.html" rel="nofollow noreferrer">http://einops.rocks/pytorch-examples.html</a> although there are differences as I understand that pytorch expects input in the form <code>(L, N, C)</code> where L=words, N=batch, C=embedding for performance reasons.</p>
  <h1>Transformer's attention needs more attention</h1>
  <pre><code>class MultiHeadAttentionNew(nn.Module):
      def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
          super().__init__()
          self.n_head = n_head
          
          self.w_qs = nn.Linear(d_model, n_head * d_k)
          self.w_ks = nn.Linear(d_model, n_head * d_k)
          self.w_vs = nn.Linear(d_model, n_head * d_v)
          
          nn.init.normal_(self.w_qs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
          nn.init.normal_(self.w_ks.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
          nn.init.normal_(self.w_vs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_v)))
          
          self.fc = nn.Linear(n_head * d_v, d_model)
          nn.init.xavier_normal_(self.fc.weight)
          self.dropout = nn.Dropout(p=dropout)
          self.layer_norm = nn.LayerNorm(d_model)
  
      def forward(self, q, k, v, mask=None):
          residual = q
          q = rearrange(self.w_qs(q), 'b l (head k) -&gt; head b l k', head=self.n_head)
          k = rearrange(self.w_ks(k), 'b t (head k) -&gt; head b t k', head=self.n_head)
          v = rearrange(self.w_vs(v), 'b t (head v) -&gt; head b t v', head=self.n_head)
          attn = torch.einsum('hblk,hbtk-&gt;hblt', [q, k]) / np.sqrt(q.shape[-1])
          if mask is not None:
              attn = attn.masked_fill(mask[None], -np.inf)
          attn = torch.softmax(attn, dim=3)
          output = torch.einsum('hblt,hbtv-&gt;hblv', [attn, v])
          output = rearrange(output, 'head b l v -&gt; b l (head v)')
          output = self.dropout(self.fc(output))
          output = self.layer_norm(output + residual)
  </code></pre>
  <p>Understand MultiHeadAttentionNew first then understand MultiHeadAttentionOld, and then maybe look at the pytorch code.</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-05-18T12:42:32.530Z
Answers.Text: Thank you for the link! It is really helpful. While I do not understand the details of some of the einops functions, it is clear enough. Since also bigger matrices are used in MultiHeadAttentionNew it really seems that you can just use bigger matrices and it gets you the same result.
  
  I'll try to write it down mathematically.
Answers.UserDisplayName: Aushilfsgott
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-05-18T23:28:44.657Z
Answers.Text: @Aushilfsgott The design of multi-head attention is for performance rather than having multiple single head attentions with the same input. Also I should have mentioned in my answer its worth really understanding einops and einsum, rather using the older pytorch functions
Answers.UserDisplayName: Tom Huntington
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-05-19T05:23:53.557Z
Answers.Text: Thanks again Tom. I got the basic principle of einsum so it is fine! The link is definitely very good for understanding the function.
Answers.UserDisplayName: Aushilfsgott
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2022-05-18T09:29:03.223Z
Answers.LastActivityDate: 2022-05-18T09:29:03.223Z
Answers.UserDisplayName: Tom Huntington
Body: <p>I am confused by the Multi-Head part of the Multi-Head-Attention used in Transformers. My question concerns the implementations in Pytorch of <a href="https://pytorch.org/docs/stable/_modules/torch/nn/modules/activation.html#MultiheadAttention" rel="noreferrer">nn.MultiheadAttention</a> and its forward method multi_head_attention_forward and whether these are actually identical to the paper. Unfortunately, I have been unable to follow along the original code of the paper. So I could not check whether the implementations in Pytorch are acutally identical to the paper.</p>
  <p>Please forgive the excessive use of illustrations. However, I hope it will improve understanding my problem.</p>
  <p>What is the correct order for calculating the Queries Q, Keys K and Values V and splitting the operation into the individual Attention-Heads? Unfortunately most explanations I found online while helpful for understanding the general principle and intuition of Multi-Head-Attention did not go into the details of the implementation.</p>
  <p>In the original paper <a href="https://arxiv.org/abs/1706.03762" rel="noreferrer">Attention is all you need</a> Multi-Head-Attention is explained as followed:</p>
  <p><a href="https://i.sstatic.net/LmHJi.png" rel="noreferrer"><img src="https://i.sstatic.net/LmHJi.png" alt="enter image description here" /></a></p>
  <p>First, according to my current understanding, if we have a sequence of vectors with 512-dimensions (like in the original Transformer) and we have <span class="math-container">$h=8$</span> Attention-Heads (again like the original), every Attention-Head attends to <span class="math-container">$512/8=64$</span> entries of the input vector used to calculate the Attention in the corresponding head. So the first Attention-Head attends to the first 64 entries, the second to the second 64 entries and so on. However, if the split is conducted before calculating Q,K,V this would refer to the first 64 entries of X (this does not seem match the explanation in the paper I believe) while in the other case it would refer to the first 64 entries of Q,K,V.</p>
  <p>In the text they say &quot;project the queries, keys and values h times with different, learned linear projections to <span class="math-container">$d_k,d_k$</span> and <span class="math-container">$d_v$</span> dimensions and since they set <span class="math-container">$d_k=d_v=d_{model}/h=512/8=64$</span>. Therefore, if we actually have single matrices for every Attention-Head h we would have</p>
  <p><span class="math-container">$W^Q_i,W^K_i,W^V_i \in \mathbb{R}^{512x64} \forall i \in h$</span>.</p>
  <p>This matches the illustration found here <a href="https://jalammar.github.io/illustrated-transformer/" rel="noreferrer">https://jalammar.github.io/illustrated-transformer/</a></p>
  <p><a href="https://i.sstatic.net/ehiUv.png" rel="noreferrer"><img src="https://i.sstatic.net/ehiUv.png" alt="enter image description here" /></a></p>
  <p>It is explained that the input X is transformed into the Queries, Keys and Values for the different attention heads by using different projection matrices which are learned during training.</p>
  <p>This seems to indicate that the split into the individual Attention-Heads is conducted after the calculation of <span class="math-container">$Q,K,V$</span> (or rather during the calculation). Since we have <span class="math-container">$h=8$</span> this leads in sum to <span class="math-container">$3*8*512*64=3*512*512$</span> learnable parameters in total (if we ignore the bias). Thus as far as the overall number of parameters is concerned we would have the same number if we would instead use three big matrices which concatenate the matrices of the individual Attention-Heads.</p>
  <p><span class="math-container">$W^Q=[W^Q_1,W^Q_2,...,W^Q_h] \in \mathbb{R}^{512x512}$</span> <br />
  <span class="math-container">$W^K=[W^K_1,W^K_2,...,W^K_h] \in \mathbb{R}^{512x512}$</span> <br />
  <span class="math-container">$W^Q=[W^V_1,W^V_2,...,W^V_h] \in \mathbb{R}^{512x512}$</span></p>
  <p>In the explanation from the same author of GPT-2 (this model has an embedding dimension of 768 and 12 Attention-Heads, instead of 512 and 8 like the original Transformer) found here <a href="https://jalammar.github.io/illustrated-gpt2/#part-2-illustrated-self-attention" rel="noreferrer">https://jalammar.github.io/illustrated-gpt2/#part-2-illustrated-self-attention</a></p>
  <p><a href="https://i.sstatic.net/I8fz9.png" rel="noreferrer"><img src="https://i.sstatic.net/I8fz9.png" alt="enter image description here" /></a></p>
  <p><a href="https://i.sstatic.net/V75eY.png" rel="noreferrer"><img src="https://i.sstatic.net/V75eY.png" alt="enter image description here" /></a></p>
  <p>the Queries Q, Keys K and Values V are first calculated by multiplying the input with one big matrix which is the concatenation of <span class="math-container">$W^Q,W^K,W^V$</span>. If the input for calculating Q, K and V is identical (which is the case for self-attention), it is clear to me that you can use <span class="math-container">$W_{concat}=[W^Q,W^K,W^V]$</span> and obtain <span class="math-container">$[Q,K,V]$</span>, since you essentially still multiply the input with each weight matrix separately.</p>
  <p>Then you can split the result to again obtain <span class="math-container">$Q,K,V$</span> as individual matrices (The image displays <span class="math-container">$q_9,k_9,v_9$</span> as an example, not the complete matrices). Then <span class="math-container">$q_9,k_9,v_9$</span> are again split into 12 vectors which results in a matrix of dimension <span class="math-container">$(12x64)$</span>.</p>
  <p>So overall here we did not use individual matrices per Attention-Head but only one larger matrix.</p>
  <p>Is this method mathematically identical to the one using individual smaller matrices per Attention-Head?</p>
  <p>It appears that this is the way the calculation is implented in Pytorch, if <span class="math-container">$d_{model}=kdim=vdim$</span> Though here unlike in the paper which used <span class="math-container">$d_k$</span> and <span class="math-container">$d_v$</span> as names, <span class="math-container">$kdim$</span> and <span class="math-container">$vdim$</span> refer to the dimension of all Attention-Heads summed up, e.g. <span class="math-container">$kdim=d_k*num_{heads}$</span>(=512 for the original Transformer).</p>
  <p>So  In the documentation of nn.modules.MultiheadAttention the model either creates three separate projection matrices to generate the Queries, Keys and Values or one big matrix (if the dimensions are identical). The following is part of the <code>_init_</code> function.</p>
  <pre><code>if self._qkv_same_embed_dim is False:
              self.q_proj_weight = Parameter(torch.empty((embed_dim, embed_dim), **factory_kwargs))
              self.k_proj_weight = Parameter(torch.empty((embed_dim, self.kdim), **factory_kwargs))
              self.v_proj_weight = Parameter(torch.empty((embed_dim, self.vdim), **factory_kwargs))
              self.register_parameter('in_proj_weight', None)
          else:
              self.in_proj_weight = Parameter(torch.empty((3 * embed_dim, embed_dim), **factory_kwargs))
              self.register_parameter('q_proj_weight', None)
              self.register_parameter('k_proj_weight', None)
              self.register_parameter('v_proj_weight', None)
  </code></pre>
  <p>If we stay in the standard case of <code>_qkv_same_embed_dim=True</code> the input is passed through a <code>nn.linear</code> as part of <code>_in_projection_packed</code> which is using <code>self.in_proj_weight</code> as the weight</p>
  <pre><code>if not use_separate_proj_weight:
          assert in_proj_weight is not None, &quot;use_separate_proj_weight is False but in_proj_weight is None&quot;
          q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)
      else:
          assert q_proj_weight is not None, &quot;use_separate_proj_weight is True but q_proj_weight is None&quot;
          assert k_proj_weight is not None, &quot;use_separate_proj_weight is True but k_proj_weight is None&quot;
          assert v_proj_weight is not None, &quot;use_separate_proj_weight is True but v_proj_weight is None&quot;
          if in_proj_bias is None:
              b_q = b_k = b_v = None
          else:
              b_q, b_k, b_v = in_proj_bias.chunk(3)
          q, k, v = _in_projection(query, key, value, q_proj_weight, k_proj_weight, v_proj_weight, b_q, b_k, b_v)
  
  def _in_projection_packed(
      q: Tensor,
      k: Tensor,
      v: Tensor,
      w: Tensor,
      b: Optional[Tensor] = None,
  ) -&gt; List[Tensor]:
      r&quot;&quot;&quot;
      Performs the in-projection step of the attention operation, using packed weights.
      Output is a triple containing projection tensors for query, key and value.
      Args:
          q, k, v: query, key and value tensors to be projected. For self-attention,
              these are typically the same tensor; for encoder-decoder attention,
              k and v are typically the same tensor. (We take advantage of these
              identities for performance if they are present.) Regardless, q, k and v
              must share a common embedding dimension; otherwise their shapes may vary.
          w: projection weights for q, k and v, packed into a single tensor. Weights
              are packed along dimension 0, in q, k, v order.
          b: optional projection biases for q, k and v, packed into a single tensor
              in q, k, v order.
      Shape:
          Inputs:
          - q: :math:`(..., E)` where E is the embedding dimension
          - k: :math:`(..., E)` where E is the embedding dimension
          - v: :math:`(..., E)` where E is the embedding dimension
          - w: :math:`(E * 3, E)` where E is the embedding dimension
          - b: :math:`E * 3` where E is the embedding dimension
          Output:
          - in output list :math:`[q', k', v']`, each output tensor will have the
              same shape as the corresponding input tensor.
      &quot;&quot;&quot;
      E = q.size(-1)
      if k is v:
          if q is k:
              # self-attention
              return linear(q, w, b).chunk(3, dim=-1)
          else:
              # encoder-decoder attention
              w_q, w_kv = w.split([E, E * 2])
              if b is None:
                  b_q = b_kv = None
              else:
                  b_q, b_kv = b.split([E, E * 2])
              return (linear(q, w_q, b_q),) + linear(k, w_kv, b_kv).chunk(2, dim=-1)
      else:
          w_q, w_k, w_v = w.chunk(3)
          if b is None:
              b_q = b_k = b_v = None
          else:
              b_q, b_k, b_v = b.chunk(3)
          return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
  
  
  def _in_projection(
      q: Tensor,
      k: Tensor,
      v: Tensor,
      w_q: Tensor,
      w_k: Tensor,
      w_v: Tensor,
      b_q: Optional[Tensor] = None,
      b_k: Optional[Tensor] = None,
      b_v: Optional[Tensor] = None,
  ) -&gt; Tuple[Tensor, Tensor, Tensor]:
      r&quot;&quot;&quot;
      Performs the in-projection step of the attention operation. This is simply
      a triple of linear projections, with shape constraints on the weights which
      ensure embedding dimension uniformity in the projected outputs.
      Output is a triple containing projection tensors for query, key and value.
      Args:
          q, k, v: query, key and value tensors to be projected.
          w_q, w_k, w_v: weights for q, k and v, respectively.
          b_q, b_k, b_v: optional biases for q, k and v, respectively.
      Shape:
          Inputs:
          - q: :math:`(Qdims..., Eq)` where Eq is the query embedding dimension and Qdims are any
              number of leading dimensions.
          - k: :math:`(Kdims..., Ek)` where Ek is the key embedding dimension and Kdims are any
              number of leading dimensions.
          - v: :math:`(Vdims..., Ev)` where Ev is the value embedding dimension and Vdims are any
              number of leading dimensions.
          - w_q: :math:`(Eq, Eq)`
          - w_k: :math:`(Eq, Ek)`
          - w_v: :math:`(Eq, Ev)`
          - b_q: :math:`(Eq)`
          - b_k: :math:`(Eq)`
          - b_v: :math:`(Eq)`
          Output: in output triple :math:`(q', k', v')`,
           - q': :math:`[Qdims..., Eq]`
           - k': :math:`[Kdims..., Eq]`
           - v': :math:`[Vdims..., Eq]`
      &quot;&quot;&quot;
      Eq, Ek, Ev = q.size(-1), k.size(-1), v.size(-1)
      assert w_q.shape == (Eq, Eq), f&quot;expecting query weights shape of {(Eq, Eq)}, but got {w_q.shape}&quot;
      assert w_k.shape == (Eq, Ek), f&quot;expecting key weights shape of {(Eq, Ek)}, but got {w_k.shape}&quot;
      assert w_v.shape == (Eq, Ev), f&quot;expecting value weights shape of {(Eq, Ev)}, but got {w_v.shape}&quot;
      assert b_q is None or b_q.shape == (Eq,), f&quot;expecting query bias shape of {(Eq,)}, but got {b_q.shape}&quot;
      assert b_k is None or b_k.shape == (Eq,), f&quot;expecting key bias shape of {(Eq,)}, but got {b_k.shape}&quot;
      assert b_v is None or b_v.shape == (Eq,), f&quot;expecting value bias shape of {(Eq,)}, but got {b_v.shape}&quot;
      return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
  
  def linear(
      input: Tensor, weight: Tensor, bias: Optional[Tensor] = None,
      scale: Optional[float] = None, zero_point: Optional[int] = None
  ) -&gt; Tensor:
      r&quot;&quot;&quot;
      Applies a linear transformation to the incoming quantized data:
      :math:`y = xA^T + b`.
      See :class:`~torch.nn.quantized.Linear`
      .. note::
        Current implementation packs weights on every call, which has penalty on performance.
        If you want to avoid the overhead, use :class:`~torch.nn.quantized.Linear`.
      Args:
        input (Tensor): Quantized input of type `torch.quint8`
        weight (Tensor): Quantized weight of type `torch.qint8`
        bias (Tensor): None or fp32 bias of type `torch.float`
        scale (double): output scale. If None, derived from the input scale
        zero_point (long): output zero point. If None, derived from the input zero_point
      Shape:
          - Input: :math:`(N, *, in\_features)` where `*` means any number of
            additional dimensions
          - Weight: :math:`(out\_features, in\_features)`
          - Bias: :math:`(out\_features)`
          - Output: :math:`(N, *, out\_features)`
      &quot;&quot;&quot;
      if scale is None:
          scale = input.q_scale()
      if zero_point is None:
          zero_point = input.q_zero_point()
      _packed_params = torch.ops.quantized.linear_prepack(weight, bias)
      return torch.ops.quantized.linear(input, _packed_params, scale, zero_point)
  </code></pre>
  <p>Then later the Queries, Keys and Values are split up into the individual Attention-Heads:</p>
  <pre><code>    #
      # reshape q, k, v for multihead attention and make em batch first
      #
      q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
      if static_k is None:
          k = k.contiguous().view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
      else:
          # TODO finish disentangling control flow so we don't do in-projections when statics are passed
          assert static_k.size(0) == bsz * num_heads, \
              f&quot;expecting static_k.size(0) of {bsz * num_heads}, but got {static_k.size(0)}&quot;
          assert static_k.size(2) == head_dim, \
              f&quot;expecting static_k.size(2) of {head_dim}, but got {static_k.size(2)}&quot;
          k = static_k
      if static_v is None:
          v = v.contiguous().view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
      else:
          # TODO finish disentangling control flow so we don't do in-projections when statics are passed
          assert static_v.size(0) == bsz * num_heads, \
              f&quot;expecting static_v.size(0) of {bsz * num_heads}, but got {static_v.size(0)}&quot;
          assert static_v.size(2) == head_dim, \
              f&quot;expecting static_v.size(2) of {head_dim}, but got {static_v.size(2)}&quot;
          v = static_v
  </code></pre>
  <p>Therefore, it appears to me that both ways of calculations should be equal and the one using the bigger matrices is just more efficient to compute. Is this correct?</p>
  <p>In this case, I ask myself why nn.MultiheadAttention requires that <code>embed_dim</code> is divisible by <code>num_heads</code>, since the split into the individual Attention-Heads is actually conducted after generating <span class="math-container">$Q,K,V$</span>. Should then not <span class="math-container">$d_q,d_k,d_v$</span> be made to be divisible by <code>num_heads</code>? Since these dimensions do not have to be equal to the dimension of the inputs?</p>
  <p>Thanks for any advice in advance!</p>
  
Tags: <transformer><pytorch><attention>
Title: When exactly does the split into different heads in Multi-Head-Attention occur?
UserDisplayName: Aushilfsgott
ThreadText: # When exactly does the split into different heads in Multi-Head-Attention occur?
  
  (Asked by: username_0 on 2022-05-17)
  
  I am confused by the Multi-Head part of the Multi-Head-Attention used in Transformers. My question concerns the implementations in Pytorch of nn.MultiheadAttention and its forward method multi_head_attention_forward and whether these are actually identical to the paper. Unfortunately, I have been unable to follow along the original code of the paper. So I could not check whether the implementations in Pytorch are acutally identical to the paper.
  Please forgive the excessive use of illustrations. However, I hope it will improve understanding my problem.
  What is the correct order for calculating the Queries Q, Keys K and Values V and splitting the operation into the individual Attention-Heads? Unfortunately most explanations I found online while helpful for understanding the general principle and intuition of Multi-Head-Attention did not go into the details of the implementation.
  In the original paper Attention is all you need Multi-Head-Attention is explained as followed:
  
  First, according to my current understanding, if we have a sequence of vectors with 512-dimensions (like in the original Transformer) and we have $h=8$ Attention-Heads (again like the original), every Attention-Head attends to $512/8=64$ entries of the input vector used to calculate the Attention in the corresponding head. So the first Attention-Head attends to the first 64 entries, the second to the second 64 entries and so on. However, if the split is conducted before calculating Q,K,V this would refer to the first 64 entries of X (this does not seem match the explanation in the paper I believe) while in the other case it would refer to the first 64 entries of Q,K,V.
  In the text they say "project the queries, keys and values h times with different, learned linear projections to $d_k,d_k$ and $d_v$ dimensions and since they set $d_k=d_v=d_{model}/h=512/8=64$. Therefore, if we actually have single matrices for every Attention-Head h we would have
  $W^Q_i,W^K_i,W^V_i \in \mathbb{R}^{512x64} \forall i \in h$.
  This matches the illustration found here https://jalammar.github.io/illustrated-transformer/
  
  It is explained that the input X is transformed into the Queries, Keys and Values for the different attention heads by using different projection matrices which are learned during training.
  This seems to indicate that the split into the individual Attention-Heads is conducted after the calculation of $Q,K,V$ (or rather during the calculation). Since we have $h=8$ this leads in sum to $3*8*512*64=3*512*512$ learnable parameters in total (if we ignore the bias). Thus as far as the overall number of parameters is concerned we would have the same number if we would instead use three big matrices which concatenate the matrices of the individual Attention-Heads.
  $W^Q=[W^Q_1,W^Q_2,...,W^Q_h] \in \mathbb{R}^{512x512}$ 
  $W^K=[W^K_1,W^K_2,...,W^K_h] \in \mathbb{R}^{512x512}$ 
  $W^Q=[W^V_1,W^V_2,...,W^V_h] \in \mathbb{R}^{512x512}$
  In the explanation from the same author of GPT-2 (this model has an embedding dimension of 768 and 12 Attention-Heads, instead of 512 and 8 like the original Transformer) found here https://jalammar.github.io/illustrated-gpt2/#part-2-illustrated-self-attention
  
  
  the Queries Q, Keys K and Values V are first calculated by multiplying the input with one big matrix which is the concatenation of $W^Q,W^K,W^V$. If the input for calculating Q, K and V is identical (which is the case for self-attention), it is clear to me that you can use $W_{concat}=[W^Q,W^K,W^V]$ and obtain $[Q,K,V]$, since you essentially still multiply the input with each weight matrix separately.
  Then you can split the result to again obtain $Q,K,V$ as individual matrices (The image displays $q_9,k_9,v_9$ as an example, not the complete matrices). Then $q_9,k_9,v_9$ are again split into 12 vectors which results in a matrix of dimension $(12x64)$.
  So overall here we did not use individual matrices per Attention-Head but only one larger matrix.
  Is this method mathematically identical to the one using individual smaller matrices per Attention-Head?
  It appears that this is the way the calculation is implented in Pytorch, if $d_{model}=kdim=vdim$ Though here unlike in the paper which used $d_k$ and $d_v$ as names, $kdim$ and $vdim$ refer to the dimension of all Attention-Heads summed up, e.g. $kdim=d_k*num_{heads}$(=512 for the original Transformer).
  So  In the documentation of nn.modules.MultiheadAttention the model either creates three separate projection matrices to generate the Queries, Keys and Values or one big matrix (if the dimensions are identical). The following is part of the _init_ function.
  if self._qkv_same_embed_dim is False:
              self.q_proj_weight = Parameter(torch.empty((embed_dim, embed_dim), **factory_kwargs))
              self.k_proj_weight = Parameter(torch.empty((embed_dim, self.kdim), **factory_kwargs))
              self.v_proj_weight = Parameter(torch.empty((embed_dim, self.vdim), **factory_kwargs))
              self.register_parameter('in_proj_weight', None)
          else:
              self.in_proj_weight = Parameter(torch.empty((3 * embed_dim, embed_dim), **factory_kwargs))
              self.register_parameter('q_proj_weight', None)
              self.register_parameter('k_proj_weight', None)
              self.register_parameter('v_proj_weight', None)
  
  If we stay in the standard case of _qkv_same_embed_dim=True the input is passed through a nn.linear as part of _in_projection_packed which is using self.in_proj_weight as the weight
  if not use_separate_proj_weight:
          assert in_proj_weight is not None, "use_separate_proj_weight is False but in_proj_weight is None"
          q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)
      else:
          assert q_proj_weight is not None, "use_separate_proj_weight is True but q_proj_weight is None"
          assert k_proj_weight is not None, "use_separate_proj_weight is True but k_proj_weight is None"
          assert v_proj_weight is not None, "use_separate_proj_weight is True but v_proj_weight is None"
          if in_proj_bias is None:
              b_q = b_k = b_v = None
          else:
              b_q, b_k, b_v = in_proj_bias.chunk(3)
          q, k, v = _in_projection(query, key, value, q_proj_weight, k_proj_weight, v_proj_weight, b_q, b_k, b_v)
  
  def _in_projection_packed(
      q: Tensor,
      k: Tensor,
      v: Tensor,
      w: Tensor,
      b: Optional[Tensor] = None,
  ) -> List[Tensor]:
      r"""
      Performs the in-projection step of the attention operation, using packed weights.
      Output is a triple containing projection tensors for query, key and value.
      Args:
          q, k, v: query, key and value tensors to be projected. For self-attention,
              these are typically the same tensor; for encoder-decoder attention,
              k and v are typically the same tensor. (We take advantage of these
              identities for performance if they are present.) Regardless, q, k and v
              must share a common embedding dimension; otherwise their shapes may vary.
          w: projection weights for q, k and v, packed into a single tensor. Weights
              are packed along dimension 0, in q, k, v order.
          b: optional projection biases for q, k and v, packed into a single tensor
              in q, k, v order.
      Shape:
          Inputs:
          - q: :math:`(..., E)` where E is the embedding dimension
          - k: :math:`(..., E)` where E is the embedding dimension
          - v: :math:`(..., E)` where E is the embedding dimension
          - w: :math:`(E * 3, E)` where E is the embedding dimension
          - b: :math:`E * 3` where E is the embedding dimension
          Output:
          - in output list :math:`[q', k', v']`, each output tensor will have the
              same shape as the corresponding input tensor.
      """
      E = q.size(-1)
      if k is v:
          if q is k:
              # self-attention
              return linear(q, w, b).chunk(3, dim=-1)
          else:
              # encoder-decoder attention
              w_q, w_kv = w.split([E, E * 2])
              if b is None:
                  b_q = b_kv = None
              else:
                  b_q, b_kv = b.split([E, E * 2])
              return (linear(q, w_q, b_q),) + linear(k, w_kv, b_kv).chunk(2, dim=-1)
      else:
          w_q, w_k, w_v = w.chunk(3)
          if b is None:
              b_q = b_k = b_v = None
          else:
              b_q, b_k, b_v = b.chunk(3)
          return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
  
  
  def _in_projection(
      q: Tensor,
      k: Tensor,
      v: Tensor,
      w_q: Tensor,
      w_k: Tensor,
      w_v: Tensor,
      b_q: Optional[Tensor] = None,
      b_k: Optional[Tensor] = None,
      b_v: Optional[Tensor] = None,
  ) -> Tuple[Tensor, Tensor, Tensor]:
      r"""
      Performs the in-projection step of the attention operation. This is simply
      a triple of linear projections, with shape constraints on the weights which
      ensure embedding dimension uniformity in the projected outputs.
      Output is a triple containing projection tensors for query, key and value.
      Args:
          q, k, v: query, key and value tensors to be projected.
          w_q, w_k, w_v: weights for q, k and v, respectively.
          b_q, b_k, b_v: optional biases for q, k and v, respectively.
      Shape:
          Inputs:
          - q: :math:`(Qdims..., Eq)` where Eq is the query embedding dimension and Qdims are any
              number of leading dimensions.
          - k: :math:`(Kdims..., Ek)` where Ek is the key embedding dimension and Kdims are any
              number of leading dimensions.
          - v: :math:`(Vdims..., Ev)` where Ev is the value embedding dimension and Vdims are any
              number of leading dimensions.
          - w_q: :math:`(Eq, Eq)`
          - w_k: :math:`(Eq, Ek)`
          - w_v: :math:`(Eq, Ev)`
          - b_q: :math:`(Eq)`
          - b_k: :math:`(Eq)`
          - b_v: :math:`(Eq)`
          Output: in output triple :math:`(q', k', v')`,
           - q': :math:`[Qdims..., Eq]`
           - k': :math:`[Kdims..., Eq]`
           - v': :math:`[Vdims..., Eq]`
      """
      Eq, Ek, Ev = q.size(-1), k.size(-1), v.size(-1)
      assert w_q.shape == (Eq, Eq), f"expecting query weights shape of {(Eq, Eq)}, but got {w_q.shape}"
      assert w_k.shape == (Eq, Ek), f"expecting key weights shape of {(Eq, Ek)}, but got {w_k.shape}"
      assert w_v.shape == (Eq, Ev), f"expecting value weights shape of {(Eq, Ev)}, but got {w_v.shape}"
      assert b_q is None or b_q.shape == (Eq,), f"expecting query bias shape of {(Eq,)}, but got {b_q.shape}"
      assert b_k is None or b_k.shape == (Eq,), f"expecting key bias shape of {(Eq,)}, but got {b_k.shape}"
      assert b_v is None or b_v.shape == (Eq,), f"expecting value bias shape of {(Eq,)}, but got {b_v.shape}"
      return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
  
  def linear(
      input: Tensor, weight: Tensor, bias: Optional[Tensor] = None,
      scale: Optional[float] = None, zero_point: Optional[int] = None
  ) -> Tensor:
      r"""
      Applies a linear transformation to the incoming quantized data:
      :math:`y = xA^T + b`.
      See :class:`~torch.nn.quantized.Linear`
      .. note::
        Current implementation packs weights on every call, which has penalty on performance.
        If you want to avoid the overhead, use :class:`~torch.nn.quantized.Linear`.
      Args:
        input (Tensor): Quantized input of type `torch.quint8`
        weight (Tensor): Quantized weight of type `torch.qint8`
        bias (Tensor): None or fp32 bias of type `torch.float`
        scale (double): output scale. If None, derived from the input scale
        zero_point (long): output zero point. If None, derived from the input zero_point
      Shape:
          - Input: :math:`(N, *, in\_features)` where `*` means any number of
            additional dimensions
          - Weight: :math:`(out\_features, in\_features)`
          - Bias: :math:`(out\_features)`
          - Output: :math:`(N, *, out\_features)`
      """
      if scale is None:
          scale = input.q_scale()
      if zero_point is None:
          zero_point = input.q_zero_point()
      _packed_params = torch.ops.quantized.linear_prepack(weight, bias)
      return torch.ops.quantized.linear(input, _packed_params, scale, zero_point)
  
  Then later the Queries, Keys and Values are split up into the individual Attention-Heads:
      #
      # reshape q, k, v for multihead attention and make em batch first
      #
      q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
      if static_k is None:
          k = k.contiguous().view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
      else:
          # TODO finish disentangling control flow so we don't do in-projections when statics are passed
          assert static_k.size(0) == bsz * num_heads, \
              f"expecting static_k.size(0) of {bsz * num_heads}, but got {static_k.size(0)}"
          assert static_k.size(2) == head_dim, \
              f"expecting static_k.size(2) of {head_dim}, but got {static_k.size(2)}"
          k = static_k
      if static_v is None:
          v = v.contiguous().view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
      else:
          # TODO finish disentangling control flow so we don't do in-projections when statics are passed
          assert static_v.size(0) == bsz * num_heads, \
              f"expecting static_v.size(0) of {bsz * num_heads}, but got {static_v.size(0)}"
          assert static_v.size(2) == head_dim, \
              f"expecting static_v.size(2) of {head_dim}, but got {static_v.size(2)}"
          v = static_v
  
  Therefore, it appears to me that both ways of calculations should be equal and the one using the bigger matrices is just more efficient to compute. Is this correct?
  In this case, I ask myself why nn.MultiheadAttention requires that embed_dim is divisible by num_heads, since the split into the individual Attention-Heads is actually conducted after generating $Q,K,V$. Should then not $d_q,d_k,d_v$ be made to be divisible by num_heads? Since these dimensions do not have to be equal to the dimension of the inputs?
  Thanks for any advice in advance!
  
  **Tags:** transformer, pytorch, attention
  
  ---
  
  ## 1 answers
  
  ### Answer by username_1 on 2022-05-18. Score: 4 | Accepted answer
  
  The queries, keys and values are calculated then chunked so that each chunk depends on (is a linear combination of) all values of the input embedding.
  As for understanding an implementation, I didn't bother with pytorch but instead understood this code http://einops.rocks/pytorch-examples.html although there are differences as I understand that pytorch expects input in the form (L, N, C) where L=words, N=batch, C=embedding for performance reasons.
  Transformer's attention needs more attention
  class MultiHeadAttentionNew(nn.Module):
      def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
          super().__init__()
          self.n_head = n_head
          
          self.w_qs = nn.Linear(d_model, n_head * d_k)
          self.w_ks = nn.Linear(d_model, n_head * d_k)
          self.w_vs = nn.Linear(d_model, n_head * d_v)
          
          nn.init.normal_(self.w_qs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
          nn.init.normal_(self.w_ks.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
          nn.init.normal_(self.w_vs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_v)))
          
          self.fc = nn.Linear(n_head * d_v, d_model)
          nn.init.xavier_normal_(self.fc.weight)
          self.dropout = nn.Dropout(p=dropout)
          self.layer_norm = nn.LayerNorm(d_model)
  
      def forward(self, q, k, v, mask=None):
          residual = q
          q = rearrange(self.w_qs(q), 'b l (head k) -> head b l k', head=self.n_head)
          k = rearrange(self.w_ks(k), 'b t (head k) -> head b t k', head=self.n_head)
          v = rearrange(self.w_vs(v), 'b t (head v) -> head b t v', head=self.n_head)
          attn = torch.einsum('hblk,hbtk->hblt', [q, k]) / np.sqrt(q.shape[-1])
          if mask is not None:
              attn = attn.masked_fill(mask[None], -np.inf)
          attn = torch.softmax(attn, dim=3)
          output = torch.einsum('hblt,hbtv->hblv', [attn, v])
          output = rearrange(output, 'head b l v -> b l (head v)')
          output = self.dropout(self.fc(output))
          output = self.layer_norm(output + residual)
  
  Understand MultiHeadAttentionNew first then understand MultiHeadAttentionOld, and then maybe look at the pytorch code.
  
  > **Comments:**
  > **username_0**: Thank you for the link! It is really helpful. While I do not understand the details of some of the einops functions, it is clear enough. Since also bigger matrices are used in MultiHeadAttentionNew it really seems that you can just use bigger matrices and it gets you the same result.  I'll try to write it down mathematically.  
  > **username_1**: @Aushilfsgott The design of multi-head attention is for performance rather than having multiple single head attentions with the same input. Also I should have mentioned in my answer its worth really understanding einops and einsum, rather using the older pytorch functions  
  > **username_0**: Thanks again Tom. I got the basic principle of einsum so it is fine! The link is definitely very good for understanding the function.  
  
  ---
  
  