💬 NLP

Masked Distillation: Internalizing the Chain-of-Thought in Language Models

本文介绍了“掩码蒸馏”（masked distillation），这是一个通过训练学生模型仅预测解题 Token，同时接收基于教师模型思维链（Chain-of-Thought）的条件反馈，从而将大语言推理模型（Large Reasoning Models）的推理能力内化至学生模型参数中的框架，进而实现直接答案生成并降低推理延迟，且不牺牲准确性。

原作者： Durgesh Kalwar, Vardhan Palod, Subbarao Kambhampati

发布于 2026-07-28

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Durgesh Kalwar, Vardhan Palod, Subbarao Kambhampati

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下这样一个世界：你最喜欢的 AI 助手就像一个才华横溢但爱唠叨的学生。当你问它一个难题时，它不会直接脱口而出答案，而是会涂涂画画写满几页笔记，画出图表，并在最终写下解决方案之前，把思考的每一个步骤都说出来。这种“大声思考”的过程，在技术领域被称为“思维链”（Chain of Thought），它让 AI 模型在处理数学和逻辑谜题时变得异常聪明。然而，这里有一个代价：那些笔记页面的生成需要耗费大量的时间和能量。这就像是请一位出租车司机送你去超市，但不仅要付他开车到超市的钱，还要付他先在城市街区里绕圈子的钱，仅仅是为了让他“思考”你应该去哪个货架。科学家们长期以来一直在思考：我们能否教会 AI 在它的“脑海中”完成所有这些思考？如果我们能做到这一点，AI 就能跳过那些冗长的、唠叨的笔记，直接给出答案，从而节省大量的计算时间和算力。

这篇名为《掩码蒸馏》（Masked Distillation）的论文深入探讨了这个问题。研究人员在处理大型推理模型（LRMs）时，提出了一个问题：他们能否通过训练，让一个较小的“学生”模型内化那个喜欢大声思考的聪明“教师”模型的思考过程。他们想看看学生模型是否能学会跳过笔记直接输出最终答案，同时依然保持同样的高水平智能。他们使用了一种被称为“掩码蒸馏”的巧妙技术来测试这一点。想象一下，老师正在解开一个谜题，而学生正在旁观。老师看到了完整的题目及其详细的思考过程，但学生只看到了题目。目标是看学生是否能学习老师的秘密逻辑，并在不需要写下步骤的情况下独立解决谜题，即便在学习过程中老师拥有所有的笔记作为引导。

团队设计了两个主要的实验。在第一个实验中，他们使用同一个 AI 模型既作为教师又作为学生，只是要求它以不同的方式进行思考。在第二个实验中，他们使用一个更大、更聪明的模型作为教师，以及一个更小、更简单的模型作为学生。他们还引入了一个“脚手架”（scaffold）的概念。你可以把它想象成自行车的辅助轮。有时，学生可能需要看到老师笔记中的一小部分（即“脚手架”）才能得出正确答案，而不是尝试从零开始。他们测试了学生需要看到多少老师的思考内容才能获得最佳效果。

结果呈现出一种“惊喜”与“并非如此简单”并存的状态。在一个名为 GSM8K 的小学数学问题数据集上，学生模型表现得非常出色。它们成功实现了思考过程的内化。完全“掩码”的学生（即一个从不写下任何笔记的学生）答对率达到了 76.0%，这比其初始能力有了巨大的飞跃，而且平均仅使用了 369.6 个 token（文本单位），而教师模型则使用了 2,398.1 个 token。这实现了 6.5 倍的效率提升！

然而，在处理一个更难的数字游戏 Countdown 时，情况发生了变化。在这里，试图在脑海中完成一切的学生（完全掩码模式）表现反而变差了，准确率降至 41.7%。事实证明，对于这类特定类型的谜题，学生需要看到老师的一部分笔记才能取得成功。但有趣的是，研究人员发现了一个“甜点位”（sweet spot）。通过让学生看到老师思考过程的一小部分（约 30%，即 alpha 为 0.3），学生的准确率飙升至 86.2%。这几乎接近了教师的水平（87.3%），但使用的 token 数量仍比教师少约 1.3 倍。这就像是给了学生恰到好处的辅助轮，让他们能够完美平衡，而不必背负整辆自行车。

