Paper: Efficient Streaming Language Models with Attention Sinks (ICLR-2024)

Reviewer comment (Experiments):
The authors focus on LLMs that are autoregressive, meaning that they generate text one token at a time, conditioned on the previous tokens. Dose the attention sink phenomenon exist in other types of models, such as bidirectional models (e.g., BERT) or non-transformers models?

Author response:
**Q: Attention Sink Phenomenon in Non-Autoregressive Models**

A: The attention sink phenomenon is not exclusive to autoregressive models, which we focused on in our study. Appendix Section H of our revised paper shows that encoder models, like BERT, also exhibit this phenomenon. Concurrent to our work, Darcet et al. observe similar attention concentration on random background patch tokens in Vision Transformers, termed "registers." This suggests that the tendency to allocate significant attention to sink tokens might be a broader characteristic of transformer architectures, irrespective of whether they are autoregressive or bidirectional.

Reviewer comment (Experiments):
Can attention sink be applied to the pre-training stage, retaining more placeholders from scratch to further improve the performance?

Author response:
**Q: Using more Attention Sinks in the Pre-Training Stage**

A: Thanks for your suggestion. We added an experiment where we pre-trained the 160M language model with two sink tokens, detailed in Appendix Section I. Our experiments indicated that adding more than one attention sink token does not significantly enhance performance. This finding suggests that a single dedicated attention sink during pre-training is sufficient to leverage the benefits of this phenomenon.

Reviewer comment (Experiments):
Does this phenomenon also exist in larger models like Falcon-40B, LLaMa2-70B?

Author response:
**Q: Attention Sink Phenomenon in Larger Models**

A: Yes, this phenomenon persists in larger models like Falcon-40B and Llama2-70B. We conducted perplexity experiments in Figure 5 to confirm this and included visualizations in Appendix Section G of our manuscript. These results demonstrate that the attention sink phenomenon exists across different model sizes, and StreamingLLM is scalable to all of them.

Reviewer comment (Experiments):
Does the Group Query Attention structure have any impact?

Author response:
**Q: Impact of Multi/Group Query Attention Structure**

A: The attention sink phenomenon was still observable in our experiments, including models like Falcon-40B using multi-query attention and Llama-2-70B using group-query attention. This suggests that different attention structures, such as group query attention, do not negate the presence of attention sinks (Appendix G). StreamingLLM effectively addresses this issue despite the architecture differences (Figure 5), demonstrating its versatility across various model architectures.

Reviewer comment (Experiments):
If you combine proposed Attention Sink with Sliding Window w re-computation, I wonder whether it performs better than re-computation.

Author response:
**Q: Combining Attention Sink with Sliding Window with Re-computation**

A: Combining attention sinks with sliding windows with re-computation does not provide additional benefits over re-computation alone. This is because the sliding window with re-computation inherently includes full attention on truncated text, meaning attention sinks are always present. This is why sliding windows with re-computation perform better than the window attention baseline. Our approach with StreamingLLM focuses on efficiency and speed, significantly reducing the computational complexity compared to this baseline.

Reviewer comment (Experiments):
StreamingLLM is compared with only very basic baselines (i.e., (a) dense attention, (b) window attention, and (c) sliding window with re-computation) without alternative ways to use LLMs with longer inputs. The authors argue StreamingLLM’s 22.2x speedup compared to (c), but it looks too extreme because there could be a sweet spot that could achieve reasonable ppl and inference speed by sliding window with appropriate strides (less frequent re-computation instead of every step).

Author response:
We appreciate your point about comparing StreamingLLM with sliding window approaches with varying strides. Our 22.2x speedup comparison was established under identical conditions to ensure a fair evaluation. StreamingLLM's linear $O(N)$ complexity contrasts sharply with the sliding window's quadratic $O(N^2)$ complexity. While adjusting stride lengths in sliding windows may offer some efficiency, it introduces context size and computational workload variability, resulting in inconsistent decoding times and greater computational demands during re-computation phases. This inconsistency is a significant drawback compared to StreamingLLM's consistent performance and computational efficiency.

Reviewer comment (Experiments):
The paper argues that StreamingLLM can handle infinite sequence length, but again the abovementioned basic baselines on PG-19 are used to show this. Experiments on benchmark datasets that truly require long context modeling are necessary.

Author response:
It is crucial to highlight that StreamingLLM is not designed to extend the context window or enhance long-term memory capabilities. Instead, its primary aim is efficient text generation from recent tokens without frequent cache refreshes. We have included further clarification on this in Appendix Section A, to articulate StreamingLLM's intended application better and differentiate it from long-context models.

Reviewer comment (Novelty):
There have been several previous works that observe position bias on attention patterns in transformers. Also, there are existing architectures, including LongFormer, BigBird, and ETC, that have global tokens for attention, though most of them are encoder-based language models. It is worth having a discussion on them.

