本文提出强化蒸馏(REDI)框架,通过两阶段训练利用正向和负向推理轨迹,显著提升小型语言模型的数学推理性能,Qwen-REDI-1.5B在公开数据上达到1.5B模型的最新水平。
Reinforcement Learning, Large Language Model, Reasoning, Supervised Learning, Efficiency
Shuyao Xu, Cheng Peng, Jiangxuan Long, Weidi Xu, Wei Chu, Yuan Qi
INFLY TECH (Shanghai) Co., Ltd., National University of Singapore, AI Institute of Fudan University
Generated by grok-3
Background Problem
近年来,大型推理模型(如DeepSeek-R1、OpenAI的o1)在复杂推理任务中展现了卓越能力,但小型模型由于资源限制难以通过大规模在线强化学习(RL)获得类似能力。模型蒸馏作为一种替代方法,通过从大型‘教师’模型生成的推理轨迹(Chain-of-Thought, CoT)中学习,为小型‘学生’模型提供了一种经济高效的推理能力提升路径。然而,当前蒸馏方法多采用拒绝采样,仅利用正确的推理轨迹(正向数据),而忽略了生成过程中产生的错误推理轨迹(负向数据),导致数据利用不充分。本文聚焦于解决这一关键问题:在离线蒸馏设置中,如何有效利用正向和负向推理轨迹,以最大化小型语言模型的推理性能,尤其是在开放数据集上的应用。
Method
本文提出了强化蒸馏(REDI),一个两阶段的离线训练框架,旨在通过正向和负向推理轨迹提升模型推理能力。
- 第一阶段:监督微调(SFT):使用仅包含正确推理轨迹的数据集(DSFT),通过最大化生成正确轨迹的对数似然(如公式 )对基础模型进行微调,建立推理基础并提供初始策略。
- 第二阶段:REDI目标优化:利用包含正向和负向轨迹对的数据集(DPref),通过一个无参考、非对称加权的损失函数 进一步优化模型,其中 控制负向轨迹的梯度权重(实验中 效果最佳),以平衡稳定性和性能。
- 核心创新与批判:REDI目标函数摒弃了DPO和SimPO中的KL正则化项,通过非对称加权避免训练崩溃,同时保持较高峰值性能。然而,方法对学习率和 的选择高度敏感,论文未充分探讨不同数据集或模型规模下的鲁棒性。此外,负向轨迹的选择策略(每个问题仅选一个错误轨迹)可能引入偏差,未必能全面代表错误模式,限制了负向信号的利用效果。
Experiment
实验基于OpenR1-Math-Raw数据集(排除难度较低的cn_k12子集),构建了包含78k正向轨迹(DSFT)和53k正负轨迹对(DPref)的训练数据,使用Qwen2.5-Math-1.5B作为基础模型。
- 设置:第一阶段对DSFT进行SFT(3或5轮),第二阶段对DPref应用REDI、DPO和SimPO进行1轮优化。评估在多个数学推理基准(MATH-500, AIME24, AMC23, Minerva, OlympiadBench)上进行,采用pass@1(16次采样)和pass@16指标。
- 结果:Qwen-REDI-1.5B(基于5轮SFT+REDI)在MATH-500上达到83.1%(pass@1),平均性能49.5%,超越了使用800k专有数据训练的DeepSeek-R1-Distill-Qwen-1.5B(48.6%),显示出数据效率优势。REDI在稳定性与性能之间优于DPO和SimPO,尤其在非对称加权()下避免了训练崩溃。
- 批判:实验设计较为全面,涵盖多个基准和消融研究,但数据集规模较小(仅131k样本),且负向轨迹选择策略单一,可能影响结果的泛化性。pass@16结果表明REDI未损害模型多样性,但其提升幅度有限,未能明确证明负向数据对知识广度的贡献。此外,实验未涉及跨领域泛化测试,数学推理外的应用效果存疑。结果虽符合预期,但对超参数(如学习率、)的依赖性过强,未提供充分的鲁棒性分析。
Further Thoughts
REDI框架通过非对称加权利用负向推理轨迹提供了一个有趣的视角,但其方法对超参数的敏感性可能限制其在不同模型或任务上的应用,未来可以探索自适应加权策略以提高鲁棒性。此外,负向数据的选择策略值得深入研究,例如是否可以通过聚类或多样性采样更全面地捕捉错误模式,从而进一步提升学习效果。另一个思考方向是REDI与在线RL的结合,论文提到REDI未损害模型未来在线RL的潜力,但未提供实际结合实验,未来可以在此方向上验证其是否能作为在线RL的预训练步骤,特别是在资源受限场景下。此外,REDI在数学推理外的应用(如自然语言推理或多模态任务)可能面临挑战,负向信号在不同领域的有效性需要进一步探索,可能与领域内错误模式的复杂性相关。最后,考虑到数据效率的优势,REDI或许可以与联邦学习结合,用于分布式环境中小型模型的推理能力提升,同时保护数据隐私。