Multi-Head Attention in Large Language Models: How Query, Key, and Value Projections Route Information Across Subspaces

The attention mechanism serves as the computational core of modern Transformer architectures, providing a mechanism for tokens to dynamically exchange information across arbitrary sequence positions without recurrent state transitions or fixed convolutional receptive fields. While single-head attention computes a single set of attention weights per token pair, modern large language models rely on Multi-Head Attention (MHA) to route information across multiple representation subspaces concurrently.

By projecting input hidden states into multiple distinct Query, Key, and Value subspaces, Multi-Head Attention prevents the model from averaging out distinct linguistic, syntactic, and semantic relationships into a single homogenized distribution.

1. Mathematical Formulation of Scaled Dot-Product Attention

The foundational building block of the Transformer attention layer is Scaled Dot-Product Attention, introduced by Vaswani et al. (2017). Given an input sequence matrix $X \in \mathbb{R}^{N \times d_{\text{model}}}$ representing $N$ tokens of hidden dimension $d_{\text{model}}$, the input is projected into three distinct representations via learnable linear weight matrices:

$$Q = X W_Q, \quad K = X W_K, \quad V = X W_V$$

Where:

• $W_Q \in \mathbb{R}^{d_{\text{model}} \times d_k}$ is the Query projection matrix.
• $W_K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ is the Key projection matrix.
• $W_V \in \mathbb{R}^{d_{\text{model}} \times d_v}$ is the Value projection matrix.
• $d_k$ is the dimensionality of queries and keys, and $d_v$ is the dimensionality of values (typically $d_k = d_v$).

The Scaled Dot-Product Equation

Scaled Dot-Product Attention computes the inner product between queries and keys, scales the resulting logit scores, applies a row-wise softmax normalization, and weights the value vectors:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$$

Input Tokens (X)
  ├──> [ * W_Q ] ──> Query (Q) ──┐
  ├──> [ * W_K ] ──> Key (K)   ──┴──> MatMul(Q, K^T) ──> Scale (1/√d_k) ──> [Mask] ──> Softmax ──┐
  └──> [ * W_V ] ───────────────────> Value (V) ───────────────────────────────────────────┴──> MatMul(A, V) ──> Head Output

Why the $\sqrt{d_k}$ Scaling Factor Is Necessary

The division by $\sqrt{d_k}$ is critical for training stability in high-dimensional vector spaces. To understand why, consider the components of a single query vector $q \in \mathbb{R}^{d_k}$ and a key vector $k \in \mathbb{R}^{d_k}$.

Assuming the individual vector elements $q_i$ and $k_i$ are independent random variables with zero mean ($\mathbb{E}[q_i] = \mathbb{E}[k_i] = 0$) and unit variance ($\text{Var}(q_i) = \text{Var}(k_i) = 1$):

1. The expected value of the inner product is zero:

$$\mathbb{E}[q \cdot k] = \mathbb{E}\left[\sum_{i=1}^{d_k} q_i k_i\right] = \sum_{i=1}^{d_k} \mathbb{E}[q_i]\mathbb{E}[k_i] = 0$$

1. Because the individual products $q_i k_i$ are independent, the variance of their sum equals the sum of their variances:

$$\text{Var}(q_i k_i) = \mathbb{E}[(q_i k_i)^2] - (\mathbb{E}[q_i k_i])^2 = \mathbb{E}[q_i^2]\mathbb{E}[k_i^2] - 0 = 1 \times 1 = 1$$ $$\text{Var}(q \cdot k) = \text{Var}\left(\sum_{i=1}^{d_k} q_i k_i\right) = \sum_{i=1}^{d_k} \text{Var}(q_i k_i) = d_k$$

1. The standard deviation of the raw dot product is therefore $\sigma = \sqrt{d_k}$.

As the head dimension $d_k$ grows (for instance, $d_k = 64$ or $d_k = 128$), the magnitude of the dot products scales proportionally with $\sqrt{d_k}$. When unscaled dot products enter the softmax function, large inputs push the exponential terms to extremes, driving the softmax outputs toward one-hot distributions.

In this saturated regime, the gradients of the softmax function with respect to its inputs approach zero ($\frac{\partial \text{softmax}(z)_i}{\partial z_j} \approx 0$). Dividing $Q K^T$ by $\sqrt{d_k}$ rescales the variance of the attention logits back to 1.0, preserving healthy gradient flow throughout backpropagation.

Causal Masking in Autoregressive Models

In decoder-only language models (such as GPT-4, LLaMA, and Claude), tokens must not attend to future sequence positions. This constraint is enforced by applying an additive upper-triangular causal mask $M \in \mathbb{R}^{N \times N}$ to the scaled logits prior to the softmax operation:

$$M_{ij} = \begin{cases} 0 & \text{if } i \ge j \\ -\infty & \text{if } i < j \end{cases}$$

$$\text{Attention}_{\text{causal}}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}} + M\right) V$$

When $M_{ij} = -\infty$, the corresponding exponential term in the softmax denominator evaluates to $e^{-\infty} = 0$, ensuring zero attention weight is assigned to subsequent tokens.

2. Multi-Head Attention Architecture

Single-head attention computes a single weighted combination of value vectors for each token. In natural language, however, a single token simultaneously participates in multiple relationships:

• Syntactic dependencies (for example, linking a verb to its subject and direct object).
• Coreference resolution (linking pronouns to preceding proper nouns).
• Positional proximity (attending to immediately preceding or subsequent tokens).
• Semantic association (linking related domain concepts across long paragraph spans).

A single attention distribution forces the model to average these disparate attention patterns into a single weighted vector, diluting specific relational signals. Multi-Head Attention overcomes this limitation by partitioning the feature space into $h$ parallel attention heads.

                  Input Hidden States X (N x d_model)
                                   │
      ┌──────────────┬─────────────┼─────────────┬──────────────┐
      ▼              ▼             ▼             ▼              ▼
   Head 1         Head 2        Head 3        Head 4         Head h
  Q1, K1, V1    Q2, K2, V2    Q3, K3, V3    Q4, K4, V4     Qh, Kh, Vh
      │              │             │             │              │
  Attention      Attention     Attention     Attention      Attention
      │              │             │             │              │
   Output 1       Output 2      Output 3      Output 4       Output h
      └──────────────┴─────────────┼─────────────┴──────────────┘
                                   │
                           Concat(head_1, ..., head_h)
                                   │
                           Linear Projection (W_O)
                                   │
                      Output States (N x d_model)

Multi-Head Equations

Instead of performing a single attention function with $d_{\text{model}}$-dimensional queries, keys, and values, Multi-Head Attention linearly projects queries, keys, and values $h$ times with distinct learnable parameter matrices:

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_h) W_O$$

$$\text{where } \text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V)$$

The parameter dimensions are structured as:

• $W_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}$
• $W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$
• $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$
• $W_O \in \mathbb{R}^{h d_v \times d_{\text{model}}}$

In standard implementations, the head dimension is set to $d_k = d_v = d_{\text{model}} / h$. For example, in a model with $d_{\text{model}} = 4096$ and $h = 32$ heads, each individual head operates in a $d_k = 128$ dimensional subspace.

Because each head operates at reduced dimensionality ($d_{\text{model}} / h$), the total computational cost of computing $h$ parallel heads is identical to that of a single full-dimensional attention head with $d_k = d_{\text{model}}$.

Fused Matrix Implementation

In production inference and training engines (such as PyTorch, vLLM, and TensorRT-LLM), the $h$ individual head projections are not executed as separate sequential kernel launches. Instead, they are fused into three large matrix multiplications:

# Batch dimension B, sequence length N, model dimension D, heads H, head_dim d_k
# Fused linear projections:
q = self.w_q(x)  # Shape: [B, N, H * d_k]
k = self.w_k(x)  # Shape: [B, N, H * d_k]
v = self.w_v(x)  # Shape: [B, N, H * d_v]

# Reshape into parallel heads:
q = q.view(B, N, H, d_k).transpose(1, 2)  # Shape: [B, H, N, d_k]
k = k.view(B, N, H, d_k).transpose(1, 2)  # Shape: [B, H, N, d_k]
v = v.view(B, N, H, d_v).transpose(1, 2)  # Shape: [B, H, N, d_v]

# Batched scaled dot-product attention across all heads simultaneously:
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
if causal_mask is not None:
    scores = scores + causal_mask
attn_weights = torch.softmax(scores, dim=-1)
context = torch.matmul(attn_weights, v)  # Shape: [B, H, N, d_v]

# Recombine heads and apply output projection:
context = context.transpose(1, 2).contiguous().view(B, N, H * d_v)
output = self.w_o(context)  # Shape: [B, N, D]

3. Mechanistic Head Specialization

