Skip to content
Arc 7: Training & Post-Training

Long-Context Training Techniques

18 min readbuilds on Positional Encodings: Sinusoidal, Learned, RoPE, and ALiBi, Training View vs. Inference View, Supervised Fine-Tuning, PagedAttention: Virtual Memory for the KV Cache and Prefill vs. Decode: The Two Phases of Inference

Extending a short-context model past its training length is a RoPE-frequency problem. Rescale the rotations, fine-tune briefly, and the same weights work on a much longer window.

Here's the thing I kept bumping into when I first started reading long-context papers. If you trained a model on 4k-token sequences and then at inference time feed it a 32k-token document, it doesn't politely degrade. It falls over. The attention scores spike in weird places, the model's reasoning devolves into word salad, and retrieval accuracy at far positions drops to essentially zero. "Just make it longer" is not a thing you can do.

By 2026 the advertised context windows are enormous: Gemini 1.5 Pro's 2M-token window, Claude's 200k standard (1M in beta), Llama 3.1's 128k. These numbers didn't come from training runs with 128k sequences from scratch. That would be absurdly expensive, because attention cost scales quadratically with sequence length. They came from a specific recipe: train the base model short, then extend its context with a set of tricks that rescale position encodings, followed by a surprisingly small amount of fine-tuning on actually-long data.

This post is about that recipe. What breaks when you try the naive thing, why RoPE is the reason all of this is tractable, and what the three main extension schemes (Position Interpolation, NTK-aware scaling, YaRN) do to the rotation frequencies. I want you to walk away knowing what knobs exist, what they do mechanically, and which one you'd reach for in practice.

There are two problems stacked on top of each other. Position encoding: the model's notion of "where is this token" breaks when it sees positions beyond its training range. And attention compute: cost scales with the square of the sequence length, so training at 128k is roughly 1024× more expensive per token than training at 4k. This post is mostly about the position-encoding problem, because that's where the field has a clean recipe. The compute problem is why the recipe is staged: you want to spend as few tokens as possible at the expensive lengths.

Prerequisites

This post assumes you've read:

If RoPE is hazy for you, the one-line version: RoPE encodes position by rotating the query and key vectors in 2D pairs, where each pair rotates at its own frequency. The dot product between a query at position mm and a key at position nn then depends on the relative offset m−nm-n, not on absolute positions. That's the property we're about to break and carefully repair.

Why naive extrapolation fails

Let's start with the obvious thing. A transformer with RoPE has a per-dimension-pair rotation frequency:

θi=b−2i/d\theta_i = b^{-2i/d}

where bb is the base (usually 1000010000) and dd is the head dimension. At position mm, dim pair ii is rotated by angle mθim \theta_i. If you've only ever seen positions m∈[0,4096]m \in [0, 4096] during training, the model has learned what dot products at those angles mean for attention. The rotation function cos⁡(mθi)\cos(m \theta_i) is smooth, so in principle you could just feed it m=10000m = 10000 and it would produce a value. It would not, however, be a value the model is calibrated to.

Position extrapolation
m = 0  |  train_len = 4,096
rotation anglecos(θ · m) for a dim pairsample q·k scorek fixed at pos 0, q at pos min-distributionout-of-distribution01k4k16kposition m (log scale)
Top: the cosine of the rotation angle for a single dim pair, sweeping m from 0 to 16k. Inside the blue band (m < train_len), the model has seen every angle many times. Past the red line, the dot-product geometry is off-distribution. Bottom: a sample q·k score as m grows. The score drifts into values the model has no training signal for.

A few things to notice. The angle itself doesn't "break" when mm crosses 4k. It's just cos⁡(mθi)\cos(m\theta_i), which is periodic and well-defined forever. But the joint distribution of rotation angles across all dim pairs at a given position is what the trained attention heads have learned to read. Out-of-distribution positions produce rotation patterns the heads have never seen, and the score you get is functionally random.

