# Fast-weight Product Key Memory

arXiv: https://arxiv.org/abs/2601.00671 · HTML: https://mlsanity.com/paper/2601.00671/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 53 references that are in the map.
- First posted on arXiv 2 Jan 2026.

## Builds on (most cited first, 30 of 42)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks](https://mlsanity.com/paper/1602.07868/) (2016)
- [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://mlsanity.com/paper/2412.06464/) (2024)
- [Gated Linear Attention Transformers with Hardware-Efficient Training](https://mlsanity.com/paper/2312.06635/) (2023)
- [Linear Transformers Are Secretly Fast Weight Programmers](https://mlsanity.com/paper/2102.11174/) (2021)
- [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://mlsanity.com/paper/2406.06484/) (2024)
- [Memory Networks](https://mlsanity.com/paper/1410.3916/) (2014)
- [Learning to (Learn at Test Time): RNNs with Expressive Hidden States](https://mlsanity.com/paper/2407.04620/) (2024)
- [Titans: Learning to Memorize at Test Time](https://mlsanity.com/paper/2501.00663/) (2024)
- [Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free](https://mlsanity.com/paper/2505.06708/) (2025)
- [Landmark Attention: Random-Access Infinite Context Length for Transformers](https://mlsanity.com/paper/2305.16300/) (2023)
- [RedPajama: an Open Dataset for Training Large Language Models](https://mlsanity.com/paper/2411.12372/) (2024)
- [Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling](https://mlsanity.com/paper/2406.07522/) (2024)
- [Using Fast Weights to Attend to the Recent Past](https://mlsanity.com/paper/1610.06258/) (2016)
- [RWKV-7 "Goose" with Expressive Dynamic State Evolution](https://mlsanity.com/paper/2503.14456/) (2025)
- [Large Memory Layers with Product Keys](https://mlsanity.com/paper/1907.05242/) (2019)
- [Test-time regression: a unifying framework for designing sequence models with associative memory](https://mlsanity.com/paper/2501.12352/) (2025)
- [Nested Learning: The Illusion of Deep Learning Architectures](https://mlsanity.com/paper/2512.24695/) (2025)
- [It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization](https://mlsanity.com/paper/2504.13173/) (2025)
- [Learning Associative Inference Using Fast Weight Memory](https://mlsanity.com/paper/2011.07831/) (2020)
- [Memory Layers at Scale](https://mlsanity.com/paper/2412.09764/) (2024)

## Built on by (most cited first, 13 of 13)

- [RoboTTT: Context Scaling for Robot Policies](https://mlsanity.com/paper/2607.15275/) (2026)
- [IndexMem: Learned KV-Cache Eviction with Latent Memory for Long-Context LLM Inference](https://mlsanity.com/paper/2605.25475/) (2026)
- [MARCH: Scaling Recurrent Memory with Content-Routed State Anchors](https://mlsanity.com/paper/2608.12435/) (2026)
- [Online Vector Quantized Attention](https://mlsanity.com/paper/2602.03922/) (2026)
- [Proteus: Incremental Memory Activation for Long-Context Sequence Modeling](https://mlsanity.com/paper/2608.16844/) (2026)
- [Safin-1: Safety from Within through Memory-Native State Evolution](https://mlsanity.com/paper/2609.00092/) (2026)
- [MemDLM: Memory-Enhanced DLM Training](https://mlsanity.com/paper/2603.22241/) (2026)
- [TokenMem: Faithful Knowledge Injection for Frozen LLMs](https://mlsanity.com/paper/2607.22625/) (2026)
- [MindMemOS: A Portable and Self-Evolving Memory Operating Layer for AI Agents](https://mlsanity.com/paper/2608.12428/) (2026)
- [Latest Exact Match Attention](https://mlsanity.com/paper/2609.25802/) (2026)
- [Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling](https://mlsanity.com/paper/2609.36529/) (2026)
- [Decision Titan: Test-Time Training for Long-Term Memory in Offline Reinforcement Learning](https://mlsanity.com/paper/2610.01513/) (2026)
- [Towards In-Parameter Memory Augmentation for Large Language Models](https://mlsanity.com/paper/2610.08630/) (2026)

## Abstract (arXiv)

Sequence modeling layers in modern language models typically face a trade-off between storage capacity and computational efficiency. While softmax attention offers unbounded storage at prohibitive quadratic cost, linear variants are more efficient but suffer from limited, fixed-size storage. We introduce Fast-weight Product Key Memory (FwPKM), a sparse fast-weight memory layer that resolves this tension. FwPKM updates sparsely activated parameters at both training and inference time using chunk-level gradient descent on a local memory-rewrite objective. This performs Test-Time Training (TTT)-style gradient updates on activated slots in a sparse memory, enabling rapid memorization and retrieval of many new key-value associations while keeping per-token compute low and fixed. Experiments show that FwPKM functions as an effective episodic memory that complements the semantic memory of standard modules, yielding significant perplexity reductions on long-context datasets. Notably, in Needle-in-a-Haystack evaluations, FwPKM generalizes to 128K-token contexts despite being trained on only 4K-token sequences.