← 返回 2026-01-15

面向卓越长链式思维推理的分布对齐序列蒸馏 Distribution-Aligned Sequence Distillation for Superior Long-CoT Reasoning

Shaotian Yan, Kaiyuan Liu, Chen Shen, Bing Wang, Sinan Fan, Jun Zhang, Yue Wu, Zheng Wang, Jieping Ye 📅 2026-01-14 👍 64 2026-07-13 08:35
SFT 代码生成 数学推理 知识蒸馏 链式思维推理 长CoT

三项蒸馏改进技术将4B小模型推理能力提升至超越32B大模型水平

前置知识

序列级蒸馏 (Sequence-Level Distillation)

序列级蒸馏是一种知识蒸馏范式,其核心思想是让小模型(学生模型)通过监督微调(SFT)学习大模型(教师模型)生成的完整响应序列。与经典的逐token对数概率蒸馏(logit distillation)不同,序列级蒸馏不要求访问教师模型的词表级别logits,而是直接将教师生成的文本序列作为训练目标。其理论基础来自 Kim & Rush (2016) 的工作,通过最小化学生与教师在序列级输出分布之间的 KL 散度来实现知识转移。实践中,这种方法等价于在教师生成的数据上做标准 SFT。

本文的核心贡献就是改进这一范式,识别了现有序列级蒸馏的三个关键缺陷并提出了解决方案。理解这一基础范式是理解全文技术路线的前提。

暴露偏差 (Exposure Bias)

暴露偏差是指训练阶段和推理阶段的输入分布不一致所导致的问题。在训练时,学生模型看到的是教师生成的前缀(teacher forcing),每个位置的输入都是「正确」的历史token;但在推理时,模型必须依赖自己自回归生成的token作为下一步输入。由于训练和推理的输入来源不同,错误会在生成过程中逐步累积,导致模型输出偏离训练分布,尤其在长序列生成时问题更加严重。

这是本文提出的第三个核心问题。论文通过实验验证了暴露偏差在长CoT蒸馏中的显著影响,并提出了混合策略蒸馏来缓解这一问题。

采样温度 (Sampling Temperature)

采样温度 $T$ 是控制语言模型生成多样性的关键超参数。给定模型在位置 $t$ 的logits $z_t$,采样概率为 $p_t = \text{softmax}(z_t / T)$。当 $T < 1$ 时分布更尖锐,生成结果更确定、更集中于高概率模式;当 $T > 1$ 时分布更平坦,增加低概率token被选中的机会,从而生成更多样化的输出。在蒸馏场景中,温度直接影响教师生成数据对教师输出分布的覆盖程度。

温度调度学习是本文第一个核心创新,通过先低温度后高温度的两阶段训练策略,兼顾了学习稳定性和分布覆盖广度。

KL 散度 (Kullback-Leibler Divergence)

KL 散度衡量两个概率分布之间的差异,定义为 $D_{KL}(P \| Q) = \sum_x P(x) \log \frac{P(x)}{Q(x)}$。在知识蒸馏中,KL 散度是最常用的损失函数,用于衡量学生模型输出分布与教师模型输出分布之间的距离。最小化 KL 散度意味着让学生模型的输出尽可能接近教师模型。在序列级蒸馏中,KL 散度是在整个响应序列空间上计算的,即 $D_{KL}(p_T(y|x) \| p_S(y|x))$。

KL 散度是序列级蒸馏的理论基础,理解它有助于把握本文从分布对齐角度改进蒸馏的核心思想。

研究动机

当前社区广泛采用的长CoT蒸馏范式——在教师模型生成的数据上做 SFT——虽然简单高效,但存在三个系统性缺陷。第一,教师分布覆盖不足:现有方法通常使用随机采样加质量过滤的方式生成训练数据,这种方式往往只能覆盖教师输出分布的一小部分模式,无法充分反映教师模型的完整序列级分布,尤其容易遗漏低概率但有价值的知识。第二,教师分布与学生学习能力不对齐:SFT 只能增加目标token在每个预测位置的似然,但无法像经典 logit 蒸馏那样在每个位置精确调整整个词表上的概率分配。当教师赋予某个 token 低概率而学生赋予高概率时,SFT 会错误地进一步推高该概率,产生误导性梯度,使学生偏离教师分布。第三,暴露偏差:训练时使用教师前缀做 teacher forcing,推理时依赖自回归生成,这种不匹配导致错误累积,论文实验表明长响应的截断率随训练序列长度显著上升。

