PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Shengyao Wang, Jiang Liu

arXiv:2609.32395v1cs.AIcs.LG

TL;DR

Transformer KV states are prefix-dependent, limiting exact reuse and making removal require recomputing later states. SMem encodes blocks independently and reads their union; at 4× training length, it retrieves beyond trained-length windows where tested transformer variants score at most 0.02.

  • Problem

    Transformer KV representations depend on the entire prefix, limiting exact reuse to contexts sharing a verbatim prefix and making removal invalidate later states.

    Source: §1 · pp. 1–2

  • Method

    SMem independently encodes each context block into memory rows and conditions generation on their union through cross-attention.

    Source: §1 · p. 2, §2 · p. 3

  • Results

    At 4× the training length, SMem reaches exact-match 0.14–0.28 at distances of 31 and 63 blocks, versus at most 0.02 for tested transformer variants.

    Source: §1 · p. 3

  • Takeaways & Limitations

    SMem provides exact composition and deletion with comparable perplexity, while enabling cached-context serving and retrieval beyond trained length in the reported settings.

    Source: §7 · p. 10

    SMem is a pretraining-time choice, not a retrofit; the composite has limited seed evidence at 410M and none at 1.5B, and far-range retrieval under the tuned recipe is near the floor.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

SMem replaces prefix-entangled context caches with independently encoded block memories, enabling exact composition, deletion, and path-independent edits. The paper reports retrieval beyond trained-length windows, faster cache operations, and perplexity comparable to matched transformers.

1 Introduction

  • Transformer KV states depend on the entire prefix, restricting exact reuse to verbatim shared prefixes and making removal require suffix recomputation.The cited production study finds 8% of requests and 18% of prefill tokens covered by prefix reuse.Source: §1 · pp. 1–2
  • SMem encodes each block independently into memory rows, then conditions generation on their union through reader cross-attention.The block-local encoder makes each block’s memory independent of other blocks.Source: §1 · p. 2, Figure 1 · p. 3
  • At fixed block indices, blocks compose exactly, deletion costs O(b) for b-token blocks, and the resulting memory is independent of the edit path for every parameter setting.These guarantees follow from the architecture rather than training; a transformer instead must re-encode at O(n) rather than O(b).Source: §1 · pp. 2–3
  • At 4× the training length, SMem reaches exact-match 0.14–0.28 at distances of 31 and 63 blocks, versus at most 0.02 for tested transformer variants.These retrieval results are under the shared recipe and beyond any trained-length window.Source: §1 · p. 3
  • SMem serves cached contexts at near-constant cost and delivers bandwidth-bound decode 1.4–1.7× faster on 34–38% fewer KV rows, with delete-then-generate 8.5–452× faster than suffix recomputation.Source: §1 · p. 3
  • At 160M–1.5B, the perplexity gap to a parameter-matched transformer ranges from −4.7% to +2.8% across two recipes and a learning-rate search.The SMem+RoPE composite matches or leads the matched transformer at 160M and 410M.Source: §1 · p. 3

2 Architecture and guarantees

SMem encodes blocks independently into memory rows and lets a block-local reader condition generation on their union. At fixed block indices, memories compose exactly, and deleting a block is exact and independent of the edit path.

Model

  • SMem encodes each block independently into b memory rows, then uses a reader with block-local self-attention and cross-attention to earlier blocks.The memory is not mutated during reading; beyond the trained block count, older blocks are folded into slot 1.Source: §2 · pp. 3–4

Structural guarantees (∀θ)

  • At fixed block indices, per-block encodings compose exactly, and deleting a present block yields the memory as if it had never been encoded.Remaining blocks keep their indices, and the resulting memory depends on the current block multiset rather than the edit path.Source: §2 · p. 4
  • Deleting a block touches only its b rows, while contiguous storage shifts later rows and index-mapped storage keeps deletion flat but adds a read-time gather.Memories and outputs were bit-identical across 100 random edit orders; Figure 2a compares deletion latency by memory size and layout.Source: §2 · p. 4, Figure 2 · p. 5

Expressivity and the cost of reading

  • Cross-attention sensitivity is independent of memory row count at fixed key/value diameters, while the cited sensitivity bound is not a measure of retrieval fidelity.The memory dependence enters through diameters bounded analytically by row normalization.Source: §2 · p. 4
  • The reader class contains a transformer as a limiting case, while a narrow reader solving a k-step retrieval chain has a per-width floor of R = k under the stated capacity assumption.The construction was verified exactly on 105 width/size/depth cells.Source: §2 · p. 4

3 The cache in operation

