An Efficient Recipe for Long Context Extension via Middle-Focused Positional Encoding

Thank you very much for taking the time to review our paper. First and foremost, we would like to express our sincere apologies for any confusion that our work brings to you. As you may have noticed, our presentation reached an average of 3.33 / 4 among three other reviewers. Specifically, reviewer srX3 appreciates the clear writing in methods and fWsk confirms the excellent presentation. Nevertheless, we deeply respect your criticism and would like to further clarify all your confusion and misunderstandings point-by-point as follows.

Q1: First, I do not understand the concept of continuity in positional encoding. The authors use a short paragraph to explain continuity (lines 69-74) without any formulation, which is very abstract. The same goes for relativity, though I understood it from another paper. After explaining these two concepts, you still need to explain why they are important (I did not see this in the paper).

A1: In short, the concept of continuity in positional encoding lies in the importance of ensuring the consistency of position indices between fine-tuning and pre-training. Specifically,

1. In lines 69-74, we've discussed the continuity of positional encoding in detail and provided relevant references. PoSE also highlights the importance of continuity in positional encoding during training.
2. In lines 75-82, we elaborate on the relativity of positional encoding and provide a theoretical proof in Appendix B.
3. In lines 38-44, we explain the roles of continuity and relativity. Our ablation experiments (Figure 5(c)) experimentally validate the importance of continuity and relativity.

Q2: Second, I do not understand why splitting the pre-trained context into three segments with different sizes can achieve continuity and relativity.

A2: In lines 105-110 of our paper, we've provided a detailed explanation of this matter. The continuity design aims to allow the middle segment to closely approximate the pre-trained context window. The relativity design aims to enable the model to learn as many relative positions as possible.

Q3: Third, I do not understand why truncated Gaussian middle sampling mitigates the "Lost-in-the-Middle" problem. Line 113 states that it reduces the interval overlap in Eq2. But what is interval overlap and why does it result in the "Lost-in-the-Middle" problem?

A3: In lines 111-115 of our paper, we note that truncated Gaussian sampling can alleviate the "lost in the middle" problem by directing the model's attention towards the middle section of the long context. The overlap of intervals refers to the repeated position indices in Equation (2). The "lost in the middle" issue is not caused by interval overlap but by a bias inherent in the model, as noted in [1]. In Appendix B and L114-115, we further provide theoretical justification that sampling the middle part of a long context with a high importance rate yields a maximization of learned relative position intervals as identified in Equations (2).

This work indeed has gone through extensive experiments, careful writing, and figure polishing. We tried to make all contents illustrative and provided detailed supplementary materials in appendices. We hope our answers together with appendices have provided sufficient explanations and clarified your confusion. We are more than willing to elaborate more details if you have further questions.

Reference:

[1] Lost in the middle: How language models use long contexts. TACL, 2024.

Thank you very much for taking the time to review our paper and acknowledging the novelty and empirical superiority of our approach. Below, we provide detailed replies to your comments and hope we can resolve your major concerns.

W1: Generalizability: While the paper shows impressive results, it is not clear how generalizable these findings are to other LLMs beyond the tested Llama2-7B model. And I'm wondering whether the method can work with PEFT tuning techniques, such as QLoRA.