本文的目标是本文的目标是从分布对齐的角度系统性地改进序列级蒸馏范式,使学生模型不仅学习教师生成的高质量数据,更能有效学习教师的完整输出分布和泛化能力。具体目标包括:(1)更好地覆盖教师的序列级输出分布;(2)找到更有利于学生学习的目标分布;(3)缓解训练-推理不匹配的暴露偏差。最终,以 Qwen3-4B-Instruct-2507 为学生、gpt-oss-120b 为教师,仅用 448K 训练样本即达到数学、代码、科学推理等多领域 SOTA 性能。

与已有工作不同的是,现有社区工作几乎全部从 SFT 视角出发,聚焦于设计启发式规则来过滤训练数据(如正确性验证、多样性控制、长度偏好等),而忽略了蒸馏的核心本质——让学生学习教师的完整输出分布以继承其泛化能力。换言之,现有方法缺乏显式的师生交互机制来约束蒸馏过程,蒸馏的本质被低估了。本文的独特切入点是回归蒸馏的理论根基(序列级 KL 散度最小化),从三个维度(教师分布表示、学生学习目标、训练-推理对齐)识别问题并逐一提出针对性的改进方案,构建了一个完整的分布对齐序列蒸馏(DASD)管线。

核心方法

本文的方法可以类比为一个精心设计的「教学过程」:首先,让学生从教师的「标准答案」(低温度、高置信度的响应)开始学习,掌握最基本、最一致的推理模式(温度调度学习的第一阶段);然后,逐步引入教师的「多元解法」(高温度、更丰富的响应),拓宽学生的知识覆盖面(第二阶段);在选择练习题时,优先挑选那些教师很擅长但学生目前薄弱的题目(散度感知采样);最后,让学生在教师监督下独立做题,教师在学生犯错的地方接管指导(混合策略蒸馏)。从技术路线看,整个管线包含三个核心创新组件和一个完整的多阶段训练流程:问题收集 → 响应采样(含散度感知采样)→ 响应过滤 → 温度调度训练 → 混合策略蒸馏。

本文最核心的创新在于将序列级蒸馏从一个「SFT 数据工程」问题重新定义为一个「分布对齐」问题。具体来说,三个创新分别对应分布对齐的三个关键维度。温度调度学习解决了「如何更好地表示教师分布」的问题——不是用单一温度采样,而是先用低温度捕捉教师的高概率核心模式,再用高温度扩展低概率模式的覆盖。散度感知采样(DAS)解决了「什么样的教师分布最有利于学生学习」的问题——通过分析教师和学生在句子级别的概率差异,发现教师高置信度而学生低概率的「教师句」与测试正确率正相关,因此优先选择富含此类句子的训练样本。混合策略蒸馏解决了「训练和推理分布不匹配」的问题——让学生自己生成响应,在犯错处由教师接管续写,构建训练数据中师生混合的 on-policy 轨迹。这三个方法相互正交,可以组合使用,共同构成了一个完整的「分布对齐」蒸馏框架。

方法步骤详情

DASD 管线的具体实施分为以下步骤。第一步,问题收集:从四个领域收集训练问题——数学推理(约105K题,来源包括 NuminaMath-CoT 和 AoPS)、代码生成(来源 OpenCodeReasoning 数据集)、科学推理(来源 OpenScience Reasoning 数据集)和指令遵循(来源 AM-DeepSeek-R1-Distilled-1.4M)。第二步,响应采样:使用教师模型 gpt-oss-120b 在每个问题上生成多个候选响应,分别在低温度(T=0.6)和高温度(T=1.0)下采样,同时应用散度感知采样策略优先保留富含「教师句」的样本。第三步,响应过滤:执行长度过滤(丢弃超过学生模型上下文长度的响应)、结构过滤(过滤包含函数调用的响应,确保每个响应包含思考过程和最终答案)、重复内容过滤(使用正则表达式和 n-gram 匹配过滤包含重复段落的响应)。最终获得 105K 低温度和 330K 高温度的训练数据。第四步,温度调度训练:先在低温度数据上训练(学习率 5e-5 衰减至 1e-5,cosine 调度,64K 上下文长度,ZeRO-3 + Liger kernels 优化,全局 batch size 64,6 个 epoch),再在高温度数据上继续训练。第五步,混合策略蒸馏:从训练集中采样 50K 问题,让前一阶段训练好的学生模型生成响应,识别出约 15K 个截断响应,在随机位置截断后由教师模型续写,经过质量过滤后保留 12.7K 混合策略数据,再训练一个 epoch。

技术新颖性

