LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

1 · 动机:proxy 之坑

1.1 历史脉络:为什么大家一直用 KL/CE

从 Leviathan 2023 把 SD 落地开始,draft model 的训练几乎一律是一个 KD(knowledge distillation)流程:把 target 当 teacher,用 KL(p∥q) 或等价的 CE 让 draft 学习 target 的输出分布。MEDUSA(2024)、EAGLE 系列(2024–2025)、DeepSeek-V3 MTP(2024)、FR-Spec(2025)——架构在变,但 loss 的骨架没变。EAGLE 在此基础上加了 hidden-state 的 MSE regression,MEDUSA 加了 LM 目标,但都是额外补丁,核心仍然是 KL。

这种共识基于一个事实:KL 和 acceptance rate α 共享同一个全局最优点。如果 q 能完美匹配 p,KL=0、α=1。所以"训 KL ≈ 训 acceptance"在理论上是没毛病的——前提是你能到达那个全局最优。

但 draft model 通常只有 target 1–5% 的参数(EAGLE-3 一层 transformer 对 70 层的 Qwen3-235B),它注定到不了 q=p。问题就是:在你必然停留的这个 suboptimal 点附近,KL 最小化和 α 最大化是两件不同的事。这就是这篇 paper 的核心攻击面。

1.2 别的方案为什么不够

注意最后一行:DistillSpec 早就指出 TV 是对的目标。但他们的实验是从一个已经会说话的 LM 上微调(初始 q 已经接近 p),所以 TV 的梯度问题被掩盖。在从头训 EAGLE/MEDUSA 这种小 head 的真实场景下,纯 TV 训练会立刻撞墙——这是 LK 这篇要解决的工程关键。

1.3 为什么这事不平凡

看上去很直白:既然 α=1−TV(p,q),那就直接最小化 TV 不行吗?或者更直接,把 α 写出来当 loss?三个障碍让这件事过去三年没人做成:

1. min(p,q) 不光滑。α=Σ min(p_i,q_i) 在 p_i=q_i 处不可微,对应 TV 中的 sign 跳变。直接梯度只携带方向(过预测 vs 欠预测),不携带幅度——欠预测一点点和欠预测很多得到一样大小的 sign 信号。
2. TV 梯度幅度消失。论文 §A.5 给出在"q 几乎均匀(随机初始化) + p 集中在 k 个 token"的 regime 下:‖∇_zTV‖ = O(√k / V)。对 V=128k 的现代 vocabulary,这个梯度比 KL 的 O(1/√k) 小好几个数量级。换句话说,从随机初始化开始,TV 推不动 logits。
3. −log α 看似优雅,但 α=Σ min(p,q) 不是 softmax 的 likelihood。它没有 NLL 那种"目标 token 一个 slot"的好分解(虽然在 p 是 point mass 时退化为 NLL,见 §B)。它的梯度推不出标准 cross-entropy 那种 q−p 形式;LK 推到最后是 (1/α)·∇TV,等于"TV 加自适应 1/α 放大"——这层自适应是关键,把 TV 的梯度消失问题在 α→0(随机初始化)时自动恢复 O(1/√k)。

所以这篇 paper 的"非平凡"含金量,不在 idea(idea 是 obvious 的:就直接优化 α 嘛),而在把 TV 梯度的病理诊断清楚 + 给出两条都能绕过的路径(混合调度 / 自适应放大)。这是工程师视角下的"算法障碍解决"。

2 · 背景速查

3 · 梯度解剖:KL vs TV(为什么直接做 TV 撞墙)

论文最 load-bearing 的不是 idea 而是这一节的诊断,值得逐字啃。

3.1 KL 的梯度:漂亮但代理

这是个被反复用过的好朋友:每个 logit 收到的力 = "我现在的概率"减"目标概率"。方向、大小都自洽。在随机初始化 + p 集中在 k token 的 regime 下,‖q−p‖ = O(1/√k) — 假设 k=20,梯度幅度 ~0.22,稳稳能动。

3.2 TV 的梯度:方向对了但小到看不见

问题三连:

• 方向对、幅度无关:s_i 只是 ±1。q_i 比 p_i 低 0.001 还是低 0.5,得到一样的 sign。
• 有 q 的尺度因子:整体被 ⊙q 调制。随机初始化时 q_i ≈ 1/V,V=128k → q_i ≈ 8e-6。‖∇TV‖ = O(√k/V) ≈ √20/128000 ≈ 3.5e-5。比 KL 小约 7000 倍。
• landscape 不光滑:在 q_i=p_i 的流形上 sign 跳变。

3.3 关键数值表(论文 Table 3 的复述)

读法:L^α_LK 在主梯度成分上和 KL 同量级(支持 token 上 −1/k),但方向和 TV 一致(对 α 直接求导)。这就是它在工程上"既能从 random init 起步,又把 α 当真目标"的来源。

4 · L^λ_LK:KL+TV 自适应混合(trust region 的味道)

4.1 公式

sg[·] 是 stop-gradient,防止 λ 自己被反向传播改写。α 用当前 batch / sequence / position 聚合后的真实 acceptance rate 估计。

4.2 直觉:课程学习而非线性混合

关键点不是"加权平均",而是自适应调度:

• 训练初期 α≈0:λ = exp(0) = 1 → 等价于 KL,享受光滑、有幅度的梯度。
• 训练后期 α 接近高水位(η=3 时,α=1 → λ=exp(−3)≈0.05):TV 接管,直接打磨 α。

这和 trust region(TRPO)的精神一致:KL 当 soft constraint 把 q 拉进 p 附近,然后在 trust region 内对真目标 TV 做下降。论文里直接给了这个对偶解读:

4.3 Worked example: 一个 token 走一遍

设词表 V=1024(简化),target p 集中在 token x=42,p₄₂=0.7,p₁₁=0.2,p₉₉=0.1,其余为 0。draft 早期 q 接近均匀,q_i≈1/1024≈9.8e-4,设当前 batch 平均 α ≈ 0.05(等价于很差的对齐)。

• 当下 λ = exp(−3 × 0.05) = exp(−0.15) ≈ 0.86 → KL 占大头(86% KL + 14% TV)。
• KL 在 z₄₂ 的梯度 ≈ q₄₂ − p₄₂ ≈ −0.699:把 logit 42 往上推。
• TV 在 z₄₂:s₄₂=sign(q−p)=−1, E_q[s] ≈ +1−2k/V ≈ 0.994, q₄₂·(s−E_q[s])/2 ≈ 9.8e-4 × (−1.994)/2 ≈ −9.8e-4。比 KL 小 700×。
• 合成:λ·∇KL + (1−λ)·∇TV ≈ 0.86×(−0.699) + 0.14×(−9.8e-4) ≈ −0.601。KL 项主导,稳稳推。

训了几千步后 α=0.6,q₄₂=0.55,p₄₂=0.7:

• λ = exp(−1.8) ≈ 0.165 → TV 接管(83.5%)。
• KL 在 z₄₂:0.55−0.7=−0.15。
• TV 在 z₄₂:q₄₂ 大了,梯度幅度也跟上来了,这时 sign-only 信号反而是优点 — 不会被 outlier 大误差带跑。

5 · L^α_LK:−log α 一条龙

5.1 直觉

把 α 写成 marginal probability of acceptance:α = Σ q(x)·β(x)。这是"draft 抽到 x 且被接受"的边缘概率。最大化 α 的最大似然写法就是最小化 −log α:

比 hybrid 更简洁:不需要混合权重、不需要 schedule。但能不能从随机初始化训得动?这就要看它的梯度。

5.2 关键关系(§A.4)

这是整篇 paper 最漂亮的一行。意思是:−log α 的梯度 = TV 梯度 × 1/α。1/α 是自适应 boost:

• α 小(随机初始化、训练早期):1/α 大,把 TV 那个消失的梯度强行放大回 O(1/√k) — 量级追上 KL。
• α 大(训练后期):1/α 接近 1,梯度恢复成纯 TV 的尺度,正好这时 q 已经接近 p,TV 本身的梯度也不再病态。

