History Replays Itself: Unleashing Long-Context Capability in Linear Attention

AI summary

Overview

Research area: Efficient long-context language modelling — specifically, linear/recurrent attention architectures (Gated DeltaNet, Mamba-2) and how to extend their effective memory without retraining.

Technical level: Intermediate. Readers should know roughly how attention works and what a recurrent "state" is; the paper's maths is confined to a short preliminary section, but the evaluation section assumes familiarity with benchmarks like RULER and LongBench.

Scope (one sentence): The paper introduces NoteBank, a training-free, plug-and-play sparse episodic memory that archives selected text spans during context encoding and replays them through the backbone's own forward pass at question time, recovering long-context retrieval in pretrained linear-attention models at near-zero added computation.

What This Paper Is About

Linear attention models compress all past context into a fixed-size recurrent state, so a single vector must both support ongoing computation and store everything from an ever-growing history. This has reinforced a view that the state's capacity sets a hard ceiling on long-context ability. The authors argue the ceiling is only apparent: if selected raw text is archived before any question is seen and later replayed through the model's normal forward pass, the recurrent state no longer has to carry the full burden of long-term storage — and the model can use distant evidence it could never have retained on its own.

Key Contributions

1. Breaking the apparent fixed-state ceiling. The paper shows that a pretrained recurrent model's usable long-context capability can substantially exceed what its fixed-state interface exposes, because poor retention through the state does not imply an inability to use distant evidence once that evidence is preserved separately.

2. Turning pretrained computation into episodic memory. NoteBank decouples evidence preservation from recurrent computation using three parts: question-independent sparse text archiving (write), addressing via the backbone's native query features, and on-demand replay through the native forward pass (read). It needs no retraining, no external embedding model, no separately trained retriever, and no separate neural reader.

3. Recovering capability at near-zero compute overhead. Across recurrent and hybrid backbones the method recovers performance from controlled million-token retrieval to natural-document and agent-memory tasks, while adding negligible analytical FLOPs and requiring orders of magnitude less persistent context storage than a matched full-attention KV cache.

4. Diagnostics that separate memory access from read adaptation. Through archive-coverage measurement, oracle routing interventions, router swaps (random, input embeddings, BM25, BGE), and frozen-versus-LoRA readers, the authors disentangle how much of the gain comes from making evidence accessible versus from adapting the reader.

Main Findings

• Million-token retrieval recovered without training: At one million tokens, NoteBank raises Gated DeltaNet's mean accuracy across six RULER retrieval tasks from 1.7% to 70.3%, with no observed significant retrieval degradation from 4K to 1M tokens. At 1M, the training-free variant exceeds the plain forward pass by 68.6, 66.6, and 37.8 percentage points on GDN, GDN-2, and Mamba-2 respectively.

• Full-prompt repetition does not substitute for this: A repetition baseline (JRT) approaches zero accuracy as context grows, providing little improvement at long contexts, whereas NoteBank maintains roughly 70% accuracy on GDN and GDN-2 (Mamba-2 declines more noticeably but retains substantial accuracy at 1M).

• Gains transfer to real-world tasks: On 12 LongBench-v1 tasks with 1.3B backbones, training-free NoteBank improves overall scores from 15.49 to 16.07 on GDN and from 6.79 to 13.43 on Mamba-2. With LoRA read adaptation, overall scores reach 23.03 and 17.10 respectively, both above the Transformer reference of 16.05 (aggregated from Behrouz et al., 2026, whose LongBench models were trained with 16K contexts).

• Works plug-and-play on larger off-the-shelf models: On MemoryAgentBench (2,100-question QA subset), Mamba-2-8B goes from 0.0 to 14.7 (+14.7) and Qwen3.5-9B from 55.7 to 56.8 (+1.1). On the 375-question QA/dialogue subset of LongBench-v2, Qwen3.5-9B goes from 46.1 to 49.3 (+3.2), with YaRN enabled for both conditions.