论文还研究了这些学生在处理新颖、未见过的谜题（分布外任务）时的表现。在数学方面，学生表现良好；但在 Countdown 谜题方面，结果则有些复杂。有趣的是，在某些情况下，接受了少量脚手架训练的学生，在面对难度增加或发生变化的谜题时，其泛化能力甚至比教师本身还要好。

一个关键发现是，教学方法至关重要。当他们尝试仅仅通过向学生展示老师的笔记并要求其模仿（这种方法称为监督微调，Supervised Fine-Tuning）来进行教学时，效果并不理想。而“掩码蒸馏”方法——它使用一种更复杂的、将学生的猜测与教师实际想法进行对比的方式——在转移“思考”技能方面要有效得多。

最后，这篇论文表明，我们不能简单地为所有问题都“抹除”掉所有的思考步骤。AI 是否能“静默思考”很大程度上取决于问题的类型以及学生模型本身对该主题已有的了解程度。对于像基础数学这样熟悉型的任务，AI 可以学会静默思考并节省大量时间。对于更难、更陌生的谜题，它可能仍然需要一点帮助——即一小部分笔记的“脚手架”——才能完成任务。研究人员总结道，这种“脚手架”是一个可以调节的强大工具，它让我们能够在 AI 的智能程度与给出答案的速度之间找到完美的平衡。

技术摘要：掩码蒸馏——将思维链内化至语言模型中

问题陈述

大型推理模型（LRMs）通过在生成最终答案之前产生长且显式的中间标记序列（思维链或 CoT），在复杂任务（数学、代码、规划）上取得了高性能。虽然这些追踪过程提高了准确性，但它们也带来了显著的成本，包括推理延迟、内存使用（KV 缓存占用）以及服务能耗，这些成本随追踪长度线性增加。

最近的研究发现，这些中间追踪过程的正确性与最终答案的正确性之间不存在因果联系，且追踪长度并不能可靠地反映问题的复杂度。这提出了一个关键问题：能否将这些中间标记中所表达的计算过程内化到语言模型的参数中，从而使其能够直接生成正确答案（或生成显著缩短的追踪），而无需在推理时显式生成推理步骤？

方法论：掩码蒸馏

作者提出了掩码蒸馏（Masked Distillation），这是一种知识蒸馏框架，旨在将“思考型”教师模型的推理过程内化到“非思考型”学生模型的参数中。

核心框架

在该框架中，学生模型被训练为仅根据输入问题预测解题标记（solution tokens）。然而，教师模型会基于问题及其自身生成的 CoT 追踪过程提供反馈。其目标是最小化学生分布（仅以输入为条件）与教师分布（以输入和中间推理过程为条件）之间的散度。

训练目标使用在策略反向 KL 散度（on-policy reverse-KL divergence）：
$L_{MD} = D_{KL}(\pi_S(\cdot | x, y_{<t}) \parallel \pi_T(\cdot | x, \text{ITs}, y_{<t}))$
其中 $\pi_S$ 是学生策略， $\pi_T$ 是教师策略， $x$ 是输入， $y_{<t}$ 是先前生成的标记， $\text{ITs}$ 是教师的中间标记。

实验设置

作者在两种设置中实例化了该框架：

自我蒸馏（Self-Distillation）： 同一个模型既作为教师（处于思考模式）又作为学生（处于非思考模式）。
双模型设置（Dual-Model Setting）： 一个更大的推理模型作为教师，监督一个独立的、较小的非思考型学生模型。

$\alpha$ -后缀脚手架（ $\alpha$ -Suffix Scaffold）

为了探索完全内化与推理成本之间的权衡，作者引入了 $\alpha$ -后缀掩码蒸馏 变体。

$\alpha = 0$ （完全掩码）： 学生模型内化整个中间追踪过程，并仅输出解题标记。
$\alpha = 1$ （非掩码）： 学生模型在输出答案前重现完整的追踪过程（等同于标准的 CoT 生成）。
$0 < \alpha < 1$ ： 学生模型内化教师追踪过程的前 $(1-\alpha)$ 部分，并在推理时输出剩余的 $\alpha$ 比例的后缀追踪以及解题标记。

这使得作者可以将中间标记视为一种“脚手架”，可以部分保留或完全内化。

核心贡献

