Paper: YaRN: Efficient Context Window Extension of Large Language Models (ICLR-2024)

Reviewer comment (Theory):
Connection to NTK stuff is very difficult to understand and somehow not convincing. I feel the only thing authors want to say is we shouldn't treat all dimensions equally, which IMO is already a good idea. Connecting to NTK stuff actually confused me quite a bit.

Author response:
Thank you very much for your valuable insights. We will improve the explanations about NTK-aware and NTK-by-parts in order to connect them better to YaRN. The original justification for NTK-aware scaling was definitively taken from the lessons learned from the Fourier Features paper and NTK theory. Without that paper, we would not have thought of the initial seed idea on YaRN. Also, we will be providing a much better baseline in the next draft, where we will compare all four methods after and before fine-tuning with the same hyperparameters and data.

## Regarding the questions:

We will improve the explanations on the connections of NTK-aware/by-parts interpolation to NTK theory and Fourier Features.

Reviewer comment (Experiments):
Part of baseline in experiments could be missing.

Author response:
Thank you very much for your valuable insights. We will improve the explanations about NTK-aware and NTK-by-parts in order to connect them better to YaRN. The original justification for NTK-aware scaling was definitively taken from the lessons learned from the Fourier Features paper and NTK theory. Without that paper, we would not have thought of the initial seed idea on YaRN. Also, we will be providing a much better baseline in the next draft, where we will compare all four methods after and before fine-tuning with the same hyperparameters and data.

Reviewer comment (Writing):
In page 4 you mentioned RoPE can't retain high frequency information, but you also say RoPE closely resembles Fourier Features (Tancik et al., 2020) in many aspects, as it is possible to define RoPE as a special 1D case of a Fourier Feature. But somehow I don't know why Fourier Feature can't learn high frequency stuff. Fourier Feature is just trying to map the basis to all sort of features and your description here is fairly confusing. Without a full proper explanation, I think the motivation will be fairly questionable then.

Author response:
The quick explanation on its motivation is that given that RoPE can be expressed similarly to Fourier Features, and we know that linearly interpolating the RoPE up to a certain scale would significantly degrade the model’s performance, even after fine-tuning. The reason behind it is explained much better in the cited Fourier Features paper, but in short if they showed that high-frequency components in Fourier Features are critical for networks to learn the embeddings properly, the same should apply to RoPE. 
This is also somewhat corroborated by a lack of extremely large context sized models that use PI. We are pretty sure that most ML labs out there would have tried and failed to train a PI model of >=64k context size as the degradation from it is unrecoverable. All open models that have bigger than 32k context either use NTK-aware or YaRN interpolation. We initially only tried to train a 64k and 128k model using YaRN and not PI because we knew that PI would not have worked at such extreme scalings.

Reviewer comment (Experiments):
Is there a training-free context window extension stuff that should be included? At least the baseline should be just using encoding scheme as in T5 that we could just change the context window size. For certain task I tried, it doesn't really affect the performance much. So this should be treated as a baseline as well.

Author response:
YaRN with Dynamic Scaling works somewhat well for up to s=4 (that is 4x the original context size) without any fine-tuning. We will add more PPL tests and ablations that include Dynamic Scaling in the next version.

Reviewer comment (Evaluation):
Can you also add the real "clock-time" of training of these 3 methods (PI, NTK, YaRN) so we will have a better understanding what's the cost of using it?

Author response:
Finally, the clock-time of training these three methods are identical if we do not care about the final results. The three interpolation methods do not add any additional compute or memory overhead during training, and they are also free during inference. The only difference is that YaRN converges much faster (and also achieves a lower final loss), and we will provide training convergence comparisons in the next version. As a quick comparison though, the context extension part of Code Llama was fine-tuned using roughly 6400 A100-hours of time, while YaRN 128k model only used 380 A100-hours of compute. We know that this comparison is not optimal as Code Llama was fine-tuned on code, so we will provide a better baseline using PI.

Reviewer comment (Evaluation):
The paper does not provide a comprehensive comparison of YaRN with other length extending methods in terms of computational efficiency and memory usage.

Author response:
This is a good point. We will add a table comparing A100-hours to hopefully address the comparison of compute efficiency between different methods.

Reviewer comment (Novelty):
In the Related Work Section, the authors mention two works, ReRoPE and LM-Infinite, which tackle the same target in terms of extending LLM sequence length, and claim that they can extend to infinite length without severe loss expansion. If that is true, the contribution of YaRN will be greatly reduced (i.e., extend to 16k). Also, I am interested in the results of these two works compared to the proposed YaRN. I know due to the attention modification, they are not supported directly by Flash Attention. However, they represent a series of length extending methods. It would be better to compare them as well even under a small but fair setting.

