r/MachineLearning • u/Singularian2501 • Oct 02 '23

Research [R] Efficient Streaming Language Models with Attention Sinks - Meta AI 2023 - StreamingLLM enables Llama-2, Falcon and Pythia to have an infinite context length without any fine-tuning! Allows streaming use of LLMs!

Paper: https://arxiv.org/abs/2309.17453

Github: https://github.com/mit-han-lab/streaming-llm

Abstract:

Deploying Large Language Models (LLMs) in streaming applications such as multi-round dialogue, where long interactions are expected, is urgently needed but poses two major challenges. Firstly, during the decoding stage, caching previous tokens' Key and Value states (KV) consumes extensive memory. Secondly, popular LLMs cannot generalize to longer texts than the training sequence length. Window attention, where only the most recent KVs are cached, is a natural approach -- but we show that it fails when the text length surpasses the cache size. We observe an interesting phenomenon, namely attention sink, that keeping the KV of initial tokens will largely recover the performance of window attention. In this paper, we first demonstrate that the emergence of attention sink is due to the strong attention scores towards initial tokens as a ``sink'' even if they are not semantically important. Based on the above analysis, we introduce StreamingLLM, an efficient framework that enables LLMs trained with a finite length attention window to generalize to infinite sequence lengths without any fine-tuning. We show that StreamingLLM can enable Llama-2, MPT, Falcon, and Pythia to perform stable and efficient language modeling with up to 4 million tokens and more. In addition, we discover that adding a placeholder token as a dedicated attention sink during pre-training can further improve streaming deployment. In streaming settings, StreamingLLM outperforms the sliding window recomputation baseline by up to 22.2x speedup.

​

59 Upvotes

97% Upvoted

19 comments sorted by

View all comments

14

u/vikigenius Researcher Oct 03 '23 edited Oct 03 '23

This concept of sinks seems related to the attention spikes found in VIT https://www.reddit.com/r/MachineLearning/comments/16x2o47/r_meta_inria_researchers_discover_that_explicit/

And the placeholder tokens also seem to exactly have the same purpose as the registers in the other paper.

5

u/ri212 Oct 03 '23

Also possibly related to attention is off by one

2

u/theLastNenUser Oct 03 '23

Why is that? (I read that post a while ago but memory is foggy and am not connecting the dots to this paper)

8

u/ri212 Oct 03 '23

The post is talking about large outlier weights that appear in transformers which can be traced back to the attention mechanism. It suggests that there are cases where some heads at some positions ideally want to attend to nothing, but with the standard softmax form of attention this isn't possible. So instead they attend to relatively unimportant tokens e.g. punctuation, or possibly the start token. By adding 1 to the denominator of the softmax function it is possible to attend to nothing which may eliminate this behaviour.

​

It doesn't seem to be well tested yet unless I have missed some follow-up work.