本文的技术新颖性体现在三个方面。首先,温度调度学习将课程学习思想引入序列级蒸馏的温度选择中,这不同于以往固定温度或仅靠质量过滤的做法,通过「先易后难」的温度课程有效平衡了学习稳定性和分布覆盖度。其次,散度感知采样是首个在序列级蒸馏中系统分析师生分布差异并据此选择训练数据的方法,与经典 logit 蒸馏需要全词表 logits 不同,DAS 只需要每个 token 的边际概率(教师侧在采样时自然获取,学生侧可本地计算),实现成本极低。再次,混合策略蒸馏与已有 on-policy 方法的关键区别在于它不需要为每个学生模型重新生成所有训练数据,而是在已有 off-policy 数据基础上,只对学生偏离教师的子集做 on-policy 修正,数据和计算成本极低(仅 12.7K 样本训练一个 epoch 即可见效)。这三个方法的正交性和互补性也是一个重要的新颖发现。

DASD-4B-Thinking 的整体训练管线
Figure 2: DASD-4B-Thinking 的整体训练管线
不同温度下教师生成数据的概率分布与训练损失对比
Figure 3: 不同温度下教师生成数据的概率分布与训练损失对比
三个模型预测概率的联合对比
Figure 5: 三个模型预测概率的联合对比
四种句子类型在正确/错误答案中的位置分布对比
Figure 6: 四种句子类型在正确/错误答案中的位置分布对比
不同 token 长度下截断响应的比率
Figure 7: 不同 token 长度下截断响应的比率
训练数据的领域分布
Figure 8: 训练数据的领域分布

实验结果

DASD-4B-Thinking 在多个高难度推理基准上达到了同等规模开源模型的 SOTA 水平,甚至超越了多个更大规模(如 32B)的模型,且仅使用了 448K 训练样本,比大多数现有开源方法少一个数量级。在数学推理方面,AIME24 得分 88.5(超越 AM-thinking-v1 的 85.3 和 Qwen3-32B 的 81.4),AIME25 得分 83.3(超越 NVIDIA-OpenReasoning-Nemotron-7B 的 78.2 和 NVIDIA-Nemotron-Ultra-253B 的 72.5)。在代码生成方面,LiveCodeBench v5 得分 69.3(超越 DeepSeek-R1-0528-Qwen3-8B 的 60.5 和 Qwen3-32B 的 65.7),LCB v6 得分 67.5(超越 Qwen3-4B-Thinking-2507 的 55.2)。在科学推理方面,GPQA Diamond 得分 68.4(与 Qwen3-32B 的 68.4 持平)。消融实验清晰展示了三个训练阶段的递进效果:低温度训练将 AIME25 从 47.4% 提升至 74.0%(+26.6%),高温度训练进一步提升至 83.0%(+9.0%),混合策略蒸馏最终达到 83.3%(+0.3%)。温度调度学习在不同教师-学生对和多领域场景下均验证有效。散度感知采样在相同数据量下 consistently 优于随机采样,且在某些情况下甚至超越两倍数据量的随机采样基线。此外,将管线扩展到 MoE 学生模型 Qwen3-30B-A3B(仅用第一阶段管线,直接复用 4B 模型的数据)同样取得了竞争力结果,DASD-30B-A3B-Thinking-Preview 在 LCB v6 上达到 72.8%,平均得分 77.3%。

不同温度设置下的性能对比
Table 1: 不同温度设置下的性能对比
多领域混合训练中不同温度设置的性能对比
Table 2: 多领域混合训练中不同温度设置的性能对比
随机采样 vs. 散度感知采样的性能对比
Table 3: 随机采样 vs. 散度感知采样的性能对比
多领域中随机采样 vs. 散度感知采样的性能对比
Table 4: 多领域中随机采样 vs. 散度感知采样的性能对比
混合策略蒸馏的消融实验
Table 5: 混合策略蒸馏的消融实验
DASD-4B-Thinking 与各基线模型在五个基准上的全面对比
Table 6: DASD-4B-Thinking 与各基线模型在五个基准上的全面对比
训练阶段的累积消融实验
Table 7: 训练阶段的累积消融实验
MoE 模型上的性能对比
Table 8: MoE 模型上的性能对比
DASD-4B-Thinking 在基准测试集上的性能表现
Figure 1: DASD-4B-Thinking 在基准测试集上的性能表现
模型参数量 vs. 性能散点图(AIME25 和 LCB v5)
Figure 9: 模型参数量 vs. 性能散点图(AIME25 和 LCB v5)
查看结构化数据
任务指标本文基线提升
数学推理 (AIME24) 平均准确率 (Pass@64) 88.5 86.0 (DeepSeek-R1-Qwen3-8B) +2.5
数学推理 (AIME25) 平均准确率 (Pass@64) 83.3 79.4 (POLARIS-4B) +3.9
代码生成 (LiveCodeBench v5) 平均准确率 (Pass@64) 69.3 63.9 (NVIDIA-OpenReasoning-Nemotron-7B) +5.4
代码生成 (LiveCodeBench v6) 平均准确率 (Pass@64) 67.5 55.2 (Qwen3-4B-Thinking-2507) +12.3
科学推理 (GPQA Diamond) 平均准确率 (Pass@64) 68.4 68.4 (Qwen3-32B) 持平(用 1/8 参数量)

