SinkLoRA: Enhanced Efficiency and Chat Capabilities for Long-Context Large Language Models

Motivation 1: Elevating Attention Scores for Initial Tokens
Prior studies have demonstrated the Attention Sink phenomenon, where certain tokens, typically the initial tokens in a sequence, receive disproportionately high attention scores during the model’s computation ³⁹. This often occurs because these tokens are visible to all subsequent tokens, leading to significant attention even when they lack semantic importance, particularly in autoregressive language models ³³.

The Sparse Shifted Attention mechanism implemented in LongLoRA ⁶ attempts to address this by shifting the high attention scores from these initial tokens to other tokens that previously received lower attention. This shift reduces the overemphasis on initial tokens. To further improve this, we need to develop a method that directly modifies the attention pattern. By applying this technique, we can effectively redistribute attention scores, thereby reducing the undue emphasis on initial tokens across different token groups.

Motivation 2: Maintaining Initial Tokens During Fine-Tuning
The concept of attention sinks is also utilized in Streaming LLM ³⁹ to improve the model’s handling of long texts. By retaining the Key-Value (KV) pairs of a few initial tokens (attention sinks) along with the most recent tokens, the model ensures stable attention scores and performance even for extended sequences. Inspired by this approach, we aim to carry this mindset from training into inference. Our research aims to modify the fine-tuning process so that initial tokens attend to all other tokens, thereby accumulating more attention scores and enhancing the model’s capacity to handle long sequences.

Motivation 3: Flexible Deployment of Inference Strategy
Efficient deployment of computationally intensive large language models (LLMs) in production environments often relies on Key-Value (KV) caching ¹⁶. KV caching stores the key-value states of previously generated tokens, significantly reducing the need for repetitive computations and thus lowering latency in autoregressive generation. However, LongLoRA ⁶ retains only the original standard self-attention mechanism during inference. To address this limitation, it is necessary to apply an optional KV cache function. This enhancement allows for a more flexible and efficient inference strategy, reducing computational overhead while maintaining model performance.

The primary obstacle in scaling Transformer models to handle longer sequence lengths lies in the self-attention mechanism, which scales quadratically with sequence length in terms of computation time and memory requirements. This quadratic computational burden has prompted significant research efforts focused on developing more efficient sparse Transformer models. Notable examples include Longformer ⁴ and BigBird ⁴¹, which utilize a combination of local, global, and sparse attention mechanisms to manage long contexts, thereby reducing the complexity to O(n). These models achieve a balance between maintaining sufficient context for understanding while managing computational load. For achieving complexity of O(n log n), several approaches have been proposed. Fixed Window Attention ⁷ employs a fixed-size window for attention, which confines the attention computation to a limited context window. Reformer ²¹ introduces locality-sensitive hashing (LSH) to approximate attention by hashing similar tokens into the same buckets, thus reducing the computational complexity. LSG Attention ⁹, adapted from BigBird, combines local, sparse, and global attention to effectively handle long contexts while minimizing computational overhead. Equipping Transformer ⁴⁰ proposes a novel reading strategy termed random access, which enables Transformers to efficiently process long documents without needing to examine every token. This method shows promising results across pretraining, fine-tuning, and inference phases, demonstrating its efficacy in handling extended contexts. Despite these advancements, the ability of these methods to manage long-context conversations, such as those required in chat applications, remains limited. This highlights an ongoing challenge in enhancing the context-handling capabilities of Transformer models for interactive and real-time applications.

Recent advancements in Large Language Models (LLMs) have significantly extended their capabilities, including handling long-context inputs. Math Word Problems (MWPs) have demonstrated notable performance in solving mathematical questions using LLMs ³⁴. Moreover, leveraging LLMs for SQL querying has shown promise in optimizing resource allocation, though it remains less efficient than traditional relational databases ⁴². LongLoRA ⁶, employing Position Interpolation ⁵, has successfully extended the context window of Llama 2 from 4096 to 32768 tokens without requiring substantial GPU or TPU resources. Meta’s Llama 3, featuring up to 70 billion parameters, represents a significant advancement in open-source LLMs, offering enhancements in computational efficiency, trust and safety tools, and collaborations with major platforms ³⁸. Open-source models such as BLOOM ²², OPT ¹⁸, and Falcon ²⁸ continue to challenge proprietary models, although models like Vicuna ²⁹ and Alpaca ¹ still lag behind their closed-source counterparts in certain aspects. Despite these advancements, effectively managing long-context interactions remains a significant challenge, necessitating ongoing research and development to address the complexities in long-context LLM applications.

