LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

[ICLR 2025] LK Losses：拒绝代理指标，直接为“接受率”而战的投机采样新方案

LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

总结

问题

方法

结果

要点

摘要

本文提出了 LK Losses，一种旨在直接优化投机采样（Speculative Decoding）中“接受率”的训练目标函数。通过引入负对数接受率损失（Likelihood-based）和带有自适应权重的混合损失（Hybrid Objective），该方法在从 8B 到 685B 参数的多种模型架构上，实现了 8-10% 的平均接受长度提升。

TL;DR

投机采样（Speculative Decoding）的效率核心在于接受率（Acceptance Rate）。长期以来，开发者习惯用 KL 散度来训练投机模型（Draft Model），但这其实是一个“不完美”的代理指标。本文提出的 LK Losses 拨乱反正，通过数学推导和实验证明：直接优化接受率不仅理论更完备，且在无需任何额外推理开销的前提下，能将主流 LLM 的推理加速性能提升 8-10%。

背景定位：代理指标的局限性

在投机采样的范式中，我们希望小模型 $q$ 尽可能预测大模型 $p$ 的行为。在理想状况下（全球最优解）， $q = p$ ，此时 KL 散度和接受率确实是统一的。然而，现实很残酷：

容量瓶颈：投机模型通常只有目标模型参数的 1%~5%，它永远不可能真正复现大模型的分布。
优化偏好：当模型容量不足时，KL 散度倾向于平均分配概率质量（Mode-covering），而接受率的提升则需要模型更精准地捕捉高频 Token（Mode-seeking）。

作者通过一个简单的混合高斯拟合（见下图）展示了直觉：总变分距离（Total Variation, TV） 寻找的解能产生最大的分布重叠区域，从而直接最大化接受率。

拟合实验对比：TV 距离相比 KL 能获得更高的接受率

动机与深度洞察：为什么不直接用 TV 距离？

既然 TV 距离直接等价于 $(1 - e x t 接受率)$ ，为什么前人不用它作为 Loss？通过梯度分析，作者揭示了三个痛点：

梯度消失：在初始化阶段，由于预测分布比较平滑，TV 的梯度幅值是 $O (k / V)$ ，在词表规模 $V > 100 k$ 的今天，这个梯度几乎为零。
不连续性：TV 损失的曲面在 $q_{i} = p_{i}$ 处不可微，优化不稳定。
忽略误差量级：TV 的梯度只看误差的方向（正负号），而不看差距有多大，这导致训练早期收敛极慢。

核心方法论：LK Losses

为了解决上述梯度问题，作者提出了两种极具工程智慧的 Loss 设计：

1. 似然优化法 ( $L_{L K}^{α}$ )

公式定义为： $L_{L K}^{α} = - lo g (e x t 接受率)$ 。 物理直觉：这相当于在 TV 梯度的基础上，乘以了一个 $1/ α$ 的缩放因子。当接受率 $α$ 很低时，梯度被自动放大；当 $α$ 接近 1 时，梯度回归正常。这巧妙地解决了 TV 损失在早期的梯度消失问题。

2. 自适应混合调度 ( $L_{L K}^{λ}$ )

这是一种更稳健的课程学习策略： $L_{L K}^{λ} = λ \cdot K L + (1 - λ) \cdot T V$ 其中 $λ$ 会随着当前 Batch 的实时接受率 $α$ 动态衰减。

初期： $α$ 低， $λ$ 接近 1，利用 KL 散度平滑的曲面快速把模型拉入“信任区域”。
后期： $α$ 提高， $λ$ 减小，模型开始专注优化接受率，做最后的冲刺。

EAGLE-3 架构与 LK Loss 结合示意图

实验结果：全方位的 SOTA 提升

实验覆盖了从 Llama-3 (8B/70B) 到巨大规模的 DeepSeek-V3 (685B) 的多组对比。

广泛的适配性：无论是并行的 MEDUSA 头，还是递归的 EAGLE 架构，LK Losses 均稳定优于 KL 散度。
低容量增益更高：在架构更简单的 MTP 或 MEDUSA 上，由于模型更难拟合完整分布，直接优化接受率带来的红利更加可观（提升 5.6%~8.3%）。
复杂领域表现：在编程 (HumanEval) 和数学 (GSM8K) 任务中，平均接受长度 $a u$ 的提升尤为明显，因为这些领域对 Token 的精确性要求极高。

不同 Loss 在主流模型上的平均接受长度对比表

深度思考与局限性

LK Losses 揭示了**“指标对齐”**在系统优化中的重要性。然而，该方法目前主要针对“接受率”这一 Token 级别的指标。作者也客观指出，未来的终极目标应该是直接优化“系统推理延迟”，但这涉及更多非确定性的硬件环境因素。

此外，对于已经极度强大的 Draft Model，LK Losses 的边际效应会递减。但对于大多数仍在努力提升小模型推理效率的场景，LK Losses 提供了一个零成本、即插即用的性能 Buff。

总结

LK Losses 告诉我们，有时候数学上的“完美代理”在有限的工程约束下并非最优解。通过深入理解梯度动态并引入自适应机制，我们可以让投机采样跑得更远。

关键词：Speculative Decoding, Acceptance Rate, KL Divergence, Total Variation, DeepSeek-V3

发现相似论文

试试这些示例

查找最近一年内针对投机采样（Speculative Decoding）中 Draft Model 训练目标改进的其他学术论文。
论文中提到的“全变分距离（Total Variation Distance）”在知识蒸馏任务中最早由谁提出用于替代 KL 散度？
是否有研究将类似的自适应混合损失函数（Hybrid Objective）应用在扩散模型（Diffusion Models）的剪枝或加速任务中？

[ICLR 2025] LK Losses：拒绝代理指标，直接为“接受率”而战的投机采样新方案

1. TL;DR

2. 背景定位：代理指标的局限性

3. 动机与深度洞察：为什么不直接用 TV 距离？

4. 核心方法论：LK Losses

4.1. 1. 似然优化法 ($L_{LK}^{\alpha}$)

4.2. 2. 自适应混合调度 ($L_{LK}^{\lambda}$)

5. 实验结果：全方位的 SOTA 提升

6. 深度思考与局限性

7. 总结

[ICLR 2025] LK Losses：拒绝代理指标，直接为“接受率”而战的投机采样新方案

LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

TL;DR

背景定位：代理指标的局限性

动机与深度洞察：为什么不直接用 TV 距离？

核心方法论：LK Losses

1. 似然优化法 (LLKα​)

2. 自适应混合调度 (LLKλ​)

实验结果：全方位的 SOTA 提升

深度思考与局限性

总结

发现相似论文

试试这些示例

1. 似然优化法 ( $L_{L K}^{α}$ )

2. 自适应混合调度 ( $L_{L K}^{λ}$ )