局限与改进

论文存在若干局限性。首先,散度感知采样目前仅在句子粒度上进行分析和应用,而更细粒度(如 token 级别或短语级别)的分布差异识别可能带来更精准的训练数据选择,但论文未探索这一方向。其次,混合策略蒸馏阶段的效果相对温和(AIME24 +0.8%, AIME25 +0.3%),论文将其归因为「已经很强的模型上仍有增益」,但也可能暗示 on-policy 修正的潜力尚未被充分挖掘。第三,评估基准全部为英文推理任务,未涉及中文或其他语言场景下的表现。第四,虽然论文声称方法具有跨架构兼容性(验证了 MoE 学生),但 MoE 实验仅使用了管线的第一阶段,完整的三阶段效果在 MoE 上尚未验证。第五,论文使用了闭源教师模型 gpt-oss-120b,虽然方法本身不依赖特定教师,但闭源教师的不可复现性限制了社区的完整复现。最后,论文未与经典的 logit 蒸馏方法做直接对比,无法量化序列级蒸馏改进相对于 logit 蒸馏的优劣。

独立分析的弱点

散度感知采样的粒度选择是一个明显的改进方向。当前以句子为单位划分「教师句」「学生句」等四类分布类型,但句子边界在数学推理中可能不是最自然的知识单元——一个数学推导步骤可能跨多个句子,而一个句子中可能包含多个推理环节。引入更细粒度(如按推理步骤或语义段落)的划分方式可能更精确地捕捉师生差异。此外,温度调度学习中的温度选择仍然依赖人工经验(0.6 和 1.0),缺少理论指导或自适应机制来确定最优温度组合,论文中也承认「实践中可根据评估表现和响应概率分布选择温度」,但这本质上仍是一个需要大量实验调优的超参数。混合策略蒸馏中教师续写的质量控制标准(「predefined quality filters」)在论文中描述模糊,其对最终效果的具体影响未被充分分析。最后,训练效率方面,64K 上下文长度的训练即使使用了 ZeRO-3 和 Liger kernels,计算成本仍然可观,论文未报告具体的 GPU 时数。

未来方向

论文作者提出了三个明确的未来方向:(1)探索分布感知的 SFT 重加权,利用教师的序列级输出概率来更精确地近似目标分布,提高蒸馏效果和数据效率;(2)进一步改进混合策略蒸馏以提升训练效率和稳定性;(3)集成知识检索和工具使用等 agentic 能力,开发更强的领域适应推理模型。基于本文成果还可以延伸以下方向:将温度调度思想与 DAS 结合设计自适应温度采样策略;探索将 DASD 管线应用于更大的学生模型(如 8B、14B)是否能带来同等比例的收益;研究散度感知采样在 RL(强化学习)训练阶段是否同样有效;以及将序列级蒸馏的分布对齐思想与 MCTS 等搜索方法结合,探索更高效的推理能力提升路径。

复现评估

本文的可复现性较好。模型权重和训练数据集均已开源在 Hugging Face 和 ModelScope 上。论文详细描述了训练配置(学习率 5e-5、cosine 调度、64K 上下文、ZeRO-3 + Liger kernels、batch size 64、6 epochs),以及各阶段的具体数据量(105K 低温度 + 330K 高温度响应,最终 448K 训练样本)。问题来源的数据集均为公开可用的开源数据集。然而,教师模型 gpt-oss-120b 为闭源模型,无法完全复现教师响应的生成过程;不过散度感知采样仅需要教师和学生的 token 级概率,社区可以使用其他开源教师模型进行替代实验。评估使用了标准化设置(温度 1.0、top-p 1.0、每题 64 次采样取平均),便于公平对比。总体而言,算力需求较高(64K 上下文训练 + 64 次采样评估),但管线设计清晰,关键参数已公开。