Skip to main content

Command Palette

Search for a command to run...

RAG as an Attention Optimizer

Why you can't just dump the whole company wiki into a 2-million-token context window.

Updated
6 min readView as Markdown
RAG as an Attention Optimizer

Why are we building this pipeline at all?

You already know what RAG is. Take the user's query, hit a vector database, grab the relevant chunks, stuff them into the prompt, ask the LLM to answer. We've all seen the tutorials.

But as engineers we should be asking why. Why build this whole pipeline? Why can't we take the entire company wiki, the codebase and the Slack history, and drop all of it into a 2-million-token context window?

It isn't just about cost. It's about a fundamental architectural bottleneck inside the model itself. It's about the math of attention.

(Yes, attention. The same thing humans have been running short of since TikTok.)

Self-attention: the engine

LLMs don't read left to right the way we do. To understand context, they use self-attention.

Read this sentence: "The bank of the river." Now this one: "The bank on the corner."

How does the model know the first bank is dirt and water, and the second is money? It doesn't look at "bank" in isolation. It looks at the words around it.

When an LLM processes a prompt, it treats every token like a node in a network and asks one question: "How much does this token matter to every other token?" It draws an invisible web of connections across the whole input to work out what each word means in context.

The bottleneck: the handshake problem

That web of connections is a massive math problem. I call it the handshake problem.

If 5 people are in a room and everyone shakes hands with everyone else, that's 10 handshakes. With 1,000 people, it's 499,500 — roughly half a million.

Drag the slider and watch the connections explode:

https://codepen.io/Devanshu-Mishra-the-scripter/pen/bNqEYLE

People / tokens (N) Handshakes: N(N−1)/2 Attention scores: N²
5 10 25
100 4,950 10,000
1,000 499,500 1,000,000
100,000 ~5 billion 10 billion

Notice the third column. A handshake is symmetric — you and I shake once and we're done. Attention isn't. Every token scores every token (including itself), and it does so in both directions: how much bank cares about river is a different number from how much river cares about bank. So the count isn't N(N−1)/2. It's N².

1,000 tokens    →     1,000,000 attention scores
100,000 tokens  → 10,000,000,000 attention scores

That's quadratic, O(N²) — the kind of complexity that gets you flagged in code review. And every one of those scores is floating-point math on a GPU, in every layer, for every token the model generates. Clever kernels like FlashAttention have made the memory side of this manageable, but the compute doesn't go away: double the context and you quadruple the work. You can't buy your way out of a quadratic.

Signal degradation: lost in the middle

Even with unlimited budget, forcing a million tokens into the context breaks something else: the attention itself. This is what the "Lost in the Middle" paper documented — models get noticeably worse at using a fact when it's buried in the middle of a long context than when it sits near the start or the end.

Here's the mental model I use. Think of attention as a pie. For each token it processes, the model has exactly one pie — 100% — to divide among everything in the prompt.

You have one paragraph of facts that actually answers the question. Watch what happens to its share of the pie as you dump more raw documents into the context:

https://codepen.io/Devanshu-Mishra-the-scripter/pen/yyMePpM

Paragraphs in context Share for the one that matters
1 100%
5 20%
50 2%
1,000 0.1%

One caveat, because the pie is a simplification: attention isn't spread evenly. The model learns where to look, and modern long-context models are good at pulling a single needle out of a haystack. But the more noise you pack around your fact — near-duplicates, unrelated docs, chunks that look relevant but aren't — the more likely the model is to weight the wrong thing. Dump the entire wiki into the prompt and you're slicing the pie into a hundred thousand crumbs, and the fact you needed gets one of them.

The result: the model gets distracted, skips the data you gave it, and starts making things up. In other words — hallucination.

RAG as an attention optimizer

This is the real purpose of RAG. RAG is not just a search engine. RAG is a compute-optimization layer for self-attention.

flowchart LR
    A["All company data<br/>(millions of documents)"] --> B["Cheap compute layer<br/>(vector search filters for relevance)"]
    B --> C["Intensive compute layer<br/>(LLM self-attention runs only on the top 5 chunks)"]

It uses a cheap computation — vector search, keyword search, metadata filters — to pick out the handful of chunks that matter, and only that handful goes into the expensive computation: the LLM's attention. You pay O(N²) on five paragraphs instead of on the whole wiki.

Attention isn't the only reason RAG exists, and I'd be cheating if I pretended it was:

  • Fresh data. The model's weights stopped learning at its training cutoff. Retrieval is how it gets to see yesterday's incident report.

  • Access control. In a multi-tenant product, the retrieval layer is also where you enforce which documents a given user or tenant is allowed to see. The model can't leak what it was never shown.

  • Citations. If the answer came from chunk #3, you can show the reader chunk #3. Try doing that with something the model "remembers".

  • Cost and latency. Every token in the prompt costs money and milliseconds. Fewer, better tokens win on both.

But the attention argument is the one that explains why the pipeline has to exist at all, even in a world of 2-million-token context windows. Bigger windows don't kill RAG. They just move the line where the cheap filter starts paying for itself.


Further reading: Attention Is All You Need (the original transformer paper) · Lost in the Middle · FlashAttention

A

Framing retrieval as attention budgeting rather than search explains something that confuses people the first time they measure it: adding more relevant context can make answers worse. Every extra passage competes for the same attention mass, so a correct chunk sitting at rank nine behind eight plausible ones is functionally invisible. It also predicts which distractors hurt most, and it is not unrelated text - near-duplicates do the damage, because an older revision of the same policy looks exactly like the right answer from the model's point of view and there is no signal telling it which copy is authoritative. That is also why lost-in-the-middle is worse in practice than the benchmarks suggest: real corpora are full of near-copies, and the position penalty lands on top of an already ambiguous set.