InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory

We sincerely appreciate your thorough review of our paper. Here are some clarifications regarding the points raised:

Q1: Novelty

Please refer to Q1, Q2, Q3 in the Global Response. We implement an LLM with token-level memories, and it takes 2 hours to process a sequence with 128K tokens. The cost of time is unaffordable. We argue that building a memory for LLMs is the simplest and most effective method to address long-context challenges. Previous algorithms have struggled to achieve optimal efficiency and effectiveness. Long-text processing is a highly application-oriented field, and therefore, implementing a simple idea with high effectiveness is more conducive to real-world deployment. We believe our empirical methods can greatly benefit LLM deployment for long sequences. The simplicity of our approach, combined with its demonstrated effectiveness, makes it particularly valuable for practical applications. By focusing on a straightforward yet powerful memory mechanism, we have developed a solution that not only improves performance but also maintains feasibility for real-world implementation. This balance between simplicity and effectiveness is crucial in bridging the gap between theoretical advancements and practical deployment in the field of long-context language models.

Q2: Representative Tokens

Yes, different layers tend to choose different representative tokens. We conducted a case study on the selection of representative tokens, and from the results, we can broadly observe that in the top layers, the model generally selects words with stronger semantic information within spans. However, sometimes the model only selects common words, such as "the" or "a", as representations for spans. This indicates that there is still room for improvement in our unit representation approach.

Q3: Memory Unit Size

The segmentation of memory units is a trade-off between effectiveness and efficiency. When choosing a smaller memory unit size, it becomes easier to split a continuous statement into several semantically discontinuous spans, which can further interfere with retrieval effectiveness. Adopting a more dynamic span segmentation strategy is a future direction to improve InfLLM's performance. When we fix the number of retrieved units and vary the size of each unit, the results are as follows. As can be seen, when the unit size increases, the overall amount of relevant content retrieved also increases, leading to better performance.

These findings suggest that larger unit sizes tend to capture longer contexts, resulting in improved performance. However, this also highlights the importance of balancing unit size with computational efficiency and the risk of including irrelevant information. Future work could explore adaptive unit sizing strategies to optimize this trade-off for different types of content and tasks.

Dear Reviewer,

Thank you once more for your insightful feedback on our paper. The discussion period is now underway, and we welcome any additional questions regarding our responses. Reviewer 5j15 has acknowledged that our explanations addressed concerns about comparisons with memory-based methods and the novelty of InfLLM. We are keen to hear your thoughts on these issues as well. If you find that our responses have sufficiently addressed your concerns, we would appreciate a consideration for an improved score. Thank you again for your dedication and time.

We sincerely appreciate your thorough review of our paper. Here are some clarifications regarding the points raised:

Q1: Evaluation on RULER

Please refer to Q5 in the Global Response.

Q2: Comparison to LLMs with 128K Context Window

Please refer to Q4 in the Global Response. We compare InfLLM with the full-attention model based on Yi-9B-200K. From the results, we can observe that InfLLM can still achieve satisfactory performance, which further proves the effectiveness of InfLLM.

Q3: Novelty

Please refer to Q1, Q2, and Q3 in the Global Response. We argue that building a memory for LLMs is the simplest and most effective method to address long-context challenges. Previous algorithms have struggled to achieve optimal efficiency and effectiveness. Long-text processing is a highly application-oriented field, and therefore, implementing a simple idea with high effectiveness is more conducive to real-world deployment. We believe our empirical methods can greatly benefit LLM deployment for long sequences.

The simplicity of our approach, combined with its demonstrated effectiveness, makes it particularly valuable for practical applications. By focusing on a straightforward yet powerful memory mechanism, we have developed a solution that not only improves performance but also maintains feasibility for real-world implementation. This balance between simplicity and effectiveness is crucial in bridging the gap between theoretical advancements and practical deployment in the field of long-context language models.

We sincerely appreciate your thorough review of our paper. Here are some clarifications regarding the points raised:

Q1: Comparison to Existing Memory-based Methods

