LK Losses: Otimização Direta da Taxa de Aceitação para Decodificação Especulativa
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
Resumo
A decodificação especulativa acelera a inferência em modelos de linguagem grandes (LLM) autoregressivos usando um modelo rascunho leve para propor tokens candidatos que são então verificados em paralelo pelo modelo alvo. O ganho de velocidade é significativamente determinado pela taxa de aceitação, no entanto, o treinamento padrão minimiza a divergência de Kullback-Leibler (KL) como um objetivo substituto. Embora a divergência KL e a taxa de aceitação compartilhem o mesmo ótimo global, modelos rascunho pequenos, com capacidade limitada, normalmente convergem para soluções subótimas onde minimizar a KL não garante maximizar a taxa de aceitação. Para resolver este problema, propomos as perdas LK, objetivos de treinamento especiais que visam diretamente a taxa de aceitação. Experimentos abrangentes envolvendo quatro arquiteturas de rascunho e seis modelos alvo, variando de 8B a 685B de parâmetros, demonstram melhorias consistentes nas métricas de aceitação em todas as configurações em comparação com o treinamento padrão baseado em KL. Avaliamos nossa abordagem nos domínios geral, de programação e matemático e relatamos ganhos de até 8-10% no comprimento médio de aceitação. As perdas LK são fáceis de implementar, não introduzem sobrecarga computacional e podem ser integradas diretamente em qualquer estrutura de treinamento de especuladores existente, tornando-as uma alternativa atraente aos objetivos de treinamento de rascunho atuais.
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.