A Controlled Study on Long Context Extension and Generalization in LLMs

Figure 1. I am surprised by the LM-Infinite result as the authors of the LM-Infinite reported an pass rate of about 80% on the task of Passkey Retrieval, which is very similar to the needle-in-the-haystack evaluation conducted in Figure 1. However, LM-Infinite fails most cases on NIAH as reported in Figure 1. Could you help explain what are potential aspects that lead to this gap?

We thank the reviewer for the question. We ran the LM-Infinite pass-key retrieval task, and found that indeed within the length LM-Infinite achieves good results, however, when it is going beyond the length, the results degrade a lot. Results are shown in Table 1.

In a concurrent literature, InfLLM[1], authors had a similar observation that with context length getting longer the accuracy of LM-Infinite decreases a lot in their pass key retrieval tasks with their Mistral-7B base model. We hypothesize that LM-Infinite might be good at attending to closer tokens within the window, hence there is a decrease when the length is getting longer. However, LM-Infinite is very good at preserving short-context ability.

Table 1: Lm-infinite passkey retrieval task:

Table 2: Results taken from InfLLM[1]

Reference:

Xiao, Chaojun, et al. "Infllm: Training-free long-context extrapolation for llms with an efficient context memory." The Thirty-eighth Annual Conference on Neural Information Processing Systems. 2024.

Line 194: What is the CLEX method here? Currently there is little introduction of it.

Thank the reviewer for the suggestion. We added some of the explanations of the CLEX method. CLEX achieved near SOTA recents and is considered a relatively recent baseline.

Line 221: "key matrix" refers to "key and query matrices"?

We thank reviewers for pointing this out, and we will revise as suggested.

Line 505 "Context extension hurts in the short term and gains in the long term". I'm not quite sure what this title means here. What are "short term" and "long term" referring to here?

We thank the reviewer for the question. Both short-term and long-term here refers to context length here. We will change to short context and long context to avoid any confusion.

Figure 3, where we analyze the average negative log-likelihood across different context positions, suggests that extending context length can impact performance at shorter contexts. To further validate our hypothesis we took short tasks from Open LLM Leaderboard and re-evaluated the short tasks.

Table 3: Short-Text Task Performance

We had the following observations:

• Performance Degradation: Most long-context extension methods exhibit a slight decrease in performance on short-text tasks compared to the base model. Our discovery, aligning with what you suggested, shows that NTK-Frozen demonstrates better performance on short-text tasks compared to methods like NTK-RoPE.
• Trade-Off Between Long and Short Contexts: The reduction in short-text performance is more pronounced in models using continuous fine-tuning methods. This indicates a potential trade-off between enhancing long-context capabilities and maintaining optimal performance on short-text tasks.
• Alignment with Original Findings: These results align with our observations in Figure 3, where we analyze the average negative log-likelihood across different context positions, suggesting that extending context length can impact performance at shorter contexts.

The presentation and organization of the paper can be improved in various aspects. e.g., using more visualizations instead of tables to summarize the results and highlight main findings; having a table to summary the key characteristics of the 8 compared methods.

We thank the reviewer for the valuable feedback on improving our paper’s presentation and organization. Currently, we have summarized our main findings in Table 1 and included a heatmap and three plots to demonstrate our results.

To enhance clarity, we will add additional visualizations that summarize perplexity trends and final results. Additionally, we will include a new table that outlines the key characteristics of the eight compared methods. These enhancements aim to better highlight our main findings and provide a comprehensive overview of each method’s features.

Thank you for your insightful suggestions.

Is there any inference speed trade-off between these compared methods? e.g., Are attention approximation methods faster? By how much? For some applications the inference speed may be critical. Providing such information will help users make informed decisions.

We thank the reviewer for raising this insightful question. In our experimental study, our primary focus was on evaluating the effectiveness of model performance, which aligns with the approach taken in much of the recent literature.

That said, we also conducted inference speed comparisons under controlled conditions using the same hardware setup. As shown in Table 1, we observed that approximate attention methods are indeed faster, achieving a speedup of approximately 1.5x to 2x compared to LLaMA when the context length is short; however, when the context length gets longer, we didn't see a significant margin. We hypothesize that the discrepancy between the theoretical FLOPs-based comparisons and the observed speedup arises due to differences in hardware characteristics and CUDA implementations of the respective methods.

Table 1: Efficiency analysis of prefill stage time cost, decoding speed, and memory usage

The prefill time cost represents the time required to generate the first token. The decoding speed (seconds / per token) is averaged over 100 token inferences at each sequence length. Memory consumption corresponds to the peak GPU memory usage during inference. All methods, except for LM-Infinite and Landmark, utilize Flash-Attention 2 for enhanced computational efficiency.