There's a more concrete way to see the failure. At position m=100,000m = 100{,}000, the fastest-rotating dim pair has gone through thousands of full revolutions. The slowest-rotating dim pair has barely moved. The "fingerprint" of a position depends on all the dim pairs together, and once you're far enough past the training length, the fingerprint space gets weird. Nearby positions become harder to distinguish, far positions produce extreme query-key dot products, and the softmax does what softmax does when you feed it weird scores.

So. We can't just extrapolate. We need to do something to the rotation frequencies themselves.

Position Interpolation: the simplest fix

The first thing that worked, from Chen et al., is almost embarrassingly simple. Instead of making the model handle larger mm, make the positions look smaller. If the original model handled m∈[0,4096]m \in [0, 4096], and you want to handle m∈[0,32,768]m \in [0, 32{,}768], divide every position by s=8s = 8 before computing the rotation:

m↦msm \mapsto \frac{m}{s}

Equivalently, divide every rotation frequency by ss:

θi↦θis\theta_i \mapsto \frac{\theta_i}{s}

Every position mm in the new 32k window now lives at a rotation angle that falls inside the original 4k range. The model sees rotations it recognizes.

The catch is that this is not free. The model was trained to distinguish position mm from position m+1m+1 using the rotation of dim pair 0 (the fastest pair). After dividing all frequencies by 8, position mm and position m+1m+1 are now at rotation angles that are 8×8\times closer together than anything the model saw during training. Fine positional distinctions get muddled. You've solved "far positions are garbage" by making "near positions slightly less crisp" everywhere.

In practice, PI works, but it requires fine-tuning. You take the base model, rescale all the RoPE frequencies, and then train on long-context data for some number of steps until the attention heads readjust. The original PI paper showed that roughly a thousand fine-tuning steps on 32k sequences are enough to recover full quality at the new length.

NTK-aware scaling: preserve the fast frequencies

People on the LocalLLaMA subreddit noticed that PI's uniform rescaling is wasteful. The fastest-rotating dim pair is the one doing the fine-grained local position work, and we don't want to mess with it. The slowest-rotating dim pair is the one that needs to be compressed to cover the new longer range. We want non-uniform rescaling.

The fix, dubbed NTK-aware because of some arguments from neural tangent kernel theory, is to scale the base bb instead of dividing each frequency:

b↦b⋅sd/(d−2)b \mapsto b \cdot s^{d/(d-2)}

When you raise the base and recompute θi=b−2i/d\theta_i = b^{-2i/d}, something nice happens. The fast frequencies (small ii) are barely affected, because b−2i/db^{-2i/d} with i=0i = 0 doesn't depend on bb at all. The slow frequencies (large ii) get compressed a lot, which is what you want for covering the extended range. High-resolution position information is preserved; the "coarse" positional information does the stretching.

Crucially, NTK-aware scaling often works without fine-tuning at all for moderate extensions. You can take a Llama-style model and run it at 2x or 4x the training length by just swapping the RoPE base, and it keeps working. That was a surprise to me when I first saw it.

YaRN: the one that actually shipped

NTK-aware is smart, but it still has issues. The fast frequencies get slightly perturbed, and for very long extensions (8× or more) the accumulated drift hurts. Peng et al.'s YaRN ("Yet another RoPE extensioN") cleans this up with three ideas stacked together.

First, YaRN makes the scaling per-dim-pair rather than global. Each pair gets classified into one of three regimes based on how many full rotations it completes over the training sequence length:

  • High-frequency pairs (many rotations over 4k): leave them alone. The model relies on these for local position, don't touch.
  • Low-frequency pairs (much less than one rotation over 4k): apply PI-style scaling θi/s\theta_i / s. These were the ones underutilized at 4k anyway.
  • Middle pairs: a smooth ramp between the two.

Second, YaRN adds a small attention temperature adjustment. When you extend context from L0L_0 to L0⋅sL_0 \cdot s, the distribution of attention logits gets a bit sharper, which is bad for the softmax at long range. Dividing the attention scores by 1+0.1ln⁡s\sqrt{1 + 0.1 \ln s} (an empirically derived factor) brings the temperature back to something the model is used to.

