Tensor Cache: Eviction-conditioned Associative Memory for Transformers

Technical Summary: Tensor Cache

Problem Statement
Autoregressive Transformer inference relies on caching Key-Value (KV) pairs to avoid recomputing prefixes. However, the retained KV state grows linearly with context length, depth, and concurrency, creating a memory bottleneck. Sliding-window caching bounds memory usage by retaining only the most recent $W$ tokens, but this approach suffers from "hard forgetting": once a token is evicted from the window, its information is entirely discarded, making relevant evidence outside the window inaccessible. Existing bounded-state methods either discard evicted tokens entirely or attempt to compress every token into a running average, which can introduce noise or fail to distinguish between recent and distant context effectively.

Methodology: Tensor Cache (TC)
Tensor Cache introduces a two-level bounded-state architecture that transforms KV eviction from a deletion event into an associative-memory write operation.

1. Two-Level Architecture:

   • L1 (First-Level Cache): A fixed-size KV ring buffer holding the most recent $W$ tokens. This supports exact softmax attention for local context.
   • L2 (Second-Level Cache): A fixed-size outer-product matrix $A$ (fast-weight memory) that stores information exclusively from tokens evicted from the L1 window.
   • Write Mechanism: When the ring buffer overwrites an entry, the displaced $(k, v)$ pair is not discarded. Instead, it is written into the L2 matrix $A$ via an outer-product update: $A \leftarrow \lambda A + \eta (k \otimes v)$. Here, $\lambda$ (decay) and $\eta$ (write rate) are learned per-head scalars.
   • Read Mechanism: Future queries $q_t$ access the compressed past by computing $r_t = q_t A$. Leveraging the linear-attention identity $q(k \otimes v) = \langle q, k \rangle v$, this operation effectively retrieves information from all evicted tokens without storing them individually.
   • Fusion: A learned scalar gate $g$ fuses the outputs of the local attention (L1) and the memory read (L2): $y_t = y_{local} + \sigma(g) m_t$.

2. Training Optimization (Parallel Scan):
   The paper identifies a critical issue with standard chunked training shortcuts. A naive approach summarizes a chunk by its mean ($\bar{k} \otimes \bar{v}$), which introduces $C^2 - C$ spurious cross-token outer products ($k_i \otimes v_j$ where $i \neq j$) that do not exist during streaming inference.

   • Solution: The authors implement a parallel weighted-sum scan that is mathematically equivalent to sequential per-token writes. This closes the gap between training and inference within float32 epsilon ($< 10^{-7}$ relative error), ensuring the model learns the correct associative dynamics without the noise of cross-token interference.

Key Contributions

• Eviction-Conditioned Associative Memory: Unlike prior outer-product memories (e.g., Infini-attention, mLSTM) that ingest every token, TC uses the outer-product matrix strictly as an L2 cache fed only by sliding-window evictions. This cleanly partitions tokens into "resident in L1" and "compressed in L2."
• Correction of Training Artifacts: The paper identifies and resolves a previously overlooked training shortcut (chunked-mean) that introduces spurious cross-token products. The proposed parallel weighted-sum scan eliminates this mismatch at the same computational cost.
• Bounded-State Efficiency: TC provides a mechanism where older context influences predictions without retaining individual evicted entries, offering a practical alternative to full KV retention.

Experimental Results
The authors evaluated TC across systems scaling, synthetic associative recall, and real-text long-context language modeling (OpenWebText and Shakespeare).

• Synthetic Associative Recall: In tasks requiring retrieval of information beyond the sliding window, TC achieved 100% accuracy, matching Full KV performance. It reduced retained inference state by 72–84% compared to Full KV while significantly outperforming Infini-attention (which degraded as gaps increased) and sliding-window baselines (which performed at chance levels).
• Long-Context Language Modeling: On OpenWebText, TC achieved the lowest mean Negative Log-Likelihood (NLL) at every evaluation context length from 1,024 to 32,768 tokens, with the single exception of $L=2,048$ where it was statistically indistinguishable from Infini-attention. At $L=32,768$, TC reached an NLL of 5.14, outperforming Full KV (6.00) and Infini-attention (5.42), while using 2.4$\times$ less peak GPU memory than Full KV.
• Throughput: TC incurred a ~30% throughput overhead compared to pure sliding-window methods (due to L2 read/update) but remained consistently faster than Infini-attention across all tested context lengths.
• Capacity Diagnostics: Experiments on raw associative memory capacity showed that retrieval performance degrades as more unrelated associations are compressed into the fixed state, confirming the bounded nature of the memory.

Significance and Claims
The paper claims that Tensor Cache successfully bridges the gap between the memory efficiency of sliding-window attention and the recall capabilities of full-context attention. By treating KV eviction as a structured memory operation rather than a destructive deletion, TC enables Transformers to access relevant evidence outside the immediate window without the linear memory cost of full KV caches. The authors emphasize that their primary contribution is not the invention of the outer-product memory itself, but its specific application as an eviction-conditioned L2 cache and the rigorous correction of training artifacts that previously hindered such methods. TC is presented as a practical, bounded-state alternative that improves the memory-quality frontier for long-context inference.