K HEADS IS ALL YOU NEED: TASK-DEPENDENT CIRCUIT GRANULARITY AND ZERO-TRAINING ROUTING IN LLMS
Authors/Creators
Description
I show that a single computational behavior – needle-in-a-haystack (NIAH) retrieval in Qwen2.5-1.5B – is supported by a causal circuit of just six attention
heads, sufficient to recover the behavior under isolation and necessary under knockout, validated against layer-matched random controls rather than raw ablation deltas. The same six-head class of circuit recovers NIAH-MultiKey retrieval with minor expansion (with partial accuracy retention at base K pending capacitymatched circuit expansion), and a five-head compression of the original circuit
produces a 100% causal steering flip on a counterfactual retrieval target. This raises an immediate question: is transformer computation generally this sparse? I test this directly by asking whether the same circuits and steering vectors generalize to other computational primitives – counting, sorting, variable tracking, and document-level QA – using the identical causal validation protocol (sufficiency, necessity, specificity, and layer-matched controls) throughout. The answer is no, in a structured way: linear steering vectors extracted from retrieval circuits fail to improve accuracy on algorithmic primitives at any tested scale, some tasks (variable tracking) resist sparsification below full dense attention entirely, and naive gradient attribution – a cheap alternative to causal search – actively destroys working sparse circuits rather than merely underperforming them. I map this boundary across four model families (Qwen2.5-0.5B/1.5B, Llama-3.2-1B, Gemma-2-2B) and a hierarchy of Python computational tasks, and use it to build a zero-training, gradient-free geometric router that detects – from a single early-layer cosine gate – whether a given input belongs to the sparse-compressible regime, validated on the official NVIDIA RULER 4K long-context benchmark with zero false activations on 100 out-of-distribution samples. The result is a precise, falsifiable answer to “how sparse is transformer computation”: sparse for retrieval, not for reasoning, and the boundary between the two is geometrically detectable at inference time without training a single parameter.
Files
k_heads_paper.pdf
Files
(271.7 kB)
| Name | Size | Download all |
|---|---|---|
|
md5:a566b62d7e2ec916ad830e361faf510b
|
271.7 kB | Preview Download |
Additional details
Software
- Repository URL
- https://github.com/skhavin/attentionheadgenome
- Programming language
- Python