Summary (Overview)

  • Key Contribution: The paper introduces NoteBank, a training-free, plug-and-play sparse episodic memory system that decouples evidence preservation from recurrent computation in linear attention models, challenging the prevailing view that fixed-size recurrent states impose a fundamental ceiling on long-context capability.
  • Major Result: At one million tokens, NoteBank raises Gated DeltaNet's mean accuracy across six RULER retrieval tasks from 1.7% to 70.3%, with no significant retrieval degradation from 4K to 1M tokens.
  • Efficiency: The approach adds ≤0.025% to analytical FLOPs and requires persistent context storage approximately 1/2,600 that of a Transformer's full MHA KV cache at matched width, depth, and precision.
  • Generality: Benefits extend across multiple architectures—from 1.3B research models (GDN, GDN-2, Mamba-2) to larger post-trained models (Mamba-2-8B, Qwen3.5-9B)—with improvements on LongBench-v1/v2 and MemoryAgentBench.
  • Key Insight: Poor long-range retention through a recurrent state need not imply an inability to use distant evidence once that evidence is independently preserved and made accessible.

Introduction and Theoretical Foundation

Background and Motivation

Long-context capability enables language models to reason over extensive documents, conversations, and interaction histories. However, full softmax attention requires quadratic computation and a key-value cache that grows linearly with sequence length. Linear attention offers an alternative by compressing history into a fixed-size memory state updated as each new token arrives.

The Core Challenge

The paper identifies a fundamental tension in linear attention models: a single recurrent state must serve two competing roles:

  1. Ongoing computation: The state must continually transform and compress history
  2. Long-term memory: Specific past evidence must remain recoverable even when relevance emerges only later

As history grows while the state remains fixed in size, continual compression risks irreversibly discarding evidence before its future relevance is known. Prior work establishes finite-state limits on exact copying (Jelassi et al., 2024), identifies recency bias and over-smoothing in selective state-space models (Wang et al., 2025b), and documents substantial associative-recall gaps in efficient attention alternatives (Arora et al., 2024b).

Theoretical Foundation

The paper builds on the mathematical framework of linear attention. For a single layer and head, the backbone's feature transformations produce query and key vectors qt,kt∈Rdk\boldsymbol{q}_t, \boldsymbol{k}_t \in \mathbb{R}^{d_k} and a value vector vt∈Rdv\boldsymbol{v}_t \in \mathbb{R}^{d_v} from input representations at token position tt. The recurrent memory update is:

St=St−1At+βtvtkt⊤,ot=Stqt,(1)\boldsymbol{S}_t = \boldsymbol{S}_{t-1} \boldsymbol{A}_t + \beta_t \boldsymbol{v}_t \boldsymbol{k}_t^{\top}, \quad \boldsymbol{o}_t = \boldsymbol{S}_t \boldsymbol{q}_t, \tag{1}

where St∈Rdv×dkS_t \in \mathbb{R}^{d_v \times d_k} serves as a compressed memory of the prefix, At∈Rdk×dk\boldsymbol{A}_t \in \mathbb{R}^{d_k \times d_k} controls memory retention and transformation, and βt\beta_t controls write strength. Standard additive linear attention sets At=I\boldsymbol{A}_t = \boldsymbol{I} and βt=1\beta_t = 1; gated and delta-rule variants modify these updates.

The problem setup considers a pretrained model Mϕ\mathcal{M}_\phi processing context X=(x1,…,xT)\boldsymbol{X} = (x_1, \dots, x_T) before receiving a question Q\boldsymbol{Q}. The model reads XX causally, retaining an inference state ST\mathcal{S}_T consisting of recurrent states and auxiliary caches. When QQ arrives, it generates a response Y∼pϕ(⋅∣ST,Q)Y \sim p_\phi(\cdot \mid S_T, Q). Access to preceding context is limited to what remains in STS_T, meaning evidence must be preserved before the model knows which information the question will require.

Methodology

NoteBank Architecture

NoteBank maintains a sparse episodic archive alongside the backbone's unchanged state updates. After processing tt context tokens, the archive is:

NBt={(zi,Ai)}i=1Nt,Ai={ρiℓ,h}(ℓ,h)∈L×H,(2)\mathcal{NB}_t = \{(\mathbf{z}_i, \mathcal{A}_i)\}_{i=1}^{N_t}, \qquad \mathcal{A}_i = \{\boldsymbol{\rho}_i^{\ell,h}\}_{(\ell,h) \in \mathcal{L} \times \mathcal{H}}, \tag{2}

where NtN_t is the number of stored notes, zi\mathbf{z}_i contains a span's original token IDs, and each address ρiℓ,h\rho_i^{\ell,h} is a compact vector derived from native query features within that span.

Write (Top-W Selection)

The context is partitioned into consecutive, non-overlapping chunks of CC tokens. Once a chunk accumulates CC tokens, the writer scores tokens by native-key distinctiveness:

