Memory as a cache: Exact context reuse and deletion by construction

Shengyao Wang; Jiang Liu

cs.AISep 26, 2026

Memory as a cache: Exact context reuse and deletion by construction

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.

Figures & tables

Panel of Figure 1 — Figure 1: SMem. Each block is encoded independently (attention masked to the block); the reader attends within the current block and cross-attends into the memory rows of earlier blocks. Composition, deletion, and path independence hold for every $\theta$ .

Figure 2: (a) Deletion latency vs memory size: re-encoding the remainder (pilot-scale model), contiguous deletion, and index-mapped deletion (flat; its read-time gather cost is in Appendix K ). (b) Serving a fully cached context by block-skip vs cold prefill (410M’s 8k prefill point is an unresolved anomaly).

arm	runs	1	3	7	15	31	63
SMem	R1 s0	.71	.57	.42	.18	.13	.12
SMem	R2, 3 seeds	.46–.74	.24–.65	.13–.43	.01–.13	.01–.11	.01–.11
SMem+RoPE, sat.	R1+R2, 6 runs	.65–.73	.47–.60	.24–.42	.07–.11	.03–.12	.05–.10
SMem	410M R2, 2 seeds	.55–.60	.38–.42	.13–.20	.01–.05	.00–.04	.00–.03
SMem+RoPE, sat.	410M R2 s0	.75	.77	.59	.15	.14	.14
transformer (learned pos.)	R1 s0, R2 3 seeds	.00–.02	.00–.01	.00	.00–.01	.00–.01	.00

Table 1: Exact-match retrieval at

4\times

the trained length by needle distance in blocks (160M unless noted, replication harness,

n{=}384

per cell; ranges span seeds and regimes as listed; sat. = block offsets saturated at the trained range). Cross-scale primary-harness numbers are in the text; full grid in Appendix J .

Figure 3: Needle retrieval at 1.5B: conditional gain by needle distance. At

4\times

the training length the transformer’s signal collapses; SMem’s survives to the farthest distance tested (grid and controls in Appendix J ).

scale	regime	transf.	SMem	$\Delta$
160M	1 (no warmup)	29.72	28.32	$-4.7\%$
160M	2 (warmup)	26.51	26.96	$+1.7\%$
410M	1 (no warmup)	21.26	20.32	$-4.4\%$
410M	2 (warmup)	19.10	19.50	$+2.1\%$
1.5B	1 (no warmup)	14.46	14.26	$-1.4\%$

Table 2: Validation perplexity on FineWeb-Edu (means over seed-paired streams; per-seed statistics in Appendices F / G ). Left: the matched pair (learned positions;

\Delta

= SMem vs transformer) at every scale and regime run. Right: the positional-scheme arms (RoPE transformer; SMem+RoPE composite), to be read against the matching rows on the left; 160M rows are three-seed means, the 410M row a single seed.

Figure 4: Perplexity gap (SMem vs learned-position transformer) by scale and regime; markers are paired means, bars span cross-seed combinations (CIs in Appendix F ; 1.5B is a single Regime-1 pair).

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Validation NLL over training at 1.5B, Regime 1 (single pair; the SMem log ends at

28.8

B tokens, its marker is the final evaluation).

KV Caching

Key-value (KV) caching reuses Transformer states to speed generation while managing memory and data movement.

See all 281 papers

May 11, 2026

Key-Value Means: Transformers with Expandable Block-Recurrent Compressed Memory

We present Key-Value Means ("KVM"), a novel block-recurrence for attention that can accommodate either fixed-size or growing state. Equipping a strong transformer baseline with fixed-size KVM attention layers yields a strong

O(N)

chunked RNN, while adding only an insignificant number of new parameters. We train a transformer with a growable KVM cache and show it performs competitively on long-context tests with only subquadratic prefill time and sublinear state growth. KVM is implementable with standard operations and without custom kernels, and supports chunk-wise parallelizable training and prefill. It provides many of the benefits of both traditional transformers (expandable context memory, chunk-wise parallelizable training and prefill) and linear RNNs in a single unified package. It can be used on every layer, saving KV-cache memory, and allowing a continuous range of choices of prefill time complexity between

O(N)

and

O(N^2)

