Chapter 3.2 β€” Long Context

Contents

  1. Two different problems that both go by the name "context length"
  2. Positional encodings revisited: what "not generalizing" actually means
  3. Positional interpolation: compressing position into a familiar range
  4. YaRN: treating frequencies unequally instead of uniformly
  5. Sliding-window attention in production: bounding the computation itself
  6. How these pieces fit together in a real long-context system
  7. Interview angle
  8. Self-check questions
  9. Sources

1. Two different problems that both go by the name "context length"

Chapter 3.1 established that attention is expensive at long sequence lengths and surveyed several ways to make it cheaper: sparsify the pattern, reformulate the algebra, or optimize the memory access pattern without changing the math at all. It would be easy to walk away from that chapter thinking "long context" is now a solved engineering problem β€” make attention cheap enough, and you can run a model at any context length you like. That conclusion misses a second, independent obstacle, and this chapter is dedicated entirely to it: even if attention were free, a model trained on sequences of length, say, 4,096 tokens has never seen a position index higher than 4,096 during training, and there is no guarantee that its positional encoding scheme will behave sensibly if you simply hand it position 50,000 at inference time.

This is worth stating as sharply as possible because it's a common point of confusion, including in interviews: the cost of attention (Chapter 3.1) is a compute and memory problem, solvable by better algorithms and better hardware utilization. The generalization of positional encodings (this chapter) is a statistical and representational problem β€” it's about whether the function the model learned, which was fit to a certain range of position values, continues to produce sensible outputs outside that range. You can have unlimited attention budget and still have a model that falls apart at long context, and you can have a model with perfect long-range positional generalization that is still computationally infeasible to run at that length. Production long-context systems need answers to both problems simultaneously, which is exactly why this chapter follows directly from the last one rather than replacing it.

2. Positional encodings revisited: what "not generalizing" actually means

Part II covered the mechanics of RoPE and ALiBi in depth, so this chapter won't re-derive them, but it's worth being precise about what "extrapolation failure" concretely looks like, since it motivates everything that follows. RoPE encodes position by rotating each query/key vector's sub-pairs of dimensions by an angle proportional to position, with different dimension-pairs rotating at different frequencies β€” some rotate quickly (completing many full cycles as position increases), others rotate slowly. This construction gives RoPE a very useful property within the range of positions seen during training: the dot product between a query at position $m$ and a key at position $n$ depends only on the relative distance $m - n$, not on the absolute positions, which is exactly the inductive bias you want for language, where relative order usually matters far more than absolute position in the document. Problems appear when $m$ and $n$ both grow far beyond anything seen in training. The high-frequency dimension-pairs cycle so quickly that, at very large positions, their rotation angles effectively become uninterpretable β€” the model never learned what a given angle "means" out at that range, because it never saw those angles during training, and worse, the model's learned query/key representations for a given relative distance may only have been trained to be meaningfully consistent within the trained window.

ALiBi, covered in Chapter 2.2, took a different approach entirely: rather than rotating vectors, it adds a fixed, distance-proportional penalty directly to the attention scores, penalizing attention to far-away tokens more heavily as distance grows. Because that penalty is a simple, monotonic function of distance with no learned frequency components to go out of range, ALiBi was specifically designed for, and empirically shown to have, much better extrapolation behavior β€” a model trained on short sequences with ALiBi can be evaluated at meaningfully longer sequences with a graceful degradation rather than a cliff. This is worth remembering as you read the rest of this chapter, because it reframes RoPE's extrapolation problem as somewhat self-inflicted: RoPE gets excellent relative-position behavior and strong empirical performance within its trained range in exchange for a fragility outside it, and the two techniques covered next β€” positional interpolation and YaRN β€” exist specifically because RoPE, not ALiBi, is what almost every major production model actually uses, which makes fixing RoPE's extrapolation behavior a practically urgent problem rather than a purely academic one.

3. Positional interpolation: compressing position into a familiar range

