YaRN: Efficient Context Window Extension of Large Language Models

In plain words

A language model reads text as a sequence and needs to know where each word sits in it. One common way to give it that sense of position is to rotate the internal numbers for each word by an angle that depends on its position. Models trained this way work well up to the length they saw in training, and then quickly get worse on longer text.

The obvious fix is to squeeze the longer text into the familiar range of positions. That works a little, but it blurs the fine detail of nearby word order, and the blur gets worse the more you squeeze. YaRN treats the rotation speeds differently depending on how fast they turn. Fast-turning parts already tell nearby words apart, so they are left alone. Slow-turning parts are the ones that would otherwise run out of range, so they are squeezed. A smooth blend sits in between. YaRN also scales the attention scores by a factor that depends on how far the context is stretched.

The result is a model that can read much longer text after only a short extra training run, or in some cases with no training at all. It keeps its original behaviour on short tasks, and it needs no change to the attention code itself, so fast attention kernels still work.

TL;DR

• YaRN extends the context of RoPE models such as Llama 2 by changing only the rotary frequencies and one attention scale factor §3.3.
• The method is the "NTK-by-parts" frequency blend plus a length-based logit scale of 0.1 ln(s) + 1 under the square root, fitted on LLaMA §3.2 §3.3 eq. 15.
• Authors report that 7B and 13B Llama 2 models fine-tuned at 64k or 128k keep short-context benchmark scores near the originals §4.4 tab. 3.
• Code is public (MIT) and matches most of the paper, but the optimizer defaults, the learning-rate decay, the gradient clipping and the dynamic variant differ from the text (see Paper vs code).

Problem

Rotary position embeddings (RoPE) encode position by rotating query and key vectors. Models pretrained with RoPE fail on sequences longer than their training length, even though RoPE is in principle a relative encoding §1 §3.2. Position Interpolation (PI) squeezes all positions into the trained range with one scale factor. It works after fine-tuning, but the squeeze removes high-frequency detail, and fine-tuned PI models degrade beyond a scale of about 8 §3.1. The "NTK-aware" base change avoids that loss but leaves the best base to be found by search, and some dimensions are pushed past their trained range §3.1 app. A.2. Fine-tuning cost is also a concern: the conclusion stresses extension under limited compute §5, and the training-efficiency appendix compares steps and data with PI app. B.2.

Method

The authors write every RoPE variant as a modified rotary function that replaces the position m with g(m) and the frequency θ_d with h(θ_d) §2.1 eq. 7. The base frequency of pair d is θ_d = b^(−2d/|D|) with b = 10000 §2.1. The wavelength λ_d = 2π/θ_d eq. 8 gives the number of rotations r(d) = L/λ_d that pair d makes over the pretrained window L §2.2 eq. 10.

Three steps lead to YaRN.

1. Position Interpolation sets g(m) = s·m with scale s = L'/L and leaves h as the identity §2.3 eq. 9.
2. NTK-by-parts keeps pairs with many rotations unchanged, divides pairs with few rotations by s, and blends in between with a ramp γ(r). The ramp is 0 for r below α, 1 for r above β, and linear in r between them §3.2 eq. 11. The frequency becomes h(θ_d) = (1 − γ) θ_d / s + γ θ_d §3.2 eq. 12 eq. 13. For Llama models the authors found α = 1 and β = 32 work well §3.2.
3. YaRN adds an attention temperature t, so that the logits become q·k / (t√|D|) §3.3 eq. 14. The authors fit √(1/t) = 0.1 ln(s) + 1 to the perplexity minimum for LLaMA 7B, 13B, 33B and 65B without fine-tuning, and report that it also fits Llama 2 7B, 13B and 70B §3.3 eq. 15. The factor is applied by scaling the rotary cos and sin tables by √(1/t), so that both q and k scale and the logits scale by 1/t. That costs nothing at inference or training and keeps FlashAttention 2 compatible §3.3.