Compressing the size of KV cache is more difficult than reducing the size of weights because they are more sensitive and dependent on model inputs. A cost-effective method for KV cache compression is token dropping ^{25, 43, 16}, which establishes an importance policy to retain significant KVs and remove insignificant ones. Jiang et al. ²⁰ and Xiao et al. ³⁹ suggest preserving tokens that are local to the current sequence position, as these are crucial for generation. For example, H₂O ⁴³ and FastGen ¹⁶ proposed reducing KV cache size by dropping tokens based on their attention scores. Similarly, SparQ ³² drops tokens according to attention score sparsity and also considers the error in the pruned value cache.

These pruning methods are generally effective for summarization tasks and zero-shot inference. Zhang et al. ⁴³ and Liu et al. ²⁵ recommend identifying a small set of influential tokens, termed heavy-hitters, to better maintain generation quality. Ge et al. ¹⁶ have empirically demonstrated that different attention heads prioritize different tokens and have developed an adaptive importance policy for evicting KVs. However, these approaches can cause significant issues since the context contained in the evicted KVs is entirely discarded.

LongLoRA. LongLoRA, introduced by Chen et al. (2023) ⁶, is an innovative fine-tuning methodology aimed at extending the context window sizes of large language models (LLMs) efficiently. Leveraging Position Interpolation ⁵, LongLoRA builds upon the low-rank adaptations introduced by LoRA ¹⁷ and incorporates Shifted Sparse Attention (S²-Attn) ⁶. Inspired by the Swin Transformer ²⁴, S²-Attn manages extended contexts by partitioning the total context into multiple groups and performing attention computations independently within each group. To ensure continuity, it shifts half of the attention heads by half the group size. This method effectively simulates long-context operations using short attention spans during training, allowing for the handling of significantly larger contexts without a substantial increase in computational overhead. Furthermore, LongLoRA is designed to be compatible with existing LLM optimization techniques and infrastructures, such as Flash-Attention2 ^{11, 10}, thereby enhancing its usability without requiring major system modifications. The approach also emphasizes efficient parameter utilization, necessitating minimal adjustments to the learnable embedding and normalization layers, which represent a small portion of the overall model parameters. This efficiency is crucial for scaling up to larger models while maintaining the practicality of LongLoRA for enhancing performance in tasks that demand deep contextual understanding.

Attention Sink. In autoregressive large language models (LLMs), an intriguing phenomenon known as the “attention sink” ³⁹ is observed, where initial tokens receive a disproportionate amount of attention scores, regardless of their semantic relevance to the task. These initial tokens, although lacking significant semantic importance, tend to accumulate high attention scores. This phenomenon arises due to the nature of the Softmax function used in the attention mechanism, which ensures that attention scores across all tokens sum to one. In situations where few tokens are strongly related to the current query, the model redistributes attention to available tokens, often defaulting to the initial ones. Given the autoregressive nature of LLMs, where each token predicts the next in sequence, initial tokens are consistently visible to nearly all subsequent tokens, inherently training them to become preferred targets for attention, thus acting as “attention sinks.” This insight highlights a fundamental challenge in the attention mechanism of autoregressive models and suggests areas for improving attention distribution strategies in language modeling.

Heavy-Hitter Oracle. The Heavy-Hitter Oracle (H₂O) ⁴³ is a SOTA method that addresses the challenge of reducing the memory footprint of the KV cache in language models. This method is based on the observation that a small subset of tokens, termed “Heavy Hitters” (H2), significantly contribute to the overall attention scores in language models. Analysis shows that these H2 tokens frequently co-occur within the text, making their presence a natural outcome of the text’s structure. Experiments indicate that removing these tokens can significantly impair model performance, underscoring their importance. The H₂O approach combines this understanding with a dynamic cache eviction policy that optimally balances the retention of recent tokens and crucial H2 tokens. This strategy ensures efficient memory use while maintaining the computational dynamics essential for robust model performance.

