Long-Context Training Techniques

Extending context requires compatible position encodings, training data with useful long-range relationships, and evaluation at the intended length.

A model can accept a longer token sequence without using its contents reliably. For a model trained on 4k-token examples, answering a question about evidence 30k tokens away introduces position patterns and dependencies that training may not have covered.

Published context extensions often reuse a shorter-context checkpoint. Llama 3.1, for example, extended training to 128k through stages. This post examines three RoPE scaling approaches: Position Interpolation, NTK-aware scaling, and YaRN.

Position encoding and compute are separate constraints. Moving from 4,096 to 131,072 tokens multiplies sequence length by 32. Dense attention’s pairwise work grows by 1,024 per sequence, or 32 per token. Other operations do not all scale quadratically, so these are attention-work ratios, not whole-training cost estimates.

Prerequisites

This post assumes you've read:

• Positional encodings and RoPE for why RoPE rotates query and key vectors as a function of position.
• Training view vs. inference view for why attention cost matters and what a forward pass is.
• Supervised fine-tuning for what "fine-tune on X" concretely means.

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 $m$ and a key at position $n$ then depends on the relative offset $m-n$, not on absolute positions. That relative-position identity remains true outside the training window; learned performance need not generalize with it.

Why naive extrapolation fails

A transformer with RoPE has a per-dimension-pair rotation frequency:

where $b$ is the base (taken as $10000$ in this example) and $d$ is the head dimension. At position $m$, dim pair $i$ is rotated by angle $m \theta_i$. If you've only ever seen positions $m \in [0, 4095]$ during training, the model has learned what dot products at those angles mean for attention. The rotation function $\cos(m \theta_i)$ is smooth, so in principle you could just feed it $m = 10000$ and it would produce a value. That calculation alone does not establish that learned attention will generalize to the new offsets.

RoPE rotations preserve query and key norms. Their dot products remain bounded by those norms even at positions beyond training. What changes is the combination of relative offsets presented to learned attention heads. Familiar periodic angles in one pair do not guarantee that the full pattern supports retrieval.

Frequency scaling changes those patterns to make an extension easier to learn.

Position Interpolation: the simplest fix

Chen et al.’s Position Interpolation maps positions into a shorter interval. To extend a 4,096-token window to 32,768 tokens, use an extension factor $s=8$:

Equivalently, divide every rotation frequency by $s$:

Positions 0 through 32,767 map into the continuous interval $[0,4096)$. These fractional positions need not match any position seen during training, and the relationships among token positions have changed.

Uniform scaling also changes nearby offsets. Adjacent positions differ by $\theta_i/8$ rather than $\theta_i$ in every pair. The model must adapt to those changed patterns; preserving the old angular range does not preserve its learned behavior.

The original PI recipe rescales the base model’s RoPE frequencies and fine-tunes it on long-context data. The paper extended the tested Llama models to 32k with up to 1,000 fine-tuning steps and evaluated them on its selected tasks.

NTK-aware scaling: preserve the fast frequencies

NTK-aware scaling changes the base instead of scaling all frequencies equally. This preserves the fastest pair while stretching lower frequencies more:

When you raise the base and recompute $\theta_i = b^{-2i/d}$, the fastest pair remains unchanged: at $i = 0$, the expression does not depend on $b$. Higher-index pairs are compressed increasingly strongly.

NTK-aware scaling was proposed as a way to obtain moderate extensions without additional training in some Llama experiments. That is an empirical result for particular checkpoints and lengths, not a guarantee for any RoPE model.

YaRN: selective scaling and attention temperature

YaRN combines selective frequency scaling, an attention-temperature adjustment, and fine-tuning. Its choices differ from changing only the RoPE base.

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): retain their original frequencies to preserve these local-position patterns.
• Low-frequency pairs (much less than one rotation over 4k): apply PI-style scaling $\theta_i / s$. Their original wavelengths exceeded the training span.
• Middle pairs: a smooth ramp between the two.

Second, YaRN changes attention temperature. In the original formulation, let $m(s)=1+0.1\ln s$ for an extension factor $s\geq1$. Scaling both rotated queries and keys by $m(s)$ multiplies their dot products by $m(s)^2$. Equivalently, divide logits by temperature $t=1/m(s)^2$. This sharpens attention.

Third, YaRN includes fine-tuning. The paper reports extensions using about 0.1% of the original pretraining data in its experiments.

YaRN is one published extension method. Llama 3.1 used its own frequency scaling and progressive training recipe; describing every selective scaling rule as YaRN obscures those differences.

Progressive context extension: the staged schedule

A progressive schedule increases the maximum sequence length in stages and checks performance before advancing. Llama 3’s report describes multiple stages on the way to 128k. A useful schedule depends on the checkpoint, compute budget, and evaluation tasks.

Longer sequences increase attention work and activation memory. Training at intermediate lengths can allocate compute more selectively.

Data strategies: why the text matters

Increasing the input limit and supplying long sequences does not establish that the model learns useful long-range dependencies.

Concatenating independent short documents creates a long sequence without necessarily supplying predictions that require distant evidence. Attention masking and boundary tokens also matter.

Fu et al. study how the composition and length distribution of training data affect long-context performance. Useful data can include:

• Genuinely long documents: books, long technical papers, extensive legal filings, software repositories processed as single units. Check whether predictions actually depend on distant information; document length alone is insufficient.
• 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) can supply the intended dependency signal.
• Answer-absent examples: include examples where the relevant information isn't in the context, so the model learns when to say "not in the document."

