只用思维链部分训练,性能损失小但效率大幅提升
Distilling the Essence: Efficient Reasoning Distillation via Sequence Truncation

- 仅对思维链部分进行知识蒸馏,避开冗长的提示和答案
- 保留50%序列长度即可维持91%的数学任务性能
- 适合资源受限场景下的高效模型压缩
将大型推理模型(LRM)的能力蒸馏到小型学生模型通常需要大量推理数据。然而,在包含提示(P)、思维链(CoT)和答案(A)的长序列上进行知识蒸馏(KD)会带来巨大的计算开销。本文研究了在不同部分(P、CoT、A)分配监督对学生产生的影响。分析表明,当提示和答案信息已包含在思维链中时,仅对思维链令牌进行选择性蒸馏仍具有效性。基于此,我们建立了一种截断协议,量化序列长度与计算效率之间的权衡。观察发现,超过某一特定长度后,更长的训练序列对下游性能提升微乎其微,但内存和浮点运算量显著增加。因此,仅使用每个训练序列前50%的令牌进行训练,可在数学基准测试上平均保留约91%的完整序列性能,同时将训练时间、内存占用和浮点运算量均降低约50%。代码已开源。
原文摘要 · Abstract (English)
Distilling the capabilities from a large reasoning model (LRM) to a smaller student model often involves training on substantial amounts of reasoning data. However, knowledge distillation (KD) over lengthy sequences with prompt (P), chain-of-thought (CoT), and answer (A) sections makes the process computationally expensive. In this work, we investigate how the allocation of supervision across different sections (P, CoT, A) affects student performance. Our analysis shows that selective KD over only the CoT tokens can be effective when the prompt and answer information is encompassed by it. Building on this insight, we establish a truncation protocol to quantify computation-quality tradeoffs as a function of sequence length. We observe that beyond a specific length, longer training sequences provide marginal returns for downstream performance but require substantially higher memory and FLOPs. To this end, training on only the first $50\%$ of tokens of every training sequence can retain, on average, $\approx91\%$ of full-sequence performance on math benchmarks while reducing training time, memory usage, and FLOPs by about $50\%$ each. Codes are available at https://github.com/weiruichen01/distilling-the-essence.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。