PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Streaming Language Models with Attention Sinks

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, Mike Lewis

arXiv:2309.17453v4cs.CLcs.AI

TL;DR

Streaming LLMs face memory costs from caching prior Key and Value states and fail to generalize beyond finite training windows. The paper introduces StreamingLLM, which preserves attention sinks alongside recent-token KV states, supports sequences up to 4 million tokens, and achieves up to 22.2× speedup over recomputation.

  • Problem

    LLMs trained with finite attention windows struggle with persistent long-sequence streaming, while caching prior Key and Value states consumes extensive memory.

  • Method

    StreamingLLM retains the KV states of initial attention-sink tokens together with a rolling cache of recent tokens, without fine-tuning.

  • Results

    StreamingLLM reliably models up to 4 million tokens across Llama-2, MPT, Falcon, and Pythia, with up to 22.2× per-token speedup over sliding-window recomputation.

  • Takeaways & Limitations

    Preserving initial attention sinks decouples the training attention-window size from actual generation length and enables efficient streaming deployment.

  • Takeaways & Limitations

    Increasing cache size does not consistently reduce language-modeling perplexity, suggesting limited use of the full available context.

Abstract

from arXiv · show

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. Code and datasets are provided at https://github.com/mit-han-lab/streaming-llm.

1 INTRODUCTION

Streaming LLM deployment is limited by KV-cache memory costs and finite-length generalization, while window attention fails after evicting initial tokens. StreamingLLM preserves those attention sinks alongside recent tokens to stabilize efficient generation on extended texts.

  • 1 INTRODUCTION: LLM streaming faces excessive KV-cache memory use and performance degradation beyond the finite attention window used during pre-training.These limitations obstruct persistent long-sequence deployments.
  • 1 INTRODUCTION: Window attention maintains constant-size recent-token caches but collapses when the sequence exceeds cache capacity and initial-token KV is evicted.Sliding-window recomputation retains strong performance but is significantly slower because it repeatedly computes quadratic attention within the window.
  • 1 INTRODUCTION: Attention sinks are initial tokens that receive unusually large attention scores despite lacking semantic significance.Softmax requires attention scores to sum to one, and globally visible initial tokens become readily trained locations for excess attention.
  • 1 INTRODUCTION: StreamingLLM keeps several initial attention-sink KVs together with recent-token KVs, enabling finite-window models to process arbitrarily long text without fine-tuning.The framework reports reliable modeling of 4 million tokens and more across Llama-2, MPT, Falcon, and Pythia.
  • 1 INTRODUCTION: Pre-training with an extra learnable sink token can preserve streaming performance using one dedicated sink instead of multiple naturally occurring initial-token sinks.The finding comes from pre-training 160-million-parameter language models from scratch.
  • 1 INTRODUCTION: StreamingLLM supports continuous operation with minimal memory use and can complement context-extension methods to increase attendable recent context.It generates coherently from tokens within the KV cache without extending the model’s context length.

2 RELATED WORK

Related work addresses long-text language modeling through length extrapolation, context-window extension, and improved utilization of long text. These directions target different limitations and do not necessarily produce competent use of lengthy contexts.

  • 2 RELATED WORK: Length extrapolation develops position-encoding methods that let models trained on shorter texts handle longer sequences during testing.Rotary Position Embeddings transform queries and keys in every attention layer for relative-position integration.
  • 2 RELATED WORK: Context-window extension expands the number of tokens processed in one forward pass but faces quadratic attention computation and substantial training memory demands.The cited work includes system-focused optimizations such as FlashAttention.
  • 2 RELATED WORK: Dense attention fails beyond the pre-training window, whereas window attention collapses beyond cache capacity and StreamingLLM nearly matches sliding-window recomputation in perplexity.Figure 3 reports these language-modeling trends on 20K-token texts across various LLMs.
  • 2 RELATED WORK: Improving long-text utilization focuses on capturing and using contextual content rather than merely increasing the available context length.The passage identifies effective use of prolonged contexts as an ongoing challenge.

3 STREAMINGLLM

