YaRN: Efficient Context Window Extension of Large Language Models
Authors: Bowen Peng
Jeffrey Quesnelle
Honglu Fan
Enrico Shippole
Venue: ICLR 2024 poster (2024)
Keywords: transformers
nlp
fine-tuning
context window
attention
rotary embedding
Primary area: generative models
TL;DR: We build on Positional Interpolation and achieve SOTA performance in extending the context window of decoder-only transformer using Rotary Embedding, both with and without finetuning.

Abstract: Rotary Position Embeddings (RoPE) have been shown to effectively encode positional information in transformer-based language models. However, these models fail to generalize past the sequence length they were trained on. We present YaRN (Yet another RoPE extensioN method), a compute-efficient method to extend the context window of such models, requiring 10x less tokens and 2.5x less training steps than previous methods. Using YaRN, we show that LLaMA models can effectively utilize and extrapolate to context lengths much longer than their original pre-training would allow, while also surpassing previous the state-of-the-art at context window extension. In addition, we demonstrate that YaRN exhibits the capability to extrapolate beyond the limited context of a fine-tuning dataset. The models fine-tuned using YaRN has been made available and reproduced online up to 128k context length.

Published as a conference paper at ICLR 2024
YARN: EFFICIENT CONTEXT WINDOW EXTENSION OF LARGE LANGUAGE MODELS
Bowen Peng'
Jeffrey Quesnelle'
Honglu Fan
Enrico Shippole
Nous Research
EleutherAl
3 University of Geneva
ABSTRACT
Rotary Position Embeddings (RoPE) have been shown to effectively encode posi- tional information in transformer-basee language models. However, these models fail to generalize past the sequence length they were trained on. We present YaRN (Yet another RoPE extensioN method). a compute-efficiet method to extend the context window of such models, requiring 10x less tokens and 2.5x less training steps than previous methods. Using YaRN, we show that L.LMM models can effectively utilize and extrapolate to context lengths much longer than their original pre-training would allow, while also surpassing previous the state-of-the-art at context window extension. In addition, we demonstrate that YaRN exhibits the capability to extrapolate beyond the limited context of a fine-tume tinettuning dataset. The models fine- tuned using YaRN has been made available and reproduced online up to 128k context length.
INTRODUCTION
Transformer- Large Language Models Vaswani al. 2017) (LLMs) have become the near- ubiquitous choice for many natural language processing (NLP) tasks where long-range abilities such as in-context learning (ICL) has been crucial. In performing the NLP tasks, the maximal length of the sequences (the context window) determined by its training processes has been one of the major limits of pretrained L.M.. Being able to dynamically extend the context window via a small amount of fine-tuning (or without fine-tuning) has become more and more desirable. To this end, the position encodings of transformers are the center of the discussions.
The original Transformer architecture used an absolute sinusoidal position encoding, which was later improved to learnable absolute position encoding (Gehring et al. 2017). Since then, relative posi- tional encoding schemes (Shaw al. 2018) have further increased the performance of Transformers. Currently, the most popular relative positional encodings are T5 Relative Bias (Roberts et al. 2019).
RoPE (Su et al. 2022), XPos (Sun al. 2022 and ALiBi (Press et al. 2022).
One reoccurring limitation with positional encodings is the inability to generalize past the context window seen during training. While some methods such ALiBi are able to do limited generalization, none are able to generalize to sequences significantly longer than their pre-trained length (Kazemnejad
al. 2023).
Some works have been done to overcome such limitation. (Chen et al. 2023) and concurrently (kaio- kendev 2023) proposed to extend the context length by slightly modifying RoPE via Position Interpolation (PI) and fine-tuning on a small amount of data. As an alternative. (bloc97 2023a) proposed the "NTK-aware" interpolation by taking the loss of high frequency into account. Since then, two improvements of the "NTK-aware" interpolation have been proposed, with different emphasis:
the "Dynamic NTK" interpolation method (emozilla 2023 for pre-trained models without fine-tuning.
the "NTK-by-parss" interpolation method (bloc97 2023b) which performs the best when fine-tuned on small amount of longer-context data.
The "NTK-aware" interpolation and the "Dynamic NTK" interpolation have already seen their presence in the open-source models such as Code Llama (Rozière et al. (203) (using "NTK-aware" interpolation) and Qwen 7B (qwe) (using "Dynamic NTK").

---

Published as a conference paper at ICLR 2024
In this paper, in addition to making complete account of the previous unpublished works the "NTK-aware", the 'Dynamic NTK" and the -NTK-b-parte interpolations, we present YaRN (Yet another RoPE extensioN method), improved method to efficiently extend the context window of models trained with Rotary Position Embeddings (RoPE) including the LLaMA (Touvron et al.,
2023a), the GPT-NeoX (Black et al., 2022), and the PaL.M (Chowdhery al. 2022) families of models.
The relationship between different methods and how they evolve into YaRN can be summarized into the following diagram:
apply ramp over
Position Interpolation
temperature
scallng of
NTK-by-parts Interpolation
RoPE
YaRN
NTK-aware Interpolation
keep the insight from NTK: scale differently!
Figure 1: An outline of the relationship between different interpolation methods.
YaRN reaches tate-of-the-art performances in context window extensions after fine-tuning on less than ~0..1% of the original pre-training data. the meantime, by combining with the inference time technique called Dynamic Scaling, the Dynamic- extension without any fine-tuning.
Dynamic-YaRN YaRN allows for more than 2x context window
BACKGROUND AND RELATED WORK
2.1 ROTARY POSITION EMBEDDINGS
The basis of our work is the Rotary Position Embedding (RoPE) introduced in (Su et al. 2022). We work on a hidden layer where the set of hidden neurons are denoted by D. Given a sequence of vectors "1, , E RIPI, following the notation of (Su et al. 2022). the attention layer first converts the vectors into the query vectors and the key vectors:
Fr
4m fg (mm, m) E R PI, kn fa(#n,n) RI.
Next, the attention weights are calculated
T Kn
softmax
D
where qm, kn are considered as column vectors so that qTd kn is simply the Euclidean inner product. In RoPE, we first assume that [D] is even and identify the embedding space and the hidden states as complex vector spaces:
R where the inner product q" becomes the real part of the standard Hermitian inner product Re(q" k). More specifically, the isomorphisms interleave the real part and the complex part
DI D1/2
((xm) )1, (Im))DI) ((xm)) 1(xm)2, ((Tm)) i(zm))p:))
(3)
((gm)). 1, (qm))p1) ((gm) + ((.m))........ ((qm))(())1))))) +i(g.))))))
(4)
To convert embeddings xm, En into query and key vectors, we are first given RR-lneer operators
Wg, W. : R + IR PI
Let 0 diag(O),...,,6) ,00 12). In complex coordinates, we define
[w (Frm m, 0) im0 W amm:
2

---

Published as a conference paper at ICLR 2024
for any linear operator W. The functions fg. fx in RoPE given by
fq fw., Sn -w.
(6)
where 0₂ = D and b = 10000. This way, RoPE associates each (complex-valued) hidden neuron with separate frequency 0a- The benefit of doing so is that the dot product between the query vector and the key vector only depends on the relative distance m nt.
In later discussions, context length interpolation usually aims modify the equation Eq. 5 To up a uniform convention for these discussions, note that a modification f'w can take the following form:
Sw (mm, m. f/w(wa (m), h()),,
where (mm) is map between real numbers and h(0) acts the entries of the diagonal matrix O uniformly by diag(h(0)), ... p112) according to a function h. and h are method-dependentt functions.
In the subsequent sections, when we introduce a new interpolation method the form Eq. 7 we only specify the functions g(m) and (0a)).
2.2 ADDITIONAL NOTATIONS
Given the pretrained maximal context length L, our goal is to extend it to L' > L either with or without finetuning. We introduce the notion of scale factor defined by 4
For the convenience of some discussions, also introduce wavelength Aa associated with the d-th hidden dimension of RoPE as follows:
2TT
Ad
2mb
(8)
0₂
The wavelength describes the length of tokens needed in order for the rotary position embedding at dimension perform full rotation (2rr..
2.3
RELATED WORK
Position Interpolation (PI) is one of the earlier works extending context lengths of RoPE proposed by
Chen et al. (2023), and concurrently kaiokendev (2023). Under the notation of Eq. 7 is setting
g(m) = 8 m, h( O,
where is the scale factor 4. We include some details in Appendix A.1
ReRoPE (Su 2023) also aims to extend the context size of existing models pre-trained with RoPE, and claims "infinite" context length without needing any fine- fine-tuning. This claim is backed by monotonically decreasing loss with increasing context length up to 16k on the Llama 2 13B model. achieves context extension by modifying the attention mechanism and thus not purely an embedding interpolation method. Since it is currently not compatible with Flash Attention (Dao) 2023) and requires two attention passes during inference, we do not consider it for comparison.
Concurrnntyy with our work, LM-Infinite (Han et al. 2023) proposes similar ideas to YaRN, but focuses on "on-the-fly" length generalization for non-fine-tuned models. Since they also modify the attention mechanism of the models, it is not an embedding interpolation method and is not immediately compatible with Flash Attention 2.
METHODOLOGY
Whereas PI stretches all RoPE dimensions equally, we find that the theoretical interpolation bound described by PI (Chen et al. 2023 insufficient at predicting the complex dynamics between RoPE and the LLM's internal embeddings. In the following subsections. with PI we have individually identified and solved, so as to give the readers the context, origin and justifications of each method which we in concert obtain the full YaRN method.
describe the main issues
3

---

Published as a conference paper at ICLR 2024
3.1 Loss OF HIGH FREQUENCY INFORMATION "NTK- INTERPOLATION
If we look at rotary position embeddings (RoPE) only from an information encoding perspective, it was shown in (Tancik al., 2020), using Neural Tangent Kernel (NTK) theory, that deep neural networks have trouble learning high frequency information if the input dimension is low and the corresponding embeddings lack high frequency components. Here we can see the similarities: a token's s positional information is one-dimensional, and RoPE expands it to an n-dimensional complex vector embedding. RoPE closely resembles Fourier Features (Tancik et al. 2020 in many aspects, as it is possible to define RoPE as special ID case of Fourier Feature.
In the case of Positional Interpolation (PI), as we strech all dimensions equally by a factor 8, removes the high frequency components of RoPE. This degradation is worsened the scaling factor s grows, and at some point, the network will not be able to recover. Previous fine-tunes (kaiokendev
2023) (Chen al. 2023) [Together.ai] 2023) Quesnelle al. 2023) using Pl were only able to achieve a scaling factor of roughly 8 = before the L.M''s 's outputs starts to degrade, even after fine-tuning.
In order to alleviate this issue, the "NTK-aware interpolation was developed in (bloc97 2023a).
Instead of scaling every dimension of RoPE equally by factor s, we spread out the interpolation pressure across multiple dimensions by scaling high frequencies less and low frequencies more. One can obtain such a transformation in many ways, but the simplest would be perform base change on the value of 0. The details are described in the Appendix A.2 and the method has seen some open-source adoption
One main issue of this "NTK-aware" scaling is that it is very difficult to determine what optimal base should be used for an intended context extension by s times. The best base to use for "NTK-aware" interpolation usually has to be found empirically, which significantly increases the difficulty and cost of obtaining a successful fine-tuned model. Despite its limitations, the observations from the NTK theory is valid and the following idea is still maintained and executed in different way in the ENTK-by-pen parrts" interpolation introduced in the next section.
3.2 LOsS OF RELATIVE LOCAL DISTANCES
INTERPOLATION
To understand why "aTK-" interpolation works better than PI and to eliminate its disadvantages, we have to take closer look RoPE. In this section, think heavily in terms of the wavelengths As defined in Eq. g in the formula of RoPE. For simplicity, we omit the subscript d in As and the reader is encouraged to think about the wavelength of an arbitrary periodic function.
In theory, as RoPE is a relative position embedding, it should be quite surprising that it fails to generalize to unseen longer context sizes. However, we can show that in practice, RoPE does not only encode relative position. One observation we can make is that given a context size L, there are some dimensions d where the wavelength is longer than the maximum context length seen during pretraining (A > I.), this suggests that some dimensions' rotary embeddings might not be distributed evenly in the rotational domain (i.c. does not perform a full rotation for the entire training context size). In such cases, we presume having unique position pairs? Pairgiimplies information remains intact in those dimensions. On the contrary, when the wavelength is short, only relative positional information is accessible the network.
implies that the absolute positional
Given these observations, we can see that it is important to not touch the dimensions that only encode relative positional information, they are crucial for the network to distinguish the relative order of nearby tokens. Meanwhile, dimensions that only encode absolute positional information should always be interpolated, as larger distances will be out of distribution. Instead of arbitrarily changing the base in "NTK-award" interpolation (which basically does something similar to what is described here), we can formulate an explicit and targeted interpolation method that takes in account all of the above.
'We note that shortly before the release of this article, Code Llama (Rozière et al. 2023 was released and uses "NTK-aware" scaling by manually scaling the base to 1M, in which they call this method RoPE "adjusted hase frequency" (ABF).
"Since the dimension never rotates fully at least once during pre- preiirainng, if we pick the first token as the anchor, every other token during pre-training has unique distance it, which the neural network can use to determine its absolute position.
4

