DL S078 · Long Context — RoPE Scaling, YaRN, ALiBi

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 78 / 80 · Module M12 · ~2 hours

The story we're starting with

Your model was trained with max sequence length 4096. You give it a 5000-token input at inference and something horrible happens: perplexity explodes, generation turns to babble, sometimes the model outputs literal noise. It's not that the model "forgets" the far context — it's that it has literally never seen position embeddings for positions 4096–5000. The rotary matrices at those positions produce values the attention pattern has never been calibrated for. It's an out-of-distribution query at the positional level.

For a while (2020–2022) the answer was "retrain the whole thing at longer context." That is expensive. Then Su et al. published RoPE in 2021, someone noticed a decade later that RoPE has a frequency-domain interpretation, and suddenly a series of clever tricks appeared: Position Interpolation (Chen et al., Meta, 2023) said "just scale positions down and fine-tune briefly." NTK-aware RoPE (a Reddit post, of all things, in mid-2023) said "don't scale uniformly — scale high-freq and low-freq differently to preserve local structure." YaRN (Peng et al., late 2023) systematized this and added a temperature correction. Result: you can take Llama-2-7B pretrained at 4k and extend it to 128k with a few hundred fine-tuning steps. 32× context for ~0.5% of the compute cost of pretraining.

Today we look at what RoPE is doing frequency-wise, why naive extrapolation breaks, and how each of the three main extension techniques (PI, NTK-RoPE, YaRN) fixes the failure. We'll also touch ALiBi — an alternative positional scheme that avoided the problem in the first place — and why the field went with RoPE anyway.

Prerequisites

• S042 · RoPE (rotary positional embeddings). Non-negotiable. If you don't have RoPE math firm, revisit that session first.
• S036 · Self-attention. You need to know how position info enters the Q·K^T product.
• S062 · Fine-tuning basics. Every extension technique benefits from a short fine-tune; you need to know what that means.

1 · RoPE, one-minute refresher

For a token at position m and a head-dim pair (2i, 2i+1), RoPE rotates the (Q or K) subvector by angle m · θ_i, where:

θ_i = 10000^(-2i/d_head)         for i = 0, 1, ..., d_head/2 - 1

• i = 0 (dim pair 0,1): θ ≈ 1.0 rad per position. High frequency — completes a full rotation every ~6 tokens.
• i = d/2 − 1: θ ≈ 1/10000 rad per position. Low frequency — near-static across thousands of tokens.

The attention score between position m and position n (in that dim pair) ends up depending on cos((m − n) · θ_i). Relative position is what determines attention. That's the whole point of RoPE.

1.1 · Why extrapolation breaks

At training time, the model sees m in [0, L_train). The attention head learns, for each frequency, what values of cos((m − n) · θ_i) correspond to "attend" vs "don't attend."

At inference, if you feed m > L_train, the high-frequency dimensions produce cosine values in ranges the model has never seen at any (m, n) combination during training. Because these rotations are periodic, they even map to values that during training corresponded to different relative distances — so the model may confuse a token 5000 tokens ago with one 20 tokens ago. Chaos.

Low-frequency dims are safe (they haven't completed a rotation yet at m=5000), but the high-frequency ones are the problem.

2 · Position Interpolation (PI)

Idea: if the model only knows positions [0, L_train), just squeeze the extended range into that box. Given a new max length L_new = s · L_train (scale factor s, say s=4 for 4k→16k), replace position m with m / s. Now m ∈ [0, L_new) maps to [0, L_train).

def rope_positions_pi(positions, s):
    return positions / s

That's it. One line.

2.1 · Why it works better than extrapolation

You're now querying the model at fractional positions it hasn't seen (m=0.25, 0.5, 0.75, …) — but crucially, still within the trained range. The cosine functions vary smoothly, so interpolation into never-seen inputs produces reasonable outputs. Extrapolation into never-seen ranges does not.

2.2 · The catch

Because positions are compressed by s, the effective per-position resolution drops. Adjacent tokens now have position difference 1/s instead of 1. High-frequency dims that used to distinguish token 5 from token 6 now barely rotate between adjacent tokens — you've lost local resolution.

Fix: a short fine-tune (1000 steps on long sequences) recovers most of the loss. PI + fine-tune works well up to ~4× extension.

3 · NTK-aware RoPE — the free-lunch trick

Someone on r/LocalLLaMA (yes, really — the original was a Reddit post) noticed that PI scales all frequencies equally, which is dumb. Low-frequency dims don't need scaling (they weren't near saturation). High-frequency dims are the ones about to break.

NTK-aware scaling interpolates only low frequencies (which is the same as PI for them) and extrapolates high frequencies (safe because high-freq rotations are naturally periodic and can extend beyond training range without OOD issues in that specific sense).

The clean way to do it: change the RoPE base from 10000 to a larger value. Specifically:

base_new = base_old * (s ** (d_head / (d_head - 2)))

