Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

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

TL;DR

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.

Abstract

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.

LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

TL;DR

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.

Abstract

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.
Paper Structure (45 sections, 29 equations, 2 figures, 4 tables)

This paper contains 45 sections, 29 equations, 2 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Draft Model Architectures
  4. Training Methodologies for Draft Models
  5. Background and Motivation
  6. Speculative Decoding
  7. Divergence Losses
  8. Motivating Example
  9. Methodology
  10. Gradient Analysis of Divergence Losses
  11. Hybrid Objective with Adaptive Blending
  12. Adaptive schedule.
  13. Likelihood-based Approach
  14. Gradient behaviour.
  15. Vocabulary Truncation
  16. ...and 30 more sections

Figures (2)

  • Figure 1: Acceptance length $\tau$ vs maximum length $K$ for EAGLE-3 draft models, trained using different objectives with Qwen3-235B-A22B-Instruct as a target model. The values were obtained on the MT-bench dataset with chain sampling and temperature = 1.
  • Figure 2: Fitting a single Gaussian to a Gaussian mixture under different objectives. Top: Loss landscapes (log-scale) over parameters $\mu$ and $\sigma$. Bottom: Resulting distributions and overlap (green). KL divergence produces a mass-covering solution ($\alpha = 50.2\%$), reverse KL exhibits mode-seeking behavior ($\alpha = 50.8\%$), while $\operatorname{TV}$ maximizes overlap ($\alpha = 60.2\%$).