PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training

We sincerely appreciate your efforts in reviewing our paper and providing valuable advice. We are pleased that you recognize our paper for addressing a significant problem, substantially reducing memory and computation requirements, and being well-written and illustrated. Regarding the issues you raised, we will reply in detail and hope that our response could address your concerns. Your kind consideration of potentially reevaluating the rating would be greatly valued if our responses address your concerns.

Q1. Section 3.1 and 3.2 are better to be placed in a different background section, or in related work.

Thank you for the suggestion. Section 3.1 introduces the background, and we will separate it into a "Preliminaries" section. Section 3.2 introduces our methodology, and we will present it as a standalone section.

Q2. Limited evaluation for long context. Either provide more non-perplexity based measures or include more baselines in the passkey retrieval test in Figure 2.

Based on your valuable feedback, we have added some new baselines for the passkey retrieval test. We believe this will make our evaluation for long context more comprehensive.

In the original version of our paper, we use the following models for the passkey retrieval test, as in Figure 2:

• Original: The original LLaMA with neither position interpolation nor fine-tuning

• PoSE-16k: the PoSE-extended 16k model

• PoSE-32k: the PoSE-extended 32k model

In the revised version of our paper, we have newly added the baselines below:

• PI-Only-16k: Apply position interpolation for 16k context on LLaMA.

• RandPos-16k: the RandPos-extended 16k model.

• Full-length-16k: the 16k model extended via full-length fine-tuning.

The evaluation results have been updated in Figure 2b. We can see that, without surprise, Full-length-16k is the best within 16k context, and PoSE-16k & 32k are very closely behind. On the contrary, PI-only-16k, RandPos-16k and Original perform poorly. We think that these models could provide a comprehensive comparison for passkey retrieval test, and hope to clarify your concerns.

Thank you very much for your valuable feedback! We sincerely appreciate your recognition of our paper as well-written and offers clear, accessible explanations; our method as effectiveness, time and memory efficient, and adaptable across various position interpolation methods. In response to your inquiries, we provide the following explanations.

Q1. Whether PoSE has genuinely extended the contextual window.

A1. We really appreciate this concern! It is verfied that PoSE has indeed extended the contextual window, as evidenced by the passkey retrieval test (Figure 2). As you mentioned, a lower PPL does not necessarily mean that the model fully understands the entire long input, as in extreme cases, one can always make the model only focus on the last $L_c$ tokens (where$L_c$ is the original context length of the model) using a sliding window approach, and achieve a low perplexity. However, this is not possible for the passkey retrieval test, which requires the model to truly see all input tokens. Our model achieved very high accuracy (>90%) in the target context window, which is strong evidence that our method can genuinely see full sequence during inference. See A3 below for more discussion with StreamingLLM and LM-Infinite.

Q2. Performance of the proposed method when the first chunk does not initiate from position 0.

A2. Thanks for your question! We would like to clarify this. PoSE's design focuses on relative distance rather than absolute distance, so whether the first chunk starts from position 0 or not does not affect the results. Let $L_c$ and $L_t$ denote the original and target context window lengths, respectively, the length of the first chunk be $l_0$, the bias term be $u_0$, the length of the second chunk be $L_c-l_0$, the bias term be $u_0$, we compare the results of two settings: v1) $u_0=0$, i.e., the first chunk starts from position 0; v2) $u_1=L_t-L_c$, i.e., the second chunk ends at position $L_t-1$. Other settings are same as Table 1. The experimental results are attached below. It can be seen that whether the first chunk starts from position 0 or not does not affect the experimental results.

We really appreciate your association with StreamingLLM. However, although they may seem to have some similarities, they are essentially very different. In Discussion1, we will explain this in detail.

Q3. Performance comparison with zero-shot methods like StreamingLLM and LM-Infinite.

A3. Thank you for this question. We would like to clarify the differences between our approach and StreamingLLM & LM-Infinite in this part. StreamingLLM and LM-Infinite are primarily designed for streaming applications, allowing the model to operate continuously on extremely long inputs without extensive memory or dependency on past data[1]. They do extend the contextual window, and can not see long context by design [1]. Our method, however, genuinely extends the context window, providing several benefits:

1) Strong Performance in passkey retrieval. When the input length exceeds $L_c$, zero-shot methods like StreamingLLM and LM-Infinite experience a significant decrease in accuracy on the passkey retrieval task [2]. In contrast, our model achieved very high accuracy (>90%) in the target context window.

2) Lower PPL for long context. By extending context window, models trained with PoSE can better model long-term dependency, gaining lower PPL for long context. For this, we sampled test data exceeding 16k from the GovReport & Proof-Pile dataset and calculated the average PPL of the first 2k, 4k, 8k, and 16k tokens using PoSE-16k, Full-Length-16k, and StreamingLLM, respectively. To maximize the ability of StreamingLLM, we set the cache size to the pre-training window length (i.e., 2048), instead of the default 256, to allow it to retain as much context information as possible. However, this also greatly reduced the running speed of StreamingLLM, so for each dataset, we sampled 10 texts. The experimental results are shown in the following table:

As shown, the performance of PoSE and Full-Length is very close, indicating that PoSE can achieve close effect as full-length fine-tuning. For the first 2k tokens, StreamingLLM achieved slightly better performance than PoSE, which is reasonable as its effect is equivalent to the original LLaMA at this stage. However, as the input becomes longer, especially for 16k context, PoSE largely outperforms StreamingLLM. Hence, we believe that PoSE is superior to zero-shot methods in long text modeling.

Discussion1. Relevance with the initial tokens of KV cache used in StreamingLLM.

A. We are very happy to share with you our observations, although it may not be directly related to our paper. The initial tokens used in StreamingLLM do not refer to tokens with position ids starting from 0, but rather those at the beginning of the sequence, which only attend to a few or no tokens. For reasons that may include serving as an Attention Sink [3], the hidden states of these initial tokens are crucial for the proper functioning of the model, while the hidden states of subsequent tokens do not have this effect. This phenomenon holds for absolute position encoding, relative position encoding, even NoPE. In a caching scenario, if the initial tokens are evicted, the model cannot function properly. In contrast, in a non-caching scenario, the hidden states of initial tokens are always recalculated and thus possess the aforementioned property. PoSE does not evict initial tokens, and there is no need to force $u_0=0$. Based on this dicussion, we can see that StreamingLLM and PoSE are orthogonal work. We hope this answers your question and welcome further discussion.

References

[1] The FAQ part of StreamingLLM. https://github.com/mit-han-lab/streaming-llm#faq

[2] Figure 4 of LM-Infinite (version 1). https://arxiv.org/pdf/2308.16137v1.pdf

[3] Efficient Streaming Language Models with Attention Sinks

Thank you for recognizing our idea as simple and effective, greatly reducing memory and time overheads, and being adaptable to various RoPE-based models and position interpolation strategies. We really appreciate your valuable feedback. In terms of your concerns, we have addressed them carefully below:

Q1: Limited theoretical analysis. Add more discussion of how manipulating position indices allows the model to generalize to longer contexts.

We genuinely appreciate the suggestion to enhance the theoretical analysis in our paper. So far, we have provided some empirical analysis in Appendix B. Based on this, we can potentially explain by analyzing the probability of each relative position being covered by a training example. Let $L_c$ and $L_t$ denote the original and target context window lengths, respectively. For full-length fine-tuning, the probability of a relative position within $L_c$ being covered is 1, and the probability of a relative position in $[L_c, L_t-1]$ being covered is 0. PoSE reduces the coverage probability of positions in $[0,L_c-1]$, while all relative positions in $[L_c,L_t-1]$ receive a certain increase in chance of being covered, making it possible to simulate long inputs with a short context window.

Another potential explanation is the Conceptual Model for Relative Position Attention recently proposed by LM-Infinite[1]. This theory suggests that for a long context, the beginning is responsible for encoding absolute position information, the ending tokens are helpful for the normal operation of the attention layer, and the middle tokens contain relatively less position information. This explains why we can skip middle tokens while receiving only marginal performance drop.

Q2: More details on experimental setups and hyperparameters

