arXiv cs.AI· Debjyoti Saha Roy·· 10 天前AI 评分22
先刻画再蒸馏:大输出空间中的机制推理
Characterize Then Distill: Mechanistic Reasoning in Large Output Spaces
AI 导读
研究揭示推理训练语言模型如何在数千至数十万个候选中零样本完成多标签选择:在 MIMIC-IV 临床编码任务上,5,651 个候选诊断码全部置于上下文时,一组小型、全局、分阶段的注意力头对每个摘要的决策既必要又充分。作者提出 MISTILL 蒸馏方法,除教师推理文本外还监督学生模型在这些决策事件上的池化注意力,在跨家族学生模型中使对比决策的因果恢复近乎翻倍。
正文 · AI 翻译
摘要:经过推理训练的语言模型可以零样本执行多标签任务,即需要从数千到数十万候选标签的集合中选出一小部分相关标签。我们探究它们是如何在机制上做到这一点的,以及该机制能否被蒸馏。我们通过将每个决策视为一个由模型自身决策边际评分的 token 级事件,使该问题变得可度量:选择标签空间粗粒度区域的 token、在该区域内选择某个标签的 token,以及输出偏离推理中先前命名的接近备选项(近似命中)的 token。归因、精确均值消融、敲入另一个示例的上下文,以及一种先对通用注意力头进行折扣的零校准,共同赋予单个注意力头因果地位。在医院出院小结的临床编码(MIMIC-IV)上,当上下文中包含全部 5,651 个候选诊断代码时,一个小的、全局的、按阶段结构化的注意力头集合,通过消融和敲入,在几乎每一份小结上都是必要且充分的;不同的注意力头家族分别关注候选区域和回看先前命名的近似命中;并且,对于推理中已确定的提及,该区域在代码出现前若干 token 就已经可以被引出,来自一个读取输入的、不相交的中间层集合。我们提出 MISTILL:与仅迁移教师推理文本的思维链蒸馏不同,它还在这些确切的决策事件上监督学生的池化注意力。在训练后发现的注意力头上读取,它在一个跨家族学生中将对比决策的因果恢复几乎翻倍,并在一个已经承载了大部分因果恢复的学生中增加了小的、种子稳定的增益,当两个目标都训练 bf16 权重时未检测到任务差异,而使用 fp32 主权重时存在任务成本。
| 评论: | 大幅修订和扩展;取代 v1。新分析(token 级决策事件、MIMIC-IV 临床编码上的注意力头级因果测试)、新蒸馏方法(MISTILL)、实验和文本;作者列表反映了本版本作者身份。58 页,6 幅图。代码:此 https URL |
| 主题: | 计算与语言(cs.CL);人工智能(cs.AI);机器学习(cs.LG) |
| 引用为: | arXiv:2606.06840 [cs.CL] |
| (或 arXiv:2606.06840v2 [cs.CL] 用于本版本) | |
| https://doi.org/10.48550/arXiv.2606.06840 arXiv 颁发的 DOI 通过 DataCite |
提交历史
来自:Debjyoti Saha Roy [查看电子邮件]
[v1]
2026年6月5日星期五 02:32:24 UTC(334 KB)
[v2]
2026年10月6日星期二 02:10:23 UTC(301 KB)
来源:arXiv cs.AI · arxiv.org