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

Authors: tong  wu
yanpeng  zhao
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 13:02:29

Abstract: Recently, many methods have been developed to extend the context length of pre-trained large language models (LLMs), but they often require fine-tuning at the target length ($\gg4K$) and struggle to effectively utilize information from the middle part of the context. To address these issues, we propose $\textbf{C}$ontinuity-$\textbf{R}$elativity ind$\textbf{E}$xing with g$\textbf{A}$ussian $\textbf{M}$iddle ($\texttt{CREAM}$), which interpolates positional encodings by manipulating position indices. Apart from being simple, $\texttt{CREAM}$ is training-efficient: it only requires fine-tuning at the pre-trained context window (e.g., Llama 2-4K) and can extend LLMs to a much longer target context length (e.g., 256K). To ensure that the model focuses more on the information in the middle, we introduce a truncated Gaussian to encourage sampling from the middle part of the context during fine-tuning, thus alleviating the ''Lost-in-the-Middle'' problem faced by long-context LLMs. Experimental results show that $\texttt{CREAM}$ successfully extends LLMs to the target length for both Base and Chat versions of $\texttt{Llama2-7B}$ with ``Never Miss A Beat''. Our code is publicly available at https://github.com/bigai-nlco/cream.

Main paper: 
 introduction: 
 Transformer-based Large Language Models (LLMs) are typically pre-trained with a fixed context window size, e.g., 4K tokens in Touvron et al. [2023a] . However, many downstream applications, including in-context learning [Huang et al., 2023 , Li et al., 2023a] and LLM agents [Qian et al., 2023 , Zheng et al., 2023] necessitate the processing of significantly longer contexts, e.g., up to 256K tokens. Recent works have proposed promising approaches that efficiently extend the context window of pre-trained LLMs by interpolating Positional Encodings (PEs) [Chen et al., 2023 , Peng and Quesnelle, 2023 , Peng et al., 2023 , Xiong et al., 2023 , Zhang et al., 2024] with a short period of fine-tuning. Unlike other techniques such as efficient transformer [Tworkowski et al., 2024 , Munkhdalai et al., 2024] and memory augmentation [Tan et al., 2024] , PE-based methods do not necessitate alterations to the model's architecture or the incorporation of supplementary modules. Consequently, PE-based methods offer the advantages of straightforward implementation and rapid adaptation, making them a practical solution for extending the operational range of LLMs in tasks involving larger context windows.Despite the simplicity and effectiveness, existing PE-based methods exhibit two significant limitations. First, prior approaches, such as positional interpolation [Chen et al., 2023] , still require fine-tuning on the target context window size, which imposes a substantial computational overhead [Zhu et al., 2023] . Secondly, though some PE methods have demonstrated potential in handling extremely long sequences, as evidenced by low sliding window perplexity scores, their performance deteriorates notably in "in-the-middle" scenarios [Liu et al., 2024] . Specifically, when the model is required to accurately retrieve and process content located in the middle of an extended context, there is a marked drop in performance on the extended window size (Figure 1 and Figure 3 ). These observations and insights underscore a fundamental question: Can we extend the context window size of pre-trained LLMs efficiently while simultaneously optimizing their effectiveness in processing "in-the-middle" content?To answer the above question, we propose CREAM, namely Continuity-Relativity indExing with gAussian Middle. CREAM is a novel PE-based fine-tuning recipe that shows both efficiency in fine-tuning and effectiveness in enhanced middle content understanding. Our key insights lie in manipulating the positional indices of long target sequences to produce shorter ones within the pre-trained context window size (Figure 2 ).In Section 2.1, we summarize two crucial ingredients of effective positional indices: continuity that produces densely connected positional indices and relativity that reveals the long-range dependencies between fragments. CREAM is a recipe designed with the best of both worlds by introducing two indexing strategies for continuity and relativity, respectively (Section 2.2). Besides, to alleviate the "Lost-in-the-Middle" challenge, we introduce truncated Gaussian distribution for middle segment sampling, enabling the LLM to prioritize the information in the middle positions, even when performing positional interpolation within the pre-trained context window size.In Section 3, we conduct comprehensive experiments to demonstrate the efficiency and effectiveness of CREAM. We continually pre-trained on Llama 2-7B with CREAM for a short period and extend the context window size from 4K to up to 256K. Furthermore, we instruction tuning on Llama 2-7B-Chat with CREAM for 100 steps and obtain promising results. We highlight our empirical advantages as:1. CREAM can not only fine-tune within the pre-training context window size, but also alleviate the issue of the model easily getting lost in the middle. e.g., CREAM-YaRN outperforms PoSE-YaRN [Zhu et al., 2023] by over 20% on average in the "Lost in the Middle" [Liu et al., 2024] task. 2. CREAM can further be enhanced by integrating novel designs on positional interpolation frequencies (such as Linear [Chen et al., 2023] , NTK [Peng and Quesnelle, 2023] , Yarn [Peng et al., 2023] , etc.), and can be extended to context window sizes of up to 256K or beyond. 3. CREAM-Chat model requires only 100 steps of instruction-tuning to achieve nearly perfect performance on the Needle-in-a-Haystack pressure test, and it outperforms existing strong baselines on LongBench [Bai et al., 2023] . 
