Parallel Attention Computations

Parallel attention computation involves projecting the original queries ($Q$), keys ($K$), and values ($V$) into $h$ different subspaces using distinct learned linear transformations $W^Q_i, W^K_i, W^V_i$ for each head $i$. For each head, attention calculation then occurs independently and simultaneously. This parallel processing is a defining characteristic of multi-head attention and a significant contributor to its effectiveness and computational profile.

For each head $i$, where $i$ ranges from 1 to $h$, we compute the scaled dot-product attention exactly as defined previously, but using the projected matrices $Q_i$, $K_i$, and $V_i$ specific to that head:

$$\text{head}_i = \text{Attention}(Q_i, K_i, V_i) = \text{softmax}\left(\frac{Q_i K_i^T}{\sqrt{d_{k_i}}}\right) V_i$$

Here:

• $Q_i = Q W^Q_i$ represents the queries projected for head $i$.
• $K_i = K W^K_i$ represents the keys projected for head $i$.
• $V_i = V W^V_i$ represents the values projected for head $i$.
• $d_{k_i}$ is the dimension of the keys (and queries) within head $i$. This scaling factor, analogous to the one used in single-head scaled dot-product attention, stabilizes gradients during training.

It's important to manage the dimensions correctly. If the input embedding dimension is $d_{\text{model}}$ and we use $h$ heads, the projections are typically designed such that the dimension for each head's keys, queries ($d_{k_i}$), and values ($d_{v_i}$) are equal: $d_{k_i} = d_{v_i} = d_{\text{model}} / h$. This division ensures that the total computational cost is similar to that of a single-head attention mechanism with the full $d_{\text{model}}$ dimension, while distributing the representational capacity across multiple heads. Furthermore, it guarantees that when the outputs of all heads are later concatenated, the resulting dimension matches the expected input dimension $d_{\text{model}}$ for subsequent layers, maintaining consistency throughout the model architecture.

Assuming an input sequence of length $N$ (number of tokens), the shapes of the matrices for head $i$ (omitting the batch dimension for simplicity) are typically:

• $Q_i$: $N \times d_{k_i}$
• $K_i$: $N \times d_{k_i}$ (since in self-attention, Keys come from the same input sequence as Queries)
• $V_i$: $N \times d_{v_i}$

Consequently, the output of the attention calculation for head $i$, denoted $\text{head}_i$, will have the shape $N \times d_{v_i}$. Since we usually set $d_{v_i} = d_{k_i} = d_{\text{model}} / h$, the output shape is $N \times (d_{\text{model}} / h)$.

Computationally, this structure is highly amenable to parallelization. Modern deep learning frameworks and hardware like GPUs excel at performing large matrix multiplications. Instead of iterating through each head sequentially, the computations for all $h$ heads can often be executed in parallel. This is typically achieved by reshaping the projected Q, K, and V tensors to include a distinct "head" dimension before the attention calculation. For instance, a tensor representing queries for a batch might be reshaped from (batch_size, seq_len, d_model) to (batch_size, num_heads, seq_len, d_k_i). The batched matrix multiplications (matmul) required for the $\frac{Q_i K_i^T}{\sqrt{d_{k_i}}}$ term, followed by the softmax and the final matmul with $V_i$, can then operate efficiently across both the batch and head dimensions simultaneously.

Independent scaled dot-product attention computations are performed for each head using its specific projected Q, K, V matrices ($Q_i, K_i, V_i$). The outputs ($\text{head}_1, ..., \text{head}_h$), each with dimension $N \times d_{v_i}$, are generated in parallel before being passed to the next stage. Note $d_{k_i}$ and $d_{v_i}$ represent $d_{\text{model}}/h$.

The primary advantage of this parallel structure extends past computational efficiency. It permits each attention head to potentially specialize and learn different types of relationships or attend to information from distinct representation subspaces simultaneously. For instance, one head might learn to focus on local syntactic dependencies (like adjective-noun agreement), while another captures longer-range semantic connections (like coreference resolution across sentences), and yet another might focus on positional relationships. A single attention mechanism would be forced to average these potentially disparate signals, which could dilute the information. Multi-head attention provides multiple independent "channels" for information flow, allowing the model to aggregate diverse relational information and ultimately build richer, more context-aware representations.

The outputs of these parallel computations, $\text{head}_1, \text{head}_2, ..., \text{head}_h$, capture different aspects of the input sequence's internal relationships. They are now ready to be combined in the next step: concatenation followed by a final linear projection.