[控场AI]
· 5 分钟阅读· 2,837 字

RCSD:强化学习加持投机解码,推理提速1.35倍

RCSD:强化学习加持投机解码,推理提速1.35倍

RCSD用强化学习对齐投机解码目标,接受长度提升35%,端到端加速1.35倍。

RCSD是一种针对投机解码草稿模型的新型训练框架。传统软标签蒸馏要求草稿模型逐位置模仿目标模型的完整分布,在高熵区域难以给出高置信预测,导致接受率低、早期拒绝风险高。RCSD转而利用推理中自然产生的接受/拒绝序列构造隐式偏好对,以强化学习信号直接优化接受长度,并引入掩码损失避免惩罚目标模型已正确对齐的token。在多个基准上,RCSD相比HASS将平均接受长度提升35%,实现1.35倍端到端推理加速且生成质量不降。其持续学习机制还能适应特定领域数据分布,适合垂直领域低延迟部署,但代码尚未公开,实际效果仍需更多验证。

投机解码(Speculative Decoding)是加速大语言模型推理的主流技术之一,其核心思路是用一个小型草稿模型快速生成候选 token,再由目标模型并行验证。然而,草稿模型的训练方式一直存在隐患:传统的软标签蒸馏迫使草稿模型模仿目标模型的完整输出分布,包括那些高熵区域——这些区域不仅拉低接受率,还增加了早期拒绝的风险。针对这一问题,研究者提出了 RCSD(Reinforcement and Continual Learning for Domain-Specific Speculative Decoding),一个将草稿模型优化与系统级目标(最大化接受长度)直接对齐的训练框架。

RCSD论文标题截图

投机解码的瓶颈:软标签蒸馏的局限

投机解码的加速效果取决于草稿模型提出的 token 被目标模型接受的比例。接受率越高,一次验证能通过的 token 越多,端到端延迟就越低。传统做法是让草稿模型通过软标签蒸馏学习目标模型的输出分布,但这相当于要求草稿模型在每一个位置上都逼近目标模型的完整概率分布。问题在于,目标模型在高熵区域的分布本身就很分散,草稿模型即使努力模仿,也很难在这些位置给出高置信度的预测,导致接受率下降。更糟糕的是,一旦早期 token 被拒绝,后续的候选序列就全部作废,加速收益大打折扣。

RCSD 的作者指出,这种训练目标与推理时的实际目标存在错位:推理时我们只关心接受长度,而不关心草稿模型是否完美复现了目标模型的分布。因此,需要一种新的训练范式,让草稿模型直接为“被接受”而优化。

软标签蒸馏是一种知识蒸馏形式:教师模型(此处为目标大模型)对每个位置输出完整的概率分布(即“软标签”),学生模型(草稿模型)通过最小化与教师分布的差异(如KL散度)来学习。与硬标签(仅取argmax的one-hot标签)相比,软标签携带了类别间相似性等更多信息,通常能让学生模型学得更好。但在投机解码场景下,软标签蒸馏要求学生模型在每一个token位置都逼近教师的完整分布,包括教师自身分布也很分散的高熵区域。这种逐位置的分布匹配目标与推理时“最大化连续接受长度”的系统目标并不一致,成为接受率提升的瓶颈。

RCSD 的核心机制:隐式偏好对与掩码损失

RCSD 的关键创新在于利用投机推理过程中自然产生的接受和拒绝序列来构造隐式偏好对。具体来说,在推理时,草稿模型提出的 token 会被目标模型验证,被接受的 token 序列和包含拒绝的序列天然形成了一组对比数据。RCSD 将这些数据作为强化学习中的偏好信号,鼓励草稿模型生成更可能被接受的 token 序列,同时抑制那些容易被拒绝的生成模式。

为了避免误伤,RCSD 还引入了掩码损失(masked loss)。在训练过程中,对于那些与目标模型对齐的 token(即目标模型本身也会高概率生成的 token),掩码机制会避免对它们施加惩罚。这样一来,草稿模型不会因为强化学习信号而偏离目标模型已经正确的预测,保证了生成质量不受影响。这种设计让 RCSD 在提升接受率的同时,不牺牲输出的准确性。