Line 484 "Perplexity and downstream tasks". This argument and Figure 4 do not fully convince me. It seems that the linear trend in Figure 4 is highly dependent on LongLora and Landmark, in the sense that the linear trend is likely to disappear without these two compared methods. It seems that when perplexity is below a certain level, e.g., below 6, the perplexity differences between models are small but the downstream performance differences are large. Thus I'm concerned with the claim that perplexity is a "general-purpose performance indicator" as suggested in the abstract, it can only indicate well within a certain region from my understanding. Please consider adding metrics such as rank correlation to strengthen the claim further.

We thank the reviewer for the question. In our revision, we will make this clear about the distance and Here are the updated results for the correlation using a non-parametric method, ken-tau correlation.

Table 1: Kendall correlation of downstream task performance and PPL

Key Findings:

1. Consistency Across Tasks: The results show a strong and statistically significant negative correlation between PPL and downstream performance for most tasks.
   • This supports the claim that lower PPL values are generally associated with better downstream task performance.
2. Task-Specific Observations:
   • The strongest correlations are observed for Needle and RULER, where Kendall's tau indicates a robust alignment between PPL and task performance rankings.
   • For Mshots, the correlation is moderate and statistically weaker, suggesting that PPL's predictive ability may vary slightly depending on the task.
3. Impact of Perplexity Range:
   • Even when perplexity values are close (e.g., below 6), PPL rankings remain a reliable indicator of downstream performance. However, the narrower range may amplify the observed performance differences, highlighting the need for nuanced interpretation.

We will add these findings to our revision.

Dear authors and area chair,

Given the new results and discussion posted by the authors, I have raised my rating to 6. My concerns about the LM-infinite comparison and inference speed trade-off are resolved.

We still have some pending discussion regarding whether there is a strong correlation between ppl and downstream performance. While the conclusion holds when considering all long-context extension methods, the correlation appears weaker when focusing on the strongest or most advanced long-context extension methods. I'm concerned with the broader applicability of the conclusion.

Reviewer j6WF

We thank Reviewer j6WF for the suggestion to further improve the paper.

We have changed our results according to the suggestion of Reviewer j6WF and softened the tone and discovery presented in our paper. We specifically mentioned our discovery is "to some extent" and also mentioned that our discovery is limited as Reviewer j6WF has mentioned in our limitation section.

For example, we wrote the introduction paragraph as,

First, while there have been suggestions that we need new ways to measure performance, our findings show that perplexity does align with how well models perform to some extent on various tasks in our controlled studies. Though some newer attention methods don't show this pattern as clearly, we generally found that when models got better at predicting text, they also got better at most other tasks we tested them on.

We rewrote our discovery paragraph as,

While prior work [1, 2] suggests that perplexity may not reliably predict long-range task performance, our analysis in Figure 2 reveals to some extent perplexity might be reliable. We observe a general correlation between perplexity and model performance across tasks. However, we also observed that approximate attention methods, including LongLora and Landmark on RULER, show minor deviations but maintain a roughly linear relationship. We hypothesize that this apparent discrepancy with previous findings may stem from their less controlled experimental conditions and noisier datasets.

References

[1] Sun, Simeng, et al. "Do long-range language models actually use long-range context?." arXiv preprint arXiv:2109.09115 (2021).

[2] An, Chenxin, et al. "L-eval: Instituting standardized evaluation for long context language models." arXiv preprint arXiv:2307.11088 (2023).

The study is limited to three base models, which may not accurately represent the performance of other, potentially larger models. Expanding the analysis to include a more diverse set of base models could strengthen the conclusions.

Thank you for highlighting the importance of evaluating larger models. In response, we have conducted additional experiments using the LLaMA-13B model and are also running experiments on the LLaMA-70B model. Once those results are available, we will update the manuscript.

Table 1: Performance of LLaMA-13B Extension Methods

Table 2: Detailed RULER Benchmark Results for LLaMA-13B at 32k

We had the following new observations based on the experiments added and our previous analysis,

• Performance Trends: With the larger LLaMA-13B model, we observe that non-extension methods like NTK-Frozen and Self-Extend show improved performance on intrinsic tasks such as Needle-in-a-Haystack compared to their performance at smaller scales.

• Continual Fine-Tuning Methods: Despite the improvements in non-extension methods, continual fine-tuning methods still outperform them within their extension range.

• Perplexity and Downstream Tasks: The correlation between perplexity and downstream task performance remains consistent, reinforcing our original conclusions. We will update the manuscript to include these new findings. The additional results provide deeper insights into how model scaling influences long-context capabilities and validate the robustness of our conclusions across different model sizes.