Third, YaRN still involves fine-tuning, but a tiny amount, on the order of 0.1% of the original pretraining tokens for a solid result. The combination of per-dim-pair scaling and temperature correction means the model starts much closer to "correct" than PI did, so there's less work for fine-tuning to do.

RoPE frequencies under four extension schemes (s = 8)
no scaling
θno scalingi=0i=63dim pairfastslowPI (÷s)i=0i=63dim pairfastslowNTK-awarei=0i=63dim pairfastslowYaRNi=0i=63dim pairfastslowno scalingfrequencies unchanged, positions > 4k fail
Four panels showing how each extension scheme reshapes RoPE's frequency curve. Grey dashed line is the original base. Color is the rescaled frequency. Watch how PI compresses everything uniformly while NTK-aware and YaRN leave the fast (left) pairs mostly alone.

The panel comparison is worth staring at. Position Interpolation shifts the entire frequency curve down. Fast pairs slowed, slow pairs slowed. NTK-aware bends the curve so the right side (slow pairs) drops much more than the left. YaRN takes this further: the leftmost pairs are visually indistinguishable from the base, and the scaling only kicks in once the pairs are slow enough to not be doing local-position work. That selectivity is why YaRN extends further with less fine-tuning.

By 2026, YaRN (or variants of it) is the default starting point for extending open-weight models. Llama 3.1's 128k context used a YaRN-style ramp plus a staged training schedule.

Progressive context extension: the staged schedule

You might imagine the training recipe is: train at 4k, then fine-tune at 128k. In practice that jumps too far. The attention patterns at 4k and at 128k are different enough that going straight from one to the other wastes compute. A typical schedule has multiple stages.

Progressive context extension
0B tokens seen  |  context: 4k
contextlength2000B tok20B tok2B tok1. Pretrainbase training at short context2. Mid-trainapply YaRN rescale + fine-tune at 32k3. Long-context tunecurated data at 128kwall-clock / compute →pretrain dwarfs everything else; long-context tuning is a tiny sliver
Three stages across the training timeline. Pretrain runs for the full token budget at 4k. Mid-train applies YaRN rescaling and fine-tunes briefly on 32k data. Long-context fine-tune is the smallest slice, curated to be actually long and dependency-rich.

The specific stage lengths vary by model, but the shape is consistent:

  1. Pretrain at 4k for the bulk of the token budget (typically trillions of tokens). This is the expensive part. Attention is cheap here because n2n^2 at n=4096n=4096 is manageable.
  2. Mid-train at 32k after applying YaRN (or PI) rescaling. Fine-tune for somewhere between a few billion and tens of billions of tokens. The model adjusts to the new RoPE frequencies; attention heads that previously didn't need to reach past a few thousand tokens learn to use the longer range.
  3. Long-context fine-tune at 128k+ with curated data. This is the smallest slice, often under a billion tokens, and it's where you make sure the model actually uses the entire window rather than just "technically accepting" long sequences.

Why stages rather than one big jump? Two reasons. First, training at 128k is quadratically more expensive per token than training at 4k, so you want to do as few steps at 128k as possible. Second, gradient dynamics at 128k are harsher. The attention patterns have so much more room to be wrong that cold-starting an optimizer there can diverge. Warming up through 32k first gives the model an attention pattern that's closer to useful before you ramp to the full window.

Data strategies: why the text matters

Here's something a lot of early long-context attempts got wrong. You can take a base model, apply YaRN, fine-tune on 128k of random text, and end up with a model that technically accepts 128k-token inputs but gets worse at using them than a 32k model would.

The failure mode is this. If your "long" data is just short documents concatenated together (say, ten 12k-word news articles mashed end-to-end to fill a 128k window), the model learns a bad lesson. Each article is self-contained. The correct answer to "what's the token after X" very rarely depends on anything more than a few thousand tokens back, because the articles don't refer to each other. The model's attention heads learn to ignore the distant context, because distant context is demonstrably useless.

What the attention heads learn
same cursor, two data regimes
trained on concat-short4 short docs, each self-containedlearned attention reacheffective: 3 blocks backpolicy never needed distant contexttrained on genuinely longone long doc, back-references reach farlearned attention reacheffective: 3 blocks backpolicy uses the whole window
Same window length, same cursor position, two very different learned attention policies. Concat-short training data teaches the heads to throw away everything past the nearest document boundary. A genuinely long document lets them keep far positions useful, because the training signal rewarded reaching back.

