**The formula is `head_dim = hidden_size / num_heads`**, with the requirement that `hidden_size` is exactly divisible by `num_heads`.
The reason is mechanical. Multi-head attention takes the per-token hidden vector of size `hidden_size` and reshapes it into a 2D block of shape `(num_heads, head_dim)` so each head can compute attention independently on its slice. The reshape requires the two output dimensions to multiply back to the input size:
```
hidden_size = num_heads * head_dim
```
Rearranging gives `head_dim = hidden_size / num_heads`. The integer divide is implicit because PyTorch's `view` and `reshape` operations require the dimensions to match exactly, no rounding allowed.
Worked examples:
- GPT-2 small: hidden_size = 768, num_heads = 12, so head_dim = 64.
- Llama-2 7B: hidden_size = 4096, num_heads = 32, so head_dim = 128.
- Llama-2 70B: hidden_size = 8192, num_heads = 64, so head_dim = 128.
The other answer choices are wrong:
- Multiplying instead of dividing would give an absurd value (768 * 12 = 9216).
- Taking sqrt would produce a non-integer in nearly every case.
- Treating head_dim as free contradicts how the reshape inside attention actually works.
**head_dim = hidden_size / num_heads**, integer divide, with `hidden_size % num_heads == 0` required.
The relationship is fixed by the multi-head attention reshape. The Q, K, V projections produce tensors of shape `(batch, seq_len, hidden_size)`. To run attention per head, the implementation views these as `(batch, seq_len, num_heads, head_dim)`, then transposes to `(batch, num_heads, seq_len, head_dim)` for the matmul. The view operation requires the trailing dimensions to factor back to hidden_size exactly.
The identity:
$$d_\text{head} = \frac{d_\text{model}}{n_\text{heads}}$$
Applied to standard configurations:
- BERT base / GPT-2 small: 768 / 12 = 64.
- Llama-2 7B / Mistral 7B: 4096 / 32 = 128.
- Llama-2 70B: 8192 / 64 = 128.
The divisibility constraint is enforced at model construction. PyTorch's `MultiheadAttention` and Hugging Face's attention implementations both raise an error if `hidden_size % num_heads != 0`.
**Why this matters for design choices.** You cannot pick hidden_size, num_heads, and head_dim independently. Most modern LLMs fix head_dim at 128 (kernel-friendly) and pick hidden_size + num_heads to satisfy the constraint. A few research models (some MLA variants, some MoE configs) add an explicit projection layer between the linear and the heads, which lets head_dim be set independently. The default and the interview-correct answer is the simple integer division.
Quick reference
TL;DR
head_dim = hidden_size / num_heads, integer divide. hidden_size must be divisible by num_heads or the per-head reshape errors out.
Imagine a 12-pack of soda you want to split evenly into smaller packs. If you want 4 packs, each gets 3 cans. If you want 6 packs, each gets 2. The arithmetic is just: 12 / number of packs = cans per pack. Multi-head attention does the same with a flat hidden vector: split the hidden_size into num_heads chunks, each chunk has head_dim entries, and the math is head_dim = hidden_size / num_heads. If your split does not divide evenly, the soda example breaks the same way the model does: leftover cans, errored reshape.
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example.
Architectures using MLA or other learned-latent KV compression schemes MLA decouples per-head dimensionality from hidden_size / num_heads via explicit projection layers; the standard formula does not apply.
Designs that add an explicit projection layer between the QKV linear and the head split Some research architectures project hidden_size to a different intermediate dimension before splitting into heads, breaking the standard identity by design.
Cross-attention setups where Q and K, V come from different feature widths Encoder-decoder cross-attention can have asymmetric Q and KV widths; the simple hidden_size / num_heads identity may not apply to both sides of the projection.
Discussions of GQA's num_kv_heads as if it changed the head_dim formula GQA only reduces num_kv_heads, not num_attention_heads; head_dim is computed from the query head count, so GQA does not change the formula.
Quick scan · 2–4 min
Everything important, quickly.
60s: state the formula head_dim = hidden_size / num_heads, note the divisibility requirement, explain that the multi-head reshape enforces it, give worked examples (GPT-2 small: 768/12 = 64, Llama-2 7B: 4096/32 = 128), mention parameter-count preservation as the original Transformer's motivation.
Real-world examples
Real products, models, and research that use this idea.
4Llama-2 70B and Llama 4 Maverick continue with head_dim = 128 across configurations.
5DeepSeek V4 MLA decouples per-head dim from the standard ratio via explicit projection fields in its config.
Interviewer follow-ups
What an interviewer would ask next. Try answering before peeking at the approach.
QWhat happens at runtime if you try to build a model with hidden_size = 800 and num_heads = 12?
Sign in or go Pro to see the approach for this follow-up.
QWhy does the original Transformer pick parameter count to preserve under num_heads variation?
Sign in or go Pro to see the approach for this follow-up.
QHow does MLA decouple head_dim from the standard relationship, and why?
Sign in or go Pro to see the approach for this follow-up.
Don't say this
Red flags and common mistakes that signal junior thinking. Click to expand.
Most common mistake
Treating head_dim as a free hyperparameter independent of hidden_size and num_heads. It is locked: head_dim = hidden_size / num_heads, with hidden_size % num_heads == 0 required.