Characterize Then Distill: Mechanistic Reasoning in Large Output Spaces

# Mechanistic Reasoning in Large Output Spaces
Source: [https://arxiv.org/html/2606.06840](https://arxiv.org/html/2606.06840)
## Characterize Then Distill: Mechanistic Reasoning in Large Output Spaces

Debjyoti Saha Roy, Byron C\. Wallace & Javed A\. Aslam Khoury College of Computer Sciences, Northeastern University \{saharoy\.d,b\.wallace,j\.aslam\}@northeastern\.edu

###### Abstract

Modern reasoning models offer surprisingly strong zero\-shot performance on challenging multi\-label tasks that require selecting a small set of relevant options from hundreds of thousands to millions of candidate labels\. We investigate how they achieve this mechanistically\. We characterize reasoning as a two\-phase process: A broad “shortlisting” of candidates followed by fine\-grained reasoning over the resulting set\. We provide evidence across a range of datasets that these steps can be isolated and are complementary\. Using this characterization, we develop a*mechanistic distillation*strategy that consistently outperforms standard distillation\. We[open\-source](https://github.com/research-anon-487/xcube/tree/reasoning)code and trained models\.

## 1Introduction

![Refer to caption](https://arxiv.org/html/2606.06840v1/x1.png)

Figure 1:LLaMA3\-70Bexhibits early*Focus*buildup & late*Confusion*reduction over CoT\. Mechanistic distillation recovers teacher trajectory while representative CoT distillations\(he\_semcot\_2025\)fail\.In very large\-scale multi\-label settings, a small set of relevant labels must be selected from hundreds of thousands to millions of candidates\(zhou\_quest\_2024;zhang\_elmo\_2025;ortego\_large\_2025\)\. Such settings arise in applications like Wikipedia category tagging\(gupta\_dual\-encoders\_2024;dahiya\_prototypical\_2025\), large\-scale e\-commerce categorization\(ortego\_large\_2025\), and document tagging in scientific repositories\(tabatabaei\_can\_2025\)\. These kinds of tasks are challenging due to the prevalence of many rare “tail” labels with few examples\.

This problem resembles the “needle\-in\-a\-haystack” challenge in long\-context evaluation, where models must retrieve or reason over sparse, critical information buried in vast distracting contexts\(kamradt2023needle;hsieh\_ruler\_2024\)\. This similarity suggests that challenging multi\-label tasks may benefit from recent effective reasoning strategies that have succeeded for such problems\. Recent reasoning models—such as OpenAI’s o\-series\(OpenAI2025a;OpenAI2025b\)and DeepSeek\-R1\(guo2025deepseek\)—aim to emulate*System 2*reasoning\(zhang\_llm\_2024;ji2025test;besta\_reasoning\_2025;xu\_towards\_2025;li\_system\_2025\)\. Reasoning has enabled strong performance on long context tasks\(modarressi\_nolima\_2025;wang\_reasoning\_2025;kuratov\_babilong\_2024;marjanovic\_deepseek\-r1\_2026\), by leveraging structured mechanisms such as long Chain\-of\-Thought \(CoT\)\(chen\_towards\_2025;yeo\_demystifying\_2025\), tree\-of\-thoughts\(yao\_tree\_2023\), and deliberative reasoning\(guo2025deepseek;guan\_deliberative\_2025\)\.

We ask:How do large reasoning models select a small set of relevant labels from millions of candidates without any task\-specific training?Past work on mechanistic interpretability\(bereska\_mechanistic\_2024;hanna\_have\_2024;geiger\_causal\_2025;du\_how\_2025\)has sought to characterize reasoning tasks like arithmetic\(nanda\_progress\_2023;stolfo\_mechanistic\_2023;chen\_how\_2025\), commonsense question answering\(chuang\_faithlm\_2025;basu\_mechanistic\_2025\), ontology reasoning\(dutta\_how\_2024\), and code generation\(michaud\_opening\_2024\)\. But the mechanisms of reasoning over very large output spaces has not, to our knowledge, been explored\.

We seek to characterize the implicit mechanisms used by large reasoning models and then systematically transfer them to smaller models\.A simpler approach would be to use standard CoT distillation for this purpose, which has worked well for other tasks\(wang\_scott\_2023;feng\_keypoint\-based\_2024;wadhwa\-etal\-2024\-investigating;chen\_skip\-thinking\_2025;chen\_unveiling\_2025;kang\_distilling\_2025;li\_llms\_2025\)\. These approaches aim to imitate the rationales from larger model rather than latent reasoning trajectories that give rise to them\(dai\_beyond\_2024;hu\_han\_2026;tan\_probing\_2026;chen\_molecular\_2026\)\. We observe that for these tasks smaller distilled models perform relatively poorly \(35–38% drops in performance; Figure[1](https://arxiv.org/html/2606.06840#S1.F1)\)\. We report that mechanistic distillation can result in stronger student model performance\.

Prompt:68\-year\-old male with progressivedyspnea,orthopnea, and leg edema\. Echo shows EF35%, consistent withheart failure\. \.\.\. mild pulmonary congestion; elevated creatinine suggestive ofcardiorenal syndrome\. Task:You are an expert medical coder\. Assign appropriate ICD codes from the provided large set \(10,000\+ labels\)\. I50\.23 — Acute on chronic systolic HFI50\.2 — Systolic \(congestive\) HFR60\.0 — Localized edemaJ18\.9 — Pneumonia I42\.9 — CardiomyopathyI13\.0 — Hypertensive heart and CKD w/ HFN17\.9 — Acute kidney failure J44\.9 — COPD \.\.\. \+9000 unrelated codes68\-year\-old male … cardiorenal syndrome \+ 10000 codes\.Pt w/ clear CHF \+ volume overload → progressive dyspnea, orthopnea, leg edemaEarly CoTs⇓\\DownarrowBroad focus areas:CHF, volume overload, dyspnea, orthopnea, leg edema68\-year\-old male … Pt w/ CHF volume overload …edemaEcho EF 35% → systolic dysfunction, HFrEFBroad focus areas:CHF, volume overload … leg edemasystolic dysfunction, HFrEF68\-year\-old male … EF 35% → HFrEFMild pulm cong \+ ↑Cr → cardiorenal synd in decomp HFrEF, need careful vol mgmtBroad focus areas:CHF, vol overload … HFrEF …cardiorenal syndrome, decompensated HFrEF68\-year\-old … CHF volume overload … edemaNot pneumonia/COPD exacerbation → true cardiogenic volume overload from systolic HFLater CoTs⇓\\DownarrowContrastive reasoning:ruled out: pneumonia, COPD×true: CHF volume overload✓… CHF vol overload … EF 35% → HFrEFNot pure cardiomyopathy or CAD → EF 35% locks in HFrEFContrastive reasoning:ruled out: cardiomyopathy, CAD×true: systolic HFrEF✓… HFrEF …Not isolated AKI from hypovolemia or nephrotoxins → true cardiorenal syndrome type 1 in setting of decompensated HFrEFContrastive reasoning:ruled out: isolated AKI, hypertensive heart alone×true: cardiorenal syndrome✓

Figure 2:LLM Reasoning\. \(*Left*\) Early CoT progressively builds focus on the broad categories\. \(*Right*\) Late CoT progressively rules out near\-miss categories until only the true signals remain\.We evaluate performance along complementary dimensions for the teacher and its distilled student \(formal definitions in Sec\.[2](https://arxiv.org/html/2606.06840#S2)\): \(1\)Reasoning focusmeasures how early CoT aligns with the dominant semantic signal in the input, capturing progressive evidence consolidation into coherent anchors\. \(2\)Near\-Miss Confusionmeasures how strongly evolved reasoning remains aligned with plausible but incorrect alternatives, capturing the effectiveness of contrastive refinement in ruling them out\.

As an illustration, consider the clinical example in Figure[2](https://arxiv.org/html/2606.06840#S1.F2), which serves as a running example throughout the paper; the task is to assign the correct ICD111Standardized medical labeling system with10410^\{4\}–10510^\{5\}codes, used as a benchmark for large\-output tasks\(liu\_lost\_2023\)\.codes\(who2025icd\)from a large label space\. An LLM exhibits*focused reasoning*: in early steps \(left panel\), it rapidly organizes evidence into coherent semantic anchors \(e\.g\., grouping symptoms under heart failure, then refining to systolic dysfunction\), effectively narrowing the search space\. In later steps \(right panel\), it performs*contrastive refinement*, explicitly ruling out closely related alternatives \(near\-misses\) such as competing diagnoses\. This yields both high*focus*and low*confusion*\. In contrast, smaller distilled models fail to maintain this structure: their reasoning drifts across loosely related concepts, lacking strong early anchoring and insufficiently separating true labels from plausible distractors\.

Ourmain contributionsare as follows:\(1\)We characterize and provide empirical, mechanistic evidence for a two\-phase process in large reasoning models for challenging multi\-label tasks—*coarse semantic filtering*, followed by*fine\-grained reasoning over a small shortlist*—and show the phases are complementary and isolatable\.\(2\)Using this characterization, we introduce a*mechanistic distillation*approach that directly supervises phase\-specific computations at the level of aggregated attention heads and consistently outperforms standard CoT distillation\.

## 2How do Reasoning Models do Large\-Scale Multi\-Labeling?

We present a two\-phase characterization of LLM reasoning in large spaces, with mechanistic evidence\.

Prompt:68\-year\-old male with progressivedyspnea,orthopnea, and leg edema\. Echo shows EF35%, consistent withheart failure\. \.\.\. mild pulmonary congestion; elevated creatinine suggestive of cardiorenal syndrome\. Task:Predict ICD\-10 codes from a large set \(10000\+ labels\)\. I50\.22 — Chronic systolic HFJ18\.9 — PneumoniaI42\.9 — CardiomyopathyN18\.3 — CKD \(stage 3\)I25\.10 — Coronary artery diseaseN17\.9 — Acute kidney failureJ44\.9 — COPD\.\.\. \+9000 unrelated codessemantic anchors\(from text \+ candidate labels\)dyspneaorthopneaEF 35%heart failurecardiorenalEarly Chain\-of\-Thought Reasoning Classicheart failuresigns:dyspnea,orthopnea, edema\. Echo: EF35%→\\rightarrowsystolic dysfunction\. Consistent with chronic systolic HF \(I50\.22\)\. Elevated creatinine→\\rightarrowcardiorenal syndrome\.Strong semantic signals identified earlyheart failuredyspneaEF 35%orthopneaMechanistic Question—Phase 1 When an*early CoT token*latches onto a strong semantic signal, what is an attention head doing internally?early CoT token:‘‘heart failure’’hhdyspneaEF 35%orthopneaWhat did it attend to?sharp focus on salient tokenshhresidual writetoward semantic anchorheart failuresystolicdysfunctionWhat did it write?a semantic update in the residual stream

Figure 3:Phase 1—While generating early CoT, we measure whether attention heads attend to salient evidence and write semantically aligned updates as𝖢𝗈𝖺𝗋𝗌𝖾𝖲𝖼𝗈𝗋𝖾\\mathsf\{CoarseScore\}\(Eq\.[1](https://arxiv.org/html/2606.06840#S2.E1)\)\.Phase 1: Coarse Semantic Filtering

Figure[3](https://arxiv.org/html/2606.06840#S2.F3)provides an example to build intuition for Phase 1\. The*top\-left panel*shows an input prompt and candidate label space, from which semantic anchors emerge\. The*bottom\-left panel*shows early CoT, where strong semantic signals \(e\.g\., heart failure, EF 35%\) are identified\. The*right panel*poses the key question:*when these early CoT tokens are produced, what are the attention heads doing internally?*Specifically, do heads \(i\) focus on relevant evidence, and \(ii\) write updates aligned with the semantic signal? This setup generalizes to large output spaces tasks\. Phase 1 captures this intuition via a*coarse score*, quantifying these properties, as detailed next\.

Early\-layer attention heads drive coarse filtering by linking salient input tokens to broad semantic anchors\.To provide evidence for this, we construct prompts with controlled semantic signals and analyze attention patterns and head contributions during CoT via mechanistic interventions, motivated by prior work showing that a small subset of heads drives most long\-context computation, while MLP layers play a limited role\(wu\_retrieval\_2024;zheng\_attention\_2024;tang\_razorattention\_2024;fu\_mixture\_2025\)\. We design*clean*prompts comprising text snippets with a strong coarse signal, followed by a long, multi\-domain label list\. For instance \(Figure[3](https://arxiv.org/html/2606.06840#S2.F3),*top left*\), a heart failure discharge summary is paired with related ICD codes \(cardiovascular/pulmonary\) and unrelated domains\. In turn, early CoT aligns with the dominant semantic signal in the input \(i\.e\., cardiovascular concepts;*bottom left*\)\.

For each head, we compute a*coarse filtering score*\(Equation[1](https://arxiv.org/html/2606.06840#S2.E1)\) on clean prompts, restricting the analysis to query positions from early CoT segments\. Here𝒬early​\(x\)\\mathcal\{Q\}\_\{\\text\{early\}\}\(x\)denotes early\-CoT query positions \(first1010–3030generated tokens\);𝜶q,⋅h\\boldsymbol\{\\alpha\}^\{h\}\_\{q,\\cdot\}the attention distribution of headhhat query positionqq, and;κ​\(𝜶q,⋅h\)\\kappa\(\\boldsymbol\{\\alpha\}^\{h\}\_\{q,\\cdot\}\)its excess kurtosis, which identifies heads with sharply peaked attention patterns\.Δh​\(q\)\\Delta^\{h\}\(q\)denotes the residual\-stream update written by headhhat positionqq\. Let𝒜​\(x\)\\mathcal\{A\}\(x\)be the set of relevant semantic anchors \(e\.g\. cardiovascular conditions like “heart failure” or “pulmonary edema”\) with embeddings𝐞​\(y\)\\mathbf\{e\}\(y\)\. Then the anchor centroid is defined as𝐜anchor​\(x\)≜meany∈𝒜​\(x\)⁡𝐞​\(y\)\\mathbf\{c\}\_\{\\mathrm\{anchor\}\}\(x\)\\triangleq\\operatorname\{mean\}\_\{y\\in\\mathcal\{A\}\(x\)\}\\mathbf\{e\}\(y\)\.

𝖢𝗈𝖺𝗋𝗌𝖾𝖲𝖼𝗈𝗋𝖾\(h\)=𝔼xq∈𝒬early​\(x\)\[κ​\(𝜶q,⋅h\)⏟sharp attention \(kurtosis\)⋅cos⁡\(Δh​\(q\),𝐜anchor​\(x\)\)⏟anchor\-aligning write\]\\boxed\{\\mathsf\{CoarseScore\}\(h\)=\\operatorname\*\{\\mathbb\{E\}\}\_\{\\begin\{subarray\}\{c\}x\\\\ q\\in\\mathcal\{Q\}\_\{\\mathrm\{early\}\}\(x\)\\end\{subarray\}\}\\Big\[\\underbrace\{\\hbox\{\\pagecolor\{NavyBlue\!20\}$\\kappa\\\!\\left\(\\boldsymbol\{\\alpha\}^\{h\}\_\{q,\\cdot\}\\right\)$\}\}\_\{\{\\color\[rgb\]\{0\.15,0\.488671875,0\.8306640625\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{0\.15,0\.488671875,0\.8306640625\}\\textbf\{sharp attention \(kurtosis\)\}\}\}\\\!\\\!\\\!\\\!\\cdot\\;\\;\\;\\;\\underbrace\{\\hbox\{\\pagecolor\{BurntOrange\!20\}$\\cos\\\!\\Big\(\\Delta^\{h\}\(q\),\\mathbf\{c\}\_\{\\mathrm\{anchor\}\}\(x\)\\Big\)$\}\}\_\{\{\\color\[rgb\]\{0\.8306640625,0\.4322265625,0\.15\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{0\.8306640625,0\.4322265625,0\.15\}\\textbf\{anchor\-aligning write\}\}\}\\Big\]\}\(1\)
High𝖢𝗈𝖺𝗋𝗌𝖾𝖲𝖼𝗈𝗋𝖾\\mathsf\{CoarseScore\}values identify heads thatattend sharply to semantic anchorsandalign the residual representation with those anchors\. As illustrated in Figure[4](https://arxiv.org/html/2606.06840#S2.F4)\(*Left*\), early\-layer heads \(e\.g\., L3H22, L2H7, and L1H12\) in theLLaMA3\-70Bmodel dominate the top ranks during early CoT generation\. Also, as illustrated in Figure[4](https://arxiv.org/html/2606.06840#S2.F4)\(*Right*\), when a top\-ranked early\-layer head is queried from the position immediately preceding an early CoT token, its attention weights become sharply peaked on a small number of semantic anchors in the prompt, such as “heart failure”, “EF 35%”, and “prior MI” \(cardiac anchors\), as well as related respiratory and renal anchors like “pulmonary edema” and “pneumonia”\. This indicates early\-layer heads rapidly detect and focus on relevant concepts\. We then ask:*Do these peaked heads causally implement coarse filtering?*

![Refer to caption](https://arxiv.org/html/2606.06840v1/fig_phase1_attn_heads_kurtosis_llama70b.pdf)
![Refer to caption](https://arxiv.org/html/2606.06840v1/inset_medical_early_head_attention_llama70b.pdf)

Figure 4:*\(Left\)*Top heads byCoarseScoreinLLaMA3\-70Bmodel\. Early\-layer heads \(L1–6\) exhibiting anchor\-focused attention and aligned residual updates during early CoT\.*\(Right\)*Attention from a top early\-layer head \(L3H22\), sharply focusing on key clinical anchors \(red: cardiac phrases like “heart failure” and “EF 35%”;orange: respiratory/renal like “pulmonary edema” & “pneumonia”\)\.![Refer to caption](https://arxiv.org/html/2606.06840v1/fig_phase1_denoising_patching.pdf)
![Refer to caption](https://arxiv.org/html/2606.06840v1/fig_phase1_noising_patching.pdf)

Figure 5:Early\-layer attention heads causally control coarse filtering\.\(*Left*\) Denoising patching progressively larger bins of top heads ranked by𝖢𝗈𝖺𝗋𝗌𝖾𝖲𝖼𝗈𝗋𝖾\\mathsf\{CoarseScore\}substantially restores reasoning focus toward semantic anchors\. \(*Right*\) Noising these heads significantly degrades focus\.To answer, we evaluated causal sufficiency/necessity using denoising/noising activation patching\(heimersheim\_how\_2024\)\. We ranked heads byCoarseScore\(Equation[1](https://arxiv.org/html/2606.06840#S2.E1)\) and grouped them into*cumulative bins*\(Top 3, Top 8, Top 16, Top 40, Top 100\)\. Within each bin, we patched the*attention patterns*of the selected heads, using randomly selected and low\-ranked \(count\-matched\) heads as negative controls\. For patching, we*corrupted*the clean prompt by perturbing the text snippet to remove coarse semantic signal \(e\.g\., replacing clinically meaningful terms like “dyspnea” or “heart failure” with neutral terms\)\. Consequently, early CoT segments show weak alignment with the \(original\) input semantics\. We measured a single metric capturing this alignment—*reasoning focus*—defined as how strongly early CoT tokens align with the dominant semantic signal in the clinical text portion of inputxx:𝖥𝗈𝖼𝗎𝗌=𝔼q∈𝒬early​cos⁡\(e​\(rq\),𝐜anchor​\(x\)\),\(2\)\\boxed\{\\mathsf\{Focus\}=\\mathbb\{E\}\_\{q\\in\\mathcal\{Q\}\_\{\\mathrm\{early\}\}\}\\cos\\\!\\big\(e\(r\_\{q\}\),\\mathbf\{c\}\_\{\\mathrm\{anchor\}\}\(x\)\\big\)\},\\;\(2\)whererqr\_\{q\}denotes CoT token at positionqq\(see Appendix[B](https://arxiv.org/html/2606.06840#A2.SSx1)for details\)\.

Figure[5](https://arxiv.org/html/2606.06840#S2.F5)\(*Left*\) shows that denoising patching—replacing attention patterns in corrupted runs with those from clean runs—progressively restores reasoning focus toward cardiovascular semantic anchors in corrupted runs \(up to 85% restoration\)\. Even the top 3 heads provided meaningful recovery, with further gains as more heads were included, demonstrating*sufficiency*\. Conversely, Figure[5](https://arxiv.org/html/2606.06840#S2.F5)\(*Right*\) shows that noising these same heads in clean runs produced diffuse early CoT traces, significantly disrupting reasoning focus \(up to 71% disruption\), with noticeable degradation even for the top few heads and further collapse as more heads were noised, demonstrating*necessity*\. In both cases the negative control heads defined earlier produced negligible effects\. Together these results demonstrate that specific early\-layer attention heads causally implement Phase 1 coarse filtering by linking salient tokens to broad semantic anchors, enabling early\-stage pruning of the label space\.

Prompt:68\-year\-old male with progressivedyspnea,orthopnea, and leg edema\. Echo shows EF35%, consistent withheart failure\. \.\.\. mild pulmonary congestion; elevated creatinine suggestive of cardiorenal syndrome\. Task:Predict ICD\-10 codes from a large set \(10000\+ labels\)\. I50\.22 — Chronic systolic HFJ18\.9 — PneumoniaI42\.9 — CardiomyopathyN18\.3 — CKD \(stage 3\)I25\.10 — Coronary artery diseaseN17\.9 — Acute kidney failureJ44\.9 — COPD\.\.\. \+9000 unrelated codesLater Fine\-Grained Chain\-of\-Thought Reasoning EF 35% confirmssystolic dysfunction \(HFrEF\)— rules outHFpEF\(EF typically≥\\geq50%\)\. Chronic orthopnea \+ peripheral edema pattern favorschronic systolic HFover acute myocarditis or flash pulmonary edema from other causes\. Elevated creatinine in low\-output state supports cardiorenal syndrome \(secondary to HF\) rather than primary renal disease as driver of fluid overload\.Contrastive differentiation: true vs near\-miss candidatesHFrEF \(EF 35%\)vs HFpEF \(EF≥\\geq50%\)Chronic systolic HFvs Acute myocarditis / flash edemaCardiorenal syndrome \(secondary\)vs Primary renal failureMechanistic Question—Phase 2 When a*later CoT token*performs fine\-grained contrastive reasoning, what is an attention head doing internally?later CoT token:‘‘HFrEF \(EF 35%\)’’hhEF 35%✓HFpEF×chronic orthopnea✓What did it attend to?iterative contrastive attention:boost own prior true shortlist\(✓\)suppress near\-miss candidates\(×\)hhresidual writewiden target–near\-missHFrEF confirmed✓HFpEF ruled out×What did it write?margin shaping:boost target\(✓\)separate near\-miss\(×\)

Figure 6:Phase 2—During later CoT, attention heads refine predictions by suppressing near\-misses, widening margins between these and the correct labels \(as captured by𝖱𝖾𝖿𝗂𝗇𝖾𝖲𝖼𝗈𝗋𝖾\\mathsf\{RefineScore\}; Equation[3](https://arxiv.org/html/2606.06840#S2.E3)\)\.Phase 2: Fine\-Grained Reasoning over a Shortlist

Figure[6](https://arxiv.org/html/2606.06840#S2.F6)conveys the intuition using the same representative example\. The*bottom\-left panel*shows later CoT, where the model contrastively differentiates among shortlisted candidates \(e\.g\., HFrEF vs\. HFpEF\)\. The*right panel*asks:*When processing later CoT tokens, what are the attention heads doing internally?*We argue that they \(i\) iteratively attend to prior shortlisted candidates, reinforcing true ones while suppressing near\-misses, and; \(ii\) write updates that widen the margin between target and competing alternatives\. We quantify this intuition via a*refine score*, defined below\.

Later\-layer attention heads drive Phase 2 iterative refinement by re\-attending to the model’s own earlier representations of the retained shortlist and widening margins between target and near\-miss subcategories\.To investigate and provide evidence we perform mechanistic interventions analogous to Phase 1\. We first ask:*Which heads iteratively loop back to their own prior representations of the retained shortlist while suppressing representations of subcategories that were previously shortlisted but have since been ruled out \(near\-misses\)?*To answer, for each headhh, we define a*Q​KQK\-refinement*score\. We mark positionqqin the CoT as*refinement*position if it contains contrastive reasoning \(like confirming “HFrEF” and ruling out “HFpEF” in the bottom\-left cardiovascular example of Figure[6](https://arxiv.org/html/2606.06840#S2.F6)\)\. At refinement positionsqq, we extract the head\-specific query vector𝐪qh\\mathbf\{q\}^\{h\}\_\{q\}\. For earlier positionsppwhere the CoT explicitly referenced shortlist candidates \(tokens like “HFrEF/systolic”\) or ruled out near\-miss subcategories \(“HFpEF/diastolic”\), we extract key vectors𝐤ph\\mathbf\{k\}^\{h\}\_\{p\}\. We then computeSimshortlist​\(q\)=𝔼p∈𝒮​cos⁡\(𝐪qh,𝐤ph\)\\text\{Sim\}\_\{\\textsf\{shortlist\}\}\(q\)=\\mathbb\{E\}\_\{p\\in\\mathcal\{S\}\}\\cos\(\\mathbf\{q\}^\{h\}\_\{q\},\\mathbf\{k\}^\{h\}\_\{p\}\)andSimnear​\(q\)=𝔼p∈𝒩​cos⁡\(𝐪qh,𝐤ph\)\\text\{Sim\}\_\{\\textsf\{near\}\}\(q\)=\\mathbb\{E\}\_\{p\\in\\mathcal\{N\}\}\\cos\(\\mathbf\{q\}^\{h\}\_\{q\},\\mathbf\{k\}^\{h\}\_\{p\}\), where𝒮\\mathcal\{S\}and𝒩\\mathcal\{N\}denote shortlist and near\-miss positions, respectively\. Then we compute the refinement score as,𝔼q​\[Simshortlist​\(q\)−Simnear​\(q\)\]\\mathbb\{E\}\_\{q\}\[\\text\{Sim\}\_\{\\text\{shortlist\}\}\(q\)\-\\text\{Sim\}\_\{\\text\{near\}\}\(q\)\]averaged over refinement positions and examples\.

𝖱𝖲​\(h\)=𝔼xq∈ℛ​\(x\)\[\(𝔼p∈𝒮​\(x\)cos⁡\(𝐪qh,𝐤ph\)⏞re\-attend to prior shortlist−𝔼p∈𝒩​\(x\)cos⁡\(𝐪qh,𝐤ph\)⏞suppress near\-miss\)⏟QK loop\-back \(where to read\)⋅⟨Δh​\(q\),∇𝐡final​\(q\)margin​\(q\)⟩⏟OV helpfulness \(what to write\)\]\\boxed\{\\mathsf\{RS\}\(h\)\\\!=\\\!\\\!\\\!\\\!\\\!\\\!\\operatorname\*\{\\mathbb\{E\}\}\_\{\\begin\{subarray\}\{c\}x\\\\ q\\in\\mathcal\{R\}\(x\)\\end\{subarray\}\}\\\!\\\!\\Big\[\\underbrace\{\\Big\(\\overbrace\{\\hbox\{\\pagecolor\{Teal\!20\}$\\operatorname\*\{\\mathbb\{E\}\}\\limits\_\{p\\in\\mathcal\{S\}\(x\)\}\{\\cos\(\\mathbf\{q\}^\{h\}\_\{q\},\\mathbf\{k\}^\{h\}\_\{p\}\)\}$\}\}^\{\\text\{re\-attend to prior shortlist\}\}\-\\overbrace\{\\hbox\{\\pagecolor\{Teal\!20\}$\\operatorname\*\{\\mathbb\{E\}\}\\limits\_\{p\\in\\mathcal\{N\}\(x\)\}\\cos\(\\mathbf\{q\}^\{h\}\_\{q\},\\mathbf\{k\}^\{h\}\_\{p\}\)$\}\}^\{\\text\{suppress near\-miss\}\}\\Big\)\}\_\{\{\\color\[rgb\]\{0\.15,0\.575,0\.575\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{0\.15,0\.575,0\.575\}\\textbf\{QK loop\-back \(where to read\)\}\}\}\\cdot\\underbrace\{\\hbox\{\\pagecolor\{purple\!20\}$\\langle\\Delta^\{h\}\(q\),\\operatorname\*\{\\nabla\}\\limits\_\{\\mathbf\{h\}\_\{\\mathrm\{final\}\}\(q\)\}\\mathrm\{margin\}\(q\)\\rangle$\}\}\_\{\{\\color\[rgb\]\{0\.7875,0\.15,0\.3625\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{0\.7875,0\.15,0\.3625\}\\textbf\{OV helpfulness \(what to write\)\}\}\}\\Big\]\}\(3\)
*Do these heads also write residual updates that widen target–near\-miss margins viaO​VOV?*To investigate, we define a*O​VOV\-helpfulness*score per head\. For each refinement positionqqwe look at every positionpp\(previous mention of shortlist token\) that this head was paying significant attention \(α\>0\.3\\alpha\>0\.3\)\. We call these position attended shortlist ‘keys’\. Then we get the head updates for these ‘key’ positions using:Δh​\(q\)=∑keysWOh​\(α⋅WVh​𝐡key\)\\Delta^\{h\}\(q\)=\\sum\_\{\\text\{keys\}\}W^\{h\}\_\{O\}\\big\(\\alpha\\cdot W^\{h\}\_\{V\}\\mathbf\{h\}\_\{\\text\{key\}\}\\big\)\. To measure how strongly the model favors correct shortlist labels over the near\-miss labels at refinement stepqq, we definemargin​\(q\)=avg​\_​logittarget⁡\(q\)−avg​\_​logitnear\_miss⁡\(q\)\\textsf\{margin\}\(q\)=\\operatorname\{avg\\\_logit\}\_\{\\text\{target\}\}\(q\)\-\\operatorname\{avg\\\_logit\}\_\{\\text\{near\\\_miss\}\}\(q\)\(see Appendix[B](https://arxiv.org/html/2606.06840#A2), Equation[7](https://arxiv.org/html/2606.06840#A2.E7)\)\. Then to measure how much an head update aligns the residual stream in a direction that favors the target logits—for example, by increasinglogit​\(systolic\)\\textrm\{logit\}\(\\texttt\{systolic\}\)relative tologit​\(diastolic\)\\textrm\{logit\}\(\\texttt\{diastolic\}\)—we compute*O​VOV\-helpfulness*score as:OVhelph​\(q\)=Δh​\(q\)⋅\(∂margin​\(q\)/∂hfinal​\(q\)\)\\mathrm\{OV\}^\{h\}\_\{\\text\{help\}\}\(q\)=\\Delta^\{h\}\(q\)\\cdot\(\\partial\\,\\textsf\{margin\}\(q\)/\\partial h\_\{\\text\{final\}\}\(q\)\), averaged over refinement positions and examples\. Finally, we compute𝖱𝖾𝖿𝗂𝗇𝖾𝖲𝖼𝗈𝗋𝖾\\mathsf\{RefineScore\}as the product of*QK*and*OV*scores \(Equation[3](https://arxiv.org/html/2606.06840#S2.E3)\)\.

High positive𝖱𝖾𝖿𝗂𝗇𝖾𝖲𝖼𝗈𝗋𝖾\\mathsf\{RefineScore\}identifies heads that re\-attend to prior shortlist representations while avoiding near\-misses, and write updates that widen the target–near\-miss margin\. As illustrated in Figure[7](https://arxiv.org/html/2606.06840#S2.F7)\(*Left*\), mid\-to\-late heads \(L14 and above\) inLLaMA3\-70Bshow strong iterative looping and discriminative updates\. Figure[7](https://arxiv.org/html/2606.06840#S2.F7)\(right\) illustrates the temporal dynamics of the top ranked heads across CoT stages \(x\-axis: Early to Late tokens\)\. These heads show increasing QK similarity to their prior shortlist \(top\-left\) while decreasing similarity to near\-miss candidates \(top\-right\), e\.g\., attending more to “EF 35% confirms systolic dysfunction \(HFrEF\)” and less to “HFpEF \(≥\\geq50%\)” or “acute myocarditis” as the CoT progresses\. Concurrently, the OV circuit \(bottom\) shows rising alignment \(cos⁡\(Δh,∇margin\)\\cos\(\\Delta^\{h\},\\nabla\\mathrm\{margin\}\)\) with widening the logit margin between correct and alternative options—“chronic systolic HF” vs\. “flash edema”, and “secondary cardiorenal syndrome” vs\. “primary renal disease”—supporting contrastive reasoning in later stages\.

![Refer to caption](https://arxiv.org/html/2606.06840v1/fig_phase2_refine_score_ranking.pdf)
![Refer to caption](https://arxiv.org/html/2606.06840v1/fig_phase2_refine_temporal_dynamics.pdf)

Figure 7:Later\-layer heads showing iterative refinement\.*\(Left\)*Mid\-to\-late attention heads ranked byQ​KQKpreference for own prior shortlist keys over near\-miss keys \+O​VOVupdates widening margins\.*\(Right\)*Over CoT, top refinement heads sharpenQ​KQKpreference toward shortlist representations, downweight near\-misses, and strengthenO​VOVcontributions for iterative margin widening\.![Refer to caption](https://arxiv.org/html/2606.06840v1/fig_phase2_denoising_patching.pdf)
![Refer to caption](https://arxiv.org/html/2606.06840v1/fig_phase2_noising_patching.pdf)

Figure 8:Later\-layer heads causally drive iterative refinement\.*\(Left/Right\)*Denoising successively larger bins of top𝖱𝖾𝖿𝗂𝗇𝖾𝖲𝖼𝗈𝗋𝖾\\mathsf\{RefineScore\}\-ranked heads sharply↓↓\\downarrow\\downarrownear\-miss confusion & noising them↑↑\\uparrow\\uparrowit\.Finally, we ask:*Do these heads causally drive iterative refinement?*Similar to Phase 1, we perform activation patching: We ranked heads by*refinement score*, formed cumulative bins \(Top 3–100\), and patched their*attention patterns*, using count\-matched random and low\-ranked heads as controls\. We corrupt the clean prompt by perturbing the text snippet to favor near\-miss subcategories—for example, in the cardiovascular \(top left panel\) prompt in Figure[6](https://arxiv.org/html/2606.06840#S2.F6), replacingprogressive→\\rightarrowacute,EF 35%→\\rightarrow55%, andcardiorenal syndrome→\\rightarrowprimary renal failure\. Consequently, later CoT segments gets steered toward the near\-miss semantics rather than the true target, e\.g\., towarddiastolicinterpretation as opposed tosystolicfor a clean prompt\. We measure a single metric capturing this effect—*near\-miss confusion*—defined as how strongly the*anchor\-conditioned late CoT*aligns with near\-miss representations:𝖢𝗈𝗇𝖿𝗎𝗌𝗂𝗈𝗇=𝔼q∈ℛrefinecos\(e\(sq\),𝐜near\(x\)\),​\(4\)\\boxed\{\\mathsf\{Confusion\}=\\mathbb\{E\}\_\{q\\in\\mathcal\{R\}\_\{\\mathrm\{refine\}\}\}\\cos\\\!\\big\(e\(s\_\{q\}\),\\,\\mathbf\{c\}\_\{\\mathrm\{near\}\}\(x\)\\big\),\}\\;\(4\)wheresq=a​\(x\)∥\[r1​…​rq\]s\_\{q\}=a\(x\)\\\|\[r\_\{1\}\\dots r\_\{q\}\]is the cumulative prefix formed by appending refinement CoT tokens to a fixed anchor text, and𝐜near​\(x\)\\mathbf\{c\}\_\{\\mathrm\{near\}\}\(x\)denotes the centroid of near\-miss categories𝒩​\(x\)\\mathcal\{N\}\(x\)\(see App\.[B](https://arxiv.org/html/2606.06840#A2.SSx2)for details\)\.

Figure[8](https://arxiv.org/html/2606.06840#S2.F8)\(*Left*\) shows that denoising patching progressively lowers near\-miss confusion at refinement positions\. Even the top 3 heads provided sharp reduction, with further improvements as more heads were included \(upto 67% reduction with the top 100 heads\), proving*sufficiency*\. Conversely, noising the same heads \(*right*panel\) steers the refinement\-CoT toward near\-miss subcategories, boosting near\-miss confusion by up to 200%, proving*necessity*\. Negative controls had negligible effects in both cases\. Together these results demonstrate that specific mid\-to\-late layer heads causally implement sharp distinctions between target and near\-miss subcategories during contrastive reasoning\.

## 3How Can We Distill the Reasoning Mechanism?

Having characterized two mechanisms for large\-scale multi\-labeling—Phase 1 coarse filtering and Phase 2 iterative margin refinement—we perform*mechanistic distillation*\. To align phase\-specific signals, we distill via teacher forcing on the teacher\-generated CoT trace\(wang\_scott\_2023;yan\_towards\_2025;chen\_unveiling\_2025;fang\_knowledge\_2026\)\. For inputxx, the teacher producesτT​\(x\)\\tau^\{T\}\(x\), and both models are run onzT​\(x\)=z^\{T\}\(x\)=\[prompt⁡\(x\);τT​\(x\)\\operatorname\{prompt\}\(x\);\\tau^\{T\}\(x\)\], allowing comparisons at aligned query positions\. Our goal is to transfer not merely the teacher’s generated tokens, but the underlying*mechanics*that produce them\. Our distillation objective mirrors this\. Because teacher and student heads may not align one\-to\-one, we identify high\-scoring head*pools*for each phase and distill three signals: \(1\)*where*the model reads, \(2\)*what*it writes into residual stream, & \(3\) the interaction between the two\.

For Phase 1 we select the top\-kkearly\-layer heads by the coarse\-filtering score \(Eq\.[1](https://arxiv.org/html/2606.06840#S2.E1)\); for Phase 2 the top\-kkmid\-to\-late heads by the refinement score \(Eq\.[3](https://arxiv.org/html/2606.06840#S2.E3)\)\. LetℋP1\{T,S\}\\mathcal\{H\}^\{\\\{\\\!T,S\\\!\\\}\}\_\{\\mathrm\{P1\}\}denote the resulting head sets\. ForM∈\{T,S\}M\\\!\\in\\\!\\\{T,S\\\}andh∈ℋP1Mh\\\!\\in\\\!\\mathcal\{H\}^\{M\}\_\{\\mathrm\{P1\}\}, we assign normalized weightswhM=exp⁡\(𝖢𝖲M​\(h\)/τ\)∑h′exp⁡\(𝖢𝖲M​\(h′\)/τ\),h′∈ℋP1M,w\_\{h\}^\{M\}=\\frac\{\\exp\(\\mathsf\{CS\}^\{M\}\(h\)/\\tau\)\}\{\\sum\_\{h^\{\\prime\}\}\\exp\(\\mathsf\{CS\}^\{M\}\(h^\{\\prime\}\)/\\tau\)\},\\quad h^\{\\prime\}\\\!\\in\\\!\\mathcal\{H\}^\{M\}\_\{\\mathrm\{P1\}\},whereτ\\taucontrols weight sharpness\. Phase 2 uses same weighting scheme, replacing𝖢𝖲\\mathsf\{CS\}with𝖱𝖲\\mathsf\{RS\}\.

Phase 1 heads route attention toward salient tokens during early reasoning\(cf\. Eq\.[1](https://arxiv.org/html/2606.06840#S2.E1)\)\. To distill this behavior, we pool attention using the head sets and weights defined above:𝜶¯q,⋅M=∑hwhM​𝜶q,⋅M,h,h∈ℋP1M,M∈\{T,S\}\.\\bar\{\\boldsymbol\{\\alpha\}\}^\{M\}\_\{q,\\cdot\}=\\sum\_\{h\}w\_\{h\}^\{M\}\\,\\boldsymbol\{\\alpha\}^\{M,h\}\_\{q,\\cdot\},\\;h\\\!\\in\\\!\\mathcal\{H\}^\{M\}\_\{\\mathrm\{P1\}\},\\;M\\\!\\in\\\!\\\{T,S\\\}\.Phase 1 heads also write updates that align the residual representation with the relevant semantic anchor\(cf\. Eq\.[1](https://arxiv.org/html/2606.06840#S2.E1)\)\. To distill this behavior, we pool write vectors using the head sets and weights defined above:Δ¯M​\(q\)=∑hwhM​ΔM,h​\(q\),h∈ℋP1M,M∈\{T,S\}\.\\bar\{\\Delta\}^\{M\}\(q\)=\\sum\_\{h\}w\_\{h\}^\{M\}\\,\\Delta^\{M,h\}\(q\),\\;h\\\!\\in\\\!\\mathcal\{H\}^\{M\}\_\{\\mathrm\{P1\}\},\\;M\\\!\\in\\\!\\\{T,S\\\}\.We then compute the anchor\-alignment scoreAM​\(q\)=cos⁡\(Δ¯M​\(q\),𝐜anchor​\(x\)\)\.A^\{M\}\(q\)=\\operatorname\{cos\}\\\!\\Big\(\\bar\{\\Delta\}^\{M\}\(q\),\\mathbf\{c\}\_\{\\mathrm\{anchor\}\}\(x\)\\Big\)\.Phase 1 coarse filtering arises from the interaction between sharp attention and anchor\-aligning writes\(cf\. Eq\.[1](https://arxiv.org/html/2606.06840#S2.E1)\)\. To distill this joint behavior, we match a pooled coarse\-filtering score combining pooled attention and pooled writes\. ForM∈\{T,S\}M\\\!\\in\\\!\\\{T,S\\\}we definePCSM​\(q\)=κ​\(𝜶¯q,⋅M\)⋅cos⁡\(Δ¯M​\(q\),𝐜anchor​\(x\)\),\\mathrm\{PCS\}^\{M\}\(q\)=\\kappa\\\!\\left\(\\bar\{\\boldsymbol\{\\alpha\}\}^\{M\}\_\{q,\\cdot\}\\right\)\\cdot\\operatorname\{cos\}\\\!\\Big\(\\bar\{\\Delta\}^\{M\}\(q\),\\mathbf\{c\}\_\{\\mathrm\{anchor\}\}\(x\)\\Big\),mirroring multiplicative structure of Eq\.[1](https://arxiv.org/html/2606.06840#S2.E1)at the pooled\-head level \(cf\. App\.[B](https://arxiv.org/html/2606.06840#A2.SSx4)\)\.*We minimize the teacher–student discrepancies in𝛂¯q,⋅M\\bar\{\\boldsymbol\{\\alpha\}\}^\{M\}\_\{q,\\cdot\},AM​\(q\)A^\{M\}\(q\), andPCSM​\(q\)\\mathrm\{PCS\}^\{M\}\(q\)via theKL,write, &interactionterms in Eq\.[5](https://arxiv.org/html/2606.06840#S3.E5)\.*

ℒattnP1\+ℒwriteP1\+ℒintP1=𝔼xq∈𝒬early​\(x\)\[KL​\(𝜶¯q,⋅T∥𝜶¯q,⋅S\)\+\(AT​\(q\)−AS​\(q\)\)2\+\(PCSS​\(q\)−PCST​\(q\)\)2\],\\hbox\{\\pagecolor\{NavyBlue\!20\}$\\mathcal\{L\}\_\{\\mathrm\{attn\}\}^\{\\mathrm\{P1\}\}$\}\+\\hbox\{\\pagecolor\{BurntOrange\!20\}$\\mathcal\{L\}\_\{\\mathrm\{write\}\}^\{\\mathrm\{P1\}\}$\}\+\\hbox\{\\pagecolor\{OliveGreen\!20\}$\\mathcal\{L\}\_\{\\mathrm\{int\}\}^\{\\mathrm\{P1\}\}$\}=\\operatorname\*\{\\mathbb\{E\}\}\_\{\\begin\{subarray\}\{c\}x\\\\ q\\in\\mathcal\{Q\}\_\{\\mathrm\{early\}\}\(x\)\\end\{subarray\}\}\\Bigl\[\\hbox\{\\pagecolor\{NavyBlue\!20\}$\\mathrm\{KL\}\\Bigl\(\\bar\{\\boldsymbol\{\\alpha\}\}^\{T\}\_\{q,\\cdot\}\\\|\\bar\{\\boldsymbol\{\\alpha\}\}^\{S\}\_\{q,\\cdot\}\\Bigr\)$\}\+\\hbox\{\\pagecolor\{BurntOrange\!20\}$\(A^\{T\}\(q\)\-A^\{S\}\(q\)\)^\{2\}$\}\+\\hbox\{\\pagecolor\{OliveGreen\!20\}$\\bigl\(\\mathrm\{PCS\}^\{S\}\(q\)\-\\mathrm\{PCS\}^\{T\}\(q\)\\bigr\)^\{2\}$\}\\Bigr\],\(5\)ℒloopP2\+ℒwriteP2\+ℒintP2=𝔼xq∈ℛ​\(x\)\[\(RT​\(q\)−RS​\(q\)\)2\+\(HT​\(q\)−HS​\(q\)\)2\+\(P​R​ST​\(q\)−P​R​SS​\(q\)\)2\]\.\\hbox\{\\pagecolor\{Teal\!20\}$\\mathcal\{L\}\_\{\\mathrm\{loop\}\}^\{\\mathrm\{P2\}\}$\}\+\\hbox\{\\pagecolor\{purple\!20\}$\\mathcal\{L\}\_\{\\mathrm\{write\}\}^\{\\mathrm\{P2\}\}$\}\+\\hbox\{\\pagecolor\{Amber\!20\}$\\mathcal\{L\}\_\{\\mathrm\{int\}\}^\{\\mathrm\{P2\}\}$\}=\\operatorname\*\{\\mathbb\{E\}\}\_\{\\begin\{subarray\}\{c\}x\\\\ q\\in\\mathcal\{R\}\(x\)\\end\{subarray\}\}\\Bigl\[\\hbox\{\\pagecolor\{Teal\!20\}$\(R^\{T\}\(q\)\-R^\{S\}\(q\)\)^\{2\}$\}\+\\hbox\{\\pagecolor\{purple\!20\}$\(H^\{T\}\(q\)\-H^\{S\}\(q\)\)^\{2\}$\}\+\\hbox\{\\pagecolor\{Amber\!20\}$\(PRS^\{T\}\(q\)\-PRS^\{S\}\(q\)\)^\{2\}$\}\\Bigr\]\.\(6\)

Phase 2 heads refine decisions by repeatedly re\-attending to previously mentioned shortlist representations while suppressing near\-miss alternatives\(cf\. Eq\.[3](https://arxiv.org/html/2606.06840#S2.E3)\)\.Phase 2 heads also write residual updates that widen the margin between target and near\-miss labels via theO​VOVcircuit\(cf\. Eq\.[3](https://arxiv.org/html/2606.06840#S2.E3)\)\.Phase 2 refinement emerges from the interaction between loop\-back retrieval and margin\-widening writes\(cf\. Eq\.[3](https://arxiv.org/html/2606.06840#S2.E3)\)\. Analogous to the Phase 1 distillation outlined above, we distill pooled Phase 2 refinement heads using signals: \(1\) loop\-back retrieval behavior, \(2\) margin\-improving write updates, and \(3\) their interaction\. Due to space constraints, we defer the full formulation and definitions ofRM​\(q\)R^\{M\}\(q\),HM​\(q\)H^\{M\}\(q\), andPRSM​\(q\)\\mathrm\{PRS\}^\{M\}\(q\)to App\.[B](https://arxiv.org/html/2606.06840#A2.SSx5); the resulting objective is shown in Eq\.[6](https://arxiv.org/html/2606.06840#S3.E6)\.

The final objective combines standard CoT distillation with our mechanistic distillation:ℒtotal=ℒCoT⏟vanilla CoT distillation\+λP1​ℒP1\+λP2​ℒP2⏟mechanistic distillation \(ours\)\\mathcal\{L\}\_\{\\mathrm\{total\}\}=\\underbrace\{\\mathcal\{L\}\_\{\\mathrm\{CoT\}\}\}\_\{\\text\{vanilla CoT distillation\}\}\+\\underbrace\{\\lambda\_\{\\mathrm\{P1\}\}\\mathcal\{L\}^\{\\mathrm\{P1\}\}\+\\lambda\_\{\\mathrm\{P2\}\}\\mathcal\{L\}^\{\\mathrm\{P2\}\}\}\_\{\\textbf\{mechanistic distillation \(ours\)\}\}whereℒCoT\\mathcal\{L\}\_\{\\mathrm\{CoT\}\}is the standard token\-level distillation loss under teacher forcing \(App\.[B](https://arxiv.org/html/2606.06840#A2.SSx6)\),ℒP1\\mathcal\{L\}^\{\\mathrm\{P1\}\}andℒP2\\mathcal\{L\}^\{\\mathrm\{P2\}\}are the Phase 1&2 mechanistic losses \(Eqs\.[5](https://arxiv.org/html/2606.06840#S3.E5),[6](https://arxiv.org/html/2606.06840#S3.E6)\), andλP1,λP2\\lambda\_\{\\mathrm\{P1\}\},\\lambda\_\{\\mathrm\{P2\}\}weight the coarse\-filtering and refinement supervision terms\.

## 4Do Mechanistically Distilled Students Reason Like the Teacher?

Setup & Baselines\.We evaluate onMIMIC\-IV\-full\(ICD\-10 coding from ICU discharge summaries\) andLF\-WikiSeeAlso\-320K\(Wikipedia related\-article recommendation\), following recent work\(johnson2023mimic;dahiya\_prototypical\_2025\); cf\. App\.[C](https://arxiv.org/html/2606.06840#A3)for details\. We use aLLaMA\\mathrm\{LLaMA\}\-70B teacher andLLaMA\\mathrm\{LLaMA\}\-7B/13B students, trained under teacher forcing on identical teacher\-generated CoT traces\. Baselines use standard CoT distillation \(token\-level cross\-entropy \+ KL; see App\.[B](https://arxiv.org/html/2606.06840#A2.SSx6)\), as well as recent strong methods: \(1\) UniCoTT\(zhuang\_unicott\_2025\): structural CoT distillation; \(2\) MoRSD\(yan\_towards\_2025\): self\-guided rationale selection; \(3\) SemCoT\(he\_semcot\_2025\): implicit\-token CoT; \(4\) CWT\(chen\_skip\-thinking\_2025\): chunk\-wise CoT training\. Our method augments standard CoT with phase\-specific supervision \(Eqs\.[5](https://arxiv.org/html/2606.06840#S3.E5),[6](https://arxiv.org/html/2606.06840#S3.E6)\):ℒCoT\+λP1​ℒP1\+λP2​ℒP2\.\\mathcal\{L\}\_\{\\mathrm\{CoT\}\}\+\\lambda\_\{\\mathrm\{P1\}\}\\mathcal\{L\}^\{\\mathrm\{P1\}\}\+\\lambda\_\{\\mathrm\{P2\}\}\\mathcal\{L\}^\{\\mathrm\{P2\}\}\.We tuneλP1,λP2\\lambda\_\{\\mathrm\{P1\}\},\\lambda\_\{\\mathrm\{P2\}\}on validation and keep training budget and optimization fixed across methods\. Inference uses greedy decoding \(temperature=0=0\), generating full CoT traces followed by final multi\-label predictions\.

Evaluation Metrics\.We evaluate distillation quality using CoT fidelity metric of teacher\-faithfulness, measured by Leakage\-Adjusted Simulatability \(LAS; App\.[D\.1](https://arxiv.org/html/2606.06840#A4.SS1)\)\(wang\_scott\_2023\)\. To assess whether student recovers the phase\-specific mechanisms from Section[2](https://arxiv.org/html/2606.06840#S2), we measure: \(a\)Reasoning Focuswhich captures Phase 1 \(coarse filtering; Eq\. \([2](https://arxiv.org/html/2606.06840#S2)\)\), \(b\)Near\-Miss Confusionwhich captures Phase 2 \(refinement; Eq\. \([2](https://arxiv.org/html/2606.06840#S2)\)\)\. To further probe, we evaluate mechanistic fidelity using theCoarseScore​\(Phase 1; Eq\.[1](https://arxiv.org/html/2606.06840#S2.E1)\) andRefineScore\(Phase 2; Eq\.[3](https://arxiv.org/html/2606.06840#S2.E3)\) from Sec\.[2](https://arxiv.org/html/2606.06840#S2), which provide mechanistic views of the above metrics\. We also report standard task metrics \(e\.g\., Macro\-F1; App\.[D\.2](https://arxiv.org/html/2606.06840#A4.SS2)\)\.

RQ1: Do mechanistically distilled students remain more*focused*and less*confused*than ones trained with vanilla CoT distillation?Table[1](https://arxiv.org/html/2606.06840#S4.T1)comparesFocusandConfusionof mechanistically distilled students \(LLaMA3\-7B,LLaMA3\-13B\) against CoT distillation baslines onLF\-WikiSeeAlso\-320KandMIMIC\-IV\-full, usingLLaMA3\-70Bas the teacher; green rows denote mechanistic distillation, with per\-dataset average gains\. Baselines achieve partial imitation \(LAS\), but lowFocusand highConfusionshow it does not recover the teacher’s coarse\-to\-fine reasoning\. In contrast, mechanistic distillation delivers largestFocusandConfusiongains, showing that transferring Phase 1/2 behaviors via our method brings students closer to the teacher’s reasoning\.

RQ2: Do mechanistically distilled students recover the teacher’s phase\-specific reasoning mechanisms?Table[1](https://arxiv.org/html/2606.06840#S4.T1)reports the constituent subscores ofCoarseScore\(Eq\.[1](https://arxiv.org/html/2606.06840#S2.E1)\) andRefineScore\(Eq\.[3](https://arxiv.org/html/2606.06840#S2.E3)\), respectively: Phase 1—sharp attention, anchor alignment; Phase 2—QK loop\-back, OV helpfulness\. Mechanistic distillation recovers both phases: students learn to focus attention on salient tokens \(high kurtosis\) and write updates that align representations with the semantic anchor, while also re\-attending to shortlisted candidates over near\-misses and writing updates that increase the target–near\-miss margin\. In contrast, baselines exhibit diffuse attention, weak anchor alignment, and poor refinement behavior\. Gains remain strong but smaller onMIMIC\-IV\-fullthan onLF\-WikiSeeAlso\-320K, likely due to noisier data, longer texts, and the complexity of clinical notes\.

Table 1:Mechanistic distillation improves reasoning quality, downstream performance, and recovers phase\-specific reasoning mechanismson theMIMIC\-IV\-fulldataset\. See Table[2](https://arxiv.org/html/2606.06840#A6.T2)in the appendix for results onLF\-WikiSeeAlso\-320K\. Results across both datasets and additional student models \(GEMMA\-3\-12B,QWEN3\-4B\) are in App\.[F](https://arxiv.org/html/2606.06840#A6), Tables[3](https://arxiv.org/html/2606.06840#A6.T3)and[4](https://arxiv.org/html/2606.06840#A6.T4)\.Model/MethodFoc↑\\uparrowConf↓\\downarrowLAS↑\\uparrowF1↑\\uparrowSharp↑\\uparrowAlign↑\\uparrowLoop↑\\uparrowHelp↑\\uparrowLLaMA3\-7B/Vanilla0\.390\.350\.690\.421\.620\.380\.270\.24LLaMA3\-7B/UniCoTT0\.480\.290\.730\.471\.850\.440\.340\.31LLaMA3\-7B/MoRSD0\.510\.270\.740\.491\.920\.460\.360\.33LLaMA3\-7B/SemCoT0\.470\.300\.720\.461\.790\.430\.320\.29LLaMA3\-7B/CWT0\.530\.260\.750\.501\.980\.470\.380\.35\\rowcolorgreen\!15LLaMA3\-7B/Mech\.0\.62▲\\blacktriangle\(\+0\.09\)0\.19▼\\blacktriangledown\(\-0\.07\)0\.76▲\\blacktriangle\(\+0\.01\)0\.52▲\\blacktriangle\(\+0\.02\)2\.18▲\\blacktriangle\(\+0\.20\)0\.55▲\\blacktriangle\(\+0\.08\)0\.46▲\\blacktriangle\(\+0\.08\)0\.41▲\\blacktriangle\(\+0\.06\)LLaMA3\-13B/Vanilla0\.450\.310\.720\.481\.750\.430\.310\.28LLaMA3\-13B/UniCoTT0\.540\.260\.750\.521\.960\.490\.380\.35LLaMA3\-13B/MoRSD0\.560\.240\.760\.542\.040\.510\.400\.37LLaMA3\-13B/SemCoT0\.530\.270\.740\.511\.910\.480\.360\.33LLaMA3\-13B/CWT0\.580\.230\.770\.552\.090\.520\.420\.39\\rowcolorgreen\!15LLaMA3\-13B/Mech\.0\.66▲\\blacktriangle\(\+0\.08\)0\.17▼\\blacktriangledown\(\-0\.06\)0\.78▲\\blacktriangle\(\+0\.01\)0\.57▲\\blacktriangle\(\+0\.02\)2\.28▲\\blacktriangle\(\+0\.19\)0\.59▲\\blacktriangle\(\+0\.07\)0\.50▲\\blacktriangle\(\+0\.08\)0\.44▲\\blacktriangle\(\+0\.05\)\\rowcolorgreen\!15Avg\. gain with mech\.𝚫\+0\.136\\boldsymbol\{\\Delta\}\{\+\}0\.136∇−0\.098\\boldsymbol\{\\nabla\}\{\-\}0\.098𝚫\+0\.033\\boldsymbol\{\\Delta\}\{\+\}0\.033𝚫\+0\.051\\boldsymbol\{\\Delta\}\{\+\}0\.051𝚫\+0\.339\\boldsymbol\{\\Delta\}\{\+\}0\.339𝚫\+0\.109\\boldsymbol\{\\Delta\}\{\+\}0\.109𝚫\+0\.126\\boldsymbol\{\\Delta\}\{\+\}0\.126𝚫\+0\.101\\boldsymbol\{\\Delta\}\{\+\}0\.101LLaMA3\-70B/Teacher0\.730\.140\.830\.642\.450\.650\.570\.52

RQ3: How do phase\-specific losses shape the reasoning dynamics?Figure[1](https://arxiv.org/html/2606.06840#S1.F1)shows coarse\-to\-fine reasoning structure inLLaMA3\-70B\(teacher\) andLLaMA3\-7B\(student\) by tracking*cumulative*metrics along the normalized CoT trajectoryt∈\[0,1\]t\\in\[0,1\]\. At each token positiont≤t∗t\\leq t^\{\*\},*Reasoning Focus*\(left panel; Eq\. \([2](https://arxiv.org/html/2606.06840#S2)\)\) measures how strongly the CoT prefix up tottaligns with the dominant semantic anchorsc​\(x\)c\(x\), capturing Phase 1 coarse semantic filtering\. Conversely, fort≥t∗t\\geq t^\{\*\},*Near\-Miss Confusion*\(right panel; Eq\. \([2](https://arxiv.org/html/2606.06840#S2)\)\) measures how strongly the evolving reasoning prefix—formed by augmenting fixed semantic anchors with progressively generated late CoT tokens—remains aligned with semantically plausible but incorrect alternatives, quantifying Phase 2 contrastive refinement\. The teacher exhibits a clear coarse\-to\-fine structure: rapid focus buildup in early positions followed by progressive reduction in confusion later\. Vanilla CoT distillation induces neither phase effectively, producing weak early focus and persistently high late\-stage confusion\. Phase\-specific ablations reveal distinct mechanistic roles—Phase\-1\-only supervision improves early focus but not late refinement, while Phase\-2\-only supervision yields limited early focus but some late discrimination\. In contrast, full mechanistic distillation \(Phase 1&2\) closely matches teacher’s trajectory onLF\-WikiSeeAlso\-320KandMIMIC\-IV\-full, showing both objectives are complementary & necessary\.

Ablation Analysis

![Refer to caption](https://arxiv.org/html/2606.06840v1/x2.png)Figure 9:Ablating phases & specific atten\./write/interaction losses in Eqs\.[5](https://arxiv.org/html/2606.06840#S3.E5),[6](https://arxiv.org/html/2606.06840#S3.E6)degrades CoT fidelity\.To better understand the contribution of our mechanistic distillation framework, we conduct an ablation study on the*LAS*metric\(wang\_scott\_2023\), which measures how faithfully the student’s generated rationale recovers the teacher’s gold answer\. Figure[9](https://arxiv.org/html/2606.06840#S4.F9)presents the results acrossMIMIC\-IV\-fullandLF\-WikiSeeAlso\-320Kdatasets\. Panel \(a\) ablates entire phases, showing that removing either Phase 1 \(coarse filtering\) or Phase 2 \(iterative refinement\) substantially degrades CoT fidelity\. Panels \(b\) and \(c\) further break down the three\-term decomposition of each phase’s loss \(attention/QK, write/OV, and their multiplicative interaction in Eqs\.[5](https://arxiv.org/html/2606.06840#S3.E5),[6](https://arxiv.org/html/2606.06840#S3.E6)\)\. These results demonstrate that all components are important, with the interaction terms consistently contributing strong complementary gains\. Overall, the full mechanistic distillation objective yields the highest teacher faithfulness, validating our phase\-specific pooled\-head distillation design\.

## 5Related Work

Recent CoT distillation improve smaller reasoning models by transferring teacher\-generated rationales, often through better rationale selection, structural supervision, or more efficient reasoning representations\. For example,zhuang\_unicott\_2025distills structured reasoning trajectories beyond linear CoT,yan\_towards\_2025improves data efficiency through rationale selection,he\_semcot\_2026replaces explicit reasoning traces with semantically aligned implicit tokens, whilexu2025softcotperforms reasoning in a continuous latent space using soft thought tokens, andchen\_skip\-thinking\_2025compresses reasoning by distilling chunk\-wise CoT execution\. Related work studies efficient reasoning supervision through rationale pruning, implicit reasoning tokens, and compressed CoT representations\(deng2023implicit;cheng2024compressed;feng\_keypoint\-based\_2024;wei2025sim;fan2025cothink;shen2025codi;kang2025c3ot\)\. However, they do not explore large space reasoning and treat reasoning traces as supervision target, focusing on*what*reasoning text to transfer\. In contrast, we study*how*reasoning is internally implemented and distill the underlying mechanisms\.

Our work is also related to mechanistic interpretability and representation\-level distillation\. Prior studies use attention decomposition to identify functional transformer components and analyze reasoning behavior\(dutta\_how\_2024;cabannes2024iteration;chen\_how\_2025\)\. Recent work further reveals specialized attention heads, including retrieval heads that enable long\-context factuality and information copying\(wu\_retrieval\_2024\), heterogeneous attention spans across heads for efficient inference\(fu\_mixture\_2025\), and retrieval heads that enable aggressive KV cache compression while preserving performance\(tang\_razorattention\_2024\)\. A comprehensive survey of attention head functions in LLMs\(zheng\_attention\_2024\)further categorizes their roles in knowledge recalling, in\-context identification, latent reasoning, and expression preparation\. We connect these directions by first characterizing reasoning in large label spaces as a two\-phase process and then distilling these phases through mechanism\-specific objectives over attention patterns and residual updates\.

## References

## Appendix ATraining Details

#### Models\.

We use publicly available open\-weight large language models as both teachers and students\. Our primary teacher model isLLaMA3\-70B, while student models include smaller open\-weight architectures such asLLaMA3\-7B\. Additional model variants used in ablations and scaling experiments are described in Sec\.[4](https://arxiv.org/html/2606.06840#S4)\.

#### Datasets\.

Experiments are conducted on publicly available large\-scale multi\-label datasets spanning clinical coding and extreme multi\-label retrieval settings\. These includeMIMIC\-IV\-full,MIMIC\-III\-full, andLF\-WikiSeeAlso\-320K\. We use the standard train/dev/test splits provided by prior work for all datasets\.

#### Training setup\.

All experiments are conducted on 8×\\timesA100–80GB GPUs using DeepSpeed ZeRO\-3 offloading for memory efficiency and distributed training\. Models are trained using mixed\-precision FP16 training with gradient accumulation\.

For mechanistic distillation, both teacher and student are run under teacher forcing on the same reasoning trajectories to enable aligned phase\-wise supervision\. Training combines the standard chain\-of\-thought distillation objective with our proposed mechanistic losses: \(i\) Phase 1 coarse semantic focusing supervision and \(ii\) Phase 2 contrastive refinement supervision\.

#### Optimization\.

We optimize all models using AdamW with cosine learning\-rate scheduling and linear warmup\. Hyperparameters are selected based on validation performance on the development split\. Gradient clipping is applied for training stability\.

#### Mechanistic supervision\.

Phase 1 supervision operates over early reasoning positions identified using the temporal reasoning trajectory described in Sec\.[2](https://arxiv.org/html/2606.06840#S2)\. Heads are ranked using𝖢𝗈𝖺𝗋𝗌𝖾𝖲𝖼𝗈𝗋𝖾\\mathsf\{CoarseScore\}and pooled to form phase\-level supervision targets\.

Phase 2 supervision operates over mid\-to\-late reasoning positions corresponding to refinement behavior\. Heads are ranked using𝖱𝖾𝖿𝗂𝗇𝖾𝖲𝖼𝗈𝗋𝖾\\mathsf\{RefineScore\}and used to supervise shortlist retrieval and margin\-widening refinement behavior\.

#### Activation patching\.

For mechanistic validation, we perform both denoising and noising activation patching over attention heads ranked by the proposed phase\-specific metrics\. Patching experiments are conducted independently from training and are used solely for causal analysis\.

#### Evaluation metrics\.

We report standard extreme multi\-label classification metrics including macro\-F1, Precision@kk, and propensity\-scored metrics where applicable\. For mechanistic evaluation, we additionally report trajectory\-based reasoning metrics including*Focus*and*Near\-Miss Confusion*\.

#### Statistical significance\.

Following prior work, we evaluate statistical significance using paired Wilcoxon signed\-rank testing\. Additional details are provided in App\.[E](https://arxiv.org/html/2606.06840#A5)\.

#### Reproducibility\.

We release anonymized code and training scripts at submission time:[https://github\.com/research\-anon\-487/xcube/tree/reasoning](https://github.com/research-anon-487/xcube/tree/reasoning)\. The repository contains preprocessing scripts, training pipelines, mechanistic analysis utilities, and evaluation commands required to reproduce the primary experiments\. Trained checkpoints will be released upon acceptance\.

## Appendix BMathematical Details

### Reasoning Focus

To quantify the progressive alignment of early reasoning with the core clinical signals, we define the reasoning focus metric\. For each inputxx, we first precompute a compact broad\-category representationc​\(x\)c\(x\)that captures the dominant high\-level semantic signals \(e\.g\., systolic heart failure, volume overload, cardiorenal syndrome\)\. This is obtained in one of two ways:

When ground\-truth labelsL​\(x\)L\(x\)are available,c​\(x\)c\(x\)is the centroid of the embeddings of the true labels \(or their ontological ancestors\):

c​\(x\)=1\|L​\(x\)\|​∑ℓ∈L​\(x\)e​\(ℓ\)\.c\(x\)=\\frac\{1\}\{\|L\(x\)\|\}\\sum\_\{\\ell\\in L\(x\)\}e\(\\ell\)\.When ground\-truth labels are unavailable,c​\(x\)=e​\(LLMsummary​\(x\)\)c\(x\)=e\\big\(\\textsc\{LLM\}\_\{\\text\{summary\}\}\(x\)\\big\), where a frozen auxiliary LLM extracts a short list of key clinical concepts fromxx\.

Letr1,r2,…,rKr\_\{1\},r\_\{2\},\\dots,r\_\{K\}denote the tokens of the early CoT segment\. We then define the focus after the firstkktokens as the average cosine similarity with the broad\-category vector:

𝖥𝗈𝖼𝗎𝗌​\(k\)=1k​∑q=1kcos⁡\(e​\(rq\),c​\(x\)\)\.\\mathsf\{Focus\}\(k\)=\\frac\{1\}\{k\}\\sum\_\{q=1\}^\{k\}\\cos\\\!\\big\(e\(r\_\{q\}\),\\;c\(x\)\\big\)\.
Intuition:This construction captures progressive semantic alignment: early CoT tokens are initially diffuse and may reflect surface\-level cues or partial hypotheses, yielding weaker alignment with the dominant signalc​\(x\)c\(x\)\. As reasoning unfolds, tokens increasingly consolidate around the core clinical concepts, causing the cumulative average to rise\. Strong models therefore exhibit a clear increasing trajectory in𝖥𝗈𝖼𝗎𝗌​\(k\)\\mathsf\{Focus\}\(k\)askkgrows, reflecting rapid convergence toward the correct semantic anchors, while weaker or distilled models show slower growth or lower saturation, indicating less coherent early\-stage reasoning\.

### Near\-Miss Confusion

To quantify how effectively late\-stage contrastive reasoning suppresses near\-miss alternatives, we define the near\-miss confusion metric\. For each inputxx, we first extract a fixed set of broad semantic anchors𝒜​\(x\)\\mathcal\{A\}\(x\)\(e\.g\., “systolic heart failure”, “volume overload”, “cardiorenal syndrome”, “EF 35%”, etc\.\)\. We then create a short fixed anchor texta​\(x\)a\(x\)that explicitly lists or summarizes these anchors \(this is the same textual representation used for the Focus metric in early reasoning\)\.

Then, letn​\(x\)n\(x\)be the centroid of the near\-miss categories𝒩​\(x\)\\mathcal\{N\}\(x\)\(e\.g\., HFpEF, pneumonia, primary renal failure, isolated AKI, etc\.\):

n​\(x\)=1\|𝒩​\(x\)\|​∑m∈𝒩​\(x\)𝐞​\(m\)\.n\(x\)=\\frac\{1\}\{\|\\mathcal\{N\}\(x\)\|\}\\sum\_\{m\\in\\mathcal\{N\}\(x\)\}\\mathbf\{e\}\(m\)\.Letr1,…,rKr\_\{1\},\\dots,r\_\{K\}be the tokens of the late \(contrastive/refinement\) CoT segment\. For eachk=1,…,Kk=1,\\dots,K, we form the cumulative prefix

sk=a​\(x\)∥\[r1​…​rk\],s\_\{k\}\\;=\\;a\(x\)\\;\\\|\\;\[r\_\{1\}\\dots r\_\{k\}\],i\.e\., we append the progressive late CoT tokens to the fixed anchor text\. The near\-miss confusion afterkklate tokens is

𝖢𝗈𝗇𝖿𝗎𝗌𝗂𝗈𝗇​\(k\)=1k​∑q=1kcos⁡\(e​\(sq\),n​\(x\)\)\.\\mathsf\{Confusion\}\(k\)=\\frac\{1\}\{k\}\\sum\_\{q=1\}^\{k\}\\cos\\\!\\big\(e\(s\_\{q\}\),\\;n\(x\)\\big\)\.
Intuition:This cumulative construction is intentional: the anchor texta​\(x\)a\(x\)alone captures only broad categories and therefore still has non\-negligible similarity to many near\-misses\. Only when the late contrastive tokens \(explicitly ruling out competitors\) are appended does the full prefixsks\_\{k\}pull away fromn​\(x\)n\(x\)\. Strong models therefore show a clear decreasing trajectory in𝖢𝗈𝗇𝖿𝗎𝗌𝗂𝗈𝗇​\(k\)\\mathsf\{Confusion\}\(k\)versuskk, while distilled models exhibit flatter or slower decay\.

### Computingavg​\_​logittarget⁡\(q\)\\operatorname\{avg\\\_logit\}\_\{\\text\{target\}\}\(q\)

At a refinement token positionqq, we compute the average target logit as follows\.

First, run the model’s forward pass up to positionqqand obtain the final residual stream representationhfinal​\(q\)h\_\{\\text\{final\}\}\(q\)\. The vocabulary logits at this position are computed by applying the model’s unembedding matrix:

logits​\(q\)=Wunembed​hfinal​\(q\)\+b,\\text\{logits\}\(q\)=W\_\{\\text\{unembed\}\}\\,h\_\{\\text\{final\}\}\(q\)\+b,
wherelogits​\(q\)∈ℝ\|𝒱\|\\text\{logits\}\(q\)\\in\\mathbb\{R\}^\{\|\\mathcal\{V\}\|\}contains scores for all vocabulary tokens\.

Next, identify the target labels in the shortlist\. These correspond to the ground\-truth positive labels \(or the correct shortlist subcategories identified by the model\), such as “acute pancreatitis” or “NBA playoff overtime rules”\. Suppose there arekksuch targets with token IDs

label\_idstarget=\[i​d1,i​d2,…,i​dk\]\.\\text\{label\\\_ids\}\_\{\\text\{target\}\}=\[id\_\{1\},id\_\{2\},\\dots,id\_\{k\}\]\.
The average target logit at positionqqis then computed as

avg​\_​logittarget⁡\(q\)=1k​∑i=1klogits​\(q\)​\[label\_idstarget​\[i\]\]\.\\operatorname\{avg\\\_logit\}\_\{\\text\{target\}\}\(q\)=\\frac\{1\}\{k\}\\sum\_\{i=1\}^\{k\}\\text\{logits\}\(q\)\[\\text\{label\\\_ids\}\_\{\\text\{target\}\}\[i\]\]\.
This quantity represents the mean logit value among the correct labels at positionqq\.

We compute an analogous quantity for the near\-miss labels \(e\.g\., “chronic pancreatitis”, “regular season tiebreaker”\)\. If their token IDs are

label\_idsnear\_miss=\[i​d1,i​d2,…,i​dm\],\\text\{label\\\_ids\}\_\{\\text\{near\\\_miss\}\}=\[id\_\{1\},id\_\{2\},\\dots,id\_\{m\}\],
then

avg​\_​logitnear\_miss⁡\(q\)=1m​∑j=1mlogits​\(q\)​\[label\_idsnear\_miss​\[j\]\]\.\\operatorname\{avg\\\_logit\}\_\{\\text\{near\\\_miss\}\}\(q\)=\\frac\{1\}\{m\}\\sum\_\{j=1\}^\{m\}\\text\{logits\}\(q\)\[\\text\{label\\\_ids\}\_\{\\text\{near\\\_miss\}\}\[j\]\]\.
Finally, we define the margin at positionqqas the scalar difference

margin​\(q\)=avg​\_​logittarget⁡\(q\)−avg​\_​logitnear\_miss⁡\(q\)\.\\textsf\{margin\}\(q\)=\\operatorname\{avg\\\_logit\}\_\{\\text\{target\}\}\(q\)\-\\operatorname\{avg\\\_logit\}\_\{\\text\{near\\\_miss\}\}\(q\)\.\(7\)
This quantity measures how strongly the model favors the correct shortlist labels over the competing near\-miss labels at that refinement step\. Lower/negative margin means the model still finds near\-misses plausible; higher/positive margin means refinement has widened the gap in favor of the target\.

### Stable Computation of Pooled Kurtosis\.

The Phase 1interaction term\(cf\. Eq\.[5](https://arxiv.org/html/2606.06840#S3.E5)\) uses the excess kurtosis of the pooled attention distribution \(Eq\.[1](https://arxiv.org/html/2606.06840#S2.E1)\)\. To ensure numerical stability, we compute kurtosis on the normalized pooled attention vector𝜶¯q,⋅\\bar\{\\boldsymbol\{\\alpha\}\}\_\{q,\\cdot\}, which is already a probability distribution produced by the softmax attention mechanism\. Let𝜶¯q,⋅=\(α1,…,αn\)\\bar\{\\boldsymbol\{\\alpha\}\}\_\{q,\\cdot\}=\(\\alpha\_\{1\},\\ldots,\\alpha\_\{n\}\)denote the pooled attention over all key positions for a queryqq\. We compute excess kurtosis using the standard moment\-based estimator

κ​\(𝜶¯q,⋅\)=1n​∑i=1n\(αi−μ\)4\(1n​∑i=1n\(αi−μ\)2\+ϵ\)2−3,μ=1n​∑i=1nαi,\\kappa\(\\bar\{\\boldsymbol\{\\alpha\}\}\_\{q,\\cdot\}\)=\\frac\{\\frac\{1\}\{n\}\\sum\_\{i=1\}^\{n\}\(\\alpha\_\{i\}\-\\mu\)^\{4\}\}\{\\left\(\\frac\{1\}\{n\}\\sum\_\{i=1\}^\{n\}\(\\alpha\_\{i\}\-\\mu\)^\{2\}\+\\epsilon\\right\)^\{2\}\}\-3,\\qquad\\mu=\\frac\{1\}\{n\}\\sum\_\{i=1\}^\{n\}\\alpha\_\{i\},whereϵ\\epsilonis a small constant \(10−610^\{\-6\}in our implementation\) added for numerical stability\.

This computation is well behaved in practice for two reasons\. First, attention vectors are normalized and bounded in\[0,1\]\[0,1\], preventing large moment magnitudes\. Second, the stabilization constant in the denominator avoids division by very small variances when the distribution approaches uniformity\. As the kurtosis term appears only in the interaction objective and the primary routing supervision is provided by the pooled\-attention KL loss \(Eq\.[5](https://arxiv.org/html/2606.06840#S3.E5)\), this statistic acts as a stable auxiliary signal preserving the coarse\-filtering structure identified in our mechanistic analysis\. Our implementation follows the standard moment\-based definition of excess kurtosis used in statistical analysis\[westfall2014kurtosis\]\.

### Phase 2 Mechanistic Distillation

Phase 2 heads refine decisions by repeatedly re\-attending to previously mentioned shortlist representations while suppressing near\-miss alternatives\(cf\. Eq\.[3](https://arxiv.org/html/2606.06840#S2.E3)\)\. To distill this loop\-back behavior, we pool query–key similarity scores using the head sets and weights defined above\. ForM∈\{T,S\}M\\\!\\in\\\!\\\{T,S\\\}we defineRM​\(q\)=∑hwhM​\(𝔼p∈𝒮​\(x\)​cos⁡\(𝐪qM,h,𝐤pM,h\)−𝔼p∈𝒩​\(x\)​cos⁡\(𝐪qM,h,𝐤pM,h\)\),h∈ℋP2M,R^\{M\}\(q\)=\\sum\_\{h\}w\_\{h\}^\{M\}\\Big\(\\mathbb\{E\}\_\{p\\in\\mathcal\{S\}\(x\)\}\\cos\(\\mathbf\{q\}^\{M,h\}\_\{q\},\\mathbf\{k\}^\{M,h\}\_\{p\}\)\-\\mathbb\{E\}\_\{p\\in\\mathcal\{N\}\(x\)\}\\cos\(\\mathbf\{q\}^\{M,h\}\_\{q\},\\mathbf\{k\}^\{M,h\}\_\{p\}\)\\Big\),\\;h\\\!\\in\\\!\\mathcal\{H\}^\{M\}\_\{\\mathrm\{P2\}\},which measures how strongly pooled Phase 2 heads re\-attend to shortlist representations while avoiding near\-miss ones\.Phase 2 heads also write residual updates that widen the margin between target and near\-miss labels via theO​VOVcircuit\(cf\. Eq\.[3](https://arxiv.org/html/2606.06840#S2.E3)\)\. To distill this behavior, we pool write vectors using the Phase 2 head sets and weights defined above:Δ¯M​\(q\)=∑hwhM​ΔM,h​\(q\),h∈ℋP2M,M∈\{T,S\}\.\\bar\{\\Delta\}^\{M\}\(q\)=\\sum\_\{h\}w\_\{h\}^\{M\}\\,\\Delta^\{M,h\}\(q\),\\;h\\\!\\in\\\!\\mathcal\{H\}^\{M\}\_\{\\mathrm\{P2\}\},\\;M\\\!\\in\\\!\\\{T,S\\\}\.We then compute the pooled OV\-helpfulness scoreHM​\(q\)=⟨Δ¯M​\(q\),∇𝐡final​\(q\)margin​\(q\)⟩,H^\{M\}\(q\)=\\Big\\langle\\bar\{\\Delta\}^\{M\}\(q\),\\;\\nabla\_\{\\mathbf\{h\}\_\{\\mathrm\{final\}\}\(q\)\}\\,\\mathrm\{margin\}\(q\)\\Big\\rangle,which measures how strongly pooled write updates increase the margin between target and near\-miss labels\.Phase 2 refinement emerges from the interaction between loop\-back retrieval and margin\-widening writes\(cf\. Eq\.[3](https://arxiv.org/html/2606.06840#S2.E3)\)\. To distill this joint behavior, we match a pooled refinement score combining the pooled QK loop\-back signal and pooled OV helpfulness\. ForM∈\{T,S\}M\\\!\\in\\\!\\\{T,S\\\}we definePRSM​\(q\)=RM​\(q\)⋅HM​\(q\),\\mathrm\{PRS\}^\{M\}\(q\)=R^\{M\}\(q\)\\,\\cdot\\,H^\{M\}\(q\),which mirrors the multiplicative structure of Eq\.[3](https://arxiv.org/html/2606.06840#S2.E3)at the pooled\-head level\.*We minimize the teacher–student discrepancies inRM​\(q\)R^\{M\}\(q\),HM​\(q\)H^\{M\}\(q\), andPRSM​\(q\)\\mathrm\{PRS\}^\{M\}\(q\)via theQK,O​VOV, andinteractionterms in Eq\.[6](https://arxiv.org/html/2606.06840#S3.E6)\.*

### Standard CoT Distillation

In Section[3](https://arxiv.org/html/2606.06840#S3)we include a standard chain\-of\-thought \(CoT\) distillation objective to transfer the teacher’s reasoning trace to the student\. Following common practice in CoT distillation, training is performed under teacher forcing on the teacher\-generated reasoning sequence\. Concretely, for an inputxx, the teacher first produces a reasoning traceτT​\(x\)=\(t1T,…,tmT\)\\tau^\{T\}\(x\)=\(t\_\{1\}^\{T\},\\dots,t\_\{m\}^\{T\}\)\. Both teacher and student are then evaluated on the shared sequencezT​\(x\)=\[prompt​\(x\);τT​\(x\)\]z^\{T\}\(x\)=\[\\mathrm\{prompt\}\(x\);\\tau^\{T\}\(x\)\], allowing token\-level distributions to be aligned at each position\.

The CoT distillation loss combines cross\-entropy on the teacher tokens with a KL divergence between the teacher and student output distributions:

ℒCoT=𝔼x​∑t=1m\[−log⁡pS​\(ttT∣z<tT\)⏟token supervision\+λKL​KL\(pT\(⋅∣z<tT\)∥pS\(⋅∣z<tT\)\)⏟distribution matching\],\\mathcal\{L\}\_\{\\mathrm\{CoT\}\}=\\mathbb\{E\}\_\{x\}\\sum\_\{t=1\}^\{m\}\\Big\[\\underbrace\{\-\\log p\_\{S\}\(t\_\{t\}^\{T\}\\mid z^\{T\}\_\{<t\}\)\}\_\{\\text\{token supervision\}\}\+\\lambda\_\{\\mathrm\{KL\}\}\\underbrace\{\\mathrm\{KL\}\\\!\\left\(p\_\{T\}\(\\cdot\\mid z^\{T\}\_\{<t\}\)\\,\\\|\\,p\_\{S\}\(\\cdot\\mid z^\{T\}\_\{<t\}\)\\right\)\}\_\{\\text\{distribution matching\}\}\\Big\],wherepTp\_\{T\}andpSp\_\{S\}denote the teacher and student token distributions, respectively\. This term provides the standard CoT supervision used in prior distillation work, while our proposed objectives focus on transferring the mechanistic behaviors of Phase 1 and Phase 2 reasoning\.

## Appendix CDataset

#### MIMIC\-IV\-full\.

MIMIC\-IV\-fullis derived from the MIMIC\-IV electronic health record \(EHR\) dataset\[johnson2023mimic\], which contains deidentified ICU patient data from 2008–2019, including structured information and free\-text clinical notes\. In the clinical coding setting, instances correspond to discharge summaries annotated with ICD\-10 codes, forming a high\-dimensional multi\-label prediction task\. The dataset comprises 122,279 discharge summaries and 7,942 unique ICD\-10 codes, exhibiting a long\-tailed and highly sparse label distribution characteristic of real\-world medical coding\.

#### LF\-WikiSeeAlso\-320K\.

LF\-WikiSeeAlso\-320Kis a large\-scale extreme multi\-label dataset constructed from Wikipedia, where each instance corresponds to an article and labels correspond to its “See Also” links, representing semantically related pages\. It contains 1,779,881 training instances and 769,421 test instances with a label space of 501,070 unique labels\. The dataset has an average of 4\.75 labels per instance and 16\.86 instances per label, reflecting moderate density alongside a long\-tailed distribution\. It models the task of related\-article recommendation, requiring selection of a small subset of relevant labels from an extremely large candidate space, and has been used in recent work such asdahiya\_prototypical\_2025\. The dataset is publicly available at[http://manikvarma\.org/downloads/XC/XMLRepository\.html](http://manikvarma.org/downloads/XC/XMLRepository.html)\.

## Appendix DEvaluation Metrics

### D\.1CoT Fidelity Metric

Followingwang\_scott\_2023, our primary metric for teacher faithfulness is*Leakage\-Adjusted Simulatability \(LAS\)*, which quantifies how well the student’s rationale helps a simulator recover the teacher’s \(gold\) answer\. This measures the degree to which the student remains faithful to the teacher’s reasoning trace\.

Formally, for a questionqq, student\-generated rationalerr, and the teacher’s gold answera∗a^\{\*\}, teacher\-faithfulness LAS is defined as

LASteacher​\(q,r,a∗\)=Acc​\(q⊕r→a∗\)−Acc​\(q→a∗\),\\text\{LAS\}\_\{\\text\{teacher\}\}\(q,r,a^\{\*\}\)=\\text\{Acc\}\(q\\oplus r\\rightarrow a^\{\*\}\)\-\\text\{Acc\}\(q\\rightarrow a^\{\*\}\),\(8\)where⊕\\oplusdenotes concatenation and Acc is the accuracy of a fine\-tuned simulator model \(typically T5\-large\)\[wang\_scott\_2023\]\.

As complementary surface\-level metrics of imitation, we report:

- •Token\-level alignment via BLEU\-4 between the teacher trace𝐜T\\mathbf\{c\}^\{T\}and student trace𝐜S\\mathbf\{c\}^\{S\}\[ramesh\_generalization\_2025\]: BLEU​\(𝐜T,𝐜S\)\.\\text\{BLEU\}\(\\mathbf\{c\}^\{T\},\\mathbf\{c\}^\{S\}\)\.\(9\)
- •Distributional alignment via perplexity of the teacher trace under the student model\[yan\_towards\_2025\]: PPL​\(𝐜T∣x\)=exp⁡\(1L​∑t=1L−log⁡pS​\(ctT∣x,𝐜<tT\)\)\.\\text\{PPL\}\(\\mathbf\{c\}^\{T\}\\mid x\)=\\exp\\left\(\\frac\{1\}\{L\}\\sum\_\{t=1\}^\{L\}\-\\log p\_\{S\}\(c^\{T\}\_\{t\}\\mid x,\\mathbf\{c\}^\{T\}\_\{<t\}\)\\right\)\.\(10\)

Higher teacher\-faithfulness LAS, higher BLEU, and lower perplexity together indicate stronger faithfulness of the student toward the teacher’s reasoning trace\.

### D\.2Task\-Specific Answer Accuracy

In addition to the fidelity metrics above, we extract the final answera′a^\{\\prime\}from the student’s generated reasoning trace and evaluate it against the gold answera∗a^\{\*\}using macro\-F1\. This metric measures the correctness of the final prediction, independent of trace\-level imitation\.

Let𝒞\\mathcal\{C\}denote the set of classes\. The macro\-F1 score is

Macro\-F1=1\|𝒞\|​∑c∈𝒞F1c,\\text\{Macro\-F1\}=\\frac\{1\}\{\|\\mathcal\{C\}\|\}\\sum\_\{c\\in\\mathcal\{C\}\}\\text\{F1\}\_\{c\},\(11\)where

F1c=2​Pc​RcPc\+Rc,Pc=TPcTPc\+FPc,Rc=TPcTPc\+FNc,\\text\{F1\}\_\{c\}=\\frac\{2P\_\{c\}R\_\{c\}\}\{P\_\{c\}\+R\_\{c\}\},\\quad P\_\{c\}=\\frac\{\\mathrm\{TP\}\_\{c\}\}\{\\mathrm\{TP\}\_\{c\}\+\\mathrm\{FP\}\_\{c\}\},\\quad R\_\{c\}=\\frac\{\\mathrm\{TP\}\_\{c\}\}\{\\mathrm\{TP\}\_\{c\}\+\\mathrm\{FN\}\_\{c\}\},andTPc,FPc,FNc\\mathrm\{TP\}\_\{c\},\\mathrm\{FP\}\_\{c\},\\mathrm\{FN\}\_\{c\}are computed over binary predictions for classccin the multi\-label setting\.

## Appendix EStatistical Significance

#### Statistical Significance via Wilcoxon Signed\-Rank Test\.

We assess statistical significance using the non\-parametric Wilcoxon Signed\-Rank Test\[demvsar2006statistical\]for comparing paired model outputs\. For metrics computed at the instance level \(e\.g\.,P@15\), we apply the test directly to the paired per\-instance scores between the base model and itsPLANT\-enhanced counterpart\. For aggregate metrics such asF1andAUC, which are reported as single values over the full test set, we first collectNNpaired scores—either from repeated evaluations \(e\.g\.,N=10N=10in1010\-fold cross\-validation\) or fromNNbootstrap resamples\. Let\{a1,a2,…,aN\}\\\{a\_\{1\},a\_\{2\},\\dots,a\_\{N\}\\\}and\{b1,b2,…,bN\}\\\{b\_\{1\},b\_\{2\},\\dots,b\_\{N\}\\\}denote the scores of the base model and thePLANT\-enhanced model, respectively\. We compute the differencedi=bi−aid\_\{i\}=b\_\{i\}\-a\_\{i\}for each pair and rank the absolute values\|di\|\|d\_\{i\}\|\(excluding zeros\), averaging ranks in the case of ties\. Each rank is assigned the sign ofdid\_\{i\}, and we compute the rank sumsW\+W^\{\+\}andW−W^\{\-\}over positive and negative differences\. The test statistic isW=min⁡\(W\+,W−\)W=\\min\(W^\{\+\},W^\{\-\}\)\.

For smallNN, statistical significance is determined using exact Wilcoxon critical values; for largerNN, we apply the normal approximation with

μ\\displaystyle\\mu=N​\(N\+1\)4,σ=N​\(N\+1\)​\(2​N\+1\)24,\\displaystyle=\\frac\{N\(N\+1\)\}\{4\},\\quad\\sigma=\\sqrt\{\\frac\{N\(N\+1\)\(2N\+1\)\}\{24\}\},z\\displaystyle z=W−μσ\.\\displaystyle=\\frac\{W\-\\mu\}\{\\sigma\}\.We reject the null hypothesis of no difference if the resultingpp\-value is less than a thresholdα\\alpha\(typically0\.050\.05\)\. In our tables, statistically significant improvements are marked using▲\\blacktriangle\. This test is readily implemented in standard libraries such asscipy\.stats\.wilcoxonin Python orwilcox\.test\(paired=TRUE\)in R\.

#### Reporting Gains with Confidence Intervals\.

We also report absolute gains along with95%95\\%confidence intervals \(CI\) using paired bootstrap resampling\. For each evaluation metric, we drawB=1000B=1000bootstrap samples from the test set and compute the differenceΔb=𝖬𝖾𝗍𝗋𝗂𝖼bPLANT−𝖬𝖾𝗍𝗋𝗂𝖼bBase\\Delta\_\{b\}=\\mathsf\{Metric\}\_\{b\}^\{\\text\{\{PLANT\}\}\}\-\\mathsf\{Metric\}\_\{b\}^\{\\text\{Base\}\}for each samplebb\. The reported gain is the meanμ^\\hat\{\\mu\}of\{Δb\}\\\{\\Delta\_\{b\}\\\}, and the CI is computed using the percentile bootstrap method by taking the 2\.5th and 97\.5th percentiles of the empirical distribution of\{Δb\}\\\{\\Delta\_\{b\}\\\}\.

We mark results as statistically significant only if the Wilcoxon signed\-rank test \(α=0\.05\\alpha\{=\}0\.05\) is passed*and*the 95% CI excludes 0\. In such cases, we annotate the score with a colored arrow:▲\\blacktrianglefor statistically significant gains and▼\\blacktriangledownfor significant drops\. If the CI includes 0, no arrow is shown\. For example,14\.714\.7▲\\blacktriangle\(\+1\.2,\[0\.6, 1\.8\]\)indicates a statistically significant gain over the base model, while70\.170\.1▼\\blacktriangledown\(\-1\.4,\[\-2\.1, \-0\.7\]\)denotes a significant drop\. In contrast,73\.873\.8\(\+0\.3,\[0\.0, 0\.6\]\)is not statistically significant and is shown without an arrow\.

#### Note on Bootstrap Resampling\.

WhenNNbootstrap resamples are used to obtain paired scores for the Wilcoxon signed\-rank test, we note that the resamples are not strictly independent due to sampling with replacement\. In such cases, we place primary emphasis on the bootstrap percentile confidence intervals for assessing the significance and magnitude of the difference; the Wilcoxon test is applied mainly for datasets where independent replicates \(e\.g\., cross\-validation folds\) are available\. Results are marked as statistically significant only when*both*the 95% CI excludes zero and the Wilcoxon test \(where applicable\) yieldsp<0\.05p<0\.05\.

## Appendix FAdditional Results

In Table[2](https://arxiv.org/html/2606.06840#A6.T2)we present the compressed results on theLF\-WikiSeeAlso\-320Kdataset, combining both performance metrics and phase\-specific mechanistic alignment scores forLLaMA3\-7BandLLaMA3\-13Bstudents\. See main paper Table[1](https://arxiv.org/html/2606.06840#S4.T1)for the correspondingMIMIC\-IV\-fullresults\.

Table 2:Mechanistic distillation improves reasoning quality and downstream performance onLF\-WikiSeeAlso\-320K\.See main paper Table[1](https://arxiv.org/html/2606.06840#S4.T1)forMIMIC\-IV\-fullresults\.Model/MethodFocus↑\\uparrowConfusion↓\\downarrowLAS↑\\uparrowmacro\-F1↑\\uparrowSharpness↑\\uparrowAnchorAlign↑\\uparrowLoopBack↑\\uparrowHelpfulness↑\\uparrowLLaMA3\-7B/Vanilla0\.410\.320\.710\.471\.840\.420\.310\.28\\rowcolorgreen\!15LLaMA3\-7B/Mech\.0\.66▲\\blacktriangle*\(\+0\.25\)*0\.16▼\\blacktriangledown\(\-0\.16\)0\.79▲\\blacktriangle*\(\+0\.08\)*0\.57▲\\blacktriangle*\(\+0\.10\)*2\.47▲\\blacktriangle*\(\+0\.63\)*0\.61▲\\blacktriangle*\(\+0\.19\)*0\.52▲\\blacktriangle*\(\+0\.21\)*0\.46▲\\blacktriangle*\(\+0\.18\)*LLaMA3\-13B/Vanilla0\.480\.290\.740\.511\.960\.470\.350\.32\\rowcolorgreen\!15LLaMA3\-13B/Mech\.0\.70▲\\blacktriangle*\(\+0\.22\)*0\.14▼\\blacktriangledown\(\-0\.15\)0\.81▲\\blacktriangle*\(\+0\.07\)*0\.60▲\\blacktriangle*\(\+0\.09\)*2\.55▲\\blacktriangle*\(\+0\.59\)*0\.65▲\\blacktriangle*\(\+0\.18\)*0\.56▲\\blacktriangle*\(\+0\.21\)*0\.49▲\\blacktriangle*\(\+0\.17\)*\\rowcolorgreen\!15Avg\. gain with mech\.𝚫\+0\.235\\boldsymbol\{\\Delta\}\{\+\}0\.235∇−0\.155\\boldsymbol\{\\nabla\}\{\-\}0\.155𝚫\+0\.075\\boldsymbol\{\\Delta\}\{\+\}0\.075𝚫\+0\.095\\boldsymbol\{\\Delta\}\{\+\}0\.095𝚫\+0\.61\\boldsymbol\{\\Delta\}\{\+\}0\.61𝚫\+0\.19\\boldsymbol\{\\Delta\}\{\+\}0\.19𝚫\+0\.21\\boldsymbol\{\\Delta\}\{\+\}0\.21𝚫\+0\.18\\boldsymbol\{\\Delta\}\{\+\}0\.18LLaMA3\-70B/Teacher0\.750\.120\.850\.672\.710\.720\.630\.58

In Table[3](https://arxiv.org/html/2606.06840#A6.T3)we present additional distillation results using two different student backbones \(QWEN3\-4BandGEMMA\-3\-12B\) from the sameLLaMA3\-70Bteacher\. See main paper Table[1](https://arxiv.org/html/2606.06840#S4.T1)for results withLLaMA3\-7BandLLaMA3\-13Bstudents\.

Table 3:Mechanistic distillation with different student backbones\.Results using Qwen3\-4B and Gemma\-3\-12B as students\. See main paper Table[1](https://arxiv.org/html/2606.06840#S4.T1)for results with Llama students\.ModelFocus↑\\uparrowConfusion↓\\downarrowLAS↑\\uparrowmacro\-F1↑\\uparrowDataset:LF\-WikiSeeAlso\-320KQWEN3\-4B\(Vanilla\)0\.360\.370\.660\.41\\rowcolorgreen\!15QWEN3\-4B\(Mech\)0\.60▲\\blacktriangle*\(\+0\.24\)*0\.18▼\\blacktriangledown\(\-0\.19\)0\.77▲\\blacktriangle*\(\+0\.11\)*0\.54▲\\blacktriangle*\(\+0\.13\)*GEMMA\-3\-12B\(Vanilla\)0\.500\.270\.760\.54\\rowcolorgreen\!15GEMMA\-3\-12B\(Mech\)0\.69▲\\blacktriangle*\(\+0\.19\)*0\.14▼\\blacktriangledown\(\-0\.13\)0\.82▲\\blacktriangle*\(\+0\.06\)*0\.62▲\\blacktriangle*\(\+0\.08\)*\\rowcolorgreen\!15Avg\. gain with mech\.𝚫\+0\.215\\boldsymbol\{\\Delta\}\{\+\}0\.215∇−0\.16\\boldsymbol\{\\nabla\}\{\-\}0\.16𝚫\+0\.085\\boldsymbol\{\\Delta\}\{\+\}0\.085𝚫\+0\.105\\boldsymbol\{\\Delta\}\{\+\}0\.105LLaMA3\-70B\(Teacher\)0\.750\.120\.850\.67Dataset:MIMIC\-IV\-fullQWEN3\-4B\(Vanilla\)0\.340\.390\.640\.39\\rowcolorgreen\!15QWEN3\-4B\(Mech\)0\.58▲\\blacktriangle*\(\+0\.24\)*0\.20▼\\blacktriangledown\(\-0\.19\)0\.75▲\\blacktriangle*\(\+0\.11\)*0\.51▲\\blacktriangle*\(\+0\.12\)*GEMMA\-3\-12B\(Vanilla\)0\.460\.300\.730\.49\\rowcolorgreen\!15GEMMA\-3\-12B\(Mech\)0\.64▲\\blacktriangle*\(\+0\.18\)*0\.17▼\\blacktriangledown\(\-0\.13\)0\.79▲\\blacktriangle*\(\+0\.06\)*0\.57▲\\blacktriangle*\(\+0\.08\)*\\rowcolorgreen\!15Avg\. gain with mech\.𝚫\+0\.21\\boldsymbol\{\\Delta\}\{\+\}0\.21∇−0\.16\\boldsymbol\{\\nabla\}\{\-\}0\.16𝚫\+0\.085\\boldsymbol\{\\Delta\}\{\+\}0\.085𝚫\+0\.10\\boldsymbol\{\\Delta\}\{\+\}0\.10LLaMA3\-70B\(Teacher\)0\.730\.140\.830\.64

In Table[4](https://arxiv.org/html/2606.06840#A6.T4)we present additional results on phase specific mechanistic fidelity using two different student backbones \(QWEN3\-4BandGEMMA\-3\-12B\) distilled from the sameLLaMA3\-70Bteacher\. See Table[1](https://arxiv.org/html/2606.06840#S4.T1)for results withLLaMA3\-7BandLLaMA3\-13Bstudents\.

Table 4:Mechanistic distillation recovers phase\-specific reasoning mechanisms with different student backbones\.Values are constituent sub\-scores ofCoarseScore\(sharpness, anchor alignment\) andRefineScore\(QK loop\-back, OV helpfulness\) – see Eqs\.[1](https://arxiv.org/html/2606.06840#S2.E1)and[3](https://arxiv.org/html/2606.06840#S2.E3)\. See main paper Table[1](https://arxiv.org/html/2606.06840#S4.T1)for results with Llama students\.Model / MethodSharpness↑\\uparrowAnchor align\.↑\\uparrowQK loop\-back↑\\uparrowOV helpfulness↑\\uparrowDataset:LF\-WikiSeeAlso\-320KQWEN3\-4B/ Vanilla1\.410\.350\.240\.22\\rowcolorgreen\!15QWEN3\-4B/ Mech\.2\.12▲\\blacktriangle*\(\+0\.71\)*0\.54▲\\blacktriangle*\(\+0\.19\)*0\.44▲\\blacktriangle*\(\+0\.20\)*0\.39▲\\blacktriangle*\(\+0\.17\)*GEMMA\-3\-12B/ Vanilla1\.920\.460\.340\.31\\rowcolorgreen\!15GEMMA\-3\-12B/ Mech\.2\.41▲\\blacktriangle*\(\+0\.49\)*0\.62▲\\blacktriangle*\(\+0\.16\)*0\.53▲\\blacktriangle*\(\+0\.19\)*0\.47▲\\blacktriangle*\(\+0\.16\)*\\rowcolorgreen\!15Avg\. gain with mech\.𝚫\+0\.60\\boldsymbol\{\\Delta\}\{\+\}0\.60𝚫\+0\.175\\boldsymbol\{\\Delta\}\{\+\}0\.175𝚫\+0\.195\\boldsymbol\{\\Delta\}\{\+\}0\.195𝚫\+0\.165\\boldsymbol\{\\Delta\}\{\+\}0\.165LLaMA3\-70B\(Teacher\)2\.710\.720\.630\.58Dataset:MIMIC\-IV\-fullQWEN3\-4B/ Vanilla1\.280\.320\.210\.19\\rowcolorgreen\!15QWEN3\-4B/ Mech\.1\.92▲\\blacktriangle*\(\+0\.64\)*0\.49▲\\blacktriangle*\(\+0\.17\)*0\.39▲\\blacktriangle*\(\+0\.18\)*0\.35▲\\blacktriangle*\(\+0\.16\)*GEMMA\-3\-12B/ Vanilla1\.680\.410\.290\.26\\rowcolorgreen\!15GEMMA\-3\-12B/ Mech\.2\.14▲\\blacktriangle*\(\+0\.46\)*0\.56▲\\blacktriangle*\(\+0\.15\)*0\.47▲\\blacktriangle*\(\+0\.18\)*0\.42▲\\blacktriangle*\(\+0\.16\)*\\rowcolorgreen\!15Avg\. gain with mech\.𝚫\+0\.55\\boldsymbol\{\\Delta\}\{\+\}0\.55𝚫\+0\.16\\boldsymbol\{\\Delta\}\{\+\}0\.16𝚫\+0\.18\\boldsymbol\{\\Delta\}\{\+\}0\.18𝚫\+0\.16\\boldsymbol\{\\Delta\}\{\+\}0\.16LLaMA3\-70B\(Teacher\)2\.450\.650\.570\.52

## Appendix GLimitations

While our mechanistic characterization and distillation framework yields consistent improvements across datasets and model scales, several limitations remain\.

#### Dependence on teacher quality\.

Our framework assumes that the teacher model exhibits strong and internally consistent reasoning behavior\. If the teacher itself produces unstable or low\-quality reasoning trajectories, the distilled student may inherit these deficiencies\. Mechanistic supervision cannot compensate for fundamentally incorrect teacher reasoning\.

#### Interpretability assumptions\.

Our characterization decomposes reasoning into two phases: coarse semantic focusing and later contrastive refinement\. Although supported by multiple mechanistic analyses \(e\.g\., activation patching, attention tracing, and temporal trajectory analysis\), these phases remain an interpretive abstraction rather than a formally identifiable ground\-truth decomposition of transformer computation\. Different interpretability methodologies may yield alternative views of the same underlying mechanisms\.

#### Restricted evaluation scope\.

We primarily evaluate reasoning over large label spaces in domains such as clinical coding and extreme multi\-label retrieval\. While these settings capture realistic “needle\-in\-a\-haystack” reasoning scenarios, additional studies across broader reasoning tasks \(e\.g\., mathematical reasoning, program synthesis, or multimodal reasoning\) would further test the generality of the proposed mechanisms\.

#### Computational overhead\.

Mechanistic analysis introduces additional computational cost beyond standard distillation, particularly for activation patching, head ranking, and phase\-specific supervision\. Although these analyses are performed offline and do not affect inference efficiency, they increase overall experimentation cost during development\.

#### Dependence on chain\-of\-thought supervision\.

Our method relies on access to teacher reasoning traces during distillation\. In settings where high\-quality reasoning traces are unavailable, restricted, or unreliable, the applicability of the proposed framework may be limited\.

#### Open\-weight model dependence\.

Our experiments are conducted using publicly available open\-weight models\. While this improves reproducibility and transparency, conclusions may not directly transfer to proprietary frontier systems whose internal architectures, training procedures, or decoding strategies are inaccessible\.

## Appendix HBroader Impacts

This work studies the mechanistic structure of reasoning in large language models and proposes a distillation framework that transfers these behaviors to smaller models\. The research has both potentially beneficial and potentially harmful implications\.

#### Positive impacts\.

Understanding how reasoning models operate over extremely large output spaces may improve the efficiency, interpretability, and accessibility of reasoning systems\. Our framework enables smaller open models to better approximate the reasoning behavior of substantially larger systems, potentially reducing computational cost and increasing accessibility for researchers and practitioners with limited resources\. Applications include large\-scale retrieval, medical coding support, recommendation systems, and scientific information organization\.

Additionally, mechanistic analyses may contribute to broader interpretability and alignment efforts by providing tools for understanding how reasoning trajectories evolve internally across transformer layers and heads\.

#### Potential risks\.

Distilled reasoning systems could also be misused in high\-stakes or safety\-critical settings where incorrect outputs may appear deceptively well\-reasoned\. Because distilled models imitate reasoning behaviors from teachers, they may inherit systematic biases, hallucinations, or unsafe reasoning patterns present in the teacher model or training data\.

Moreover, improvements in efficient reasoning over large candidate spaces could potentially be adapted for harmful applications such as large\-scale surveillance, profiling, automated misinformation organization, or malicious retrieval systems\.

#### Mitigations\.

To support responsible use, we rely exclusively on publicly available datasets and open\-weight models released under their respective licenses\. Our released artifacts are intended solely for research and reproducibility purposes\. We encourage future work on mechanistic auditing, robustness analysis, and safety evaluation of distilled reasoning systems before deployment in high\-stakes settings\.