Author response:
Thank you for pointing out the need for a more detailed comparison with related works. We have added Appendix Section B to include a comprehensive discussion on architectures like LongFormer, BigBird, and ETC that employ global tokens for attention. This enhanced comparison delineates the unique aspects of our approach, emphasizing its simplicity and compatibility with existing transformer architectures while also recognizing the foundational concepts established by these earlier models.

Reviewer comment (Theory):
Could you elaborate more on why (1) window attention blows up and (2) attention sink exists? Current explanations are mostly empirical and not theoretically grounded. Specifically, attention scores could be skewed toward earlier tokens due to the autoregressive nature of language models, but why are very few (less than 4) tokens enough to cover most attention scores?

Author response:
The issue of window attention blowing up is linked to the dynamics of the SoftMax function used in attention mechanisms. As detailed in Appendix Section F, more than half of attention scores in most layers and heads are allocated to the first token. Over half of the SoftMax function's denominator is removed when these initial tokens are evicted. This significant alteration in the attention score distribution deviates from the model's pre-training conditions, leading to a drastic change in the SoftMax distribution. Consequently, this results in model performance degradation, as the attention mechanism no longer functions as it was trained to.

The attention sink phenomenon primarily arises from the characteristics of the SoftMax function used in attention mechanisms. SoftMax necessitates that all attended tokens’ attention scores add up to one, forcing the model to distribute some attention across contextual tokens even when the current token embedding already contains ample self-relevant information for prediction or lacks a strong correlation with many preceding keys. As a result, the model learns to allocate excessive attention scores to certain tokens.

Given the sequential nature of autoregressive language modeling, initial tokens are consistently visible to all subsequent tokens. In contrast, later tokens have visibility restricted to fewer subsequent tokens. This visibility pattern predisposes initial tokens to become attention sinks as they are capable of absorbing unnecessary attention from a wider range of tokens.

We empirically observe that the number of tokens required to function as effective attention sinks are relatively small. A few tokens are typically sufficient to absorb excess attention. The precise reason why this limited number of tokens is adequate remains an area for future theoretical research.

Reviewer comment (Evaluation):
Visualization on a short sequence (length 16) in Figure 2 might not be enough. Does the same phenomenon occur for longer sequences, which is the main interest? Some important tokens with high attention scores might appear for each text (though not shown if averaging per position over different texts). Are there any quantitative values on how much attention sink occupies the attention scores, especially for longer inputs?

Author response:
The attention sink phenomenon persists across all sequence lengths. For clarity in visualization, we focused on shorter sequences in the paper’s figure. In Appendix Section E of our updated manuscript, we include attention map visualizations for longer sequences, illustrating this phenomenon's consistency. Appendix Section F also includes quantitative results on the ratio of attention sinks’ attention scores on long sequences with 4096 tokens. 

Thank you for your suggestion! We conducted experiments to quantify the attention sinks' score ratio in the Llama models in Appendix Section F. We plotted attention allocated by the 4096th token to the initial token in every layer. Notably, the attention scores for the first token are significantly high, often exceeding half of the total attention, except for the two bottom layers. This observation empirically substantiates the preferential focus on the first token by most layers and heads, irrespective of other tokens' distances within the sequence. Such a trend underscores the critical role of the initial tokens in a sequence, as their removal greatly impacts language model performance due to a large portion of the denominator in the SoftMax function being removed.

Reviewer comment (Evaluation):
In Figure 3, MPT has different patterns from other LLMs. Is it because MPT uses ALiBi and others use RoPE? Any explanation?

Author response:
The general behavior of the MPT model aligns with that observed in RoPE models, particularly in how perplexity sharply rises when the input length surpasses the pre-training length in dense attention scenarios and similarly when it exceeds the cache size in window attention contexts.

However, the specific perplexity curve pattern in MPT differs due to a combination of factors, including variations in pre-training approaches, model architecture, training datasets, positional encoding, and activation functions. A comprehensive analysis of these discrepancies and their underlying causes is complex and falls outside the scope of our current work. We acknowledge this as an area for future research to explore in more detail.

Reviewer comment (Experiments):
Regarding pre-training with a sink token, reducing 4 -> 1 might not be important considering the large size of the cache. Why is a model with a sink token consistently better than the vanilla one? Would the same conclusion hold for larger LMs beyond the 160M-parameter one used due to academic budget constraints?

Author response:
Our pre-training experiments aimed to test whether introducing a globally visible sink token would direct excessive attention to this token, thus eliminating other attention sinks. This experiment validated our hypothesis, showing that a single token is enough to recover the streaming perplexity, and the model no longer relies on other initial tokens. While our study used 160M LM due to budget constraints, we expect our pre-training solution to generalize to larger models. Future research with more resources can test our findings in larger LLMs.

Reviewer comment (Evaluation):
In Figure 10, why do two options have different memory usage for 4096 cache size?

Author response:
The slightly higher memory usage in sliding windows with recomputation is attributed to the need for intermediate variables for the full QxK quadratic attention computation. In contrast, StreamingLLM efficiently computes only the qxK attention, reducing overall memory requirements.

