Where an Attention Head Spends Its Bytes

An attention head has three matrices. Two of them, the query and key projections, decide where a token looks. The third, the value projection, decides what it brings back. On a device the three compete for the same byte budget. We shrank each path to the same file sizes and measured the loss. Bytes taken from the value path cost the most, and Needle 3's head shape follows from that.

The head

For a token $x$ , an attention head computes a query, and for every context token a key and a value:

$q = W_q x, \qquad k_i = W_k x_i, \qquad v_i = W_v x_i.$

The query and keys have width $d_{qk}$ . The values have width $d_v$ . The scores compare the query with each key, the softmax turns the scores into weights, and the weights average the values:

$w_i = \mathrm{softmax}_i\!\left(\frac{q \cdot k_i}{\sqrt{d_{qk}}}\right), \qquad \text{out} = W_o \sum_i w_i\, v_i.$

$d_{qk}$ only ever appears inside a dot product that becomes one number per position. $d_v$ is the width of what is carried into the residual stream, the only thing the block writes. Needle uses grouped-query attention with 16 query heads and 2 key-value heads, so both widths cost bytes in $W_q, W_k$ and in $W_v, W_o$ , but they buy different things: one buys routing precision, the other buys content.

One attention head, three budgets. Queries and keys only produce the attention scores; values are the content those scores select. Cutting the same number of bytes from each part costs very different amounts of loss.

Routing is cheap, content is not

Two results from the attention-only study say what to expect. The routing matrices $Q$ and $K$ settle early in training at low stable rank: an attention pattern needs few directions to pick out the positions that matter. The matrices that write content into the residual stream keep gaining rank for as long as training continues. The output of a head is a weighted average of value vectors, so it lives inside the shape those vectors span; a narrower $d_v$ is a smaller shape, and every token's update has to fit inside it. A narrower $d_{qk}$ only makes the weights slightly less sharp.

So the prediction is that bytes spent on $d_v$ do more than bytes spent on $d_{qk}$ .

The experiment

The base model is the 20-layer, $d = 1024$ Needle stack with 16 heads and 2 key-value heads, trained for 5,000 steps on the SYNTH mixture, without the engram so that the attention heads are the only thing that changes. The full head has $d_{qk} = d_v = 64$ . Three series shrink it:

cut queries and keys, values fixed at 64: $d_{qk} = 48, 32, 16$ ;
cut values, queries and keys fixed at 64: $d_v = 54, 44, 34$ ;
cut both by the same amount: $58/58$ and $52/52$ .

The widths were chosen so that each series lands at the same three file sizes, about 65, 60 and 54 MB at 16 bits, and every run used the same optimiser, schedule and seed. The change column is against the full head.

The same file size, bought three ways. At every size, taking the bytes out of the query-key path costs the least loss. Cutting both is never the best of the three, and at 60 MB and below the value path is the most expensive place to cut.

head, $d_{qk}$ / $d_v$	file size	removed	validation CE	loss change
64 / 64, the full head	71.0 MB		2.096
48 / 64, queries and keys cut	65.5 MB	7.8%	2.101	+0.3%
32 / 64	60.0 MB	15.5%	2.117	+1.0%
16 / 64	54.5 MB	23.3%	2.174	+3.7%
58 / 58, both cut	65.5 MB	7.8%	2.120	+1.1%
52 / 52	60.0 MB	15.5%	2.143	+2.2%
64 / 54, values cut	65.3 MB	8.1%	2.114	+0.8%
64 / 44	59.5 MB	16.2%	2.156	+2.9%
64 / 34	53.8 MB	24.2%	2.185	+4.3%

At every size the cheapest bytes to give up are in the routing. Taking a quarter of the query-key width away removes 7.8% of the file for 0.3% more loss; taking a sixth of the value width away removes 8.1% for 0.8%, nearly three times as much. Cutting the routing all the way to 16 dimensions still costs less than cutting the values to 34, and the same attention patterns are being computed with a quarter of the numbers. Cutting both together is never the best of the three.

Growing the values

The other direction is the useful one. Keep $d_{qk} = 48$ and widen only the values, in an 18-layer model trained for 3,000 steps:

Growing the value path alone. With queries and keys fixed at 48, doubling the value width from 64 to 128 takes the loss from 2.022 to 1.891, a larger gain than most whole-architecture changes we have measured at this scale.

$d_v$	file size	added	validation CE	loss change
64	185 MB		2.022
96	227 MB	+23%	1.982	−2.0%
128	270 MB	+46%	1.891	−6.5%

Doubling the value width takes 6.5% off the loss for 46% more file. For comparison, in the same series, moving the two engram memory sites to different layers changed the loss by under 1%, and making the engram tables three and a half times larger, 45% more file, took off 2.2%. Value width is the strongest single knob we measured at this scale.

What Needle 3 does

Needle 3 ships with $d_{qk} = 48$ and $d_v = 64$ , the first cut in the table: 7.8% of the file for 0.3% of the loss, the cheapest cut we measured. The engine gets a second benefit from it that the training run did not see. A 48-wide key is three 16-byte vector loads, so scoring a key against a query is three dot-product instructions and the whole key cache is 25% smaller than it would be at 64, which on a phone is cache that stays resident.

Width 64 for the values was the budget at the time. The table above says that if there is one number in the shape to raise next, it is that one.

References

Vaswani et al., Attention Is All You Need, 2017.
Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, 2023.
Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, 2019.
Bhojanapalli et al., Low-Rank Bottleneck in Multi-head Attention Models, 2020.
DeepSeek-AI, DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model, 2024.
Ndubuaku et al., A Controlled Study of Attention-Only Transformers, 2026.