We will update the manuscript to include these new findings . The additional results provide deeper insights into how model scaling influences long-context capabilities and validate the robustness of our conclusions across different model sizes.

While the paper acknowledges limitations due to fixed hyperparameters, a more in-depth exploration of how different hyperparameter settings might affect the results would be beneficial.

Thank you for highlighting the importance of exploring the impact of hyperparameter settings on our results. We added following empirical studies to address your concerns.

Continual fine-tuned methods

We agree that hyperparameters can significantly influence the performance of different context extension methods, particularly approximate attention methods. We sweep traditional training hyperparameters, such as batch size, and learning rate. Results are shown below,

Table 1: Perplexity Results of LongLora on PG19 and Proof-file

We made the following observations:

• High Sensitivity: Approximate attention methods like LongLoRA are highly sensitive to hyperparameter settings. Small changes in learning rate or training steps led to significant fluctuations in performance.

• Robustness to Hyperparameters: NTK and YaRN methods demonstrated robustness to changes in hyperparameter settings. Their performance remained stable across a wide range of configurations.

• Optimization Challenges: Training times were more predictable for NTK and YaRN and generally shorter because fewer hyperparameter adjustments were needed. For LongLora, achieving optimal performance requires careful tuning, which can be computationally intensive and time-consuming.

Inference Time Optimization

Additionally, we experimented with different hyperparameter settings during inference, such as: scaling factors and some particular hyperparameter tailored to methods. Specifically, for Self-Extend, we follow the empirical rule proposed by the authors for selecting hyperparameters of neighbor tokens and group size, using 32k as the target length to set our hyperparameters of the size of neighbor tokens(window size) and group size. Table 2 below shows the PPL w.r.t different combination of hyperparameters.

Table 2: Perplexity Results of Self-Extend with different group and window size

Despite not requiring fine-tuning, we found that Self-Extend is sensitive to hyperparameters during inference time, when the input context gets longer(32k): their performance varies significantly based on choices like group sizes. Small changes in hyperparameters can lead to considerable fluctuations in model performance during inference. This sensitivity can affect the reliability of these methods in practical applications where consistent performance is necessary.

We will add those into our analysis of the manuscript.

The generalization of findings to longer contexts beyond 32k (e.g., 128k and 1m) is not addressed, which is a significant limitation given the focus on long-context models.

We sincerely thank the reviewer for highlighting the importance of evaluating model performance on contexts longer than 32k tokens.

In our study, we define generalization as the model's ability to perform well across all tasks that extend beyond the training context length. Specifically, we have evaluated our models on tasks where the input lengths exceed 32k tokens, such as NIAH , Perplexity (PPL) and RULER, with sequences up to 64k tokens. We will revise our writing to highlight this in our manuscript. In our original paper, we found that NTK-Dynamic yields the best performance beyond 32k.

To further evaluate the generalization, we evaluated sequences up to 128k tokens on NTK, which in our submission draft works the best in generalization. Results are shown below.

Table 1: Generalization of NTK beyond 32k on RULER

Table 2: Generalization of NTK beyond 32k on NIAH

We found that our conclusion holds for 64k tokens. However, when the context length increased to 128k tokens (4x the fine-tuned length of 32k), we noticed a decrease in performance. This indicates that even for the best generalized methods we discovered in our controlled setting, their generalization becomes weaker when the context length is much larger than the fine-tuned length.

The insights provided in this paper have been discussed in previous studies, and I did not gain any new takeaways from it.

We appreciate the reviewer's candid feedback regarding the novelty of our insights. While we acknowledge that some of our conclusions align with findings from previous studies, we believe that our work offers several distinct contributions that differentiate it from existing research:

1. Controlled Experimental Protocol:

   • Consistent base models, datasets, and metrics ensure fair comparisons.

   • Eliminated confounding variables, enhancing credibility and reproducibility.

2. Clear Mathematical Connections:

   • Explicit mathematical relationships between methods unify different approaches.

   • Provided deeper insights into their performance and interrelations.

3. Comprehensive Evaluation:

   • Extensive tasks include intrinsic metrics (e.g., perplexity) and extrinsic benchmarks (e.g., LongBench, RULER).

   • Tested various model sizes (e.g., LLaMA-13B, adding LLaMA-70B) for stronger generalizability.

   • Standardized benchmarking framework enables meaningful comparisons and best practice identification.

While prior studies may have explored similar themes, our work differentiates itself through the combination of a controlled experimental setup, mathematical summarization contributions, and a comprehensive evaluation framework. We believe that quantifying these properties across multiple approaches and presenting them within a standardized benchmark adds significant value to the field.

