Reproduce “Draft-based Approximate Inference for LLMs” on OpenResearch

About this paper

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

The "Draft-based Approximate Inference" framework improves efficiency for large language models by employing small draft models to provide lookahead predictions, which enables more accurate estimation of token and Key-Value (KV) pair importance for compression and dropping. This approach yields higher accuracy and substantial reductions in latency and memory usage on long-context tasks compared to previous approximate inference methods.

Key insights

• Lookahead tokens generated by a small draft model provide reliable approximations of future outputs, which can be leveraged for more accurate importance estimation in approximate inference for larger target models.
• There is a strong correlation between importance scores computed by draft models' attention activations and those from target models, justifying the use of draft model attention for efficient prompt compression.
• Cascading prompt compression (SpecPC) with KV cache dropping (SpecKV) synergistically enhances both accuracy and efficiency, as initial prompt reduction creates a more optimal input for subsequent KV cache management.

Summary by alphaXiv.