When exactly does the split into different heads in Multi-Head-Attention occur?

Question

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!

Answer 1

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.