• Almost free in FLOPs: At 1M tokens, NoteBank adds only 0.0171%–0.0250% of full-model analytical FLOPs (the abstract states ≤0.025%, the introduction states less than 0.026% at layer level), with routing accounting for just 0.50–1.72 GFLOPs. Relative overhead decreases with context length as fixed costs are amortised.

• Compact persistent storage: GDN at 1M tokens adds 24.4 MiB of GPU addresses and helper tensors plus approximately 1.0 MiB of CPU text and metadata; total persistent context storage including backbone states is 74.5 MiB, approximately 1/2,600 of a theoretical BF16 full MHA KV cache (192 GiB) matched in context length, width, and depth.

• Modest measured runtime cost: On an H100 with batch size 1 and 32 generated tokens, latency increases by 10.6%–13.2% across 4K, 16K, 64K, and 1M; at 1M tokens latency rises from 125.11 to 140.64 s, and peak GPU memory including model weights from 2.93 to 3.27 GiB.

• Preserved evidence can remain unretrieved: Oracle note selection with archive, reader, and replay budget fixed raises GDN only from 75.0 to 75.7 at 1M, but Mamba-2 from 40.4 to 61.6, revealing much larger routing headroom for Mamba-2.

• Routing is the main lever, and the native router is competitive: On RULER at 1M, native routing reaches 59.8 (GDN) and 37.6 (Mamba-2), versus random 1.7/0.0, mean-pooled input embeddings 30.8/20.8, and BGE 45.4/34.0; BM25 scores higher at 64.0/52.7, and on LongBench raises Mamba-2 from 14.25 to 18.66 with a fixed LoRA reader. Native routing takes 26 ms for query encoding and routing on GDN at 1M versus 87 ms for CPU BM25 in the authors' implementation.

• Read adaptation and memory access are separate effects: On GDN at 1M RULER, LoRA without NoteBank leaves the score at 1.7 while training-free NoteBank reaches 59.8, and adding LoRA gives 58.4 — so long-range retrieval recovery does not require adaptation. On LongBench, adaptation improves 15.49 → 22.72 without NoteBank and 16.07 → 23.03 with it, a difference of only 0.31 points, meaning adaptation accounts for most of GDN's LongBench improvement in this setting.

Methodology in Plain English

The design rests on a simple separation of duties: the recurrent state keeps doing its normal job, and a separate, cheap archive holds verbatim text.

Write (during context processing). The context is split into consecutive, non-overlapping chunks of a fixed token count, with the recurrent state carried across chunk boundaries as usual. Once a chunk is complete, each token gets a "distinctiveness" score: how far its key vector sits from the chunk's average key, averaged over recurrent layer–head channels. Candidate short windows are scored by summing these values, and the highest-scoring window's central sentence becomes the seed of a note, expanded with neighbouring sentences up to a token limit. Overlapping sentences are excluded from later selections, so the writer stores at most a fixed number of notes per chunk. Crucially, this happens before any question is seen.

Address (cheap index entries). For each stored span and each layer–head channel, the writer takes the position with the highest distinctiveness in that span and stores a compacted version of that position's native query vector as the span's address. No external embedder is involved — the addresses are the model's own features.

Retrieve (when the question arrives). The question is processed from a copy of the retained inference state to produce native query features. Each layer–head channel scores every stored note by the maximum cosine similarity between its address and any question position's query, then votes for its top note. The notes with the most votes are retrieved (earlier-inserted notes win ties), concatenated in original document order, and the temporary question-processing state is thrown away.

Read (replay). Starting from a fresh copy of the retained state, the backbone replays the retrieved text through its ordinary forward computation, then processes the question again to generate the answer. Because this is the model's native forward pass, standard parallel algorithms apply without customisation. Optionally, lightweight LoRA adapters can be trained on the read stage with pretrained weights frozen and adapters disabled during writing and retrieval; the training-free variant requires no such step.

Cost model. Archive size grows linearly with history at a rate set by the write and token budgets; matching is exhaustive but linear; the replay budget stays constant. That is why the arithmetic overhead shrinks relative to backbone work as contexts get longer.

Why This Matters

