K Heads Is All You Need

Task-Dependent Circuit Granularity and Zero-Training Routing in LLMs

Overview · Key Results · Method: JCCS · Code · Claims · Citation

---

The computational budget required to execute a task is a function of the task's cognitive complexity - and it is discoverable through joint causal search.

6 heads out of 336 (1.78%) form a causally necessary and sufficient sparse circuit for long-context retrieval. For compositional code generation, it scales to 168 heads - but the entire first half of the network can be zeroed without any performance loss.

---

Overview

This is the fourth and final paper in the ISA Research Program. It asks the most precise version of the ISA question:

What is the minimum causally necessary and sufficient set of components to execute a specific task?

Prior work established that a fixed head-level ISA does not generalize, that the residual stream has conserved geometry, but that this geometry is not causally sufficient to steer computation. This paper delivers the causal answer using JCCS (Joint Causal Circuit Search).

Key findings:

• 6/336 heads form a causally sufficient + necessary sparse circuit for NIAH (long-context retrieval)
• The computational budget scales with cognitive complexity, replicating across Qwen, Llama, and Gemma
• The entire first half of the transformer can be zeroed - all computation lives in the top 50% of depth
• EAP fails non-gracefully: adding attribution-ranked heads to a working circuit dropped multi-key accuracy 73.3% → 26.7%
• A zero-training cosine gate at Layer 4 routes tasks correctly with zero false positives on 100 OOD prompts

---

Key Results

The Computational Budget Law

Cognitive Tier	Qwen K	Llama K	Gemma K	Sparsity
Syntax + Algorithmic Primitives	16 heads	16 heads	12–16 heads	~92–97%
Local Semantics + Type Dispatch	64 heads	64 heads	64 heads	~70–88%
Compositional Multi-Step Chains	168 heads	256 heads	104 heads	50%

NIAH Sparse Circuit (Qwen2.5-1.5B)

25-component JCCS circuit
→ compressible to 6 pure retrieval heads (1.78% of 336 total)
→ Knockout:  100% → 0%
→ Oracle steering retention: 95%

The EAP Double Indictment

Working 6-head sparse circuit:   Multi-Key accuracy = 73.3%
After adding EAP expansion heads: Multi-Key accuracy = 26.7%  ← worse than baseline

Attribution scoring does not fail gracefully. Causal patching is irreplaceable.

Zero-Training Router (Layer 4 Cosine Gate)

Task	Cosine Similarity	Decision
NIAH Single-Key (calibration)	baseline	ROUTE
NIAH Multi-Key	+0.781	✅ ROUTE
NIAH Multi-Value	+0.743	✅ ROUTE
Variable Tracking	+0.372	❌ ABSTAIN
Common Words	+0.441	❌ ABSTAIN
100 OOD Neutral	-	0 false positives

---

Method: JCCS

Joint Causal Circuit Search = greedy ACDC over attention heads and MLPs simultaneously, validated by CoAx specificity tagging.

Standard EAP/ACDC searches over attention heads only, or treats MLP blocks as atomic. JCCS searches jointly:

For each (head, MLP) component:
  1. Knock out component
  2. Measure task metric drop
  3. Greedily include if drop > threshold
  4. Validate with CoAx specificity tag
     (does ablating this component affect ONLY the target task?)

CoAx validation prevents false positives: a component that appears causal for NIAH but also degrades Arithmetic is tagged as shared infrastructure, not task-specific.

---

Code

Repository Structure

kheads/
├── paper/
│   └── k-heads-paper.pdf              # preprint
├── circuit/
│   ├── phase0/                        # Kill-switch ablation (attention head level)
│   │   ├── 01_extract_candidates.py
│   │   ├── 02_ablation_harness.py
│   │   ├── 03_evaluate_kill_switch.py
│   │   └── 04_robustness_check.py
│   ├── phase0b/                       # Extended ablation + EAP baselines
│   ├── phase1/                        # CFG engine + IOI sanity + NIAH closure
│   │   ├── 01_cfg_engine.py           # Core position-indexed EAP engine
│   │   ├── 02_cfg_ioi_sanity.py       # G1: IOI head recovery on GPT-2
│   │   ├── 03_cfg_niah_discovery.py   # NIAH CFG extraction
│   │   └── 04_cfg_closure_test.py     # G2: NIAH closure test (EAP failed -> JCCS)
│   └── phase2/                        # JCCS full pipeline
│       └── outputs/
│           └── 05_coax_discovery_raw.md
└── claims.md                          # All claims mapped to code + outputs

Requirements

pip install torch>=2.0 transformers>=4.40 einops

Models used: Qwen/Qwen2.5-1.5B · meta-llama/Llama-3.2-1B · google/gemma-2-2b

Running

# Phase 0: Kill-switch ablation on attention heads
python circuit/phase0/02_ablation_harness.py

# Phase 1: CFG engine + IOI sanity
python circuit/phase1/01_cfg_engine.py
python circuit/phase1/02_cfg_ioi_sanity.py

# Phase 1: NIAH closure test (validates attribution engine before JCCS)
python circuit/phase1/04_cfg_closure_test.py

---

Claims

See claims.md for every experimental claim mapped to its paper section, code file, output location, and result status.

---

ISA Research Program

#	Paper	Repo	Key Finding
1	HeadGenome	attentionheadgenome	Static weight geometry is blind; Sink heads are causally active
2	Necessary But Not Sufficient	necstuff	Head ISA falsified; Residual ISA confirmed (ρ = 0.9644)
3	Shape of Computation	taskgeometry	Computation = trajectory not point; geometry is real but not causally sufficient
4	K Heads Is All You Need	this repo	Causal circuit budget scales with cognitive complexity; EAP is unreliable

---

Citation

@article{khavin2027kheads,
  title={K Heads Is All You Need: Task-Dependent Circuit Granularity and Zero-Training Routing in LLMs},
  author={Khavin, S.},
  journal={Preprint},
  year={2027}
}

---

_{Independent mechanistic interpretability research · S. Khavin · Age 20}