Module 10: Long Context

Introduction

Your GPT was trained on sequences up to some length \(L\) — say 2048 tokens. Ask it to read 8192 and something strange happens: the output degrades not gradually but sharply, often into gibberish, the moment you pass the length it saw in training. The model did not “run out of memory.” It ran out of positions it recognizes.

Context extension is the set of tricks that let a model trained to length \(L\) work at length \(L' > L\) without retraining from scratch — usually with only a short fine-tune, sometimes with none at all. Every method in this module is a single idea applied to the Rotary Position Embeddings you built in m04: don’t show the model new position-angles it has never seen; rescale the long sequence so it reuses the angles it already knows.

Why it matters for LLMs:

  • The context-length race — 2k → 4k → 32k → 128k → 1M tokens — was won mostly by these cheap post-hoc tricks, not by training at full length from the start.
  • They cost almost nothing: Position Interpolation and YaRN extend a model 8–16× with a few billion fine-tuning tokens, versus the trillions used to pretrain it.
  • They build directly on RoPE (m04) and pair with the KV-cache and efficient attention (m08, m09) that make a long context affordable to serve.

What You’ll Learn

After this module, you can:

  • Explain why RoPE breaks past the training length — the low-frequency pairs rotate to angles the model never saw.
  • Implement Position Interpolation (PI) from scratch: squeeze positions by the scale factor \(s = L'/L\).
  • Implement NTK-aware scaling: raise the frequency base so fast pairs are preserved and slow pairs are stretched.
  • Implement YaRN: interpolate each frequency pair by how many rotations it completes, plus an attention-softmax temperature.
  • Drive a ScaledRotaryEmbedding that swaps between all four behaviors with one argument.

Prerequisites

This module requires familiarity with:

  • Module 04: Embeddings — Rotary Position Embeddings, the theta_i = base^{-2i/d} frequency schedule, and the rotate-don’t-add idea we are about to rescale.
  • Module 09: Efficient Attention — the KV-cache and GQA that make serving the extended context practical.

Intuition: The Context Wall

Recall RoPE from m04. Each 2-D feature pair \(i\) of a query or key is rotated by an angle \(m\,\theta_i\), where \(m\) is the token’s position and \(\theta_i = \text{base}^{-2i/d}\) is that pair’s fixed rotation speed. Fast pairs (small \(i\)) spin almost a full turn every token; slow pairs (large \(i\)) creep around over thousands of tokens.

During training on lengths up to \(L\), the model only ever sees position angles in the band \([0,\ (L-1)\,\theta_i]\). For a fast pair that band already wraps around the circle many times — position \(L{+}100\) looks like some earlier angle it has seen, so extrapolation is harmless. For the slowest pair, though, \((L-1)\,\theta_i\) is a small fraction of a single turn. Push past \(L\) and that pair swings into fresh, never-trained angular territory — and that is the pair whose signal the model leans on for long-range position. That is the wall.

There are two ways over it. Either squeeze the long sequence back into the trained band (interpolate), or stretch the slow pairs’ clock so the same band now spans more tokens. Drive the slowest pair and watch plain RoPE walk right off the edge while a scaling method folds it back inside:

NoteKey Insight

Only the slowest pairs hit the wall, and they hit it hard: at position \(L'\) they sit at an angle the model was never trained on. The fast pairs are fine — they wrapped the circle many times during training, so a bit more rotation is nothing new. Every method below is a way of leaving the fast pairs alone while pulling the slow pairs back into the trained band.

The Math: Three Ways to Rescale RoPE

All three methods start from RoPE’s per-pair frequencies \(\theta_i = \text{base}^{-2i/d}\) and produce a new set. The scale factor is \(s = L'/L\).

1. Position Interpolation (PI). The bluntest fix: divide every position by \(s\). Position \(L'-1\) then lands on the angle the model saw at \((L'-1)/s \approx L-1\). Because \(m\,\theta_i = (m/s)\,\theta_i\) when you fold the \(1/s\) into the frequency, this is identical to dividing every frequency by \(s\):

\[ \theta_i^{\text{PI}} = \frac{\theta_i}{s}. \]

Simple and effective, but it squeezes the fast pairs too — and those didn’t need help, so you lose some local, high-frequency resolution.

2. NTK-aware scaling. Instead of moving the positions, raise the frequency base so the stretch is spread unevenly across pairs — almost nothing for the fast pairs, a lot for the slow ones:

\[ \text{base}' = \text{base}\cdot s^{\,d/(d-2)}, \qquad \theta_i^{\text{NTK}} = (\text{base}')^{-2i/d}. \]

The exponent is chosen so the fastest pair (\(i=0\)) is left exactly alone while the slowest pair is interpolated by roughly \(s\). No fine-tuning is needed for modest extensions, which is why it spread through the open-source community before it had a paper.

3. YaRN (“NTK-by-parts”). Make the per-pair choice explicit. For pair \(i\), count how many full rotations it completes within the original context:

\[ \lambda_i = \frac{2\pi}{\theta_i}, \qquad r_i = \frac{L}{\lambda_i}. \]

A pair that spins many times (\(r_i > \beta\)) is doing local work — leave it alone (extrapolate). A pair that has not even completed one turn (\(r_i < \alpha\)) is doing long-range work — interpolate it fully. Ramp linearly between:

\[ \gamma_i = \operatorname{clamp}\!\left(\frac{r_i - \alpha}{\beta - \alpha},\,0,\,1\right), \qquad \theta_i^{\text{YaRN}} = (1-\gamma_i)\,\frac{\theta_i}{s} + \gamma_i\,\theta_i. \]

With \(\alpha=1,\ \beta=32\) (the LLaMA values), YaRN keeps the fast pairs, fully interpolates the slow pairs, and blends the handful in between — plus a small attention temperature we cover below. It is the method behind most 128k-token open models.

Here is the whole story in one picture: the ratio of each pair’s rescaled frequency to its original. PI drops every pair by the same factor \(1/s\); NTK and YaRN leave the fast pairs near 1 and pull only the slow pairs down.

TipTry This
  1. Watch the shapes diverge. PI is a flat line at \(1/s\) — every pair squeezed equally. NTK curves smoothly from 1 down. YaRN sits on top of NTK for the fast pairs (ratio ≈ 1), then drops to the PI line (\(1/s\)) for the slow pairs — the best of both.
  2. Crank \(s\) to 16. The dashed \(1/s\) floor sinks toward zero; PI drags every pair down with it (blurring local detail), while NTK/YaRN keep the fast pairs pinned near 1.

Code: Scaling from Scratch

long_context.py builds all three schemes as transforms of RoPE’s inverse frequencies. The base schedule is exactly m04’s:

import torch
from long_context import rope_inv_freq

theta = rope_inv_freq(head_dim=64, base=10000.0)
print(f"θ has one entry per pair: {theta.shape}")
print(f"fastest pair θ₀   = {theta[0]:.4f}   (≈ 1 turn / token)")
print(f"slowest pair θ₃₁  = {theta[-1]:.2e}  (≈ 1 turn / {2*3.14159/theta[-1]:.0f} tokens)")
θ has one entry per pair: torch.Size([32])
fastest pair θ₀   = 1.0000   (≈ 1 turn / token)
slowest pair θ₃₁  = 1.33e-04  (≈ 1 turn / 47117 tokens)

Position Interpolation divides every frequency by the scale; NTK raises the base; YaRN blends per pair. Each reduces to plain RoPE when scale=1:

from long_context import (
    linear_interpolation_inv_freq,
    ntk_inv_freq,
    yarn_inv_freq,
)

s = 8.0
pi = linear_interpolation_inv_freq(64, 10000.0, s)
ntk = ntk_inv_freq(64, 10000.0, s)
yarn = yarn_inv_freq(64, 10000.0, s, original_max_position=2048)

print(f"{'pair':>4} {'plain':>10} {'PI':>10} {'NTK':>10} {'YaRN':>10}")
for i in [0, 8, 16, 24, 31]:
    print(f"{i:>4} {theta[i]:>10.2e} {pi[i]:>10.2e} {ntk[i]:>10.2e} {yarn[i]:>10.2e}")
pair      plain         PI        NTK       YaRN
   0   1.00e+00   1.25e-01   1.00e+00   1.00e+00
   8   1.00e-01   1.25e-02   5.85e-02   1.00e-01
  16   1.00e-02   1.25e-03   3.42e-03   1.89e-03
  24   1.00e-03   1.25e-04   2.00e-04   1.25e-04
  31   1.33e-04   1.67e-05   1.67e-05   1.67e-05

Read the top row: pair 0 is untouched by NTK and YaRN (still 1.00) but divided by 8 under PI. Read the bottom row: the slowest pair is interpolated by all three. That is the whole design — preserve fast, interpolate slow — expressed as numbers.

ScaledRotaryEmbedding wraps this into a drop-in replacement for m04’s rotary layer. Pick a method and a scale; it precomputes the cos/sin tables out to the extended length and rotates a query or key tensor by position:

from long_context import ScaledRotaryEmbedding

rope = ScaledRotaryEmbedding(
    head_dim=64, scale=4.0, method="yarn", original_max_position=2048
)

q = torch.randn(1, 8, 8192, 64)   # (batch, heads, seq=8192 > 2048, head_dim)
q_rot = rope(q)
print(f"Rotated a {q.shape[-2]}-token sequence: {tuple(q_rot.shape)}")
print(f"YaRN attention-temperature factor: {rope.attention_scale:.4f}")
Rotated a 8192-token sequence: (1, 8, 8192, 64)
YaRN attention-temperature factor: 1.1386

RoPE’s signature property — attention depends only on the relative offset — survives the rescaling within the extended range. demonstrate_context_extension makes the wall quantitative, tracking the slowest pair’s angle at the target length relative to the trained band:

from long_context import demonstrate_context_extension

ratio = demonstrate_context_extension(
    head_dim=64, train_len=2048, target_len=8192
)
print(f"\nPlain RoPE asks the model to read angles {ratio:.1f}× outside its band.")
==============================================================
CONTEXT EXTENSION - slowest RoPE pair
==============================================================
  head_dim=64, base=10000, train_len=2048, target_len=8192  (scale s=4)

  Trained max angle for the slowest pair: 0.273 rad
  Angle at target_len / trained max  (1.0 = stays in trained band):
    none  (plain RoPE, extrapolates)     4.00x
    linear (Position Interpolation)      1.00x
    ntk   (NTK-aware base)               1.00x
    yarn  (NTK-by-parts)                 1.00x

Plain RoPE asks the model to read angles 4.0× outside its band.

YaRN: Interpolate Per Dimension

The heart of YaRN is that ramp \(\gamma_i\) — the decision, pair by pair, of keep versus interpolate. It splits the pairs into three zones by how many rotations they complete within the training length. Step through them:

Beyond the frequencies, YaRN adds one more correction: it multiplies the attention logits by a temperature \(\sqrt{1/t} = 0.1\ln(s) + 1\) before the softmax. Stretching positions slightly flattens the attention distribution; this factor re-sharpens it and recovers a bit of the perplexity lost to interpolation. ScaledRotaryEmbedding folds it into the cos/sin tables, so a YaRN layer is still drop-in:

from long_context import yarn_attention_scale

for scale in (1, 2, 4, 8, 16):
    print(f"s = {scale:>2} →  attention scale √(1/t) = {yarn_attention_scale(scale):.4f}")
s =  1 →  attention scale √(1/t) = 1.0000
s =  2 →  attention scale √(1/t) = 1.0693
s =  4 →  attention scale √(1/t) = 1.1386
s =  8 →  attention scale √(1/t) = 1.2079
s = 16 →  attention scale √(1/t) = 1.2773
NoteKey Insight

NTK-aware scaling and YaRN are the same instinct — don’t touch the fast pairs — at two levels of precision. NTK bakes the taper into a single raised base; YaRN makes the keep/interpolate decision explicit per pair (with a ramp) and adds the softmax temperature. When you see a model card say “128k context via YaRN,” this ramp is what it means.

Measuring It: Needle in a Haystack

Extending the positions is only half the job — you have to check the model can still use the far context. The standard probe is needle-in-a-haystack: hide a single fact (the “needle”) at a controlled depth inside a long filler document (the “haystack”), then ask a question only that sentence answers. Sweep the needle’s depth and the total length, and you get a 2-D grid of pass/fail.

  • Plain RoPE past \(L\): the grid turns red as soon as the needle sits beyond the training length — the model literally cannot address that position.
  • PI / NTK / YaRN: the grid stays green much further out, with YaRN typically holding longest before the far-depth cells start to fade.

Perplexity tells the same story more cheaply: measured on held-out long documents, it stays flat under a good scaling method and explodes the moment plain RoPE crosses \(L\). The interactive at the top of this lesson is exactly that crossing, one frequency pair at a time.

WarningExtension is not free capability

Rescaling positions lets the model attend to far tokens; it does not teach it to reason over them. A short fine-tune at the target length (a few billion tokens) is what turns “can address position 100k” into “can actually use position 100k.” YaRN needs the least of this fine-tuning, which is much of why it won.

Common Pitfalls

When extending context, watch out for:

  1. Forgetting scale = 1 must be identity. A correct implementation leaves a model untouched inside its training length. If your PI/NTK/YaRN frequencies differ from plain RoPE at scale=1, you have a bug — the tests assert this.
  2. Interpolating the fast pairs (plain PI). PI’s uniform squeeze blurs local, high-frequency position information. NTK/YaRN exist precisely to spare those pairs; reach for them past ~4× extension.
  3. Extending without any fine-tuning and expecting magic. NTK-aware buys a modest zero-shot extension; larger jumps (8–16×) need a short fine-tune at the new length or quality still sags.
  4. Dropping YaRN’s attention temperature. The frequency ramp alone recovers most of the quality, but the \(\sqrt{1/t}\) softmax scaling is part of the method; omitting it leaves perplexity measurably higher.
  5. Mismatched train/inference scaling. The scale and base used at inference must match what the model was fine-tuned with. A model tuned for YaRN at 8× will misbehave if you serve it with plain RoPE, and vice versa.

Exercises

Exercise 1: PI folds positions back into the band

import torch
from long_context import linear_interpolation_inv_freq, rope_inv_freq

# Show that Position Interpolation at scale s makes position L' land on roughly the
# trained angle at (L')/s. Compare the slowest pair's angle at L'=8192 under PI(s=4)
# against plain RoPE's angle at L=2048.

# Your implementation here:
# pi = linear_interpolation_inv_freq(64, 10000.0, 4.0)
# ...

Exercise 2: The NTK taper

import torch
from long_context import ntk_inv_freq, rope_inv_freq

# Confirm NTK-aware scaling leaves the fastest pair (index 0) unchanged while
# interpolating the slowest pair. Print the per-pair ratio ntk/plain for a few
# indices and check it decreases from ~1 toward ~1/s.

# Your implementation here:

Exercise 3: Build a YaRN ramp by hand

import torch
from long_context import yarn_ramp

# For head_dim=64, base=10000, L=2048, compute gamma and report: how many pairs are
# kept (gamma≈1), how many are fully interpolated (gamma≈0), and how many are in the
# ramp in between? Change L to 512 and watch the boundaries move.

# Your implementation here:

Summary

Key takeaways:

  1. RoPE breaks past the training length because of the slow pairs — they reach position-angles the model never saw, while the fast pairs (which wrap the circle many times) extrapolate harmlessly.
  2. Every method rescales RoPE’s frequencies \(\theta_i\) — the scale factor is \(s = L'/L\), and all methods reduce to plain RoPE at \(s=1\).
  3. Position Interpolation divides every frequency by \(s\) — simple and effective, but it squeezes the fast pairs it didn’t need to.
  4. NTK-aware raises the base, \(\text{base}\cdot s^{d/(d-2)}\), tapering the stretch so fast pairs are preserved and slow pairs interpolated — often with no fine-tuning.
  5. YaRN interpolates per pair by rotation count (\(\gamma_i\) ramp, \(\alpha=1\), \(\beta=32\)) and adds a \(\sqrt{1/t}=0.1\ln s + 1\) attention temperature — the method behind most long-context open models.
  6. Extension enables addressing, not comprehension — a short fine-tune at the target length is what turns reachable far positions into usable ones, and needle-in-a-haystack / long-document perplexity are how you measure it.

What’s Next

You can now stretch a from-scratch GPT far past its training length. The frontier keeps going: Mixture of Experts grows a model’s knowledge without growing its per-token compute (sparse FFNs and top-\(k\) routing), and fast inference (quantization, speculative decoding) makes serving all of it cheap. Both reuse the efficient-attention and KV-cache machinery from m09 that a long context depends on.

Going Deeper

Core Papers:

Practical Resources: