SOAR Weekly Champion Notes 09 | Vanilla Xiaozhang: LK^λ loss and improving speculative decoding acceptance rate

Editor's note: loss and acceptance rate

In this week's SOAR 2026 weekly contest, Vanilla Xiaozhang took the weekly title again. This technical share focuses on the core training objective of speculative decoding. Starting from the nature of the token-level expected acceptance rate α, it breaks down the central problem with using KL divergence as the loss in EAGLE's original training: with a small-capacity draft model, the local optimum of KL is not the same as the local optimum of α. From the overlooked key identity that maximizing acceptance rate is equivalent to minimizing TV distance, to the LK^λ adaptive hybrid loss that fixes TV's weak early-training gradients, to validating the approach in EAGLE draft training for MiniCPM-SALA, this share clearly reveals the logic that aligns the training loss with the true optimization objective. For developers interested in speculative decoding training and deployment optimization for large models, it is a reference with both theoretical derivation and engineering practice value.

Glad to take the Week 8 title again with my teammates.

01 What exactly is the training objective of speculative decoding?

The core mechanism of speculative decoding is: the draft model predicts several tokens, and the target model verifies them in parallel in a single forward pass. Whether each draft token is accepted follows the classic importance sampling acceptance rule — given the target distribution p and the draft distribution q, the probability that token x is accepted is min(1, p(x)/q(x)). Taking the expectation over q gives the closed form:

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

α is the token-level expected acceptance rate; the speculative decoding paper uses it directly to compute the wall-clock speedup. What draft training for speculative decoding really wants to maximize is α.

So why do people train drafts with KL divergence instead of optimizing α directly?

02 EAGLE's original loss

EAGLE's original training loss (Li et al., arXiv:2401.15077) consists of two terms:

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) # output distribution cross-entropy

The first term keeps the draft's intermediate representations close to the target's; the second keeps the draft's output token distribution consistent with the target's. EAGLE-3 further drops the first term and keeps only the fit at the output-distribution level, which is essentially −Σ p_target · log p_draft. In the EAGLE code it is conventionally called plogp loss — that is the variable name, though strictly it is p·log q (KL and CE differ only by a constant term H(p) that does not depend on q).

KL is a very reasonable training objective — it shares a global optimum with α: when q ≡ p, KL reaches its minimum of 0 and α = 1. The problem is that "sharing a global optimum" is actually of little use: the capacity of a small draft model is far from enough to make q = p, so training inevitably stops at some local optimum, and KL's local optimum and α's local optimum are not the same point.

03 An identity that keeps getting overlooked

LK Losses (Samarin et al., arXiv:2602.23881, 2026.02) starts from an identity that has long existed but is rarely used for training (first seen in Leviathan et al., 2023):

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

where TV(p, q) = (1/2) · Σ_x |p(x) − q(x)| is the total variation distance. The proof takes only one line of algebra:

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

This identity shows that maximizing the acceptance rate is equivalent to minimizing the TV distance. TV is continuous and convex, its computation involves only sums of absolute values with no log or division, and it is numerically more stable than KL. Swapping KL for TV as the training loss directly:

L = TV(p, q)

The paper includes a toy problem: fitting a mixture-of-Gaussian target with a single-Gaussian draft with 2 parameters (mean μ, standard deviation σ). After KL optimization converges, the acceptance rate is 50.2%; after TV optimization converges, it is 60.2% — with the same model capacity, simply changing the loss function yields +10pp in acceptance rate.

04 LK^λ: adaptive mixing

Using TV alone has a small engineering problem — early in training, when q and p barely overlap, TV approaches its upper bound of 1, and the gradient signal is weak and "flat." KL, by contrast, provides denser gradients when the distributions differ greatly (the penalty diverges logarithmically).

The paper's compromise is an adaptive mixture (LK^λ):

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

sg is stop-gradient — λ itself does not participate in backpropagation. Intuitively:

Early in training, α ≈ 0 and λ ≈ 1, so the loss degenerates to KL, providing dense gradients for warm-up

Later in training, α rises and λ → 0, so the loss transitions smoothly to TV and aligns directly with the acceptance rate

The entire switch is controlled by α itself, with no schedule to tune by hand.

05 Observations from practice

In engineering terms, LK^λ is almost a zero-cost switch — its computational complexity is on the same order as KL, with no extra forward or backward pass. SpecForge (SGLang's official EAGLE training framework) has a ready-made implementation in PR #492. We switched the loss from plogp to LK^λ in EAGLE draft training for MiniCPM-SALA, and the acceptance length did improve, with a trend consistent with the paper.

But LK^λ is not "magic." The improvement reported in the paper's ablations (acceptance length +4~10%, covering 4 draft architectures, 6 targets, and 8B~685B parameter scales) is a steady and consistent gain.

In hindsight, the training loss was separated from the metric we actually care about by KL as a proxy, and for a long time nobody saw a problem — until someone wrote out that line of algebra.

✒️ About the authors

Wang Zhiyao: team lead, a fourth-year undergraduate in artificial intelligence at Huazhong University of Science and Technology.

Wu Xianjin: teammate, a fourth-year undergraduate in artificial intelligence at Huazhong University of Science and Technology.

🌟 Friends of the OpenBMB community:

The SOAR 2026 Sparse Operator Acceleration Grand Prix is in full swing, with new challenges every week and a live leaderboard waiting for you. If you want to demonstrate your technical strength in model inference acceleration, kernel-level optimization, and quantization, we sincerely invite you to enter the competition, compete alongside Cao Yi and other outstanding participants, and jointly push the inference limits of MiniCPM-SALA! Click the image below or visit the official competition site at https://soar.openbmb.cn/ to learn more.