The most direct fix for RoPE's extrapolation problem, introduced by Chen et al. at Meta AI, starts from an observation about what actually goes wrong versus what doesn't. The failure mode isn't that RoPE's math breaks at large positions β€” the rotation is perfectly well-defined at any position mathematically β€” it's that the model's learned parameters were only ever exposed to, and only ever tuned to make good predictions using, angles corresponding to positions up to the training length $L$. Positional interpolation's core idea is disarmingly simple: instead of letting position indices run up to some new, larger length $L'$ that the model has never seen, rescale them down by the ratio $L/L'$ so that the largest position index the model ever actually computes with is still $L$, the value it was trained on. Concretely, if you want to extend a model trained at 4,096 tokens out to 16,384 tokens, you don't feed RoPE the raw position indices 0 through 16,383; you multiply every position index by $4096/16384 = 0.25$ before computing rotation angles, so a token at true position 16,000 gets treated, for the purposes of its rotation angle, as if it were at position 4,000 β€” comfortably inside the range the model already understands.

The tradeoff this buys is exactly what the name suggests: you're trading resolution for range. Squeezing four times as many true positions into the same range of angles means that two tokens which are, say, 40 positions apart in the compressed representation might actually be 160 positions apart in the real document β€” the model's fine-grained sense of "how far apart are these two tokens" gets coarser as the compression ratio grows. Empirically, this coarsening turns out to be something the model can adapt to reasonably well with a small amount of additional fine-tuning at the new context length, because the model is still operating entirely within the range of angles it already knows how to interpret; it just has to learn a new (or adjusted) mapping from "this compressed distance" to "this much relevance," rather than having to make sense of angles it has genuinely never encountered. The method's real appeal is that this fine-tuning is cheap β€” a comparatively small number of additional training steps at the new length β€” rather than requiring a full retraining run, which is precisely why interpolation-based context extension became a standard, practical recipe rather than a purely theoretical fix.

4. YaRN: treating frequencies unequally instead of uniformly

Positional interpolation's uniform rescaling is simple, but it treats every dimension-pair's rotation frequency exactly the same way, and that uniformity is actually a mismatch with how RoPE's frequencies behave. Recall that RoPE assigns each dimension-pair a different frequency, typically spanning a wide range from high frequencies (which complete many full rotations within even a short sequence) to low frequencies (which barely complete a fraction of a single rotation even at the model's full trained length). A high-frequency dimension-pair, having already cycled through many full rotations within the training range, has effectively already "seen" a dense sample of possible angles β€” extrapolating it further doesn't ask the model to interpret an unfamiliar angle so much as continue a pattern it has already sampled thoroughly. A low-frequency dimension-pair, by contrast, might only sweep through a small fraction of a full rotation even across the entire training range, meaning the model has only ever seen a narrow sliver of that dimension-pair's possible angles β€” pushing that dimension further out is a much riskier extrapolation, since there's comparatively little training signal about what larger angles for that dimension-pair should even encode.

This frequency-dependent reasoning first appeared as NTK-aware scaling, an interpolation scheme proposed in community write-ups before being formalized in the literature, which reasons from Neural Tangent Kernel theory about which frequencies a network can learn to extrapolate versus which it can't, and interpolates each frequency by a different amount accordingly rather than applying positional interpolation's single uniform factor. NTK-aware scaling is the direct ancestor of the more refined scheme below β€” YaRN's own name is a nod to the fact that it's one of several "NTK-aware"-style methods that followed the original idea.

YaRN (Yet another RoPE extensioN method), introduced by Peng, Quesnelle, Fan, and Shippole, takes this observation and turns it into the core design principle: rather than applying positional interpolation's single uniform rescaling factor to every frequency, YaRN interpolates high-frequency dimensions less aggressively (since they extrapolate reasonably well on their own) and interpolates low-frequency dimensions more aggressively (since they need the most help), with a smooth transition, calibrated by frequency, between these two regimes. This targeted approach avoids a specific weakness of uniform interpolation: because every dimension gets the same compression under plain positional interpolation, high-frequency dimensions β€” which didn't actually need much help β€” get needlessly compressed too, unnecessarily blurring exactly the fine-grained relative-position information that RoPE's high frequencies are good at encoding. YaRN also incorporates a complementary adjustment to the attention temperature to compensate for a subtle shift in the distribution of attention scores that interpolation introduces, which the authors found further improved quality. The practical upshot, reported in the paper, is that YaRN achieves strong context extension results with substantially less fine-tuning data and fewer training steps than uniform positional interpolation requires to reach comparable quality β€” a direct payoff of respecting the fact that RoPE's frequency bands are not interchangeable.

