Neural Mastery
← Practice
Multi-Head Attention: Split Headseasy

Multi-Head Attention: Split Heads

easy Inside an LLM 10–12 min

NumPy allowed · plain Python earns 10 bonus points

Attention step 1/4

Tokens attend to tokens

Q · Kᵀ / √d → softmax → weighted values, computed in the page.
Open lesson →
01tokens
02Q K V
03scores
04attention
modelslearnrelationships
models32%
learn38%
relationships30%

Q · Kᵀ / √d scores normalized with softmax

Implement the head-splitting operation in multi-head attention, dividing the embedding dimension into parallel attention sub-spaces.

Task

Implement split_heads(x, num_heads) that splits vector x (length d_model) into num_heads chunks of size d_model // num_heads.

Function Signature

split_heads(x: list[float], num_heads: int) -> list[list[float]]

Examples

Example 1: Split 8-dim into 2 heads
Input: x = [1, 2, 3, 4, 5, 6, 7, 8] num_heads = 2
Output: [[1, 2, 3, 4], [5, 6, 7, 8]]
Example 2: Split 4-dim into 4 heads
Input: x = [1, 2, 3, 4] num_heads = 4
Output: [[1], [2], [3], [4]]

Constraints

  • len(x) must be divisible by num_heads, else raise ValueError.
  • Return a list of num_heads sub-vectors.
Python3Saved ✓
def split_heads(x: list[float], num_heads: int) -> list[list[float]]:
"""Split an embedding vector into num_heads equal-sized head vectors.
Args:
x: Embedding vector of length d_model (list of floats).
num_heads: Number of attention heads. Must divide len(x) evenly.
Returns:
list[list[float]]: List of num_heads vectors, each of length d_model // num_heads.
Raises:
ValueError: If len(x) is not divisible by num_heads.
"""
# Your implementation here
pass

Case 1: Split 8-dim into 2 heads
Input: x = [1, 2, 3, 4, 5, 6, 7, 8] num_heads = 2
Expected: [[1, 2, 3, 4], [5, 6, 7, 8]]
Case 2: Split 4-dim into 4 heads
Input: x = [1, 2, 3, 4] num_heads = 4
Expected: [[1], [2], [3], [4]]