. It can also be implemented in a hybrid solution in tandem with LRNN layers in place of traditional attention, to supplement the LRNN with improved sublinear memory growth context length usage and long context decoding. We release our code at https://github.com/featherless-ai/KVM-paper and trained models at https://huggingface.co/collections/featherless-ai/kvm-paper under the Apache 2.0 license.

Daniel Goldstein, Eugene Cheah

Featherless AI · Eleuther AI

Transformer Inference Self-Attention KV Caching Long-Context Language Modeling Recurrent Transformers Attention Mechanisms

Ask PDF

Jul 30, 2026

CoMem: Reusing Transformer Depth across Queries with Persistent Intermediate Residuals

Repeated queries over shared documents repeatedly execute the same lower transformer layers. We introduce CoMem, which makes split depth j an explicit reusable-context axis: write one depth-j residual per token, select a bounded chunk set, and resume only layers [j:L). Among document-reuse systems we are aware of, CoMem jointly makes split depth a tunable serving axis and isolates it with a matched j=0 endpoint. On Qwen3-8B, j=12 reduces selected-pack Read from 931.9 to 664.4 ms (1.403x), with a 3.12-point RULER cost (95% CI [2.36, 3.93]); a continuous-prefix oracle recovers the full gap. The resulting depth axis quantifies a quality-latency-storage trade-off; a separate same-adapter, Write-inclusive pipeline is 2.74x faster. Equal-latency raw replay leads by 11.56 points with BM25, directly measuring an applicability boundary of prepaid depth rather than hiding it. CoMem stores 8 KiB/token versus 144 KiB/token for a protocol-aligned same-Qwen3 CacheBlend-style diagnostic; the cohorts and adaptation budgets are not matched. A context-position factorization identifies missing lower-layer document context as the dominant tested multikey error, and a 32-token overlap raises 92.5 to 98.5 without increasing persistent bytes or per-query Read. CoMem opens transformer depth as a measurable, tunable dimension for repeated-query long-context serving.

Hanzuo Liu, Xuan Qi, Chunyu Liu +6

Tsinghua University · Tencent

Memory-Augmented Language Models Long-Context Language Model Inference LLM Inference Acceleration

Ask PDF

May 21, 2026

Tensor Cache: Eviction-conditioned Associative Memory for Transformers

Autoregressive Transformer KV caches grow linearly with context length; sliding-window caching bounds memory but discards evicted tokens entirely, so relevant evidence outside the window becomes inaccessible. We introduce Tensor Cache, a two-level cache that pairs sliding-window softmax attention as a first-level cache (L1) with a fixed-size outer-product fast-weight memory as a second-level cache (L2) fed by KV pairs evicted from the window. Recent tokens remain in exact local attention; evicted pairs are compressed into a per-layer matrix

A

and read by future queries through a single matrix multiplication, exploiting the linear-attention identity

q_t(k_i \otimes v_i)=\langle q_t,k_i\rangle v_i

. A learned scalar gate fuses the L1 and L2 outputs, and per-head decay and write-rate parameters are trained end-to-end. The outer-product memory and the read identity are well-known; our contribution is their use as an L2 cache fed exclusively by sliding-window evictions, plus identifying that the common chunked-mean training shortcut

A\!\leftarrow\!λA\!+\!η(\bar k\!\otimes\!\bar v)

silently introduces

C^2{-}C

spurious cross-token outer products per chunk, and closing the gap with a parallel weighted-sum scan equivalent to per-token writes within float32 epsilon. Across systems scaling, controlled associative recall, long-context language modeling, and memory-capacity diagnostics, Tensor Cache improves the memory--quality frontier over bounded-state baselines.

Kabir Swain, Sijie Han, Daniel Karl I. Weidele +2

Massachusetts Institute of Technology, Cambridge, MA, USA · University of Toronto, Toronto, Canada · IBM Research, Cambridge, MA, USA

Memory-Augmented Neural Networks Self-Attention KV Caching KV-Cache Eviction KV-Cache Management Fast Weights

Ask PDF

Abstract

KV Caching

Explore similar work

Key-Value Means: Transformers with Expandable Block-Recurrent Compressed Memory

CoMem: Reusing Transformer Depth across Queries with Persistent Intermediate Residuals

Tensor Cache: Eviction-conditioned Associative Memory for Transformers