5. Sliding-window attention in production: bounding the computation itself

Everything so far in this chapter has been about making a model's positional understanding generalize further, while implicitly assuming Chapter 3.1 has already made attention itself affordable at that length. Production systems sometimes take a more blunt, architecturally direct approach instead: bound the attention computation so that it never needs to scale with the nominal context length at all, no matter how long the conversation or document nominally is. Mistral 7B is a concrete, widely deployed example of this. Its attention layers use a sliding window: each token attends only to a fixed-size window of the most recent tokens (a window of a few thousand tokens in Mistral's design), regardless of how much text precedes that window in the overall sequence. Within a single layer, this bounds attention's cost to the same $O(n \cdot w)$ pattern as Longformer's local component from Chapter 3.1 β€” but the motivation and framing here is slightly different: this isn't primarily positioned as an approximation strategy for making a fixed-context model cheaper, it's the model's native design for handling sequences that are notionally much longer than any single layer's direct receptive field.

The key idea that makes this workable rather than crippling is that receptive field compounds across layers, the same way it does in a convolutional network with a fixed kernel size. A token in layer 1 can only see $w$ tokens directly behind it, but a token in layer 2 is built from layer-1 representations that themselves each saw $w$ tokens behind them, so a layer-2 token's effective receptive field is roughly $2w$; after $k$ layers of sliding-window attention, information can in principle have propagated across roughly $k \times w$ tokens of the original sequence, even though no single layer ever directly attended beyond its own window. Mistral's design leans on having enough layers that this compounded receptive field comfortably covers its intended operating context length. The tradeoff is genuine, though, and worth being honest about: this is indirect information flow, mediated through several transformations, rather than the direct token-to-token comparison that full attention provides, and there's no guarantee that a specific piece of information many layers'-worth of distance away survives that indirect path as cleanly as a direct attention link would have preserved it. In practice, sliding-window attention plus enough depth has proven to be a workable, efficient recipe for a range of production use cases, but it is a real design compromise, not a free lunch β€” it's the same fundamental tradeoff Longformer made in Chapter 3.1, adopted here as a primary architectural choice rather than as a retrofit onto an existing model.

6. How these pieces fit together in a real long-context system

It's worth stepping back and being explicit about how the techniques in this chapter compose with each other and with Chapter 3.1, because a real production long-context system typically isn't choosing just one of these ideas β€” it's combining them. A model might use RoPE, extended via YaRN-style frequency-aware interpolation to handle position indices well beyond its original training range, computed efficiently via FlashAttention so that the raw compute and memory cost of full attention over that extended length stays tractable, and in some deployments further bounded by a sliding window so that even FlashAttention's linear-in-memory cost doesn't grow without bound as conversations get arbitrarily long. None of these techniques substitutes for another: fixing positional generalization does nothing to reduce the $O(n^2)$ compute of full attention over a long sequence, and making attention cheap does nothing to fix a RoPE-based model's confusion about position indices it has never encountered. Treating "long context" as a single problem with a single fix is a mistake that shows up often in casual discussion of this area, and being able to cleanly separate "this is a positional-generalization issue" from "this is a compute/memory issue" β€” and to name the specific technique that addresses each β€” is one of the more reliable signals of genuine depth on this topic in a technical interview.