In this section, we apply the H₂O algorithm to LongLoRA. The LongLoRA codebase⁶ only provides a full attention solution for the inference stage. We incorporate the H₂O eviction system ⁴³ by creating a new class to manage the KV cache values in memory, thus optimizing the inference process. This integration is illustrated in Figure 3.

Perplexity Result.

Analyzing the data from Table 2 and Table 3, we observe the following:

Table 2 presents the perplexity evaluation on the proof-pile test split using three attention mechanisms: Full-Attn, S²-Attn, and SF-Attn for 7B and 8B model sizes. The 7B-Llama2 model, with a training context length of 8192, shows that the SF-Attn mechanism consistently outperforms the other attention mechanisms across all target context lengths (4096, 6144, and 8192). Specifically, the perplexity values for SF-Attn are 8.60, 8.17, and 7.85, respectively, indicating more efficient performance.

Similarly, the 8B-Llama3 model, trained with the same context length, exhibits lower perplexity scores when using the SF-Attn mechanism: 5.11, 4.86, and 4.66 for target context lengths of 4096, 6144, and 8192, respectively. This demonstrates that the SF-Attn mechanism provides a notable improvement in perplexity over both Full-Attn and S²-Attn.

Table 3 summarizes the maximum context length we can fine-tune for various model sizes using a single 2x A100 machine. Both Llama2 and Llama3 models were evaluated using the same training settings. The 7B-Llama2 model shows perplexity results of 8.73, 8.55, 8.30, 8.13, and 8.05 for evaluation context lengths of 2048, 4096, 8192, 12288, and 16384, respectively. The 8B-Llama3 model achieves perplexity values of 5.11, 5.02, 4.89, 4.77, and 4.72 for the same evaluation context lengths.

These results indicate that the SF-Attn mechanism outperforms both Full-Attn and S²-Attn mechanisms in terms of perplexity, especially for larger context lengths. This validates the efficiency and effectiveness of the H₂O algorithm within the LongLoRA framework.

Passkey Retrieval Result.

Analyzing the data from Table 4 and Figure 5, we observe the following:

Table 4 presents the topic retrieval evaluation with LongChat. This task involves retrieving target topics from very long conversations with context lengths of 3k, 6k, 10k, 13k, and 16k. Our model, fine-tuned on Llama3 8B, achieves comparable performance to the state-of-the-art LongChat-13B with a lower fine-tuning cost. Specifically, our model maintains an accuracy of 1.00 for context lengths up to 10k and shows a slight decrease to 0.98 and 0.96 for 13k and 16k context lengths, respectively. This demonstrates our model’s robustness and efficiency in handling long-context retrieval tasks.

Figure 5 compares the passkey retrieval accuracy between Llama2 7B, LongLoRA 7B, and our 7B model fine-tuned on a context length of 32,768. Our model shows no retrieval accuracy degradation up to 33k or 36k, surpassing the context length limits of LongLoRA, which are 30k and 34k. This indicates that our model can handle longer context lengths with higher accuracy compared to existing models.

LongBench Result.

Analyzing the data from Table 6, we observe that GPT-3.5-Turbo consistently outperforms other models across various tasks, achieving the highest average score (44.0). This indicates its superior overall performance. Notably, our model, Ours-7B, secures the second highest average score (38.8), demonstrating competitive performance. Specifically, Ours-7B excels in the Code task with a leading score of 59.2, surpassing GPT-3.5-Turbo (54.1).

Despite these strengths, there are areas for improvement, particularly in the Synthetic task where Ours-7B scores 22.3, significantly lower than GPT-3.5-Turbo’s 37.8. This highlights the need to enhance our model’s capabilities in synthetic tasks to further boost its overall performance. Overall, while GPT-3.5-Turbo remains the top performer, our model shows promising results and potential for optimization.

H₂O Inference Result

In this analysis, we examine the performance and efficiency of various KV cache methods for the LLaMA-3-8B model, as depicted in Figure 5 and Table 6.

