PaperlayerView exact revision on arXiv

Source-linked AI summary

LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

Alexander Samarin, Sergei Krutikov, Anton Shevtsov, Sergei Skvortsov, Filipp Fisin, Alexander Golubev

arXiv:2602.23881v2cs.LGcs.CL

TL;DR

Capacity-limited draft models can converge where KL minimization does not guarantee maximal speculative-decoding acceptance. The paper introduces LK losses that directly optimize acceptance rate, consistently improving acceptance metrics across models and architectures, with gains of up to 10% in average acceptance length.

  • Problem

    Capacity-limited draft models may reach suboptimal points where minimizing KL divergence does not guarantee maximizing speculative-decoding acceptance rate.

  • Method

    LK losses directly target acceptance rate through a negative-log-acceptance objective and a hybrid loss that shifts from KL toward direct optimization.

  • Results

    LK losses consistently improve acceptance metrics across six target models and four draft architectures, with gains of up to 10% in average acceptance length.

  • Takeaways & Limitations

    LK losses provide a drop-in alternative to standard draft-training objectives without computational overhead during training.

  • Takeaways & Limitations

    Vocabulary truncation can make KL divergence infinite and require optimizing KL against masked target probabilities, turning it into a proxy of a proxy.

Abstract

from arXiv · show

Speculative decoding accelerates autoregressive large language model (LLM) inference by using a lightweight draft model to propose candidate tokens that are then verified in parallel by the target model. The speedup is significantly determined by the acceptance rate, yet standard training minimizes Kullback-Leibler (KL) divergence as a proxy objective. While KL divergence and acceptance rate share the same global optimum, small draft models, having limited capacity, typically converge to suboptimal solutions where minimizing KL does not guarantee maximizing acceptance rate. To address this issue, we propose LK losses, special training objectives that directly target acceptance rate. Comprehensive experiments across four draft architectures and six target models, ranging from 8B to 685B parameters, demonstrate consistent improvements in acceptance metrics across all configurations compared to the standard KL-based training. We evaluate our approach on general, coding and math domains and report gains of up to 8-10% in average acceptance length. LK losses are easy to implement, introduce no computational overhead and can be directly integrated into any existing speculator training framework, making them a compelling alternative to the existing draft training objectives.

1. Introduction

Speculative decoding accelerates LLM inference by having a lightweight draft model propose tokens for parallel verification, with efficiency largely determined by acceptance rate. LK losses directly optimize acceptance rate because capacity-constrained draft models can reach suboptimal solutions where KL minimization does not guarantee higher acceptance.

  • Motivation: Speculative decoding uses a lightweight draft model to propose multiple tokens that the target model verifies in one forward pass, preserving the target distribution.Its efficiency is largely determined by the expected probability that drafted tokens are accepted.
  • Motivation: Draft models typically contain only 1–5% of the target model’s parameters and commonly converge to suboptimal solutions.At those solutions, minimizing KL divergence provides no guarantee of maximizing acceptance rate.
  • LK losses: LK losses introduce two training-objective variants that directly target acceptance rate in speculative decoding.One minimizes negative log-acceptance rate, while the hybrid gradually shifts from KL toward direct acceptance optimization during training.
  • Contributions: LK losses yield consistent improvements in acceptance metrics across multiple target models and draft architectures, demonstrating model- and architecture-agnostic behavior.The paper also releases training datasets and draft model weights to support reproducibility and further research.

2. Related Work

Related work improves speculative decoding through integrated drafting architectures, autoregressive token prediction, and reduced draft-head overhead. Speculator training is typically framed as knowledge distillation with KL or cross-entropy, while newer studies explore alternative objectives and selective distillation targeting acceptance rate.

  • Drafting Architectures: Standalone target-model copies are straightforward draft models but often bottleneck performance and require same-family pretrained LMs or large-scale training for high acceptance rates.Early speculative decoding works used smaller, standalone target-model versions as draft models.
  • Drafting Architectures: Later methods integrate drafting into the target model, including MEDUSA’s parallel heads and autoregressive speculators such as multi-stage MLPs, EAGLE, and DeepSeek-V3’s MTP module.MEDUSA predicts draft tokens independently, whereas autoregressive designs better model dependencies between draft tokens and their context; MTP requires no separate post-training.
  • Efficiency: FR-Spec reduces draft-head latency by truncating large token vocabularies to a small subset of high-frequency tokens learned from training data.The method addresses LM-head overhead observed even in powerful approaches such as EAGLE.
  • Training Objectives: Speculator training is predominantly knowledge distillation, using KL divergence or equivalently cross-entropy to make draft distributions approximate the target distribution.Draft models are treated as students and target models as teachers.
  • Training Objectives: Related work combines KL with target-model LM loss or hidden-state regression, and investigates reverse KL, TV distance, and selective distillation to address training–inference mismatch.DistillSpec reports that divergence choice depends strongly on the task and KD data, while TV is theoretically linked to directly maximizing acceptance rate.