dtℓ,h=1−⟨k^tℓ,h,μ‾C(t)ℓ,h⟩,dt=1∣L×H∣∑(ℓ,h)∈L×Hdtℓ,h,(3)d_t^{\ell,h} = 1 - \left\langle \hat{\boldsymbol{k}}_t^{\ell,h}, \overline{\boldsymbol{\mu}}_{\mathcal{C}(t)}^{\ell,h} \right\rangle, \qquad d_t = \frac{1}{|\mathcal{L} \times \mathcal{H}|} \sum_{(\ell,h) \in \mathcal{L} \times \mathcal{H}} d_t^{\ell,h}, \tag{3}

where C(t)\mathcal{C}(t) contains positions in the chunk containing tt, and μ‾C(t)ℓ,h\overline{\boldsymbol{\mu}}_{\mathcal{C}(t)}^{\ell,h} is the normalised mean of keys in that chunk. Hats denote ℓ2\ell_2 normalisation.

For each selected text segment with token positions IiI_i, the address is constructed by selecting the token with highest key-distinctiveness score per layer-head channel:

ti∗,ℓ,h=arg⁡max⁡t∈Ii dtℓ,h,ρiℓ,h=compact(qti∗,ℓ,hℓ,h)∈Rdρ.(4)t_i^{*,\ell,h} = \underset{t \in I_i}{\arg\max} \, d_t^{\ell,h}, \qquad \boldsymbol{\rho}_i^{\ell,h} = \text{compact}\left(\boldsymbol{q}_{t_i^{*,\ell,h}}^{\ell,h}\right) \in \mathbb{R}^{d_\rho}. \tag{4}

Addressing (Top-R Retrieval)

When the question arrives, it is processed from a copy of the retained inference state. Each layer-head channel scores every stored note by maximum cosine similarity and votes for its highest-scoring note:

siℓ,h=max⁡1≤τ≤∣Q∣⟨u^τℓ,h,ρ^iℓ,h⟩,Vi=∑(ℓ,h)∈L×H1[i=arg⁡max⁡j sjℓ,h].(5)s_i^{\ell,h} = \max_{1 \leq \tau \leq |\boldsymbol{Q}|} \left\langle \hat{\boldsymbol{u}}_\tau^{\ell,h}, \hat{\boldsymbol{\rho}}_i^{\ell,h} \right\rangle, \qquad V_i = \sum_{(\ell,h) \in \mathcal{L} \times \mathcal{H}} \mathbf{1}\left[i = \underset{j}{\arg\max} \, s_j^{\ell,h}\right]. \tag{5}

Up to RR notes with the most votes are retrieved, with earlier-inserted notes preferred in ties.

Readout with Native Replay

The backbone reads retrieved text through its native forward computation:

SR=Fϕ(Sread,R),Y∼pϕ(⋅∣SR,Q).(6)\mathcal{S}_{\boldsymbol{R}} = F_\phi(\mathcal{S}_{\text{read}}, \boldsymbol{R}), \quad \boldsymbol{Y} \sim p_\phi(\cdot | \mathcal{S}_{\boldsymbol{R}}, \boldsymbol{Q}). \tag{6}

Replay introduces at most R×BR \times B tokens and traverses the full backbone. Optionally, lightweight LoRA adapters can be trained while keeping pretrained weights frozen.

Memory and Computational Overhead

Memory: With J=∣L×H∣J = |\mathcal{L} \times \mathcal{H}| indexed layer-head channels and at most WW notes per completed CC-token chunk:

NT≤W⌊T/C⌋,StorageNB=O(NT(B+Jdρ)).(7)N_T \leq W \lfloor T/C \rfloor, \quad \text{Storage}_{\mathcal{NB}} = O(N_T(B + J d_\rho)). \tag{7}

Computation: Writing adds O(T)O(T) work; addressing requires one additional forward pass over the question plus exhaustive matching at cost O(∣Q∣NTJdρ)O(|\boldsymbol{Q}| N_T J d_\rho); readout adds at most R×BR \times B replay tokens.

Empirical Validation / Results

Training-Free Long-Range Retrieval

The paper evaluates six RULER-NIAH tasks (three single-needle and three multi-key/multi-query/multi-value tasks) across document lengths from 4K to 1M tokens, with 100 samples per task.

Key Results (Figure 3):

  • At 1M tokens, NB-TF exceeds Plain by 68.6, 66.6, and 37.8 percentage points on GDN, GDN-2, and Mamba-2, respectively
  • Plain and JRT (full-prompt repetition) approach zero accuracy as context grows
  • NB-TF maintains roughly 70% accuracy on GDN and GDN-2 at 1M tokens

Real-World Tasks (LongBench)

Table 1: LongBench-v1 Results (1.3B backbones)