Window attention fails when evicting initial-token KV states because those tokens act as attention sinks that help preserve the normal attention-score distribution. StreamingLLM combines these sink KVs with a rolling cache, while dedicated sink-token pre-training can further stabilize streaming.

  • The Failure of Window Attention and Attention Sinks: Window attention’s perplexity spikes after the sequence exceeds the cache size because removing initial-token KV states disrupts attention computation.The SoftMax denominator loses a considerable contribution from initial tokens, shifting attention scores away from normal inference.
  • The Failure of Window Attention and Attention Sinks: Initial tokens serve as attention sinks despite limited semantic relevance, because autoregressive training makes them visible to nearly all subsequent tokens.The model consequently learns to allocate substantial attention to these globally visible positions.
  • The Failure of Window Attention and Attention Sinks: Four initial tokens generally suffice to restore streaming perplexity, whereas one or two are insufficient and additional tokens yield diminishing returns.Table 2 evaluates cache configurations on 400K tokens from the concatenated PG19 test set.
  • The Failure of Window Attention and Attention Sinks: Reintroducing four initial tokens alongside 1020 recent tokens restores perplexity, and replacing them with linebreak tokens achieves comparable recovery.These results are measured on the first 65K-token PG19 book.
  • Rolling KV Cache with Attention Sinks: StreamingLLM keeps four attention-sink KVs together with a rolling KV cache of recent tokens to stabilize attention while retaining current language-modeling context.Its cache is conceptually divided into attention sinks and a rolling cache; positional information is assigned within the cache, with transformations handled for RoPE or ALiBi.
  • Pre-training LLMs with Attention Sinks: Pre-training with a learnable Sink Token produces stable streaming perplexity using only that sink token, unlike vanilla attention, which requires several initial tokens.Zero Sink provides only a partial improvement and still relies on other initial tokens.

4 EXPERIMENTS

Experiments evaluate StreamingLLM across model families, long-text language modeling, streaming question answering, sink-token pre-training, and efficiency. The results show stable long-context performance, competitive streaming accuracy, and substantially lower decoding latency than recomputation.

  • 4.1 LANGUAGE MODELING ON LONG TEXTS ACROSS LLM FAMILIES AND SCALES: StreamingLLM reliably handles more than 4 million tokens across Llama-2, Falcon, Pythia, and MPT model families and scales.Perplexity remains stable throughout the concatenated PG19 test set, with fluctuations attributed to book transitions.
  • 4.1 LANGUAGE MODELING ON LONG TEXTS ACROSS LLM FAMILIES AND SCALES: StreamingLLM matches the sliding-window recomputation oracle in perplexity on texts spanning 20K tokens.Dense attention fails beyond the pre-training window, while window attention struggles after evicting initial tokens.
  • 4.2 PRE-TRAINING WITH ATTENTION SINKS: Adding a dedicated sink token preserves similar convergence and standard NLP benchmark performance while improving streaming performance.A sink-token-trained model achieves satisfactory streaming performance with one sink token, unlike the vanilla model, which requires multiple tokens.
  • 4.3 RESULTS ON STREAMING QUESTION ANSWERING WITH INSTRUCTION-TUNED MODELS: StreamingLLM achieves comparable accuracy to the one-shot sample-by-sample baseline in concatenated ARC question answering.Dense attention produces out-of-memory errors, while window attention yields low accuracy after the input exceeds the cache size.
  • 4.4 ABLATION STUDIES: Four initial tokens are sufficient as attention sinks, while adding more produces only marginal effects.One or two initial tokens are insufficient for stable streaming perplexity.
  • 4.4 ABLATION STUDIES: Increasing cache size does not consistently reduce perplexity, suggesting that models may not fully use the available context.The result identifies a limitation in exploiting extensive contexts.
  • 4.5 EFFICENCY RESULTS: 22.2× per-token speedup is achieved over sliding-window recomputation while maintaining a similar memory footprint.StreamingLLM latency grows linearly with cache size, whereas recomputation latency rises quadratically.

5 CONCLUSION

The conclusion presents StreamingLLM as a way to support long streaming generation by retaining attention sinks with recent tokens. It reports efficient modeling of texts up to 4 million tokens and improved streaming deployment with dedicated sink-token pre-training.

  • 5 CONCLUSION: StreamingLLM enables fixed-window LLMs to handle unlimited texts without fine-tuning by combining attention sinks with recent tokens.The framework decouples the pre-training window size from actual text generation length.
  • 5 CONCLUSION: StreamingLLM efficiently models texts of up to 4 million tokens and can improve streaming performance when models are pre-trained with a dedicated sink token.The conclusion frames this as a route toward streaming deployment of LLMs.

REPRODUCIBILITY STATEMENT

The paper states that its findings are reproducible through released code, datasets, model references, and experimental details.

  • REPRODUCIBILITY STATEMENT: Released code, datasets, model references, hyperparameters, training protocols, and evaluation methods support reproduction of the presented results.The authors state that all findings can be reproduced.

IMPACT STATEMENT