A separate inference-time variant, Dynamic Scaling, recomputes the factor s from the current length on each forward pass, as max(1, l'/L) where l' is the current length §3.4. Used with the NTK-aware base change, this is called Dynamic NTK. The authors note that cached rotary keys must be stored before the rotation, since the rotation changes when s changes §3.4.

# YaRN inverse frequencies  [§3.2] [eq. 10] [eq. 11] [eq. 12] [eq. 13]
# Input:  head_dim |D|, base b, pretrained length L, scale s = L'/L
# Output: inv_freq (|D|/2,)

def yarn_inv_freq(head_dim, base, pretrained_len, s, alpha=1, beta=32):
    # RoPE frequency per pair, theta_d = b^(-2d/|D|)  [§2.1]
    theta = base ** (-2 * arange(head_dim // 2) / head_dim)
    # Rotations over the pretrained window, r(d) = L / lambda_d  [eq. 10]
    rotations = pretrained_len * theta / (2 * pi)
    # Ramp: 0 below alpha, 1 above beta, linear in r  [eq. 11]
    gamma = clip((rotations - alpha) / (beta - alpha), 0, 1)
    # Blend between interpolated and unchanged frequency  [eq. 12] [eq. 13]
    inv_freq = (1 - gamma) * theta / s + gamma * theta
    return inv_freq

# Length-based attention factor  [§3.3] [eq. 15]
# Input:  scale s = L'/L
# Output: sqrt(1/t), applied to both q and k

def attention_factor(s):
    # The official code returns 1 for s <= 1
    # [code: scaled_rope/modeling_llama_yarn.py#L122-L125]
    if s <= 1:
        return 1.0
    return 0.1 * log(s) + 1

# Rotary tables built once per context length  [§3.3]
# Input:  length T, inv_freq (|D|/2,), factor sqrt(1/t)
# Output: cos_table, sin_table (T, |D|/2)

def rotary_tables(length, inv_freq, factor):
    # Angle for position m and pair d is m * inv_freq_d  [§2.1] [eq. 7]
    angles = outer(arange(length), inv_freq)
    # Multiplying both tables by the factor scales q and k  [§3.3]
    cos_table = cos(angles) * factor
    sin_table = sin(angles) * factor
    return cos_table, sin_table

# Dynamic Scaling at inference  [§3.4]
# Input:  current length l', pretrained length L
# Output: scale s used for this forward pass

def dynamic_scale(current_len, pretrained_len):
    # Recomputed each forward pass, so it can fall back toward 1  [§3.4]
    return max(1, current_len / pretrained_len)

Hyperparameters

Values in the paper are stated in the text unless noted. "Code" lists the official repository defaults or settings in train.sh, which the paper does not mention. The code paths are relative to the repository at code_commit.

Results

Author-reported. Setups: Llama 2 7B and 13B fine-tuned at s=16 for 400 steps on PG19 64k chunks, with 200 more steps at s=32 for 128k §4.1. Perplexity is a sliding window (S = 256) over ten Proof-pile documents of at least 128k tokens, with the test split only §4.2. Baselines are PI, NTK-aware, NTK-by-parts and published Together and Code Llama models §4.2 app. B.3.

Long-sequence perplexity tab. 1

Open LLM Leaderboard benchmarks, LLaMA 7B, 2k × 16, 400 steps tab. 2

Open LLM Leaderboard benchmarks, Llama 2 tab. 3

Training cost in A100 GPU-hours tab. 4. The paper reports these as given; it does not describe how the other rows were measured.

Ablation at 32k, LLaMA 7B, perplexity by evaluation length tab. 5 (fine-tuned rows, 400 steps, scale 2k × 16)

The same table also has non-fine-tuned rows for all four methods, with YaRN giving the lowest perplexity at 4096 for scale 2k × 2 and the lowest at 8192 and 16384 for scale 2k × 8 tab. 5.

Training efficiency, Llama 2 7B, 4096 to 8192 tab. 6

Comparison with open long-context models tab. 7

GovReport perplexity, 50 documents of at least 16k tokens, window 32k tab. 8

Passkey retrieval tab. 9. Ten trials per context size in the appendix, 50 in the main text for 32k §4.3 app. B.5.

The paper also reports that the 7B and 13B models at 128k pass passkey retrieval with over 99% accuracy across the whole window app. B.5. Table 10 (Open LLM results with Together and Code Llama rows) is not reproduced here app. B.6.

Headline claims in the abstract are that YaRN needs about 10x fewer tokens and 2.5x fewer training steps than earlier methods (abstract). The body supports only the step and data comparison: the YaRN run uses 400 steps against PI's 1000 for similar results app. B.2 tab. 6. The body does not derive the token ratio.

Paper vs code

Repository: jquesnelle/yarn at 995db5b575e75230b3384d658f8b944c9662f775. The paper's authors include the repository owner, and the code carries the MIT licence. Searched first: links in the abstract and the paper itself.

Matches:

• The ramp bounds α = 1 and β = 32 §3.2 are the code defaults beta_slow=1, beta_fast=32 modeling_llama_yarn.py.
• Frequency blend: the code computes inv_freq = interp * (1 − mask) + extrap * mask. Here mask equals the paper's γ, so the result is the paper's h(θ) modeling_llama_yarn.py §3.2 eq. 12 eq. 13.
• Attention factor 0.1 ln(s) + 1 eq. 15 is applied to both cos and sin tables modeling_llama_yarn.py modeling_llama_yarn.py.

Differences:

• Ramp shape. The paper defines γ as linear in the rotation count r §3.2 eq. 11. The code computes the bounds from the log of the rotation count, then makes the ramp linear in the dimension index modeling_llama_yarn.py. Since the dimension index is a log of r, the code's ramp is linear in log r, not in r. The bounds are also rounded to whole dimensions by floor and ceil modeling_llama_yarn.py. This is my reading of the formulas. I did not run the code to confirm it.
• Optimizer. The paper says β2 = 0.95 and no weight decay §4.1. The default path calls torch.optim.AdamW without betas or weight decay finetune.py. That gives PyTorch's defaults, which are β2 = 0.999 and weight decay 0.01 in current PyTorch. The DeepSpeed configs set β2 = 0.95 zero2.json zero3.json, but they run only with --deepspeed.
• Learning-rate decay. The paper states a linear warmup of 20 steps §4.1 and says nothing about decay. The default schedule decays linearly to zero over the run finetune.py. The 70B run in train.sh uses a constant schedule with lr 1e-5 train.sh. The paper does not mention the 70B run.
• Gradient clipping. The paper does not mention clipping. The code clips only when args.grad_norm is a float finetune.py. The flag is store_true, so its value is a bool, and the clip never runs finetune.py.
• Global batch. The paper says 64 §4.1. The code defaults to batch 1 per process with 8 accumulation steps finetune.py, so the global batch depends on the process count set in accelerate. The paper does not say how many GPUs it used.
• Ablation config. The 32k LLaMA 7B ablation §4.1 has no matching command in train.sh. The 8k ablations there use Llama 2 7B train.sh.
• Dynamic scaling. The paper uses s = max(1, l'/L) for each forward pass, with L the pretrained length §3.4. In modeling_llama_yarn.py, the rotary cache is rebuilt only when a sequence is longer than the cache so far, and the scale is then seq_len / max_position_embeddings, not the pretrained length modeling_llama_yarn.py. Shorter inputs later reuse the larger-scale cache, so the scale does not fall back toward 1. The standalone file divides by the pretrained length LlamaDynamicYaRNScaledRotaryEmbedding.py. The config path dynamic-yarn never passes finetuned=True, so it starts as plain RoPE modeling_llama_yarn.py modeling_llama_yarn.py.
• Perplexity protocol. The sliding window S = 256 matches the paper §4.2 perplexity.py.
• Extra models. train.sh also runs Mistral 7B and Llama 2 70B train.sh. The paper does not describe these runs.

Pitfalls

• The attention factor is fitted, not derived. The formula 0.1 ln(s) + 1 comes from a fit on LLaMA 7B to 65B without fine-tuning, and the authors report that it also fits Llama 2 §3.3. The paper does not say whether it holds for other architectures, and the code applies it to any RoPE model with the same constant modeling_llama_yarn.py.
• α and β are per-family choices. The paper says they should be tuned case by case and gives values only for the Llama family §3.2. Other families need their own search.
• Dynamic scaling is ambiguous. The text and the code disagree on the denominator and on whether the scale can fall back §3.4 modeling_llama_yarn.py. An implementer should decide which one they want and test it. The paper's Dynamic-YaRN result is on Llama 2 with no fine-tuning, on a single GovReport sample app. B.7.
• Cached rotary keys go stale. If you cache rotated keys during dynamic scaling, the cache is wrong once s changes. Cache the keys before the rotation §3.4.
• The s = 16 models stop at 64k. They show perplexity above 10 at 131072 tab. 1 and at 98304 tab. 7. The s = 32 model, trained for 200 more steps on 64k data, reaches 128k §4.2 tab. 1.
• Perplexity can understate retrieval. The authors note that the 128k Code Llama model keeps strong passkey accuracy despite rising perplexity, and suggest that perplexity alone does not show effective context length app. B.5.
• Small evaluation samples. Perplexity uses ten Proof-pile documents §4.2. Passkey uses 50 trials at 32k and 10 at 64k and 128k §4.3 app. B.5. Differences of a few tenths of a point, or a few percentage points, are within noise, but the paper does not report variance for these numbers.
• Benchmark drops are averages. The authors report an average 0.49% drop from s = 16 to s = 32 §4.4. The per-task drops in tab. 3 are larger for some tasks.
• Mismatched table reference. The text cites the 32k benchmark table as tab:open_llm, but that label is on the appendix table. The 32k table is labelled tab:open_llm_32k §4.4 app. B.6. When you cite these tables, use the numbers above.
• Hidden compute. YaRN needs no runtime cost: the rotary tables are precomputed §4.5. The fine-tuning cost is in the GPU-hours of tab. 4, which are the authors' own runs. The paper does not state the hardware count.
• The default code does not match the paper's optimizer. To reproduce the paper's fine-tune, set β2 = 0.95 and no weight decay explicitly, and choose between the linear decay and a constant schedule. The default path sets none of these.

Lineage

• Builds on: RoPE, arXiv:2104.09864, the rotary embedding this work modifies §2.1. No card here.
• Builds on: Position Interpolation, arXiv:2306.15595 §2.3 app. A.1. No card here.
• Builds on: NTK-aware base change, NTK-by-parts and Dynamic Scaling, which the paper attributes to unpublished community posts rather than to arXiv §3.1 §3.2 §3.4. Not an arXiv paper.
• Builds on: FlashAttention 2, arXiv:2307.08691, which the implementation uses without change §3.3 §4.1. No card here.
• Builds on: LLaMA, used for the ablation and the attention-factor fit §3.3 §4.1. Card: LLaMA.
• Compared against: Code Llama, card Code Llama, and its NTK-aware variant tab. 4 tab. 7 tab. 9. The paper also compares with Together's PI model tab. 7.
• Compared against: LM-Infinite, arXiv:2308.16137, named as concurrent work that is not compared §2.3.
• Superseded by: not known from content in this wiki.