Please refer to Q1, Q2, and Q3 in the global response. Regarding the articles you mentioned, Unlimiformer employs token-level memory. From our experiments, token-level memory requires significant time for building the retrieval index. Even accelerated with Faiss, it still takes us 2 hours to process a 128K sequence. Focused Attention heavily relies on training unit representations, and the experimental results show that even with training, Focused Attention fails to achieve accurate retrieval on simple tasks like Passkey. Similarly, Infinite Attention uses linear attention mechanism to memorize distant tokens. However, this memory module suffers from catastrophic forgetting due to the limited memory state size, and experimental results demonstrate that this method cannot achieve general distant context memorization. Infinite Attention requires separate fine-tuning for different tasks, making it very impractical. In summary, InfLLM can achieve efficient and general context memory, enhancing LLMs' long-text processing capabilities without requiring training.

Q2: Comparison to Well-trained Long LLM

Please refer to Q4 in the Global Response. InfLLM can still achieve comparable performance with Yi-9B-200K, demonstrating InsLLM's effectiveness. Regarding the performance drop in the Choice task, we believe this is due to the nature of the task, which requires retrieving four separate context sections while encoding a short span (the four options). InfLLM's retrieval during the pre-filling stage is conducted on a block basis, making it challenging to retrieve information for four semantically distinct options. Therefore, further research into how to efficiently improve the granularity of InfLLM's retrieval during the pre-filling stage is a crucial area for future investigation.

Q3: Evaluation on RULER

Please refer to Q5 in the Global Response.

Q4: Ablation Study for Block-level Design

We implement an LLM with token-level memories, and it takes 2 hours to process a sequence with 128K tokens. The time cost is unaffordable, and thus we can not present the results for token-level memories. As for performance drops in decoding only settings, it indicates that prefilling stage is quite important for long sequence processing. With only sparse attention in the prefilling stage, LLMs cannot capture the semantic similarity between distant tokens. Similar phenomenon is also observed in MInference.

Q5: Latency Across Different Context Windows

In our current implementation, during the memory lookup process, we need to compute the dot product between the query and all unit representations. Therefore, in terms of complexity, our current implementation still has O(n^2) complexity. Recent research has shown that different attention heads exhibit significant sparsity (MInference). We can potentially combine InfLLM with these findings to further reduce inference complexity. In this paper, we emphasize InfLLM's context length extrapolation capability, which allows LLMs trained on shorter texts to better handle longer documents. Further reducing InfLLM's computational complexity is a direction for our future work.

We sincerely appreciate your thorough review of our paper. Here are some clarifications regarding the points raised:

Q1: Comparison to Existing Memory-based Methods

Please refer to Q1, Q2, and Q3 in the global response. We are the first to construct a block-level training-free memory module for LLMs. Training-based methods require substantial computational resources and high-quality long-text pre-training and SFT data, which demand significant human effort. Moreover, in many real-world applications, such as ongoing conversations or LLM-driven agents, we need to process streaming input. In these scenarios, training-based methods are constrained by the maximum length used during their training process. Therefore, training-free context extrapolation is a crucial topic. Existing memory-based methods are mainly designed for pre-training language models from scratch. In contrast, InfLLM can be applied in a plug-and-play manner to all transformer-based methods, enhancing their ability to process long texts. Additionally, InfLLM is the first work to design a memory module specifically suited for training-free context extrapolation, implementing a comprehensive and efficient inference solution. Therefore, in this paper, we focus on the performance of memory-based methods in length extrapolation. In our revision, we will further discuss the differences between our work and existing memory-based methods.

As for the comparison with training-based methods, please refer to Table 2 in the original paper and Table 1 in the global response. We compared InfLLM with full-attention models, which can be considered the upper bound of long-context LLMs. The experimental results show that InfLLM can perform comparably to these models on InfiniteBench, demonstrating the effectiveness of InfLLM.

Q2: Efficiency in Offloading Mechanism

As stated in our paper, due to the semantic continuity of long text sequences, the memory units required for adjacent text spans are often highly similar. Therefore, utilizing a GPU cache can effectively reduce the communication between GPU and CPU. Even we employ the random strategy to manage GPU cache, the missing rate is quite low. Therefore, our findings are that the offloading mechanism should be widely-used for long-context inference, which will not incur much additional time costs.

