Efficient Streaming Language Models with Attention Sinks

The KV cache of all previous tokens keeps growing, using more memory and slowing decoding; and models degrade on sequences longer than their pre-training attention window.

Window attention (cache only the last \(L\) tokens’ keys and values) is efficient but collapses as soon as the first tokens are evicted. Recomputing the KV states of the recent window for each new token works well but costs \(O(TL^2)\), far too slow.

Initial tokens that receive a surprisingly large share of attention in most layers and heads regardless of their meaning, because softmax forces attention weights to sum to one: when a query matches nothing, the excess attention has to go somewhere.

In autoregressive models the first tokens are visible to every later token, while later tokens are visible only to those after them, so the initial ones are the easiest to train as a place to dump unneeded attention.

Their keys took a large part of the softmax denominator in most layers; removing them shifts the attention distribution far from what the model saw in training, and perplexity spikes.

Replacing the first four tokens with line-break tokens and keeping them alongside the recent window restored perplexity almost as well as the original tokens: the model attends to them because of where they are.

It keeps the keys and values of a few initial tokens (four suffice) as attention sinks, plus a rolling window of the most recent tokens, evicting everything in between. No fine-tuning is needed.

By position within the cache, not in the original text: with cached tokens 0–3 and 6–8 while decoding token 9, positions are 0–7. For RoPE, keys are cached before rotation and rotated by their cache position at each step; ALiBi gets a contiguous bias.

Models like Llama-2 weren’t pre-trained with a consistent first token in every sample (the start token is added before text is chunked), so they learned to spread sink attention over several initial tokens; one or two don’t fully restore perplexity.

Prepend a learnable sink token to every training sample. In 160M-parameter models it became a dedicated attention sink: the sink token plus the recent window gave stable streaming perplexity, with no loss in convergence or benchmark accuracy.

Letting attention sum to less than one (equivalent to a zero key and value token) helped only somewhat; the model still relied on other initial tokens. A learnable sink token worked much better.

Llama-2, MPT, Falcon and Pythia models (2.8B–70B) kept stable perplexity over 4 million tokens of PG19 books, matching the sliding-window-with-recomputation oracle, while dense attention failed past the pre-training length and window attention past the cache size.

No: increasing the cache size didn’t consistently lower perplexity, suggesting the models don’t fully use the context they get. StreamingLLM also doesn’t extend the context window: tokens evicted from the cache are forgotten.

Its per-token decoding latency grows linearly with cache size versus quadratically for recomputation, giving up to 22.2× speedup per token with a similar memory footprint.