Figure 5 presents a comparative analysis between the baseline model utilizing a full cache, the H₂O method, and the “Local” strategy, which employs the most recent KV embeddings. The evaluation was conducted on the OpenBookQA dataset. Notably, the Heavy-Hitter Oracle (H₂O) method sustains high accuracy levels even with a reduced KV cache budget, demonstrating its effectiveness in preserving model performance. Conversely, the Local method exhibits a marked decline in accuracy as the KV cache budget diminishes.

Table 6 details the inference hours required for LLaMA-3-8B under different KV cache budgets. The H₂O method substantially reduces inference time compared to the Full method across all cache budgets. For instance, at a 100% KV cache budget, the H₂O method decreases inference time from 4.7 hours (Full) to 4.4 hours. This reduction is even more significant at lower KV cache budgets, underscoring the efficiency of the H₂O method in environments with limited resources.

In conclusion, the H₂O method provides an optimal balance between maintaining high accuracy and reducing inference time, thereby representing a valuable strategy for optimizing large language models such as LLaMA-3-8B. This method’s ability to offer substantial computational savings while preserving performance highlights its potential for broader applications in resource-constrained settings.

Ablation on SF-Attn training steps

To evaluate the effectiveness of different components in our SF-Attn training methodology, we conducted a series of ablation studies. Table 7 presents the perplexity results for various training methods using the validation split of the PG-19 dataset. The SF-Attn method alone yields a perplexity of 8.64, establishing a baseline for comparison. Introducing global attention to S²-Attn results in a perplexity of 8.91, indicating that global attention alone does not significantly improve performance. Combining S²-Attn with the Segmentation and Reassembly (S&R) Algorithm results in a perplexity of 8.86, showing a moderate improvement over S²-Attn alone. The baseline S²-Attn method alone has a higher perplexity of 9.09 compared to the SF-Attn baseline, demonstrating the need for additional enhancements to achieve better performance. These results indicate that while the SF-Attn method alone is effective, combining it with additional techniques like the S&R Algorithm can yield improvements in perplexity.

Further ablation studies were conducted to analyze the impact of individual components of the Segmentation and Reassembly (S&R) Algorithm. Table 8 provides perplexity results for these experiments. The baseline S²-Attn method has a perplexity of 9.09. When the model only applies the "shift up" step, the perplexity increases to 9.35, indicating that this step alone is not sufficient for improving performance. Applying only the "shift back" step results in a perplexity of 9.60345459, showing that this step alone is also insufficient. The complete Segmentation and Reassembly Algorithm achieves a perplexity of 8.86, demonstrating the effectiveness of combining both shifting steps. These ablation studies highlight the importance of the Segmentation and Reassembly Algorithm in enhancing the performance of the SF-Attn training method. By integrating both the "shift up" and "shift back" steps, the S&R Algorithm effectively reduces perplexity and improves model performance.

Ablation on SF-Attn Varation

specifically focusing on the Sparse Fixed Attention (SF-Attn) and its variations. Figure 7 illustrates the different attention patterns: Sparse Fixed Attention, Stride Attention, and Random Attention. These variations represent different strategies for distributing attention across the token sequence. Sparse Fixed Attention uses a fixed sparse pattern, Stride Attention distributes attention in a stride pattern, and Random Attention allocates attention randomly.

Table 9 presents the results of fine-tuning a Llama2 7B model on sequences with varying target context lengths (4096, 6144, and 8192) and evaluating on the PG19 validation set. The performance is measured across different settings of the attention mechanism. Key observations include that Full Attention (Full-Attn) shows consistent performance but is computationally expensive. S²-Attention (S²-Attn) performs slightly better than Full Attention, particularly at longer context lengths. Sparse Fixed Attention maintains competitive performance and is more efficient. Stride and Random Attention also perform well but slightly lower than Sparse Fixed Attention. The SF-Attention method, combining elements of these strategies, achieves the best balance with the lowest perplexity scores, indicating it is the most efficient and effective attention mechanism for handling long sequences.