For s=4, d_head=128, this is roughly 10000 * 4^(128/126) ≈ 10000 * 4.06 ≈ 40600.

Effect on θ_i = base^(-2i/d_head):

• Small i (low freq): θ_i barely changes — like uniform scaling.
• Large i (high freq): θ_i shrinks noticeably — high-freq rotations "slow down", so at extended positions they don't exceed their training-range angles as fast.

Zero fine-tuning required. Works up to ~2× extension without fine-tuning, ~4× with a very short fine-tune. This is what made "extend Llama to 8k in an afternoon" viable.

4 · YaRN — the systematic version

YaRN (Yet another RoPE Note) formalized the mess. Two contributions:

4.1 · Piecewise frequency scaling

Don't apply one uniform base change. Split RoPE frequencies into three bands:

• High freq (period much less than L_train): extrapolate (no scaling).
• Low freq (period much more than L_train): interpolate (PI-style scaling by s).
• Mid freq: smoothly blend between the two using a ramp function.

This is finer-grained than NTK's implicit smooth curve, and empirically better for large extension factors (8×–32×).

4.2 · The temperature correction

Here's the subtle one. When you interpolate positions, attention scores get compressed too (all m − n differences shrink). This changes the effective softmax temperature — attention becomes softer and the model's "sharpness" of picking one token to attend to drops. YaRN fixes this by scaling the attention logits by a temperature t = 0.1 · ln(s) + 1:

logits = (Q · K^T) * t / sqrt(d_head)

That's a small correction (t ≈ 1.14 for s=4) but it's the last few points of perplexity recovery.

4.3 · Recipe

• Apply YaRN's frequency scaling to RoPE.
• Multiply attention logits by the temperature correction.
• Fine-tune for a few hundred steps on ~1B tokens of long sequences (Books3, arxiv full-papers).
• You now have a model with 32k+ effective context at a rounding-error training cost.

5 · ALiBi — the road not taken

Before RoPE dominated, Press et al. proposed ALiBi (Attention with Linear Biases). Instead of adding position info to Q/K, add a fixed linear bias to the attention scores based on distance:

attention_scores[i, j] -= m_h * |i - j|

where m_h is a per-head slope. That's it. No positional embedding, no rotation, no fancy math. Each head learns its own preferred "attention decay" rate.

Long-context property: because the bias is a smooth linear function of distance, ALiBi extrapolates to sequences longer than training with almost no degradation. It was long-context-friendly by design.

Why the field went RoPE anyway: ALiBi's linear decay is aggressive — attention scores at large distances go to −∞ quickly. That's fine for local tasks but hurts long-range copying and needle-in-haystack retrieval. RoPE + YaRN, once figured out, dominates on retrieval-heavy long-context benchmarks. So ALiBi lives on in some architectures (BLOOM, MPT) but Llama-family + all the extension tricks won mindshare.

6 · Empirical numbers to remember

• PI (no fine-tune): breaks past ~1.5× extension. With ~1000 steps fine-tune: usable to 4×.
• NTK-aware (no fine-tune): usable to ~2×, mostly good to 4× with brief fine-tune.
• YaRN (no fine-tune): usable to ~4× with mild degradation. With fine-tune on ~1B long-context tokens: 8×–32× with near-zero perplexity increase.
• Needle-in-a-haystack at 32k: YaRN-tuned Llama-2 recovers > 90% needle recall; PI-only recovers ~50%.

7 · War stories

8 · Diagram — where each technique lives on the RoPE frequency axis

9 · The 2024–2025 long-context frontier

YaRN was 2023's peak. Since then, three things happened: (a) proprietary models pushed to 1M–10M tokens; (b) new techniques went beyond RoPE-scaling; (c) evaluations got harder so we learned how much of the marketing was real.

9.1 · LongRoPE (Microsoft, Feb 2024, arXiv:2402.13753)

Ding et al. asked: why apply the same scale factor s to every RoPE dim? LongRoPE runs an evolutionary search over per-dimension scale factors, discovering that different dims want different rescalings. Result: Mistral-7B extended to 2,048k context (2M tokens) with only a short fine-tune, keeping short-context performance. This is the algorithm behind Phi-3-mini's 128k variant.

9.2 · Ring Attention & Blockwise Parallel (Liu et al., 2023–2024)

Ring Attention (arXiv:2310.01889) is the systems answer, not an algorithmic one. Attention memory is O(seq²), so a 1M-token forward pass on one GPU is impossible. Ring Attention shards the sequence across N devices, has each device compute attention against its local chunk, then rotates K/V blocks around a ring so every device eventually sees every K/V — with communication overlapped by compute. Combined with Blockwise Parallel Transformer for the MLP, this is how Gemini 1.5's 2M and Llama 4 Scout's 10M contexts are physically trainable.

9.3 · StreamingLLM and attention sinks (Xiao et al., ICLR 2024, arXiv:2309.17453)