3. Background and Motivation

Speculative decoding drafts tokens with q and verifies them sequentially using p, so efficiency depends on acceptance. Because acceptance is equivalent to minimizing total variation distance, divergence choice matters especially when draft-model capacity is limited.

  • Speculative sampling: The draft model q proposes K tokens conditioned on context c, which the target model p verifies in parallel and accepts sequentially.The first rejected token terminates the accepted sequence and discards all subsequent drafts.
  • Acceptance rate: Acceptance rate is the expected acceptance probability for each draft token, and its global optimum is achieved when q = p.The acceptance probability depends on the target-to-draft probability relationship for each token prefix.
  • Divergence objectives: α = 1−TV(p, q), making acceptance maximization strictly equivalent to minimizing Total Variation distance.Forward and reverse KL divergences are instead proxy objectives for acceptance.
  • Capacity limitation: Limited draft-model capacity makes the gap between direct acceptance optimization and proxy divergence minimization especially pronounced.This motivates focusing training on the distributional behavior most relevant to acceptance rather than matching the full target distribution.
  • Divergence behavior: Forward KL spreads probability broadly, reverse KL collapses toward the dominant mode, whereas TV maximizes distributional overlap and achieves substantially higher acceptance.TV focuses on tokens carrying the major probability mass under the target distribution, which suits capacity-limited draft models.

4. Methodology

The methodology develops LK objectives that target acceptance rate directly, using gradient analysis to combine KL’s stable optimization with TV’s direct objective. It also presents an acceptance-based alternative and handles truncated draft vocabularies without modifying the target distribution.

  • Objective design: LK objectives directly target acceptance rate instead of using KL divergence as a proxy, based on differences in divergence-gradient optimization.The method is motivated by gradient analysis of how divergence choices guide training.
  • Gradient analysis: TV directly targets acceptance rate but has sign-only gradients, non-differentiable points, and extremely small initialization gradients of O(k/V).For vocabulary sizes exceeding 100k, these properties make pure TV optimization impractical from random initialization.
  • Hybrid objective: LK(p, q) = λ · KL(p∥q) + (1 −λ) · TV(p, q), combining KL’s early-training stability with TV’s later direct acceptance-rate optimization.λ = 1 recovers standard KL training, while λ = 0 gives pure TV optimization.
  • Adaptive schedule: The adaptive schedule drives λ toward 1 as α approaches 0 and toward a small value as α approaches 1, transitioning from KL-dominated to TV-dominated optimization.Ablations empirically confirm that the adaptive schedule outperforms pure TV, KL, and fixed-weight mixtures.
  • Alternative formulation: LK performs TV optimization with adaptive gradient scaling, matching KL gradient magnitude early while matching TV’s gradient direction.The 1/α factor amplifies gradients when acceptance is low, addressing vanishing gradients.
  • Vocabulary truncation: LK losses handle vocabulary truncation naturally because tokens outside the draft vocabulary contribute min(pi, 0) = 0 to acceptance and require no target-probability modification.This preserves optimization against the original target distribution rather than an approximation.

5. Experimental Settings

Experiments test LK losses across six target models, four speculator architectures, three task domains, and both greedy and stochastic decoding. Training and evaluation use standardized datasets, optimization settings, sampling procedures, and acceptance-based metrics.

  • Target models: Six target models span three orders of magnitude and include both dense and mixture-of-experts architectures.The targets range from Llama-3.1-8B-Instruct to DeepSeek-V3.
  • Speculator architectures: Four speculator architectures are evaluated, with three trained for Llama-3.1-8B and EAGLE-3 used for larger models.DeepSeek-V3 uses its native MTP module, which is fine-tuned from pretrained weights.
  • Training setup: 660K prompts from Infinity-Instruct-0625 are paired with responses generated by each target model, using 8K input sequences.Training data matches the distribution encountered during inference.
  • Loss configurations: The comparison includes forward KL, negative log-acceptance, and LK losses, with adaptive scheduling and exponential head weighting γ = 0.8.For Llama-3.1-8B with EAGLE-3, total variation and Lλ losses are also explored.
  • Evaluation protocol: Evaluations cover MT-Bench, HumanEval, and GSM8K using full datasets under greedy decoding T = 0 and stochastic sampling T = 1.All draft models use chain sampling to isolate training-objective effects from inference-time search optimization.
  • Evaluation metrics: τ = K × # accepted tokens / # drafted tokens + 1 is the primary evaluation criterion because it drives speculative-decoding speedup.The bonus token sampled after verification ensures at least one token per speculation round.

