LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

📄 arXiv: 2602.23881v1 📥 PDF

作者: Alexander Samarin, Sergei Krutikov, Anton Shevtsov, Sergei Skvortsov, Filipp Fisin, Alexander Golubev

分类: cs.LG, cs.CL

发布日期: 2026-02-27

💡 一句话要点

提出LK损失函数，直接优化推测解码中的接受率，提升LLM推理速度。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 推测解码 大型语言模型 加速推理 接受率优化 损失函数设计

📋 核心要点

  1. 现有推测解码训练依赖KL散度，但小模型优化KL散度不一定能最大化接受率，导致推理效率受限。
  2. 论文提出LK损失函数，直接优化推测解码的接受率，避免了KL散度作为代理目标带来的次优解。
  3. 实验表明，在多种模型和任务上，LK损失函数能显著提升接受长度，最高可达8-10%。

📝 摘要（中文）

推测解码通过使用轻量级的草稿模型来提议候选token，然后由目标模型并行验证，从而加速自回归大型语言模型（LLM）的推理。加速效果很大程度上取决于接受率，但标准训练通常最小化Kullback-Leibler（KL）散度作为代理目标。虽然KL散度和接受率具有相同的全局最优值，但容量有限的小型草稿模型通常会收敛到次优解，其中最小化KL散度并不能保证最大化接受率。为了解决这个问题，我们提出了LK损失函数，这是一种直接针对接受率的特殊训练目标。在四种草稿架构和六种目标模型（参数范围从8B到685B）上的综合实验表明，与基于标准KL的训练相比，所有配置的接受指标都得到了持续改进。我们在通用、编码和数学领域评估了我们的方法，并报告了平均接受长度高达8-10%的增益。LK损失函数易于实现，不引入计算开销，并且可以直接集成到任何现有的推测器训练框架中，使其成为现有草稿训练目标的引人注目的替代方案。

🔬 方法详解

问题定义：推测解码旨在加速自回归LLM的推理过程。核心问题在于，草稿模型（draft model）的训练目标（通常是最小化KL散度）与最终的推理效率指标（接受率）之间存在gap。小型草稿模型容量有限，最小化KL散度并不一定能最大化接受率，导致推理速度提升有限。

核心思路：论文的核心思路是直接优化接受率，而不是使用KL散度作为代理目标。通过设计新的损失函数，使得草稿模型的训练目标与最终的推理效率指标对齐，从而提高草稿模型的质量，进而提升推测解码的整体效率。

技术框架：整体框架仍然是标准的推测解码流程，包括草稿模型生成候选token，然后目标模型验证这些token。论文的主要贡献在于草稿模型的训练阶段，通过引入LK损失函数来替代传统的KL散度损失。训练完成后，草稿模型和目标模型即可用于推测解码推理。

关键创新：最重要的技术创新点在于LK损失函数的设计。与传统的KL散度损失不同，LK损失函数直接基于接受率进行优化。具体来说，LK损失函数鼓励草稿模型生成更容易被目标模型接受的token序列，从而提高整体的接受率。

关键设计：LK损失函数的具体形式未知，论文中可能包含其数学定义和实现细节。关键参数可能包括损失函数的权重系数、梯度裁剪策略等。网络结构方面，论文使用了多种草稿模型架构进行实验，表明LK损失函数具有较好的泛化能力。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，与基于KL散度的标准训练相比，LK损失函数在各种草稿模型架构和目标模型上都取得了显著的性能提升。在通用、编码和数学领域，平均接受长度提高了8-10%。这些结果表明，LK损失函数是一种有效的草稿模型训练方法，可以显著提高推测解码的效率。

🎯 应用场景

该研究成果可广泛应用于各种需要加速LLM推理的场景，例如在线对话系统、机器翻译、文本生成等。通过提高推测解码的效率，可以降低推理延迟，提升用户体验，并降低计算成本。未来，该方法有望进一步推广到其他加速LLM推理的技术中。

📄 摘要（原文）

Speculative decoding accelerates autoregressive large language model (LLM) inference by using a lightweight draft model to propose candidate tokens that are then verified in parallel by the target model. The speedup is significantly determined by the acceptance rate, yet standard training minimizes Kullback-Leibler (KL) divergence as a proxy objective. While KL divergence and acceptance rate share the same global optimum, small draft models, having limited capacity, typically converge to suboptimal solutions where minimizing KL does not guarantee maximizing acceptance rate. To address this issue, we propose LK losses, special training objectives that directly target acceptance rate. Comprehensive experiments across four draft architectures and six target models, ranging from 8B to 685B parameters, demonstrate consistent improvements in acceptance metrics across all configurations compared to the standard KL-based training. We evaluate our approach on general, coding and math domains and report gains of up to 8-10% in average acceptance length. LK losses are easy to implement, introduce no computational overhead and can be directly integrated into any existing speculator training framework, making them a compelling alternative to the existing draft training objectives.