SMem makes block memories reusable and deletable by construction, with fast cached-context serving and retrieval beyond the trained window under the shared recipe. Across the reported comparisons, perplexity remains comparable to matched transformers, while recipe, read-time resources, and the scope of deletion qualify the results.

3.1 Reuse: an exact fast path for cached contexts

  • Fully cached contexts run the final block alone at 3.1–6.2 ms; bandwidth-bound decode uses 34–38% fewer KV rows and reaches 1.4–1.7× throughput.The reuse applies to any repeated block under different surrounding content, without reconstruction error.Source: §3.1 · p. 5

3.2 Retrieval that survives length under the shared recipe, unlike every transformer variant tested

  • At 4× training length, SMem reaches 0.14–0.28 exact match at 31 and 63 blocks, while the learned-position transformer scores 0.00–0.02.The Block-Attention-style arm scores 0.00 at 4×, and transformer extensions recover nothing beyond 15 blocks.Source: §3.2 · p. 6, Figure 3 · p. 7

3.3 Deletion as cache management

  • Deleting the needle block returns retrieval to the never-planted rate in all 40 cells (≤0.003); delete-then-generate is 8.5–452× faster than suffix recomputation.The flat deletion curve and decode gain use different memory layouts.Source: §3.3 · p. 7

4 What the move costs: matched-parameter pretraining and a recipe– architecture interaction

  • Across scales and recipes, the perplexity gap stays within −4.7% to +2.1% (up to +2.8% with learning-rate search), and recipe effects exceed architecture differences.The tuned recipe favors the transformer, whereas the shared-recipe comparisons favor SMem at the reported scales.Source: §4 · p. 8
  • SMem+RoPE leads SMem and the matched transformer at 160M and closes 29–59% of SMem’s gap to the RoPE transformer.The 410M Regime-2 composite result is from a single seed.Source: §4 · p. 9

5 Mechanism: the encoding tax and the price of reading

  • Block-local encoding costs +0.156 ± 0.034 nats overall and +1.34 ± 0.06 on retrievable-rare predictions; four reader rounds remove the pilot-scale deficit.The comparison isolates the cross-block component of the encoding tax.Source: §5 · pp. 9–10
  • On retrieval chains, minimal R·w grows roughly linearly with chain length k (OLS exponent 1.27; rank-based estimate 1.00).At widths w ≤4, formed cells succeed at R = k and minimal R·w equals k; width substitution appears at w ≥8.Source: §5 · p. 10

6 Limitations

  • SMem is a pretraining-time choice, not a retrofit; exact deletion covers context in memory, not knowledge learned into weights.The tuned-recipe far-range signal sits near the floor; deletion timing results also use different layouts from decode benchmarks.Source: §6 · p. 10, Ethics statement · p. 11

A Architecture and training details

SMem uses block-local encoding and a block-local reader with cross-attention to earlier-block memory; its architecture and verification gates specify how memory is built, extended, and checked.

A Architecture and training details

  • The encoder uses bidirectional attention within each block, while the reader attends only within its current block and cross-attends to strictly earlier-block memory.The reader also has a learned null row, which block 1 reads alone.Source: A Architecture and training details · pp. 13–14
  • Learned absolute positions index tokens by block and within-block offset; beyond trained length, changing slot assignments require re-encoding affected blocks, while SMem+RoPE re-encodes none.Within trained length, cached rows are never re-encoded.Source: A Architecture and training details · p. 14
  • The Block-Attention-style control uses two streams over shared weights, and the encoding-tax comparison instead uses a causal decoder-only model without a reader.That comparison's encoder mask differs from SMem's bidirectional within-block encoder.Source: A Architecture and training details · p. 14

B Verification gates

  • G5–G8 report exact composition, reader invariance, bit-exact deletion, and path independence under the stated deterministic-kernel and canonical-ordering conditions.G5 reports fp64 discrepancy < 10^-12; G7 and G8 report identical tensors and reads or bit-identical memories and outputs, respectively.Source: B Verification gates · pp. 14–15

C Theory statements and arguments

The theory section qualifies a reader-containment argument, sketches a round lower bound for a synthetic task, and gives an explicit matching construction verified across tested settings.

Containment

  • The reader class contains an L-layer transformer with R=L rounds only under the stated relaxation; this does not establish strict containment or apply to SMem's block-local reader.The construction disables cross-attention and removes the window restriction; the synthetic testbed uses a different parameterization.Source: C Theory statements and arguments · p. 15

Round lower bound

  • For PCk, the round-elimination argument sketches R = Ω(k) for narrow readers, but is not presented as a fully formalized theorem.Its refinement gives R ≥ k under a one-clean-retrieval-per-slot-per-round assumption, until width leaves that assumption.Source: C Theory statements and arguments · p. 15

