Reproduce “SparQ Attention: Bandwidth-Efficient LLM Inference” on OpenResearch

About this paper

Luka Ribar, Ivan Chelombiev, Luke Hudlass-Galley, Charlie Blake, Carlo Luschi, Douglas Orr

SparQ Attention speeds language-model inference by estimating which parts of the key-value cache matter for each new token and selectively fetching them, while keeping the full cache available. On Llama 2 13B’s SQuAD task with seven unrelated passages, it achieves 74.9% exact-match accuracy at one-eighth of dense attention transfers, versus 80.8% for dense attention; this is not an eightfold speedup.

Key insights

• Query vectors in LLMs exhibit heavy-tailed distributions, enabling effective approximation of dot products by considering only a few salient components.
• Attention scores are inherently sparse, meaning only a small subset of tokens contributes significantly to the output, which can be exploited for selective data retrieval.
• Maintaining a running mean value vector effectively compensates for the 'missing mass' of attention from less important tokens, preserving overall approximation accuracy without information loss.

Summary by alphaXiv.