SOAR 周冠军笔记09 | 香草小张： LK^λ 损失与推测解码接受率提升

2026-05-17

编者按·损失与接受率

本周 SOAR 2026 周赛中，香草小张再次拿下本周冠军。本次技术分享聚焦推测解码的核心训练目标，从 token 级期望接受率 α 的本质出发，拆解了 EAGLE 原始训练采用 KL 散度作为损失的核心问题 —— 小容量 draft 模型下，KL 的局部最优与 α 的局部最优并不一致。从「最大化接受率等价于最小化 TV 距离」这一被忽略的关键恒等式，到解决 TV 训练早期梯度偏弱问题的 LK^λ 自适应混合损失，再到该方案在 MiniCPM-SALA 的 EAGLE draft 训练中落地验证，这份分享清晰揭示了训练损失与真实优化目标的对齐逻辑。对于关注大模型推测解码训练与落地优化的开发者而言，这是一份兼具理论推导与工程实践价值的参考。

很高兴再次和队友拿下第八周冠军。

01 推测解码的训练目标到底是什么

推测解码（Speculative Decoding）的核心机制是：draft 模型预测若干 token，target 模型在单次 forward 中并行 verify。每个 draft token 是否被接受，依据的是经典的 importance sampling 接受规则——给定 target 分布 p 与 draft 分布 q，token x 被接受的概率是 min(1, p(x)/q(x))。在 q 上求期望就得到 closed form：

α(p, q) = Σ_x min(p(x), q(x))

α 就是 token 级的期望接受率，speculative decoding 论文里直接拿它算 wall-clock 加速比。推测解码 draft 训练真正希望最大化的，就是 α。

那为什么大家训练 draft 时不直接优化 α，而是用 KL 散度？

02 EAGLE 的原始 loss

EAGLE（Li et al., arXiv:2401.15077）的原始训练 loss 由两项组成：

L = L_reg + w_cls · L_cls, w_cls = 0.1L_reg = SmoothL1(h_draft, h_target) # hidden state regressionL_cls = CE(p_target, p_draft) # 输出分布的 cross-entropy

第一项约束 draft 中间表示与 target 中间表示接近，第二项约束 draft 的输出 token 分布与 target 一致。EAGLE-3 进一步把第一项去掉，只保留输出分布层面的拟合，本质上就是 −Σ p_target · log p_draft，EAGLE 代码里惯称 plogp loss——变量名如此，严格说是 p·log q（KL 与 CE 只差一个不依赖 q 的常数项 H(p)）。

KL 是一个非常合理的训练目标——它和 α 共享全局最优解：当 q ≡ p 时 KL 取最小值 0，同时 α = 1。问题在「共享全局最优」这件事其实没什么用：小 draft 模型的容量远不足以让 q = p，训练过程必然停在某个局部最优，而 KL 的局部最优和 α 的局部最优 并不是同一个点。

03 一个会被反复忽略的恒等式

LK Losses（Samarin et al., arXiv:2602.23881, 2026.02）的出发点是一个早已存在但很少被用于训练的恒等式（最早见于 Leviathan et al., 2023）：

α(p, q) = 1 − TV(p, q)

其中 TV(p, q) = (1/2) · Σ_x |p(x) − q(x)| 是 total variation distance。证明只用一行代数：

Σ_x min(p, q) = Σ_x [p − (p − q)_+] = 1 − Σ_x (p − q)_+ = 1 − TV(p, q)

这条恒等式说明：最大化 acceptance rate 等价于最小化 TV distance。TV 连续、凸，计算只涉及绝对值求和，不含 log 或除法，数值上比 KL 稳定。直接把 KL 换成 TV 作为训练 loss：

L = TV(p, q)

论文里有个 toy problem：用一个 2 个参数（均值 μ、标准差 σ）的单高斯 draft 去拟合 mixture-of-Gaussian target。KL 优化收敛后接受率 50.2%，TV 优化收敛后接受率 60.2%——同样的模型容量，仅仅换 loss 函数就拿到 +10pp 接受率。

04 LK^λ：自适应混合

只用 TV 在工程上有个小问题——训练早期 q 与 p 几乎不重叠时，TV 接近上界 1，梯度信号弱且「扁平」。KL 在分布差异大时反而能给出更密集的梯度（penalty 随对数发散）。

论文给的折中是一个自适应混合（LK^λ）：

L_LK^λ = λ · KL(p ‖ q) + (1 − λ) · TV(p, q)λ = exp(−η · sg[α])

sg 是 stop-gradient——λ 本身不参与反传。直观地：

训练初期 α ≈ 0，λ ≈ 1，loss 退化为 KL，提供密集梯度做 warm-up

训练后期 α 上升，λ → 0，loss 平滑过渡到 TV，直接对齐接受率

整条切换由 α 自身控制，没有需要手调的 schedule。

05 落地观察

LK^λ 在工程上几乎是零成本切换——计算复杂度与 KL 同阶，不引入额外 forward / backward。SpecForge（SGLang 官方 EAGLE 训练框架）的 PR #492 有现成的实现。我们在 MiniCPM-SALA 的 EAGLE draft 训练中把 loss 从 plogp 换成 LK^λ，acceptance length 确实有提升，趋势和论文一致。

但 LK^λ 不是「魔法」。论文 ablation 里给出的提升幅度（接受长度 +4~10%，覆盖 4 种 draft 架构、6 个 target、8B~685B 参数规模）是稳定但一致的改进。

回头看，training loss 和真正关心的 metric 之间隔着 KL 这层代理，大家用了很久也没觉得有问题——直到有人把那行代数写出来。

✒️ 作者介绍

王致尧：队长，华中科技大学人工智能专业大四本科生。

武显瑾：队友，华中科技大学人工智能专业大四本科生。

🌟 OpenBMB 社区的小伙伴们：

SOAR 2026 稀疏算子加速大奖赛正火热进行中，每周都有新挑战，实时榜单虚位以待 。如果你想在模型推理加速、算子级优化、量化技术上展示你的技术实力 ，我们诚挚邀请你入场参赛，与曹议及更多优秀的小伙伴同台竞技，共同挑战 MiniCPM-SALA 的推理极限！点击下图或访问大赛官网 https://soar.openbmb.cn/ 了解更多。