Efficient Streaming Language Models with Attention Sinks

Xiao et al., 2024, Efficient Streaming Language Models with Attention Sinks (The Twelfth International Conference on Learning Representations (ICLR 2024))

Subscriber deck — free during early access

4 cards in “1 Introduction” · whole deck

1 Introduction

  1. What two challenges arise when running a language model on an endless input stream such as a day-long chat?
    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.
  2. Why are window attention and sliding-window recomputation unsatisfactory for streaming?
    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.
  3. What is an attention sink?
    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.
  4. Why do the first tokens, in particular, become attention sinks?
    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.

Source

  • Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han and Mike Lewis. 2024. Efficient Streaming Language Models with Attention Sinks. In The Twelfth International Conference on Learning Representations (ICLR 2024). https://arxiv.org/abs/2309.17453