量化模型走到哪里,就在哪里训练:面向低比特推理的在策略蒸馏
Train Where the Quantized Model Goes: On-Policy Distillation for Low-Bit Reasoning
精选理由
提出面向量化模型自身轨迹的蒸馏方法,显著恢复低比特模型的数学与代码推理能力。
AI 摘要
量化感知蒸馏(QAD)恢复了低于 3 比特量化所损失的大部分短答案问答性能,但数学推理和代码推理仍受到严重影响。较长的生成往往会退化为重复循环,在未完成解答前耗尽解码预算。我们将这一差距归因于量化放大的暴露偏差:QAD 在固定的语料前缀上进行训练,而量化引起的偏差会沿着模型自身的自回归轨迹不断累积。为解决这一失配问题,我们引入在策略蒸馏(OPD)阶段,将教师监督置于量化模型实际到达的位置。从一个 QAD 检查点出发,学生模型通过部署时所使用的量化前向路径进行生成,并由冻结的全精度教师模型针对其自身前缀提供反馈,将密集的词元级指导与任务验证器奖励相结合。在四个模型、2.79 和 1.88 有效比特的设置下,OPD 将 MATH-500 上的平均 BF16 性能保持率从 35% 提升至 70%,将 HumanEval 上的平均 BF16 性能保持率从 66% 提升至 91%,同时保持短答案性能不变;在预算匹配的对比中,其推理能力增益显著超过继续进行教师强制 QAD 所取得的增益。通过结合 QAD 稳定的低比特初始化与 OPD 对在策略推理能力的恢复,我们的框架提供了一种全面的低于 3 比特解决方案,在保留广泛能力的同时恢复长程推理能力。