5.3 Worked example:同 CE 不同 α

这是一个最能击穿"KL 训练 = α 训练"幻觉的对比。设 V=4,target p = (0.50, 0.30, 0.15, 0.05)。

读这个表:q_A 的 CE 最低(KL 也最低,因 KL=CE−H(p)),但α 反而比 q_B 低 0.05。q_B 牺牲了一点 CE,换来 acceptance rate 提升。q_C 最铺得均匀,CE 居中但 α 最差。这就是"KL 最小 ≠ α 最大"的具体数值证据。LK 的两个 loss 都会在 q_A 和 q_B 之间偏向后者。

6 · Vocabulary truncation:KL 多了一层 proxy,LK 没这个问题

EAGLE-3 + FR-Spec 配置下,draft 的 LM head 只输出一个截断词表(比如 32k,target 词表 128k),其它 token 的 q_i=0。这对 KL 是致命的:p_i>0 而 q_i=0 → KL=∞。常规做法是把 target 也截:p̃ = softmax(m⊙z_p),把截断外的 logit 设 −∞,再去做 KL(p̃∥q)。

这就是论文一句犀利的吐槽:"makes KL a proxy of a proxy"(代理的代理)。一层是 KL≠α,另一层是 p̃≠p。

LK 怎么处理?从 α=Σmin(p_i,q_i) 看,被截掉的 token 上 q_i=0,贡献 min(p_i, 0)=0,自然被忽略 — 不需要修改 p。L^α_LK 和 hybrid 的 TV 部分都直接对原始 p 做优化,不引入额外近似。这是 LK 落地 EAGLE-3+FR-Spec 时的一个免费红利。

7 · 公式串与数值直觉

7.1 推导链条(§A 复述)

1. softmax Jacobian:∂q_i/∂z_q,j = q_i(δ_ij − q_j)。
2. KL:∂KL/∂z_q,j = −Σ_i (p_i/q_i)·q_i(δ_ij − q_j) = −p_j + q_j·Σp_i = q_j − p_j。
3. TV:链式 sign 经过 softmax Jacobian 得到 ½q_j(s_j − E_q[s])。
4. −log α:∂(−log α)/∂z = (1/α)·(−∂α/∂z) = (1/α)·∂TV/∂z(因 α=1−TV)。

7.2 在 p 是 point mass 的退化情况

若 p(x*)=1,其它为 0,则 α = Σ min(p_i, q_i) = min(1, q(x*)) = q(x*)(因 q 是概率)。所以:

这是个不错的 sanity check:当 target 退化为硬标签,LK 等于普通的 cross-entropy 训练。它是 NLL 在软标签(distillation)情况下的"正确" generalization——比直接做 soft-label CE 更直接对 α。

7.3 数值敏感性表(η 选择)

读法:η 越大,KL 越早退场。EAGLE-3 / MTP 这种α 涨得快的架构默认 η=3 即可;MEDUSA 因为头之间独立、α 涨得慢,论文用 η=10 强行推快 TV 接管。这是个有点反常识但可解释的调参逻辑——你想让 TV 在 α 还没涨太高的时候就接管,因为再不接管它就不打磨了。

8 · 实验关键结果

8.1 跨架构(LLaMA-3.1-8B,Table 1 节选)

关键观察:

• 纯 TV 全面输给 KL —— 印证 §3 的 TV 梯度消失分析,这不是抽象数学,是真训练曲线。
• 容量越小,LK 增益越大:MEDUSA + 7.8%, MLP + 8.3%, EAGLE-3 + 3.8% — 这与"limited capacity → proxy gap 越严重"的理论一致。
• 固定 λ=0.5 几乎没提升 — 自适应 schedule 是非平凡的。

8.2 跨 target scale(Table 2)

大模型(尤其 MoE)+小 dense draft 的容量差越悬殊,LK 的相对提升越大 — Qwen3-235B 的 +8.2% 是论文最亮的数字。DeepSeek-V3 是 fine-tune MTP(不是从头训),也提升 5.6%,印证 LK "在 q,p 距离大时都能赢" 的普适性。

