Published August 24, 2026 | Version v1

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