To validate the effectiveness of our cache, we conducted an ablation study: running InfLLM without the GPU cache. The experimental results demonstrate that for encoding a 100K sequence, the addition of a GPU cache reduces our time costs from 21.5s to 18.8s.

Thank you for your comments. I would like to re-emphasize the following points:

1. Training-Free Context Extrapolation: You mentioned that training-free is "a slightly different setting." I cannot agree with this view. As we discussed in the global response, training-free context extrapolation is a challenging and crucial task requiring further efforts and is parallel to training-based context extension. As large models continue to grow, the ability to extrapolate the context window size without training has become a crucial area of focus, receiving significant attention [1-6]. The most relevant works to our paper are those that apply sliding window attention in a training-free manner, such as LM-Infinite (NAACL 2024 Outstanding Paper) and StreamingLLM (ICLR 2024), both of which have garnered considerable attention and follow-up. All contributions in our paper are aimed at advancing the issue of training-free context extrapolation, not just improving the efficiency of LLMs in long text inference. We agree that memory-based models are widely applied in language models, but their application without training remains a significant, largely unexplored challenge.

2. Importance of Training-Free Extrapolation: As stated in our introduction and global response, training-free context extrapolation is a crucial and challenging task. We wish to further emphasize its importance:

   • It holds potential for real-world applications requiring unlimited streaming inputs. As mentioned, any training-based model encounters challenges with out-of-distribution scenarios caused by unseen context lengths. Even a 128K LLM struggles in these streaming contexts.
   • Training on long texts can adversely affect model performance on shorter texts [7]. Many studies have shown that training on long texts impacts short-text performance, a problem avoided in training-free settings.
   • Continuing to pre-train on long texts is not the optimal approach. Creating diverse, high-quality long-text alignment datasets remains extremely difficult. For instance, even after extended pre-training of the Llama3-base model on long texts, we still lack the high-quality alignment data needed to achieve results comparable to Llama3-instruct.

   To summarize, training-free context extrapolation is an important method for most researchers to get a well-performing long-text model. At the same time, all the above points cannot be solved by training-based methods.

3. Formula of Our Approach: Given the importance and challenges of training-free context extrapolation, suggesting a change in the formulas of InfLLM may not be a good choice. Our contributions are specifically within the context of training-free extrapolation, a key research direction parallel to training-based methods.

4. Model Performance: We compared InfLLM with full-attention models in the original paper and global response. Full-attention model can be regarded as the upper bound of long-text models. Our experiments on InfiniteBench have shown that our approach can achieve comparable performance with full-attention methods, even without training. Thus, we argue that our comparative experiments are sufficiently robust, demonstrating that our method approaches the effectiveness of full-attention models in a training-free setup. Therefore, compared to training-based memory method, InfLLM can achieve at least comparable performance with no additional training and can be applied in a plug-and-play manner for all transformer-based models, demonstrating its practicability.

5. RULER Experiment: The RULER experiment was conducted using Llama3 and focused on multi-step reasoning tasks in long texts, which explains the relatively lower accuracy.

6. Offloading Experiment: We have supplemented our paper with additional experiments on offloading. Here we present the relationship between the size of GPU cache and inference time costs of Passkey Retrieval. From the results, we can observe that the offloading mechanism will not incur much additional time costs with limited cache size.

   Size of GPU Cache	All	32	16	0
   Time (s)	28.80	28.89	29.33	33.50

I appreciate your attention to these points and look forward to your reconsideration based on these clarifications. We will add the discussion in the revision.

[1] Efficient streaming language models with attention sinks. ICLR 2024. [2] LM-Infinite: Zero-Shot Extreme Length Generalization for Large Language Models. NAACL 2024 (Outstanding Paper). [3] LLM maybe LongLM: Self-extend llm context window without tuning. ICML 2024. [4] Training-Free Long-Context Scaling of Large Language Models. ICML 2024. [5] Scaling laws of rope-based extrapolation. ICLR 2024. [6] A human-inspired reading agent with gist memory of very long contexts. ICML 2024. [7] LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens. ICML 2024.