[KongchangAI]
· 3 min read· 1,924 words

NoteBank: Replaying a Few Text Snippets Boosts Linear Attention Long-Context Recall from 1.7% to 70%

NoteBank: Replaying a Few Text Snippets Boosts Linear Attention Long-Context Recall from 1.7% to 70%

KAUST proposes NoteBank, a training-free module that restores long-context retrieval in linear attention models.

A KAUST preprint, History Replays Itself, proposes NoteBank, a training-free sparse episodic memory module that attaches to Gated DeltaNet and Mamba-2 without changing backbone weights. It stores distinctive text snippets during encoding, addresses them with the model's native query vectors when a question arrives, and replays the top notes before answering. On million-token RULER needle-in-a-haystack tasks, GDN's average accuracy jumps from 1.7% to 70.3%, with storage about 1/2600 of a full KV cache.

Linear attention and state-space models process long inputs cheaply because they replace a growing KV cache with a fixed-size state. The trade-off: the state forgets. A new KAUST paper, History Replays Itself (arXiv:2610.15715), proposes NoteBank—a training-free add-on module that stores a few raw text snippets alongside the state and feeds them back to the model when a question arrives.

NoteBank paper source screenshot

What NoteBank Is

NoteBank is a sparse episodic memory module that can be attached directly to pretrained recurrent and hybrid language models such as Gated DeltaNet (GDN) and Mamba-2. It does not change the backbone's weights or architecture, nor does it require an external embedding model, a separately trained retriever, or a standalone reader network. The model decides what to store, retrieves it itself, and reads it itself.

The authors' framing is clear: the recurrent state is asked to do two things at once—carry the current computation and serve as long-term storage. As context grows, the state is continuously overwritten, and evidence can be lost before anyone realizes it mattered. NoteBank splits these two responsibilities apart: the recurrent state keeps doing what it does, while a small verbatim archive handles recall.

Why This Matters

Models with linear attention layers are increasingly showing up in hybrid designs, including Kimi Linear and the Qwen3-Next series, because they cut the memory and compute overhead of long contexts. Their weak spot is retrieval: finding a specific fact deep in the input.

On six RULER needle-in-a-haystack tasks at one million tokens, the paper reports GDN's average accuracy jumping from 1.7% to 70.3%. From 4K to 1M tokens, retrieval performance shows no clear degradation. Persistent storage is roughly 1/2600 that of a full multi-head attention KV cache at the same width, depth, and precision (74.5 MiB vs. 192 GiB in Figure 1 of the paper). Extra compute is at most 0.025% of the analytical FLOPs.

The trade-off between linear attention and KV cache

Standard Transformer attention computes dot products against the Key and Value of all previous tokens when generating each token, so the KV cache grows linearly with context length. Under long contexts, memory usage and per-step compute inflate together. Linear attention rewrites the attention computation using associative kernels or state recurrence, compressing historical information into a fixed-size state matrix that only needs updating each step; computational complexity drops from O(n²) to O(n), and memory becomes constant. State-space models (SSMs) like the Mamba series follow the same recurrence paradigm. The cost: a fixed state cannot losslessly retain all historical detail, and early information gets overwritten by later updates. This is the core tension NoteBank tries to remedy—the state is good for fast computation but not for long-term precise storage. Once the two responsibilities are separated, linear models' long-context retrieval weakness has a chance to be substantially alleviated.

The Workflow: Write, Address, Replay

The whole pipeline has three stages, all built on features the model has already computed.

Write

The context is processed in chunks of C tokens, with the recurrent state passed across chunks as usual. Within each chunk, the writer scores each token by how "unique" its key vector is: 1 minus the cosine similarity between that token's key and the chunk's average key, then averaged over layers and heads. Windows with high total uniqueness are selected, expanded to sentence boundaries, and kept within a token budget to be stored as notes. This step happens before any question appears, so selection can only rely on a generic signal like uniqueness.

Why uniqueness scoring works

NoteBank's write policy rests on an assumption: in long documents, key vectors of frequently recurring phrases and common connectives tend to be similar, while tokens that genuinely carry unique information (proper nouns, dates, numbers, rare terms) have keys that deviate from the local average direction. Using 1 minus cosine similarity as the uniqueness score is essentially measuring how much a token "deviates" semantically from its surroundings. This signal comes entirely from features the model has already computed, requiring no additional embedding model or supervision labels. The cost is that it cannot anticipate what future questions will need, so it can only prioritize storing snippets that "look distinctive." For tasks where information is spread evenly and the answer depends on aggregating multiple ordinary sentences, this uniqueness-based selection will miss key evidence—one of the limitations the paper acknowledges.

Address

Each note gets a compact address on each layer-head channel: the native query vector at the note's most unique token. When a question arrives, the question is run starting from a copy of the saved state to obtain its own query vectors. Each channel scores all notes by maximum cosine similarity and casts a vote for the one it considers the best match. The top R notes by vote count are retrieved, with ties broken in favor of earlier notes.

Independent voting per layer-head channel

In multi-head attention or linear attention structures, the query spaces learned by each layer and each head may attend to different aspects of information. Rather than compressing note addresses into a single vector, NoteBank stores a native query vector for each note on every layer-head channel, effectively recording multiple "retrieval clues" for the same note. When a question arrives, each channel independently computes its cosine similarity against all note addresses and casts one vote for the best-matching note. The final top R are selected by vote count. The benefit of this multi-channel voting mechanism is robustness: even if some heads are insensitive to the current question, others may still hit the right note. The cost is that address storage grows linearly with the number of layers and heads. Breaking ties in favor of earlier notes is a stabilizing choice biased toward temporal order, avoiding randomness affecting retrieval results.