提出的框架： 引入了掩码蒸馏，将 CoT 推理过程内化到语言模型参数中，实现直接答案生成。
脚手架分析： 定义并评估了 $\alpha$ -后缀参数，建立了一个受控的轴线，用于调节准确性与推理成本之间的权衡。
全面评估： 在两个领域进行了广泛实验：
- 数学： GSM8K（分布内）、MATH-500 和 AIME-25（分布外）。
- Countdown： 一个数字谜题搜索任务，包含目标范围和搜索深度的变化（OOD 分割）。
消融研究： 将使用反向-KL 的掩码蒸馏与在教师追踪上进行监督微调（SFT）以及使用显式“教学标记”作为脚手架进行了对比。

结果

1. 完全内化的任务依赖性

完全内化推理（ $\alpha=0$ ）的能力高度依赖于任务以及学生在预训练期间对该领域的先验暴露程度：

GSM8K（数学）： 完全内化是成功的。在自我蒸馏设置下，学生模型达到了 76.0% 的准确率（教师为 87.9%），同时减少了 6.5 倍 的推理标记。在双模型设置下，准确率从 15.69%（基准）提升至 45.26%，并实现了 5 倍的标记减少。
Countdown： 在没有脚手架的情况下，完全内化会失败。在自我蒸馏中，完全掩码的学生模型准确率降至 41.7%（低于基准模型的 56.2%）。在双模型设置中，准确率从 0% 提升至 34.1%，但仍远低于教师（87.3%）和非掩码变体（81.1%）。

2. $\alpha$ -后缀脚手架的有效性

提供少量的脚手架（即输出一部分追踪过程）可以显著恢复性能，尤其是在完全内化失效的任务上：

Countdown（自我蒸馏）： 设置 $\alpha = 0.3$ 将准确率从 41.7% 提高到 86.2%，基本匹配了教师（87.3%），同时与非掩码变体相比减少了约 1.3 倍 的推理标记。
Countdown（双模型）： $\alpha = 0.5$ 变体实现了 84.2% 的准确率，优于非掩码基准（81.1%），同时使用了更少的标记，找到了帕累托最优解。
GSM8K： 小规模脚手架（ $\alpha=0.3$ ）恢复了大部分准确率差距（83.2% vs 87.9%），且标记成本约为非掩码变体的一半。

3. 对分布外（OOD）的泛化能力

目标范围偏移（Target-Range Shift）： 使用各种脚手架范式训练的模型可以顺利迁移到目标数字范围发生变化的 OOD 任务中。
搜索深度偏移（Search-Depth Shift）： 在 Countdown 的自我蒸馏设置中，带有脚手架训练的学生（特别是 $\alpha=0.5$ ）在增加搜索深度时比完全内化模型表现出更好的泛化能力（32.8% vs 25.7%）。
双模型局限性： 在更难的数学基准测试（MATH-500, AIME-25）的双模型设置中，教师与学生之间的容量差距过大；无论是完全内化还是脚手架都无法恢复 OOD 性能，所有变体都坍塌到了基准学生的水平。

4. 训练目标

Reverse-KL vs. SFT： 论文发现，在教师追踪上进行 监督微调（SFT） 无法有效利用脚手架。在 Countdown 中，使用脚手架的 SFT 准确率仅约为 38%，而使用相同脚手架的反向-KL 掩码蒸馏达到了约 67%。
教学标记（Teaching Tokens）： 提供显式的“教学标记”作为脚手架虽然比完全掩码有所提升，但效果不如 $\alpha$ -后缀方法。

重要性与主张

论文声称，后缀脚手架（suffix scaffolding） 提供了一种可控机制，用于平衡准确性和推理效率。其最佳运行点取决于具体任务以及教师与学生之间的容量差距。

内化是可能的，但并非普适： 模型可以内化其熟悉领域（如 GSM8K）的推理过程，但在处理新型搜索任务（如 Countdown）时，如果不保留部分显式追踪，则难以实现内化。
成本-性能权衡： 少量的显式推理（例如 $\alpha=0.3$ ）可以恢复几乎所有的全量 CoT 生成性能，同时提供显著的推理成本节省（1.3 倍至 2 倍的标记减少）。
泛化： 脚手架可以提高特定 OOD 场景下的泛化能力（例如自我蒸馏中的搜索深度偏移），这表明保留部分追踪有助于模型在某些语境下比完全内化模型或完整追踪模型更好地导航复杂的搜索空间。

作者得出结论，虽然完全内化是一个极具吸引力的效率目标，但 $\alpha$ -后缀方法提供了一种实用的、可调的替代方案，能够在降低延迟的同时保持高水平的性能。