preliminaries: 
 Problem Formulation. Given an LLM with a pre-trained context window size N , our goal is to unlock the inference capacity of the LLM on the testing data D test with an extended context window size L (where L > N ) by efficiently learning from a small-scale training data D train with a maximum sequence length N . We expect the extended model to perform reasonably well in long-context evaluation. Continuity in Positional Encoding. Transformer-based language models typically encode positional indices sequentially as {0, 1, • • • , N -1}. Traditional length extension methods [Chen et al., 2023 , Peng and Quesnelle, 2023 , Peng et al., 2023] directly fine-tune on the target length L with an updated positional index. This approach preserves the continuity of all absolute positions and learns all position indices within [0, L -1], thereby successfully extending to the target length. Furthermore, PoSE [Zhu et al., 2023] attributed their superior performance over RandPos [Ruoss et al., 2023] to the ensured continuity of segments during fine-tuning.Relativity in Positional Encoding. Relative positional encoding (RPE) [Shaw et al., 2018] has been proposed as an effective positional encoding method, where only the relative positions between two tokens are considered. Similar to prior works [Ruoss et al., 2023 , Zhu et al., 2023 , Wu et al., 2024] , our work focuses on rotary positional encoding (RoPE) [Su et al., 2024] , which is one of the most prominent RPE methods and has been widely applied to LLMs including the recent Llama family [Touvron et al., 2023b ,a, AI@Meta, 2024] . In RoPE, only the relative distances between position pairs (|j -i|; 0 ≤ i < j ≤ L -1) are learned during fine-tuning (Appendix A). Due to this property, we can manipulate the position indices such that all relative positions between [0, L -1] are learnable within the pre-trained window size.2.2 Proposed Recipe: Continuity-Relativity indExing with gAussian Middle (CREAM)In the following section, we start by introducing our design of dividing the context window N to learn relative positional information. Then, we propose two strategies that target continuity and relativity, respectively. Lastly, we propose a novel truncated Gaussian sampling method to enhance the middle part of the long context. The overall framework is depicted in Figure 2 .Context division. We first discuss the motivations behind our design of the context length. First, prior works [Han et al., 2023 , Xiao et al., 2023] observed that a significant amount of attention score is allocated to the beginning tokens of a sequence, which can potentially encode absolute positional information even without explicit positional encoding [Kazemnejad et al., 2024] . Secondly, the starting and ending tokens of long contexts can be treated as two pointers that localize the middle indices with the help of relative encodings. Therefore, we divide the pre-trained context window into three segments. The detailed ablation results are shown in Section 3.6. Definition 2.1. Given the pre-trained context window size N and target extended length L, the position set of {Head, M iddle, T ail} is defined as follows:EQUATIONwhere L h and L t denote the length of the head and tail segments, P s and P e denote the start and end position index of the middle segment.The relative positions among the three segments in each sample are calculated in pairs, i.e., {|j -i|; ∀i, j ∈ {Head, M iddle, T ail}}.The formed relative distance union D r learned by the model is given by:[0, max(L h -1, Pe -Ps, Lt -1)]∪[Ps -L h +1, Pe]∪[L-Lt -Pe, L-1-Ps]∪[L-Lt -L h +1, L-1]. (2)Given that not all samples possess the same values for L h , P s , P e , and L t , as fine-tuning progresses, the union D r in Equation (2) can encompass the entire range [0, L -1], facilitating the model to learn all relative positions within the target length L.Two segmentation strategies. For the sake of continuity, we set the L h and L t to a very small value k, where 0 < k ≪ N . Specifically, we use k = 32 in our experiments. This choice allows the middle segment to closely approximate the pre-trained context window. To maintain relativity, we divide N equally into three parts and fix the L h and L t to N/3, enabling the model to learn as many relative positions as possible. In our fine-tuning process, both types of examples are sampled with equal probability to maintain balance.Truncated Gaussian Middle Sampling To better focus the training process on the middle part of the long context, we introduce a truncated Gaussian function. This approach reduces the interval overlap in Equation ( 2) and directs the model's attention toward the middle section of the long context. In Appendix B, we provide theoretical justifications of our truncated Gaussian design, indicating that the maximization of |D r | holds for middle positions in[N, L/2) ∪ (L/2, L -N ].Formally, given the probability density function (PDF) of a Gaussian distribution:f (x) = 1 σ √ 2π exp - (x -µ) 2 2σ 2 ,where µ is the mean and σ is the standard deviation. The corresponding cumulative distribution function (CDF) is:EQUATIONwhere E(•) is the error function. To calculate the CDF value within the truncated interval, we use a sufficiently large number (e.g. 1000) of equally spaced x values from the given interval [1, L/N ]:EQUATIONBy substituting Equation (4) into Equation (3), the cumulative distribution function (CDF) curve is derived within the truncated interval. For sampling from this truncated Gaussian distribution, the inverse transform method is employed, as demonstrated in Equation ( 5):EQUATIONwhere u ∼ Uniform(0, 1), round(•) represents rounding to the nearest integer. Finally, we can get:EQUATIONwhere L m denotes the length of the middle segments. In summary, the overall sampling flow of our algorithm is presented in Algorithm 1. 