BackboneMethodSingle-doc QAMulti-doc QASummarisationFew-shotOverall
Transformer†Plain13.4018.978.4323.4016.05
GDNPlain6.887.0213.9934.0815.49
GDNJRT7.726.6314.7933.5615.67
GDNNB-TF6.667.3313.5236.7616.07
GDNNB-LoRA19.0620.4811.3741.2223.03
Mamba-2Plain2.222.1011.2711.566.79
Mamba-2JRT2.391.899.218.345.12
Mamba-2NB-TF5.985.3616.5125.8813.43
Mamba-2NB-LoRA17.8415.547.4627.5217.10

†Reference results from Behrouz et al. (2026), trained with 16K contexts.

Plug-and-Play on Off-the-Shelf LLMs

Table 2: Results on Larger Models Without Training

BenchmarkBackbonePlain+ NB-TFΔ
MemoryAgentBenchMamba-2-8B0.014.7+14.7
MemoryAgentBenchQwen3.5-9B55.756.8+1.1
LongBench-v2Qwen3.5-9B46.149.3+3.2

Efficiency Analysis

  • FLOPs: At 1M tokens, NoteBank adds only 0.0171%–0.0250% to full-model analytical FLOPs
  • Storage: At 1M tokens, GDN adds 24.4 MiB of GPU addresses plus ~1.0 MiB CPU text/metadata; total persistent context storage is 74.5 MiB (~1/2,600 of a matched MHA KV cache)
  • Latency: NoteBank increases latency by 10.6%–13.2% (at 1M tokens: 125.11 → 140.64 s)
  • Memory: Peak GPU memory rises from 2.93 to 3.27 GiB at 1M tokens

Routing Strategy Comparison

Table 3: Effect of Routing Strategy (LoRA readers fixed)

Routing SignalAuxiliary RetrieverRULER @ 1M (GDN)RULER @ 1M (Mamba-2)LongBench (GDN)LongBench (Mamba-2)
RandomNone1.70.021.9613.83
Input embeddingNone30.820.822.4016.70
Native (NB)None59.837.623.0314.25
BM25Lexical index64.052.723.5018.66
BGEExternal encoder45.434.023.7717.24

Memory Access vs. Read Adaptation

Table 4: Separating Memory Access from Read Adaptation (GDN)

Memory AccessRULER @ 1M (Frozen)RULER @ 1M (LoRA)LongBench (Frozen)LongBench (LoRA)
w/o NoteBank1.71.715.4922.72
w/ NoteBank59.858.416.0723.03

Key finding: Long-range retrieval recovery does not require read adaptation (59.8 vs. 58.4 with LoRA on RULER), while adaptation accounts for most LongBench gains on GDN.

Theoretical and Practical Implications

Challenging the Fixed-State Ceiling

The paper's central theoretical contribution is demonstrating that the long-context capability of pretrained recurrent models can substantially exceed what their conventional fixed-state memory interface exposes. This reframes the debate about linear attention limitations: poor long-range retention through the recurrent state need not imply an inability to use distant evidence.

Decoupling Computation from Memory

NoteBank establishes a principled separation of concerns:

  • Recurrent state handles ongoing computation
  • Sparse episodic archive handles evidence preservation
  • Native replay handles evidence utilisation

This decoupling is achieved without modifying backbone weights or architecture, requiring no external embedding model, separately trained retriever, or separate neural reader.

Practical Implications

  1. Plug-and-play deployment: NoteBank works with existing pretrained checkpoints, making it immediately applicable to deployed models
  2. Near-zero overhead: The negligible compute and storage costs make it feasible for production long-context applications
  3. Architecture-agnostic: Benefits extend across gated delta networks, state-space models, and hybrid architectures
  4. Routing headroom: The analysis reveals substantial room for improvement in retrieval quality, suggesting future work on better addressing mechanisms

Limitations and Open Questions

  • Routing headroom exists, particularly for Mamba-2 (oracle routing raises score from 40.4 to 61.6 at 1M tokens)
  • Read adaptation benefits are task-dependent (NB-LoRA underperforms NB-TF on summarisation)
  • The paper notes that answer containment in oracle notes is a proxy diagnostic, not a guaranteed performance upper bound

Conclusion

The paper introduces NoteBank, a training-free, plug-and-play sparse episodic memory that separates evidence preservation from recurrent computation in linear attention models. By archiving selected text spans before questions are known, addressing them with native features, and reading retrieved evidence through native forward computation via replay, NoteBank recovers substantial long-context capability at negligible cost.

Key takeaways:

  1. The fixed-state ceiling is not fundamental: Pretrained recurrent models can use distant evidence far beyond what their recurrent states expose
  2. Native computation suffices: The backbone itself can write, address, and read memory without external components
  3. Near-free capability recovery: 37.8–68.6 percentage-point gains at 1M-token retrieval with ≤0.025% FLOPs overhead and 1/2,600 the storage of matched MHA KV cache
  4. Broad applicability: Works across architectures and scales, from 1.3B research models to 8-9B post-trained models

Future directions identified by the authors include improving writing, routing, and reading mechanisms, with selective training of these components potentially yielding further gains in evidence preservation, access, and utilisation.

Related papers