QK Circuits

The raw score being tested

For each head, we compute the raw pre-softmax attention score from the [ANS] query to a source token key. This isolates the QK ranking mechanism before attention probabilities are normalized.

q_ans = (tok_embed[ANS] + pos_embed[10]) @ W_Q.T
k_source = (tok_embed[d] + pos_embed[pos]) @ W_K.T
score = q_ans @ k_source / sqrt(d_head)

Mean QK score by digit for all heads — Mean raw QK score from `[ANS]` to each digit, averaged over the five number positions.

Headline metrics

Head	Corr digit-score	Score range	Argmax digit	Position dependence
0	0.758	27.28	9	0.0203
1	-0.490	36.24	1	0.1284
2	0.885	30.12	9	0.0122
3	0.995	45.76	9	0.0084

Position dependence is the largest absolute deviation from the mean digit-score curve across the five number positions. The tiny values for Heads 2 and 3 mean their QK scores are mostly digit-based, not position-based.

Per-head QK score maps

Head 0 QK digit heatmap — Head 0 prefers high digits, but the curve is less clean than Head 3.

Head 1 QK digit heatmap — Head 1 does not rank digits by max value; its digit-score correlation is negative.

Head 2 QK digit heatmap — Head 2 has a secondary high-digit preference.

Head 3 QK digit heatmap — Head 3 gives an almost monotonic score ramp from digit 0 to digit 9.

QK metric bars — Head 3 has the strongest digit-score correlation and the largest score range.

Special token QK scores — Special-token scores matter for heads that attend away from number tokens, especially Head 1.

Interpretation

Head 3's QK circuit is a clean digit-ranking mechanism: higher source digits receive higher raw scores from [ANS], and the score curve barely changes across number positions.

This explains the attention result: once softmax is applied over the five number tokens and specials, the largest digit tends to dominate Head 3's attention distribution. Head 2 and Head 0 also prefer high digits, but less cleanly. Head 1 is not a max-selector.

Head metrics CSV Mean digit scores CSV Special scores CSV Metadata JSON