Evaluate retrieval and reasoning separately. A model that finds one exact-match fact may still fail to combine evidence across documents or recognize that an answer is absent.

Worked example

Greg Kamradt’s needle-in-a-haystack test probes retrieval of an inserted fact from a long input. 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.

Real experiments must specify the model, prompts, haystack, lengths, repetitions, and scoring. A dip in the middle can reveal positional sensitivity, but its shape and size depend on that setup. Needle retrieval is only one probe; add multi-document reasoning and answer-absent cases.

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, s, orig_ctx_len=4096, base=10000.0,
               beta_fast=32.0, beta_slow=1.0):
    """Original YaRN-style linear ramp in dimension-pair index."""
    if s < 1 or not (0 < beta_slow < beta_fast):
        raise ValueError("Require s >= 1 and 0 < beta_slow < beta_fast")
    d = 2 * freqs.numel()
    def correction_dim(rotations):
        return d * math.log(orig_ctx_len / (2 * math.pi * rotations)) / (2 * math.log(base))
    low = max(0, math.floor(correction_dim(beta_fast)))
    high = min(d - 1, math.ceil(correction_dim(beta_slow)))
    indices = torch.arange(freqs.numel(), device=freqs.device)
    ramp = ((indices - low) / max(high - low, 1e-3)).clamp(0, 1)
    return (1 - ramp) * freqs + ramp * (freqs / s)
 
def yarn_logit_multiplier(s):
    if s < 1:
        raise ValueError("Require s >= 1")
    return (1 + 0.1 * math.log(s)) ** 2
 
 
# --- 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}")

The ramp is linear in dimension-pair index between frequency-derived boundaries. It is not a linear interpolation in the number of rotations. Use the base, original context length, and rotary dimension of the checkpoint you are extending.

The helper also gives the factor for multiplying attention logits before softmax. Some implementations instead scale RoPE’s cosine and sine tables by $m(s)$, thereby scaling both queries and keys. Apply the correction once, not both ways.

Misconceptions

"The advertised window measures usable context." A window limit states how much input the interface accepts. Evaluate performance at the lengths and tasks you need, including retrieval, evidence combination, and absent-answer cases. RULER and LongBench probe more than a single exact-match needle.

"Context extension requires retraining from scratch." Position Interpolation and YaRN demonstrate extensions of existing checkpoints. Their training costs are specific to the tested models and extension factors, not a fixed percentage for every target window.

"Attention cost grows quadratically per token." Dense pairwise attention is quadratic per sequence and linear in sequence length per token. Prefill, decoding, KV-cache storage, and model projections have different costs; see prefill vs. decode.

"YaRN fully solves long context." It supplies a frequency-scaling and temperature recipe. It does not guarantee position generalization or solve missing training dependencies, retrieval errors, or multi-hop reasoning.

What's next

The next post covers synthetic data and distillation: how teacher outputs become supervision, how to validate what transfers, and why data replacement can lose rare cases.

Additional reading (and watching)

• Chen, S., Wong, S., Chen, L., & Tian, Y. (2023). Extending Context Window of Large Language Models via Positional Interpolation. arXiv:2306.15595. The paper that introduced PI as a principled way to extend context beyond the training length with minimal fine-tuning.

• bloc97. (2023). NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation. r/LocalLLaMA. The original online write-up of the NTK-aware base-scaling trick. An informal proposal for changing the RoPE base.

• Peng, B., Quesnelle, J., Fan, H., & Shippole, E. (2023). YaRN: Efficient Context Window Extension of Large Language Models. arXiv:2309.00071. The per-dim-pair ramp plus attention temperature correction. Includes experiments extending Llama models with selective frequency scaling and attention-temperature adjustment.

• Meta AI. (2024). The Llama 3 Herd of Models. arXiv:2407.21783. Documents the staged long-context training schedule and frequency rescaling used for Llama 3.1's 128k context.

• Fu, Y., Panda, R., Niu, X., Yue, X., Hajishirzi, H., Kim, Y., & Peng, H. (2024). Data Engineering for Scaling Language Models to 128K Context. arXiv:2402.10171. Examines training data composition and length distributions for context extension.

• Kamradt, G. (2023). Needle In A Haystack — Pressure Testing LLMs. GitHub. The original NIAH evaluation methodology and code. A reproducible starting point for a single-fact retrieval probe.

• Liu, N. F., Lin, K., Hewitt, J., et al. (2023). Lost in the Middle: How Language Models Use Long Contexts. arXiv:2307.03172. Identified the U-shaped recall pattern where information in the middle of the context is recalled worse than at the ends.

• Hsieh, C.-P., Sun, S., Kriman, S., et al. (2024). RULER: What's the Real Context Size of Your Long-Context Language Models?. arXiv:2404.06654. A broader long-context benchmark that goes well beyond needle-in-a-haystack into multi-hop retrieval and aggregation tasks.

• Bai, Y., Lv, X., Zhang, J., et al. (2024). LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding. arXiv:2308.14508. Task-level long-context evaluation across summarization, QA, and code.

• Ding, Y., Zhang, L. L., Zhang, C., et al. (2024). LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens. arXiv:2402.13753. Uses evolutionary search for per-dimension rescaling factors and evaluates very long context extensions.

• Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864. The original RoPE paper, for reference on where $\theta_i = b^{-2i/d}$ comes from.

• Quesnelle, J. (2023). YaRN reference implementation. Dimension-index ramp and attention scaling used in the accompanying model implementation.