Context Length Extension

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Why Models Fail at Longer Contexts
5. Position Interpolation
6. NTK-Aware Scaling
7. YaRN: Yet Another RoPE Extension
8. Sliding Window + Attention Sinks
9. Ring Attention (Distributed Context)
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Explain why RoPE fails at positions beyond training length.
2. Derive position interpolation and its effect on attention patterns.
3. Explain NTK-aware scaling and frequency decomposition.
4. Describe how ring attention distributes long sequences across GPUs.
5. Analyze the quality-context tradeoff for extended models.

Notation

• $L_{\text{train}}$ — training context length
• $L_{\text{target}}$ — desired context length ($L_{\text{target}} > L_{\text{train}}$)
• $\alpha = L_{\text{target}}/L_{\text{train}}$ — extension ratio
• $\theta_{\text{base}} = 10000$ — RoPE base frequency

Core Intuition

A model trained with max position 4096 has never seen RoPE rotations at position 8000. At test time, positions beyond 4096 produce attention patterns the model has never encountered — leading to catastrophic perplexity degradation. Extension methods modify the positional encoding to map longer sequences into the "familiar" range, often requiring minimal fine-tuning.

Why Models Fail at Longer Contexts

RoPE at position $p$: Rotation angle $\theta_i = p/\theta_{\text{base}}^{2i/d}$.

At training: Model sees angles up to $\theta_{\max} = L_{\text{train}}/\theta_{\text{base}}^{2i/d}$.

At inference with $p > L_{\text{train}}$: Angles exceed anything seen in training. High-frequency dimensions (small $i$) extrapolate worst — they rotate past their training range.

Result: Attention scores become random at extended positions → perplexity explosion.

Position Interpolation

Idea (Chen et al., 2023): Scale positions to fit within training range:

Position 8000 in a 4K→8K extension becomes position 4000 (within training range).

Effect: All rotation angles stay within $[0, \theta_{\max}]$. But adjacent tokens now have smaller angular differences → attention is less sharp.

Fine-tuning needed: ~1000 steps on extended-length data to adapt to the compressed positional signal. Without fine-tuning: moderate quality degradation.

NTK-Aware Scaling

Problem with uniform interpolation: High-frequency dimensions (which encode local position) get compressed — the model loses fine-grained positional information.

NTK-aware scaling: Scale the base frequency instead:

Effect: Low-frequency dimensions (global position) are scaled more; high-frequency dimensions (local position) are preserved.

Intuition: Like adjusting the "zoom" of a Fourier representation — extend the range without losing resolution at fine scales.

Advantage over interpolation: Better preserves local attention patterns; less fine-tuning needed.

YaRN: Yet Another RoPE Extension

YaRN (Peng et al., 2023): Combines NTK scaling with attention temperature adjustment:

where $t > 1$ is a temperature factor that compensates for the changed attention distribution.

Additionally: Applies different scaling ratios to different frequency bands:

• High-frequency (local): no scaling (preserve local attention).
• Low-frequency (global): full interpolation (extend range).
• Mid-frequency: smooth interpolation between extremes.

Result: Extends 4K → 128K with only 400 steps of fine-tuning. Highest quality among RoPE extensions.

Sliding Window + Attention Sinks

StreamingLLM (Xiao et al., 2023): For infinite-length streaming without any modification to RoPE:

Cache structure: $\{x_1, \ldots, x_s\} \cup \{x_{t-w+1}, \ldots, x_t\}$ (first $s$ "sink" tokens + last $w$ tokens).

Attention sink phenomenon: The first few tokens (regardless of content) receive high attention scores. They act as "sinks" that absorb attention mass. Removing them causes perplexity explosion even for nearby tokens.

Explanation: Softmax must sum to 1. When no token is strongly relevant, attention defaults to a "dump" position (the first tokens serve this role through training).

Constant memory: Only need $s + w$ KV-cache entries regardless of context length.

Ring Attention (Distributed Context)

Problem: Even with extensions, a single GPU can't hold KV-cache for 1M+ tokens.

Ring attention (Liu et al., 2023): Distribute the sequence across $N$ GPUs in a ring:

• Each GPU holds $T/N$ tokens of the KV-cache.
• Queries are broadcast around the ring.
• Each GPU computes partial attention with its local keys.
• Results are combined using online softmax (same as FlashAttention).

Complexity per GPU: $O((T/N)^2 d)$ compute, $O(T/N)$ memory.

Enables: 1M+ token contexts by distributing across 8–64 GPUs.

Common Pitfalls

Pitfall 1. Extending context without any fine-tuning. Even the best extension methods degrade without at least a few hundred steps on extended-length data.

Pitfall 2. Assuming longer context = better performance. "Lost in the middle" phenomenon: models with 128K context often fail to retrieve information from the middle of long prompts.

Pitfall 3. Extending too aggressively in one step. 4K → 128K (32x extension) degrades significantly. Progressive extension (4K → 16K → 64K → 128K) works better.

Summary

• Models fail at longer contexts because RoPE angles exceed training range.
• Position interpolation: Scale positions down; simple but loses local resolution.
• NTK-aware scaling: Scale base frequency; preserves high-frequency local information.
• YaRN: Best quality; frequency-dependent scaling + temperature compensation.
• StreamingLLM: Infinite context with fixed memory using sink tokens + sliding window.
• Ring attention: Distribute across GPUs for million-token contexts.

Exercises

Exercise 1. Compute the RoPE rotation angle at position 8192 for dimension $i=0$ with $\theta_{\text{base}}=10000, d=128$. Is it within training range for $L_{\text{train}}=4096$?

Exercise 2. Derive the NTK-aware $\theta'_{\text{base}}$ for extending from 4K to 32K with $d=128$.

Exercise 3. Explain the "lost in the middle" phenomenon: why do models attend to the beginning and end but not the middle of long contexts?

Exercise 4. Compute the memory requirements for ring attention with $T=1M$ tokens across 8 GPUs (model: 7B, $d_k=128, L=32, G=8$).

Exercise 5. Design an experiment to determine the optimal number of "attention sink" tokens $s$ for a given model.