Inferencia Aproximada Basada en Borradores para Modelos de Lenguaje de Gran Escala

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

Resumen

La optimización de la inferencia para modelos de lenguaje de gran contexto (LLM, por sus siglas en inglés) es cada vez más importante debido a la complejidad cuadrática en cómputo y lineal en memoria de los Transformers. Los métodos de aproximación existentes, como la eliminación de la caché clave-valor (KV), la atención dispersa y la compresión de prompts, suelen basarse en predicciones aproximadas de la importancia de los tokens o pares KV. Proponemos un marco novedoso para la inferencia aproximada de LLM que aprovecha modelos de borrador pequeños para predecir con mayor precisión la importancia de los tokens y los pares KV. Específicamente, presentamos dos instancias de nuestro marco propuesto: (i) SpecKV, que utiliza una salida de borrador para evaluar con precisión la importancia de cada par KV y así mejorar la eliminación de la caché KV, y (ii) SpecPC, que emplea las activaciones de atención del modelo de borrador para identificar y descartar tokens de prompt no importantes. Hasta donde sabemos, este es el primer trabajo que utiliza modelos de borrador para acelerar la inferencia aproximada de LLM, ampliando su utilidad más allá del decodificado especulativo tradicional sin pérdidas. Fundamentamos nuestros métodos con análisis teóricos y empíricos, y demostramos una fuerte correlación entre los patrones de atención de los modelos de borrador y los modelos objetivo. Experimentos exhaustivos en benchmarks de contexto largo muestran que nuestros métodos logran consistentemente una mayor precisión que las líneas base existentes, manteniendo las mismas mejoras en el uso de memoria, latencia y rendimiento. Nuestro código está disponible en 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.