It's worth naming one more limitation explicitly, since it's neither a positional-generalization issue nor a compute/memory issue, and none of the techniques above touch it: even a model that handles a given context length correctly in the "can it attend there at all" sense documented above often uses information in the middle of a long context markedly worse than information near the beginning or end. Liu et al. documented this directly, finding a consistent U-shaped performance curve on tasks requiring retrieval from a specific position in a long document β€” accuracy is highest when the relevant information sits at the very start or very end of the context, and drops substantially for the same information placed in the middle, in models that handle the raw context length without any architectural failure. This is a genuinely distinct problem from everything else in this chapter: positional interpolation and YaRN are about whether the model can represent and attend to a position at all, sliding-window attention and FlashAttention are about whether attending there is affordable, and "lost in the middle" is about whether the model actually makes good use of information it can technically see and afford to look at β€” a training-and-behavior question more than an architectural one, and one that a system relying purely on making context longer, without addressing retrieval quality within that context, will still run into.

With attention now both affordable (Chapter 3.1) and extendable to positions beyond training (this chapter), the next chapter turns to an entirely different axis of scaling: not making a fixed amount of computation cheaper or more general, but making the model itself larger β€” specifically, growing its parameter count dramatically β€” without paying a proportional cost in compute per token. That's the province of Mixture of Experts, which replaces the dense feedforward layer that every token has passed through unconditionally so far with a much larger bank of specialized feedforward networks, only a handful of which any given token actually uses.

7. Interview angle

