Inferenza Approssimata Basata su Bozze per Modelli Linguistici di Grande Dimensione

Draft-based Approximate Inference for LLMs

June 10, 2025

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

cs.AI

Abstract

L'ottimizzazione dell'inferenza per i modelli linguistici di grandi dimensioni (LLM) con contesti lunghi sta diventando sempre più importante a causa della complessità computazionale quadratica e della complessità di memoria lineare dei Transformer. I metodi di approssimazione esistenti, come l'eliminazione della cache chiave-valore (KV), l'attenzione sparsa e la compressione del prompt, si basano tipicamente su previsioni approssimative dell'importanza dei token o delle coppie KV. Proponiamo un nuovo framework per l'inferenza approssimata degli LLM che sfrutta piccoli modelli di bozza per prevedere con maggiore precisione l'importanza dei token e delle coppie KV. Nello specifico, introduciamo due istanze del nostro framework proposto: (i) SpecKV, che utilizza un output di bozza per valutare accuratamente l'importanza di ciascuna coppia KV per un'eliminazione più efficace della cache KV, e (ii) SpecPC, che utilizza le attivazioni di attenzione del modello di bozza per identificare e scartare i token del prompt non importanti. Per quanto ne sappiamo, questo è il primo lavoro che utilizza modelli di bozza per accelerare l'inferenza approssimata degli LLM, estendendo la loro utilità oltre il tradizionale decoding speculativo senza perdita. Motiviamo i nostri metodi con analisi teoriche ed empiriche, e mostriamo una forte correlazione tra i pattern di attenzione dei modelli di bozza e quelli target. Esperimenti estesi su benchmark a contesto lungo dimostrano che i nostri metodi raggiungono costantemente una maggiore accuratezza rispetto alle baseline esistenti, preservando allo stesso tempo i miglioramenti nell'uso della memoria, nella latenza e nel throughput. Il nostro codice è disponibile all'indirizzo 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.