---

Published as a conference paper at ICLR 2024
In other words,
if the wavelength À is much smaller than the context size L, we do not interpolate;
if the wavelength À is equal to bigger than the context size L, want to only interpolate and avoid any extrapolation (unlike the previous "NTK-aware" method);
dimensions in-between can have bit of both, similar to the "NTK-aware" interpolation.
As a result, is more convenient to introduce the ratio r 4 between the original context size I. and the wavelength A. This ratio represents the number of rotations certain RoPE dimension makes given a fixed pretrained context length L. In the d-th hidden state, the ratio r depends on d in the following way:
r(d)
(10)
A
2 TP1
34
In order to define the boundary of the different interpolation strategies above, we introduce two extra parameters a, All hidden dimensions d where r(d) < a are those where we linearly
interpolate by a scale (exactly like PI, avoiding any extrapolation), and the d where r(d) > B are those where we do not interpolate at all. Define the ramp function 7 to
0,
if < a if
(())
1
> 3
(1I)
otherwise.
B
With the help of the ramp function, the "NTK-by-pen parrtppa method can be described as follows.
Definition 1 The "NTK-by WNK--by-artss interpolation is modification of RoPE using Eq. with the following functions
g(m) =
(12)
h(aa) = (1-
8u + (F((d))
(13)
The values of a and 3 should be tuned on a case-by-case basis. For example, we have found experimentally that for the Ilama family of models, good values for and 8 are and 8 32.
Using the techniques described in this section, a variant of the resulting method was released under the name "NTK-by-parts" interpolation (bloc97, (2023b). This improved method performs better than the previous PI (Chen al., (023) and "NTK-aware" 3. interpolation methods, both with non-fine-tune tuned models and with fine-tuned models, shown in (bloc97 2023b) and Section 4.2
3.3 YARN
In addition to the previous interpolation techniques, we also observe that introducing a temperature on the logits before the attention softmax has uniform impact perplexity regardless of the data sample and the token position over the extended context window (See Appendix[A... (A.)) More precisely, instead of Eq. 2 we modify the computation of attention weights into
T kin
softmax
(14)
WDDI
The reparametrization of RoPE as a of 2D matrices has a clear benefit on the implementation of this attention scaling: we can instead use "length scaling" trick which scales both qm and k, by constant factor / lft by simply scaling the complex rotary position embeddings by the same amount. With this, YaRN can effectively alter the attention mechanism without modifying its code. Furthermore, it has zero overhead during both inference and training, as rotary position embeddings are generated in advance and are reused for all forward passes. Combining with the "NTK-by- by---ts-arts' interpolation, have the YaRN method.
'The interpolation by linear ramp on á may have alternatives, such harmonic mean over 0a/s and lg converted from linear interpolation wavelengths. The choice of h here was for the simplicity of implementation, but both would work.
5

---

Published as a conference paper at ICLR 2024
Definition 2 By the "YaRN method", we refer combination of the attention scaling in Eq. 14 and the "NTK- -y-K-by- interpolation introduced in Section 3.2
For LLaMA and Llama 2 models, we recommend the following values:
V = 0.1 0.11nss)++1
(15)
The equation above found by fitting vItt at the lowest perplexity against the scale extension by various factors using the "NTK-by- by---ts-arts' method (Section B.2) on LLaMA 7b, 13b, 33b and 65b models without fine-tuning. We note that the same values of also apply fairly well to Llama 2 models (7b, 13b and 70b). It suggests that the property of increased entropy and the temperature constant / may have certain degree of "universality" and may be generalizable across some models and training data.
The YaRN method combines all our findings and surpasses all previous methods in both fine-tuned and non-fine-tuned scenarios. Thanks its low footprint, YaRN allows for direct compatibility with libraries that modify the attention mechanism such as Flash Attention Dao 2023
3.4 DYNAMIC SCALING "DYNAMIC NTK" INTERPOLATION
In lot of use cases, multiple forward-passer are performed with varying sequence lengths from to the maximal context size. A typical example is the autoregressive generation where the sequence lengths increment by 1 after each step. There are two ways of applying an interpolation method that uses scale factor (including Pl, "NTK-aware", "NTK-by-parts" and YaRN);
1. Throughout the whole inference cycle, the embedding layer is fixed including the scale factor s = L'/L where L' is the fixed number of extended context size.
2. In each forward-pass, the position embedding updates the scale factor s max(1,1'/L) where (' is the sequence length of the current sequence.
The problem of (1) is that the model may experience a performance discount at a length less than and an abrupt degradation when the sequence length is longer than L'. But by doing Dynamic Scaling as (2), allows the model to gracefully degrade instead of immediately breaking when hitting the trained context limit L'. We call this inference-tie time method the Dynamic Scaling method. When it is combined with INTK-aware interpolation, we call it "Dynamic NTK" interpolation. first appeared in public as reddit post in (emozilla 2023).
One notable fact is that the "Dynamic NTK" interpolation works exceptionally well on models pre- trained on L without any finetuning (L' L). This is supported by the experiment in Appendix B.7
Often in the repeated forward-passes, the kv-caching (Chen 2022) is applied so that we can reuse the previous key-value vectors and improve the overall efficiency. We point out that in some implementations when the rotary position embeddings are cached, some care has to be taken in order to modify it for Dynamic Scaling with kv-caching. The correct implementation should cache the kv-embeddings before applying rotary position embeddings, as the RoPE of every token changes when s changes.
EXPERIMENTS
4.1 TRAINING
We broadly followed the training and evaluation procedures as outlined in Chen et al. (203))
For training the 128k context window size models, extended the Llama 2 (Touvron et al. 2023b 7B and 13B parameter models. No changes were made to the LLaMA model architecture other than the calculation of the embedding frequencies described in Section 3.3 with 16 and 32.
We used a learning rate x 10 with no weight decay and a linear warmup of 20 steps along with AdamW (Loshchilov and Hutter 2019) 31 0.9 and 32 0.95. For the 16 model,

---

Published as a conference paper at ICLR 2024
fine-tuned for 400 steps with global batch size 64 using PyTorch Paszke al. 2019) Fully Sharded Data Parallelism (Zhao et al. (0233) and Flash Attention 2 (Dao 2023) on the PG19 dataset (Rae et al. 2020) chunked into 64k segments bookended with the BOS and EOS token. For
32 followed the same procedure, but due to compute constraints, we started from the finished 8 = 16 checkpoint and trained for only an additional 200 steps. Note that the 8 32 model is also trained with 64k context data, but we show that it is able to extrapolate to a context size of 128k in Section (.2
For the ablation studies, we used the LLaMA 7B model. has the same architecture the newer Llama 2 models except for a shorter pretrained context window size which reduces compute requirements and allows for faster training and evaluations. The training procedure similar the 128k models, but we chunk the PG19 dataset into 32k segments instead, and train using 16 for 400 steps. As shown in Figure 2 YaRN converges faster compared to other interpolation techniques during training and consistently has lower loss.
5.5
NTKaaaare
NiKkky-prrt YaRN
5.0
4.5
4.0
198
2.2
3.5
3.0
2.0
M
1.8
50
100
150
200
250
300
350
400
50
100
350
200
250
300
350
400
Traiiing Steps
Training Steps
Figure 2: Training loss curves for the LLaMA 7B model extended 32k context size using different interpolation techniques. The graph on the right is zoomed in.
4.2 LONG SEQUENCE LANGUAGE MODELING
To evaluate the long sequence language modeling performances, we use the GovReport (Huang et al.) 2021) and Proof-pile (Azerbayev et al. 2022) datasets both of which contain many long sequence samples. For all evaluations, the test splits of both datasets were used exclusively. All perplexity evaluations were calculated using the sliding window method from Press et al. (202)) with S = 256, which takes in account the entire documents' perplexity contribution, even if the context window of the model is shorter.
First, we select 10 random samples from Proof-pile with at least 128k tokens each and evaluate the perplexity of each of these samples when truncated 2k steps from sequence length of 2k tokens through 128k tokens. Table [ssows the long sequence performance of fine-tuned Llama 2 8 = 16 and s 32 models. We demonstrate that YaRN able generalize and extrapolate to unseen context lengths and benefit from transfer learning, since the 32 model was only further trained for 200 steps using the s 16 checkpoint with 64k data and is able to extrapolate 128k context.
In order to further confirm the effectiveness of YaRN, we compare all four interpolation methods in Figure 3 on the left and Table 5 from Appendix [B.Ias as an ablation study. YaRN consistently outperforms (has lower perplexity than) other methods in both non fine-tuned and fine-tuned scenarios when using the same number of training steps. We also demonstrate that YaRN has better training efficiency compared to PI in Appendix B.2 More comparisons against open models can be found in Appendix B.3
4.3 PASSKEY RETRIEVAL
The passkey retrieval task as defined in Mohtashami and Jaggi (2023) measures model's ability to retrieve simple passkey (i.e., a five-digit number) from amongst a large amount of otherwise meaningless text. For our evaluation of the fine-tuned 32k LLaMA 7B models, we performed
"ILLMAA models have pretrained context size of 2k tokens, while Llama 2 models have 4k.
7

---

Published as a conference paper at ICLR 2024
Model Size
Extension Method
Fine- tuned
Training
Extension Scale
Evaluation Context
Window Size
Steps
8192
16384
32768
65536
131072
7B
YaRN
400
4k x 16
3.51
2.99
2.65
2.42
10
7B
YaRN
400+200
4k x 32
3.56
3.04
2.70
2.45
2.37
13B
YaRN
/
400
4k x 16
3.25
2.79
2.50
2.29
10'
13B
YaRN
400+200
4k x 32
3.29
2.83
2.53
2.31
2.24
Table 1: Sliding window perplexity (S = 256) of ten 128k Proof-pile documents over Llama 2 models extended via YaRN. We show successful context size extrapolation and transfer learning from 64k to 128k given only 64k context training data.
6.0
100
H
5.5
NaKayysarts TAHN
Passkey Retreval Accuracy (%a
5.0
3.5
NTKewaree NTKby-arts
3.0
5000
10000
15000
20000
25000
30000
30000
Context
Context Size
Figure 3: Sliding window perplexity (S 256) of ten 128k Proof-pile documents and passkey retrieval accuracy at different prompt lengths for finetuned LLaMA 7B models fine-tuned 32k context for 400 steps using different interpolation techniques. YaRN outperforms other interpolation methods given the same training budget.
50 iterations of the passkey retrieval task with the passkey placed random location uniformly distributed across the evaluation context window on different prompt lengths ranging from 2k to 32k. YaRN achieves higher scores compared to other interpolation methods when given similar training budget, as seen in Figure on the right. More results and comparisons for Llama 2 models are shown in Appendix B.5
4.4 STANDARDIZED BENCHMARKS
The Hugging Face Open LLM Leaderboard (Hugging Face 2023) compares multitude of LLMs across a standardized set of four public benchmarks. Specifically, we use 25-shot ARC- Challenge (Clar al., 2018), 10-shot HellaSwag (Zellers et al. 2019). 5-shot MMLU (Hendrycks et al. 2021) and O-shot TruthfulQA (Lin et al. 2022).
To test the degradation of models' short context performance under context extension, we evaluated our Llama 2 and 32k I.LaMA 7B models using this suite and compared it to established scores for the baselines. The results are summarized in Table and Table More results for Llama models shown in Appendix B.6
observe that there minimal performance degradation between the YaRN models and their respective Llama 2 baselines. Some variance to be expected as the PG19 dataset (Rae al. 2020) we used for fine-tuning is very different from the original pre- mre-trannng datased used for LLaMA and average 0.49% drop in scores between the
Llama models. We also observe that there YaRN s = 16 and s = 32 models and can conclude that the the iterative extension from 64k to 128k results in negligible performance loss.
8

---

Published as a conference paper at ICLR 2024
Extension
Fine- tuned
Extension Scale
ARC-c
Hellaswag
MMLU
TruthfulQA
Method
None
X
51.0
77.8
35.7
34.3
PI
2k x 16
44.8
70.2
25.9
34.1
NTK-aware
2k
47.4
73.9
27.7
32.6
NTK-by-parts
2k x 16
48.5
76.6
32.7
33.4
YaRN
2k x 48.1
77.2
30.0
35.1
Table 2: Performance of context window extensions methods, fine-tuned for 400 steps, the Hugging Face Open LLM benchmark suite compared with original LLaMA baselines.
Model Size
Extension Method
Fine- tuned
Extension Scale 8
ARC-c
Hellaswag
MMLU
TruthfulQA
7B
None
53.1
77.8
43.8
39.0
7B
YaRN
4k x 16
52.3
78.8
42.5
38.2
7B
YaRN
/
4k x 32
52.1
78.4
41.7
37.3
13B
None
59.4
82.1
55.8
37.4
13B
YaRN
4k X 16
58.1
82.3
52.8
37.8
13B
YaRN
32
58.0
82.2
51.9
37.3
Table 3: Performance of YaRN on the Hugging Face Open L.M benchmark suite compared with original
Llama haselines.
4.5 COMPUTATIONAL EFFICIENCY
Given that rotary position embeddings are cached during training and inference when the context window size is fixed to a preset length L. modifying the interpolation on rotary position embeddings incurs additional computational memory cost compared to previous context extension methods, which is the case for all four interpolation methods outlined in this work. YaRN converges the fastest during training compared other methods, thus is the most computationally efficient, shown in Table 4
Model
Model
Extension
Extension Scale s
Effective Context
Training Time in GPU-Hours (A 101))
Size
Name
Method
7B
LLaMA YaRN
YaRN
2k x 16
32k
128
7B
Llama 2 YaRN
YaRN
4k 16
64k
256
7B
Llama 2 YaRN
YaRN
4k x 32
128k
256
128
7B
CCen et al.
2023
PI
2k
16k
640
7B
Together.a al
2023
4k x
32k
7B
(Xiong al.
2025
NTK-
4k 44.2
23 50k
64000
7B
(Rozière el al. 2023)
NwK-aware
4k x 88.6
RJ 100k
6400
Table 4: Comparison of training time A100-hours different open and closed models using different extension methods.
5 CONCLUSION
In conclusion, we have shown that YaRN improves upon all existing RoPE interpolation methods and can act as a drop-in replacement to PI, with no downsides and minimal implementation effort. The fine-tuned models preserve their original abilities multiple benchmarks while being able to attend to a very large context size. Furthermore, YaRN allows efficient extrapolation with fine- tuning shorter datasets and take advantage of transfer learning for faster convergence, both of scenarios. Finally, we have shown the effectiveness of extrapolation with YaRN where it is able to "train short, and test long".
which are crucial under compute-constraing

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY
To aid in reproducibility, we provide, as supplementary material, the entirety of of the code used to train the YaRN models in Table 7 well the evaluation code that produced Figure 7 and Tables 6 TO [ and D The code also contains implementations of various extension methods referenced throughout the paper. For training YaRN, we used the publicly available PG19 dataset (Rae al. 2020) tokenized to contiguous chunks of 64k tokens.
REFERENCES
Introducing Qwen-7B: Open foundation and human-aligned models (of the state-of-the-arts). URL https://github http://eco//////eeeeeeeeeee http///wennn//i///ee 7Btbl///main/teeeeeeeeeeee memo.ma
Z. Azerbayev, E. Ayers. and B. Piotrowski. Proof-pile, 2022. URL
com/
chanoi=-areebbaye//proo--
Black, S. Biderman, Hallahan, Q. Anthony, L. Gao, L. Golding, H. He, C. Leahy, K. McDonell, J. Phang, M. Pieler, U. S. Prashanth, S. Purohit, L. Reynolds, J. Tow, B. Wang, and S. Weinbach. GPT-NeoX-20EB An open-source autoregressive language model, 2022. arXiv: 2204.06745.
bloc97. NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) con- text size without any fine-tuning and minimal perplexity degradation, 2023a. URL
htpps://w.rreddttccoo www. rwwdiriddit.o wtw.ddittcoo/ee/
httpa//aaaa////aa/////////////// sca scaled @eaaled_rope__llne htt.:ppe.aa.llssslaaaaaamoooeoao _o
bloc97. Add NTK-Aware interpolation "by parts" correction, 2023b. URL https://githu:
com htap///uueenallllssseeeeee-oo//p
C. Chen.
Transformer Inference Arithmetic, 2022.
URL https::/kipp.11//bll
S. Chen, S. Wong, L. Chen, and Y. Tian. Extending context window of large language models via positional interpolation, 2023. arXiv: 2306. 15595.
A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, Gehrmann, P. Schuh, K. Shi, S. Tsvyashchenko, J. Maynez, Rao, Barnes, Y. Tay, N. Shazeer, V. Prabhakaran, E. Reif, N. Du, B. Hutchinson, R. Pope, J. Bradbury, J. Austin, M. Isard, G. Gur-Ari, P. Yin, Duke, A. Levskaya, S. Ghemawat, S. Dev, H. Michalewski,
X. Garcia, V. Misra, K. Robinson, L. Fedus, D. Zhou, D. Ippolito. D. Luan, H. Lim, B. Zoph, A. Spiridonov, R. Sepassi, D. Dohan, S. Agrawal, M. Omernick, A. M. Dai, S. Pillai, M. Pellat, A. Lewkowycz, E. Moreira, R. Child, O. Polozov, K. Lee, Z. Zhou, X. Wang, B. Saeta, M. Diaz, O. Firat, M. Catasta, J. Wei, K. MeiMMMMieeeer Hellstern, D. Eck, Dean, S. Petrov, and N. Fiedel. PaLM: Scaling language modeling with pathways, 2022. arXiv: 2204.02311.
Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and Tafjord. Think you have solved question answering? try ARC, the A12 Reasoning Challenge, 2018. arXiv: 1803.05457.
T. Computer. Redpajama: An open source recipe reproduce Ilama training dataset, 2023. URL
T. Dao. Flashattention-2: Faster attention with better parallelism and work partitioning, 2023. arXiv: 2307.08691.
emozilla. Dynamically Scaled RoPE further increases performance of long context LLaMA with fine-tuning, 2023. URL http:::pp//.cc...... www. wtwpddttcco////oaaaaaaaaaacaaaaa htwaaaaa/aMaaM///mats..... hinrpprddynel callaalllll..... Salllyy_calllle sww.sscaleddaeee rope wttpropeeeeeteeeeeeeeeeeeeee
J. Gehring, M. Auli, D. Grangier, D. Yarats, and Y. N. Dauphin. Convolutional sequence to sequence learning. 2017. arXiv: 1705.03122.
C. Han, Q. Wang. W. Xiong, Y. Chen, H. Ji, and S. Wang. LM-Infinite: Simple on-the fly length generalization for large language models, 2023, arXiv:
10

---

Published as a conference paper at ICLR 2024
D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR), 2021.
L. Huang, S. Cao, N. Parulian, H. Ji, and L. Wang. Efficient attentions for long document summariza- tion. In Proceedings of the 2021 Conference of the North American Chapter the Association for Computationnl Linguistics: Human Language Technologies, pages 1419-1436. Association for Computational Linguistics, June 2021.
Hugging Face. Open L.LM Leaderboard, 2023. URL http:tss//uggggggfccccc//////// HuggingF BaaggginFacceH//ppeeee
eade rboard
kaiokendev. Things I'm learning while training superhot., 2023. URL htt.:tts///kaokkeddeae.
github. ic/t 1eltttieexxtetee
A. Kazemnejad, I. Padhi, K. N. Ramamurthy, P. Das, and S. Reddy. The impact of positional encoding on length generalization in transformers, 2023. arXiv:
S. Lin, J. Hilton, and O. Evans. TruthfulQA: Measuring how models mimic human falsehoods. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 3214-3252, May 2022.
I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
A. Mohtashami and M. Jaggi. Landmark attention: Random-access infinite context length for transformers, 2023. arXiv: 2305.16300.
A. Paszke, S. Gross, Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Yang. Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. PyTorch: imperative style, high- seer-oerformeece deep learning library. In NeurlPS, pages 8024-8035, 2019.
O. Press, N. Smith, and M. Lewis. Train Short, Test Long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022.
Quesnelle, E. Shippole, and "Kaiokendev". Llongma: Scaling rotary embeddings htt.gggggggaaae...c/ocoooooooodd http//onceetofiinn/.
through linear positional interpolation. https::/huuginggfaee ca LLongMA-2-7b- 2023.
J. W. Rae, A. Potapenko, S. M. Jayakumar, C. Hillier, and T. P. Lillicrap. Compressive transformers for long-range sequence modelling. In International Conference on Learning Representations, 2020.
A. Roberts, C. Raffel, K. Lee, M. Matena, N. Shazeer, J. Liu, Narang, W. Li, and Y. Zhou. Exploring the limits of transfer learning with unified text-to-text transformer. Technical report, Google, 2019.
B. Rozière, J. Gehring, F. Gloeckle, S. Sootla, I. Gat. X. E. Tan, Adi, J. Liu, T. Remez, J. Rapin, A. Kozhevnikov, I. Evtimov, J. Bitton, M. Bhatt, C. C. Ferrer, A. Grattafiori, W. Xiong, A. Défossez, J. Copet, F. Azhar, H. Touvron, L. Martin, N. Usunier, T. Scialom, and G. Synnaeve. Code Llama: Open foundation models for code, 2023. arXiv: 12950..
P. Shaw, J. Uszkoreit, and A. Vaswani. Self-attention with relative position representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Compu- fational Linguistics: Human Language Technologies, Volume 2 (Short Papers), pages 464 468-468, New Orleans, Louisiana, June 2018. Association for Computational Linguistics..
J. Su. Rectified rotary position embeddings.
https://githuu/ http://wwwb.cccm//oo///////erogo 2023.
J. Su, Lu. S. Pan, A. Murtadha, B. Wen, and Y. Liu. RoFormer: Enhanced transformer with rotary position embedding, 2022. arxiv: 2104.09864.
11

---

Published as a conference paper at ICLR 2024
Y. Sun, L. Dong, B. Patra, S. Ma, S. Huang, A. Benhaim, V. Chaudhary, X. Song, and F. Wei. A anntth-xtrappolttable transformer, 2022. arXiv: 2212.10554.
M. Tancik, P. P. Srinivasan, B. Mildenhall, S. Fridovich-Keil, N. Raghavan, U. Singhal, R. Ra- mamoorthi, J. Barron, and R. Ng. Fourier features let networks learn high frequency functions in low dimensional domains. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS'20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.
Together.i. LLaMA-2-7B-32X, 2023.
URL https://huggingface
Togetherccopptter//aaaee LLLMA-2 2-7B-32K
H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample. LLaMA: Open and efficient foundation language models, 2023a. arXiv: 2302.13971.
H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, D. Bikel, L. Blecher,
Ferrer, M. Chen, G. Cucurull, D. Esiobu, J. Fernandes, J. Fu, W. Fu, B. Fuller, C. Gao, V. Goswami, N. Goyal, A. Hartshorn, S. Hosseini, R. Hou, H. Inan, M. Kardas, Lachaux, T. Lavril, J. Lee, D. Liskovich, Y. Lu, Y. Mao, X. Martinet, T. Mihaylov, P. Mishra, Molybog, Y. Nie, A. Poulton, J. Reizenstein, R. Rungta, K. Saladi, A. Schelten, R. Silva, E. M. Smith, R. Subramanian, X. E. Tan, B. Tang, R. Taylor, A. Williams, J. X. Kuan, P. Xu, Z. Yan, Zarov, Y. Zhang, A. Fan, M. Kambadur, S. Narang, A. Rodriguez, R. Stojnic, S. Edunov, and T. Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023b.
Kerkez, M. Khabsa, I. Kloumann, A. Korenev, P. S. Koura, M.-A.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
W. Xiong, J. Liu, I. Molybog, H. Zhang, Bhargava, R. Hou, L. Martin, R. Rungta, K. A. Sankarara- man, B. Oguz, M. Khabsa, H. Fang, Y. Mehdad, S. Narang, K. Malik, A. Fan, S. Bhosale, S. Edunov, M. Lewis, S. Wang, and H. Ma. Effective long-context scaling of foundation models, 2023.
R. Zellers, A. Holtzman, Bisk, A. Farhadi, and Choi. HellaSwag: Can machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Y. Zhao, A. Gu, Varma, L. Luo, C.-C. Huang, M. Xu, L. Wright, H. Shojanazeri, M. Ott. S. Shleifer, A. Desmaison, C. Balioglu, B. Nguyen, G. Chauhan, Y. Hao, and S. Li. PyTorch FSDP: Experiences on scaling fully sharded data parallel, 2023, arXiv: 2304. 11277
12

---

Published as a conference paper at ICLR 2024
A ADDITIONAL DETAILS ON INTERPOLATION METHODS
A.1 POSITION INTERPOLATION
As mentioned in Section 2.2 PI is one of the earlier works extending context lengths of RoPE. We include some extra details here:
While direct extrapolation does not perform well on sequences w1,-. WL with L larger than the pre-trained limit, they discovered that interpolating the position indicies within the pre-trained limit
works well with the help of small amount of fine-tuning. Specifically, given pre-trained language model with RoPE, they modify the RoPE by
mL
fw (Im: m. 0) fw Em: L'
(16)
where L' >> L a new context window beyond the pre-trained limit. With the original pre-trained model plus the modified RoPE formula, they fine-tuned the language model further several orders of magnitude fewer tokens (a few billion in (Chnn et al. (2023) and successfully acheived context window extension.
A.2 DETAILS OF "NTK-AWARE" INTERPOLATION
In Section 3.1 we introduce a change of basis from b to b' in the definition of "NTK-aware" interpolation method.
Precisely, following the notations set out in Section 2.1 Eq. we define the "NTK-aware" interpola- tion scheme as follows:
Definition The "NTK-aware" interpolation a modification of RoPE using Eq. 7 the following functions, given as the scale factor.
gmm) m
(17)
h(0a) b' 5'-200 2d/|DD
(18)
where
IPI
6' b.s s 10122
(19)
L'
(20)
2023a). this method performs much better at extending the context size of non-fine-tuned models compared to PI (Chen et al. 2023). However, major disadvantage of this method is that given it is not just an interpolation scheme, some dimensions are slightly extrapolated to Yout-of-bounde values, thus fine-tuning with "NTK-aware" interpolation (bloc97 2023a) yields inferior results to PI (Chen et al. 2023). Furthermore, due to the "out-of-bound" values, the theoretical scale factor does accurately describe the true context extension scale. In practice, the scale value s has to be set higher than the expected scale for a given context length extension.
Given the results from (bloc97
The mathematical derivation of the base change is the following:
Recall that our goal is to spread out the interpolation pressure across the hidden dimensions using a base-change instead of scaling the frequencies by fixed factor s. The property we want to guarantee is that: The lowest frequency needs to be scaled as much as linear positional scaling and the highest frequency stay constant.
We introduce new base b' such that the last dimension matches the wavelength of linear interpolation with scale factor s. Since the original RoPE method skips odd dimensions in order to concatenate both cos( 2næ and sin( 2ng components into a single embedding, the last dimension d e D is [D] 2.
A
X
The new base D' can be chosen so that
DI-2 6' 1o
1DI-2
=..
(21)
Solving for b' yields
IP1 B' = b
2
(22)
13

---

Published as a conference paper at ICLR 2024
A.3 THE IMPACT PRE-SOFTMAX SCALING OF YARN ON PERPLEXITY
In Section we mention the impact of the factor inside the softmax computation of attention weights. Here fix 896 16k-token documents from RedPajama (Computer 2023 and calculate
their perplexity scores with different scaling Vt. The result is in Figure 4 For comparison, recall that our recommended factor in this case (s 8) is given by the following.
0.1 + 1 1.208.
(23)
attention pre-softmax scaling final perplexity 56 16k-token documents
(5=8)
9-
8-
7
perplexity
24
1/vt
Figure 4: Fix 8, compare the LLaMA perplexity on 896 16k-token documents over different scaling 1/vt. The shaded area represents standard deviation (68%).
To show the impact of the factor 1/vl on different token positions, we cut each 16k-token document into chunks of 2048 tokens, and further plot the mean perplexity change comparing 1 in percentages
ppl(() ppl(t = 1) ppl(t 1)
(24)
of each chunk. The plot is shown in Figure 5
'We choose RedPajama because it is the open-source dataset closest the training datasct of T. L.L.MA far af.
14

---

Published as a conference paper at ICLR 2024
Perplexity change percentage (from 1/vt=1 to higher) of various segments of tokens in 16k-token samples
Perplexity change percentage
0%
10%%
-20%
204048
2048 4096
6094-6144
-00%
-192
8192 10240
12208 -4336
1/vt
Figure 5: Fix 8 8, compare the mean perplexity change percentages ppl(f) ppl(t = at different ppl(t = 1)
segments of token positions 896 16k-token documents different scaling 1 1/vt.
To further demonstrate the best values across all samples over different token positions, we plot the sample counts with minimal perplexity at a given 1/vž for each of the position segments over the 16k-token range in Figure
The sample counts with minimal perplexity at given 1/Vt
(different lines represent different position segments at the 16k-token samples)
0-2048
2048 4098
4096
500
6144-8195 8192
10240
10240 12288
12288
Counts
14336 16384
200
100
120
1.30
1.35
1/vt
Figure 6: The sample counts (out of the 896 samples) with minimal perplexity at given 1/V2 for given segment of token positions over the 16k-token range.
We observe that:
for suitable i, sample may obtain better perplexity scores across the extended context window;
the best value of is mostly consistent across different samples and different positions.
We remark that this finding is consistent for different values of s and the best value of follows our recommended formula (Eq. 15) closely.
15

---

Published as a conference paper at ICLR 2024
B ADDITIONAL TABLES AND CHARTS
B.1 ABLATION STUDY
Extension
Fine- tuned
Training Steps
Extension Scale
Evaluation Context Window Size
Method
2048
4096
8192
16384
32768
None
4.05
PI
x
4.36
3.90
NTK-aware
X
2k x 2
4.08
5.97
WTK-by-parts
2k x
4.12
3.71
YaRN
X
2k x 2
4.07
3.67
PI
2k 2k x 4
7.09
6.39
6.18
NTK-aware
X
4.27
3.84
10
NTK-by-parts
2k x
4.39
4.03
4.11
YaRN
X
2k x 4
4.19
3.77
3.65
PI
2k x
10
10
10
10
NTK-aware
2k x
4.64
4.27
4.24
10a
NTK-by-parc
2k 8
4.98
4.91
5.33
5.79
YaRN
x x 8
4.37
3,95
3.81
3.33
x 16
10a
10a
10
10a
10a
NTK-aware
2k x 16
5.23
5.02
5.22
6.85
10t
WTK-by-pand
X
2k x 16
6.04
7.54
10
> 10'
> 10'
YaRN
2k x 16
4.61
4.24
4.18
3.66
3.45
PI
Dynamic 4.05
3,90
6.18
10a
10a
NTK-aware
Dynamic 4.05
5.97
10a
10a
> 10a
NTK-by-part
Dynamic
4.05
3.71
4.11
5.79
> 10'
YaRN
Dynamic
4.05
3.67
3.65
3.33
3.45
PI
400
2k x 16
5.70
4.95
4.64
3.97
3.57
NTK-aware
400
2k x 16
4.39
3.92
3.73
3.21
8.49
NTK-by-parts
400
2k x 16
4.14
3.75
3.62
3.12
2.81
YaRN
400
2k x 16
4.19
3.77
3.30
3.09
2.77
Table 5: Sliding window perplexity (S 256) of 128k Proof-pile documents the L.LMMA 7B model extended different methods.
16

---

Published as a conference paper at ICLR 2024
B.2 TRAINING EFFICIENCY OF YARN
Table shows a side-by-side comparison of the Ilama 2 7B model extended from 4096 to 8192 context length via PI (LLongMA-2 7 and YaRN. Note that the PI model was trained using the methodology in Chen et al. 2023), while YaRN used the same methodology but 2.5x less training steps and data, described in Section 4 Even if YaRN was only fine-tuned for 400 steps compared to Pl's 1000 steps, we obtain similar results to PL.
Model Size
Extension
Fine-
Training
Extension Scale
Evaluation Context Window Size
Method
tuned
Steps
2048
4096
6144
8192
7B
None
4.00
3.58
7B
4k x
4.30
3,84
3.83
3.65
7B
YaRN
X
4k x 2
4.03
3.61
3.60
3,49
7B
PI
1000
4k
3.92
3,51
3.51
3.34
7B
YaRN
400
4k x 2
3.91
3.50
3.51
3.35
Table Sliding window perplexity (S 256) of ten 128k Proof-pile documents over the Llama-2 7b model extended via Pl and YaRN with different training steps. YaRN obtains comparable results to Pl using much less training steps.
B.3 COMPARING THE PERPLEXITY OF VARIOUS METHODS OVER A SLIDING WINDOW
We further evaluated the Llama 2 models fine-tuned using YaRN at the scale factor s = 16, 32 and compared them against a few long-context open-source models fine-tuned from Llama-2 such as Together.ai (Together.ai (023) and "NTK-aware" Code Llama (Rozière et al. 2023). The results are summarized in Table (with a more detailed plot in Figure
Model Size
Model
Context Window
Extension Method
Evaluation Context Window Size
Name
8192
32768
65536
98304
131072
7B
Together
32k
PI
3.50
2.64
102
> 10
104
7B
Code Llama
100k
NTK
3.71
2.74
2.55
2.54
2.71
7B
YaRN (s = 16)
64k
YaRN
3.51
2.65
2.42
> 10a
> 10
7B
YaRN (s = 32)
128k
YaRN
3.56
2.70
2.45
2.36
2.37
13B
Code Llama
100k
NTK
3.54
2.63
2.41
2.37
2.54
13B
YaRN (s 16)
64k
YaRN
3.25
2.50
2.29
>> 10a
> 10a
13B
YaRN (s = 32)
128k
YaRN
3.29
2.53
2.31
2.23
2.24
Table 7: Sliding window perplexity (S 256) of ten 128k Proof-pile documents truncated to evaluation context
window size
We observe that the model exhibits strong performance across the entire targeted context size, with YaRN interpolation being the first method successfully extend the effective context size of Llama to 128k.
Furthermore, in Appendix B.4 we show the results of the average perplexity on 50 untruncated GovReport documents with at least 16k tokens per sample evaluated on the setting of 32k maximal context window without Dynamic Scaling Table Similar to the Proof-pile results, the GovReport results show that fine-tuning with YaRN achieves good performance on long sequences.
Table 7 summarizes the results and visualized and more detailed view is presented in Figure here.
I. 7B (Quesnelle et al. 2023 is fine-tuned from Ilama 2 7B, trained at 8k context length with Pl using the RedPajama dataset (Computer 2023
17

---

Published as a conference paper at ICLR 2024
3.8
CodaLlama-138
1
3.6
Yorn Yor-LLama-2-130-222
AggethercomppueelLLMMA2-------
3.4
CodaLlama-7b.ht
Yom--läma-2-76-66n
3.2
Tam-LLama-22778
Perpicetty leploxety(lowcr// is
3.0
2.4
2.2
20000
40000
60000
80000
100000
120000
Window
Figure 7: Sliding window perplexity (S 256) of a 1.28M-token Proof-pile documents truncated to the context window size of the fine-tuned model
B.4 GOVREPORT EVALUATIONS
In Section 4.2
we mention the evaluation on GovReport documents. The evaluation results are
detailed Table below.
Model
Model
Context Window
Extension Method
Perplexity
Size
Name
7B Together
32k
PI
3.67
7B Code Llama
100k
NTK
4.44
7B
YaRN (s 16)
64k
YaRN
3.59
7B
YaRN (s 32)
128k
YaRN
3.64
13B
Code Llama
100k
NTK
4.22
13B
YaRN (s 16)
64k
YaRN
3.35
13B
YaRN (s 32)
128k
YaRN
3.39
Table 8: Sliding window perplexity (S 256) of 50 long GovReport documents with fixed context window
size of 32k
B.5 PASSKEY RETRIEVAL
For our evaluation of the 64k and 128k models, performed 10 iterations of the passkey retrieval task with the passkey placed at a random location uniformly distributed across the evaluation context window different context window sizes ranging from 8k to 128k. Both 7b and 13b models fine-tuned using YaRN at 128k context size passes the passkey retrieval task with very high accuracy (> 99%) within the entire context window size.
18

---

Published as a conference paper at ICLR 2024
Model Size
Model
Scaling Factor
Context Window
Training Data Context
Extension Method
Passkey Context
Passkey Accuracy
Name
7B
Together
32k
32k
PI
32k
100%
7B
Code Llama
88.6
100k
16k
NTK
112k
94.3%
7B
YaRN
16
64k
64k
YaRN
64k
96.3%
7B
YaRN
32
128k
64k
YaRN
128k
99.4%
13B
Code Llama
88.6
100k
16k
NTK
128k
99.4%
13B
YaRN
16
64k
64k
YaRN
64k
97.5%
13B
YaRN
32
128k
64k
YaRN
128k
99.4%
Table 9: Passkey retrieval performance of various models. The passkey context denotes the maximum tested context window size where the accuracy passkey retrieval >= 80%, and the passkey accuracy is the average accuracy of passkey retrieval on all context sizes tested that were smaller equal than passkey context size.
Here we can observe that the lowest perplexity point alone does not provide comprehensive depiction on the "effective context size" that an LLM can attend to. While the Code Llama 13b model exhibits increasing perplexity above 100k context lengths, it was still able to accurately retrieve the passkey at a context length of 128k. This suggest that while output of Code Llama might start to degrade in quality above 100k context size, it is still able to maintain strong retrieval capabilities.
In addition, as YaRN with s = 32 was trained for 200 more steps than YaRN with 16 while having higher passkey accuracy with similar perplexity, we hypothesize that perplexity may not be great indicator of whether an LLM is able to attend to all tokens and does not exhaustively determine long context performance. This also suggests that the YaRN models with 16 might be relatively undertrained for the passkey retrieval task.
B.6 STANDARDIZED BENCHMARKS
To test the degradation of model performance under context extension, we evaluated our models using this suite and compared it to established scores for the Llama 2 baselines as well as publicly available PI and "NTK-aware" models.
Model
Model
Context Window
Extension Method
ARC-c
Hellaswag
MMLU
TruthfulQA
Size
Name
7B
Llama
4k
None
53.1
77.8
43.8
39.0
7B
Together
32k
PI
47.6
76.1
43.3
39.2
7B
Code Llama
100k
NTK-a -a
39.9
60.8
31.1
37.8
7B
YaRN (s 16)
64k
YaRN
52.3
78.8
42.5
38.2
7B
YaRN (s = 32)
128k
YaRN
52.1
78.4
41.7
37.3
13B
Llama 2
4k
None
59.4
82.1
55.8
37.4
13B
Code Llama
100k
NTK-a
40.9
63.4
32.8
43.8
13B
YaRN (s = 16)
64k
YaRN
58.1
82.3
52.8
37.8
13B
YaRN = 32)
128k
YaRN
58.0
82.2
51.9
37.3
Table 10:
Performance of context window extensions methods the Hugging Face Open LLM benchmark suite compared with original Llama 2 basclines
19

---

Published as a conference paper at ICLR 2024
B.7 DYNAMIC SCALING ON MODELS WITHOUT ANY FINE-TUNING
We first recall from Section 3.4 that the Dynamic Scaling technique is an inference-time technique that dynamically update the factor in interpolation methods such as Pl, "NTK-by-parts" and YaRN. We choose the original Llama 2, fix a sample in GovReport and calculate its perplexity on a sliding window of 256 tokens using RoPE, Dynamic-Pi and Dynamic- Maaamic-YaRN.
The comparison between RoPE, Dynamic-Pl and Dynamic-YaRN on Llama-2 7b
Dynamic-P1
Dynamic-YaRN
5.25
Perplexity
5.00
4.50
4.25
2000
3000
4000
5000
6000
Context length (# of tokens)
Figure 8: The comparison between RoPE, Dynamic-Pl and Dynamic- YaRN using Llama 2 on long GovReport sample. This model has not been finetuned for long context.
Since the original maximal context length of Llama is 4096, observe that Dynamic Scaling effec- tively extend the inference length and Dynamic- YaRN achieves better performance than Dynamic- PII The resulting chart is in Figure S
We see that
Dynamic Scaling effectively prevents the blow-up of perplexity score beyond pretrained
context window;
Dynamic-Yare outperforms Dynamic-PI terms of long-range perplexity on pretrained Llama-2 without any finetuning.
20


Reviews:
[{"review_id": "MzBqlMzn3Q", "reviewer": "Reviewer_x5Kh", "summary": "This paper proposes a way of modifying RoPE embeddings so that they extrapolate better to sequences longer than the model was originally trained on. The YaRN method leverages a recent (unpublished) method, called \"NTK-by-parts\" [1], as well as modifying the default temperature value of the attention softmax. The NTK-by-parts method is based on the following observation: Performing position interpolation [2] on the high-frequency (small wavelength) RoPE rotation angles could make it such that neighboring tokens have rotation angles that are very difficult for the model to distinguish, thus potentially confusing the model regarding the relative positions of those tokens. Thus, the NTK-by-parts method functions by determining the amount of position interpolation to perform based on the frequency of the RoPE rotation angles: no interpolation is used for high-frequency angles, while standard interpolation is used for the low-frequency angles (and \"middle\" frequency angles are handled with some interpolation). Once these updates to the RoPE embeddings (and attention temperature) are made, the whole model is fine-tuned on a small amount of data (~0.1% of the pre-training corpus).\n\nThe YaRN method is shown to perform relative well relative to baseline RoPE extension methods across various experiments (e.g., long sequence language modeling, and several standard LLM benchmark tasks from HuggingFace Open LLM leaderboard).\n\n- [1] https://github.com/jquesnelle/yarn/pull/1\n- [2] S. Chen, S. Wong, L. Chen, and Y. Tian. Extending context window of large language models via\npositional interpolation, 2023. arXiv: 2306.15595.", "questions": "- In equation (17), should you be using $b'$ or $b$? It seems more natural to use $b$, and would be confusing to use $b'$, which already includes some scaling modifications.\n- For equations 19 and 20, isn't this equivalent to $g(m) = (1-\\gamma(r(d))) * (1/s) + \\gamma(r(d))$, and $h(\\theta_d) = \\theta_d$? Perhaps this would be simpler, to make this look more like the position interpolation equations (which only modify $g(m)$), as opposed to the \"NTK-aware\" equations which only modify $h(\\theta_d)$.\n- Can you show results of YaRN with and without fine-tuning?\n- Can you clarify how the sliding window evaluation method works (Press et al., 2022)? I read the reference but was still confused, so adding a simple explanation in the paper would be very helpful. When reporting on performance with evaluation context window size $N$, does this mean you measure perplexity on predicting token N+1? Or you measure the average perplexity across all $N$ tokens (so each token on average sees N/2 context)?\n- When you say in your plots \"Yarn-Llama-2-7b-64k\", this means you used s=16? What fine-tuning data and sequence lengths did you use for both s=16 and s=32?\n- Are the only results in your paper that use dynamic scaling those in Figure 5 of appendix B.4?\n- Why do you compare with Code-Llama, which is specialized for code?\n- Do you have intuition for why temperature scaling is necessary? Does it improve performance with regular RoPE embeddings, for short documents? Or just for long documents when using YaRN? Can you add more explanation regarding the intuition/reasons for this method, and how it performs outside the context of YaRN?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The proposed method appears to perform well relative to a few baseline RoPE extension methods.\n- The proposed method is relatively well-motivated --- it seems like a reasonable idea to not interpolate for high-frequency rotation angles.\n- The proposed method has already been used in open-source LLMs in industry.", "weaknesses": "- 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).\n- I think the experimental results could be much more thorough, and much more clearly presented.\n  - Including 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.).\n  - Replacing tables with plots (similar to Figure 4 in appendix, but with clearer+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.\n  - Adding 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.).\n  - Choosing better baselines. Why are Together and Code Llama chosen as the main baselines? Why is PI not included in every table?\n\nOverall, 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!", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dXjxuXuPel", "reviewer": "Reviewer_t3XG", "summary": "This paper presents YaRN, a method to efficiently extend the context window of large language models pretrained with Rotary Position Embeddings (RoPE). The key contributions are:\n\n- Identifying issues with existing interpolation methods like Position Interpolation (PI) and proposing solutions:\n    - \"NTK-aware\" interpolation to preserve high frequency information\n    - \"NTK-by-parts\" interpolation to maintain local distance relationships \n    - Dynamic scaling for inference-time interpolation\n- Introducing an attention temperature hyperparameter for entropy control\n- Combining these techniques into YaRN, which extends context 10x more efficiently than prior work  \n\nThe method is evaluated by fine-tuning LLaMA models and testing on long-context perplexity, retrieving passkeys, and standard benchmarks. Results show YaRN matches or exceeds prior context extension techniques.", "questions": "- For the temperature parameter, did you try other hyperparameters or learning it? Why is attention temperature most effective?\n\n- You mention the temperature provides a uniform perplexity improvement over the context window. Can you elaborate why you think this is the case?\n\n- Can you include results quantifying the impact of each proposed technique independently?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper clearly identifies limitations of prior work on extending context for RoPE models and proposes principled solutions grounded in theory. This shows strong technical understanding.\n\n- The YaRN method is innovative in how it combines multiple solutions in a synergistic way to efficiently extend context. The temperature parameter for controlling entropy is also novel.\n\n- The experiments comprehensively evaluate perplexity on long sequences, passkey retrieval, and standardized benchmarks. The results demonstrate YaRN's capabilities for few-shot context extension and extrapolation. \n\n- The writing is clear and well-structured. The background builds intuition and the methodology explains each component of YaRN in a logical progression.", "weaknesses": "- The ablation study evaluating the impact of each proposed technique in isolation would help validate their necessity in YaRN.\n\n- More analysis and intuition explaining why the temperature parameter works would strengthen that contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2ZNnByd51d", "reviewer": "Reviewer_HiJG", "summary": "The paper presents YaRN (Yet another RoPE extensioN method), a compute-efficient method to extend the context window of large language models (LLMs) trained with Rotary Position Embeddings (RoPE). YaRN requires 10x less tokens and 2.5x less training steps than previous methods. The method combines attention scaling, \"NTK-by-parts\" interpolation, and dynamic scaling to achieve state-of-the-art performance in context window extensions. YaRN is compatible with libraries that modify the attention mechanism, such as Flash Attention, and allows for efficient extrapolation with fine-tuning on shorter datasets.", "questions": "- 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?\n- In what aspects or tasks does the PI method have advantages over YaRN?\n- In what aspects or tasks does the NTK method have advantages over YaRN?\n- Under what circumstances should we use PI? Under what circumstances should we modify the base of rope?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- YaRN achieves state-of-the-art performance in context window extensions with significantly fewer training steps and tokens compared to previous methods.\n- The method is compatible with libraries that modify the attention mechanism, such as Flash Attention 2, making it more versatile for various use cases. It provides a drop-in replacement for Position Interpolation (PI) with no downsides and minimal implementation effort.\n- YaRN allows for efficient extrapolation with fine-tuning on shorter datasets, enabling faster convergence in compute-constrained scenarios. The paper also demonstrates the capability of YaRN to extrapolate beyond the limited context of a fine-tuning dataset, showcasing its ability to generalize to unseen context lengths.\n- The experimental design of the paper is scientific. The authors evaluate YaRN on various benchmarks, including long sequence language modeling, passkey retrieval, and standardized benchmarks. The results consistently show that YaRN outperforms previous methods in context window extensions, supporting the author's claims.", "weaknesses": "- The paper does not provide a comprehensive comparison of YaRN with other length extending methods in terms of computational efficiency and memory usage.\n- 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.\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tvyg4PH3CY", "reviewer": "Reviewer_swin", "summary": "This paper proposed an encoding scheme that could change the after training context window size of a model without sacrificing performance.", "questions": "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.\n\n2. Is there a training-free context window extension stuff that should be included? \n    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.\n\n3. 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?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Motivated and written well.\n2. An important problem to work on.", "weaknesses": "1. 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.\n\n2. Part of baseline in experiments could be missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)