Author response:
Given the significant compute needs of running long context benchmarks, these alternate methods that do not support Flash Attention are currently low priority for us, as they have no chance to achieve 64k or 128k context sizes with a “reasonable” single-node server machine (eg. 8xA100s). Furthermore, YaRN is unique in a sense that the higher the context size extension is, the better it performs compared to other methods. It would not be appropriate to compare a fine-tuned YaRN model to non fine-tuned models at lower context sizes as the differences will either be minimal compared to other methods that support fine-tuning or simply be unfair to methods that do not support efficient fine-tuning. The results would either be inconclusive (difference is too insignificant) or YaRN would destroy them as we enjoy the benefit of a very quick fine-tune.

Reviewer comment (Evaluation):
Please provide more detailed data points or loss curves between PI and YaRN along with training tokens. Is it possible that with the same tokens, PI will converge faster than YaRN?

Author response:
We will add a training loss plot including PI, NTK-aware, NTK-by-parts, and YaRN for different training steps.

Reviewer comment (Experiments):
In Table 5, for the PASSKEY retrieval experiment, why not keep all the settings the same? Now, 7B-PI is trained on 32k and 7B-YaRN is trained on 64k, it is difficult to compare and draw reasonable conclusions. Also, how do you explain that the PI method can get 100% ACC in terms of key retrieval, while YaRN can only get 96.3% ACC (see row 1 and 3 in Table 5), does this mean that YaRN may lose some ability compared to the PI method? Another question is, is the total number of tokens trained in all experiments in Table 5 the same?

Author response:
This is a very good point and we really appreciate this. We are working on better PASSKEY experiments fixing the extended context length across different methods. By re-running the experiments over models fine-tuned over the same context length, we no longer observe such a drop. We suspect that the original drop from PI 100% to YaRN 96.3% is mostly affected by the difference of fine-tuned context length 32k vs 64k.

Reviewer comment (Evaluation):
In what aspects or tasks does the PI method have advantages over YaRN?

Author response:
The PI and NTK methods are easier to implement, whereas YaRN seems to outperform both in all situations with better data efficiency.

Reviewer comment (Evaluation):
In what aspects or tasks does the NTK method have advantages over YaRN?

Author response:
The PI and NTK methods are easier to implement, whereas YaRN seems to outperform both in all situations with better data efficiency.

Reviewer comment (Experiments):
The ablation study evaluating the impact of each proposed technique in isolation would help validate their necessity in YaRN.

Author response:
We will include more ablations and improved the rigorousness in terms of control variables. The major new ones are described above.

Reviewer comment (Theory):
More analysis and intuition explaining why the temperature parameter works would strengthen that contribution.

Author response:
In terms of why the temperature works, we originally did have a theory touching the entropy of attention weights and supported by limited experiments. But as we expand the ablation, we found that the entropy of attention weights in deeper layers behave in a way contradicting our theory. Unfortunately, we are still trying to figure out the theory side.

Reviewer comment (Experiments):
For the temperature parameter, did you try other hyperparameters or learning it? Why is attention temperature most effective?

Author response:
For other hyperparameters, we are not aware of many other choices inside RoPE. As stated above, why attention temperature is the most effective is still open. The temperature parameters can in theory be learned is a very good point and we have thought about it during our work. Since it amounts to uniform scaling on q and k, we believe the model eventually learns the temperature scaling after sufficient fine-tuning, but hard-coding it has the benefit of fine-tuning on less data (as suggested in Table 1)

Reviewer comment (Evaluation):
You mention the temperature provides a uniform perplexity improvement over the context window. Can you elaborate why you think this is the case?

Author response:
We do not have a mathematical justification. From the experiments, all we can conclude is that the improvement coming from the temperature scaling seems to be independent of the exact batch of data used during forward pass. If we freeze all other parameters and only consider the temperature scaling, it might suggest that the model is undertrained on this specific variable. A more rigorous explanation remains open and is rather intriguing.

Reviewer comment (Experiments):
Can you include results quantifying the impact of each proposed technique independently?

Author response:
We will include a more comprehensive table including PI, NTK-aware, NTK-by-parts and YaRN with different fixed context length, fine-tuned and non-fine-tuned.

Reviewer comment (Presentation):
I found the paper relatively difficult to follow. I think the method could be presented in a much simpler and more direct manner. The "NTK-aware interpolation" could likely be moved to the appendix, as it is not part of the YaRN method. The background section could be significantly shortened (currently 1.5 pages).

Author response:
Thank you for your in-depth review of our paper! Thanks to your suggestions, we will shorten the NTK-aware interpolation section and provide a diagram in order to link the four methods outlined in the paper. We will also move a significant portion of the background section to the appendix.

Reviewer comment (Experiments):
I think the experimental results could be much more thorough, and much more clearly presented.

Author response:
In the case of experimental results, we do agree that it is currently not thorough enough. We would also like to point out that even if YaRN is much more efficient than previous methods, fine-tuning these models and running evaluations at extreme context lengths require significant amounts of computing resources. We are currently finishing the training run for four new LLaMA 7B models using the same training hyperparameters and data, each extended using a different interpolation method (PI, NTK-aware, NTK-by-parts, YaRN). We will be adding a table including all four, and compare their PPL before and after fine-tuning, and also provide PPL for when Dynamic Scaling is enabled/disabled. We will also be comparing these four models on passkey and standard benchmarks in order to establish a better baseline.

