Inferência Aproximada Baseada em Rascunho para LLMs

Draft-based Approximate Inference for LLMs

June 10, 2025

Autores: Kevin Galim, Ethan Ewer, Wonjun Kang, Minjae Lee, Hyung Il Koo, Kangwook Lee

cs.AI

Resumo

A otimização da inferência para modelos de linguagem de grande escala (LLMs) de contexto longo é cada vez mais importante devido à complexidade quadrática de computação e linear de memória dos Transformers. Métodos de aproximação existentes, como o descarte do cache de chave-valor (KV), atenção esparsa e compressão de prompt, geralmente dependem de previsões grosseiras da importância de tokens ou pares KV. Propomos uma nova estrutura para inferência aproximada de LLMs que utiliza modelos de rascunho pequenos para prever com maior precisão a importância de tokens e pares KV. Especificamente, introduzimos duas instanciações da nossa estrutura proposta: (i) SpecKV, que utiliza uma saída de rascunho para avaliar com precisão a importância de cada par KV para um descarte mais eficaz do cache KV, e (ii) SpecPC, que usa as ativações de atenção do modelo de rascunho para identificar e descartar tokens de prompt irrelevantes. Até onde sabemos, este é o primeiro trabalho a utilizar modelos de rascunho para acelerar a inferência aproximada de LLMs, estendendo sua utilidade além da decodificação especulativa tradicional sem perdas. Fundamentamos nossos métodos com análises teóricas e empíricas, e mostramos uma forte correlação entre os padrões de atenção dos modelos de rascunho e alvo. Experimentos extensos em benchmarks de contexto longo demonstram que nossos métodos consistentemente alcançam maior precisão do que as linhas de base existentes, mantendo as mesmas melhorias no uso de memória, latência e taxa de transferência. Nosso código está disponível em https://github.com/furiosa-ai/draft-based-approx-llm.

English

Optimizing inference for long-context Large Language Models (LLMs) is increasingly important due to the quadratic compute and linear memory complexity of Transformers. Existing approximation methods, such as key-value (KV) cache dropping, sparse attention, and prompt compression, typically rely on rough predictions of token or KV pair importance. We propose a novel framework for approximate LLM inference that leverages small draft models to more accurately predict the importance of tokens and KV pairs. Specifically, we introduce two instantiations of our proposed framework: (i) SpecKV, which leverages a draft output to accurately assess the importance of each KV pair for more effective KV cache dropping, and (ii) SpecPC, which uses the draft model's attention activations to identify and discard unimportant prompt tokens. To the best of our knowledge, this is the first work to use draft models for approximate LLM inference acceleration, extending their utility beyond traditional lossless speculative decoding. We motivate our methods with theoretical and empirical analyses, and show a strong correlation between the attention patterns of draft and target models. Extensive experiments on long-context benchmarks show that our methods consistently achieve higher accuracy than existing baselines, while preserving the same improvements in memory usage, latency, and throughput. Our code is available at https://github.com/furiosa-ai/draft-based-approx-llm.