[Submitted on 5 Jun 2026 (v1), last revised 6 Nov 2026 (this version, v5)]
Abstract:Reasoning-trained language models can select, zero-shot, a small set of relevant labels from thousands to hundreds of thousands of candidates. We ask how they do it mechanistically, and whether the mechanism can be distilled. We treat each decision as a token-level event scored by the model’s own margin: the token that picks a coarse region of the label space, the tokens that pick a label within it, and the token where the output departs from a near-miss named earlier. Attribution, exact mean-ablation, knock-in and a null calibration give individual attention heads causal standing. On clinical coding of MIMIC-IV discharge summaries, with all 5,651 candidate codes in context, a small, global, phase-structured set of heads is necessary and sufficient on essentially every summary; distinct head families read the candidate region and the earlier near-miss, and blocking what they read removes their effect; for mentions elicitable in the reasoning, a disjoint mid-layer set makes the region elicitable several tokens before the code. The structure recurs in a second reasoner. We introduce MISTILL, which, unlike chain-of-thought distillation, also supervises the student’s pooled attention at these decision events. Read on heads found after training, it nearly doubles the causal recovery of the contrastive decision in a cross-family student and adds a small, seed-stable gain in one that already carries most of it, with no detected task difference at matched precision; under matched selection it exceeds a frozen random-head pool on both recovery and task score. On two corpora with human-annotated evidence (one read zero-shot, one in domain), a MISTILL student’s decision heads attend to each code’s evidence more, and more specifically, than a chain-of-thought-distilled student’s. A pre-registered evidence-flip test detected no improvement in evidence-dependent decisions through the tested routes.
Submission history
From: Debjyoti Saha Roy [view email]
[v1] Fri, 5 Jun 2026 02:32:24 UTC (334 KB)
[v2] Tue, 6 Oct 2026 02:10:23 UTC (301 KB)
[v3] Fri, 23 Oct 2026 01:55:43 UTC (419 KB)
[v4] Thu, 29 Oct 2026 05:38:28 UTC (437 KB)
[v5] Fri, 6 Nov 2026 05:17:16 UTC (461 KB)