A1: Thank you for the constructive comments. Our method has strong generalization capabilities and can be applied to other LLMs without modifying any parameters. To verify this, we conducted experiments on Baichuan2-7B. The experimental results are shown in the table below. (Incidentally, as shown in the table in A2 of reviewer srX3's response, the experimental results on LLaMa3-8B also support this conclusion.)

Through two sets of experiments, we demonstrate that CREAM can work seamlessly with PEFT techniques like LoRA [1] and QLoRA [2] without requiring additional modifications. This compatibility arises because CREAM does not alter any data formats or model structures during fine-tuning; it only adjusts position indices. Detailed experimental results are presented in the following two tables.

We will include the training curves of these two methods in the revised version.

W2: Lack of Comparative Analysis: The paper could benefit from a more detailed comparative analysis with other contemporary methods. For example, in Table 2, the paper only compares one model (LongChat-v1.5-7B-32k) with the same 32k context window. I suggest that in such comparision, peer models should at least take the same length of input, otherwise, it cannot prove the effectiveness of the proposed model over other Long LLMs with 32k+ context windows. I suggest the author compare with models like Mistral-7B-Instruct-v0.2, LongLoRA etc.

A2: Thank you for the advice for improving our work. Per your suggestion, we have added experimental results of all three versions of Mistral-7B-Instruct (v0.1, v0.2, v0.3) on Longbench. Since the 7B-Instruct-32K model of LongLoRA is not publicly available, we could not add corresponding results and are more than willing to add the comparison once it's open-sourced. The detailed results of Mistral are presented in the table below.

As shown, it is evident that our method outperforms all three versions of Mistral-7B-Instruct. This further confirms the effectiveness of CREAM.

Q1: What are the computational overheads associated with implementing CREAM, and how does it scale with larger context sizes?

A1: Our method does not incur any additional computational overhead beyond standard fine-tuning. This is because CREAM does not modify the model structure itself; the only change is to the positional indices. Additionally, as the context window length increases, the computational overhead will increase similarly to standard fine-tuning.

However, even with a larger context window, it is still possible to fine-tune with a 4K context length, significantly reducing computational overhead and memory usage. As as shown in the table in A2 of reviewer srX3's response, fine-tuning LLaMA3-8B with a 4K context length (despite its pre-training context length of 8K) remains remarkably effective.

Q2: Are there any specific hyperparameter settings in CREAM that are particularly sensitive, and how were these chosen?

A2: Our method requires only three preset hyperparameters: the mean $\mu$ and variance $\sigma$ of the truncated Gaussian sampling, and the length of the head and tail sections $k$.

For the mean $\mu$, it depends on the expansion factor $\alpha$, specifically $\mu = (1 + \alpha) / 2$.

For the length $k$, we follow the settings of StreamingLLM[3], which uses a value of 4. Since we have expanded by a factor of 8, we set $k = 4 \times 8 = 32$. This also depends on the expansion factor.

For the variance $\sigma$, our setting should ensure higher sampling frequency in the middle and lower (but non-zero) frequency at the ends. The experimental results for different $\sigma$ values in the table in A2 of reviewer fWsk's response can be used as a reference.

Limitations: The discussion in Appendix F regarding the limitations of the CREAM method could be misleading. The author implies that CREAM provides enhanced performance compared to other methods concerning the "Lost in the Middle" issue but acknowledges that the problem is not fully solved due to the nature of decoder-only models. For a genuine limitations section, it is essential to delve deeper into the specific constraints of the CREAM approach and how these limitations might impact its application and effectiveness. A more explicit acknowledgment of the method's shortcomings and a clear explanation of why these issues cannot be overcome with the current model design would significantly improve the section's clarity and usefulness.

A: Thank you for noticing our discussed limitations and providing suggestions. We would like to explicitly elaborate the limitations of CREAM as two-fold. First, our method did not involve any adjustments to the model architecture. As discussed in our research question as in L34-35, this work aims to reach an efficient and effective optimality based on a pre-trained model. Consequently, it encounters the decoder-only limit mentioned in [4]. Second, our methods follow prior works and fine-tune pre-trained models on a small dataset considering the efficiency. Nevertheless, our approach might benefit from further enhancement of pre-trained dataset. We will provide a more detailed discussion of the limitations in our paper.

We hope our answers have resolved your concerns. If you have any other concerns, please feel free to let us know. Thanks again for your review.

Reference:

[1] LoRA: Low-Rank Adaptation of Large Language Models. ICLR, 2022.

[2] Qlora: Efficient finetuning of quantized llms. NeurIPS, 2023.

[3] Efficient Streaming Language Models with Attention Sinks. ICLR, 2024.

[4] Lost in the middle: How language models use long contexts. TACL, 2024.

Q2: Say you had an even more limited amount of compute, to the point where finetuning on the base context length was prohibitively expensive (e.g., maybe you want to finetuned Llama3-8b but can't fit 8192 context at training time). Do you believe CREAM would be useful in this setting? Do you have any results (or speculation!) that suggest when the method may begin to break down with decreasing maximum training length?

A2: Thank you for providing the experimental setup, which further highlights the superiority of our method. Following your instructions, we fine-tuned LLaMa3-8b using a 4K context window size. The experimental results are presented in the table below.

The results presented in the table demonstrate that our method is well-suited to this setting and performs surprisingly well.

Q3: (Minor) What is the context length used in Figure 1? The x axis only reports the position of keys.

A3: This is consistent with the setting in [5] and corresponds to an approximate length of 5K. Results with longer contexts (10K) are in Table 1.

Suggestions/comments: Please explain in the text what the occasional highlighting of numbers represents (e.g. in Table 2). It would be helpful for readability to bold the best number in each column across tables.

A: Thank you for the insightful suggestions. The highlighted numbers in Table 1 and Table 2 emphasize the advantages of our results compared to other methods. In Table 2, we have bolded the average results. Following your suggestion, we will also bold the best results in each column.

We hope our answers have resolved your concerns. If you have any other concerns, please feel free to let us know. Thanks again for your review.

Reference:

[1] Yi: Open foundation models by 01. arXiv preprint arXiv:2403.04652, 2024.

[2] Qwen2 technical report. arXiv preprint arXiv:2407.10671, 2024.

[3] Internlm2 technical report. arXiv preprint arXiv:2403.17297, 2024.

[4] Data Engineering for Scaling Language Models to 128K Context. ICML, 2024.

[5] Lost in the middle: How language models use long contexts. TACL, 2024.

Thank you very much for your valuable suggestions and acknowledgment of our well-motivated method. Below we address your questions point-by-point and hope we can resolve your major concerns.

W1: The method is fairly similar to POSE in concept and performance. While I think the idea of emphasizing middle training is reasonable, well-explored, and clearly effective (e.g. in Figure 1), this has also been explored as a post-training or data selection correction in other work. Given the slight performance degradation from POSE on short-context tasks (i.e. in Table 4), I'm concerned that this might not be worth the tradeoff practically to do at training time.

A1: We would like to highlight the major contributions and differences compared with PoSE, other post-training, and data selection correction as follows.

1. For the concept, PoSE does not simultaneously leverage the benefits of continuity and relativity, nor does it emphasize the importance of intermediate context. These aspects are our core contributions.
2. For the performance, our method significantly outperforms PoSE across various tasks. For example, it achieves an average improvement of 15.9% on LongChat Lines compared to PoSE. In the Lost in the Middle task, our different interpolation methods (Linear, NTK, Yarn) also show an average improvement of about 10%.
3. In post-training, previous methods usually require a large amount of data and computational resources, with the training context length being much greater than the context length during pre-training. For instance, Yi-1.5[1] uses 10 billion tokens and requires upsampling of long sequences. Qwen-2[2] and InternLM2[3] extend their training context length from 4K to 32K in the final stages of pre-training. Therefore, we can conclude that CREAM not only requires less data and shorter training context lengths but also demonstrates impressive performance.
4. For data selection correction, such as [3] using data engineering techniques to expand context. However, their method requires fine-tuning on at least 500M tokens to achieve satisfactory performance in the Needle-in-a-Haystack test (Figure 3). In contrast, our method only requires $seqlen(4096) \times bsz(32) \times steps(1000) = 125M$ tokens. Therefore, their training data requirement is not only four times larger than ours but also requires meticulous data selection. Additionally, our method can be applied to instruction fine-tuning, which they did not mention.

It is important to emphasize that on short-context tasks (i.e. Table 4), our method performs just as well as PoSE. In our manuscript, we overlooked the minor differences between ours and PoSE in Table 4 (<1%) since our focus and superiority lie in addressing the challenges of long contexts. The minor performance turbulence results from the sampling of training data (a de facto strategy as in PoSE and others). To further validate this, we tested model performance with five different random seeds used to sample the training data. The mean$\pm$std results are shown in the table below.

Based on the experimental results in the table, we can be more confident that the performance of CREAM and PoSE on short contexts is actually comparable. Thank you for the careful review. We would update and add the discussion in our revision.

W2: Evaluation settings can differ subtly across papers, and so it would be better for the authors to reproduce baseline results rather than citing them from literature, where possible. This is most critical, I think, in Table 4, where the differences between methods are relatively small and all results are compared to the same baseline method, which is not an infeasibly expensive or unavailable model to run.

A2: Thank you for your suggestion. We've followed your advice and reproduced the results of LLaMa-2-7b-hf, as presented in Table 4. The results are detailed in the table below:

The bolded results in the table are our reproductions, and they are very close to the original results. Therefore, we believe our final conclusions and claims still stand across all tasks. Thank you for the advice. We will add the above results to improve the soundness of our work.

Q1: In Table 5: why does the 256k-extrapolated version of CREAM underperform the 128k-extrapolated version on short contexts?

A1: In Table 5, the 128K results are not ours but are cited from PoSE. Since a lower perplexity is better, the results in Table 5 indicate that our method has the potential to extend to 192K or even 256K. Furthermore, compared to PoSE, our performance is significantly superior.

We sincerely thank you for your constructive comments simple, but clever, effective and effcient and acknowledgement of our approach with very strong performances that encompass a broad range of tasks. Below, we provide detailed replies to your comments and hope we can resolve your major concerns.

W1：No error bars shown on results. In most cases the results are quite strong, but the error bars would be helpful — particularly in some of the closer comparisons with PoSE.

A1: Thank you for the helpful suggestion. In our manuscript, we fixed all random seeds of all experiments to guarantee the reproducibility of our reported results. Per your advice, we have further conducted repeated experiments with five randomly generated seeds on the linear extension experiment. The detailed results with mean$\pm$std are reported in the table below.

By comparing the mean and variance in the table, it can be observed that CREAM and PoSE have comparable capabilities in handling short texts. We will add the results in our revision.

W2: The solution of using a truncated Gaussian approach lacks motivation. The "lost in the middle" problem is clear, however the solution of using a truncated Gaussian to force more focus on the middle seems brittle. Could the parameters of the truncated Gaussian be learned from data? The approach works well based on the results shown in Figure 5a, however.

A2: Thank you for pointing this out. Our motivations for truncated Gaussian are three-fold:

1. Intuitive Explanation. The issue of "lost in the middle" highlights that the performance of LLMs is often strong at the beginning and end, but weak in the middle[1]. As discussed in L111-114, a straightforward idea to address this is to guide LLMs to focus more on the middle part relative to the beginning and end of the context. This approach results in a "reverse U" shape curve, similar to a truncated Gaussian distribution curve.
2. Theoretical support. In Appendix B and L114-115, we provide theoretical justification that sampling the middle part of a long context with a high importance rate yields a maximization of learned relative position intervals as identified in Equations (2).
3. Empirical observation. As you mentioned, it has proven to be quite effective for all demonstrated long context tasks, encompassing retrieval, lost-in-the-middle, LongBench, etc.

Could the parameters of the truncated Gaussian be learned from data?

Great question. In our truncated Gaussian sampling, as shown in Equation (3), the only hyperparameters are the mean $\mu$ and the variance $\sigma$.

1. For $\mu$, it depends on the expansion factor. For example, when expanding from 4K to 32K, the expansion factor is 8, so $\mu = (1+8)/2$. 

2. For $\sigma$, learning from the data is flexible, but since the sampling process is discrete, the gradient cannot be back-propagated, making the learning cost high. Additionally, we conducted experiments with five different $\sigma$ values, and the results in the table show that the current choice ($\sigma=3$) indeed performs the best.

Q: The authors mention that CREAM only needed to be trained for 100 in some cases. Can the authors provide more insight on why the results are very good despite so little training? This seems like a significant achievement, but was not discussed in detail.

A: Thank you for the careful review. For the Base model (Llama2-7b), we conducted 1,000 steps of continual pre-training on the to ensure a fair comparison with prior works as in [5,6]. For the Chat model (Llama2-7b-chat), we performed only 100 steps of Instruction Fine-Tuning (IFT). Indeed, we attempted to extend the IFT to 400 steps, but the performance deteriorated; the detailed results are shown in the table below. This is because we used the instruction dataset ShareGPT as a substitution for the original IFT used for Chat model (the original dataset is not publicly accessible). This observation is in line with previous studies that have pointed out that instruction fine-tuning on different datasets may lead to catastrophic forgetting and impact performance[2,3,4]. To balance the effectiveness of position index learning and the degradation of the LLM's IFT performance, we limited the IFT to 100 steps. Like you have observed, 100 steps of finetuning yield effective long-context performance.

Limitations: The amount of focus on the middle context seems fixed in this approach. While the results for Gaussian truncation appear promising, it's not clear if this direct approach to solving the "lost in the middle" problem is universal. In short, how do we know how much to reweight the importance of the middle context, and does it vary by dataset or task?

A: Sorry about confusion. We would like to clarify that the middle context our method focuses on is NOT fixed. According to Algorithm 1, each sample requires two rounds of sampling to obtain the final position index, so the intermediate context position index varies between samples. Specifically, we first sample a factor $\alpha$ from a truncated Gaussian distribution to determine the interval where the current $P_e$ is located. Then, we uniformly sample the final $P_e$ from this interval, which is the ending position index of the intermediate section. The sample flow was discussed in L122-125. We would further add more clarification to explain our strategy.

As mentioned above, our method does NOT require meticulously designed weight distributions; instead, it automatically assigns weights using truncated Gaussian sampling. Moreover, in all the experiments presented in our paper, all parameters remained consistent, yet the method still performed effectively across various tasks. This demonstrates the generality of our approach, which does not change depending on the dataset or task.

We hope the above response can resolve your questions and concerns. Please let us know if there is any further question! Thanks again for your review.

Reference:

[1] Lost in the middle: How language models use long contexts. TACL, 2024.

[2] CoachLM: Automatic Instruction Revisions Improve the Data Quality in LLM Instruction Tuning. ICDE, 2024.

[3] Lima: Less is more for alignment. NeurIPS, 2024.

[4] From Quantity to Quality: Boosting LLM Performance with Self-Guided Data Selection for Instruction Tuning. NAACL 2024.

[5] Extending Context Window of Large Language Models via Positional Interpolation

[6] PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training, ICLR 2024.