Matching upper construction

  • An explicit PCk construction achieves R = k at any w and exact accuracy on all 105 tested cells; trained-testbed width substitution appears at w=8.The 105 cells span w ∈ {1, 2, 4}, n′ ∈ {8, 12, 16, 24, 32}, and k ∈ {1, 2, 3, 4, 6, 8, 12}.Source: C Theory statements and arguments · p. 15

D Compute accounting

SMem's compute accounting reports lower per-token FLOPs and fewer decode KV rows, alongside implementation-specific throughput trade-offs; a separate leak test checks the attention masks.

D Compute accounting

  • 2.5–4.6% cheaper per token, SMem's reported FLOPs are lower because self-attention is block-local and its reader is shallower.The accounting counts each multiply-add as 2 FLOPs and excludes embeddings and the identical LM head.Source: D Compute accounting · pp. 15–16
  • 34–38% fewer KV rows at n=T=1024, SMem stores D(n + b) rows versus baseline L n; at 1k–8k decode contexts the reduction is 34–41%.At full reuse, the encoder pass is skipped; at partial reuse, savings accrue per cached block.Source: D Compute accounting · p. 16

E Leak tests

  • 0.0 logit change before position 799, the leak test found no effect after replacing every token from position 800 onward with random tokens.The test ran at initialization in every run across both regimes; the masks are structural.Source: E Leak tests · p. 16

F Per-seed results and paired statistics

Across the reported paired comparisons, SMem's overall perplexity results vary by scale and recipe, while its retrievable-rare NLL is at or below the transformer values shown.

160M Regime 1

410M and 1.5B Regime 1

  • At 410M Regime 1, paired differences were −0.78 and −1.10 ppl (t(1) = −5.9, p = 0.11); the 1.5B single pair had final NLL 2.6715 vs 2.6571.The 1.5B comparison lists transformer first and SMem second.Source: F Per-seed results and paired statistics · p. 16

Stratified results

G The recipe control (Regime 2) and architectural controls

The recipe control shows that SMem’s perplexity comparison with transformers depends on recipe, while the SMem+RoPE composite matches or outperforms the matched learned-position transformer in the reported 160M and 410M Regime-2 results. Architectural controls and a symmetric learning-rate search further probe whether these findings reflect reader locality or tuning differences.

G The recipe control (Regime 2) and architectural controls

  • The recipe-control decision rule followed Regime-1 results but was registered before the trigger runs; the quality prediction was registered before pretraining.This chronology distinguishes the post-hoc-motivated control from its prospectively registered decision rule.Source: G The recipe control (Regime 2) and… · p. 16
  • The 410M trigger variants beat the shared-recipe transformer mean by 9.4% and 6.1%, both exceeding the registered 1% re-base bar.The variants used warmup-2000 at peak 3 × 10^-4 and no warmup at 4.5 × 10^-4, respectively.Source: G The recipe control (Regime 2) and… · p. 16
  • The warmup comparison is not a single-factor test because warmup spans 12.8% of the horizon and 98.8% of the baseline’s total learning-rate budget.The warmup trigger run is also reused as seed 0 of the re-based transformer arm.Source: G The recipe control (Regime 2) and… · p. 17
  • At 160M and 410M under the re-based recipe, SMem has higher perplexity than the matched transformer by +0.445 and +0.407, respectively.At 160M the paired difference is +0.445 ppl; at 410M the listed paired differences are +0.418 and +0.395.Source: G The recipe control (Regime 2) and… · p. 17
  • At 160M and 410M under Regime 2, SMem+RoPE matches or beats the learned-position transformer and improves on SMem by 2.4% and 2.1%, respectively.At 410M it trails the RoPE transformer by 5.5%, closing 29% of SMem’s distance to it; at 160M it closes 35% under Regime 2 and 59% under Regime 1.Source: G The recipe control (Regime 2) and… · p. 18
  • A depth-matched Block-Attention-style arm averages 29.66 ppl, at parity with the learned-position transformer and 1.33 ppl behind SMem.Its largest seed spread among the reported 160M arms is 1.2 ppl.Source: G The recipe control (Regime 2) and… · p. 18

Appendix J.

  • The CEPE-style full-prefix reader averages 29.56 ppl, 9.7% behind SMem, and its roll gap of 0.12 indicates little dependence on memory.The comparison fixes parameters and depth split, but the full-prefix reader already sees the memory’s encoded tokens at T=1024.Source: Appendix J · pp. 18–19
  • Across the symmetric learning-rate search, SMem remains 2.4–2.8% above the learned-position transformer, while the recipe shifts both arms together.Across the tested rates, the arms’ validation perplexities move by 10–11%; the retrievable-rare stratum reaches parity at the two highest rates.Source: Appendix J · p. 19

