Pertes LK : Optimisation directe du taux d'acceptation pour le décodage spéculatif

LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

February 27, 2026

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

cs.AI

Résumé

Le décodage spéculatif accélère l'inférence des grands modèles de langage (LLM) autorégressifs en utilisant un modèle "brouillon" léger pour proposer des tokens candidats qui sont ensuite vérifiés en parallèle par le modèle cible. L'accélération obtenue dépend significativement du taux d'acceptation, mais l'apprentissage standard minimise la divergence de Kullback-Leibler (KL) comme objectif indirect. Bien que la divergence de KL et le taux d'acceptation partagent le même optimum global, les petits modèles brouillons, ayant une capacité limitée, convergent généralement vers des solutions sous-optimales où minimiser la divergence de KL ne garantit pas de maximiser le taux d'acceptation. Pour résoudre ce problème, nous proposons les pertes LK, des objectifs d'apprentissage spéciaux qui ciblent directement le taux d'acceptation. Des expériences complètes sur quatre architectures de modèles brouillons et six modèles cibles, allant de 8B à 685B de paramètres, démontrent des améliorations constantes des métriques d'acceptation dans toutes les configurations par rapport à l'apprentissage standard basé sur la divergence de KL. Nous évaluons notre approche dans les domaines généraux, du codage et des mathématiques et rapportons des gains allant jusqu'à 8-10% dans la longueur d'acceptation moyenne. Les pertes LK sont faciles à implémenter, n'introduisent aucune surcharge computationnelle et peuvent être intégrées directement dans tout cadre d'apprentissage de spéculateur existant, ce qui en fait une alternative convaincante aux objectifs d'apprentissage de modèles brouillons actuels.

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.