Thank you for your suggestion. We use learning rate $2e^{-5}$and a linear scheduler, with 10 warmup steps. We use AdamW optimizer with its default hyperparameters setup. We have added this information to the Training Procedure in Section 4.1.

Q3: Provide some failure cases that this method might leads to some missing details or inferior answers to the standard fine-tuned models.

Compared to Full-length fine-tuning, PoSE achieves very close performance in terms of both perplexity and passkey retrieval accuracy. Inspired by your valuable suggestions, we have identified some instances where full-length outperforms PoSE, but we have also found cases where PoSE outperforms full-length. Therefore, we believe that the differences observed in these individual cases do not imply systematic distinctions between the two approaches. We would like to develop more rigorous methods to analyze the difference between PoSE and Full-length fine-tuning in the future.

Reference

[1] LM-Infinite: Simple On-the-Fly Length Generalization for Large Language Models

Q4. How is full-length fine-tuning conducted.

A4. In our paper, full-length fine-tuning is conducted through interpolation followed by training, which is a routine approach adopted in the PI [1] and YaRN [2] papers. This approach has shown to have the best performance compared to not fine-tuning or not interpolating. For instance, when extending the window size from 2k to 16k, we first use Linear, NTK, or YaRN interpolation to extend the sequence length to 16k, and then fine-tune the model on the 16k sequence length. The same interpolation strategy is also applied to RandPos for fair comparison. Thank you for your suggestion, we have made this clearer in the Baseline Methods part of Section 4.1

Reference

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

[2] YaRN: Efficient Context Window Extension of Large Language Models

[3] (reddit) NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation

[4] (reddit) Dynamically Scaled RoPE further increases performance of long context LLaMA with zero fine-tuning

[5] LooGLE: Can Long-Context Language Models Understand Long Contexts? (Table 5)

Thank you for evaluating our method as neat and novel, reducing the memory requirement for context window extension, and demonstrating great potential to extend to 128k. We appreciate your suggestions very much! Below, we will address your questions carefully.

Q1. Fine-tuning with PoSE degrades performance on standard benchmarks (Table 3).

A1. Thanks for your valuable feedback! Actually, it is common for context-extended models to have a slight performance degradation on short-context tasks compared to the original model. Even GPT-4 has this issue[5]. In Table 3, not only PoSE but also Full-length show a decrease in performance compared to the Original. However, in most cases, the performance degradation is marginal (<5%), which we believe is negligible compared to the benefits of an extended context window.

Q2. PoSE is worse than full-length fine-tuning (Table 1). If the user really cares about the performance, finetuning with the target sequence length still gives the best result.

A2. Thanks for your concern! Our work is a trade-off between performance and efficiency. While we acknowledge that full-length fine-tuning yields better results, its training cost is prohibitively high. In comparison, PoSE sacrifices 1% of performance (Table 1) but offers significant memory and time advantages (Figure 3). Therefore, we believe that PoSE is meaningful for most researchers who cannot afford high training budget, as well as large companies aimed at reducing cost for extending to extreme lengths, such as 256k. In either way, PoSE can promisingly promote the development of the long context field.

Q3. Add comparison between PoSE and applying NTK directly without training.

A3. We sincerely agree with this and have experimented with all three PI strategies. Since we observed the same trend, we only reported the results of Linear interpolation in Table 1. For Linear / NTK / YaRN interpolation, we compared four scenarios: Full-length, PoSE, PI-only, and Original. PI-only means we only apply position interpolation without fine-tuning. Original means the original LLaMA model with neither PI nor fine-tuning. The table below shows the performance of the 16k models from Linear / NTK / YaRN interpolation, under the same settings as Table 1. We can see that for all strategies, the performance follows the same trend: Full-length ≈ PoSE > PI-only >> Original. We also noticed that the NTK method suffers from a significant increase in ppl at 16k, mainly because NTK cannot effectively extend the context window by a scaling factor $\alpha$ [3] [4] . YaRN alleviates this issue, achieving progressively decreasing ppl as the context window grows. We appreciate your suggestion and have included these results in Appendix C.