We will revise the manuscript to more clearly highlight these unique contributions. By emphasizing these aspects, we aim to better convey the novelty and significance of our work to the reader.

What is the performance of new models, such as Qwen2.5, in this experiment? We thank the reviewer for bringing up Qwen-2.5. This model has demonstrated outstanding performance across various benchmarks and is widely recognized as one of the top lightweight models in many scenarios.

Evaluation of Qwen-7B:

• Context Length Support: Qwen-7B supports an input context length of up to 128k tokens.

• Extension Method: While there is currently no technical report detailing their context extension method, based on the previous Qwen-2 report, we hypothesize that they utilize a context extension technique similar to NTK-RoPE.

• Experimental Results:

  We conducted experiments to evaluate Qwen-7B using our standardized benchmarks. The results are as follows:

Table 1: Performance of Qwen-2.5-7b on Long-Context Tasks

Table 2: Performance and generalization of Qwen-2.5-7b on RULER

We had the following observations:

• Superior Performance: Qwen-7B achieves the best performance on our long-context tasks compared to other models of similar size.

• Generalization Ability: The model demonstrates strong generalization to longer contexts beyond its training range, aligning with the trends observed in our study. We hypothesize that this is partially due to better and more continual fine-tuning data and better training recipe.

In addition to adding the 13B model, we also added the 70B model results, shown below. We observed a similar trend in Phi, Llama-7b, Llama-13B base models. This further shows that our analysis and discovered could be applied to larger models.

Table 1: Performance of LLaMA-70B Extension Methods

Table 2: Detailed RULER Benchmark Results for LLaMA-70B

Table 3 : Performance of LLaMA-7B, LLaMA-13B, LLaMA-70B Extension Methods

We sincerely thank the reviewer for their insightful comments and valuable suggestions. We have addressed each of your concerns in two separate threads.

Based on my experience, increasing the model size in long-context downstream tasks such as LongBench yields some interesting conclusions that are inconsistent with the 7B model. I hope you can add some simple experimental results from the 13B model to discuss this point.

We have conducted additional experiments using the LLaMA-13B model and are also running experiments on the LLaMA-70B model. Once those results are available, we will update the manuscript.

Table 1: Performance of LLaMA-13B Extension Methods

Table 2: Detailed RULER Benchmark Results for LLaMA-13B

We had the following new observations based on the experiments added and our previous analysis,

• Performance Trends: With the larger LLaMA-13B model, we observe that non-extension methods like NTK-Frozen and Self-Extend show improved performance on intrinsic tasks such as Needle-in-a-Haystack compared to their performance at smaller scales.

• Continual Fine-Tuning Methods: Despite the improvements in non-extension methods, continual fine-tuning methods still outperform them within their extension range.

• Perplexity and Downstream Tasks: The correlation between perplexity and downstream task performance remains consistent, reinforcing our original conclusions. We will update the manuscript to include these new findings. The additional results provide deeper insights into how model scaling influences long-context capabilities and validate the robustness of our conclusions across different model sizes.

We will update the manuscript to include these new findings . The additional results provide deeper insights into how model scaling influences long-context capabilities and validate the robustness of our conclusions across different model sizes.

While enhancing the model's long-context capabilities, we also care about the impact of different methods on short-text performance and the extent of knowledge forgetting. I hope you can add some results from the common Open LLM Leaderboard to discuss this point.

We appreciate your suggestion to assess the impact of long-context extensions on short-text performance and knowledge retention. To address this, we evaluated our models on several benchmarks from the Open LLM Leaderboard, focusing on tasks that measure short-text understanding and knowledge.

Table 3: Short-Text Task Performance

We had the following observations:

• Performance Degradation: Most long-context extension methods exhibit a slight decrease in performance on short-text tasks compared to the base model.

• Trade-Off Between Long and Short Contexts: The reduction in short-text performance is more pronounced in models using continuous fine-tuning methods. This indicates a potential trade-off between enhancing long-context capabilities and maintaining optimal performance on short-text tasks.

• Alignment with Original Findings: These results align with our observations in Figure 3, where we analyze the average negative log-likelihood across different context positions, suggesting that extending context length can impact performance at shorter contexts.

We will incorporate these findings into the revised manuscript, discussing the implications for knowledge retention and the balance between long-context capabilities and short-text performance.

The additional experiments provide valuable insights into how model scaling and context extension methods affect long and short-text tasks. We believe these enhancements strengthen our conclusions and contribute meaningfully to the field. Please let us know if there are any further concerns or suggestions