Fu et al. made this point crisply. The quality of long-context fine-tuning data correlates much more strongly with downstream needle-in-a-haystack performance than the quantity. What you want is:

  • Genuinely long documents: books, long technical papers, extensive legal filings, software repositories processed as single units. The longer the document, the more often the correct prediction depends on something far back.
  • Synthetic long-range dependencies: curated data where the answer to something at position 100k literally requires information from position 5k. These are harder to get at scale, but high-quality synthetic examples (retrieval-augmented question answering, multi-document reasoning) help enormously.
  • Shuffled negatives: include examples where the relevant information isn't in the context, so the model learns when to say "not in the document."

Without long-range-dependency data, the model learns a local attention policy masquerading as a long-context one. Benchmarks that measure "model accepts 128k tokens without crashing" pass. Benchmarks that measure "model can find a fact 90% of the way through 128k tokens" fail catastrophically.

Worked example

The standard way to probe whether a model is actually reading its context is Greg Kamradt's needle-in-a-haystack test. You take a random fact (a "needle"), insert it at various depths inside a long document of filler text ("haystack"), and ask the model questions that require the needle. You sweep over both context length and needle depth, and you plot the result as a 2D heatmap. Green cells mean the model found the needle. Red cells mean it didn't.

The before-and-after for YaRN fine-tuning looks like this:

Needle-in-a-haystack recall
base model (trained at 4k, extrapolated to 128k)
document depth (% of context)5%15%25%35%45%55%65%75%85%95%1k2k4k8k16k32k64k128kcontext length (tokens, k = 1,000)8888887944171717676767713610101051515163284444141416025222363636582322236363657222224141415520222515151541922267676752172228888885116222recall %0100base training boundaryYaRN tune boundary
Heatmap of needle recall across context lengths and needle depths. The base model (pre-YaRN) fails sharply past its 4k training boundary; inside 4k it shows the classic 'lost in the middle' dip. After YaRN rescaling plus fine-tuning on 32k data, the model holds up nearly out to 128k.

A couple of things to track as the heatmap morphs. The base model's failure mode isn't just "everything past 4k goes red." Inside 4k, needles in the middle of the document are recalled worse than needles at the very start or end. That's the "lost in the middle" effect, and it shows up in both base and extended models to some degree. Once you go past 4k, though, recall collapses almost uniformly.

After YaRN and 32k fine-tuning, recall recovers across the 32k window and holds up reasonably out to 128k. The model was only trained on 32k sequences in the fine-tune. There's no 128k-specific training data. But because YaRN's frequency ramp was designed for an 8× extension past the training length, inference at 128k stays in the scheme's intended range.

What this heatmap doesn't show: YaRN doesn't fix every long-context failure. The model still makes mistakes on retrieval that requires combining multiple far-apart pieces of information, and it can still hallucinate when the answer isn't in the context at all. NIAH is a blunt test. Recall is necessary, but nowhere near sufficient.

Implementation sketch

The rescaling itself is a few lines of code. You don't need to rewrite attention; you just change the frequencies used to build the RoPE rotation table.

import math
import torch
 
