Research map / Large language models / Reasoning and chain of thought

LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Speculative decoding · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes LK losses, a negative log-acceptance loss and an adaptive KL-to-TV hybrid, that train speculative decoding draft models to directly maximize acceptance rate, improving average acceptance length over KL training across six targets (8B-685B) and four draft architectures.

Problem

Speculative decoding draft models are trained with KL divergence as a proxy for acceptance rate, but capacity-limited drafts converge to suboptimal points where minimizing KL does not maximize acceptance rate.

Approach

Two objectives that directly target acceptance: -log of the marginal acceptance rate (equivalent to TV optimization with 1/alpha gradient scaling), and a hybrid lambda*KL + (1-lambda)*TV where lambda = exp(-eta*sg[alpha]) shifts from KL to TV as acceptance improves, acting as a trust-region-like curriculum.

Key results

• On Llama-3.1-8B EAGLE-3 at T=1, the hybrid LK loss (eta=3) raises average acceptance length across domains by +3.9% over KL (4.01 vs 3.86 mean tau).
• On Qwen3-235B-A22B with EAGLE-3 at T=1, LK hybrid gives +8.2% mean tau gain over KL (4.08 vs 3.77); GPT-OSS 120B gains +7.7%.
• MEDUSA and MLP speculators on Llama-3.1-8B at T=1 show average improvements of 7.8% and 8.3% respectively; DeepSeek-V3 MTP gains +5.6% at T=1 over KL fine-tuning.

Surprising findings

• Pure TV loss trains far worse than KL (e.g., EAGLE-3 Llama-3.1-8B MT-Bench tau 2.81 vs 3.75 at T=0), attributed to vanishing gradients from random initialization.
• Constant lambda=0.5 hybrid gives almost no gain over KL, indicating the adaptive curriculum is necessary.

Limitations

• Evaluation uses only chain sampling; tree-based drafting and block verification not evaluated, so transfer of gains to those schemes is untested.
• Only one adaptive scheduler form and fixed exponential head weighting were explored; alternatives are left to future work.
• Gains are modest for large EAGLE-3 drafts at T=0 (e.g., +0.5% on Llama-3.3-70B) and the largest gains are on T=1 settings.

Experimental setup

Datasets: Infinity-Instruct-0625 MT-Bench HumanEval GSM8K

Benchmarks: MT-Bench HumanEval GSM8K

Metrics: average acceptance length tau end-to-end speedup (tokens per second ratio)

Baselines: KL divergence training TV loss MEDUSA MLP speculator RedHatAI EAGLE-3 checkpoints yuhuili EAGLE-3 checkpoints DeepSeek-V3 MTP original

Builds on: EAGLE-3 MEDUSA FR-Spec DistillSpec speculative sampling

Closest papers by meaning

• A Theoretical Perspective for Speculative Decoding Algorithm (NeurIPS 2024 · similarity 0.92)
• Speculative Decoding with CTC-based Draft Model for LLM Inference Acceleration (NeurIPS 2024 · similarity 0.92)
• SPEED-Bench: A Unified and Diverse Benchmark for Speculative Decoding (ICML 2026 · similarity 0.91)
• SpecForge: A Flexible and Efficient Open-Source Training Framework for Speculative Decoding (ICML 2026 · similarity 0.91)
• Cactus: Accelerating Auto-Regressive Decoding with Constrained Acceptance Speculative Sampling (ICLR 2026 · similarity 0.91)
• Performance-Driven Policy Optimization for Speculative Decoding with Adaptive Windowing (NeurIPS 2026 · similarity 0.91)