🤖 machine learning

The Query Knows What to Forget: A Second Erase Direction for Linear Attention

This paper introduces Query-derived Erase Direction (QED), a novel mechanism for linear attention that adds a second, query-orthogonal erase vector to resolve state interference and significantly extend usable context length beyond training windows.

Original authors: Dhruman Gupta, Aritra Das, Debayan Gupta

Published 2026-08-17

📖 8 min read🧠 Deep dive

CC BY 4.0

Original authors: Dhruman Gupta, Aritra Das, Debayan Gupta

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to remember a story you heard a long time ago. Your brain has a limited amount of space to hold these memories. If you try to stuff too many stories into that small space, the details start to blur together, and you might mix up the ending of one tale with the beginning of another. This is a bit like how modern AI models handle long conversations. They use a special type of memory called "linear attention" to keep track of everything said so far without running out of space. Instead of saving every single word in a giant, growing pile, they compress the information into a fixed-size "state," like a magical notepad that updates itself as new words arrive.

However, there's a catch. As the conversation gets longer, the old memories on this notepad start to interfere with each other. When the AI tries to find a specific piece of information (like a name mentioned 10,000 words ago), the noise from all the other words makes it hard to retrieve the right answer. Scientists have been trying to fix this by making the AI smarter about what to delete. They've built systems that can "erase" old, irrelevant info to make room for new info, but until now, these erasers have been a bit clumsy. They only know what to delete based on the current word being processed, not on what the AI is actually looking for.

This paper introduces a clever new trick called QED (Query-derived Erase Direction) to fix that clumsiness. The researchers, working with models that have 340 million parameters, discovered that by letting the AI's "search question" (the query) help decide what to erase, the model can hold onto important information much longer. They trained their models on a massive dataset of 15 billion tokens and found that this new method allows the AI to remember details from much further back in the conversation than before. Specifically, on a test called S-NIAH-1, the new method doubled the amount of text the model could reliably read through, pushing the useful limit from about 8,000 words to 16,000 words. The paper suggests this is a significant step forward for making AI better at reading long documents, though the authors note that some specific parts of their new tool might need more testing to see if they are strictly necessary.

The Problem: The "One-Way" Eraser

To understand the breakthrough, let's look at how these AI models currently work. Imagine the AI's memory is a busy library where books (information) are constantly being swapped out. When a new book arrives, the librarian (the model) has to decide which old book to throw away to make space.

In the models used before this paper (like Gated DeltaNet-2), the librarian only looks at the new book (the "key") to decide what to toss. It's like saying, "This new book is about cats, so I'll throw away the old book about cats." This works okay, but it has a blind spot.

The problem is that the AI doesn't just look at the new book; it also has a specific question it's trying to answer (the "query"). Maybe the question is, "Who is the main character?" The librarian needs to find the answer to that specific question. But the old "one-way" eraser only knows what to delete based on the new book's topic, not the question being asked.

The authors realized that the "interference"—the noise that confuses the AI—lives exactly where the question is looking. If the question is looking for a specific detail, but the librarian only erases based on the new book's topic, the confusing noise might stay right in the path of the question. It's like trying to find a needle in a haystack while the librarian only removes hay that looks like the needle's color, ignoring the fact that the needle is actually hidden under a pile of straw.

The Solution: The "Two-Way" Eraser

The paper proposes a solution called QED. Instead of just one eraser, they add a second one that is guided by the question itself.

Think of it this way: The AI now has a "Key-Directed Eraser" (the old one) and a new "Query-Directed Eraser."

The Key-Directed Eraser still does its job: it looks at the new information and clears out the old stuff that matches it.
The Query-Directed Eraser looks at what the AI is currently asking for. If the question is "Who is the villain?", this eraser specifically targets and removes old, confusing information that might be hiding the answer to that question.

Crucially, this new eraser is designed to be "orthogonal" to the old one. In plain English, this means it cleans up the parts of the memory that the old eraser can't reach. It's like having a vacuum cleaner that can suck up dust from the corners of the room that the broom misses.

The researchers tested this by training models on 15 billion tokens of text. They found that when they let the "question" guide the eraser, the model became much better at finding specific facts in long texts.

