Structured Memory for Edge Language Models: Persistent Context and Corpus Retrieval via O(1) SSM State Injection

AI summary

Overview

Research area: Efficient language-model inference — retrieval-augmented generation (RAG), state-space models (SSMs), and persistent memory for edge deployment.

Technical level: Advanced. The paper includes a formal theorem and proof, recurrence-level notation, and storage/bandwidth analysis, though the core idea is explained clearly enough for readers with a machine-learning background.

Scope: The paper introduces PRECOG, a mechanism that pre-encodes document corpora as SSM hidden states and injects a retrieved state directly at query time to remove context-ingestion cost at prefill, plus SMC, a hierarchical persistent-memory scheme built on the same state-injection primitive, both demonstrated on the 1.2B-parameter TENNs-LLM.

What This Paper Is About

Conventional retrieval-augmented generation forces the model to re-read every retrieved document token-by-token before it can answer, and Transformer backbones additionally accumulate a key-value (KV) cache that grows with each generated token. The authors' problem is that on bandwidth-constrained edge hardware this prefill cost dominates total query latency — the paper reports roughly 27 seconds before the first response token for a single 512-token chunk. Their goal is to eliminate context re-ingestion entirely by exploiting a property unique to state-space models: a fixed-size, position-agnostic hidden state that completely summarizes everything the model has read, so it can be saved and reloaded as an initial condition.

Key Contributions

1. PRECOG (Pre-Computed Context Injection): a retrieval mechanism that pre-encodes document chunks offline as SSM hidden states and injects the best-matching state at query time, reducing context-ingestion cost at prefill from O(L_context) to O(1) in retrieved-context length.

2. Theorem 1 (PRECOG–RAG equivalence): a formal proof that PRECOG and in-context RAG produce identical state trajectories under autoregressive SSM dynamics, making "PRECOG matches in-context RAG" a mathematical guarantee rather than an empirical hope (with a sufficient-statistic lemma and an FP16 quantization bound of roughly 2^-10 · ||h|| per element).

3. SMC (Structured Memory Consolidation): a hierarchical persistent memory that partitions stored states into cognitive-domain clusters, offers a fidelity-vs-storage dial (K = N_c, K = N_c/k, or K = 1), consolidates short-term episodic states into long-term semantic memory via exponential moving average, and provides O(1) session initialization.

4. An edge deployment demonstration on TENNs-LLM: a 1.2B-parameter gated-SSM language model with a 192 KB total hidden state, evaluated on SQuAD v1.1 and deployed on a neuromorphic edge processor (Appendix H), together with storage and bandwidth comparisons against Transformer KV-cache RAG.

Main Findings

• Prefill latency collapse: PRECOG reduces prefill latency from ~27 s to <6 ms on the edge deployment target — a ~4500× speedup that "crosses the threshold from unusable to interactive." Figure 3 reports time-to-first-token of 585 ms at 19 tok/s with UFS 4.0 storage.

• Answer quality preserved: on a randomly sampled 1,000-question subset of SQuAD v1.1 dev using the official evaluation script, in-context RAG scored EM 58.2 / F1 73.6; PRECOG top-1 scored EM 58.0 / F1 73.4. The gap of 0.2 EM and 0.2 F1 is within the predicted FP16 quantization bound.

• Top-k composition is a heuristic, not exact: PRECOG top-3 (softmax-weighted state composition) scored EM 56.4 / F1 71.8, a 1.6 F1 drop, which the authors attribute to the non-exactness of the composition rule rather than the state-injection mechanism.

• Small query-time overhead: PRECOG adds ~6 ms total versus zero-context generation, decomposed as ~5 ms for the sentence-encoder forward pass over the query on CPU, <1 ms for FAISS top-k search in DRAM, ~50 µs for the flash-to-DRAM transfer of the 192 KB state on UFS 4.0 at 4.2 GB/s, and <1 ms for 24 vector copies into the recurrent buffer.