8.3 一些容易被略过的细节

• vLLM 默认实现是greedy 抽 draft token(即使 T>0),这违反了 SD 的理论假设。作者打了 patch 实现真正的 rejection sampling(基于 vllm-project/vllm#20459)。这点足够独立成 paper—— evaluation 设置下 T=1 的 acceptance 数字,在 unpatched vLLM 上是被低估的(见 §D)。
• Loss 跨 head 用 γ=0.8 exponential decay,前面的 head 权重大。理由:earlier head 的 acceptance rate 决定整段是否能继续,影响 τ 的杠杆远大于后头的。
• MTP 微调里 hybrid loss 自带"为后头 head 补 KL"的副作用:后头 α 低 → λ 大 → KL 主导 → 拉回到 trust region;前头 α 已高 → λ 小 → TV 精修。同一条公式同时处理两类 head。

9 · 与同类工作对比

核心差异化:

• vs CE/KL distillation:同样易实现、同样无开销,但目标对齐。这是真正的 drop-in replacement。
• vs MSE on hidden / EAGLE-3 multi-layer KD:这些是架构 / 特征层面的努力,正交于 LK。论文显示的实验是直接套到 EAGLE-3 上,所以这两条是叠加的关系而非互斥。
• vs DistillSpec:DistillSpec 第一个意识到 TV 才是对的目标,但他们没解决 from-scratch 训练的梯度消失问题。LK 的两条公式是工程层面的"修复方案"。
• vs MARS-类阈值方法:那些是事后挽救(reject 掉低 α head),LK 是事前(从训练目标就对齐)。两者甚至可以叠加。

10 · 局限 / 个人 take / 待验证

• "τ" 不等于 wall-clock speedup。论文只报 τ,没给端到端 latency 数字。τ 涨 8% 在 LM head 不主导的场景下大致 = 8% throughput,但在 FR-Spec 截断词表的场景下要看 head 计算占比。
• η 的调参依然是 manual(EAGLE-3 用 3,MEDUSA 用 10)。这暗示 LK 的"自适应"并不全自动,而是把"调温度"的负担从 schedule 函数式上转移到了 η 这个常数。Future work 里作者也提到 learnable per-head loss aggregation 会更好。
• α 是 batch 聚合估计,带方差。早期 batch 小时 sg[α] 噪声会让 λ 抖动,论文没讨论 batch size 的稳健性。
• top-k / top-p 部署没纳入。生产里目标分布常带 top-p 截断,这会让 α 的实际分布跟训练时定义有 gap。这一条作者在 future work 里直接承认。
• K=6/7 之外的延展性。Figure 1 显示 K 越大 LK 优势越明显,但论文没系统扫描 K∈{2,3,…,15};EAGLE-3 论文里 K 上去后 τ 边际下降,LK 是否真改这条曲线值得验证。
• p 的来源:训练用 target sample 自己生成的 660k Infinity-Instruct,这是 on-policy。如果换成静态语料(如 ShareGPT),LK vs KL 的差距会怎么变?未做对照。

需要验证的几个问题:

1. L^α_LK 在更大 V(比如 200k+)和更深 draft(2-3 层 transformer)上,1/α 放大是否足够?会不会触发数值不稳定(α 接近 0 时 1/α 爆炸)?
2. 把 LK 嫁接到 RL post-training 里的 draft model(参考 ReSpec / NeMo-RL),acceptance length 抖动场景下 sg[α] 估计的方差会不会让 λ 跑飞?
3. 对 reverse KL(mode seeking)直接也做一次 hybrid:λ·rKL + (1−λ)·TV 是否优于 λ·KL + (1−λ)·TV?论文没尝试。
4. p 是 point mass 时退化为 NLL,中间区间(p 是 top-p=0.9 截断的)是否退化为某个加权 NLL?是否解释了 LK 在 GSM8K(高 confidence target)上提升相对小?
5. Loss landscape 在 p≈q 附近 KL 与 TV 的 Hessian 性质对比 — 是否能解释为什么 hybrid 略胜 likelihood?

11 · Memory points