Context Extension: Why It's Hard, and What Works

Modern LLMs advertise eye-popping context windows: 128K, 1M, 10M tokens. Almost none of them were originally pretrained at those lengths. Behind the marketing is a serious engineering effort — context extension — that takes a model trained at, say, 8K and pushes it to 128K without retraining from scratch.

Why it's hard

Positional encodings are calibrated to the range the model saw during training. A model that learned RoPE rotations across positions 0-8191 has no real signal for what position 100,000 should look like. Naively extending the position scheme produces nonsense — attention patterns become noisy at unfamiliar relative distances, and the model stops referring back to far-away tokens reliably.

What works

Position Interpolation (Chen et al. 2023): for RoPE, interpolate the rotation frequencies so that "position 100,000" in the extended model aligns with "position 8,191" in the original. Then fine-tune briefly on long sequences.
NTK-aware scaling (Reddit/EleutherAI): scale different RoPE frequency dimensions differently. High-frequency dims unchanged (preserves local patterns), low-frequency dims rescaled (accommodates new range).
YaRN (Peng et al. 2023): refines NTK-aware scaling with attention temperature adjustments. Llama 3 used YaRN-style scaling to extend from 8K to 128K with relatively cheap fine-tuning.
Sliding window + global tokens: at architecture-design time, restrict attention to a local window with a few global tokens. Mixtral and Gemma 3 use this; LLaMA 4 Scout's 10M context relies on similar tricks.

Code

Position Interpolation — the core idea·python

# Original: positions 0..L-1 give rotation angles θ_p = p * inv_freq
# Position interpolation: rescale so that effective position is p * (L_old / L_new)
def rope_with_pi(x, L_old, L_new, base=10000.0):
    seq_len, d_head = x.shape[1], x.shape[-1]
    inv_freq = 1.0 / (base ** (torch.arange(0, d_head, 2).float() / d_head))
    scale = L_old / L_new
    t = torch.arange(seq_len, dtype=inv_freq.dtype) * scale
    # ... rest is identical to standard RoPE
    # The key change is multiplying positions by scale < 1
    # so that 'position L_new' in the new range maps to
    # 'position L_old' in the original RoPE range.

Context Extension: Why It's Hard, and What Works

Why it's hard

What works

Code

External links

Exercise

Progress

Comments 0