Context-window extension

Overview

Models trained with rotary position embeddings (RoPE) degrade on sequences longer than the length they saw in training. The shared goal of every method here is to keep the positions a model sees inside a range it handles well, or to make it handle the new range, at the lowest cost in short-context quality 2309.00071 §1 2309.00071 §3.2. The invariant all variants keep: the rotary function is unchanged in form, and only its frequencies and one attention scale change 2309.00071 §2.1 2309.00071 §3.3.

Only one paper with a card is available so far: YaRN. Its text names the earlier methods it builds on, Position Interpolation and the NTK-aware base change, but no cards exist for them yet, so this page cannot compare them from their own content 2309.00071 §3.1 2309.00071 §2.3.

Variants

Position Interpolation (named, no card)

Scales every position by s = L'/L so that the new length fits in the trained range, with no change to the frequencies 2309.00071 §2.3 2309.00071 eq. 9. The paper reports that fine-tuned PI models degrade past a scale of about 8, because the squeeze removes high-frequency detail 2309.00071 §3.1. Reference: arXiv:2306.15595.

NTK-aware base change (named, no card)

Changes the RoPE base instead of the positions. The paper attributes it to an unpublished community post, not to arXiv, and says the best base must be searched 2309.00071 §3.1. It can push some dimensions past their trained range 2309.00071 app. A.2.

NTK-by-parts

Interpolates only the slow-turning dimensions and leaves the fast ones unchanged, with a ramp between them. Its frequency is h(θ_d) = θ_d for the fast pairs and θ_d / s for the slow pairs 2309.00071 §3.2. The paper gives it as a baseline. With fine-tuning on LLaMA 7B its benchmark scores are close to YaRN's, and it is ahead on MMLU and ARC-c but behind on TruthfulQA 2309.00071 tab. 2. On the 32k ablation it is behind YaRN at the longest evaluation length 2309.00071 tab. 5.

YaRN

Uses the same per-dimension ramp as NTK-by-parts, with bounds α = 1 and β = 32 for Llama models, and adds an attention temperature. Logits become q·k / (t√|D|), and √(1/t) = 0.1 ln(s) + 1 is applied to both query and key by scaling the rotary tables 2309.00071 §3.2 2309.00071 §3.3 2309.00071 eq. 14 2309.00071 eq. 15. Card: YaRN.

Reported results, all author-reported: 7B and 13B Llama 2 models fine-tuned for 400 steps at s = 16 keep short-context benchmark scores near the originals. A further 200 steps at s = 32 reach 128k, with passkey accuracy of 99.4% 2309.00071 §4.1 2309.00071 tab. 3 2309.00071 tab. 9.

Dynamic scaling (Dynamic NTK and Dynamic YaRN)

Recomputes s on each forward pass as max(1, l'/L), where l' is the current length and L the pretrained length 2309.00071 §3.4. Needs no fine-tuning. The paper's Dynamic-YaRN result is on Llama 2 without fine-tuning, on one GovReport sample 2309.00071 app. B.7. Cached rotated keys become wrong when s changes, so keys must be cached before the rotation 2309.00071 §3.4.

When to use which

Note on the 128k result: the s = 16 models stop at 64k, and only the s = 32 model reaches 128k 2309.00071 tab. 1. Authors also suggest that perplexity alone does not show effective context length 2309.00071 app. B.5.

Implementation sketch

Frequencies with the ramp, then the length-based factor. Single-pass RoPE is the same as before, with the new inverse frequencies and the factor on cos and sin. Python-flavoured.

# YaRN inverse frequencies  [arxiv:2309.00071 §3.2] [eq. 11] [eq. 12]
# 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  [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 interpolated and unchanged frequencies  [eq. 12] [eq. 13]
    return (1 - gamma) * theta / s + gamma * theta

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

def attention_factor(s):
    # Fitted, not derived, on LLaMA  [§3.3]
    if s <= 1:
        return 1.0
    return 0.1 * log(s) + 1

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

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

Contradictions

• Dynamic scaling denominator and fall-back. The text divides by the pretrained length and lets the scale fall back toward 1 2309.00071 §3.4. The released model code divides by the model's maximum position embeddings and rebuilds its cache only for longer inputs, so the scale does not fall back 2309.00071 §3.4. This is a difference between the paper's text and its code, not between two papers. The card reports it in full.
• Ramp shape. The text makes the ramp linear in the rotation count r 2309.00071 eq. 11. The code makes it linear in the dimension index, which is linear in log r 2309.00071 §3.2. This is the card author's reading of the formulas, not a run of the code.
• Optimizer. The text gives β2 = 0.95 and no weight decay 2309.00071 §4.1. The default code path sets neither.

Other papers: none known. Only one card lists this concept.

Papers

• YaRN: Efficient Context Window Extension of Large Language Models: NTK-by-parts frequency blend, length-based attention factor, dynamic scaling, and fine-tune results for Llama 2 at 64k and 128k.