experimental setup: 
 Extended Models We use Llama-2-7B and Llama-2-7B-Chat [Touvron et al., 2023a] as the base models and extend their pre-trained context window size of 4K to a target context length of 32K. The extended models are referred to as CREAM-Base and CREAM-Chat, respectively. Note that, though the target context length is 32K, we do not have to fine-tune CREAM on 32K token long text (see Section 2.2).Require: Pre-trained context window size N , extended context window size L, training sample size S, mean µ, variance σ and hyperparameter k. 1: Generate enough x equally spaced according to Equation (4). 2: Substitute x into Equation (3) to derive the truncated Gaussian CDF F (x). 3: for i = 0 to S -1 do 4:Sample L h ∼ DiscreteUniform({k, N/3}), and let Lt = L h , Lm = N -L h -Lt.Sample u ∼ Uniform(0, 1), and substitute it into Equation (5) to get α.Calculate the start and end position ids Ps, Pe of the middle part according to Equation (6). 7:Get position set Pi = {0, 1, . . . , L h , Ps, . . . , Pe, L -Lt . . . , L -1}, where |P | = N . 8: end for 9: return P = {P0, P1, . . . , PS-2, PS-1}.Benchmarks We conduct long-context LLM evaluation of CREAM-Base on LongChat-Lines [Pal et al., 2023] and Lost-in-the-Middle [Liu et al., 2024] . Ideally, fine-tuning should not disrupt what the base model has learned, so we further evaluate CREAM-Base on the language modeling task and the evaluation benchmark [Beeching et al., 2023] adopted by Llama2. Additionally, we assess the CREAM-Chat model with Needle-in-a-Haystack1 and LongBench [Bai et al., 2023] . Unless otherwise specified, we use linear interpolation to adapt LLMs to a longer context length.Baselines As far as we know, RandPos [Ruoss et al., 2023] and PoSE [Zhu et al., 2023] are similar to our approach in that they manipulate position indices to enable fine-tuning on the pre-trained length for context expansion. Therefore, these two methods serve as the baselines for our primary comparisons. More details about the experimental setup can be found in the Appendix C.We evaluate the long-context understanding capabilities of the CREAM-Base model on two tasks: LongChat-Lines2 [Pal et al., 2023] (Figure 3 ) and "Lost in the Middle" [Liu et al., 2024] (Table 1 ). Figure 3 : Results (%) on LongChat-Lines. Each length consists of 50 samples. All results are finetuned on Llama-2-7B with 4K length data through linear position interpolation. Refer to Appendix E for ablated results using NTK [Peng and Quesnelle, 2023] and Yarn [Peng et al., 2023] .CREAM-Base performs best in retrieving information from long contexts of varying lengths. We extend the context window size up to 32K and compare CREAM with the Llama 2-7B [Touvron et al., 2023a] , RandPos [Ruoss et al., 2023] , and PoSE [Zhu et al., 2023] . As the context window size increases, the performance of all models drops, but CREAM always performs best except for the window size of 3.6K (see Figure 3 ). In terms of the average performance over all context window sizes, CREAM outperforms PoSE by 16%, demonstrating its good long-context understanding ability.CREAM-Base alleviates the Lost-in-the-Middle issue. Lost-in-the-Middle is an observation that LLMs are generally good at retrieving relevant information appearing at the beginning/end of the input context [Liu et al., 2024] . To validate the effectiveness of our middle-focused truncated Gaussian 1 , where the cyan shading indicates middle segments. We find that: regardless of the chosen interpolation method, CREAM always outperforms PoSE by a large margin. e.g., CREAM-Linear surpasses PoSE-Linear by 21.2% when the relevant information is placed at 18.We conduct long-context evaluations of CREAM-Chat on two tasks:• Needle in A Haystack (Figure 10 ) This task is a test that places an answer (i.e., Needle) at any position of a long context window (i.e., Haystack) and requires a model to retrieve the correct answer given a question-answer pair. We follow Wu et al. [2024] and use the GPT (GPT-3.5-Turbo-0125) score as the evaluation metric. • LongBench (Table 2 ) Bai et al. [2023] is a more realistic benchmark because it covers real-world application scenarios like long-context QA and summarization. Moreover, it is specifically designed for Chat models. CREAM-Chat makes best use of the extended context window size. We present results on Long-Bench in Table 2 . CREAM-Chat again surpasses strong baseline models, demonstrating its better use of extended context size. In terms of the average performance over all tasks, it outperforms the second best model, i.e., LongChat-v1.5-7B-32k [Li et al., 2023b] , by 1.6%, though it is only tuned on a very small amount of data and for only 100 steps.To demonstrate that CREAM can be directly combined with PEFT techniques (such as LoRA [Hu et al., 2022] and QLoRA [Dettmers et al., 2023] ), requiring no additional modifications. We conducted experiments on LLaMa-2-7B-Chat using the identical dataset and settings. The experimental results are presented in Table 3 . The results indicate that models fine-tuned using LoRA and QLoRA achieve performance nearly equivalent to those fine-tuned with full parameter. , we perform the classic language modeling evaluation, i.e., perplexity evaluation, on GovReport [Huang et al., 2021] and Proofpile [Zhangir Azerbayev, 2022] . Since a lower perplexity does not necessarily imply better model performance on downstream tasks [Zhang et al., 2024 , Hu et al., 2024 , Arora et al., 2024 , Park et al., 2024] , we further conduct evaluation on the standard natural-language-understanding (NLU) benchmark [Beeching et al., 2023] . This also lets us know whether fine-tuning hurts the NLU ability of the pre-trained base model.Both CREAM and PoSE demonstrate the lowest perplexity. We apply different positional interpolation methods to RandPos [Ruoss et al., 2023] , PoSE [Zhu et al., 2023] , and CREAM and report their perplexities in Table 4 . We find that: CREAM and PoSE have a similar perplexity in different settings and both outperform RandPos. This occurs primarily because the position indices used during RandPos fine-tuning are discontinuous, which creates an inconsistency with the pre-training stage.CREAM has nearly the same NLU abilities as the pre-trained base model. Ideally, fine-tuning should not adversely affect the original capabilities of the pre-trained base model. Our evaluation of CREAM confirms this, i.e., CREAM nearly retains all NLU abilities of the base Llama2-7B (see Table 5 ). Interestingly, CREAM improves over Llama2-7B on ARC-C and HellaSwag. This is because Extending the context length to 256K. We push the limit and extend the context length of Llama-2-7B up to 256K. Following Zhu et al. [2023] , we evaluate the extended model by calculating the average perplexity over 20 samples from PG-19 [Rae et al., 2019] and Book3 [Presser, 2020] . 3Since the PG-19 test set does have enough samples that are longer than 256K, we select a subset of samples from the PG-19 training set.We experiment with target context lengths 64K, 96K, 128K, 192K, and 256K and apply different positional interpolation methods to the extended model (see Table 6 ). The results of PoSE [Zhu et al., 2023] in Table 6 are based on fine-tuning LLaMa 1-7B with 2K data length, and are provided for reference only. Surprisingly, the increase of the target context length brings little to no perplexity increase, demonstrating the stability of CREAM across different target context lengths, even when the target context is extremely long.To validate the effectiveness of our modeling choices, we further conduct an ablation study of three main components of CREAM: truncated Gaussian sampling, fixed start and end segments, and the trade-off between continuity and relativity.Truncated Gaussian sampling versus Uniform sampling. We use truncated Gaussian sampling to encourage CREAM to make better use of the middle part of the context. As a comparison, we replace it with the Uniform sampling (see Figure 5 (a)). We observe that the Uniform sampling always leads to worse retrieval performance, suggesting the effectiveness of the truncated Gaussian sampling.Fixing the head and tail segments is crucial for good retrieval performance. We compare our choice of fixing the head and tail segments with three alternatives: (i) removing both the head and tail segment, (ii) fixing only the head segment, and (iii) fixing only the tail segment (see Figure 5 (b)). We find that: removing the head and tail segments leads to the worst performance; it results in a complete failure (i.e., zero score) for the context size 32K. Keeping either head or tail segments performs slightly better than removing both but underperforms our default choice of fixing both. We suppose that this is because fixing both gives rise to better relativity information, a finding that is consistent with that of Han et al. [2023] .Maintaining a good balance between continuity and relativity is necessary. We encourage continuity by setting the head and tail segment lengths to k = 32 and elicit relativity by letting k = N/3 (see Section 2.2). To balance the two desired properties, we randomly choose k = 32 and k = N/3 with an equal probability during fine-tuning. Here we compare three scenarios: (1) enforce only continuity, (2) enforce only relativity, and (3) balance continuity and relativity (see Figure 5(c) ). We find that balancing continuity and relativity gives rise to the best performance, thus justifying our modeling choice.In our implementation of truncated Gaussian sampling, as illustrated in Equation (3), the only hyperparameters are the mean µ and the variance σ. The mean µ is determined by the expansion factor. The variance σ is adaptable based on data, we conducted experiments with five different values of σ. The results, as presented in Figure 6 , indicate that the current selection (σ = 3) yields optimal performance.  
