LongRoPE2: Near-Lossless LLM Context Window Scaling

Thank you for your thoughtful review and for recognizing our contributions. We greatly appreciate your acknowledgment of our New RoPE OOD Hypothesis and the role of needle-PPL-guided search in validating this hypothesis through empirical results. We are also glad that you found our extensive experiments in Sections 4.2 and 4.3 valuable in demonstrating the effectiveness of LongRoPE2 and our key design choices, such as the real critical dimension and mixed context window training. Please let us know if there are any specific aspects where we can provide more details. Thank you again for your time and constructive evaluation!

Q1: Clarification on LLaMA3.1-8B long-context evaluation numbers, and the "overclaim" comments

Response: We appreciate your feedback and would like to clarify the following points:

1. 65.1 is the En.MC score of the instruct version, not LLaMA3.1-8B.: As noted in Table 2 of the LLaMA3.1 tech report, the 65.1 score is for the instruct-tuned version (a detail that can be overlooked). Compared to the fair baseline, LLaMA3.1-8B, our model achieves a higher score (46.72 vs. 45.85) on En. MC. Moreover, our model consistently outperforms LLaMA3.1-8B across several long-context benchmarks. Here are additional results:

InfiniteBench and LongBench:

LOFT:

2. Our RoPE scaling method is designed to improve broad long-context capabilities, not to optimize for any specific benchmark like RULER. The use of needle data for search is not designed to fit RULER but to better control long-range token dependency distances in long documents. E.g., We used only the simplest number needle synthesis method. Our extensive experiments proved the superiority over other methods (e.g., NTK, YaRN) across diverse benchmarks.

Q2: What are the criteria for you to choose the evaluation tasks?

Response: Our selection follows two key principles:

1. Effectiveness for evaluating a pre-trained LLM rather than a chat LLM. Since our method extends a pre-trained LLM without post-training, we prioritize tasks aligned with this setup: (i) completion-based tasks, such as few-shot learning and code completion and En.MC in InfiniteBench. (ii) QA tasks with few-shot examples, such as various text-retrieval QA tasks in LOFT.

2. Comprehensive long-context evaluation. To evaluate multiple aspects of long-context performance, we include tasks covering RULER, needle-in-a-haystack retrieval, real-world text QA, high-difficulty KV retrieval, multi-choice QA, few-shot learning, and code completion, as detailed in our evaluation section. We believe this selection fairly reflects the strengths of our method and provides a well-rounded assessment.

Additional results on chat-based sub-tasks. For your reference, we provide additional results on chat-based LongBench tasks. As shown below, We achieve the highest average score, even surpassing LLaMA3.1-8B.

Q3: The clarifications on the main contributions and claims:

Response: We are grateful for your questions and would like to clarify that our main contribution is not the discovery of insufficient training in high-frequency RoPE, but rather the introduction of a new RoPE OOD hypothesis. This hypothesis explains why existing RoPE rescaling methods, such as NTK and YaRN, often result in suboptimal long-context performance. This contribution has been acknowledged by the other two reviewers.

Regarding the mixed context window training, we would like to emphasize that the key difference between our approach and those used in LLaMA3.1 and LongAlign is the use of two distinct RoPE scaling factors: a short factor for short contexts and a long factor for long contexts. This dual-factor approach is essential in significantly recovering short-context performance, which we have shown through extensive experiments.

Here, we perform an additional comparison with LLaMA-3.1’s mixed training. As shown below, we significantly improve short-text performance.

We hope these responses address your concerns and clarify any confusion, and we will incorporate them in the revisions. Thank you again for your valuable feedback and suggestions, and we kindly ask you to consider re-evaluating our work.

Response: Thank you for your valuable feedback and for recognizing the strengths of our work. We appreciate the opportunity to address your concerns.

1. Affordable evolutionary search computational cost: we acknowledge that evolutionary search introduces additional costs. To further clarify its feasibility, we conduct additional experiments to evaluate the search cost when scaling from 128k (current context window length) to 1024k. Using vllm0.7.3 as the inference engine and running on an 8*A100(80GB) server setup, we measured the total search time. As shown in the table below, even when scaling to 1M tokens, the total search time remains manageable at 240 hours (10 days). Moreover, this is a one-time offline process, and the search time can be linearly reduced by increasing the number of GPUs, due to the nature of evolutionary search. Therefore, it is practical for LLM pretraining teams.

2. KV cache recalculation occurs only in specific cases and has minimal overhead: We acknowledge that KV cache recomputation is required when transitioning from the short context window (using the short factor, i.e., original RoPE) to the long context window (using the long factor). However, this recomputation does not occur in every inference. It happens only when the input length is within the short context window, but the total length (input+generated tokens) exceeds it for the first time. After this one-time recomputation, no further recomputation is needed for the rest of the generation.

In most general inference scenarios, this situation is relatively uncommon, as prompts and completions typically either remain within the short context window or start in the long context mode from the beginning. To quantify the cost, we measured KV recomputation time on a 4x80GB A100 GPU (with vllm 0.7.3) for Phi-3-mini and LLaMA3-8B, comparing it against normal decoding time:

The numbers in () indicate the amount of decoded tokens corresponding to the time spent on KV cache recomputation. These results indicate that the additional recomputation cost is equivalent to generating only ~15 (phi3-mini) and ~25(llama3-8b) tokens, which is negligible in the context of long-context generation.

Q2: Does mixed training cause interference between short/long contexts?

Response: Thank you for your insightful question. While it’s true that mixed context window training applies two RoPE scaling factors simultaneously during mid-training, which could introduce interference, our empirical results suggest that this "interference" plays a constructive role and hence does not degrade performance. In fact, it not only recovers short-context performance but also enhances long-context performance. To better illustrate this, we refer to Table 7 from our original paper.

A possible explanation for this surprising improvement is that the so-called “interference” actually plays a constructive role in learning. Specifically, the short-context window helps preserve position modeling for non-interpolated positions (e.g., LLaMA3’s native positions 0, 1, 2, ..., 8191), while the long-context window primarily facilitates the adaptation for newly interpolated positions (e.g., LLaMA3’s new positions like 1/16, 2/16, ..., 17/16). This training strategy effectively constrains the model’s adaptation to interpolated positions while maintaining consistency with the original position modeling - a concept similar to the KL divergence constraint in PPO, which prevents large deviation from the original policy model.

We appreciate this insightful question, which has prompted further reflection and discussions.