The shift operation of Sparse Shift Attention ⁶ is inspired by the Swin Transformer ²⁴. It aims to facilitate more information exchange between different groups of tokens. However, this approach also leads to an information leaking problem.

The information leaking issue likely arises because the shift operation allows tokens from different segments to share information too freely. While the intention is to enhance the model’s ability to integrate information from various parts of the sequence, it inadvertently causes tokens to access information that they should not be aware of at that particular stage of processing. This premature exposure to information can disrupt the model’s learning process, leading to overfitting or incorrect associations.

Moreover, the shift operation might blur the distinct boundaries between token groups, causing the model to lose track of the specific context within each group. This can result in a degradation of performance, as the model might struggle to maintain a coherent understanding of the local context, which is essential for accurately interpreting and generating long sequences.

In summary, while the shift operation aims to improve information exchange, it inadvertently causes information leakage by allowing tokens to prematurely access and integrate information from different segments. This undermines the model’s ability to maintain distinct contextual boundaries and can negatively impact its performance on tasks requiring precise contextual understanding.

The method described in the following section details the successful implementation of globally operative sink attention tokens, enhancing the model’s ability to handle long sequences efficiently.

The successful results observed from implementing sink attention tokens can be attributed to several key factors. First, by designating specific tokens as sink attention tokens that attend to all other tokens in the sequence, the model can effectively capture and integrate information from across the entire sequence. This global attention mechanism ensures that critical information is not lost, even in very long sequences.

Furthermore, by ensuring that all tokens in the sequence also attend to the sink attention tokens, the model can maintain a coherent understanding of the sequence context. This bidirectional attention flow allows the model to reinforce important information at multiple stages, enhancing the overall comprehension and retention of context.

The efficient complexity of O(n log n) achieved through this method, due to the relatively small number of sink attention tokens, ensures that this enhanced capability does not come at the cost of significantly increased computational overhead. This balance between maintaining comprehensive attention and computational efficiency is likely a key factor in the improved performance observed in models utilizing this technique.

In summary, the implementation of globally operative sink attention tokens allows the model to maintain a robust and coherent understanding of long sequences, ensuring critical information is attended to throughout the processing stages, thus leading to improved performance on tasks requiring extensive context comprehension.

As the results shown in the H₂O inference test indicate, the H₂O method can maintain accuracy even when the KV cache budget is reduced by half. However, the time savings achieved with this method are significant, reducing inference time by a factor of 1.5. This capability allows for flexible deployment of inference methods based on the available computing resources. The ability to adjust the KV cache budget without compromising accuracy ensures that large language models can be efficiently and effectively utilized in resource-constrained environments, optimizing both performance and computational efficiency.

As shown in the example of chat ability comparison in Section C, we observe the responses from two versions of the LongLoRA model to a story-related question.

The user input describes a narrative involving an old fisherman named Tom, who ventures farther out to sea to a place known as the “Blue Deep” in hopes of finding fish after several unsuccessful weeks. The question asks why Tom decided to venture farther out to sea than he ever had before.

The LongLoRA-7B model’s response focuses on Tom’s desperation and worry about not catching any fish, emphasizing his repeated return with an empty net and his determination to support his family. This response accurately captures the key elements of Tom’s motivation as described in the narrative.

On the other hand, the SinkLoRA-7B model provides a more detailed explanation, highlighting the prolonged disappearance of fish and Tom’s resultant decision to take the risk of venturing to the “Blue Deep.” It also mentions Tom’s hope to find fish there, providing a more comprehensive understanding of his motivations.

The reason for the differing levels of detail and context in the responses can be attributed to the training and fine-tuning differences between the two models. The SinkLoRA-7B model might have been trained on a more diverse dataset or undergone additional fine-tuning to better understand and generate contextually rich responses. This additional training could enable it to provide more nuanced and detailed answers, capturing the subtleties of the narrative more effectively.

In conclusion, both models successfully identify Tom’s primary motivation, but the SinkLoRA-7B model offers a more thorough and contextually rich explanation. This comparison underscores the effectiveness of the chat capabilities in understanding and accurately responding to narrative-based questions. The observed differences highlight the importance of extensive training and fine-tuning in enhancing model performance and response quality.