Reviewer comment (Novelty):
There is almost similar concurrent work: LM-Infinite (https://arxiv.org/abs/2308.16137). How does this work compare with yours?

Author response:
LM-Infinite (Han et al., 2023) offers valuable theoretical insights into LLM length generation failure, identifying three out-of-distribution factors. Inspired by this analysis, their approach involves employing a ""$\Lambda$-shaped" attention pattern and reconfiguring position encoding distances to enhance length generalization in LLMs. This approach bears a resemblance to our methodology. 

Distinctly, our work uncovers the "attention sink" phenomenon by looking into the attention maps of LLMs, wherein Transformer models assign high attention scores to initial tokens with minor semantics. This phenomenon extends beyond the scope of length generalization failure, indicating a more pervasive issue in Transformer models. We observe this "attention sink" behavior not only in auto-regressive language models but also in encoder Transformers such as BERT (see Appendix Section H) and Vision Transformers (Darcet et al., 2023), suggesting its broader prevalence in Transformer architectures. To mitigate the “attention sink” phenomenon, we propose the introduction of a learnable sink token during pre-training, and we support our findings with extensive ablation studies. We believe our work and LM-Infinite contribute unique insights to the same topic, providing valuable empirical evidence and theoretical perspectives that can inform future developments in Transformer model architecture.

We have added this discussion in Appendix Section B in our revised paper.

[1] Darcet, T., Oquab, M., Mairal, J., & Bojanowski, P. (2023). Vision Transformers Need Registers. *ArXiv*. /abs/2309.16588

[2] Han, C., Wang, Q., Xiong, W., Chen, Y., Ji, H., & Wang, S. (2023). LM-Infinite: Simple On-the-Fly Length Generalization for Large Language Models. *ArXiv*. /abs/2308.16137

Reviewer comment (Novelty):
The proposed attention is reminiscent of methods that introduce sparse attention patterns such as Sparse Transformer but there is no in-depth discussion that draws a connection.

Author response:
**Q: Missing Discussion with Related Works**

A: We value your feedback on contextualizing our work with prior literature. In response, we have added Appendix Section B to include a comprehensive discussion that contrasts our attention sink mechanism with sparse attention patterns found in models like Sparse Transformer, LongFormer, and BigBird. This enhanced comparison delineates the unique aspects of our approach, emphasizing its simplicity and compatibility with existing transformer architectures while also recognizing the foundational concepts established by these earlier models.

Reviewer comment (Evaluation):
Even though attention sinks maintains the perplexity levels in check for extremely large sequence lengths, the paper does not study in detail on how good utilization of the context the model is doing.

Author response:
**Q: Utilization of the Context**

A: We acknowledge the importance of evaluating how effectively the model utilizes context, especially when introducing attention sinks. In our initial experiments, we observed satisfactory context utilization within the cache limits, shown as the language modeling performance is comparable with the oracle baseline, i.e., sliding window with re-computation.

To analyze the extent of context utilization beyond these cache limits, we conducted further experiments, detailed in Appendix Section C, to assess context utilization when dealing with cache-exceeding information. The findings corroborate that StreamingLLM’s capacity to utilize context is limited to the current cache, with no ability to leverage evicted tokens. It's crucial to emphasize that StreamingLLM's primary aim is not to augment the context window or boost long-term memory retention. Instead, its strength lies in generating coherent text from recent tokens without necessitating cache refreshes.

Moreover, our experiments demonstrate that StreamingLLM can be integrated with recent advancements in context extension methods. With context-extended models, we can use larger caches to maintain a broader range of recent tokens in StreamingLLM. For a demonstration of StreamingLLM's integration with models like LongChat-7B-v1.5-32K and Llama-2-7B-32K-Instruct, please refer to Figure 9 in our paper.

Reviewer comment (Experiments):
Have you considered evaluating on a benchmark for long-range tasks such as LongBenh, MuLD or Long-range arena? It would help quantify limitations of the proposed attention in terms of context utilization.

Author response:
**Q: Long-Range Benchmark Evaluation**

A: Following your recommendation, we have evaluated StreamingLLM on long-range tasks using LongBench (Bai et al., 2023), which includes diverse NLP challenges. Our focus was on assessing how StreamingLLM, particularly using the Llama-2-7B-chat model with a maximum context length of 4k, performs on tasks like single-document QA (NarrativeQA and Qasper), multi-document QA (HotpotQA and 2WikiMQA), and summarization (GovReport and MultiNews).

The results in Appendix Section D of our revised manuscript show that StreamingLLM's performance is on par with the text truncation baseline. However, it's crucial to note that StreamingLLM is not primarily designed to enhance the utilization of long texts. Its strength lies in efficiently utilizing recent information without the need for cache refreshes or extensive memory. Therefore, while we acknowledge this limitation, addressing it falls outside the scope of our current work and is a potential area for future research.

[1] Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., Dong, Y., Tang, J., & Li, J. (2023). LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding. *ArXiv*. /abs/2308.14508