related works: 
 Efficient Transformers and Extra Memory FoT [Tworkowski et al., 2024] addresses the limitations of local attention in transformers by integrating memory attention layers, which enable large models to learn from a wide context while reducing interference. Infini-attention [Munkhdalai et al., 2024] incorporates compressed memory into the standard attention mechanism and integrates masked local attention and long-term linear attention mechanisms within a single Transformer block.LLoCO [Tan et al., 2024] employs LoRA in conjunction with context compression, retrieval, and parameter-efficient fine-tuning to learn context offline. Although these methods can successfully extend the long context window of LLMs, they either require modifications to the attention mechanism or the addition of extra modules for assistance. In contrast, CREAM does not require these operations and can be directly applied to a pre-trained model.Positional Interpolation Chen et al. [2023] first proposed extending the context window through positional interpolation, which linearly reduces the input position indices to match the original context window size, thereby preventing catastrophic high attention scores from completely disrupting the self-attention mechanism. Subsequently, various methods (such as NTK [Peng and Quesnelle, 2023] , ABF [Xiong et al., 2023] , and EABF [Zhang et al., 2024] ) emerged that modify the base frequency of rotary positional encoding to achieve positional interpolation. YaRN [Peng et al., 2023] introduced a segmented interpolation method, applying different positional interpolations to different dimensions. LongRoPE [Ding et al., 2024] identifies and utilizes two forms of non-uniformity in positional interpolation through search, and introduces a progressive expansion strategy for positiona interpolation. Moreover, CREAM can be combined with any positional interpolation method.Positional Encoding RandPos [Ruoss et al., 2023] first modified position indices so that the model leverages the relativity of positions, enabling it to extend to the target length with fine-tuning over shorter lengths. PoSE [Zhu et al., 2023] then emphasized the importance of continuous segments, dividing the training length into two parts to further enhance the interpolation effect. CREAM utilizes both relativity and continuity, and it also better enables the model to focus on the middle part of the context. 