H Zero-shot results

Across the reported scales, SMem’s clearest zero-shot advantage is on LAMBADA, while differences on ARC-e, PIQA, HellaSwag, and WinoGrande are small. The roll counterfactual shows that readers’ dependence on their own context’s memory grows with scale.

H Zero-shot results

  • LAMBADA favors SMem by 0.027–0.029 at every scale, while no per-scale difference on ARC-e, PIQA, HellaSwag, or WinoGrande exceeds about one item-sampling standard error.At 160M, 410M, and 1.5B, LAMBADA scores are T/S .155/.182, .231/.260, and .348/.376.Source: H Zero-shot results · p. 19

I Roll counterfactual (the roll gap)

  • Rolling each sequence onto a neighboring sequence’s memory raises final-checkpoint NLL by 0.83–1.07 in Regime 1 and 0.85–0.97 in Regime 2.The measured gap grows with scale and estimates how much the reader depends on its own context’s memory.Source: I Roll counterfactual (the roll gap) · p. 19

J Needle retrieval: protocol, grids, and probes

The needle probes test retrieval across controlled block distances and evaluation lengths, including beyond the training window. They find far-range retrieval in SMem, exact retrieval loss after deleting the evidence block, and distance-sensitive behavior when keys repeat; open-book QA results are less conclusive.

J Needle retrieval: protocol, grids, and probes

  • The protocol plants an 8-token key followed by a 1-token value at a controlled block distance and evaluates exact match and conditional gain.Conditional gain compares the value’s log probability under matched and mismatched keys, canceling presence bias.Source: J Needle retrieval: protocol, grids,… · pp. 19–20
  • At 4× length, saturated block offsets retain retrieval for the SMem+RoPE composite, whereas direct rotary extension scores 0.00 at all distances.The saturated composite retrieves at range in all seven reported runs; the unsaturated offsets are beyond the trained range.Source: J Needle retrieval: protocol, grids,… · pp. 20–21
  • At T′=4096 and distances 31 and 63 blocks, SMem exact match is 0.14–0.28, while the learned-position baseline is 0.00–0.02.The primary-harness summary reports SMem values of 0.14/0.15, 0.25/0.22, and 0.28/0.28 at 160M, 410M, and 1.5B; the replication grid reports far-distance SMem retrieval and near-zero baseline scores.Source: J Needle retrieval: protocol, grids,… · pp. 20, 22
  • Deleting the needle block reduces exact match from 0.03–0.82 to at most 0.003 in all 40 cells, while deleting a neighboring block leaves accuracy within 0.04 of intact.After deletion, the value’s log probability is within 0.0–0.7 nats of the never-planted reference.Source: J Needle retrieval: protocol, grids,… · p. 22
  • With identical keys at distances 1 and 15, SMem selects the near value at 0.50–0.66 and the far value at 0.01–0.11 across scales.The probe measures distance preference but does not separate recency from slot position.Source: J Needle retrieval: protocol, grids,… · p. 22
  • In the natural-text ablation at 1.5B, source-block removal raises accuracy by +0.12 to +0.36 across distance bins, compared with +0.06 to +0.22 for the baseline.A transformer restricted to the last 1024 tokens ties or beats SMem inside its window and is blind beyond it.Source: J Needle retrieval: protocol, grids,… · p. 22
  • Open-book QA is harness-sensitive: every 95% interval for the TriviaQA difference-in-differences includes zero, while SQuAD-v1 exact match favors SMem in three reported comparisons.The section identifies needle, distance, and recency probes—not QA—as its capability evidence.Source: J Needle retrieval: protocol, grids,… · p. 22

K Serving measurements

SMem reduces serving costs for block skipping, bandwidth-bound decode, and deletion, but the fastest deletion and decode configurations use different memory layouts. The layout choice trades deletion latency against read-time overhead.

