掩码自蒸馏：将思维链内化到语言模型中

摘要

大型推理模型在推理时会在生成最终答案之前产生冗长、显式的中间步骤链。这些中间轨迹主导了延迟、内存占用和服务成本，尽管最终答案的正确性与轨迹的正确性并无因果关系，且轨迹长度也不是问题复杂度的可靠指标。这引出一个显而易见的问题：这些中间token所表达的计算能否被内化到语言模型的参数中，使其能够以更短的中间轨迹产生答案？我们提出掩码自蒸馏（masked self-distillation），一个基于知识蒸馏的后训练框架，其中同一模型的副本分别被实例化为教师和学生，学生模型被训练以内化全部或部分中间轨迹，从而在推理时更加高效。我们改变学生被训练去内化的中间轨迹比例，在完全内化与完全不内化之间插值。我们在两个推理领域开展受控实验：数学与图着色。我们使用掩码自蒸馏框架对Qwen3-4B与8B模型进行后训练。我们的结果表明，该方法可在多种领域和模型规模上提升任务性能的同时提高推理效率。我们系统分析了后训练模型中的效率提升能否泛化到OOD（分布外）问题。我们发现掩码自蒸馏模型在域内OOD问题上泛化良好，且该训练不会使学生模型在域外问题上产生灾难性遗忘。此外，我们的消融实验表明，监督微调能训练模型产生更短的轨迹，但以牺牲泛化为代价，这凸显了掩码自蒸馏中在线（on-policy）训练的重要性。

掩码自蒸馏：将思维链内化到语言模型中：论文配图 — 图 1：掩蔽蒸馏框架概述。第 1 阶段，中间令牌收集（上）：对于每个输入问题 x∼𝒟x\sim\mathcal{D}，我们对老师 πT\pi^{T} 的响应进行采样。响应的形式为 <<think>>(ITs) <<\\backslashthink>⁣<><answer>>(STs)<<\\backslashanswer>>，我们从中提取中间标记 (ITs) 和解决方案标记 (STs)。第 2 阶段，训练（底部）：不思考的学生 πθS\pi^{S}_{\theta} 仅以 xx 为条件，而冻结的教师以 xx 以及第 1 阶段的 IT 的第一个 (1−α)(1-\alpha) 部分为条件。参数 α\alpha 控制学生在推理时被监督发出的中间痕迹的量：训练学生重现最后一个 α\alpha 部分教师的中间轨迹以及解决方案标记。通过在学生采样令牌上计算的策略反向 KL 损失 𝔼y∼P[log⁡P/Q]\mathbb{E}_{y\sim P}[\log P/Q]，将学生的下一个令牌分布 PP 与教师的分布 QQ 进行匹配；梯度仅更新学生。在 α=0\alpha=0 时，变体简化为完全掩蔽蒸馏，在 α=1\alpha=1 时变体简化为非掩蔽蒸馏。当老师和学生是同一个模型时，这就是自蒸馏设置；当老师是一个较大的思维模型，学生是一个较小的非思维模型时，这就是双模型设置。