先刻画后蒸馏：大规模输出空间中的机制推理

# 大型输出空间中的机制推理
来源：https://arxiv.org/html/2606.06840
## 先表征再蒸馏：大型输出空间中的机制推理

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

###### 摘要

现代推理模型在极具挑战性的多标签任务上表现出令人惊讶的强大零样本性能，这些任务需要从数十万到数百万个候选标签中选出少量相关选项。我们研究它们实现这一目标的机制原理。我们将推理表征为两个阶段：先广泛“筛选”候选，然后对得到的集合进行细粒度推理。我们在多个数据集上提供证据，表明这两步可被分离且功能互补。基于这一表征，我们开发出一种*机制蒸馏*策略，其效果始终优于标准蒸馏。我们开源 (https://github.com/research-anon-487/xcube/tree/reasoning) 代码和训练好的模型。

## 1 引言

参见图注

图1：LLaMA3-70B 在思维链早期表现出*聚焦*积累，后期出现*混淆*减少。机制蒸馏能够恢复教师轨迹，而代表性思维链蒸馏 (he_semcot_2025) 则失败。在超大规模多标签场景中，需要从数十万到数百万个候选标签中选出少量相关标签 (zhou_quest_2024; zhang_elmo_2025; ortego_large_2025)。这类场景出现在诸如 Wikipedia 类别标注 (gupta_dual-encoders_2024; dahiya_prototypical_2025)、大规模电子商务分类 (ortego_large_2025) 以及科学资料库中的文档标注 (tabatabaei_can_2025) 等应用中。这类任务因存在大量样本稀少的“长尾”标签而极具挑战性。

此问题类似于长上下文评估中的“大海捞针”挑战，模型需要从大量干扰上下文中检索或推理出稀疏的关键信息 (kamradt2023needle; hsieh_ruler_2024)。这种相似性表明，具有挑战性的多标签任务可能受益于近期在这些问题上取得成功的有效推理策略。最近的推理模型——如 OpenAI 的 o 系列 (OpenAI2025a; OpenAI2025b) 和 DeepSeek-R1 (guo2025deepseek)——旨在模仿*系统 2* 推理 (zhang_llm_2024; ji2025test; besta_reasoning_2025; xu_towards_2025; li_system_2025)。推理通过利用诸如长思维链 (chen_towards_2025; yeo_demystifying_2025)、思维树 (yao_tree_2023) 和审慎推理 (guo2025deepseek; guan_deliberative_2025) 等结构化机制，在长上下文任务上取得了强大性能 (modarressi_nolima_2025; wang_reasoning_2025; kuratov_babilong_2024; marjanovic_deepseek-r1_2026)。

我们提问：大型推理模型如何在没有任何任务特定训练的情况下，从数百万个候选标签中选出少量相关标签？以往关于机制可解释性的研究 (bereska_mechanistic_2024; hanna_have_2024; geiger_causal_2025; du_how_2025) 已尝试表征算术 (nanda_progress_2023; stolfo_mechanistic_2023; chen_how_2025)、常识问答 (chuang_faithlm_2025; basu_mechanistic_2025)、本体推理 (dutta_how_2024) 和代码生成 (michaud_opening_2024) 等推理任务。但据我们所知，针对很大输出空间的推理机制尚未被探索。

我们寻求表征大型推理模型使用的隐式机制，然后将其系统地迁移到较小模型上。一个更简单的方法是为此使用标准思维链蒸馏，这在其他任务上效果良好 (wang_scott_2023; feng_keypoint-based_2024; wadhwa-etal-2024-investigating; chen_skip-thinking_2025; chen_unveiling_2025; kang_distilling_2025; li_llms_2025)。这些方法旨在模仿较大模型的推理理由，而非产生它们的潜在推理轨迹 (dai_beyond_2024; hu_han_2026; tan_probing_2026; chen_molecular_2026)。我们观察到，对于这些任务，较小的蒸馏模型表现相对较差（性能下降 35-38%；图1 (https://arxiv.org/html/2606.06840#S1.F1)）。我们报告称，机制蒸馏可以带来更强的学生模型性能。

提示：68岁男性，进行性呼吸困难、端坐呼吸、下肢水肿。超声显示射血分数35%，符合心力衰竭……轻度肺淤血；肌酐升高提示心肾综合征。  
任务：你是经验丰富的医学编码员。从提供的大型集合（10,000+ 标签）中分配适当的ICD编码。  
I50.23 — 慢性收缩性心衰急性发作  
I50.2 — 收缩性（充血性）心衰  
R60.0 — 局限性水肿  
J18.9 — 肺炎  
I42.9 — 心肌病  
I13.0 — 高血压性心脏病合并CKD伴心衰  
N17.9 — 急性肾衰竭  
J44.9 — 慢阻肺  
… +9000个无关编码  
68岁男性……心肾综合征 + 10000个编码。  
患者有明显心衰+容量超负荷 → 进行性呼吸困难、端坐呼吸、下肢水肿  
早期思维链 ⇓  
广泛关注区域：心衰、容量超负荷、呼吸困难、端坐呼吸、下肢水肿  
68岁男性……患者有心衰、容量超负荷……水肿  
超声EF 35% → 收缩功能障碍，射血分数降低的心衰  
广泛关注区域：心衰、容量超负荷……下肢水肿、收缩功能障碍、射血分数降低的心衰  
68岁男性……EF 35% → 射血分数降低的心衰  
轻度肺淤血+肌酐↑ → 失代偿性射血分数降低心衰中的心肾综合征，需谨慎容量管理  
广泛关注区域：心衰、容量超负荷……射血分数降低心衰……心肾综合征、失代偿性射血分数降低心衰  
68岁男性……心衰、容量超负荷……水肿  
非肺炎/慢阻肺急性加重 → 真性心源性容量超负荷来自收缩性心衰  
后期思维链 ⇓  
对比推理：已排除：肺炎、慢阻肺 ×  
真实诊断：心衰容量超负荷 ✓  
……心衰、容量超负荷……EF 35% → 射血分数降低的心衰  
非单纯心肌病或冠心病 → EF 35% 锁定为射血分数降低的心衰  
对比推理：已排除：心肌病、冠心病 ×  
真实诊断：收缩性射血分数降低心衰 ✓  
……射血分数降低心衰……  
非由低血容量或肾毒性药物引起的孤立急性肾损伤 → 真性1型心肾综合征，背景为失代偿性射血分数降低心衰  
对比推理：已排除：孤立急性肾损伤、单纯高血压性心脏病 ×  
真实诊断：心肾综合征 ✓

图2：LLM 推理。（*左*）早期思维链逐步将关注集中在广泛类别上。（*右*）后期思维链逐步排除近似匹配的类别，直至仅留下真实信号。我们沿着互补维度评估教师模型及其蒸馏学生模型的性能（正式定义见第2节 (https://arxiv.org/html/2606.06840#S2)）：(1) *推理聚焦* 衡量早期思维链与输入中主导语义信号的对齐程度，捕捉将证据逐步整合为连贯锚点的过程。(2) *近似匹配混淆* 衡量演进的推理与看似合理但错误的替代方案保持对齐的程度，捕捉通过对比提炼排除它们的有效性。

作为说明，考虑图2 (https://arxiv.org/html/2606.06840#S1.F2)中的临床示例，该示例贯穿全文；任务是从大标签空间中分配正确的ICD¹¹标准化医学标签系统，包含10⁴–10⁵个编码，用作大型输出任务的基准 (liu_lost_2023)。编码 (who2025icd)。LLM表现出*聚焦式推理*：在早期步骤（左图），它迅速将证据组织成连贯的语义锚点（例如，将症状归因于心衰，然后细化至收缩功能障碍），有效缩小了搜索空间。在后期步骤（右图），它进行*对比提炼*，明确排除紧密相关的备选方案（近似匹配），如竞争诊断。这同时产生了高*聚焦*和低*混淆*。相比之下，较小的蒸馏模型未能维持这种结构：它们的推理在不同松散相关的概念间漂移，缺乏强有力的早期锚定，且未能将真实标签与合理的干扰项充分分离。

我们的**主要贡献**如下：(1) 我们为大型推理模型在具有挑战性的多标签任务中的两阶段过程提供了表征和实证的机制证据——先是*粗粒度语义过滤*，然后是对*少量候选的细粒度推理*——并表明这两个阶段是互补且可分离的。(2) 利用这一表征，我们引入了一种*机制蒸馏*方法，该方法直接监督聚合注意力头层面上的阶段特定计算，并且始终优于标准思维链蒸馏。

## 2 推理模型如何进行大规模多标签分类？

我们提出LLM在大空间中推理的两阶段表征，并提供机制证据。

提示：68岁男性，进行性呼吸困难、端坐呼吸、下肢水肿。超声显示射血分数35%，符合心力衰竭……轻度肺淤血；肌酐升高提示心肾综合征。  
任务：从大型集合（10000+ 标签）预测ICD-10编码。  
I50.22 — 慢性收缩性心衰  
J18.9 — 肺炎  
I42.9 — 心肌病  
N18.3 — 慢性肾病（3期）  
I25.10 — 冠状动脉疾病  
N17.9 — 急性肾衰竭  
J44.9 — 慢阻肺  
… +9000个无关编码  
语义锚点（来自文本+候选标签）  
呼吸困难 端坐呼吸 EF 35% 心力衰竭 心肾  
早期思维链推理  
典型心衰体征：呼吸困难、端坐呼吸、水肿。超声：EF35% → 收缩功能障碍。符合慢性收缩性心衰（I50.22）。肌酐升高 → 心肾综合征。  
早期即识别出强语义信号  
心衰 呼吸困难 EF 35% 端坐呼吸  
机制问题——阶段1  
当一个*早期思维链词元*锁定一个强语义信号时，注意力头在内部做了什么？  
早期思维链词元：“心衰”  
h h  
呼吸困难 EF 35% 端坐呼吸  
它关注了什么？  
对显著词元的锐利聚焦  
h h  
残差写入  
朝向语义锚点  
心衰 收缩功能障碍  
它写入了什么？  
残差流中的语义对齐更新

图3：阶段1——在生成早期思维链时，我们测量注意力头是否关注显著证据并写入语义对齐的更新，作为CoarseScore (CoarseScore公式1 (https://arxiv.org/html/2606.06840#S2.E1))。阶段1：粗粒度语义过滤

图3 (https://arxiv.org/html/2606.06840#S2.F3) 提供了一个建立阶段1直觉的示例。*左上方面板*显示输入提示和候选标签空间，从中涌现出语义锚点。*左下面板*显示早期思维链，其中识别出强语义信号（例如，心衰，EF 35%）。*右面板*提出关键问题：*当这些早期思维链词元被生成时，注意力头在内部做了什么？*具体来说，头部是否（i）聚焦于相关证据，以及（ii）写入与语义信号对齐的更新？这种设置可泛化到大型输出空间任务。阶段1通过一个*粗粒度分数*捕捉这一直觉，量化这些属性，如下详述。

**早期层注意力头通过将显著的输入词元与广泛的语义锚点链接来驱动粗粒度过滤。** 为提供证据，我们构建具有受控语义信号的提示，并通过机制干预分析思维链期间的注意力模式和头部贡献，受先前研究的启发，这些研究表明一小部分注意力头驱动了大部分长上下文计算，而MLP层作用有限 (wu_retrieval_2024; zheng_attention_2024; tang_razorattention_2024; fu_mixture_2025)。我们设计*干净*提示，包含具有强粗粒度信号的文本片段，后接一个长通多领域标签列表。例如（图3 (https://arxiv.org/html/2606.06840#S2.F3)，*左上*），一份心衰出院小结与相关ICD编码（心血管/肺部）及无关领域配对。相应地，早期思维链与输入中的主导语义信号对齐（即心血管概念；*左下*）。

对于每个头部，我们在干净提示上计算一个*粗粒度过滤分数*（公式1 (https://arxiv.org/html/2606.06840#S2.E1)），将分析限制在早期思维链片段的查询位置。这里Qearly(x) 表示早期思维链查询位置（前10–30个生成词元）；αₕₕ,h 表示头部 h 在查询位置 q 的注意力分布；κ(αₕₕ,h) 是其超额峰度，用于识别具有尖峰注意力模式的头部。Δᵠ(q) 表示头部 h 在位置 q 写入的残差流更新。设A(x) 为相关语义锚点的集合（例如，“心衰”或“肺水肿”等心血管疾病），其嵌入为 e(y)。则锚点质心定义为 canchor(x) ≜ mean_{y∈A(x)} e(y)。

公式1（见原文）

高CoarseScore值识别出那些*关注语义锚点*且*将残差表示与该锚点对齐*的头部。如图4 (https://arxiv.org/html/2606.06840#S2.F4)（*左图*）所示，LLaMA3-70B模型中的早期层头部（例如，L3H22、L2H7和L1H12）在早期思维链生成期间主导了排名。此外，如图4 (https://arxiv.org/html/2606.06840#S2.F4)（*右图*）所示，当从恰好位于早期思维链词元之前的查询位置查询一个排名靠前的早期层头部时，其注意力权重尖锐地集中在提示中的少量语义锚点上，如“心衰”、“EF 35%”和“既往MI”（心脏锚点），以及相关的呼吸和肾脏锚点，如“肺水肿”和“肺炎”。这表明早期层头部能迅速检测并聚焦于相关概念。然后我们提问：*这些尖峰头部是否因果地实现了粗粒度过滤？*

参见图注
参见图注

图4：（*左图*）LLaMA3-70B模型中按CoarseScore排名的顶部头部。在早期思维链期间展现锚点聚焦注意力和对齐残差更新的早期层头部（L1-6）。（*右图*）来自顶级早期层头部（L3H22）的注意力，尖锐地聚焦于关键临床锚点（红色：心脏短语如“心衰”和“EF 35%”；橙色：呼吸/肾脏短语如“肺水肿”和“肺炎”）。
参见图注
参见图注

图5：早期层注意力头因果地控制粗粒度过滤。（*左图*）逐步去噪修补按CoarseScore排名靠前的头部的大块区间，显著恢复了向语义锚点的推理聚焦。（*右图*）向这些头部添加噪声会严重降低聚焦。为了回答这个问题，我们评估了因果充分性/必要性。