• Storage trade-off in PRECOG's favour: a single 512-token chunk costs ~16 MB for a hypothetical Transformer KV-cache versus 192 KB for PRECOG — an 85× premium per chunk. A 10,000-chunk corpus consumes 1.9 GB for PRECOG versus ~160 GB for KV-cache storage; keys alone are ~3.8 MB for 10K chunks. Figure 2 reports the Llama-3.2-1B KV cache growing at 32 KB/token (16 layers, GQA, FP16) with crossover at L = 6 tokens.

• Injection is exact, not approximate: because the SSM update map depends only on (h, x) with no explicit position index, a state pre-computed over a chunk is identical to the state the model would reach by running that chunk at the start of a query.

• Transformers cannot do this: under rotary position encoding the KV cache is position-entangled — a cache pre-computed at positions 0…L−1 is invalid at positions τ…τ+L−1 for τ ≠ 0 — and recomputing it negates any prefill savings.

• Memory horizon is inherited, not lost: the closed-form unrolling shows context token c_t contributes β(c_t) · ∏ α(c_s) to the final state, decaying geometrically. Empirically, TENNs-LLM's effective memory length is dominated by tokens within the most recent ~256 positions of a 512-token chunk.

• Persistent memory is compact: with M = 5 cognitive domains and J_max = 4 sub-clusters, total SMC semantic-memory footprint is under 4 MB — three orders of magnitude smaller than the episodic store at typical chunk volumes.

Methodology in Plain English

The authors start from an architectural observation rather than a training trick. In a state-space model, the model's running memory is a fixed-size block of numbers that gets updated token by token, and that update only depends on the current token and the current memory — never on where in the sequence the token sits. That means the memory block after reading a document is a complete summary of that document, and copying it into the model at query time is mathematically the same as having made the model read the document.

Building on this, they split each document corpus into chunks (512 tokens in their experiments), run each chunk through the model once offline in inference mode, capture the resulting 24-layer hidden state, and store it at FP16 (192 KB per entry, regardless of chunk length). Each chunk also gets a compact 384-dimensional key from a sentence encoder (all-MiniLM-L12-v2). Keys live in DRAM in a FAISS index for nearest-neighbour search; the much larger states stay on flash and are demand-loaded only when retrieved.

At query time the system encodes the query with the same sentence encoder, retrieves the top-k chunks by cosine similarity (k = 3 by default), and writes the top-1 state directly into the model's recurrent buffer as the initial condition. The model then processes only the query tokens and samples via top-p sampling (p = 0.9). For multi-chunk fusion, the top-k states are combined with a softmax-weighted average — a heuristic they explicitly flag as no longer exact.

The persistent-memory component (SMC) reuses the same machinery. Conversation chunks carry metadata headers (timestamp, speaker, optional GPS), are routed into one of five cognitive-domain clusters and then into typed sub-clusters, and each sub-cluster maintains a single semantic state updated by exponential moving average. A single integer parameter K controls how many states per chunk are kept: all of them for lossless episodic recall, every k-th for a tunable middle ground, or only the final state for a semantic summary. Session initialization then writes the relevant semantic state straight into the model with no context tokens ingested.

Validation is a controlled comparison on SQuAD v1.1: identical TENNs-LLM weights, identical FP16 inference precision, differing only in retrieval and state-injection logic, across in-context RAG, PRECOG top-1, and PRECOG top-3. Three further ablations appear in the appendices — injection depth on SQuAD v1.1, top-k composition on HotpotQA-distractor, and chunk-length sensitivity on Natural Questions — with no PRECOG-specific training in any condition.

Why This Matters

Impact on research. The paper reframes a retrieval efficiency question as an algebraic one. Prior work on SSM state caching and composition — State Soup and PICASO, both cited — requires training-time modifications: PICASO trains a learned composition function and State Soup learns context-mixing weights. PRECOG claims to need none of that, operating on a stock SSM fine-tuned only for next-token prediction, with the equivalence proven rather than learned. It also identifies a regime the authors say prior work does not analyse: edge deployment where context-ingestion latency dominates, versus data-center GPUs where retrieval methods are typically evaluated.