What if you don't want a bigger context — you want an infinite stream (chat that never ends, log monitor)? Naive sliding-window attention destroys quality after the window slides. StreamingLLM discovered the attention sink: the first few tokens of any sequence absorb enormous attention mass, and if you evict them the model collapses. Fix: always keep the first 4 tokens ("sinks") + the recent window. You get infinite-length generation with bounded memory and near-zero quality loss. Every serving framework (vLLM, TGI, SGLang) now supports this as a mode.

9.4 · Alternative architectures — Mamba, Jamba, RWKV, Griffin

Transformer attention is O(L²). State-space models (Mamba, Gu & Dao, arXiv:2312.00752) do sequence mixing with a linear-time recurrence and hit transformer-level quality on many benchmarks. Jamba (AI21, Mar 2024, arXiv:2403.19887) interleaves Mamba and attention layers in a MoE, shipping a 256k model that fits on one 80GB GPU. RWKV-v6 (2024) is a fully linear-attention transformer. Griffin (DeepMind, arXiv:2402.19427) mixes local attention with a gated linear recurrence.

Where they stand mid-2025: pure Mamba slightly under-performs transformers on complex recall (in-context copy tasks); hybrid Jamba-style models are the pragmatic sweet spot. Long context without the O(L²) tax is real but still not universally winning.

9.5 · Llama 4's iRoPE, and the "no positional encoding on some layers" trick

Meta's Llama 4 Scout (Apr 2025) reaches 10M tokens using iRoPE — interleaved layers where some layers use RoPE and others use no positional encoding at all. The intuition: NoPE layers force the model to learn position implicitly via causal masking, which extrapolates better; RoPE layers give it precise relative position where needed. Combined with per-layer temperature scaling on the NoPE layers, they claim near-lossless retrieval to 10M on Needle-in-Haystack.

9.6 · Long-context evaluation, honestly (RULER, LongBench-v2, HELMET)

Hsieh et al.'s RULER (arXiv:2404.06654) showed a dirty truth: many models advertising 128k context degrade sharply past their effective context length (often 32k or less). Their 13-task benchmark (multi-key needles, variable tracking, aggregation) is now the reference. LongBench-v2 (2024) and HELMET (2024) add real-world tasks: long-doc QA, code repository understanding, many-shot in-context learning.

Mid-2025 leaderboard flavor:

• Advertised vs effective RULER context often differs by 2–4×.
• Claude 3.5 Sonnet (200k advertised) ~ 128k effective.
• Gemini 1.5 Pro (2M) ~ 512k–1M effective, best-in-class at true long-context reasoning.
• Llama 3.1 405B (128k) ~ 64k effective.

Run RULER on any model before you trust its context window in production.

9.7 · Further reading

• Position Interpolation: https://arxiv.org/abs/2306.15595
• YaRN: https://arxiv.org/abs/2309.00071
• LongRoPE: https://arxiv.org/abs/2402.13753
• Ring Attention: https://arxiv.org/abs/2310.01889
• StreamingLLM (attention sinks): https://arxiv.org/abs/2309.17453
• Mamba: https://arxiv.org/abs/2312.00752
• Jamba: https://arxiv.org/abs/2403.19887
• Griffin: https://arxiv.org/abs/2402.19427
• RULER: https://arxiv.org/abs/2404.06654
• Lost in the Middle: https://arxiv.org/abs/2307.03172

10 · Try it yourself

1. Take a 4k-trained model (TinyLlama-1.1B is fine).
2. Apply YaRN scaling to 32k. Measure perplexity on a book (Project Gutenberg, ~30k tokens).
3. Run RULER's niah_single_1 at 4k, 16k, 32k. Does recall stay >90%?
4. Do a 500-step fine-tune on long sequences (grab a subset of RedPajama-Books). Re-measure. Which task recovers fastest — needle retrieval or summarization?
5. Compare against StreamingLLM mode (sliding window + sinks). Which produces coherent output at 100k+ tokens?

11 · Recap

Position encoding is what limits a model's context window — not attention itself. RoPE-based extension (PI → NTK → YaRN → LongRoPE) buys you 4–32× with brief fine-tuning. Beyond that you need systems tricks (Ring Attention) or a different architecture (Mamba, Jamba, iRoPE). Streaming forever is a separate problem StreamingLLM solved by preserving attention sinks. And always: run RULER before believing the marketing number.

🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §1–§4 (RoPE → PI → NTK → YaRN) in 24 hours. Revisit alongside the RULER paper on day 7.

Next session (S079): you built a model. Gave it eyes, tools, evidence, MoE efficiency, long context. Before shipping, someone needs to try to break it. Red-team & safety eval — the boring, essential last mile.

Sticky note (keep on your desk): Positions past L_train are OOD for the model, not for the math. RoPE-scaling drags them back into the trained frequency band.

Previous: ← DL S077 · Next: DL S079 → Red-Team & Safety Eval