{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Module 10: Long Context" ], "id": "222414c3-a035-4291-a974-c27de7792535" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [], "id": "b326d4a6-6c59-48c6-a214-1cdc14c76eaf" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "9c00fea8-89e0-4626-908b-43f15e9dab28" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Introduction\n", "\n", "Your GPT was trained on sequences up to some length $L$ — say 2048\n", "tokens. Ask it to read 8192 and something strange happens: the output\n", "degrades not gradually but sharply, often into gibberish, the moment you\n", "pass the length it saw in training. The model did not “run out of\n", "memory.” It ran out of **positions it recognizes**.\n", "\n", "**Context extension** is the set of tricks that let a model trained to\n", "length $L$ work at length $L' > L$ *without retraining from scratch* —\n", "usually with only a short fine-tune, sometimes with none at all. Every\n", "method in this module is a single idea applied to the Rotary Position\n", "Embeddings you built in m04: **don’t show the model new position-angles\n", "it has never seen; rescale the long sequence so it reuses the angles it\n", "already knows.**\n", "\n", "Why it matters for LLMs:\n", "\n", "- The context-length race — 2k → 4k → 32k → 128k → 1M tokens — was won\n", " mostly by these cheap post-hoc tricks, not by training at full length\n", " from the start.\n", "- They cost almost nothing: **Position Interpolation** and **YaRN**\n", " extend a model 8–16× with a few billion fine-tuning tokens, versus the\n", " trillions used to pretrain it.\n", "- They build directly on RoPE (m04) and pair with the KV-cache and\n", " efficient attention (m08, m09) that make a long context affordable to\n", " serve.\n", "\n", "### What You’ll Learn\n", "\n", "After this module, you can:\n", "\n", "- Explain **why** RoPE breaks past the training length — the\n", " low-frequency pairs rotate to angles the model never saw.\n", "- Implement **Position Interpolation** (PI) from scratch: squeeze\n", " positions by the scale factor $s = L'/L$.\n", "- Implement **NTK-aware** scaling: raise the frequency *base* so fast\n", " pairs are preserved and slow pairs are stretched.\n", "- Implement **YaRN**: interpolate each frequency pair by how many\n", " rotations it completes, plus an attention-softmax temperature.\n", "- Drive a `ScaledRotaryEmbedding` that swaps between all four behaviors\n", " with one argument.\n", "\n", "### Prerequisites\n", "\n", "This module requires familiarity with:\n", "\n", "- [Module 04: Embeddings](../m04_embeddings/lesson.qmd) — Rotary\n", " Position Embeddings, the `theta_i = base^{-2i/d}` frequency schedule,\n", " and the rotate-don’t-add idea we are about to rescale.\n", "- [Module 09: Efficient\n", " Attention](../m09_efficient_attention/lesson.qmd) — the KV-cache and\n", " GQA that make serving the extended context practical.\n", "\n", "## Intuition: The Context Wall\n", "\n", "Recall RoPE from m04. Each 2-D feature pair $i$ of a query or key is\n", "rotated by an angle $m\\,\\theta_i$, where $m$ is the token’s position and\n", "$\\theta_i =\n", "\\text{base}^{-2i/d}$ is that pair’s fixed rotation speed. Fast pairs\n", "(small $i$) spin almost a full turn every token; slow pairs (large $i$)\n", "creep around over thousands of tokens.\n", "\n", "During training on lengths up to $L$, the model only ever sees position\n", "angles in the band $[0,\\ (L-1)\\,\\theta_i]$. For a **fast** pair that\n", "band already wraps around the circle many times — position $L{+}100$\n", "looks like some earlier angle it has seen, so extrapolation is harmless.\n", "For the **slowest** pair, though, $(L-1)\\,\\theta_i$ is a small fraction\n", "of a single turn. Push past $L$ and that pair swings into fresh,\n", "never-trained angular territory — and that is the pair whose signal the\n", "model leans on for *long-range* position. That is the wall.\n", "\n", "There are two ways over it. Either **squeeze the long sequence back into\n", "the trained band** (interpolate), or **stretch the slow pairs’ clock**\n", "so the same band now spans more tokens. Drive the slowest pair and watch\n", "plain RoPE walk right off the edge while a scaling method folds it back\n", "inside:" ], "id": "db68d640-0c71-4eb2-8256-521a12fe05c9" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "d29daf9d-0423-444c-9e31-710d39572328" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "eadb74b8-986e-48fc-a0bd-94df21394f3b" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "17b31f86-5b8f-473d-9dd5-a090f700a5f5" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "8281fc5c-15f4-4174-9d42-7923606f6a42" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "2836b10c-5999-4932-bb80-929f29382af6" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "9eda5e94-d5b7-4454-8e63-4d37a8350157" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> Only the *slowest* pairs hit the wall, and they hit it hard: at\n", "> position $L'$ they sit at an angle the model was never trained on. The\n", "> fast pairs are fine — they wrapped the circle many times during\n", "> training, so a bit more rotation is nothing new. Every method below is\n", "> a way of leaving the fast pairs alone while pulling the slow pairs\n", "> back into the trained band.\n", "\n", "## The Math: Three Ways to Rescale RoPE\n", "\n", "All three methods start from RoPE’s per-pair frequencies $\\theta_i =\n", "\\text{base}^{-2i/d}$ and produce a new set. The scale factor is\n", "$s = L'/L$.\n", "\n", "**1. Position Interpolation (PI).** The bluntest fix: divide every\n", "position by $s$. Position $L'-1$ then lands on the angle the model saw\n", "at $(L'-1)/s \\approx\n", "L-1$. Because $m\\,\\theta_i = (m/s)\\,\\theta_i$ when you fold the $1/s$\n", "into the frequency, this is identical to dividing every frequency by\n", "$s$:\n", "\n", "$$\\theta_i^{\\text{PI}} = \\frac{\\theta_i}{s}.$$\n", "\n", "Simple and effective, but it squeezes the *fast* pairs too — and those\n", "didn’t need help, so you lose some local, high-frequency resolution.\n", "\n", "**2. NTK-aware scaling.** Instead of moving the positions, raise the\n", "frequency base so the stretch is spread *unevenly* across pairs — almost\n", "nothing for the fast pairs, a lot for the slow ones:\n", "\n", "$$\\text{base}' = \\text{base}\\cdot s^{\\,d/(d-2)}, \\qquad\n", "\\theta_i^{\\text{NTK}} = (\\text{base}')^{-2i/d}.$$\n", "\n", "The exponent is chosen so the fastest pair ($i=0$) is left exactly alone\n", "while the slowest pair is interpolated by roughly $s$. No fine-tuning is\n", "needed for modest extensions, which is why it spread through the\n", "open-source community before it had a paper.\n", "\n", "**3. YaRN (“NTK-by-parts”).** Make the per-pair choice explicit. For\n", "pair $i$, count how many full rotations it completes within the original\n", "context:\n", "\n", "$$\\lambda_i = \\frac{2\\pi}{\\theta_i}, \\qquad r_i = \\frac{L}{\\lambda_i}.$$\n", "\n", "A pair that spins many times ($r_i > \\beta$) is doing *local* work —\n", "leave it alone (extrapolate). A pair that has not even completed one\n", "turn ($r_i < \\alpha$) is doing *long-range* work — interpolate it fully.\n", "Ramp linearly between:\n", "\n", "$$\\gamma_i = \\operatorname{clamp}\\!\\left(\\frac{r_i - \\alpha}{\\beta - \\alpha},\\,0,\\,1\\right),\n", "\\qquad\n", "\\theta_i^{\\text{YaRN}} = (1-\\gamma_i)\\,\\frac{\\theta_i}{s} + \\gamma_i\\,\\theta_i.$$\n", "\n", "With $\\alpha=1,\\ \\beta=32$ (the LLaMA values), YaRN keeps the fast\n", "pairs, fully interpolates the slow pairs, and blends the handful in\n", "between — plus a small attention temperature we cover below. It is the\n", "method behind most 128k-token open models.\n", "\n", "Here is the whole story in one picture: the **ratio** of each pair’s\n", "rescaled frequency to its original. PI drops every pair by the same\n", "factor $1/s$; NTK and YaRN leave the fast pairs near 1 and pull only the\n", "slow pairs down." ], "id": "49562027-e0da-4f67-85b9-2c9c46f083c4" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "3d67ad6c-11b7-4705-a076-dddc01739672" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "a4773978-d430-4723-87f7-0cc2a8c133f0" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "d06e5793-4ec1-46d9-832a-138d38464ea8" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Try This**\n", ">\n", "> 1. **Watch the shapes diverge.** PI is a flat line at $1/s$ — every\n", "> pair squeezed equally. NTK curves smoothly from 1 down. YaRN sits\n", "> *on top of* NTK for the fast pairs (ratio ≈ 1), then drops to the\n", "> PI line ($1/s$) for the slow pairs — the best of both.\n", "> 2. **Crank $s$ to 16.** The dashed $1/s$ floor sinks toward zero; PI\n", "> drags every pair down with it (blurring local detail), while\n", "> NTK/YaRN keep the fast pairs pinned near 1.\n", "\n", "## Code: Scaling from Scratch\n", "\n", "`long_context.py` builds all three schemes as transforms of RoPE’s\n", "inverse frequencies. The base schedule is exactly m04’s:" ], "id": "2a5fdca4-da24-42bf-8e15-cb67c06c4bb1" }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "θ has one entry per pair: torch.Size([32])\n", "fastest pair θ₀ = 1.0000 (≈ 1 turn / token)\n", "slowest pair θ₃₁ = 1.33e-04 (≈ 1 turn / 47117 tokens)" ] } ], "source": [ "import torch\n", "from long_context import rope_inv_freq\n", "\n", "theta = rope_inv_freq(head_dim=64, base=10000.0)\n", "print(f\"θ has one entry per pair: {theta.shape}\")\n", "print(f\"fastest pair θ₀ = {theta[0]:.4f} (≈ 1 turn / token)\")\n", "print(f\"slowest pair θ₃₁ = {theta[-1]:.2e} (≈ 1 turn / {2*3.14159/theta[-1]:.0f} tokens)\")" ], "id": "9d80abca" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Position Interpolation divides every frequency by the scale; NTK raises\n", "the base; YaRN blends per pair. Each reduces to plain RoPE when\n", "`scale=1`:" ], "id": "cd9e22d2-4d51-43dc-a2d4-d3bbf7c088c0" }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "pair plain PI NTK YaRN\n", " 0 1.00e+00 1.25e-01 1.00e+00 1.00e+00\n", " 8 1.00e-01 1.25e-02 5.85e-02 1.00e-01\n", " 16 1.00e-02 1.25e-03 3.42e-03 1.89e-03\n", " 24 1.00e-03 1.25e-04 2.00e-04 1.25e-04\n", " 31 1.33e-04 1.67e-05 1.67e-05 1.67e-05" ] } ], "source": [ "from long_context import (\n", " linear_interpolation_inv_freq,\n", " ntk_inv_freq,\n", " yarn_inv_freq,\n", ")\n", "\n", "s = 8.0\n", "pi = linear_interpolation_inv_freq(64, 10000.0, s)\n", "ntk = ntk_inv_freq(64, 10000.0, s)\n", "yarn = yarn_inv_freq(64, 10000.0, s, original_max_position=2048)\n", "\n", "print(f\"{'pair':>4} {'plain':>10} {'PI':>10} {'NTK':>10} {'YaRN':>10}\")\n", "for i in [0, 8, 16, 24, 31]:\n", " print(f\"{i:>4} {theta[i]:>10.2e} {pi[i]:>10.2e} {ntk[i]:>10.2e} {yarn[i]:>10.2e}\")" ], "id": "dd27864f" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Read the top row: **pair 0 is untouched by NTK and YaRN** (still `1.00`)\n", "but divided by 8 under PI. Read the bottom row: the **slowest pair is\n", "interpolated by all three**. That is the whole design — preserve fast,\n", "interpolate slow — expressed as numbers.\n", "\n", "`ScaledRotaryEmbedding` wraps this into a drop-in replacement for m04’s\n", "rotary layer. Pick a `method` and a `scale`; it precomputes the cos/sin\n", "tables out to the extended length and rotates a query or key tensor by\n", "position:" ], "id": "d17fbfe0-1fec-4b18-90b2-714e3cabe524" }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Rotated a 8192-token sequence: (1, 8, 8192, 64)\n", "YaRN attention-temperature factor: 1.1386" ] } ], "source": [ "from long_context import ScaledRotaryEmbedding\n", "\n", "rope = ScaledRotaryEmbedding(\n", " head_dim=64, scale=4.0, method=\"yarn\", original_max_position=2048\n", ")\n", "\n", "q = torch.randn(1, 8, 8192, 64) # (batch, heads, seq=8192 > 2048, head_dim)\n", "q_rot = rope(q)\n", "print(f\"Rotated a {q.shape[-2]}-token sequence: {tuple(q_rot.shape)}\")\n", "print(f\"YaRN attention-temperature factor: {rope.attention_scale:.4f}\")" ], "id": "b9d98e15" }, { "cell_type": "markdown", "metadata": {}, "source": [ "RoPE’s signature property — attention depends only on the *relative*\n", "offset — survives the rescaling within the extended range.\n", "`demonstrate_context_extension` makes the wall quantitative, tracking\n", "the slowest pair’s angle at the target length relative to the trained\n", "band:" ], "id": "e9e39efc-9134-4189-a1c9-036015bafd27" }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "==============================================================\n", "CONTEXT EXTENSION - slowest RoPE pair\n", "==============================================================\n", " head_dim=64, base=10000, train_len=2048, target_len=8192 (scale s=4)\n", "\n", " Trained max angle for the slowest pair: 0.273 rad\n", " Angle at target_len / trained max (1.0 = stays in trained band):\n", " none (plain RoPE, extrapolates) 4.00x\n", " linear (Position Interpolation) 1.00x\n", " ntk (NTK-aware base) 1.00x\n", " yarn (NTK-by-parts) 1.00x\n", "\n", "Plain RoPE asks the model to read angles 4.0× outside its band." ] } ], "source": [ "from long_context import demonstrate_context_extension\n", "\n", "ratio = demonstrate_context_extension(\n", " head_dim=64, train_len=2048, target_len=8192\n", ")\n", "print(f\"\\nPlain RoPE asks the model to read angles {ratio:.1f}× outside its band.\")" ], "id": "c593467c" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## YaRN: Interpolate Per Dimension\n", "\n", "The heart of YaRN is that ramp $\\gamma_i$ — the decision, pair by pair,\n", "of *keep* versus *interpolate*. It splits the pairs into three zones by\n", "how many rotations they complete within the training length. Step\n", "through them:" ], "id": "8fdf853b-cedf-4844-b54c-a39d9aa336db" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "e2807d86-fb07-4aa4-90ff-1ad574fbdd60" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "392b1cf5-50ea-4a1d-b347-fdcc5b785ae7" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "83595e4b-a3c3-4be7-b59f-d275aeef3e1a" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Beyond the frequencies, YaRN adds one more correction: it multiplies the\n", "attention logits by a temperature $\\sqrt{1/t} = 0.1\\ln(s) + 1$ before\n", "the softmax. Stretching positions slightly flattens the attention\n", "distribution; this factor re-sharpens it and recovers a bit of the\n", "perplexity lost to interpolation. `ScaledRotaryEmbedding` folds it into\n", "the cos/sin tables, so a YaRN layer is still drop-in:" ], "id": "b7d41147-8b64-4d94-93ab-857be4a50fed" }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "s = 1 → attention scale √(1/t) = 1.0000\n", "s = 2 → attention scale √(1/t) = 1.0693\n", "s = 4 → attention scale √(1/t) = 1.1386\n", "s = 8 → attention scale √(1/t) = 1.2079\n", "s = 16 → attention scale √(1/t) = 1.2773" ] } ], "source": [ "from long_context import yarn_attention_scale\n", "\n", "for scale in (1, 2, 4, 8, 16):\n", " print(f\"s = {scale:>2} → attention scale √(1/t) = {yarn_attention_scale(scale):.4f}\")" ], "id": "a134743b" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> NTK-aware scaling and YaRN are the same instinct — *don’t touch the\n", "> fast pairs* — at two levels of precision. NTK bakes the taper into a\n", "> single raised base; YaRN makes the keep/interpolate decision explicit\n", "> per pair (with a ramp) and adds the softmax temperature. When you see\n", "> a model card say “128k context via YaRN,” this ramp is what it means.\n", "\n", "## Measuring It: Needle in a Haystack\n", "\n", "Extending the *positions* is only half the job — you have to check the\n", "model can still *use* the far context. The standard probe is\n", "**needle-in-a-haystack**: hide a single fact (the “needle”) at a\n", "controlled depth inside a long filler document (the “haystack”), then\n", "ask a question only that sentence answers. Sweep the needle’s depth and\n", "the total length, and you get a 2-D grid of pass/fail.\n", "\n", "- **Plain RoPE past $L$**: the grid turns red as soon as the needle sits\n", " beyond the training length — the model literally cannot address that\n", " position.\n", "- **PI / NTK / YaRN**: the grid stays green much further out, with YaRN\n", " typically holding longest before the far-depth cells start to fade.\n", "\n", "Perplexity tells the same story more cheaply: measured on held-out long\n", "documents, it stays flat under a good scaling method and **explodes**\n", "the moment plain RoPE crosses $L$. The interactive at the top of this\n", "lesson is exactly that crossing, one frequency pair at a time.\n", "\n", "> **Extension is not free capability**\n", ">\n", "> Rescaling positions lets the model *attend* to far tokens; it does not\n", "> teach it to *reason* over them. A short fine-tune at the target length\n", "> (a few billion tokens) is what turns “can address position 100k” into\n", "> “can actually use position 100k.” YaRN needs the least of this\n", "> fine-tuning, which is much of why it won.\n", "\n", "## Common Pitfalls\n", "\n", "When extending context, watch out for:\n", "\n", "1. **Forgetting `scale = 1` must be identity.** A correct\n", " implementation leaves a model untouched inside its training length.\n", " If your PI/NTK/YaRN frequencies differ from plain RoPE at `scale=1`,\n", " you have a bug — the tests assert this.\n", "2. **Interpolating the fast pairs (plain PI).** PI’s uniform squeeze\n", " blurs local, high-frequency position information. NTK/YaRN exist\n", " precisely to spare those pairs; reach for them past ~4× extension.\n", "3. **Extending without any fine-tuning and expecting magic.** NTK-aware\n", " buys a modest zero-shot extension; larger jumps (8–16×) need a short\n", " fine-tune at the new length or quality still sags.\n", "4. **Dropping YaRN’s attention temperature.** The frequency ramp alone\n", " recovers most of the quality, but the $\\sqrt{1/t}$ softmax scaling\n", " is part of the method; omitting it leaves perplexity measurably\n", " higher.\n", "5. **Mismatched train/inference scaling.** The `scale` and `base` used\n", " at inference must match what the model was fine-tuned with. A model\n", " tuned for YaRN at 8× will misbehave if you serve it with plain RoPE,\n", " and vice versa.\n", "\n", "## Exercises\n", "\n", "### Exercise 1: PI folds positions back into the band" ], "id": "5aae08b4-b0d4-43ad-8b69-6337968929d8" }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [ "import torch\n", "from long_context import linear_interpolation_inv_freq, rope_inv_freq\n", "\n", "# Show that Position Interpolation at scale s makes position L' land on roughly the\n", "# trained angle at (L')/s. Compare the slowest pair's angle at L'=8192 under PI(s=4)\n", "# against plain RoPE's angle at L=2048.\n", "\n", "# Your implementation here:\n", "# pi = linear_interpolation_inv_freq(64, 10000.0, 4.0)\n", "# ..." ], "id": "50620dbc" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 2: The NTK taper" ], "id": "3d931b9f-22b5-4828-b1d4-dc0f4c7f90fd" }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "import torch\n", "from long_context import ntk_inv_freq, rope_inv_freq\n", "\n", "# Confirm NTK-aware scaling leaves the fastest pair (index 0) unchanged while\n", "# interpolating the slowest pair. Print the per-pair ratio ntk/plain for a few\n", "# indices and check it decreases from ~1 toward ~1/s.\n", "\n", "# Your implementation here:" ], "id": "c6968063" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 3: Build a YaRN ramp by hand" ], "id": "e743340c-4d1f-444c-b746-1b01066983c5" }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [], "source": [ "import torch\n", "from long_context import yarn_ramp\n", "\n", "# For head_dim=64, base=10000, L=2048, compute gamma and report: how many pairs are\n", "# kept (gamma≈1), how many are fully interpolated (gamma≈0), and how many are in the\n", "# ramp in between? Change L to 512 and watch the boundaries move.\n", "\n", "# Your implementation here:" ], "id": "d15f9230" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Summary\n", "\n", "Key takeaways:\n", "\n", "1. **RoPE breaks past the training length because of the slow pairs** —\n", " they reach position-angles the model never saw, while the fast pairs\n", " (which wrap the circle many times) extrapolate harmlessly.\n", "2. **Every method rescales RoPE’s frequencies $\\theta_i$** — the scale\n", " factor is $s = L'/L$, and all methods reduce to plain RoPE at $s=1$.\n", "3. **Position Interpolation divides every frequency by $s$** — simple\n", " and effective, but it squeezes the fast pairs it didn’t need to.\n", "4. **NTK-aware raises the base**, $\\text{base}\\cdot s^{d/(d-2)}$,\n", " tapering the stretch so fast pairs are preserved and slow pairs\n", " interpolated — often with no fine-tuning.\n", "5. **YaRN interpolates per pair** by rotation count ($\\gamma_i$ ramp,\n", " $\\alpha=1$, $\\beta=32$) and adds a $\\sqrt{1/t}=0.1\\ln s + 1$\n", " attention temperature — the method behind most long-context open\n", " models.\n", "6. **Extension enables addressing, not comprehension** — a short\n", " fine-tune at the target length is what turns reachable far positions\n", " into usable ones, and needle-in-a-haystack / long-document\n", " perplexity are how you measure it.\n", "\n", "## What’s Next\n", "\n", "You can now stretch a from-scratch GPT far past its training length. The\n", "frontier keeps going: **Mixture of Experts** grows a model’s knowledge\n", "without growing its per-token compute (sparse FFNs and top-$k$ routing),\n", "and **fast inference** (quantization, speculative decoding) makes\n", "serving all of it cheap. Both reuse the efficient-attention and KV-cache\n", "machinery from m09 that a long context depends on.\n", "\n", "### Going Deeper\n", "\n", "**Core Papers:**\n", "\n", "- [Extending Context Window of Large Language Models via Positional\n", " Interpolation](https://arxiv.org/abs/2306.15595) — Chen et al. (2023),\n", " Position Interpolation.\n", "- [YaRN: Efficient Context Window Extension of Large Language\n", " Models](https://arxiv.org/abs/2309.00071) — Peng et al. (2023),\n", " NTK-by-parts + attention temperature.\n", "- [RoFormer: Enhanced Transformer with Rotary Position\n", " Embedding](https://arxiv.org/abs/2104.09864) — Su et al. (2021), the\n", " RoPE this all builds on.\n", "\n", "**Practical Resources:**\n", "\n", "- [NTK-Aware Scaled\n", " RoPE](https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/)\n", " — bloc97 (2023), the community post that introduced NTK-aware scaling.\n", "- [Extending the RoPE](https://blog.eleuther.ai/yarn/) — EleutherAI’s\n", " annotated walkthrough of PI, NTK, and YaRN.\n", "- [Lost in the Middle: How Language Models Use Long\n", " Contexts](https://arxiv.org/abs/2307.03172) — Liu et al. (2023), why\n", " long-context *quality* (not just length) is the real bar." ], "id": "d8282fa3-39b9-4e49-8cc5-dea12822ec34" } ], "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "name": "python3", "display_name": "Python 3 (ipykernel)", "language": "python", "path": "/opt/buildhome/.asdf/installs/python/3.11.14/share/jupyter/kernels/python3" }, "language_info": { "name": "python", "codemirror_mode": { "name": "ipython", "version": "3" }, "file_extension": ".py", "mimetype": "text/x-python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.14" } } }