Research in mechanistic interpretability and natural language processing has uncovered how individual attention heads specialize during pre-training.

Positional, Syntactic, and Lexical Heads

In an empirical analysis of multi-head attention mechanisms, Voita et al. (2019) and Clark et al. (2019) demonstrated that attention heads naturally partition into specific functional categories:

1. Positional Heads: Heads that consistently attend to fixed relative offsets (such as token $i-1$, token $i+1$, or the initial sequence delimiter token <s> / [CLS]).
2. Syntactic Heads: Heads that track grammatical relationships, such as pointing from direct objects to their governing verbs, or connecting noun modifiers to their head nouns.
3. Coreference and Anaphora Heads: Heads that link pronouns (e.g., "it", "they", "she") to candidate antecedent noun phrases across multi-sentence contexts.

Syntactic Head (Verb -> Object):
  "The engineer [built] [the distributed cluster] yesterday"
                  ▲              │
                  └──────────────┘ (High attention weight)

Positional Head (Previous Token):
  "Large" <── "Language" <── "Models" <── "Scale"

Induction Heads and In-Context Learning

In A Mathematical Framework for Transformer Circuits (Anthropic, 2021), Elhage et al. showed that two-layer attention circuits form "induction heads."

An induction head detects sequence repetitions of the form [A][B] ... [A] and attends back to [B], copying it to complete the pattern. This two-step circuit consists of:

• Layer 1 (Previous-Token Head): Writes information about token $i-1$ into the residual stream at position $i$.
• Layer 2 (Induction Head): Query vector at position $j$ (matching [A]) matches Key vector at position $i$ (which carries information that [A] was followed by [B]), routing Value [B] into the next-token prediction.

Induction heads represent the primary mechanistic driver behind in-context few-shot learning in large language models.

Head Redundancy and Pruning

Despite the necessity of multiple heads during training, research indicates substantial redundancy in trained models. Voita et al. (2019) demonstrated that applying $L_0$ regularization allowed pruning up to 80% of attention heads in machine translation models without significant BLEU score degradation.

However, multi-head diversity remains essential during optimization: having numerous parallel heads allows gradient descent to explore multiple functional sub-circuits simultaneously before settling into specialized configurations.

4. Computational and Memory Bottlenecks

While Multi-Head Attention enables expressive representational capacity, it introduces substantial computational and memory demands.

Computational Complexity

For a sequence of length $N$ and model dimension $d_{\text{model}}$:

| Operation | Flops Count per Layer | Computational Regime | | :--- | :--- | :--- | | Q, K, V Projections ($3 \times X W$) | $6 N d_{\text{model}}^2$ | Compute-bound (BLAS Level 3 GEMM) | | Attention Scores ($Q K^T$) | $2 N^2 d_{\text{model}}$ | Quadratic in sequence length $N$ | | Softmax Normalization | $O(h N^2)$ | Memory bandwidth-bound | | Context Aggregation ($A V$) | $2 N^2 d_{\text{model}}$ | Quadratic in sequence length $N$ | | Output Projection ($W_O$) | $2 N d_{\text{model}}^2$ | Compute-bound (BLAS Level 3 GEMM) |

In the prefill phase (processing long prompt sequences), the $O(N^2 d_{\text{model}})$ quadratic scaling of $Q K^T$ and $A V$ dominates execution time.

The KV Cache Memory Bottleneck

During autoregressive token-by-token generation (decoding phase), past Key and Value activations are cached in GPU High Bandwidth Memory (HBM) to avoid recomputing past representations at each generation step.

For a standard Multi-Head Attention model with $L$ layers, $h$ heads, head dimension $d_k$, and batch size $B$, the KV cache memory footprint for a context of $N$ tokens is:

$$\text{KV Cache Size} = 2 \times B \times N \times L \times h \times d_k \times \text{BytesPerElement}$$

For a 70-billion parameter model (such as LLaMA-2 70B with $L=80$, $h=64$, $d_k=128$, FP16 precision):

• Key cache per token: $80 \times 64 \times 128 \times 2 \text{ bytes} = 1.31 \text{ MB}$
• Value cache per token: $80 \times 64 \times 128 \times 2 \text{ bytes} = 1.31 \text{ MB}$
• Total KV cache per token: $2.62 \text{ MB}$ per sequence

At a sequence length of $N = 4096$ tokens and a batch size of $B = 32$, the KV cache alone requires:

$$\text{Memory} = 32 \times 4096 \times 2.62 \text{ MB} \approx 343.4 \text{ GB of GPU VRAM}$$

In autoregressive decoding, loading these large Key and Value matrices from HBM into GPU SRAM at every single token step causes the memory bandwidth of the GPU to saturate long before its compute ALUs reach capacity.

5. Architectural Evolutions: MHA, MQA, GQA, and MLA

To mitigate the memory bandwidth bottleneck of standard Multi-Head Attention in production serving, modern LLMs employ modified attention projection variants.

Multi-Head Attention (MHA)      Grouped-Query Attention (GQA)      Multi-Query Attention (MQA)
  Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8          Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8            Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8
  │  │  │  │  │  │  │  │           └──┬──┘ └──┬──┘ └──┬──┘ └──┬──┘            └───┬───┴───┬───┘
  K1 K2 K3 K4 K5 K6 K7 K8             K1      K2      K3      K4                    K1
  V1 V2 V3 V4 V5 V6 V7 V8             V1      V2      V3      V4                    V1
  (8 Query, 8 Key/Value)           (8 Query, 4 Key/Value Groups)       (8 Query, 1 Key/Value)

1. Multi-Query Attention (MQA)

Proposed by Shazeer (2019), Multi-Query Attention retains $h$ distinct Query heads but uses a single shared Key head and a single shared Value head ($h_{kv} = 1$).

• KV Cache Reduction: Reduces KV cache size and memory traffic by a factor of $h$.
• Trade-off: Can lead to minor quality degradation on complex reasoning and long-context retrieval tasks due to reduced Key/Value capacity.

2. Grouped-Query Attention (GQA)

Introduced by Ainslie et al. (2023), Grouped-Query Attention groups Query heads into $g$ partitions, with each partition sharing a single Key and Value head ($1 < g < h$).

• For example, in LLaMA-3 70B ($h=64$ query heads, $g=8$ KV groups), each group of 8 Query heads shares one Key head and one Value head.
• KV Cache Reduction: Slashes KV cache memory footprint by an $8\times$ factor while matching the downstream task performance of full Multi-Head Attention. GQA has become the de facto standard for open-weight foundation models.

3. Multi-Head Latent Attention (MLA)

Pioneered in DeepSeek-V2 (2024) and DeepSeek-V3, Multi-Head Latent Attention compresses Key and Value projections into a shared low-rank latent vector $c^{KV} \in \mathbb{R}^{d_c}$ (where $d_c \ll h \cdot d_k$) before caching:

• Only the compressed latent vector is stored in the KV cache during generation.
• During attention computation, the latent representation is uncompressed on-the-fly or mathematically folded into the Query projection matrix via matrix associativity, achieving extreme KV cache compression without sacrificing multi-head expressivity.

6. Summary Comparison

| Attention Variant | Query Heads ($h_q$) | Key/Value Heads ($h_{kv}$) | KV Cache Memory per Token | Typical Adoption | | :--- | :--- | :--- | :--- | :--- | | Multi-Head Attention (MHA) | $h$ | $h$ | $2 \times L \times h \times d_k$ | Transformer (2017), GPT-3, LLaMA-1 | | Multi-Query Attention (MQA) | $h$ | $1$ | $2 \times L \times 1 \times d_k$ | PaLM, StarCoder, Falcon | | Grouped-Query Attention (GQA) | $h$ | $g$ ($1 < g < h$) | $2 \times L \times g \times d_k$ | LLaMA-2/3, Mistral, Gemma-2 | | Multi-Head Latent Attention (MLA) | $h$ | Low-rank compressed ($d_c$) | $L \times (d_c + d_{R})$ | DeepSeek-V2, DeepSeek-V3 |

Multi-Head Attention remains the foundational mechanism that enabled Transformers to replace recurrent neural networks. By decoupling sequence routing across multiple parallel geometric subspaces, MHA provides the mathematical foundation upon which modern scaling laws, in-context learning circuits, and efficient serving architectures are built.

Sources

• Attention Is All You Need (Vaswani et al., NeurIPS 2017)
• Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned (Voita et al., ACL 2019)
• What Does BERT Look At? An Analysis of Attention's Mechanisms (Clark et al., ACL 2019)
• A Mathematical Framework for Transformer Circuits (Elhage et al., Anthropic 2021)
• Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019)
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., EMNLP 2023)
• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (DeepSeek AI, 2024)