Transformer memory restructured to enable exact reuse and deletion
Memory as a cache: Exact context reuse and deletion by construction
Artificial IntelligenceMachine Learning
Summary
Current transformer models combine every token’s memory with everything before it, making it hard to reuse parts or remove sections without recalculating everything after. The authors introduce SMem, which stores memory in separate blocks that can be combined exactly and updated efficiently. This allows removing or reusing parts of the memory quickly without affecting the rest, and keeps prediction quality close to standard models. SMem also works well with common positional encoding methods to maintain accuracy.
What this means in practice
- •For natural language processing engineers: Accelerate long-context text generation by updating memory blocks without full recomputation, enabling faster responses in applications like chatbots and summarization.
- •For database system developers: Implement efficient context caching that allows exact deletion and reuse of stored blocks, reducing latency for dynamic query-driven content generation.
Authors
Shengyao Wang, Jiang Liu
Abstract
The KV cache of a transformer entangles every token's representation with its entire prefix: a passage encoded once cannot be reused under a different prefix or removed without recomputing everything after it, so exact cache reuse is limited to shared prefixes. We present SMem, an architecture whose context representation is a cache by construction. A block-local encoder maps each block to memory rows independently of other blocks, and a reader conditions generation on their union through cross-attention. For every parameter setting, memory composes exactly at fixed block indices, deleting a block is an exact $O(b)$ update for $b$-token blocks, and the memory state is independent of the edit path. At $4\times$ the training context, under the shared recipe, SMem retrieves planted needles beyond any trained-length window (exact match 0.14-0.28 at distances of 31 and 63 blocks), where learned-position, RoPE, and Block-Attention-style transformers all score at most 0.02. A fully cached context is served by computing one block alone at a near-constant 3.1-6.2 ms, whereas cold prefill grows with context; batched decode stores 34-38% fewer KV rows and runs 1.4-1.7$\times$ faster when bandwidth-bound; and deletion beats suffix recomputation by 8.5$\times$ at 512 blocks and 452$\times$ at 4096 blocks (32-256$\times$ the trained length, probing the cost model rather than a served regime). The cost is a perplexity gap of -4.7% to +2.8% (negative favors SMem) against a parameter-matched transformer with the same positional scheme, at 160M-1.5B on FineWeb-Edu across two recipes and a learning-rate search. SMem also composes with RoPE: at 160M and 410M the composite matches or leads the matched transformer and closes 29-59% of SMem's gap to a RoPE transformer. Dropping prefix entanglement thus keeps perplexity comparable while making the cache exactly composable and editable.