LLM을 위한 초안 기반 근사 추론

Draft-based Approximate Inference for LLMs

June 10, 2025

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

cs.AI

초록

긴 문맥을 처리하는 대규모 언어 모델(LLM)의 추론 최적화는 Transformer의 이차 계산 복잡도와 선형 메모리 복잡도로 인해 점점 더 중요해지고 있다. 기존의 근사 방법들은 키-값(KV) 캐시 삭제, 희소 주의 메커니즘, 프롬프트 압축 등 토큰 또는 KV 쌍의 중요성을 대략적으로 예측하는 데 의존한다. 본 연구에서는 작은 드래프트 모델을 활용하여 토큰과 KV 쌍의 중요성을 더 정확하게 예측하는 새로운 근사 LLM 추론 프레임워크를 제안한다. 구체적으로, 우리는 제안된 프레임워크의 두 가지 구현을 소개한다: (i) SpecKV는 드래프트 출력을 활용하여 각 KV 쌍의 중요성을 정확히 평가하여 더 효과적인 KV 캐시 삭제를 가능하게 하고, (ii) SpecPC는 드래프트 모델의 주의 활성화를 사용하여 중요하지 않은 프롬프트 토큰을 식별하고 제거한다. 우리가 아는 한, 이는 드래프트 모델을 근사 LLM 추론 가속화에 사용한 첫 번째 연구로, 기존의 무손실 스펙큘레이티브 디코딩의 유용성을 확장한다. 우리는 이론적 및 실증적 분석을 통해 제안 방법의 동기를 설명하고, 드래프트 모델과 타겟 모델의 주의 패턴 간의 강한 상관관계를 보여준다. 긴 문맥 벤치마크에 대한 광범위한 실험을 통해 우리의 방법이 기존 베이스라인보다 더 높은 정확도를 일관되게 달성하면서도 메모리 사용량, 지연 시간, 처리량에서 동일한 개선을 유지함을 보여준다. 우리의 코드는 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.