隐式偏好对借鉴了强化学习从人类反馈(RLHF)中的偏好优化思路(如DPO),但不需要额外的人工标注:在投机解码的每次推理验证中,被目标模型接受的token序列天然构成“正例”,包含拒绝的序列构成“负例”,二者形成对比。RCSD将这种在线产生的对比信号作为偏好数据,直接优化草稿模型生成更易被接受的序列。相比需要奖励模型的RLHF流程,隐式偏好对省去了显式奖励建模,训练更轻量,且信号与推理时的实际接受行为完全一致。

实验结果:接受长度提升35%,端到端加速1.35倍

研究者在多个基准测试上评估了 RCSD 的效果。与现有的 HASS 方法相比,RCSD 将平均接受长度提升了 35%,在实际工作负载下实现了 1.35 倍的端到端推理加速,且生成质量没有下降。接受长度的提升意味着每次验证能通过更多 token,从而减少了目标模型的验证轮次,这是加速的直接来源。你可能没注意到,RCSD 的训练框架具有持续学习(Continual Learning)的特性,能够适应特定领域的数据分布,这对于 domain-specific 场景尤为重要。

从系统角度看,1.35 倍的端到端加速在推理成本敏感的部署环境中意义显著。投机解码本身已经能带来加速,而 RCSD 在此基础上进一步挖掘了草稿模型的潜力,且不需要修改目标模型或推理架构,属于纯训练侧的优化,落地成本较低。

HASS(Harmonizing Accuracy and Efficiency for Speculative Decoding)是投机解码草稿模型训练的代表性基线方法,其思路同样关注草稿模型与目标模型的对齐质量,但主要仍基于分布匹配类的训练目标。RCSD在论文中选择HASS作为主要对比对象,是因为两者都试图解决软标签蒸馏的局限,而RCSD通过强化学习式偏好优化进一步将训练目标与接受长度直接对齐,从而在相同推理架构下取得更高的接受率和端到端加速。

对领域特定部署的启示

RCSD 的“Domain-Specific”定位值得关注。在实际应用中,企业往往会在特定领域数据上微调目标模型,草稿模型也需要随之适配。RCSD 的持续学习机制允许草稿模型在推理过程中不断从接受/拒绝信号中学习,逐步适应目标模型在特定领域的行为。这种在线适应能力对于法律、医疗、金融等垂直领域的低延迟推理服务具有实用价值。

不过,RCSD 目前仍属于学术研究阶段,其代码和模型尚未公开,实际部署效果还需在更多硬件和模型组合上验证。但其所揭示的方向——用强化学习对齐投机解码的系统级目标——为后续研究提供了清晰的思路。

持续学习(Continual Learning)指模型在不遗忘已学知识的前提下,持续从新数据或新信号中更新参数。在RCSD中,草稿模型在部署后的推理过程中不断接收接受/拒绝反馈,并以此微调自身,逐步适配目标模型在特定领域数据上的行为分布。这与传统的“训练一次、部署冻结”范式不同,属于在线适应。其挑战在于避免灾难性遗忘和保持训练稳定性,RCSD通过掩码损失等机制约束更新方向,使其在适应新领域的同时不破坏已有的对齐能力。

总结

RCSD 通过强化学习和持续学习重构了投机解码中草稿模型的训练方式,用隐式偏好对和掩码损失替代传统的软标签蒸馏,在多个基准上实现了 35% 的接受长度提升和 1.35 倍的端到端加速。这项工作提醒我们,推理加速不仅靠架构创新,训练目标的对齐同样关键。对于关注大模型推理效率的开发者和研究者,RCSD 提供了一个值得跟进的新思路。

#投机解码#强化学习#大模型推理加速#草稿模型训练#RCSD#持续学习#接受长度
分享:

相关推荐