Self-attention makes every token look at every other token. For n tokens that's an n×n matrix of scores, so compute and memory both grow with n² — double the context and you quadruple the cost. That single fact is why a 200k-token window is hard, and the entire long-context playbook is a set of tricks to dodge that wall. I built an interactive page where you drag sliders and watch each trick bite, all simulated client-side on the attention patterns with no model calls. Here's the tour.
The quadratic wall
Full attention scores every query token against every past key token. Context 4k is 16M scores; 128k is 16.8 billion. It simply doesn't scale, and everything below is a way around it.
scores = Q @ K.T # shape (n, n) <- the n×n blowup
scores = scores.masked_fill(~causal, -inf) # no peeking at the future
attn = softmax(scores, dim=-1) @ V # cost & memory ~ O(n²)
Sliding-window (local) attention
The cheapest fix: let each token attend only to the last w keys — a diagonal band instead of a full triangle. Cost drops from O(n²) to O(n·w), which is linear in context. "But then it can't see far!" — it can, through depth: each layer shifts the window a little, so L stacked window layers give an effective reach of about L·w. Mistral and Longformer use exactly this.
def sliding_window_mask(n, w):
mask = zeros(n, n, dtype=bool)
for i in range(n):
lo = max(0, i - w + 1)
mask[i, lo : i + 1] = True # causal AND within the window
return mask # ~ n*w True cells, not n²
Attention sinks — make streaming stable
A pure rolling window seems perfect for endless chat: just drop old tokens. It fails badly. Once the window slides past the very first tokens, the model destabilises and perplexity explodes. Why? Softmax must sum to 1, so it dumps its leftover probability onto the earliest tokens — they quietly become an attention sink. StreamingLLM's fix is almost free: always keep the first ~4 tokens visible to everyone. With those pinned you can stream millions of tokens, stably.
def streaming_mask(n, w, sink=4):
mask = sliding_window_mask(n, w)
mask[:, :sink] = True # first `sink` keys stay visible to everyone
return mask
# evict them -> scores blow up. keep ~4 -> stream forever, stably.
Watching this in the demo is the part that stuck with me: flip the sinks off, drag the "tokens streamed" slider past the window size, and the perplexity curve rockets from a flat ~8 to 260 the instant the sinks get evicted.
Sparse / dilated attention
Local windows are blind to distant-but-important tokens. Sparse attention keeps the cheap local band and adds a few extra links: global tokens everyone can see, plus strided (dilated) hops that skip across the sequence. Information can now travel the whole context while the pattern stays ~O(n·w). This is the Longformer / BigBird recipe — local + global + strided.
def sparse_mask(n, w, stride=64, globals=(0,1,2,3)):
mask = sliding_window_mask(n, w)
for i in range(n):
for j in range(0, i + 1, stride): # dilated long-range hops
mask[i, j] = True
mask[:, list(globals)] = True # a few always-visible tokens
return mask
RoPE scaling — stretch past the trained length
A model trained at 4k tokens breaks if you feed it position 30,000 — it never saw those rotary angles. The fix is don't extrapolate, interpolate. Rotary embeddings (RoPE) encode position as a rotation angle; squeeze the new positions back into the range the model already knows and the angles stay familiar. NTK-aware scaling and YaRN refine this, and a little fine-tuning recovers quality. This is how a 4k model becomes a 32k+ model.
scale = train_len / target_len # e.g. 4096 / 32768 = 1/8
pos = arange(target_len) * scale # positions now live in [0, 4096)
q, k = apply_rope(q, k, pos) # rotary embeddings on scaled positions
Lost in the middle — the catch none of these fix
Here's the stubborn one. Even when a fact fits in the context, models retrieve it far better when it sits near the start (primacy) or the end (recency). Bury the needle in the middle and accuracy sags into a U-shape — and it sags deeper as the context grows. It's a genuine limitation, not a bug you can patch away. A long window is a capacity, not a guarantee that every token gets read.
# accuracy is U-shaped in position -> exploit primacy & recency:
ordered = [most_relevant] + middling_chunks + [second_most_relevant]
# also: smaller top-k (less to get lost in) + re-rank the needle to an edge.
In the demo you drag a needle across the context and the retrieval curve dips into that U — a concrete reminder that "just use a longer context" is not a strategy on its own.
These tricks compose. Real long-context models mix local windows for cheap detail, a few global/sink tokens for stability and long-range glue, position interpolation to stretch the trained length, and hardware kernels like FlashAttention so the full n×n matrix is never even materialised. What they can't fully fix is lost-in-the-middle — which is exactly why retrieval systems still place the most important evidence at the edges of the prompt.
Drag the sliders and watch each trick bite:
https://dev48v.infy.uk/ai/days/day37-long-context-tricks.html
Top comments (0)