K Serving measurements

  • 3.1–6.2 ms: block-skip costs remain in this range across lengths, while cold prefill grows with context.The block-skip comparison covers batch 1; the 410M cold-prefill result at 8k matched 4k and was an unresolved anomaly.Source: K Serving measurements · p. 23
  • 1.44–1.67×: bandwidth-bound batched decode favors SMem after borderline cells are dropped; SMem stores 34–38% fewer KV rows.Batch-1 cells are launch-bound and show 1.03–1.21×; recomputing K/V is compute-bound and reaches 0.09–0.20× materialized throughput in the reported batched cells.Source: K Serving measurements · p. 23
  • 8.5× at 160M/512 blocks near the end to 452× at 160M/4096 blocks near the start: delete-then-generate beats suffix KV recomputation.At 1.5B and 4096 blocks, SMem takes 125–126 ms versus 13.9–29.9 s (111–237×).Source: K Serving measurements · p. 23
  • The index-mapped layout’s naive gather reaches 1.8–2.8× the contiguous layout’s per-token time in four of five batch-16 cells, so decode speed depends on layout.The flat-deletion layout and decode advantage were measured under different layouts in this implementation.Source: K Serving measurements · pp. 23–24
  • At 65,536 blocks, contiguous deletion takes 3.0 ms, while index-mapped deletion takes 0.014–0.027 ms but incurs a 3.3 ms read gather.At that size, the contiguous-layout gather is 0.003 ms, illustrating the deletion-time versus read-time trade-off.Source: K Serving measurements · p. 24

L Memory-suppression decomposition and the deletion task

The experiments separate cross-block memory effects from encoding costs and reader contributions. They also report synthetic resource-relation results with substantial budget and search qualifications.

L Memory-suppression decomposition and the deletion task

  • Deleting every evidence block reduces target accuracy from 0.58–0.61 to 0.361–0.363, near the own-block-only level; the control stratum changes by at most 0.001.In the suppression decomposition, accuracy is 0.597 with intact memory, 0.365 with own-block memory only, and 0.189 without cross-attention.Source: L Memory-suppression decomposition and… · p. 24

M Encoding-tax measurements

  • From-scratch matched 83.6M models show a mean encoding tax of 0.156 nats on tax tokens and +1.337 nats on retrievable rare tokens, versus −0.138 nats on non-retrievable rare tokens.Across three seed pairs, the retrievable-to-overall concentration ratio averages 8.8×.Source: M Encoding-tax measurements · p. 24
  • With Pythia-410M under an inference-time 64-token block-causal mask, the tax is +3.04 nats when the target occurs in an earlier block and +0.57 elsewhere.The model was not trained with this mask; full-attention sanity NLL is 3.29.Source: M Encoding-tax measurements · p. 24
  • With four reader rounds, the cross-block accuracy deficit is gone (∆≤0 at every width), while an equal-FLOPs extra encoder layer changes accuracy by only +0.002.The bare block-local model’s deficit is +15.1, +19.9, and +20.8 points at widths 64, 128, and 256, respectively.Source: M Encoding-tax measurements · p. 24

N Synthetic resource-relation grids

  • Among formed synthetic cells at widths ≤4, minimal successful rounds equal k; the k=6 search produced zero formations within its approximately 19 GPU-hour budget.The k=6 result is budget-censored and uninformative about the floor; under worst-case censoring, the logarithmic alternative is not excluded.Source: N Synthetic resource-relation grids · p. 25

O Pilot-scale split ablation

At pilot scale, adding reader depth improves perplexity monotonically, while encoder deepening changes it within noise and smaller blocks help at fixed compute.

O Pilot-scale split ablation

  • Perplexity falls from 6.40 to 5.79 to 5.52 as reader depth increases from 2 to 4 to 6 layers at a fixed 2-layer encoder.At fixed compute, b=16 gives 5.74 perplexity versus 5.86 for b=64; the production encoder-to-reader splits follow this ablation and stay fixed across scales.Source: O Pilot-scale split ablation · p. 25

P Training practices and precedents

The pilot-phase training practices pair diagnosed formation failures with targeted fixes, while several listed practices are described as standard practice.

P Training practices and precedents

  • Long warmup near 300 steps can prevent synthetic retrieval-circuit formation, while small curriculum-vocabulary steps with optimizer resets address multi-hop formation collapse (Hoffmann et al., 2024).The passage says the warmup constraint binds the synthetic setting; pretraining-scale SMem tolerates standard warmup.Source: P Training practices and precedents · p. 25
  • For recurrent curricula, random-depth truncation with a straight-through initial state addresses early-round starvation; active-id logit masking prevents inactive-vocabulary drift.The passage also recommends standard 0.02 initialization with AdamW and balancing embedding norms when positional signal is too weak.Source: P Training practices and precedents · p. 25
  • The pilot practices also target residual-branch gradient attenuation with zero-initialized branch outputs, stratum-hidden effects with decomposition, and small-model fp16 overflow with bf16.The passage recommends fp32 evaluation alongside bf16.Source: P Training practices and precedents · p. 25
Loading 2609.32395v1…