conclusion: 
 We proposed Continuity-Relativity indExing with gAussian Middle (CREAM), a simple yet effective method to extend the context of large language models. CREAM achieves a trade-off between continuity and relativity, enabling the model to exploit positional relativity (i.e., fine-tuning within the pre-trained length), while preserving text continuity (i.e., remaining as close as possible to the pre-trained state). Furthermore, by employing truncated Gaussian sampling, the model can concentrate more on the middle positions during fine-tuning. Experimental results demonstrate that CREAM outperforms other methods on both Base and Chat models and effectively mitigates the issue of "lost in the middle". 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [NA] Justification: All our experiments are seeded, the results are unique, there are no error conditions, and therefore they are not needed.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: For each experiment, we provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: We conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification:We use all open source datasets and models, and do not involve data or models with a higher risk of abuse, so there is no need to do so.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: The creators and original owners of assets (e.g., code, data, models), used in the paper, are properly credited and the license and terms of use are explicitly mentioned and properly respected Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: We're not introducing new assets, so they're not needed. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: We don't involve any subjects, there are no risks, so there's no need for it. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
: 
 Figure 1 : Results of applying different position interpolation methods to the "Lost-in-the-Middle" task on CREAM and PoSE [Zhu et al., 2023] . We can see that CREAM outperforms PoSE [Zhu et al., 2023] at every position, with a particularly improvement in the middle. 
 
 Appendix: We provide a simple background proof on the relative positional encoding performed by Rotary Position Embedding (RoPE) Su et al. [2024] . Given two embedding vectors x q , x k ∈ R d corresponds to query and key at positions (m, n) ∈ [0, L), where d is embedding dimension, their encoding counterparts can be defined as:EQUATIONwhereEQUATIONis the rotary matrix, Θ = {θ i = 10000 -2(i-1)/d , i = [1, 2, . . . , d/2]} is pre-defined rotation angles.Then the self attention score can be obtained with:q T m k n = ⟨f q (x q , m), f k (x k , n)⟩ = Re   d/2-1 i=0 x q[2i:2i+1] x * k[2i:2i+1] e i(m-n)θi   := g(x m , x n , m -n) (9)where x * represents the conjugate complex of x, g is the derived attention function of RoPE. As seen, RoPE only depends on the relative distances between and encodes the relative position information.Theorem B.1. If N ≪ L, the spanning size |D r | of the relative position union in Equation (2) reaches its maximum iff. one of the following groups of inequalities satisfies:EQUATION) or (L + L h )/2 -1 < P s < P e < L -L t -max(L h -1, P e -P s , L t -1), (11) where max |D r | = max(L h -1, P e -P s , L t -1) + 2N .Proof. Denote four intervals in Equation (2) as S i , i = 1, . . . , 4. According to the inequality of inclusion-exclusion principle for the cardinality of the union of n sets:EQUATION)where the equality holds iff. all sets are pairwise disjoint. That isS i ∩ S j = ∅, ∀i ̸ = j (13)Given intervals as in Equation ( 2), we haveEQUATION)where MAX = max(L h -1, P e -P s , L t -1). The above inequalities can be simplified to Equations ( 10) and ( 11).Lemma B.2. Under mild assumptions that L-L t ≈ L, L+L h ≈ L, the maximization in Theorem B.1 holds for all (P s , P e ) ∈ [N, L/2) ∪ (L/2, L -N ].Proof. Given thatEQUATIONthe inequalities in Equations ( 10) and ( 11 thus, we have:L ≥ S m + L h + L t > L -N/2 (17)Furthermore, as N L → 0, we have:L h + S m + L t → L (18)Model Hyperparameters We fine-tune all models by optimizing the causal language modeling objective. A learning rate of 2 × 10 -5 with a linear scheduler is adopted, incorporating 10 warmup steps. We use the AdamW Loshchilov and Hutter [2018] optimizer with the hyperparameter configurations specified by PyTorch Paszke et al. [2019] . To speed up fine-tuning, we resort to DeepSpeed4 ZeRO stage 1 and Flash Attention-2 Dao [2023] . We perform fine-tuning on two A100-80G GPUs with a total batch size of 32 and run inference on a single A100-80G GPU. For CREAM-Base, we fine-tune it for 1,000 steps on a dataset derived from Pile Gao et al. [2020] ; for CREAM-Chat, we fine-tune it for 100 steps on ShareGPT Zheng et al. [2024] . To ensure fair comparison, we follow the fine-tuning and inference configurations established by Zhu et al. [2023] .For training the Base model, we directly utilize The Pile data provided by Zhu et al. [2023] , and select samples with token lengths exceeding 4K. For training the Chat model, we filter the ShareGPT data from public datasets 5 . Specifically, we used the Vicuna prompt template to sequentially concatenate the ShareGPT data until each data point comprises at least 4K tokens. Then, we select 3.2K data points to train for 100 steps. Particularly, during the instruction tuning process, we mask the USER part and allow the model to calculate the loss only on the ASSISTANT part. We utilize two A100-80G machines with a global batch size of 32, fully utilizing the available memory. Running 1,000 steps for the Base model takes approximately 6 hours, while running 100 steps for the Chat model takes approximately 2 hours.Our proposed method exhibits strong generalization capabilities and can be applied to other large language models (LLMs) without the need for parameter modification. To validate this, we conducted experiments on Baichuan2-7B, with the corresponding results presented in Table 7 .Furthermore, we fine-tuned LLaMa3-8B using a context window size of 4K tokens, with the experimental outcomes shown in Table 8 .The results in Tables 7 and 8 clearly demonstrate the transferability of our method to different models, underscoring its robustness. Of particular note is that despite LLaMa3-8B having a native context length of 8K tokens, fine-tuning on training data with a 4K context window yielded unexpectedly strong performance. The interpolation methods using NTK and Yarn are presented in Figures 7 and 8 . As can be seen, CREAM performs the same as the Linear method for interpolation, still outperforming other methods. The result of NTK at 26K-32K is zero, which is due to the inherent properties of NTK, a finding that is aligns with Zhu et al. [2023] . The results of each subtask in Tables 2 are shown in Tables 10 and 11 .It is noteworthy that, to provide further evidence of the efficacy of our model, we have specifically chosen 12 tasks from the four categories outlined in Zhang et al. [2024] for comparison purposes. As delineated in Table 9 , we are able to attain superior performance on LongBench in comparison to EABF Zhang et al. [2024] , even with shorter training lengths and less data.When extending the context beyond the pre-trained length, there is an inevitable loss of information due to position interpolation, particularly when fine-tuning is restricted to the pre-trained length. However, in comparison to previous methods such as RandPos Ruoss et al. [2023] and PoSE Zhu et al. [2023] , CREAM has effectively mitigated the issue of "Lost-in-the-Middle" by introducing truncated Gaussian sampling. Additionally, as discussed in reference Liu et al. [2024] , decoder-only models are prone to inherently exhibiting a U-shaped performance curve on this task. Therefore, completely solving this problem remains challenging. Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: We provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.Justification: There are no potential positive or negative social impacts involved in our work, so there is no need for that.• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).https://github.com/gkamradt/LLMTest_NeedleInAHaystackPasskey retrieval[Mohtashami and Jaggi, 2023] is another similar task for evaluating long-context LLMs, but it is too simplistic to reflect model performance at different context window sizes, so we use the dataset provided byPal et al. [2023], which closely aligns with the task described inLi et al. [2023b]We use sliding window for calculation, with a window size of 32,768 and a sliding step size of4,096.https://github.com/microsoft/DeepSpeedhttps://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered 


