---
title: LK Losses for Speculative Decoding
url: https://www.emergentmind.com/topics/lk-losses
type: topic
---

# LK Losses for Speculative Decoding

Searching arXiv for the main paper and core speculative decoding references.
arxiv_search(query="2602.23881 speculative decoding LK losses Leviathan", max_results=10)
Searching arXiv for speculative decoding and acceptance-rate references.
arxiv_search(query="speculative decoding Leviathan acceptance total variation", max_results=10)
arxiv_search(query="2602.23881", max_results=5)
arxiv_search(query="\"LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding\"", max_results=5)
LK losses are training objectives for speculative decoding that directly optimize the acceptance rate, and consequently the average acceptance length, of a draft or speculator model against a fixed target model instead of relying on Kullback–Leibler divergence as a proxy objective. Introduced for lossless speculative sampling, they are formulated as drop-in replacements for standard KL-based draft training and comprise two variants: an adaptive hybrid KL–TV objective and a likelihood-based objective that minimizes the negative log acceptance. The method is motivated by the observation that, although KL divergence and acceptance rate share the same global optimum at $q=p$, limited-capacity draft models often converge to regimes in which minimizing KL does not maximize acceptance. Across four draft architectures and six target models ranging from $8\text{B}$ to $685\text{B}$ parameters, LK losses are reported to improve acceptance metrics consistently, with gains of up to $8$–$10\%$ in average acceptance length, without computational overhead [2602.23881].

## 1. Speculative decoding setting and acceptance metrics

Speculative decoding accelerates autoregressive large language model inference by using a lightweight draft model to propose candidate tokens that are then verified in parallel by the target model. In the standard lossless speculative sampling setting, the target model defines a next-token distribution $p(x \mid c, x_{<i})$ and the draft model defines $q(x \mid c, x_{<i})$, where $c$ is the context and $x_{<i}=(x_1,\dots,x_{i-1})$ is the drafted prefix. Draft tokens are sampled from $q$ and then verified with the rejection-sampling criterion introduced for speculative decoding by Leviathan et al. [2211.17192].

For a drafted token $x_i$, the verifier uses the per-token acceptance probability
$$
\beta(x_i \mid c, x_{<i}) = \min\!\left(1,\frac{p(x_i \mid c, x_{<i})}{q(x_i \mid c, x_{<i})}\right).
$$
Tokens are verified in parallel but accepted sequentially; the first rejection terminates the accepted chunk and discards later drafts. The expected acceptance probability at position $i$ is
$$
\alpha_i = \mathbb{E}_{x \sim q(\cdot \mid c, x_{<i})}\!\left[\beta(x \mid c, x_{<i})\right]
= \sum_{x \in \mathcal{V}} \min\!\bigl(q(x \mid c, x_{<i}), p(x \mid c, x_{<i})\bigr).
$$
When the conditioning is clear, the paper writes
$$
\alpha = \sum_x \min(p_x,q_x).
$$

A central identity is the equivalence between acceptance and total variation distance:
$$
\mathrm{TV}(p,q)=\frac{1}{2}\sum_x |p_x-q_x|,
\qquad
\alpha = 1-\mathrm{TV}(p,q).
$$
The global optimum is therefore attained at $q=p$, where $\alpha=1$ and $\mathrm{TV}(p,q)=0$ [2602.23881].

