YaRN: Efficient Context Window Extension of Large Language Models

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.

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.

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.

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.

Thank you so much for the detailed review and the encouraging rating! We would like to first point out that we have significantly revised the paper and added in many new ablations, charts and tables. We are wrapping up with the new version but it is still not fully ready at the time of this comment. We will include a new chart:

• Training loss curves for the LLaMA 7B model extended to 32k context size using different interpolation techniques. and a few new tables:
• Sliding window perplexity (S = 256) of ten 128k Proof-pile documents over the LLaMA 7b model across PI, NTK-aware, NTK-by-parts, YaRN with various context length, fine-tuned and non-fine-tuned.
• Performance of PI, NTK-aware, NTK-by-parts, YaRN on the Hugging Face Open LLM benchmark suite compared with original LLaMA 7B baselines.
• Comparison of training time in A100-hours for different models and extension methods.

Regarding the weakness:

• This is a good point. We will add a table comparing A100-hours to hopefully address the comparison of compute efficiency between different methods.
• 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.
• We will add a training loss plot including PI, NTK-aware, NTK-by-parts, and YaRN for different training steps.

Regarding the questions:

• 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.
• The PI and NTK methods are easier to implement, whereas YaRN seems to outperform both in all situations with better data efficiency.

Thank you for the thoughtful comments. We would like to first point out that we have significantly revised the paper and added in many new ablations, charts and tables. We are wrapping up with the new version but it is still not fully ready at the time of this comment. We will include a new chart:

• Training loss curves for the LLaMA 7B model extended to 32k context size using different interpolation techniques. and a few new tables:
• Sliding window perplexity (S = 256) of ten 128k Proof-pile documents over the LLaMA 7b model across PI, NTK-aware, NTK-by-parts, YaRN with various context length, fine-tuned and non-fine-tuned.
• Performance of PI, NTK-aware, NTK-by-parts, YaRN on the Hugging Face Open LLM benchmark suite compared with original LLaMA 7B baselines.
• Comparison of training time in A100-hours for different models and extension methods.

Regarding the weakness:

• We will include more ablations and improved the rigorousness in terms of control variables. The major new ones are described above.
• 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.

Regarding the questions:

• 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)
• 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.
• 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.

On the Weaknesses of the Paper

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.

Answers to Questions

• For equation 17, it is a typo, we will fix it.
• Yes, both interpretations of the NTK-by-parts equations for eq 19 and 20 are equivalent, we can change it but the original NTK-by-parts equation was derived to what is included in the paper.
• We will provide that comparison in the next draft.
• The sliding window evaluation is basically the standard perplexity evaluation, but we slide the window across the entire context of the document in order to evaluate its average perplexity across the entire document. For non-sliding window evaluation, the document is usually truncated to the context size from the beginning, but we lose the perplexity contribution of the later tokens in the document. By sliding the context window across the document and iteratively truncating and averaging different parts of the document, we can take in account the perplexity of the entire document, even if the context size is shorter.
• Both s=16 and s=32 models were trained using a dataset of 64k context. The s=32 model extrapolates to 128k context without any issues even if it had only seen 64k context during training.
• We will provide more dynamic scaling results in the next draft.
• 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.
• This result is only backed by empirical evidence, but we are trying to figure out why temperature scaling significantly improves PPL in the case of RoPE. We currently do not have a satisfactory answer to this question.