Train Where the Quantized Model Goes: On-Policy Distillation for Low-Bit Reasoning

TL;DR

提出了在低比特推理中使用策略蒸馏(OPD)的方法,显著提高了数学和代码推理性能。

cs.LG 🔴 高级 2026-09-23 6 次浏览
Yuanteng Chen Zhilei Liu Peisong Wang Yuantian Shao Chuangyi Li Weining Wang Shuang Qiu Gang Li Jing Liu Jian Cheng
量化 蒸馏 低比特 推理 深度学习

核心发现

方法论

研究采用了量化感知蒸馏(QAD)结合策略蒸馏(OPD)的方法。首先通过QAD恢复模型的广泛能力,然后通过OPD在模型实际生成的轨迹上进行教师监督,结合任务验证器奖励以提高推理性能。

关键结果

  • 在MATH-500数据集上,OPD将BF16性能保留率从35%提高到70%,在HumanEval上从66%提高到91%。
  • 在四个模型中,OPD显著超过了继续使用教师强制QAD的推理增益。
  • OPD在2.79和1.88有效比特下提高了数学和代码推理性能,同时保持短形式性能。

研究意义

该研究为低比特量化模型提供了全面解决方案,恢复了长形式推理能力,同时保持了广泛的能力。它解决了传统QAD在长生成任务上的性能缺陷,为学术界和工业界提供了新的思路。

技术贡献

技术贡献包括将策略蒸馏应用于量化模型的实际生成轨迹,结合任务验证器奖励,显著提高了长形式推理能力。与现有方法相比,提供了新的理论保证和工程可能性。

新颖性

首次将策略蒸馏应用于量化模型的推理恢复,与传统的教师强制蒸馏方法相比,提供了更有效的解决方案。

局限性

  • OPD需要额外的计算资源和时间来进行策略蒸馏,可能对资源有限的环境不友好。
  • 在极端低比特情况下,恢复效果可能不如高比特模型。

未来方向

未来工作可以探索如何进一步减少策略蒸馏的计算成本,以及在更广泛的任务中应用该方法。

AI 总览摘要

量化模型在低比特情况下的推理能力通常受到限制,尤其是在长形式生成任务中。现有的量化感知蒸馏(QAD)方法虽然能恢复部分性能,但在数学和代码推理上仍有显著缺陷。

本文提出了一种新的策略蒸馏(OPD)方法,通过在模型实际生成的轨迹上进行教师监督,结合任务验证器奖励,显著提高了低比特模型的推理性能。实验结果表明,OPD在多个数据集上显著提高了性能,同时保持了短形式任务的能力。

这种方法为低比特量化模型提供了全面解决方案,恢复了长形式推理能力,同时保持了广泛的能力。未来的研究可以进一步优化该方法的计算成本,并在更多任务中验证其有效性。

深度分析

研究背景

低比特量化模型在推理任务中面临性能下降的问题,尤其是在长形式生成任务中。量化感知蒸馏(QAD)是恢复性能的一种方法,但在长形式推理上仍有显著缺陷。

核心问题

低比特量化模型在长形式推理任务中容易进入重复循环,耗尽解码预算而无法完成解决方案。这是由于量化放大了曝光偏差,模型在训练时没有在实际生成的轨迹上进行训练。

核心创新

提出了策略蒸馏(OPD)方法,通过在模型实际生成的轨迹上进行教师监督,结合任务验证器奖励,显著提高了推理性能。与传统QAD方法相比,OPD提供了更有效的解决方案。

方法详解

  • �� 使用QAD初始化模型,恢复广泛能力。
  • �� 在模型实际生成的轨迹上进行策略蒸馏(OPD)。
  • �� 结合任务验证器奖励,提供密集的令牌级指导。

实验设计

实验在Qwen3-0.6B、Qwen3-1.7B、Qwen3-4B和Falcon3-1B模型上进行,使用MATH-500和HumanEval数据集进行评估。通过比较BF16性能保留率来验证方法的有效性。

结果分析

OPD显著提高了数学和代码推理性能,在多个数据集上恢复了长形式推理能力,同时保持了短形式任务的能力。

应用场景

该方法可用于低比特量化模型的推理恢复,适用于资源受限的环境中部署高效模型。

局限与展望

OPD需要额外的计算资源和时间来进行策略蒸馏,可能对资源有限的环境不友好。

通俗解读 非专业人士也能看懂

想象你在厨房做饭,量化模型就像一个厨师,他用低比特的工具来做菜。传统的QAD方法就像给厨师一本食谱,但他只能在固定的步骤上得到指导。策略蒸馏(OPD)就像让厨师在实际做菜时得到指导,帮助他在每一步都做出正确的选择。

简单解释 像给14岁少年讲一样

想象一下你在玩游戏,你的角色有一个很酷的技能,但因为装备不好,技能效果很差。量化模型就像这个角色,QAD方法就像给角色一个新装备,但在长时间的战斗中效果不佳。OPD方法就像给角色一个实时指导,让他在战斗中发挥出最佳效果!

术语表

量化感知蒸馏 (Quantization-aware Distillation)

一种通过教师模型监督量化模型的方法,以恢复其性能。

用于恢复低比特量化模型的广泛能力。

策略蒸馏 (On-Policy Distillation)

在模型实际生成的轨迹上进行教师监督的方法。

用于提高低比特模型的长形式推理能力。

曝光偏差 (Exposure Bias)

模型在训练时使用固定前缀,但在推理时使用自身生成的前缀。

导致量化模型在长生成任务中性能下降。

任务验证器 (Task Verifier)

用于评估模型生成结果正确性的方法。

结合OPD以提高推理性能。

BF16

一种用于深度学习模型的浮点数格式,提供较高精度。

作为性能比较的基准。

开放问题 这项研究留下的未解疑问

  • 1 如何进一步减少策略蒸馏的计算成本,以便在资源有限的环境中应用。
  • 2 在更广泛的任务中验证策略蒸馏的有效性。

应用场景

近期应用

低比特推理恢复

可用于资源受限环境中部署高效模型,恢复推理能力。

远期愿景

广泛任务应用

探索在更多任务中应用策略蒸馏的方法,进一步提高模型性能。

原文摘要

Quantization-aware distillation (QAD) restores much of the short-form question-answering performance lost to sub-3-bit quantization, yet leaves mathematical and code reasoning substantially impaired. Long generations often degenerate into repetitive loops, exhausting the decoding budget without completing a solution. We trace this gap to quantization-amplified exposure bias: QAD trains on fixed corpus prefixes, while quantization-induced deviations compound along the model's own autoregressive trajectories. To address this mismatch, we introduce an on-policy distillation (OPD) stage that places teacher supervision where the quantized model actually goes. Starting from a QAD checkpoint, the student generates through the quantized forward path used at deployment and receives feedback from a frozen full-precision teacher on its own prefixes, combining dense token-level guidance with task-verifier rewards. Across four models at 2.79 and 1.88 effective bits, OPD raises average BF16 performance retention from 35% to 70% on MATH-500 and from 66% to 91% on HumanEval while preserving short-form performance, with reasoning gains substantially exceeding those of continued teacher-forced QAD in matched-budget comparisons. By coupling QAD's stable low-bit initialization with OPD's on-policy reasoning recovery, our framework provides a comprehensive sub-3-bit solution that preserves broad capabilities while restoring long-form reasoning.

cs.LG cs.AI