Linear Projections for Q, K, V per Head

Relying on a single self-attention mechanism can constrain the model's ability to capture the multifaceted relationships present within a sequence. A single attention calculation might average diverse relational signals, potentially losing specific, valuable information. Multi-Head Attention addresses this by performing attention computations multiple times in parallel, allowing the model to jointly attend to information from different perspectives or 'representation subspaces'.

The first step in enabling this parallel processing is to generate distinct Query ($Q$), Key ($K$), and Value ($V$) vectors for each attention "head". Instead of feeding the raw input embeddings directly into each head's attention calculation, we apply separate learned linear transformations to the input sequence for each head.

Let the input sequence representation be a matrix $X \in \mathbb{R}^{n \times d_{model}}$, where $n$ is the sequence length and $d_{model}$ is the dimension of the input embeddings (and the model's hidden state size). For a multi-head attention mechanism with $h$ heads, we introduce $h$ sets of learnable weight matrices:

• Query weight matrices: $W^Q_i \in \mathbb{R}^{d_{model} \times d_k}$ for $i=1, ..., h$
• Weight matrices: $W^K_i \in \mathbb{R}^{d_{model} \times d_k}$ for $i=1, ..., h$
• Value weight matrices: $W^V_i \in \mathbb{R}^{d_{model} \times d_v}$ for $i=1, ..., h$

Here, $d_k$ represents the dimension of the Keys and Queries for each head, and $d_v$ represents the dimension of the Values for each head. A common and effective choice, as used in the original Transformer paper ("Attention Is All You Need"), is to set $d_k = d_v = d_{model} / h$. This ensures that the total computational cost across all heads remains roughly similar to a single attention mechanism operating on the full $d_{model}$ dimension, while distributing the representational capacity.

For each head $i$, the corresponding Query, Key, and Value matrices ($Q_i$, $K_i$, $V_i$) are computed by projecting the input $X$ using these weight matrices:

$$Q_i = X W^Q_i$$ $$K_i = X W^K_i$$ $$V_i = X W^V_i$$

The resulting matrices $Q_i, K_i \in \mathbb{R}^{n \times d_k}$ and $V_i \in \mathbb{R}^{n \times d_v}$ now represent the input sequence transformed into the specific subspace relevant for head $i$.

Flow showing input embeddings $X$ being projected independently for each attention head $i$ using distinct weight matrices ($W^Q_i, W^K_i, W^V_i$) to produce head-specific Query ($Q_i$), Key ($K_i$), and Value ($V_i$) matrices. The projection typically reduces dimensionality from $d_{model}$ to $d_k$ or $d_v$.

Why Use Linear Projections?

These independent linear transformations are significant for several reasons:

1. Learning Different Subspaces: Each set of projection matrices $(W^Q_i, W^K_i, W^V_i)$ effectively defines a different linear subspace. Projecting the input embeddings into these different subspaces allows each attention head to focus on different types or aspects of relationships within the sequence. For instance, one head might learn to attend to short-range syntactic dependencies, while another might focus on longer-range semantic associations.
2. Increasing Representational Capacity: Using $h$ sets of projection matrices gives the model more parameters and thus greater capacity to learn complex patterns compared to a single set of projections. Even though each head operates on a reduced dimension ($d_k, d_v$), the combined effect after integrating information from all heads can be more powerful than a single head operating at the full $d_{model}$ dimension.
3. Enabling Parallel Specialization: Since each head $i$ receives its own projected versions $Q_i, K_i, V_i$, it can learn specialized attention patterns relatively independently of the other heads during training. This encourages diversity in the learned representations.

In practice, these linear projections are implemented using standard feed-forward layers (often linear layers or dense layers in frameworks like PyTorch or TensorFlow) without bias terms. The weight matrices $W^Q_i, W^K_i, W^V_i$ are initialized randomly and learned jointly with the rest of the network parameters during the training process via backpropagation.

Once these head-specific $Q_i, K_i, V_i$ matrices are computed, they become the inputs for the parallel scaled dot-product attention computations, which we will examine next. This projection step is fundamental to how Multi-Head Attention achieves its ability to analyze input sequences from multiple perspectives simultaneously.