The paper describes StreamingLLM’s practical deployment relevance, adoption, related architectural connections, and scope limitations. It emphasizes continuous short-term-context applications while distinguishing them from long-term-memory tasks.

  • IMPACT STATEMENT: StreamingLLM has been adopted by LLM serving solutions including TensorRT-LLM, Intel Extension for Transformers, HuggingFace Transformers, and MLC LLM.
  • IMPACT STATEMENT: StreamingLLM is particularly suited to continuous applications such as multi-round dialogue that rely on recent interactions.It avoids frequent cache refreshes during extended operation.
  • IMPACT STATEMENT: StreamingLLM does not extend context windows or provide long-term memory, limiting its suitability for long-document question answering and summarization.The method is positioned for short-term-memory settings such as daily conversations and short-document QA.
  • IMPACT STATEMENT: Dedicated sink tokens relate attention sinks in autoregressive models to attention concentration phenomena such as registers in Vision Transformers.The paper distinguishes sink tokens as initial autoregressive tokens rather than intermediate-layer global information holders.

C ACCURACY ON STREAMEVAL WITH INCREASING QUERY-ANSWER LINE DISTANCE

StreamingLLM retains accuracy across increasing query-answer distances, but performance eventually collapses when the distance exceeds cache capacity. Its effectiveness therefore depends on the information retained in the cache.

  • StreamingLLM retains accuracy as query-answer line distance increases, until the distance surpasses cache capacity.Accuracy eventually drops to zero beyond the cache capacity.
  • StreamingLLM’s performance depends on the information available within its cache.The supplied evidence identifies cache capacity as the boundary for reliable long-distance answering.

D LONG-RANGE BENCHMARK EVALUATION

On LongBench, StreamingLLM’s performance depends on preserving crucial initial prompt information. Increasing the retained initial segment to match the truncation baseline restores comparable performance.

  • StreamingLLM 4+3496 underperforms the default truncation baseline because crucial initial input prompt information is lost.The baseline retains 1750 initial and 1750 final tokens, whereas 4+3496 retains only four initial tokens.
  • StreamingLLM 1750+1750 restores performance to the level of the text truncation baseline.This configuration aligns the attention sink count with the baseline’s 1750 initial tokens.
  • The LongBench evaluation covers single-document QA, multi-document QA, and summarization.The evaluated tasks include NarrativeQA, Qasper, HotpotQA, 2WikiMQA, GovReport, and MultiNews.

F QUATITATIVE ANALYSIS OF ATTENTION SINKS IN LONG INPUTS

Long-input analyses show that initial tokens receive disproportionately high attention in Llama-2 models. In lengthy sequences, the first token often receives more than half of total attention in most layers.

  • In 4096-token inputs, the first token often receives more than half of the total attention in most layers.The measurement uses the 4096th token’s attention toward the first token, averaged across 256 sequences.
  • Llama-2-70B likewise assigns much higher attention scores to initial tokens than to the rest in most layers.This extends the attention-sink observation from Llama-2-7B to the 70B model.

H ATTENTION SINKS IN ENCODER TRANSFORMERS

Attention sinks also appear in encoder-style Transformers: BERT-base-uncased assigns disproportionately high attention to the ubiquitous [SEP] token. The paper relates this pattern to register-like attention spikes reported in Vision Transformers.

  • The paper suggests that attention sinks may extend beyond autoregressive decoder-only language models to other Transformer architectures.This proposal is based on shared Transformer structure and SoftMax attention mechanisms.
  • BERT-base-uncased assigns disproportionately high attention scores to the [SEP] token in most layers.The [SEP] token functions as an omnipresent focal point for attention in the visualization.
  • The paper relates Vision Transformer register tokens to the attention-sink phenomenon observed in language models.The cited comparison attributes Vision Transformer attention spikes to random background patch tokens acting as registers for global image information.

I USING MORE SINK TOKENS IN THE PRE-TRAINING STAGE

Adding one or two dedicated sink tokens during pre-training preserves baseline-like pre-training loss, but additional sink tokens do not materially improve benchmark or streaming performance.

  • Pre-training loss: One or two sink tokens produce pre-training loss curves closely resembling the vanilla baseline.Figure 15 compares models with 0, 1, and 2 sink tokens.
  • Streaming performance: Additional sink tokens do not enhance streaming performance, while a single sink token is adequate for improving it.The analysis reports that the model may rely on both tokens for stability, but adding more does not further improve overall language-model performance.
  • Benchmark performance: A second sink token does not yield substantial improvements across most benchmark tasks.Table 9 reports zero-shot accuracy across seven NLP benchmarks.
Loading 2309.17453v4…