让大模型学会跳过推理步骤直接给答案,提升效率。
Masked Distillation: Internalizing the Chain-of-Thought in Language Models

- 用教师模型指导学生模型直接预测答案,屏蔽中间推理过程。
- 在GSM8K和Countdown任务上,缩短推理路径后准确率仍保持90%以上。
- 适合追求低延迟推理的部署场景,尤其适合资源受限设备。
大型推理模型(LRMs)在生成最终答案前会输出冗长的中间推理步骤,这些步骤显著增加延迟、内存占用和部署成本,但中间步骤的正确性与最终答案无因果关系,且步骤长度并不能可靠反映问题复杂度。这引发一个关键问题:能否将这些中间计算内化到语言模型参数中,使模型能直接生成答案(或仅用极短推理链)?本文提出「掩码蒸馏」(Masked Distillation),一种知识蒸馏框架:学生模型仅根据问题预测答案,教师模型则基于问题和自身完整思维链对学生的答案进行反馈。该框架包含两种设置:(i) 自蒸馏,同一模型在思考模式下作教师,在非思考模式下作学生;(ii) 双模型蒸馏,由更大模型作为教师监督独立的小型非思考学生。我们通过调整学生被监督的中间令牌长度,实现从完全内化(仅输出答案)到无内化(输出完整推理链)之间的连续过渡。在两个推理任务(GSM8K,小学数学题;Countdown,数字谜题搜索任务)上进行了受控实验,验证了该方法的有效性。
原文摘要 · Abstract (English)
Large Reasoning Models (LRMs) produce long, explicit chains of intermediate steps before generating a final answer at inference time. These intermediate traces dominate latency, memory usage, and serving cost, even though the final answer correctness is not causally related to the trace correctness and the trace length is not a reliable indicator of the problem complexity. This raises a natural question: can the computation expressed in these intermediate tokens be internalized into the parameters of a language model, enabling it to produce answers directly (or with much shorter intermediate traces)? We introduce \textit{masked distillation}, a knowledge-distillation framework in which a student LLM is trained to predict only the solution tokens conditioned on the question, while a reasoning teacher provides feedback on the student's responses after conditioning on the question and its own CoT trace. We instantiate this framework in two settings: (i) a \textit{self-distillation} setting, in which the same model serves as the teacher in thinking mode and as the student in non-thinking mode, and (ii) a \textit{dual-model} setting, in which a larger reasoning teacher supervises a separate smaller non-thinking student over the solution tokens. By treating intermediate tokens as a scaffold which reasoning models use to fit over the solution tokens, We additionally vary the length of intermediate-token scaffolding the student is supervised on, interpolating between full internalization (the student emits only the solution) and no internalization (the student emits the full trace before the answer). We evaluate the framework through controlled experiments on two reasoning domains: GSM8K (grade-school arithmetic) and Countdown (a number-puzzle search task).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。