6. Evaluation Results

Across target models and draft architectures, LK losses consistently improve average acceptance length over KL training. The hybrid objective with adaptive scheduling performs best, while gains are especially large for lower-capacity drafts and substantial capacity or architectural mismatches.

  • Overall results: Both LK loss types improve over the KL baseline across all tested draft architectures and sampling temperatures for LLaMA-3.1-8B.The evaluation uses average acceptance length τ as its primary metric.
  • Objective and scheduler effects: The hybrid LK objective with an adaptive scheduler achieves the highest acceptance lengths, while the likelihood-based objective also outperforms KL.With η = 1, the difference between LK objectives narrows considerably, particularly under greedy sampling.
  • Objective and scheduler effects: Constant hybrid weights λ = 0.5 make the objective inferior to every other LK setting, showing that curriculum behavior is necessary for effective training.The advantage over KL nearly disappears with constant weights.
  • Objective and scheduler effects: Pure TV loss performs substantially worse than other objectives because its gradients face severe optimization difficulties when the draft distribution is far from the target.The hybrid approach uses KL gradients to guide optimization into the trust region.
  • Architecture and capacity: 7.8% and 8.3% average improvements occur for MEDUSA and MLP speculators, respectively, compared with 3.8% for EAGLE-3 under stochastic sampling.The pattern supports greater benefits from direct acceptance-rate optimization for lower-capacity draft models.
  • Architecture and capacity: +8.2% is the largest reported gain for Qwen3-235B, while GPT-OSS 120B improves by +7.7% compared with +3.8% for GPT-OSS 20B.The strongest improvements occur with large MoE targets and much smaller dense drafts.
  • Architecture and capacity: An extra 5.6% gain over KL is achieved for DeepSeek-V3 with stochastic sampling after fine-tuning an already KL-improved MTP model.This result extends LK’s advantage beyond randomly initialized draft distributions.

7. Conclusion … A.4. Negative Log Acceptance Rate Gradient

LK losses directly optimize acceptance rate, addressing the gap between KL’s shared global optimum and the limitations of capacity-constrained draft models. The conclusion reports consistent empirical gains, while the appendix derives the relevant gradients and shows that −log α optimization is equivalent to TV optimization with adaptive scaling.

  • 7. Conclusion: LK losses directly target acceptance rate, addressing KL-based training’s failure to reach the shared optimum with capacity-limited draft models.The approach is designed as a replacement objective for standard KL minimization.
  • 7. Conclusion: TV-distance optimization outperforms the proxy objectives in a relevant example, while the theory identifies caveats in gradient-based TV minimization.The analysis supports practical acceptance-oriented objective design.
  • 7. Conclusion: 10% gains in average acceptance length are observed across task domains, with larger improvements for low-capacity architectures.Experiments cover six target models from 8B to 685B parameters and four draft architectures, with no training computational overhead.
  • 7. Conclusion: Future work should examine alternative schedulers, learnable or data-dependent per-head aggregation, and direct optimization of draft acceptance length τ.The current study uses a specific adaptive scheduler and fixed exponential aggregation across draft heads.
  • 7. Conclusion: The evaluation uses chain sampling and sequential acceptance logic, motivating tests of LK-trained speculators with tree sampling, block verification, and other inference mechanisms.The passage notes that transfer of per-position acceptance improvements to more complex schemes remains to be systematically verified.
  • A. Gradient Derivations: The appendix derives gradients by treating p as fixed, q = softmax(zq) as the draft distribution, and zq as the optimized logits.The derivations begin from the softmax Jacobian and its Kronecker-delta definition.
  • A.2. KL Divergence Gradient: The forward KL gradient is obtained by differentiating only the term that depends on zq, while the TV gradient uses elementwise multiplication.These are the appendix’s stated intermediate gradient relationships.
  • A.4. Negative Log Acceptance Rate Gradient: Optimizing −log α is equivalent to optimizing TV with an adaptive learning rate proportional to the inverse acceptance rate.The appendix states this as the key relationship for the negative log acceptance rate gradient.

A.5. Gradient Magnitude Analysis · B. Connection to Negative Log-Likelihood · C. Acceptance Rate as Densities Overlap