What They Found (and What They Didn't)

The results were quite promising, but the authors were careful not to overhype it. Here is what the data actually showed:

Long-Context Superpower: On a test called S-NIAH-1 (which checks if a model can find a "needle in a haystack" of text), the new method improved accuracy significantly. For example, at a length of 16,000 tokens (words), the standard model got about 21.8% correct, while the new QED model got 39.8% correct. At 32,000 tokens, the jump was even more dramatic, going from 9.4% to 15.8%.
Doubling the Usable Length: The authors state that this improvement effectively doubles the usable context length. Before, the model started to get confused around 8,000 tokens; with QED, it stayed reliable up to 16,000 tokens.
No Magic Perplexity: Interestingly, the new method didn't make the AI better at writing or predicting the next word in a sentence (a metric called "validation loss"). The loss stayed exactly the same. This suggests that QED doesn't make the AI "smarter" in a general sense; it just makes it better at retrieving specific memories from its compressed state.

The "Ablation" Mystery: What Parts Matter?

To figure out exactly why QED worked, the researchers ran experiments where they turned off different parts of the new system. This is like taking apart a car engine to see which bolt is actually making the car go faster.

They tested three main components:

The Direction: Using the "question" (query) to guide the eraser.
The Gate: A switch that decides when to erase (so it doesn't delete things the AI still needs).
The Projection: A mathematical rule to ensure the new eraser doesn't accidentally mess up the old eraser's work.

Here is the twist: The results were a bit mixed.

The Direction is Key: Every time they used the "question" to guide the eraser, the model got better. This part seems to be the real hero.
The Gate and Projection are Unclear: When they turned off the "gate" or the "projection," the results were inconsistent. Sometimes the model did better, sometimes worse, depending on which random seed (starting point) they used for training. The authors suggest that while these parts might help, the experiments didn't prove they are strictly necessary. It's possible that for the specific size of model they used, the "projection" wasn't actually doing much because the model naturally avoided the problems it was designed to fix.

The Bottom Line

The paper concludes that QED is a successful way to help AI models remember things in long conversations. By letting the "question" help decide what to forget, the model can clear out the noise that usually blocks its memory.

However, the authors are honest about the limits of their findings. They didn't prove that the "gate" or the "projection" are essential; they only proved that using the "query" to guide the eraser is the most consistent improvement. They also noted that the strength of this new eraser might need to be tweaked depending on the specific task, as changing the "strength" of the eraser at test time gave different results for different models.

In short, the paper suggests that if you want an AI to read a whole novel and remember the name of the villain in the last chapter, you should give it a "query-guided eraser." But you might still need to do more testing to figure out exactly how to tune the other knobs on that new machine.

Technical Summary: The Query Knows What to Forget

Problem Statement
Linear attention mechanisms replace the sequence-length-dependent cache of softmax attention with a fixed-size recurrent state, offering constant memory and computation per token. However, as context length grows, multiple stored associations must share this bounded state, leading to interference that degrades retrieval reliability. While recent architectures like Gated DeltaNet-2 (GDN-2) have improved state selectivity by decoupling erase and write gates, a fundamental mismatch remains: the erase vector is derived from the current token's key, whereas the retrieval interference is measured by the query. Consequently, content visible to the query (and causing interference) may lie in parts of the state that the key-derived erase step cannot directly modify.

Methodology: Query-derived Erase Direction (QED)
The authors propose QED, a modification to the GDN-2 update rule that introduces a second erase direction derived from the query.

Mechanism: The standard GDN-2 erase vector $a_t = b_t \odot k_t$ is augmented with a gated term derived from the query $q_t$ . The new erase vector is defined as:
$a_t = b_t \odot k_t + \lambda_h P^\perp_{k_t}(b_t \odot q_t)$
where $b_t$ is the channel-wise erase gate, $\lambda_h$ is a learned per-head strength, and $P^\perp_{k_t}$ projects the query-derived term orthogonal to the key $k_t$ .
Theoretical Justification: In the fast-weight view, a standard rank-one delta update (directed by the key) cannot alter the component of the state orthogonal to the key. Since the query selects content based on its own direction, interference often resides in the query-orthogonal-to-key subspace. QED allows the model to cancel this specific "old-state content" by projecting the correction along the query direction while maintaining the rank-one update structure and stability conditions of GDN-2.
Cost: The modification preserves the rank-one update and incurs negligible additional computational cost.

Experimental Setup
The authors trained 340M-parameter models (24 layers, 1024 hidden size) on 15 billion tokens of FineWeb-Edu data. The study utilized two random seeds and five experimental arms to isolate the contributions of the mechanism's components:

Baseline: GDN-2 (no correction).
QED (full): Full mechanism with gating, projection, and query-derived direction.
QED-nogate: Query direction and projection, but no gating.
QED-nogate-noproj: Query direction without gating or projection.
QED-rand: A control using a fixed random unit vector instead of the query for the direction.

Evaluation included language modeling perplexity, long-context retrieval (RULER S-NIAH tasks), and document recall (Based recall tasks: SWDE, FDA, SQuAD-completion).

Key Results

Long-Context Retrieval: QED significantly improves retrieval accuracy beyond the training window (2048 tokens). On the S-NIAH-1 task, QED improved accuracy by up to 39 percentage points at extrapolation lengths (4K–32K) compared to the baseline, effectively doubling the usable context length.
Document Recall: QED improved FDA recall scores by 3.4 and 3.7 points across seeds. No significant changes were observed in SQuAD-completion or SWDE (the latter being a null result).
Language Modeling: The mechanism targets retrieval specifically; validation loss (perplexity) remained unchanged across all arms, with variations within 0.0034 nats.
Ablation Insights:
- Direction: Using the query to define the erase direction is consistently beneficial. The QED-rand control (random direction) failed to match the full QED performance, confirming the necessity of the query-derived direction.
- Gating and Projection: The benefits of the erase gate ( $b_t$ ) and the key-orthogonal projection ( $P^\perp$ ) were inconclusive, as their effects varied between training seeds. Analysis suggests that at the current model scale, the learned corrections naturally have negligible alignment with the key, rendering the projection constraint effectively inactive (a safeguard that was not strictly necessary for stability in these specific runs).

Significance and Claims
The paper claims that QED resolves the mismatch between reading and editing in linear attention by allowing the model to "forget" content specifically selected by the query. This enables more effective cancellation of interference in the recurrent state.

Practical Impact: The method improves retrieval at every length past the training window without degrading language modeling quality.
Modest Scope: The authors explicitly state that while the query-derived direction is the primary driver of improvement, the necessity of the gating and projection sub-components remains undecided due to inconsistent results across runs. They emphasize that further work is required to determine the optimal correction strength ( $\lambda_{max}$ ) and to test the stability constraints of the projection at larger model scales.