Breaking the Linear Attention Barrier
A conversation with an advanced AI on why your transformer slows down with every token — and what we can do about it.
If you have ever run a large language model and noticed that it gets progressively slower as the conversation goes on, you have already felt the pain of linear attention scaling. It is not a bug. It is a fundamental property of how transformers work at inference time — and fixing it is one of the most important open problems in practical AI today. I had a long, detailed discussion with an advanced AI system about this exact problem, and I want to share what came out of it in the clearest terms I can manage.
Let’s start at the root of the problem, before jumping to solutions.
Why Does Inference Get Slower Over Time?
When a transformer generates each new token, it needs to look back at every single previous token it has produced or received. That’s by design — attention is what makes these models so context-aware. But the cost of that backward look grows with every step. For a sequence of length t, generating the next token costs, roughly:
Cost ≈ L × t × d_h
# L = layers · t = sequence length · d_h = head dimension
People often assume that KV caching — where you store previous key-value pairs to avoid recomputing them — solves this. It helps enormously, but it doesn’t solve the underlying issue. The cache removes recomputation, but you still need to scan all those stored keys every time you generate a new token. As the context grows into the tens of thousands of tokens, that scan becomes the bottleneck. The memory bus, not the processor, ends up being the limiting factor.
“The real question isn’t ‘can we avoid recomputing attention?’ — KV caching already does that. The real question is ‘do we actually need to attend to every single past token to generate the next one?’”
The answer, it turns out, is almost certainly no. Research consistently shows that attention mass is highly sparse — most of the “importance weight” during generation concentrates on a small fraction of past tokens. The challenge is figuring out which fraction, cheaply and reliably, before you compute the full attention.
The Core Paradox of Token Selection
Here is where it gets philosophically tricky. Token importance is not a fixed property. It depends on the current query — the token being generated right now. A date mentioned two pages back might be completely irrelevant for most of the response, but becomes the most critical piece of information the moment you start generating a sentence that references it. Importance is also different across layers: lower layers care about local syntax, while higher layers are where long-range reasoning happens.
The fundamental paradox is this: to know which tokens are unimportant, you often need to run attention first. This is why simple heuristics like “drop the oldest tokens” or “skip punctuation” are dangerous. They’ll work most of the time, and fail silently at the worst possible moment.
Any real solution has to approximate token importance cheaply, without first computing the full attention it’s trying to avoid. That’s a hard constraint, and it’s what separates naive ideas from robust ones.
The Solution Landscape
There are several families of approaches that researchers and engineers have explored, each with genuine strengths and genuine blind spots. I’ll go through each one in terms of what it’s actually doing and where it breaks down.
Reactive Pruning — Following the Heavy Hitters
The simplest approach is to track which tokens have historically received a lot of attention mass. If a token keeps getting attended to, it’s probably important — keep it around. Tokens that consistently receive near-zero attention can be evicted from the cache. This is sometimes called “heavy hitter” or “attention sink” tracking.
The appeal is obvious: you’re using information you’ve already computed, so there’s no overhead. You also don’t need to retrain or modify the model. The weakness is that this is entirely backward-looking. A token that nobody has needed yet might be the most important thing in the context for what comes next — what you might call a “sleeper token.” Reactive pruning has no way to anticipate this.
Static Sparse Architectures
Another family of approaches bakes sparsity directly into the model architecture. Instead of attending over everything, the model uses patterns like sliding windows (only attend to the last N tokens), block-sparse layouts, or strided patterns that skip tokens at regular intervals. Longformer and BigBird are well-known examples of this idea.
These are very efficient and predictable, but they’re architectural choices made before training. They can miss rare long-range dependencies that don’t fit the predefined pattern, and they’re not adaptive — the same pattern applies regardless of what’s actually in the context.
Learned Salience Prediction
A more sophisticated approach trains a lightweight predictor that looks at a token as it enters the context and assigns it a score: how useful is this token going to be in the future? You store one scalar per token, and at generation time, you only attend to the high-scoring ones.
This is conceptually elegant — it’s proactive rather than reactive. But it requires careful training (the target, “future usefulness,” isn’t directly observable), and the consequences of a wrong prediction are severe: once you decide a token’s score is low and stop loading it, you can’t recover that information if you were wrong. The errors are silent and irrecoverable unless you build in explicit safeguards.
Retrieval-Augmented Selection
Rather than predicting importance, you can treat the context like a database and query it. When generating a new token, run a fast approximate similarity search over stored key blocks to find the most relevant ones, then only run exact attention on that retrieved subset. The cost becomes proportional to the retrieved set size, not the total context length.
This is genuinely query-dependent, which is exactly what you want. The problem is that retrieval errors are silent — if the similarity search misses a relevant block, the model just proceeds without it. And you need to maintain an additional indexing structure, which adds implementation complexity.
Hierarchical Memory
Perhaps the most ambitious approach is to reorganize the context itself into multiple levels of resolution: raw recent tokens at the bottom, compressed chunk summaries in the middle, and high-level global summaries at the top. Attention operates across these levels simultaneously, like reading at multiple zoom levels at once.
This scales gracefully and mirrors something vaguely like how humans actually process long narratives. The difficulty is in the compression: building a good lossy summary of an arbitrary chunk of text that preserves everything the future generation process might need is itself a hard machine learning problem. Training these systems is unstable.
A Hybrid Proposal: The Scout-Based Cascade
After going through all of these, the discussion converged on a hybrid architecture that tries to take the strengths of each approach while hedging against their individual failure modes. The core idea is a two-stage cascade — a “scout” followed by exact attention.
// At each decode step t:
Stage 1 — Scout
Run lightweight Q/K projections (d' << d_h)
Attend over block representatives only
Output: top-k candidate blocks
Cost: O(t/B × d') ← much cheaper than full attention
Stage 2 — Exact Attention
Attend only to:
selected blocks from Stage 1
recent window (always kept)
protected tokens (prompt, system)
Cost: O(L × kB) where kB << t
The scout runs on compressed, low-dimensional representations of blocks of tokens rather than individual tokens. It’s cheap enough that you can afford to run it over the entire context, and its job is simply to identify which blocks are worth examining more carefully. The exact attention in the second stage then operates only on that filtered subset, keeping the total cost sublinear in context length.
What makes this more robust than a pure prediction approach is that block selection scores combine two signals:
Proactive Signal — Predicted Salience: A learned score assigned when the block enters the cache, estimating its likely future relevance. This is the forward-looking component that enables early eviction decisions.
Reactive Signal — Heavy-Hitter History: An accumulated score tracking how much attention mass this block has actually received so far. This corrects for cases where the predictor was wrong and a block keeps turning out to be relevant.
Combined Score: The final block score is a weighted blend of both: score = α × salience + (1−α) × heavy_hitter. Neither signal alone is sufficient, but together they cover each other’s blind spots.
An important subtlety here is that this system operates at block granularity, not individual token granularity. Grouping tokens into blocks of a fixed size before scoring and pruning has several practical advantages: it’s more stable, more memory-bandwidth friendly (you load or skip whole cache lines rather than scattered tokens), and it bounds the worst-case error — if a block is incorrectly dropped, you lose a small contiguous chunk rather than a randomly scattered set of tokens.
Layer-Wise Sparsity and Protected Tokens
One of the more interesting nuances that came up in the discussion is that you shouldn’t apply the same pruning budget uniformly across all layers. Lower layers in a transformer handle local coherence — syntax, nearby word relationships, local structure. They rarely need to reach far back into context to do their job. Higher layers are where the model integrates long-range semantic information and handles abstract reasoning. The sparsity schedule should reflect this: lower layers get a narrow local window with almost no global retrieval, while higher layers get the full benefit of the block selection mechanism.
There’s also a set of tokens that should be unconditionally protected from pruning regardless of their scores: the initial prompt and any system instructions, structural delimiters, and the most recent window of tokens. These are cheap to protect (they’re a small, fixed set) and essential for coherent generation. Pruning the system prompt to save memory would be a spectacular own goal.
Note: The bottleneck in long-context inference is usually not floating point computation — it’s memory bandwidth. Loading the KV cache for tens of thousands of tokens takes time. This means effective pruning must reduce actual memory reads, not just skip multiplications. Block-level granularity and cache-aligned access patterns matter as much as the selection logic itself.
Where Things Can Still Go Wrong
No pruning system is perfect, and it’s worth being clear-eyed about the failure modes. The most dangerous scenario is a rare constraint that appears early in the context and is never referenced again until very late — something like “never mention X” or “always format output as Y.” These tokens accumulate almost no attention mass and look completely unimportant to a reactive system. A salience predictor might catch them, but only if it was trained on examples of this kind of delayed relevance.
Topic shifts are another hazard. If the conversation changes direction significantly, previously irrelevant blocks might suddenly become critical, while blocks that were being heavily attended to might become noise. A system that relied too heavily on recent attention history won’t adapt quickly enough.
The proposed mitigations are pragmatic rather than perfect: mixing proactive and reactive signals dilutes the failure of either alone, block granularity limits the damage of any single wrong decision, and periodic “refresh” steps — where you temporarily relax the pruning constraints and run broader attention — can catch cases where important information has been gradually crowded out.
The Bigger Picture
Stepping back from the implementation details, what this discussion really clarified for me is that the goal isn’t token pruning — it’s memory system design. Standard transformer attention is a flat, undifferentiated memory: every past state is equally accessible, equally expensive to query, and has no internal organization. That’s elegant and powerful, but it scales badly.
What we’re really trying to build is a structured, hierarchical memory system that approximates the behavior of full attention at a fraction of the cost. The scout-based cascade is one concrete step in that direction: it imposes a lightweight structure (blocks, scores, selection) that enables selective retrieval without abandoning the query-dependence that makes attention so powerful in the first place.
“Future large language models will not rely on flat, uniformly dense memory. They will implement structured, hierarchical, and selective memory systems — closer in spirit to how biological memory works than to the original attention mechanism.”
Whether that takes the form of scout cascades, retrieval-augmented attention, hierarchical compression, or something we haven’t thought of yet, the direction seems clear. The linear scaling wall is real, and the community is actively finding ways around it. The most promising paths are the ones that preserve the query-dependence of attention while adding intelligent structure to what gets loaded and when.
I’ll be writing more about specific implementations and recent papers in this space. If you have thoughts, disagreements, or questions — especially if you’ve been working on inference optimization yourself — I’d genuinely love to hear from you.
This post is based on a structured discussion with an advanced AI system about inference optimization in autoregressive transformers. The ideas here draw from ongoing research in sparse attention, KV cache compression, and efficient long-context inference.