📘

Multi-Head Attention

MediumAttention Mechanisms 📚 LLM 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
Concatenate heads: (h, n, d/h) → (n, d)

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

Input:

Line 1: n d h
Next n lines: Q matrix
Next n lines: K matrix
Next n lines: V matrix

Output: Multi-head attention output (n, d), values rounded to 4 decimal places.

Example:

Input:

Output:

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

Reasoning:

Split $Q$ , $K$ , $V$ column-wise into $h = 2$ heads of $d/h = 2$ features (head 1 = columns 0–1, head 2 = columns 2–3): $Q_1 = K_1 = \begin{bmatrix}1&0\\0&1\end{bmatrix}$ , $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}$ .
Head 1 scores: $Q_1K_1^T/\sqrt{2} = \begin{bmatrix}0.7071&0\\0&0.7071\end{bmatrix}$ . Row-wise softmax gives $\frac{e^{0.7071}}{e^{0.7071}+e^{0}} = \frac{2.0281}{3.0281} = 0.6698$ , so the weights are $\begin{bmatrix}0.6698&0.3302\\0.3302&0.6698\end{bmatrix}$ .
Head 2: $Q_2K_2^T$ is also the identity, so it gets the same weights.
Weighted values: head 1 row 0 is $0.6698\cdot[1,2] + 0.3302\cdot[5,6] = [2.3210, 3.3210]$ and head 2 row 0 is $0.6698\cdot[3,4] + 0.3302\cdot[7,8] = [4.3210, 5.3210]$ . Row 1 gives $[3.6790, 4.6790]$ and $[5.6790, 6.6790]$ .
Concatenate the two heads for each token and round to 4 decimals: $\begin{bmatrix}2.3210&3.3210&4.3210&5.3210\\3.6790&4.6790&5.6790&6.6790\end{bmatrix}$ .

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

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

0/0

Run code to see test results.