Impact on research. The paper challenges a widely held assumption — that a linear-attention model's long-context ability is bounded by its recurrent state size — and reframes the problem as one of access rather than capacity. It also supplies a diagnostic toolkit (archived-versus-retrieved coverage, oracle routing, router swaps, frozen-versus-adapted readers) that future memory-augmentation work can reuse. Its training-free framing means results can be reproduced on public checkpoints without new pretraining runs.

Real-world applications:

• Long-document question answering and analysis — legal, scientific, or financial documents far beyond a backbone's trained window, where the specific supporting sentence must survive to question time.
• Agent memory over interaction histories — the MemoryAgentBench results target exactly this: retaining what happened earlier in a long agent trajectory and recalling it on demand.
• Codebase-level reasoning — retrieving relevant code spans from a large repository and replaying them into the model's own computation rather than into a separate reader.
• Multi-turn conversation and personalisation — preserving specific user-stated facts across very long dialogues, since notes are written without knowing the future query.

Industry relevance. The overhead profile is the selling point: negligible analytical FLOPs, 74.5 MiB of persistent context storage at 1M tokens versus a 192 GiB matched KV cache, and 10.6%–13.2% measured latency overhead. For deployments already using hybrid recurrent architectures (the paper cites Griffin, Jamba, Nemotron-H, MiniMax-01, Qwen3-Next/Qwen3.5, and Kimi Linear), this is an interface that can be bolted onto existing checkpoints rather than a retraining project.

Future Directions

• Better routing. Oracle selection shows large headroom, especially for Mamba-2 (40.4 → 61.6 at 1M), and BM25 as an alternative router reaches 64.0/52.7 on RULER where the native router reaches 59.8/37.6. The paper explicitly leaves the retrieval quality–cost trade-off open.
• Selective training of write, route, or read. The authors note that training-free effectiveness establishes compatibility, while selectively training these components may improve evidence preservation, access, and utilisation.
• Closing the gap between archived and retrieved evidence. Archive coverage closely tracks retrieval for GDN, but Mamba-2 leaves substantially more stored values unretrieved; understanding why is an open question.
• Stability and scaling of long-context inference in state-space models. The paper attributes Mamba-2-8B's complete failure on MemoryAgentBench (0.0) to long-context state instability and notes that recurrent-state rescaling was necessary for Mamba-2 to avoid numerical failure, flagging this as a separate problem NoteBank does not solve.

Target Audience

Most useful to researchers and engineers working on efficient long-context language models: linear-attention and state-space model developers, practitioners building retrieval or memory layers for agents and document QA, and anyone evaluating whether a memory-augmentation idea needs retraining. It is also a good read for students interested in how a simple architectural decoupling — storage outside the recurrent state, reading through the native forward pass — can be validated with careful controlled interventions rather than a single benchmark number.

Authors’ abstract

Linear attention promises the efficiency needed for long-context language modelling, yet its fixed-size recurrent state must simultaneously support ongoing computation and preserve information from an ever-growing history. This coupling has reinforced a prevailing view that the state's memory capacity defines the model's long-context capability. We challenge this presumed ceiling: \emph{recurrent computation need not bear the full burden of long-term storage}. We introduce NoteBank, a sparse episodic memory that archives selected text before future queries are known, uses the backbone's native features for writing and addressing, and reads retrieved evidence through its native forward computation via replay. This \emph{training-free, plug-and-play} interface requires no external embedding model, separately trained retriever, or separate neural reader. At one million tokens, it raises Gated DeltaNet's mean accuracy across six RULER retrieval tasks from 1.7% to 70.3%, with no observed significant retrieval degradation from 4K to 1M tokens. This recovery adds $\leqslant$0.025% to analytical FLOPs, with persistent context storage approximately 1/2,600 that of a Transformer's full MHA KV cache at matched width, depth, and precision. Experiments extend from three 1.3B recurrent backbones to Mamba-2-8B and post-trained Qwen3.5-9B, demonstrating gains on LongBench v1/v2 and MemoryAgentBench. Together, these results show that sparse evidence storage and native replay can substantially expand the long-context capability of pretrained linear attention models while retaining the efficiency of recurrent computation.