def base_rope_freqs(d_head: int, base: float = 10000.0) -> torch.Tensor:
    """
    Original RoPE frequencies: theta_i = base^(-2i/d) for dim pairs i = 0..d/2-1.
    Shape: (d_head / 2,)
    """
    i = torch.arange(d_head // 2, dtype=torch.float32)
    return base ** (-2 * i / d_head)
 
def pi_scale(freqs: torch.Tensor, s: float) -> torch.Tensor:
    """
    Position Interpolation: divide every frequency by s.
    Equivalent to mapping m -> m/s before rotating.
    """
    return freqs / s
 
def yarn_scale(
    freqs: torch.Tensor,
    s: float,
    orig_ctx_len: int = 4096,
    alpha: float = 1.0,
    beta: float = 32.0,
) -> torch.Tensor:
    """
    YaRN per-dim-pair ramp:
      - high-freq pairs (many wavelengths inside orig_ctx_len) stay at base.
      - low-freq pairs (<<1 wavelength inside orig_ctx_len) get PI-scaled.
      - middle pairs get a smooth linear blend.
    """
    wavelens = 2 * math.pi / freqs              # wavelength per dim pair
    ratio = orig_ctx_len / wavelens             # wavelens per orig ctx
    # ramp: 1 = keep base, 0 = full PI scale
    keep = torch.clamp((ratio - alpha) / (beta - alpha), min=0.0, max=1.0)
    return keep * freqs + (1 - keep) * (freqs / s)
 
# --- Take a 4k-trained model and rescale to 32k (s = 8) ---
base = base_rope_freqs(d_head=128)
pi = pi_scale(base, s=8.0)
yarn = yarn_scale(base, s=8.0)
 
# Print the first few and last few pairs so the pattern is visible.
for i in [0, 1, 16, 32, 48, 62, 63]:
    print(f"i={i:3d}  base={base[i]:.6f}  PI={pi[i]:.6f}  YaRN={yarn[i]:.6f}")
# i=  0  base=1.000000  PI=0.125000  YaRN=1.000000  (fast pair: YaRN untouched)
# i=  1  base=0.865964  PI=0.108246  YaRN=0.865964
# i= 16  base=0.100000  PI=0.012500  YaRN=0.100000
# i= 32  base=0.010000  PI=0.001250  YaRN=0.002808  (middle pair: YaRN blending in)
# i= 48  base=0.001000  PI=0.000125  YaRN=0.000125
# i= 62  base=0.000133  PI=0.000017  YaRN=0.000017  (slow pair: full PI scale)
# i= 63  base=0.000115  PI=0.000014  YaRN=0.000014

The yarn_scale function is the only one that matters in practice. PI is in there for comparison. The actual mechanics of attention don't change; you're just building a different table of rotation angles. The entire extension lives in those 30 lines.

The one subtlety I skipped: YaRN also adjusts the attention temperature by a factor that depends on ss. In real implementations that shows up as a multiplier on the attention logits before softmax, not as a change to the RoPE table. I left it out here to keep the code focused on the frequency math.

Misconceptions

"Longer advertised context equals longer effective context." Almost never. A model advertised at 128k might have effective usable context closer to 32k depending on how it was trained and what you're trying to do. Retrieval tasks like needle-in-a-haystack are the most forgiving, because finding a single fact only requires one good attention lookup. Reasoning tasks that need to combine multiple pieces of context fall apart much earlier. The RULER benchmark and LongBench exist because advertised-vs-effective is a real gap, and you should never trust a context-length claim without seeing actual task numbers.

"You need to retrain from scratch to get long context." You don't. The whole point of YaRN and its cousins is that you keep the weights from the short-context pretraining and add an extension on top. The fine-tuning budget for 128k is often well under 1% of the pretraining budget. This is what makes 128k and 1M context windows economically possible.

"Attention's quadratic cost is the real bottleneck." For training, yes, which is exactly why progressive extension matters. Training at 128k costs 1024× more attention FLOPs per token than training at 4k, so you'd never want to do it for very long. At inference time, though, the bottleneck is usually KV cache bandwidth and memory, not attention FLOPs. A 128k-token KV cache for a 70B model is tens of gigabytes of data that has to be read every decode step. Long-context inference is a different engineering problem than long-context training, and the PagedAttention and prefill-vs-decode posts are where that story lives.

"YaRN fully solves long context." It fully solves the position-encoding part of long context. It doesn't solve training data quality, the "lost in the middle" attention behavior, or the fact that models still struggle with multi-hop reasoning over long contexts. YaRN unlocks the door. What the model does once it's through the door depends on a lot of other things.

What's next

The next post covers synthetic data and distillation, where we look at using a stronger model to manufacture training data for a weaker one, when that works surprisingly well, when it collapses, and how it connects to the long-context data problem we just ran into.

Additional reading (and watching)