Decision: The paper introduces CREAM, a method for extending the context window of large language models that emphasizes middle context enhancement, showing promising results with efficient training and broad applicability.  CREAM manipulates position indices and employs a truncated Gaussian distribution to focus on the middle part of the context, effectively addressing the "Lost-in-the-Middle" problem.  The proposed method demonstrates superior performance over existing methods like PoSE and RandPos across various tasks. The method shows robustness and has the potential to significantly enhance LLMs' ability to process long contexts.  The authors well addressed the reviewers' concerns in the rebuttal.  Reviewers are overall positive about the paper except for one who suggests to reject but seems not very confident.

Reviews: review: paper_summary: paper_summary: The paper proposes a method for extending the context window of pretrain large language models. The approach, CREAM, relies on modifying position indices to interpolate the positional encodings. Despite the often computationally expensive nature of such work, CREAM can extend to very long context windows while only needing to train at the original pretrained context window. Additionally, in their approach to the context length problem, the authors propose a solution that focuses explicitly on learning at the middle of the context — a span that often under-performs in long context models. The experimental results cover a wide variety of problems, both in terms of tasks, as well as context length challenges. Overall the results look very promising, and show a strong method for addressing a challenging task of extending the context length of LLMs.
strength_weakness: strength_weakness: - 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.
- 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.
questions: questions: 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.
limitations: 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?
paper_summary: paper_summary: The paper proposes a new method, CREAM, that better enables extrapolation to longer contexts via finetuning at the base context length. The method uses positional embedding interpolation with the embeddings divided into three areas of interest and uses a truncated Gaussian  to sample the position used for the middle segment, to increase training of the middle context. They show that this is at least as effective as prior length extrapolation methods that use finetuning at the base context length, and that CREAM-trained models suffer less from the lost in the middle problem.
strength_weakness: strength_weakness: 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. 

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.
questions: questions: Q1. In Table 5: why does the 256k-extrapolated version of CREAM underperform the 128k-extrapolated version on short contexts?

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?

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

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.
limitations: limitations: Yes, no concerns here.
paper_summary: paper_summary: The paper presents CREAM (Continuity-Relativity indExing with gAussian Middle), an innovative approach to extend the context window of Large Language Models (LLMs) without the need for extensive fine-tuning at the target length. The authors address the "Lost in the Middle" problem, which plagues long-context LLMs by causing a performance drop when retrieving information from the middle of the context. CREAM achieves this by manipulating position indices and introducing a truncated Gaussian sampling method to focus on the middle part of the context during fine-tuning.
strength_weakness: strength_weakness: 1. **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. 

2. **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.
questions: questions: 1. What are the computational overheads associated with implementing CREAM, and how does it scale with larger context sizes?
2. Are there any specific hyperparameter settings in CREAM that are particularly sensitive, and how were these chosen?
limitations: 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.
paper_summary: paper_summary: The paper introduces CREAM (Continuity-Relativity indExing with gAussian Middle), an efficient method for extending the context window of large language models (LLMs) to handle longer contexts without the need for extensive fine-tuning. CREAM manipulates position indices for shorter sequences within the pre-trained context window, using two key strategies: continuity, and relativity. Additionally, CREAM employs a truncated Gaussian distribution to enhance the model's focus on the middle part of the context, mitigating the "Lost-in-the-Middle" problem. Comprehensive experiments demonstrate CREAM's efficiency and effectiveness, showing improved performance over baselines like RandPos and PoSE on various tasks.
strength_weakness: strength_weakness: I think the presentation of this paper is a significant problem. Actually, I cannot understand the method itself nor the intuitions behind it. 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). Second, I do not understand why splitting the pre-trained context into three segments with different sizes can achieve continuity and relativity. 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?

I may not be the ideal reader for this paper, but if other reviewers feel the same way, this paper may have been hastily completed.
questions: questions: See Weaknesses
limitations: limitations: None

