LK Losses: Ottimizzazione Diretta del Tasso di Accettazione per il Decodifica Speculativa
LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding
February 27, 2026
Autori: Alexander Samarin, Sergei Krutikov, Anton Shevtsov, Sergei Skvortsov, Filipp Fisin, Alexander Golubev
cs.AI
Abstract
La decodifica speculativa accelera l'inferenza dei grandi modelli linguistici (LLM) autoregressivi utilizzando un modello draft leggero per proporre token candidati, che vengono poi verificati in parallelo dal modello target. Lo speedup è significativamente determinato dal tasso di accettazione, eppure l'addestramento standard minimizza la divergenza di Kullback-Leibler (KL) come obiettivo surrogato. Sebbene la divergenza KL e il tasso di accettazione condividano lo stesso ottimo globale, i piccoli modelli draft, avendo capacità limitata, convergono tipicamente verso soluzioni subottimali in cui minimizzare la KL non garantisce la massimizzazione del tasso di accettazione. Per affrontare questo problema, proponiamo le perdite LK, obiettivi di addestramento speciali che mirano direttamente al tasso di accettazione. Esperimenti completi su quattro architetture draft e sei modelli target, che vanno da 8B a 685B di parametri, dimostrano miglioramenti consistenti nelle metriche di accettazione in tutte le configurazioni rispetto all'addestramento standard basato su KL. Valutiamo il nostro approccio in domini generali, di programmazione e matematici, e riportiamo guadagni fino all'8-10% nella lunghezza media di accettazione. Le perdite LK sono facili da implementare, non introducono overhead computazionale e possono essere integrate direttamente in qualsiasi framework di addestramento speculatore esistente, rendendole un'alternativa convincente agli obiettivi di addestramento draft attuali.
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.