Efficient Streaming Language Models with Attention Sinks

Introduction

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. StreamingLLM employs "attention sinks" to extend the processing capabilities of LLMs to infinite text lengths without additional training.

‍

Challenge in Existing Methods

Traditional methods of managing language models in streaming applications face two pivotal challenges:

1. Excessive Memory Usage: During decoding, Transformer-based LLMs cache the key and value states of all previous tokens, leading to high memory consumption and latency.
2. Limited Text Length Handling: Performance deteriorates as the input length exceeds the model’s training constraints.

To better understand these challenges, consider the figure below, which compares the performance of different attention mechanisms as text length increases.

The Concept of Attention Sinks

In traditional models, the first token often receives a disproportionate amount of attention—dubbed an "attention sink." This phenomenon arises because the SoftMax function, used in calculating attention scores, ensures that these scores sum up to one across all tokens. When many tokens aren't strongly relevant to the context, the model still needs to "dump" attention somewhere—often on the first token, simply because it's globally visible from any point in the sequence. This recognition leads to an essential strategy in StreamingLLM: always retain the attention sink tokens in the key-value (KV) cache.

The Recipe of StreamingLLM

StreamingLLM introduces a straightforward yet effective recipe for managing LLMs in streaming contexts:

• Maintain Attention Sinks: Always include several initial tokens as attention sinks in the KV cache.
• Utilize a Sliding Window KV: This approach helps stabilize the model’s behavior over extended texts.
• Redefine Positional Context: Use positions relative to the cache instead of absolute positions from the original text, enhancing relevance and coherence in generated responses.

Results

StreamingLLM shows remarkable stability and efficiency:

• Dense Attention Failure: Traditional dense attention mechanisms fail when the input surpasses the pre-training window size, leading to significant performance drops.
• Window Attention Collapse: This method fails when the input length exceeds the cache size, especially as initial tokens are evicted.
• Stable Performance by StreamingLLM: Unlike other models, StreamingLLM maintains performance without perplexity spikes, even tested up to 4 million tokens.

Further, StreamingLLM provides a considerable speedup on hardware implementations, such as NVIDIA’s TRT-LLM and mobile devices like the iPhone. The longer the input, the more pronounced the efficiency gains—showing StreamingLLM’s scalability and practicality.