LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

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.

Large language models can generate useful text and code, but producing tokens one by one makes them expensive and slow to run. Speculative decoding speeds this up by using a smaller draft model to propose several next tokens, which the larger target model then checks in parallel. The speedup depends heavily on how often the target model accepts the draft model's suggestions. Standard training methods usually minimize KL divergence, which encourages the draft model to imitate the target model's probability distribution, but this can be an imperfect proxy for acceptance rate when the draft model is much smaller. We introduce LK losses, a family of training objectives that optimize acceptance rate more directly. These losses can replace KL-based training objectives without changing the inference procedure or adding computational overhead. Across several target models, draft architectures, and task domains, LK losses consistently improve acceptance length, with gains of up to about 8-10%. This makes speculative decoding more effective and can help reduce the cost and latency of serving large language models.