Multi-Head Attention

mediumFREEAttention MechanismsLLM 1: Foundations

Implement multi-head attention by splitting Q, K, V into multiple heads.

Given Q, K, V of shape (n, d) and number of heads h:

Split each into h heads: reshape (n, d) → (n, h, d/h) → (h, n, d/h)
Apply scaled dot-product attention per head: $\text{softmax}\left(\frac{Q_i K_i^T}{\sqrt{d/h}}\right) V_i$ , with the softmax taken over each row (the scale uses the per-head width $d/h$ , not $d$ )
Concatenate heads: (h, n, d/h) → (n, d)

Assume d is divisible by h. No linear projections needed — just split and concat.

Input: multi_head_attention(Q, K, V, h) where Q, K, V are NumPy arrays of shape (n, d) and h is the number of heads.

Output: Multi-head attention output as an (n, d) array; the tests print each value to 4 decimal places.

Example

Input

Output

[2.3210 3.3210 4.3210 5.3210]
[3.6790 4.6790 5.6790 6.6790]

Reasoning

Inputs: $Q$ , $K$ , $V$ are $2 \times 4$ (two tokens, $d = 4$ features) and h $= 2$ , so each head owns $d/h = 2$ consecutive features: head 1 the first two columns, head 2 the last two. That gives $Q_1 = K_1 = I$ , $Q_2 = K_2 = \begin{bmatrix}0&1\\1&0\end{bmatrix}$ , $V_1 = \begin{bmatrix}1&2\\5&6\end{bmatrix}$ , $V_2 = \begin{bmatrix}3&4\\7&8\end{bmatrix}$ .
Scores: each head scales by $\sqrt{d/h} = \sqrt{2}$ , not $\sqrt{d}$ . Both $Q_1K_1^T$ and $Q_2K_2^T$ equal $I$ , so both heads score $\begin{bmatrix}0.7071&0\\0&0.7071\end{bmatrix}$ .
Row softmax: $\frac{e^{0.7071}}{e^{0.7071}+e^{0}} = \frac{2.0281}{3.0281} = 0.6698$ , each token puts $0.6698$ on itself and $0.3302$ on the other.
Weighted values: token 1 gets $0.6698\cdot[1,2] + 0.3302\cdot[5,6] = [2.3210, 3.3210]$ from head 1 and $0.6698\cdot[3,4] + 0.3302\cdot[7,8] = [4.3210, 5.3210]$ from head 2; token 2 gets $[3.6790, 4.6790]$ and $[5.6790, 6.6790]$ .
Output: each token's two head results sit side by side, rounded to 4 decimal places: [2.3210 3.3210 4.3210 5.3210] and [3.6790 4.6790 5.6790 6.6790].
Notice head 2's values are exactly head 1's plus 2: $V_2 = V_1 + 2$ and every row of weights sums to 1.

Constraints

d is divisible by h
1 <= h <= d, 1 <= n <= 10
No projection matrices — just split/concat
Round to 4 decimal places

Learn this in

Transformer ArchitectureLLM Study Plan CNNs & Sequence ModelsML Study Plan Machine TranslationNLP Study Plan Attention MechanismsNLP Study Plan

More from LLM 1: Foundations

Transformer Block Forward Passhard Feed-Forward Networkeasy Layer Normalizationeasy KV-Cache Attentionhard Causal Attention Maskmedium

Concept primer

Background Knowledge

The concept of attention mechanisms is crucial in understanding the problem of multi-head attention. Attention mechanisms allow a model to focus on specific parts of the input data that are relevant for a particular task. This is particularly useful in sequence-to-sequence models, such as machine translation, where the model needs to attend to different parts of the input sequence to generate the output sequence. The scaled dot-product attention is a specific type of attention mechanism that calculates the attention weights by taking the dot product of the query and key vectors and applying a scaling factor.

In the context of multi-head attention, the idea is to split the input vectors (Q, K, V) into multiple heads, allowing the model to jointly attend to information from different representation subspaces at different positions. This is achieved by splitting the input vectors into multiple heads, applying scaled dot-product attention to each head, and then concatenating the outputs. The number of heads (h) is a hyperparameter that needs to be set, and the dimensionality of the input vectors (d) should be divisible by the number of heads.

The mathematical formulation of scaled dot-product attention involves calculating the attention weights using the formula: $A = \text{softmax}(\frac{Q \cdot K^T}{\sqrt{d}})$ , where $Q$ , $K$ , and $V$ are the query, key, and value vectors, respectively, and $d$ is the dimensionality of the vectors. The output of the attention mechanism is then calculated as: $O = A \cdot V$ . In the context of multi-head attention, this process is applied to each head separately.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Split the input vectors (Q, K, V) into multiple heads
Apply scaled dot-product attention to each head
Concatenate the outputs from each head This approach requires a good understanding of tensor manipulation and attention mechanisms.

Step-by-Step Strategy

To implement the solution, follow these steps:

Read the input values: n, d, h, Q, K, V
Split each of Q, K, V into h heads by reshaping the tensors
Apply scaled dot-product attention to each head
Concatenate the outputs from each head
Round the output values to 4 decimal places

Common Pitfalls

Some common pitfalls to watch out for include:

Incorrect tensor reshaping and concatenation
Forgetting to apply the scaling factor in the scaled dot-product attention
Not handling the case where d is not divisible by h (although this is not a concern in this specific problem)

Time & Space Complexity

The time complexity of the solution is O(n^2 * d) due to the matrix multiplication involved in the scaled dot-product attention. The space complexity is O(n * d) for storing the input vectors and the output. Note that these complexities assume that the number of heads (h) is a constant factor.

solution.py

Python 3.13

Test results