In addition to adding the 13B model, we also added the 70B model results, shown below. We observed a similar trend in Phi, Llama-7b, Llama-13B base models. This further shows that our analysis and discovered could be applied to larger models.

Table 1: Performance of LLaMA-70B Extension Methods

Table 2: Detailed RULER Benchmark Results for LLaMA-70B

Table 3 : Performance of LLaMA-7B, LLaMA-13B, LLaMA-70B Extension Methods

The so-called "approximate attention" method includes many quite different methods and can be misleading. It isn't very clear to include the streaming-LLM methods / longlora / landmark into one class. MinInference-style work shows promising performance now, the conclusion that 'approximate attention is poor' might be misleading.

Thank you for pointing out that the term "approximate attention" in our manuscript encompasses a wide range of methods, which may lead to confusion. We agree that grouping streaming-LLM methods, LongLoRA, and Landmark attention under a single category is not sufficiently precise and could be misleading. In the revised manuscript, we will update the categorization of these methods to more accurately reflect their differences and unique characteristics.

Regarding MinInference-style work[1, 2, 3], we acknowledge that recent developments have shown promising performance in this area. We will include a discussion of these methods in our related work and limitations sections to provide a more comprehensive overview. Additionally, we would like to mention that this line of research is orthogonal to our study and can be applied in conjunction with approaches evaluated in this submission whereas the main focus of this study is effectiveness aligning with recent long-context studies[4, 5].

We will also revisit our conclusion that "approximate attention is poor" to specify that this observation is specific to our experimental setting, particularly when applying approximate attention methods during fine-tuning. This clarification should prevent any potential misunderstandings about the general effectiveness of approximate attention methods.

self-extend seems to be more like NTK-F than 'approximate attention'.

Thank you for highlighting that the Self-Extend method is more akin to NTK-Frozen methods than to "approximate attention" methods. In the revision, we will reclassify Self-Extend under RoPE-based methods and clearly differentiate it from other approaches. This adjustment will improve the clarity of our methodology section and ensure that each method is appropriately categorized.

The contribution of the work might be limited if the arguments were not sufficiently robust.

We appreciate your concern regarding the robustness of our arguments. To ensure the validity of our conclusions, we have conducted experiments across different model sizes and families within a controlled setting. This comprehensive approach enhances our confidence in the robustness of our findings.

We recognize that some existing research presents slightly contradictory results using similar methods. We hypothesize that these discrepancies may stem from the use of unfair or biased base models, or from exhaustive hyperparameter searches that are not standardized across studies. In the revised manuscript, we will discuss these factors in more detail and outline how our methodology addresses them. This added context should strengthen the reliability of our conclusions and clarify the contributions of our work.

References

[1] Leviathan, Yaniv, Matan Kalman, and Yossi Matias. "Fast inference from transformers via speculative decoding." International Conference on Machine Learning. PMLR, 2023.

[2] Xia, Heming, et al. "Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding." arXiv preprint arXiv:2401.07851 (2024).

[3] Yang, Nan, et al. "Inference with reference: Lossless acceleration of large language models." arXiv preprint arXiv:2304.04487 (2023).

[4] Gao, Tianyu, et al. "How to train long-context language models (effectively)." arXiv preprint arXiv:2410.02660 (2024).

[5] Hsieh, Cheng-Ping, et al. "RULER: What's the Real Context Size of Your Long-Context Language Models?." arXiv preprint arXiv:2404.06654 (2024).

More evaluation of short texts is required to check whether these methods perform well in keeping the original performance. (for example, NTK by parts may behave better on this than NTK-rope.)

We thank the reviewer for the suggestion. This question was raised by Reviewer ZKTb. We took short tasks from Open LLM Leaderboard and re-evaluated the short tasks.

To address this, we evaluated our models on several benchmarks from the Open LLM Leaderboard, focusing on tasks that measure short-text understanding and knowledge.

Table 3: Short-Text Task Performance

We had the following observations:

• Performance Degradation: Most long-context extension methods exhibit a slight decrease in performance on short-text tasks compared to the base model. Our discovery, aligning with what you suggested, shows that NTK-Frozen demonstrates better performance on short-text tasks compared to methods like NTK-RoPE.
• Trade-Off Between Long and Short Contexts: The reduction in short-text performance is more pronounced in models using continuous fine-tuning methods. This indicates a potential trade-off between enhancing long-context capabilities and maintaining optimal performance on short-text tasks.
• Alignment with Original Findings: These results align with our observations in Figure 3, where we analyze the average negative log-likelihood across different context positions, suggesting that extending context length can impact performance at shorter contexts. We will incorporate these findings into the revised manuscript, discussing the implications for knowledge retention and the balance between long-context capabilities and short-text performance