The analysis compares KL, TV, and LK gradient behavior when q is diffuse and p is concentrated, showing that LK avoids TV’s vanishing gradients while targeting acceptance rate. It also connects acceptance rate to NLL for point-mass targets and to density overlap in continuous settings.

  • A.5. Gradient Magnitude Analysis: Early training is modeled by approximately uniform q over vocabulary V and concentrated p over a support S of size k ≪ V.Specifically, q_i ≈ 1/V, while p_i ≈ 1/k on S and p_i ≈ 0 elsewhere.
  • A.5. Gradient Magnitude Analysis: For support tokens, the KL gradient is approximately −1/k, whereas outside-support tokens have gradient approximately 1/V.These magnitudes follow directly from q − p in the diffuse-q, concentrated-p regime.
  • A.5. Gradient Magnitude Analysis: The TV gradient is approximately −2/V on support tokens and approximately 2k/V^2 outside the support, which is nearly zero.Its expectation term is approximately 1−2k/V in this regime.
  • A.5. Gradient Magnitude Analysis: The TV gradient vanishes as vocabulary size V grows, while LK’s 1/α factor restores O(1/k) magnitude in the diffuse-q regime.In this regime, α ≈ k/V, and LK directly targets acceptance rate.
  • B. Connection to Negative Log-Likelihood: When p is a point mass at token x*, acceptance rate simplifies to precisely the negative log-likelihood of x* under q.The condition is p(x*) = 1 and p(x) = 0 for x ≠ x*.
  • C. Acceptance Rate as Densities Overlap: For continuous distributions, acceptance rate generalizes to the overlap represented by the total area under the minimum of both density curves.This follows from the continuous-distribution generalization of the acceptance-rate expression.

D. Rejection Sampling with Greedy Draft Tokens

The section identifies a mismatch in vLLM’s non-zero-temperature rejection sampling: greedy draft tokens are combined with temperature-scaled target logits. This can systematically reduce acceptance when the target distribution is diffuse, motivating an evaluation patch aligned with the acceptance rate optimized by LK losses.

  • Rejection-sampling mismatch: At non-zero temperatures, proper rejection sampling requires both p(x) and q(x) to represent the actual sampling distributions, but vLLM samples draft tokens greedily while using temperature-scaled target logits.The standard criterion accepts sampled draft tokens with probability min(1, p(x)/q(x)).
  • Greedy draft tokens: Under greedy sampling, the draft selects x∗ = arg maxx q(x), effectively substituting q(x∗) = 1 in the acceptance criterion.
  • Acceptance-rate impact: When the target distribution is diffuse or temperature is high, p(x∗) can be small even when the draft identifies the most likely token, causing systematically low acceptance rates.Acceptance remains effective when the target is confident and agrees with the draft.
  • Objective alignment: Because LK losses optimize the true acceptance rate α, greedy draft evaluation at temperature = 1 creates a mismatch between the training and evaluation objectives.
  • Objective alignment: The vLLM patch makes evaluation faithfully measure the quantity optimized during training.

E. Draft Model Architecture Details

Draft architectures vary by target type and speculator family: EAGLE-3 uses transformer-based heads, DeepSeek-V3 retains its native MTP module, and MLP Speculator and MEDUSA use independent MLP heads.

  • EAGLE-3: EAGLE-3 uses one transformer layer per head, mirroring dense target architectures and processing token embeddings with aggregated intermediate hidden states.For MoE targets, it instead uses a single dense transformer block rather than an MoE block.
  • DeepSeek-V3: DeepSeek-V3 is trained by fine-tuning its native Multi-Token Prediction (MTP) module while preserving the original architecture.
  • MLP Speculator and MEDUSA: MLP Speculator and MEDUSA use one MLP layer per head, with independent heads for each speculative position.MEDUSA predicts all positions in parallel from the same hidden state without token-level autoregression, unlike EAGLE-3’s shared weights across positions.

F. Full Experimental Results

The full experimental results compare speculative decoding objectives using average acceptance length and end-to-end speedup across three benchmarks, inference temperatures, and public checkpoints. The evaluation uses batch size 1 in a low-latency setting and distinguishes models trained from scratch from Hugging Face checkpoints.

  • Evaluation metrics: The study reports average acceptance length τ and end-to-end speedup relative to baseline inference without speculative decoding.These metrics are evaluated for speculative decoding methods trained with different objectives.
  • Evaluation setup: The evaluation covers MT-Bench, HumanEval, and GSM8K with batch size 1 in a low-latency setting at temperatures T =0 and T =1.Public Hugging Face checkpoints are compared using the same inference pipeline.
  • Checkpoint details: Results use the checkpoint revision available during the experiments, before a subsequent weight update released during the review period.One checkpoint also differs architecturally through a wider MLP and smaller vocabulary, producing a slightly larger parameter count.
Loading 2602.23881v2…