The paper treats per-step acceptance $\alpha_i$ as the fundamental optimization target and uses average acceptance length per speculation round, $\tau$, as the principal evaluation metric:
$$
\tau = K \cdot \frac{\#\text{accepted\_tokens}}{\#\text{drafted\_tokens}+1},
$$
where $K$ is the maximum draft length and the “$+1$” accounts for the guaranteed post-verification target token. Under an independence approximation used only for intuition, the expected number of accepted tokens in a chain of length $K$ would be $\sum_{i=1}^{K}\prod_{j=1}^{i}\alpha_j$; the training procedure nevertheless optimizes the per-position $\alpha_i$ directly [2602.23881].

## 2. Why KL divergence is a proxy rather than the target

Standard speculator training minimizes $\mathrm{KL}(p\|q)$. Under infinite draft capacity, this is consistent with maximizing acceptance, because
$$
\arg\max_q \alpha(p,q)=\arg\min_q \mathrm{KL}(p\|q)=\{q:q=p\}.
$$
The paper’s claim is not that KL is incorrect in the limit, but that it is only a proxy objective when the draft cannot match the target exactly due to limited capacity, architectural mismatch, or truncated vocabularies [2602.23881].

The acceptance objective depends on overlap,
$$
\alpha=\sum_x \min(p_x,q_x),
$$
whereas forward KL is mode-covering. This distinction is operationally important for small drafts: probability mass can be arranged in a way that reduces divergence without increasing the overlap that determines verifier acceptance. The exact identity $\alpha=1-\mathrm{TV}(p,q)$ implies that maximizing acceptance is equivalent to minimizing total variation, not KL. The paper further invokes the Pinsker-type inequality
$$
\mathrm{TV}(p,q)^2 \le \frac{1}{2}\mathrm{KL}(p\|q),
$$
to show that large KL implies large TV, but small KL does not tightly imply small TV.

The argument is sharpened at the gradient level. For draft logits $z_q$ with $q=\mathrm{softmax}(z_q)$, the forward-KL gradient is
$$
\nabla_{z_q}\mathrm{KL}(p\|q)=q-p.
$$
For total variation, if $s_i=\mathrm{sign}(q_i-p_i)$ and $\mathbb{E}_q[s]=\sum_i q_i s_i$, then
$$
\nabla_{z_q}\mathrm{TV}(p,q)=\frac{1}{2}q\odot\bigl(s-\mathbb{E}_q[s]\bigr).
$$
The TV gradient depends on error sign rather than error magnitude and is non-smooth along $q_i=p_i$. In the early-training regime with large vocabularies and diffuse $q$, the paper states that pure TV gradients are tiny, with norm $O(\sqrt{k}/V)$ when $p$ is concentrated on an effective support of size $k$, whereas KL gradients remain well scaled at $O(1/\sqrt{k})$. This is the core failure mode behind the paper’s critique of pure direct TV optimization from random initialization [2602.23881].

A common misconception is therefore that the exact identity $\alpha=1-\mathrm{TV}(p,q)$ makes pure TV an immediately practical replacement for KL. The reported ablations reject that conclusion: pure TV substantially underperforms because its gradients vanish early, while KL remains stable but mismatched to the ultimate acceptance objective [2602.23881].

## 3. Formal definition of the LK objectives

The paper proposes two LK objectives. The first is a hybrid KL–TV loss:
$$
L_{LK}^{\lambda}(p,q)=\lambda \cdot \mathrm{KL}(p\|q) + (1-\lambda)\cdot \mathrm{TV}(p,q).
$$
Its mixing coefficient is adaptive rather than fixed:
$$
\lambda = \exp\!\bigl(-\eta \cdot \mathrm{sg}[\alpha]\bigr),
$$
with $\eta>0$ and $\mathrm{sg}[\cdot]$ the stop-gradient operator. The coefficient is computed per position using batch-aggregated $\alpha$ for that head. As acceptance is small, $\lambda \to 1$ and training is KL-dominated; as acceptance rises, $\lambda \to 0$ and the loss becomes TV-dominated. The schedule is intended to retain KL’s optimization stability early and shift toward direct acceptance optimization later [2602.23881].

The second variant is a likelihood-based objective:
$$
L_{LK}^{\alpha}(p,q)=-\log \alpha,
\qquad
\alpha=\sum_{x\in\mathcal{V}}\min(p_x,q_x).
$$
This objective directly maximizes the marginal probability of acceptance. In the degenerate case of a point-mass target $p(x^\star)=1$, it reduces to the standard negative log-likelihood,
$$
L_{LK}^{\alpha}(p,q)=-\log q(x^\star).
$$
That reduction connects the objective to conventional language-model training while preserving its acceptance-centered meaning [2602.23881].

The likelihood-based gradient reveals why it avoids TV’s vanishing-gradient pathology:
$$
\nabla_{z_q}L_{LK}^{\alpha}(p,q)=\frac{1}{\alpha}\nabla_{z_q}\mathrm{TV}(p,q).
$$
The paper interprets this as TV optimization with adaptive gradient scaling by $1/\alpha$. In the early regime with diffuse $q$ and peaked $p$, where $\alpha \approx k/V$, the resulting gradient norm is approximately $1/\sqrt{k}$, matching KL’s magnitude while following TV’s direction.

Neither objective is convex in $z_q$. KL is smooth, whereas TV and $L_{LK}^{\alpha}$ are non-smooth. The paper nevertheless treats both as workable with the adaptive blend and scaling mechanisms above. At $q=p$, all objectives attain the same global optimum: $\alpha=1$, $\mathrm{TV}(p,q)=0$, and KL is minimized [2602.23881].

## 4. Training procedure, implementation, and truncated vocabularies

LK losses are trained against a frozen target model. The inputs are the target logits $z_p$ and probabilities $p$, the draft logits $z_q$ and probabilities $q=\mathrm{softmax}(z_q)$, a dataset of prompt–continuation pairs generated by the target model, and $K$ speculative heads. For each training step, the procedure computes $p_i$ and $q_i$ at each drafted position $i$, evaluates
$$
\alpha_i=\sum_x \min\bigl(p_i(x),q_i(x)\bigr),
$$
constructs either $L_{LK}^{\lambda}$ or $L_{LK}^{\alpha}$ per head, and aggregates across heads with exponential decay
$$
L=\sum_{i=1}^{K}\gamma^{i-1}L_i.
$$
Gradients are backpropagated into draft parameters only; the target is not updated [2602.23881].

The reported implementation choices are specific. The optimizer is AdamW with $(\beta_1,\beta_2)=(0.9,0.95)$, learning rate $4\times 10^{-4}$, cosine schedule, warmup of $100$ steps, and gradient clipping at $0.5$. The primary regime is stochastic decoding with $T=1$, with additional evaluation at $T=0$ for greedy decoding. Aggregation decay is set to $\gamma=0.8$ to prioritize early heads. The paper describes $L_{LK}^{\lambda}$ with $\eta=3$ as the strongest default, while noting that MEDUSA benefits from larger $\eta$, for example $\eta=10$, because its acceptance improves more slowly. Fixed-weight blends, such as $\lambda=\mathrm{const}$, are reported to underperform adaptive scheduling [2602.23881].

A practical implementation point concerns truncated draft vocabularies, as in FR-Spec or EAGLE-3 heads. For KL training, $\mathrm{KL}(p\|q)$ can become infinite when $q_i=0$ but $p_i>0$, so a standard workaround masks target logits and replaces $p$ with
$$
\tilde{p}=\mathrm{softmax}(m\odot z_p).
$$
LK losses avoid this approximation. When a token lies outside the draft vocabulary, $q_i=0$ and $\min(p_i,q_i)=0$, so it simply contributes nothing to $\alpha$ or TV. The paper presents this as a practical advantage over KL-based training [2602.23881].

Another implementation issue is evaluation at nonzero temperature. The authors state that at $T>0$ they patched vLLM to implement correct rejection sampling of draft tokens. Otherwise, greedy draft sampling yields
$$
\alpha_{\mathrm{greedy}}=p(x^\star),
\qquad
x^\star=\arg\max q,
$$
which underestimates acceptance at high temperature. This point addresses a measurement pitfall rather than a change in the training objective itself [2602.23881].

## 5. Empirical performance across architectures, targets, and domains

The empirical study spans six target models—LLaMA-3.1-8B-Instruct, LLaMA-3.3-70B-Instruct, GPT-OSS-20B, GPT-OSS-120B, Qwen3-235B-A22B-Instruct, and DeepSeek-V3—and four draft architectures: EAGLE-3, MEDUSA, a multi-stage MLP speculator, and DeepSeek-V3 MTP. Training uses $660\text{K}$ prompts from Infinity-Instruct-0625 with target responses generated by each target model. Evaluation covers MT-bench, HumanEval, and GSM8K under chain sampling at $T=0$ and $T=1$ [2602.23881].

The principal reported metric is average acceptance length $\tau$. Throughput or speedup is not directly measured; instead, $\tau$ is treated as the driver of speed. Across all configurations, LK losses improve acceptance metrics relative to KL-based training. The paper’s selected highlights include the following stochastic-decoding results at $T=1$:

- For LLaMA-3.1-8B with EAGLE-3, the mean $\tau$ across MT-bench, HumanEval, and GSM8K rises from $3.86$ under KL to $4.01$ under $LK^{\lambda}(\eta=3)$, a $+3.9\%$ gain; $LK^{\alpha}$ reaches $3.98$.
- For LLaMA-3.3-70B with EAGLE-3, $\tau$ rises from $4.50$ to $4.66$, a $+3.5\%$ gain.
- For GPT-OSS-120B with EAGLE-3, $\tau$ rises from $2.46$ to $2.65$, a $+7.7\%$ gain.
- For Qwen3-235B with EAGLE-3, $\tau$ rises from $3.77$ to $4.08$, a $+8.2\%$ gain.
- For DeepSeek-V3-685B with fine-tuned MTP, $\tau$ rises from $4.43$ to $4.68$, a $+5.6\%$ gain [2602.23881].

The gains are larger for smaller drafts. Under stochastic sampling, MEDUSA and the MLP speculator show average improvements of roughly $+7.8\%$ and $+8.3\%$, respectively, whereas EAGLE-3 shows $+3.8\%$. This supports the paper’s central thesis that direct optimization matters most when draft capacity is limited relative to the target [2602.23881].

The ablations clarify the role of each design choice. On LLaMA-3.1-8B with EAGLE-3, pure TV substantially underperforms; for example, at $T=1$ on MT-bench, TV yields $2.67$ versus $3.39$ for KL. A fixed-weight hybrid with $\lambda=0.5$ also underperforms the adaptive schedule. $LK^{\alpha}$ is generally stronger than KL, and $LK^{\lambda}(\eta=3)$ is reported as the best setting overall [2602.23881].

## 6. Practical interpretation, limitations, and relation to prior work

The paper situates LK losses against the background of KL-based knowledge distillation, which is standard for MEDUSA, EAGLE, and MTP-like heads. In that literature, draft training aligns distributions but only indirectly affects verifier acceptance. LK losses replace that indirect objective with losses that target the overlap quantity $\sum_x \min(p_x,q_x)$ explicitly [2602.23881].

The most direct practical guidance concerns where the method helps. The reported gains are strongest for low-capacity drafts such as MEDUSA and MLP, for large target models including MoE targets and models in the $120\text{B}$–$235\text{B}$–$685\text{B}$ range, and for mismatched regimes such as pretrained MTP modules trained only for early positions. Truncated vocabularies are another favorable regime because LK objectives handle them without masking the target distribution [2602.23881].

The main limitations are optimization-related rather than conceptual. Pure TV from random initialization suffers from vanishing gradients and should be avoided unless $q \approx p$ and the vocabulary is small. Over-rapid decay of $\lambda$ can move training into a non-smooth TV-dominated regime prematurely, whereas insufficient decay preserves too much KL guidance and leaves acceptance suboptimal. The paper recommends monitoring per-head $\alpha_i$ during training: if early stagnation appears, increasing $\eta$ or mixing in $LK^{\alpha}$ is suggested. Later heads often have lower acceptance, and the adaptive schedule responds by increasing KL guidance there [2602.23881].

The relation to inference-side accelerations is complementary rather than competitive. The paper evaluates only chain sampling to isolate the training objective’s contribution, but it argues that gains in per-position $\alpha_i$ should transfer to blockwise or tree-style speculative schemes because their acceptance is governed by the same overlap quantity. This suggests that LK losses are best understood as a training-layer modification for speculative systems, not as an alternative decoding algorithm [2602.23881].

A final point of interpretation follows from the paper’s “no computational overhead” claim. LK losses use the same target and draft tensors as KL training. The additional operations—$\min(p,q)$, absolute differences for TV, and the scalar $\alpha=\sum_x \min(p_x,q_x)$—are $O(|\mathcal{V}|)$, like KL, and require no extra forward passes, caches, or sampling steps. The paper therefore presents the method as a direct objective replacement: acceptance is optimized more faithfully, especially in capacity-limited regimes, while retaining the training footprint of existing speculative distillation pipelines [2602.23881].

Source: https://www.emergentmind.com/topics/lk-losses