Dev.to AI 🤖 Ai 👁 0 📖 5 min read

NoteBank: How Replaying a Few Saved Snippets Restores Long-Context Recall in Linear-Attention Models

Linear-attention and state-space models are cheap to run on long inputs because they carry a fixed-size state instead of a growing KV cache. The catch is that the state forgets. A new KAUST paper, "History Replays Itself

Linear-attention and state-space models are cheap to run on long inputs because they carry a fixed-size state instead of a growing KV cache. The catch is that the state forgets. A new KAUST paper, "History Replays Itself" (arXiv:2610.15715), proposes NoteBank, a training-free add-on that stores a few raw text snippets on the side and replays them through the model when a question arrives. Here is how it works and what the numbers do and don't show.

What NoteBank is

NoteBank is a sparse episodic memory bolted onto pretrained recurrent and hybrid language models, such as Gated DeltaNet (GDN) and Mamba-2. It does not change the backbone's weights or architecture. It also does not use an external embedding model, a separately trained retriever, or a separate reader network. The model itself decides what to save, finds it again, and reads it.

The authors' framing is that a recurrent state is asked to do two jobs at once: carry the ongoing computation and act as long-term storage. As context grows, the state is continually overwritten, so evidence can be lost before anyone knows it will matter. NoteBank splits the jobs. The recurrent state keeps doing what it does, and a small archive of verbatim text handles recall.

Why it matters

Models with linear attention layers are increasingly used in hybrid designs, including Kimi Linear and the Qwen3-Next family, because they cut memory and compute for long contexts. Their weak spot is retrieval: finding a specific fact buried far back in the input.

On six RULER needle-in-a-haystack tasks at one million tokens, the paper reports GDN's mean accuracy going from 1.7% to 70.3% with NoteBank. It reports no significant retrieval degradation from 4K to 1M tokens. Persistent storage is described as roughly 1/2,600 of a full multi-head-attention KV cache at matched width, depth and precision (74.5 MiB vs. 192 GiB in their Figure 1). The extra compute is reported as at most 0.025% of analytical FLOPs.

How it works

The procedure has three stages, all built on features the model already computes.

1. Write. The context is processed in chunks of C tokens, with the recurrent state carried across chunks as usual. Within each chunk, the writer scores every token by how distinctive its key vector is: one minus the cosine similarity between the token's key and the chunk's mean key, averaged over layers and heads. Windows with high total distinctiveness are chosen, expanded to sentence boundaries within a token budget, and stored as notes. This happens before any question is known, which is why selection has to rely on a generic signal like distinctiveness.

2. Address. Each note gets a compact address per layer-head channel: the native query vector at the note's most distinctive token. When a question arrives, it is run from a copy of the saved state to get its own query vectors. Each channel scores every note by maximum cosine similarity and votes for its best match. The notes with the most votes, up to R of them, are retrieved. Ties go to earlier notes.

3. Replay. The retrieved notes are concatenated in original document order and fed through the unchanged backbone, starting from a fresh copy of the saved state. Then the question is processed and the answer generated. In effect, the model re-reads the relevant passages right before answering, so the evidence is fresh in its 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)      # distinctiveness 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 use replayed evidence on downstream tasks. The headline retrieval results are training-free. The adapters are only enabled during replay and answering, so writing and addressing are unchanged.

Costs scale as you'd expect. Archive size grows linearly with context length at a rate set by the note budgets. Matching a question against the archive is exhaustive, so it also grows linearly. The replay budget is constant, at most R times the note length in extra tokens.

Real-world implications

The practical idea is that a cheap recurrent model can behave like it has a much larger memory if you give it a place to put verbatim evidence. That is attractive for long-running agents, long documents and chat histories, where holding a full KV cache is expensive. The paper reports gains beyond synthetic retrieval: improvements on LongBench v1/v2 and MemoryAgentBench, on backbones from three 1.3B research models up to Mamba-2-8B and a post-trained Qwen3.5-9B.

There are caveats worth keeping in mind:

  • Needle tests are the best case. RULER-style tasks reward finding a distinctive snippet. Questions needing evidence spread thinly across a document, or aggregation over many passages, are harder for a sparse archive that is written before the question is known.
  • Storage and compute claims are analytical. The paper's FLOP figures are layer-level analytical estimates and, as the authors state, do not measure runtime. Exhaustive matching and replay add real latency that FLOP counts don't capture.
  • It is a preprint. The results have not gone through peer review or independent replication yet.
  • It's closer to retrieval than to a new architecture. The novelty is in using the backbone's own keys and queries for writing and addressing, and its own forward pass for reading. That removes extra components, but the core pattern, store snippets and re-read them, will feel familiar to anyone who has built retrieval-augmented systems.

For builders, the useful takeaway is the design principle: you may not need to make the recurrent state bigger or retrain the model to fix long-range recall. Keeping a small raw-text archive next to the state and replaying from it can recover a lot of lost capability at modest cost.

Takeaway

NoteBank separates "remember everything" from "compute with the current context." It writes distinctive spans while reading, finds them using the model's own query features, and replays them before answering. On the reported benchmarks that lifts a 1.3B Gated DeltaNet from near-zero to about 70% mean accuracy at one million tokens with tiny storage overhead. Treat it as a promising, training-free technique to test on your own workloads, especially ones where the answer lives in one distinctive passage, rather than a solved long-context problem.

Sources

📰 Read the original article on Dev.to AI

Originally published by Dev.to AI. Aggregated on AIWithGhost for educational purposes — full credit and traffic to the original publisher.