Replay

The retrieved notes are concatenated in original document order and reprocessed starting from the unmodified backbone (from a copy of the saved state). The question is then processed and an answer generated. In effect, the model re-reads the relevant passages before answering, keeping the evidence fresh in the state rather than faded.

A rough sketch of the control flow:

# Write (during context encoding)
for chunk in chunks(context, C):
    state = backbone.update(state, chunk)
    notes += select_notes(chunk, keys, top_w=W)      # uniqueness scoring
    addresses += make_addresses(notes, queries)

# Read (when the question arrives)
q_feats = backbone.query_features(question, state_copy)
top = vote_over_channels(q_feats, addresses, top_r=R)
replay = concat_in_document_order(top)
answer = backbone.generate(state_copy, replay + question)

The paper optionally trains lightweight LoRA adapters to help the model exploit replayed evidence on downstream tasks. The core retrieval results need no training. The adapters are enabled only during the replay and answer stages; writing and addressing stay unchanged.

Costs grow as expected: the archive size grows linearly with context length, at a rate determined by the note budget; matching the question against the archive is exhaustive, hence also linear; the replay budget is constant, at most R times the note length in extra tokens.

Practical Implications and Limitations

The practical takeaway: give a cheap recurrent model a place to keep verbatim evidence, and it can behave as if it had a much larger memory. This is attractive for long-running agents, long documents, and chat histories—scenarios where maintaining a full KV cache is expensive. The paper reports improvements not just on synthetic retrieval: gains also appear on LongBench v1/v2 and MemoryAgentBench, spanning three 1.3B research models up to Mamba-2-8B and a trained Qwen3.5-9B.

A few points to keep in mind:

  • Needle-in-a-haystack is the best case. RULER-style tasks reward finding a single distinctive piece of text. If the information a question needs is sparsely scattered across the document, or requires aggregating multiple passages, a sparse archive written before the question appears will struggle more.
  • Storage and compute claims are analytical. The paper's FLOP numbers are layer-level analytical estimates, and the authors note they do not measure runtime. Exhaustive matching and replay add real latency that FLOP counts cannot capture.
  • This is a preprint. The results have not been peer-reviewed or independently reproduced.
  • It is closer to retrieval augmentation than a new architecture. The novelty lies in using the backbone's own keys and queries for writing and addressing, and its own forward pass for reading. This eliminates extra components, but the core pattern of "store snippets and re-read them" is familiar to anyone who has built retrieval-augmented systems.

For developers, the useful takeaway is a design principle: you may not need to enlarge the recurrent state or retrain the model to fix long-range recall. Keeping a small verbatim archive alongside the state and replaying from it can restore a large amount of lost capability at modest cost.

The RULER benchmark and needle-in-a-haystack tasks

RULER is a synthetic benchmark suite specifically for evaluating long-context capabilities, with its most classic task being "needle in a haystack": planting a specific fact in a long stretch of hundreds of thousands of tokens of irrelevant text, then asking the model to retrieve that fact based on a question. Such tasks test precise retrieval ability, not reasoning or multi-hop aggregation. NoteBank's jump from 1.7% to 70.3% on RULER shows that its write-address-replay mechanism can genuinely give linear models basic precise recall. But as the paper notes, needle-in-a-haystack is the best case: the information to be retrieved is a distinctive, concentrated, local piece of text, which happens to align well with selecting segments by uniqueness and voting by query similarity. If the answer requires aggregating across multiple scattered passages, or if the question involves information that was not conspicuous during the write phase, the sparse archive's recall will drop noticeably. The high RULER score should therefore be understood as the method's applicability boundary rather than comprehensive proof of general long-context ability.

Summary

NoteBank separates "remembering everything" from "computing with the current context." It writes distinctive snippets at read time, finds them using the model's own query features, and replays them before answering. On the reported benchmarks, this takes a 1.3B Gated DeltaNet from near zero to about 70% average accuracy at one million tokens, with minimal storage overhead. Think of it as a training-free technique worth testing on your own workloads—especially ones where the answer hides in a single distinctive passage—rather than a solved long-context problem.

References

One-line Summary

KAUST proposes NoteBank, a training-free module that restores long-context retrieval in linear attention models.

Paragraph Summary

A KAUST preprint, History Replays Itself, proposes NoteBank, a training-free sparse episodic memory module that can be attached to linear attention and state-space models such as Gated DeltaNet and Mamba-2. It leaves backbone weights untouched and instead separates the recurrent state's computational role from its long-term storage role: during context encoding, a small number of raw text snippets are selected by token uniqueness scoring and stored in a note bank; when a question arrives, the model's own native query vectors address the notes and vote for the most relevant ones; finally, those notes are replayed in document order before generating an answer. On the million-token RULER needle-in-a-haystack task, GDN's average accuracy rises from 1.7% to 70.3%, persistent storage is about 1/2600 of a full KV cache, and extra compute is at most 0.025% of analytical FLOPs. The paper also notes the results are not yet peer-reviewed, needle-in-a-haystack is the best case, and storage and compute figures are analytical estimates rather than measured latency.

#NoteBank#linear attention#long context retrieval#Gated DeltaNet#state space models#KV cache#RULER benchmark
Share:

Related articles