Real-world applications (drawn from the paper's framing):

• On-device assistants that hold persistent user preferences across sessions without re-ingesting prior conversation at session start.
• Appliances and devices that accumulate operational logs and interaction history and must answer queries within a bounded memory budget.
• Edge question answering over a static document corpus — the 10,000-chunk, 1.9 GB state store described as "tractable on a phone" where the hypothetical 160 GB KV-cache store is infeasible.
• Location- and time-aware conversational memory, since SMC chunk headers can bake in timestamps, speaker identifiers, and optional GPS.

Industry relevance. The work is affiliated with BrainChip Inc. and targets a neuromorphic edge processor, and the PRECOG mechanism is the subject of a related pending patent application. The design point is explicit: PRECOG is the right choice when ingestion latency is the binding constraint, not when storage is, since its storage footprint is ~200× that of raw-text RAG.

Future Directions

• Tape-out verification. The neuromorphic deployment numbers combine measured FPGA throughput with simulated 12 nm power figures; the authors state that a tape-out verification is left to future work.
• Restoring exactness beyond top-1. Theorem 1 is exact only for single-chunk top-1 injection. The top-k softmax-weighted composition used at retrieval time is a heuristic with no analogous guarantee, and empirical degradation is observed as k grows — an exact multi-chunk composition rule is an open problem.
• Memory horizon limits. PRECOG inherits the backbone's finite effective memory length exactly, so it cannot recover information the model itself would forget at matched chunk length. Whether longer effective memory can be engineered into the recurrence without breaking the time-translation invariance that makes injection exact is unaddressed.
• Generalization beyond TENNs-LLM. The authors state the time-translation-invariance assumption holds for any selective-SSM backbone, including Mamba and Mamba-2 and related linear-recurrent architectures, but the reported results are all on TENNs-LLM.

Target Audience

Researchers and engineers working on efficient inference, retrieval-augmented generation, state-space and linear-recurrent architectures, and on-device or edge AI deployment. It is most valuable to readers who already understand the mechanics of SSM recurrences and KV caches and want to see how a structural property of recurrence can be turned into an algorithmic advantage — and to practitioners weighing the storage-versus-latency trade-off when designing retrieval for hardware with tight memory bandwidth. Readers without that background can follow the problem framing, the SQuAD comparison, and the deployment numbers, but will need to work through Appendix A and the recurrence notation to assess the central claim.

Authors’ abstract

Retrieval-augmented generation (RAG) imposes a prefill cost proportional to retrieved context length, and -- with Transformer backbones -- a KV-cache that grows with each generated token. State-Space Models (SSMs) avoid the second cost by construction; we eliminate the first, collapsing prefill from $O(L_{context})$ to $O(1)$ per query. We introduce PRECOG (Pre-Computed Context Injection), a retrieval mechanism that exploits a property unique to SSMs: the fixed-size, position-agnostic recurrent hidden state is a complete summary of everything the model has read. PRECOG pre-encodes document corpora offline as SSM hidden states and injects the best-matching state directly at query time, bypassing in-context re-ingestion entirely. The same state-injection mechanism enables SMC (Structured Memory Consolidation): a hierarchical persistent memory with cognitive-domain clustering, an adjustable fidelity-vs-storage dial, and $O(1)$ session initialization, which consolidates short-term episodic states into long-term semantic memory and fuses both with retrieved corpus states at query time. We demonstrate the system on TENNs-LLM, a 1.2B-parameter gated-SSM language model with a 192 KB hidden state. PRECOG matches in-context RAG answer quality, reducing prefill latency from $\sim$27 s to $&lt;$6 ms on edge hardware -- a $\sim$4500$\times$ speedup that crosses the threshold from unusable to interactive. The mechanism is architecturally impossible for Transformer KV-caches, which are position-entangled and grow linearly with context length.