Reviewer comment (Experiments):
The experimental results should include more baselines throughout all the experimental results: Standard RoPE embeddings (e.g., Llama-2), ALiBi, T5 relative position embeddings, absolute position embeddings, Position Interpolation (PI), NTK-aware interpolation, YaRN (with and without fine-tuning, with and without temperature scaling, with static vs. dynamic scaling, different scale factors, etc.).

Author response:
For the comparison against other positional embeddings such as T5, ALiBi, absolute, etc. The hard truth is that we simply do not have the required resources in order to perform such extensive tests. In order to compare them fairly to RoPE and avoid confounding factors such as different preexisting base models being pre-trained on different datasets (eg. LLaMA's RoPE vs BLOOM's ALiBi), this experiment would each require us to pretrain and then fine-tune many models from scratch using each of the positional embeddings, and we simply do not have the resources to pretrain a model from scratch. For now, our work will only focus on RoPE, which is the most popular and most used positional embedding of all public models.

Reviewer comment (Presentation):
The tables in the experimental results should be replaced with plots (similar to Figure 4 in appendix, but with clearer and more baselines). The tables are more difficult to interpret, and the plots show results for short sequences as well as long sequences, which is helpful for understanding the performance of the method.

Author response:
We will try to provide better tables and better plots, and also move the current comparisons against Together and Code Llama to the appendix. Those were chosen as there were no other publicly available long context models that we could compare against our 64k and 128k context models. The new PPL comparisons will include an extensive ablation study comparing all different methods and different scaling factors before and after fine-tuning.

Reviewer comment (Experiments):
The paper should add detailed ablations for YaRN, as mentioned above (with and without fine-tuning, with and without temperature scaling, with static vs. dynamic scaling, different scale factors, etc.).

Author response:
The new PPL comparisons will include an extensive ablation study comparing all different methods and different scaling factors before and after fine-tuning.

Reviewer comment (Experiments):
The choice of baselines could be improved. The current choice of Together and Code Llama as the main baselines is questionable, and PI is not included in every table.

Author response:
We will try to provide better tables and better plots, and also move the current comparisons against Together and Code Llama to the appendix. Those were chosen as there were no other publicly available long context models that we could compare against our 64k and 128k context models. The new PPL comparisons will include an extensive ablation study comparing all different methods and different scaling factors before and after fine-tuning.  
...  
- We will compare against a better baseline in the next draft, we did not have any other publicly available models to compare against and we initially could not waste any compute in training an inferior PI model at 64k or 128k.

Reviewer comment (Experiments):
Overall, given the issues with the presentation of the method and (most importantly) the experimental results, I have currently chosen "marginal reject" (I was torn between "marginal accept" and "marginal reject"). While I think the community could benefit from seeing the proposed idea (in particular, the part about only doing interpolation for low-frequency RoPE angles), I think there are currently too many open questions related to the results, and how this method compares with baselines (and with itself, with ablations), to accept the paper in the current form to a peer-reviewed conference. Open to being swayed though, given the potential of this line of work!

Author response:
Thank you for your in-depth review of our paper! Thanks to your suggestions, we will shorten the NTK-aware interpolation section and provide a diagram in order to link the four methods outlined in the paper. We will also move a significant portion of the background section to the appendix.

In the case of experimental results, we do agree that it is currently not thorough enough. We would also like to point out that even if YaRN is much more efficient than previous methods, fine-tuning these models and running evaluations at extreme context lengths require significant amounts of computing resources. We are currently finishing the training run for four new LLaMA 7B models using the same training hyperparameters and data, each extended using a different interpolation method (PI, NTK-aware, NTK-by-parts, YaRN). We will be adding a table including all four, and compare their PPL before and after fine-tuning, and also provide PPL for when Dynamic Scaling is enabled/disabled. We will also be comparing these four models on passkey and standard benchmarks in order to establish a better baseline.

For the comparison against other positional embeddings such as T5, ALiBi, absolute, etc. The hard truth is that we simply do not have the required resources in order to perform such extensive tests. In order to compare them fairly to RoPE and avoid confounding factors such as different preexisting base models being pre-trained on different datasets (eg. LLaMA's RoPE vs BLOOM's ALiBi), this experiment would each require us to pretrain and then fine-tune many models from scratch using each of the positional embeddings, and we simply do not have the resources to pretrain a model from scratch. For now, our work will only focus on RoPE, which is the most popular and most used positional embedding of all public models.

We will try to provide better tables and better plots, and also move the current comparisons against Together and Code Llama to the appendix. Those were chosen as there were no other publicly available long context models that we could compare against our 64k and 128k context models. The new PPL comparisons will include an extensive ablation study comparing all different methods and different scaling factors before and after fine-tuning.

- We will compare against a better baseline in the next draft, we did not have any other publicly available models to compare against and we initially could not waste any compute in training an inferior PI model at 64k or 128k.