"A team wants to extend their model's context window from 8K to 32K tokens. What are the distinct problems they need to solve, and are they the same problem?" A strong answer separates them explicitly: they need attention to remain computationally tractable at 32K tokens (Chapter 3.1's problem β€” addressed via FlashAttention for exact computation, or sparse/linear attention if even that isn't sufficient), and they separately need their positional encoding scheme to behave sensibly at position indices it never saw during training (this chapter's problem β€” addressed via positional interpolation or YaRN if they use RoPE, largely already addressed if they use ALiBi). A strong candidate notes these require different techniques and solving one doesn't solve the other.

"Why does RoPE need interpolation-based fixes for long context but ALiBi generally doesn't?" A strong answer contrasts the mechanisms: RoPE encodes position via frequency-based rotations that were only ever trained on angles up to some maximum position, so novel, very large positions produce angles the model never learned to interpret. ALiBi instead adds a simple, monotonic, distance-proportional penalty to attention scores with no learned frequency components, so there's no analogous "unfamiliar angle" failure mode, giving it much better native extrapolation β€” the price RoPE pays for its excellent relative-position behavior and empirical performance within-range is more fragility out-of-range.

"Explain why YaRN treats high- and low-frequency RoPE dimensions differently, rather than applying uniform positional interpolation." A strong answer explains that RoPE's dimension-pairs span a range of rotation frequencies, and high-frequency pairs complete many rotations even within the original training length, so they've effectively sampled a dense range of angles already and extrapolate reasonably on their own, whereas low-frequency pairs sweep through only a small fraction of a rotation even at full training length, so extending them is a much riskier extrapolation needing more aggressive interpolation. Applying one uniform compression factor to all frequencies, as in plain positional interpolation, unnecessarily blurs the fine relative-position information carried by the high-frequency dimensions that didn't need the help.

"What's the tradeoff of sliding-window attention as used in Mistral, and how does information from outside the window ever reach a given token?" A strong answer explains that each layer's direct attention is bounded to a fixed window, but receptive field compounds across depth the way it does in a CNN β€” a token's effective information horizon after $k$ layers is roughly $k$ times the window size β€” so information from further away can still reach a token, but only via several indirect transformations rather than one direct attention comparison, which is a real (if often acceptable) fidelity tradeoff compared to full attention.

"If compute and memory for attention were literally free, would long-context problems be solved?" A strong answer says no, and explains why: positional encoding generalization is a separate, statistical problem about whether the model's learned representations are meaningful at position values outside its training distribution, and no amount of cheaper compute changes what the model actually learned. This is exactly the distinction the chapter is built around, and naming it directly is what signals real understanding rather than a surface-level "long context = expensive attention" association.

8. Self-check questions

  1. Explain, in your own words, why "attention is expensive at long context" and "positional encodings don't generalize to long context" are two separate problems, using RoPE as the concrete example for the second.
  2. Why does RoPE guarantee that the dot product between a query and key depends only on their relative distance, and why does this property break down, in practice, at position indices far outside the training range?
  3. Walk through the mechanics of positional interpolation: what specific quantity gets rescaled, by what factor, and why does this keep the model operating within a range of angles it already understands?
  4. What is the resolution-versus-range tradeoff introduced by positional interpolation, and why is a small amount of fine-tuning typically enough to adapt a model to it?
  5. Explain why YaRN interpolates different RoPE frequency bands by different amounts, and why this is expected to require less fine-tuning data than uniform interpolation to reach the same quality.
  6. In Mistral's sliding-window attention, how does a token's effective receptive field grow with network depth, and what specific fidelity tradeoff does this indirect information flow introduce compared to full attention?
  7. If a production system uses RoPE with YaRN-style extension, FlashAttention, and a sliding window all together, explain what problem each of the three components is separately solving.

9. Sources

  • Press, O., Smith, N. A., & Lewis, M. (2021). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (ALiBi). ICLR 2022. arXiv:2108.12409. https://arxiv.org/abs/2108.12409
  • Chen, S., Wong, S., Chen, L., & Tian, Y. (2023, Meta AI). Extending Context Window of Large Language Models via Positional Interpolation. arXiv:2306.15595. https://arxiv.org/abs/2306.15595
  • Peng, B., Quesnelle, J., Fan, H., & Shippole, E. (2023). YaRN: Efficient Context Window Extension of Large Language Models. ICLR 2024. arXiv:2309.00071. https://arxiv.org/abs/2309.00071
  • Jiang, A. Q., et al. (2023, Mistral AI). Mistral 7B. arXiv:2310.06825. https://arxiv.org/abs/2310.06825
  • Liu, N. F., Lin, K., Hewitt, J., Paranjape, A., Bevilacqua, M., Petroni, F., & Liang, P. (2023). Lost in the Middle: How Language Models Use Long Contexts. TACL 2024. arXiv:2307.03172. https://arxiv.org/abs/2307.03172

Self-check quiz

Test your understanding of this chapter with a short quiz β€” a mix of concept questions and quick calculations.

Why are "attention is expensive at long context" and "positional encodings don't generalize to long context" two separate problems?

Explanation: You can have unlimited attention budget and still have a model that falls apart at long context, or a model with perfect positional generalization that's still infeasible to run β€” production systems need answers to both.

Why does RoPE need interpolation-based fixes for long context, while ALiBi generally doesn't?

Explanation: RoPE trades excellent within-range relative-position behavior for fragility out-of-range; ALiBi's unbounded, monotonic penalty has no analogous "unfamiliar angle" failure mode.

What is the core idea behind positional interpolation?

Explanation: Positions are rescaled by L/L' before computing RoPE angles, so the model always operates on angles it already knows how to interpret β€” trading positional resolution for extended range.

Why does YaRN interpolate high-frequency and low-frequency RoPE dimensions differently, rather than uniformly?

Explanation: Uniform interpolation needlessly compresses high-frequency dimensions that didn't need help, blurring fine relative-position information β€” YaRN calibrates the interpolation strength by frequency instead.

A model trained with context length $L = 4096$ is extended to $L' = 16384$ using positional interpolation. What scaling factor $L/L'$ is applied to every position index before computing RoPE angles?

Explanation: L/L' = 4096/16384 = 0.25.

Using that same scaling factor of 0.25, a token at true position 16,000 gets treated, for its rotation angle, as if it were at what position?

Explanation: 16,000 Γ— 0.25 = 4,000 β€” comfortably inside the range the model already understands.

Mistral-style sliding-window attention uses a window of $w = 4096$ tokens per layer. After $k = 8$ layers, roughly how large is a token's effective receptive field ($k \times w$)?

Explanation: 8 Γ— 4096 = 32,768 β€” receptive field compounds across layers the way it does in a CNN with a fixed kernel size.

A sliding-window model uses a window of $w = 2000$ tokens and has $k = 12$ layers. What is the approximate effective receptive field, in tokens?

Explanation: 12 Γ— 2000 = 24,000 tokens, propagated indirectly through k layers of transformations.