LK损失：投机解码中直接优化接受率的训练目标 LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding
提出直接优化接受率的LK损失函数，在6个目标模型上提升8-10%的平均接受长度
前置知识
投机解码（Speculative Decoding）
投机解码是一种加速自回归大语言模型推理的技术。其核心思想是使用一个轻量级的草稿模型（draft model）快速生成K个候选token，然后由目标模型（target model）在一次前向传播中并行验证这些候选token。根据接受-拒绝采样算法，每个草稿token以概率 β(x) = min(1, p(x)/q(x)) 被接受，其中p是目标分布，q是草稿分布。第一个被拒绝的token会终止后续所有token的接受。这种方法在保持目标模型输出分布不变的前提下实现了显著加速。
本文的核心改进正是针对投机解码中的草稿模型训练目标，理解投机解码的基本机制是理解本文动机和方法的前提。
接受率（Acceptance Rate）
接受率α是投机解码效率的核心指标，定义为第i个草稿token被接受的期望概率：α_i = Σ_x min(q(x|c,x_{<i}), p(x|c,x_{<i}))。平均接受长度τ = K × (接受token数/草稿token数) + 1，表示每轮投机解码平均生成的token数。接受率的全局最优解在q=p时达到，此时α=1。接受率直接决定了投机解码的加速比，是本文优化的核心目标。
本文的核心创新就是直接优化接受率而非使用代理目标，理解接受率的数学定义和重要性是理解本文贡献的关键。
KL散度（Kullback-Leibler Divergence）
KL散度是衡量两个概率分布差异的非对称度量。前向KL散度 KL(p∥q) = Σ p_i log(p_i/q_i) 是当前草稿模型训练中最常用的目标函数。KL散度的梯度为 ∇_{z_q} KL(p∥q) = q - p，具有平滑且良好缩放的优化性质。当p和q完美匹配时KL=0，此时接受率也达到最大值1。然而，KL散度的全局最优并不意味着在有限容量模型的局部最优处也能最大化接受率。
本文指出KL散度作为接受率的代理目标在小容量草稿模型上存在根本性缺陷，这是论文的核心动机之一。
Total Variation距离（TV距离）
TV距离定义为 TV(p,q) = (1/2) Σ |p_i - q_i|，是衡量两个分布差异的对称度量。TV距离与接受率有直接关系：α = 1 - TV(p,q)，因此最大化接受率严格等价于最小化TV距离。TV距离的梯度为 ∇_{z_q} TV = (1/2) q ⊙ (s - E_q[s])，其中 s_i = sign(q_i - p_i)。然而TV距离的梯度存在严重问题：不光滑、在随机初始化时梯度范数为 O(√(k/V))，对于大词汇表V来说梯度极其微小。
TV距离是唯一与接受率有直接等价关系的散度度量，但其优化困难促使作者提出了LK损失来解决这些问题。
草稿模型架构（Draft Model Architectures）
文中评估了四种草稿模型架构：EAGLE-3使用浅层transformer块进行自回归预测，共享权重并融合目标模型各层的隐藏状态；MEDUSA在目标模型最后一层附加并行解码头，假设token间条件独立；MLP投机器使用多阶段MLP扩展MEDUSA；DeepSeek-V3的MTP模块是原生多token预测模块，可直接用于投机解码。这些架构的参数量通常只有目标模型的1-5%，面临严重的容量约束。
不同架构的容量差异直接影响了训练目标的选择效果，本文证明LK损失对低容量架构的改进更为显著。
研究动机
当前投机解码的草稿模型训练普遍采用KL散度作为目标函数，将其视为接受率优化的代理目标。虽然KL散度和接受率共享相同的全局最优解（当草稿分布完美匹配目标分布时），但实际中草稿模型仅有目标模型1-5%的参数量，必然收敛到次优解。在这些次优解处，最小化KL散度并不能保证最大化接受率。论文通过一个直观的高斯混合拟合示例说明了这个问题：当用单高斯去拟合多模态高斯混合分布时，前向KL散度产生覆盖模式的解（α=50.2%），反向KL产生模式搜索的解（α=50.8%），而TV距离通过最大化分布重叠获得显著更高的接受率（α=60.2%）。此外，DistillSpec等工作注意到TV距离理论上应该是正确的优化目标，但结论是散度选择高度依赖于具体任务和数据。这些观察表明，需要一种能直接优化接受率且在实践中可训练的目标函数。
本文的目标是本文的目标是提出直接优化接受率的训练目标函数（LK损失），替代传统的KL散度代理目标。具体而言，作者希望：（1）设计能够直接最大化接受率α的损失函数；（2）解决TV距离优化中的梯度消失和非光滑问题；（3）使新损失函数易于实现、无额外计算开销，可直接集成到现有的投机解码训练框架中；（4）在多种草稿架构和目标模型上验证方法的通用性和有效性。
与已有工作不同的是，本文的独特切入角度在于将关注点从分布对齐（KL散度优化分布差异）转向接受率直接优化。此前的工作如DistillSpec探索了不同散度类型但主要针对预训练语言模型作为外部草稿器的场景，而本文专注于从随机初始化训练的原生草稿模块。更重要的是，作者通过梯度分析揭示了KL和TV损失在优化动力学上的根本差异：KL提供平滑但间接的梯度信号，TV提供直接但病态的梯度信号。基于这一分析，作者提出了两种创新方案——自适应混合目标和基于似然的目标，巧妙结合了两者的优势，既保持了优化的稳定性又直接瞄准接受率最大化。
核心方法
本文的方法基于对不同散度损失梯度行为的深入分析。核心直觉是：KL散度虽然只是接受率的代理目标，但提供平滑、缩放良好的梯度；TV距离虽然直接等价于接受率，但梯度存在病态特性。因此，作者提出在训练的不同阶段动态混合这两种损失——早期依靠KL梯度稳定训练，后期逐渐转向TV直接优化接受率。此外，作者还提出了一种基于似然的替代方案，通过负对数接受率损失自动实现TV优化的自适应梯度缩放。整体技术路线是：先通过理论分析揭示问题本质，再设计两种互补的损失函数，最后在多种架构和模型上进行实验验证。
本文的核心创新点是提出两种LK损失变体，直接优化接受率而非使用代理散度。第一个关键创新是混合目标 LλLK = λ·KL(p∥q) + (1-λ)·TV(p,q)，配合自适应调度器 λ = exp(-η·sg[α])，实现从KL到TV的平滑过渡。当接受率α低时λ趋近1以KL为主，当α高时λ趋近小值以TV为主，类似于信任域方法。第二个关键创新是负对数接受率损失 LαLK = -log α，其梯度 ∇LαLK = (1/α)∇TV 实现了TV优化的自适应梯度放大，在α趋近0时自动放大梯度解决TV的梯度消失问题。与现有方法的本质区别在于：标准训练最小化KL作为代理，DistillSpec探索不同散度但未解决小容量模型的根本问题，而LK损失首次实现了对接受率的直接且可行的优化。
方法步骤详情
方法实现包含以下关键步骤：（1）**梯度分析**：推导KL梯度 ∇KL = q-p、TV梯度 ∇TV = (1/2)q⊙(s-E_q[s]) 和 LαLK梯度 ∇LαLK = (1/α)∇TV，分析在随机初始化（q近似均匀分布、p集中于k个token）时的梯度范数行为，证明KL为 O(1/√k)、TV为 O(√(k/V))（大词汇表V下趋近0）、LαLK恢复为 O(1/√k)。（2）**混合目标设计**：定义 LλLK = λ·KL + (1-λ)·TV，使用自适应调度 λ = exp(-η·sg[α])，其中sg表示停止梯度操作防止反向传播，η是控制过渡速度的超参数（默认η=3），α按草稿位置和批次维度聚合。（3）**词汇表截断处理**：LK损失天然处理词汇表截断，截断词汇外的token对接受率贡献 min(p_i, 0)=0，无需修改目标分布p，而KL训练需要将截断token的logit设为-∞来重新定义p̃，引入额外近似误差。（4）**训练配置**：使用660K提示词生成训练语料，输入序列长度8K，batch size 64，学习率4×10^{-4}余弦调度，AdamW优化器，100步预热，梯度裁剪0.5，草稿头间使用指数权重衰减γ=0.8优先早期位置。
技术新颖性
本文的技术新颖性体现在多个层面。首先，从理论角度，作者首次系统分析了不同散度损失在随机初始化训练场景下的梯度行为差异，揭示了TV距离梯度在大词汇表下为 O(√(k/V)) 的根本性问题，这一分析为损失函数设计提供了理论基础。其次，混合目标的自适应调度器设计借鉴了强化学习中的信任域思想，但将其创新性地应用于散度混合，通过接受率驱动的课程学习实现了从代理优化到直接优化的平滑过渡。第三，负对数接受率损失 LαLK = -log α 的梯度恰好是TV梯度乘以 1/α 的自适应缩放因子，这一优雅的数学关系将TV的直接优化目标与KL级别的梯度缩放完美结合。第四，LK损失在词汇表截断场景下的天然适用性解决了EAGLE-3等架构使用截断词汇表时KL训练面临的无限散度问题，这在实际工业应用中具有重要价值。
实验结果
本文在6个目标模型（参数量从8B到685B）和4种草稿架构上进行了全面实验，核心发现如下：（1）LK损失在所有配置下均优于KL基线。在LLaMA-3.1-8B上，EAGLE-3架构的混合LK损失（η=3）在T=1时将平均接受长度从3.39提升到3.48（+2.7%），MEDUSA从1.72提升到1.85（+7.6%），MLP从2.13提升到2.19（+2.8%）。（2）低容量架构受益更大。MEDUSA和MLP在T=1下平均改进分别为7.8%和8.3%，而EAGLE-3为3.8%，这与理论分析一致：容量有限的模型从直接接受率优化中获益更多。（3）大MoE模型改进最显著。GPT-OSS 120B在T=1下提升+7.7%，Qwen3-235B提升+8.2%，推测大模型与小草稿器之间的容量差距使得KL代理目标更为不足。（4）DeepSeek-V3的MTP微调在T=1下获得+5.6%提升，证明LK损失在非随机初始化场景下同样有效。（5）纯TV训练表现远差于所有其他目标，验证了梯度分析的正确性。（6）固定权重混合（λ=0.5）几乎丧失了LK的优势，确认了自适应课程的必要性。（7）自适应调度器中η=3是较为稳健的默认值，MEDUSA需要更大的η=10。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| LLaMA-3.1-8B + EAGLE-3 (MT-Bench, T=1) | 平均接受长度 τ | 3.48 | 3.39 (KL) | +2.7% |
| LLaMA-3.1-8B + EAGLE-3 (HumanEval, T=1) | 平均接受长度 τ | 4.52 | 4.31 (KL) | +4.9% |
| LLaMA-3.1-8B + MEDUSA (MT-Bench, T=1) | 平均接受长度 τ | 1.85 | 1.72 (KL) | +7.6% |
| LLaMA-3.1-8B + MLP (HumanEval, T=1) | 平均接受长度 τ | 2.62 | 2.16 (KL) | +21.3% |
| GPT-OSS 120B + EAGLE-3 (HumanEval, T=1) | 平均接受长度 τ | 2.44 | 2.27 (KL) | +7.5% |
| Qwen3-235B + EAGLE-3 (HumanEval, T=1) | 平均接受长度 τ | 4.42 | 4.09 (KL) | +8.1% |
| DeepSeek-V3-0324 + MTP (GSM8K, T=1) | 平均接受长度 τ | 5.51 | 5.23 (KL) | +5.4% |
| LLaMA-3.3-70B + EAGLE-3 (GSM8K, T=1) | 平均接受长度 τ | 5.01 | 4.89 (KL) | +2.5% |
局限与改进
本文的局限性包括以下几个方面。首先，实验仅在链式采样（chain sampling）下评估，未系统验证LK训练的草稿器在树状采样（tree-based drafting）等更复杂推理方案下的表现，作者也承认这是自然的下一步工作。其次，自适应调度器仅探索了指数形式 λ = exp(-η·sg[α]) 这一种参数化方式，未充分研究替代调度策略的影响。第三，草稿头间的聚合使用固定的指数权重衰减 γ=0.8，未探索可学习或数据依赖的聚合策略。第四，训练数据使用固定的660K提示词和10个epoch，未研究数据多样性与训练轮数之间的权衡。第五，虽然作者声称LK损失无额外计算开销，但自适应调度器需要计算每个位置的接受率α来更新λ，这增加了少量计算成本。此外，论文主要关注单次投机步的接受率改进，端到端加速比（wall-clock speedup）的结果被放在附录中，部分模型的加速提升相对接受率提升有所衰减。
独立分析的弱点
本文存在以下可改进的弱点：（1）**调度器设计单一**：仅探索了指数衰减形式的自适应调度器，虽然借鉴了信任域思想但未充分理论分析其最优性，可能存在的问题是不同草稿位置、不同任务类型需要不同的过渡策略，一个全局统一的调度器可能不是最优选择。（2）**树状采样未验证**：树状采样是实际工业应用中的主流方案，但本文仅在链式采样下评估，LK损失对树状采样中的token选择和验证效率的影响尚不明确。（3）**超参数敏感性**：η的选择对结果有显著影响（η=0.7到η=10的差异明显），但论文未提供系统的超参数调优指导，MEDUSA需要更大的η=10而其他架构使用默认η=3，这种架构相关的超参数选择降低了方法的即插即用性。（4）**缺乏理论收敛保证**：虽然有启发性的梯度分析，但混合目标和自适应调度的收敛性缺乏严格理论证明，特别是在非凸优化景观中的行为。（5）**评估温度固定**：训练和评估均使用temperature=1，未探索不同温度设置下的泛化能力。
未来方向
基于本文成果可延伸的未来研究方向包括：（1）**直接优化平均接受长度τ**：作者在结论中提出这是更有前景的方向，因为τ直接对应实际加速比，相比优化逐位置接受率α更接近最终目标。（2）**树状采样和块验证**：将LK损失扩展到树状采样和块验证等更复杂的推理方案，验证改进是否能迁移到这些场景。（3）**可学习的调度策略**：替代固定的指数调度器，设计数据驱动或元学习的混合权重策略。（4）**自适应草稿长度**：根据接受率动态调整最大草稿长度K，而非使用固定值。（5）**与其他优化技术结合**：将LK损失与FR-Spec的词汇表截断、AdaSPEC的选择性蒸馏等技术结合。（6）**更大规模验证**：在千B级模型和多种草稿架构上进行更广泛的验证，特别是探索LK损失在不同容量匹配下的最优策略。（7）**多任务和领域自适应**：研究LK损失在不同任务领域（代码、数学、对话）的改进差异及其原因。
复现评估
本文的复现条件较为优越。作者承诺开源训练数据集（HuggingFace nebius/infinity-instruct-completions）和草稿模型权重（HuggingFace nebius/lk-speculators），这大大降低了复现门槛。训练使用标准的Transformer训练流程，超参数设置清晰（batch size 64, 学习率 4×10^{-4}, AdamW优化器等）。硬件方面，需要能够运行8B到685B目标模型的GPU集群来生成训练数据和评估，这对算力有一定要求。LK损失的实现本身非常简单，核心仅需几行代码计算混合损失和自适应权重，无额外计算开销。评估使用vLLM框架，作者还提供了针对非零温度正确拒绝采样的补丁。总体而言，论文的复现难度中等，主要挑战在于算力需求（特别是大模型评估）和vLLM补丁的集成。
论文图表
上排展示不同损失函数（前向KL、反向KL、TV距离）的损失景观（对数尺度），参数为均值µ和标准差σ。下排展示使用各损失优化后的拟合分布与目标分布的重叠情况（绿色区域）。前向KL产生覆盖模式的解（α=50.2%），反向KL产生模式搜索的解（α=50.8%），TV距离最大化重叠（α=60.2%）。
这是论文的动机性示例，用简单直观的方式展示了为什么不同损失函数在容量受限时会产生不同质量的解，为LK损失的设计提供了理论依据。