Pérdidas LK: Optimización Directa de la Tasa de Aceptación para el Decodificado Especulativo

LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

February 27, 2026

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

cs.AI

Resumen

La decodificación especulativa acelera la inferencia en modelos de lenguaje autoregresivos (LLM) utilizando un modelo de borrador ligero para proponer tokens candidatos que luego son verificados en paralelo por el modelo objetivo. La aceleración depende significativamente de la tasa de aceptación, sin embargo, el entrenamiento estándar minimiza la divergencia de Kullback-Leibler (KL) como objetivo sustituto. Si bien la divergencia KL y la tasa de aceptación comparten el mismo óptimo global, los modelos de borrador pequeños, con capacidad limitada, generalmente convergen a soluciones subóptimas donde minimizar la KL no garantiza maximizar la tasa de aceptación. Para abordar este problema, proponemos las pérdidas LK, objetivos de entrenamiento especiales que se enfocan directamente en la tasa de aceptación. Experimentos exhaustivos con cuatro arquitecturas de borrador y seis modelos objetivo, que van desde 8B hasta 685B de parámetros, demuestran mejoras consistentes en las métricas de aceptación en todas las configuraciones en comparación con el entrenamiento estándar basado en KL. Evaluamos nuestro enfoque en dominios generales, de programación y matemáticos, y reportamos ganancias de hasta 8-10% en la longitud media de aceptación. Las pérdidas LK son fáciles de implementar, no introducen sobrecarga computacional y pueden integrarse directamente en cualquier framework existente de entrenamiento de especuladores, lo que las convierte en una alternativa convincente a los objetivos de entrenamiento de borradores existentes.

English

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.