Comments: title: Rebuttal by Authors
Rebuttal by Authors
Rebuttal by Authors
comment: Thank you for your thorough review of our rebuttal and your encouraging feedback. Your expertise has significantly contributed to improving the quality of our work, and we are dedicated to incorporating your suggestions in our revision process. If you are satisfied with our rebuttal, we would be extremely grateful if you could consider increasing the final score.
Thank you for the response. You have answered each of my questions and provided clarity on some points I misunderstood. I appreciate the additional analyses you’ve run — they provide insight and more confidence in your findings.
We would like to express our sincere gratitude for your thorough review of our paper. Your expertise has greatly contributed to enhancing the quality of our work, and we are committed to incorporating your suggestions during the revision process. Thank you once again for acknowledging our efforts.
Thanks for the detailed response! 

> W1

I appreciate the multiple random seeds for the short-context performance eval (and would love to see this extended to report error bars on the long-context evals as well). My major concern was short-context regression relative to POSE, and looking at the results over multiple random seeds shows that there is not a significant difference in short-context performance. 

> Q1

Totally my bad, somehow missed that this table was perplexity!

> we fine-tuned LLaMa3-8b using a 4K context window

I'm also glad to see how well the method performs in this even more length-constrained scenario.

Given the rebuttal, I have raised my score 5->6. Thanks for the nice paper!
>**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.

| AVG Length          | 2000 | 2700 | 3300 | 4000 | 5200 | 6500 | 7800 | 8800 | 9700 | 11000 | 12000 | 14000 | 17000 | 19000 | 24000 | 28000   | 32000 |
| :------------------ | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :---- | :---- | :---- | :---- | :---- | :---- | :------ | :---- |
| LLaMa3-CREAM-Linear | 0.98 | 0.96 | 0.98 | 1.00  | 0.92 | 0.96 | 0.96 | 0.94 | 0.86 | 0.92  | 0.92  | 0.92  | 0.86  | 0.84  | 0.70   | 0.60 | 0.48  |

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.
>**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.

| Model              | Single-Doc QA | Multi-Doc QA | Summarization | Few-shot Learning | Code Completion | Synthetic Tasks | Macro |
| :----------------- | :------------ | :----------- | :------------ | :---------------- | :-------------- | :-------------- | :---- |
| Llama2-7B-chat-4k* | 24.9          | 22.6         | 24.7          | 60                | 48.1            | 5.9             | 31.0  |
| 100 steps          | 34.8          | 31.1         | 27.2          | 65.1              | 50.4            | 7.0             | 35.9  |
| 400 steps          | 30.9          | 23.5         | 27.1          | 62.4              | 35.6            | 3.4             | 30.5  |

>**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.