Technical Summary: Structured Memory for Edge Language Models via O(1) SSM State Injection

1. Problem Statement

The practical deployment of Retrieval-Augmented Generation (RAG) on edge devices is currently constrained by two structural limitations of the dominant Transformer architecture:

1. Linear Prefill Cost: In-context RAG requires re-ingesting retrieved context tokens token-by-token during the prefill phase. This imposes a latency and energy cost proportional to the context length ($O(L_{context})$), which is prohibitive on bandwidth-constrained edge hardware.
2. Position-Entangled KV Caches: Transformer Key-Value (KV) caches incorporate positional encodings (rotary or absolute), making them non-reusable across different contexts or positions. Furthermore, the cache size grows linearly with context length, creating a storage bottleneck that scales poorly with corpus size.

State-Space Models (SSMs) offer a theoretical alternative by compressing prior context into a fixed-size, position-agnostic recurrent hidden state. However, existing RAG implementations still rely on token-by-token ingestion, failing to leverage the unique algebraic properties of SSMs that could eliminate the prefill cost entirely.

2. Methodology

The paper introduces PRECOG (Pre-Computed Context Injection) and SMC (Structured Memory Consolidation), two mechanisms that exploit the time-translation invariance of SSMs to bypass context ingestion.

PRECOG: Pre-Computed Context Injection

PRECOG reframes RAG by pre-computing the hidden states of retrieved documents offline and injecting them directly at query time.

• Offline Indexing: Document corpora are partitioned into chunks. An SSM language model processes each chunk to capture the final hidden state ($\bar{h}^{(i)}$) for every layer. These fixed-size states (192 KB for the 1.2B parameter model) are stored on flash storage, paired with lightweight sentence-encoder keys for retrieval.
• Query-Time Injection: Upon receiving a query, the system retrieves the top-$k$ matching chunks via their keys. Instead of re-ingesting the text, the system loads the corresponding pre-computed hidden states and injects them as the initial recurrent state ($\bar{h}_{init}$) of the model. The model then processes only the query tokens.
• Theoretical Guarantee (Theorem 1): The paper proves that for any initial state $h_0$, context $c$, and query $q$, the state trajectory $S(h_0, c \oplus q)$ is algebraically identical to $S(S(h_0, c), q)$. Because SSM updates depend only on the current state and token (not absolute position), injecting a pre-computed state is mathematically equivalent to having processed the context tokens sequentially. This guarantees that PRECOG matches the answer quality of in-context RAG exactly (modulo quantization noise).

SMC: Structured Memory Consolidation

SMC extends the state-injection mechanism to persistent, on-device memory for accumulating user preferences, interaction history, and logs.

• Hierarchical Clustering: Interaction chunks are routed into cognitive-domain clusters (e.g., Emotional, Temporal, Spatial, Factual) and further into sub-clusters based on recall patterns (episodic vs. semantic).
• Fidelity-Storage Dial: A parameter $K$ controls the retention of states per chunk:
  • $K=N_c$: Lossless episodic memory (all states retained).
  • $K=1$: Semantic memory (only the final state retained, identical to PRECOG chunks).
• O(1) Session Initialization: At the start of a session, the system identifies the dominant domain and injects the corresponding consolidated semantic state directly into the SSM. This initializes the model with a history of interactions without re-ingesting any tokens, maintaining $O(1)$ latency regardless of accumulated history length.

3. Key Contributions

1. Algorithmic Refactoring of RAG: The paper demonstrates that context ingestion cost can be reduced from $O(L_{context})$ to $O(1)$ by injecting pre-computed SSM hidden states, a mechanism structurally impossible for Transformers due to position entanglement.
2. Exact Equivalence Proof: Theorem 1 provides a formal guarantee that PRECOG produces identical state trajectories to in-context RAG, ensuring that the speedup does not come at the cost of answer quality.
3. Unified Memory Substrate: The same state-injection primitive unifies static corpus retrieval (PRECOG) and dynamic persistent memory (SMC), allowing edge models to fuse retrieved knowledge and user history in a single query step.
4. Edge Deployment Feasibility: By instantiating these methods on TENNs-LLM (a 1.2B-parameter gated-SSM with a 192 KB total hidden state), the authors demonstrate that state-level retrieval is practical on neuromorphic edge hardware, where the small state footprint fits within on-chip memory constraints.

4. Results

The system was evaluated on the SQuAD v1.1 development set using a 1.2B parameter TENNs-LLM backbone.

• Quality: PRECOG (top-1 injection) achieved an Exact Match (EM) of 58.0 and F1 of 73.4, matching the in-context RAG baseline (EM 58.2, F1 73.6) within the bounds of FP16 quantization noise. This empirically validates the theoretical equivalence.
• Latency: On edge hardware (simulated neuromorphic processor), PRECOG reduces the prefill latency from ~27 seconds (required for in-context ingestion of a 512-token chunk) to <6 ms for the retrieval and injection overhead itself. The total Time to First Token (TTFT), including the time to ingest the query tokens, is approximately 585 ms. The ~4,500× speedup cited in the paper specifically refers to the elimination of the ~27 s context-ingestion phase, not the total query latency.
• Storage Efficiency: For a 512-token chunk, the PRECOG state requires 192 KB, whereas a comparable Transformer KV cache requires ~16 MB (an 85× reduction). For a 10,000-chunk corpus, PRECOG requires 1.9 GB of storage versus 160 GB for a hypothetical KV-cache approach, making the latter infeasible on consumer devices.
• Ablations:
  • Injection Depth: Injecting only the bottom 12 layers (96 KB) preserved 99% of the full-layer F1 score, suggesting potential for further storage optimization.
  • Chunk Length: Performance matches in-context RAG for chunks up to ~600 tokens (the model's effective memory horizon). Beyond this, PRECOG degrades gracefully as the SSM naturally forgets distant tokens, mirroring the behavior of in-context RAG with the same backbone.

5. Significance and Claims

The paper claims that PRECOG and SMC represent a fundamental shift in how retrieval and memory are handled for edge language models. The significance lies in:

• Architectural Exclusivity: The mechanism relies on the time-translation invariance of SSMs, making it "architecturally impossible" for Transformer KV-caches.
• Edge Viability: It solves the "prefill bottleneck" that currently prevents interactive RAG on bandwidth-constrained devices, enabling a shift from cloud-dependent inference to on-device, low-latency interaction.
• Design Principle: The authors posit that as SSMs close the quality gap with Transformers, their structural properties (fixed-size, position-agnostic states) will enable a new class of retrieval, caching, and memory-extension algorithms that are inaccessible to attention-based architectures.

The paper remains modest regarding limitations, acknowledging that PRECOG inherits the SSM's finite memory horizon (cannot recover information the backbone would forget) and that the top-$k$ composition heuristic (for multi-chunk retrieval) is an approximation without the exactness guarantee of top-1 injection. Additionally, the neuromorphic deployment numbers are based on FPGA measurements and 12 nm power simulations, pending physical tape-out verification.