Concatenation and Final Linear Projection

After computing the scaled dot-product attention in parallel for each of the $h$ heads, we are left with $h$ distinct output vectors, often denoted as $head_1, head_2, ..., head_h$. Each $head_i$ is the result of applying the attention mechanism using projected versions of the original queries, keys, and values.

$head_i = \text{Attention}(XW_i^Q, XW_i^K, XW_i^V)$

where $X$ represents the input sequence embeddings (or the output from the previous layer), and $W_i^Q$, $W_i^K$, $W_i^V$ are the parameter matrices for the $i$-th head. Each $head_i$ has a dimension of $d_v$, where $d_v = d_{model} / h$. This division ensures that the computational cost across all heads remains comparable to a single head attention with the full $d_{model}$ dimension.

However, the subsequent layers in the Transformer, such as the position-wise feed-forward network and the residual connections, expect a single input tensor with the model's hidden dimension, $d_{model}$. Therefore, we need a mechanism to consolidate the outputs from these parallel heads back into a single representation.

Combining Head Outputs: Concatenation

The first step in integrating the information captured by the different heads is straightforward: concatenation. The output vectors from all $h$ heads are concatenated along their feature dimension.

If each $head_i$ is a matrix of shape (sequence length, $d_v$), the concatenation operation stacks these matrices side-by-side, resulting in a new matrix of shape (sequence length, $h \times d_v$). Since we defined $d_v = d_{model} / h$, the resulting concatenated matrix has the dimensions (sequence length, $d_{model}$).

$\text{Concat}(head_1, head_2, ..., head_h) \in \mathbb{R}^{\text{seq\_len} \times (h \cdot d_v)} = \mathbb{R}^{\text{seq\_len} \times d_{model}}$

This operation effectively aggregates the insights learned from the different representational subspaces that each head focused on.

Flow showing outputs from individual attention heads being concatenated and then projected through a final linear layer.

Final Linear Projection

While concatenation aggregates the head outputs, the resulting tensor simply places the specialized representations next to each other. To allow these representations to interact and be combined effectively, and to ensure the output dimension matches the required $d_{model}$ for subsequent layers (especially for residual connections), a final linear projection is applied.

This projection is implemented using another learned weight matrix, $W^O$, with dimensions $(h \times d_v, d_{model})$, which simplifies to $(d_{model}, d_{model})$. The concatenated output is multiplied by this weight matrix:

$\text{MultiHeadOutput} = \text{Concat}(head_1, ..., head_h) W^O$

This final linear transformation plays an important role:

1. Mixing Information: It allows information learned across different heads (representing different subspaces) to be combined and integrated. The linear layer acts as a learned combination function.
2. Dimensionality Matching: It ensures the output tensor has the exact $d_{model}$ dimension required by the rest of the Transformer block architecture, enabling operations like residual connections (Add & Norm).

In essence, the multi-head mechanism first splits the model's representation capacity into multiple subspaces ($h$ heads, each with dimension $d_v$), allows each head to specialize in attending to different aspects of the input within its subspace, and then uses concatenation followed by a linear projection ($W^O$) to merge these specialized views back into a single, richer representation of dimension $d_{model}$.

This structure allows the model to capture a variety of relational patterns (e.g., short-range dependencies, long-range dependencies, syntactic relationships) simultaneously, which would be more difficult for a single attention mechanism. The parameters of the projection matrices ($W_i^Q, W_i^K, W_i^V$ for each head) and the final output projection matrix ($W^O$) are all learned end-to-end during model training.