mlsanityData to 7 Oct 2026

Fast-weight Product Key Memory

arXiv 2601.00671 ↗First posted 2 Jan 2026cs.CL cs.AIcited by 11

Not on the map. Why?

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.

Key facts. Cited by 11 papers in the map (data through 7 Oct 2026). Builds on 42 of its 53 references that are in the map. First posted on arXiv 2 Jan 2026.

Builds on

42 of 53 references.

References · 42

  1. 2017Attention Is All You Needcited by 59,932
  2. 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,616
  3. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
  4. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,551
  5. 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 1,979
  6. 2020GLU Variants Improve Transformercited by 1,724
  7. 2024Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Dualitycited by 1,283
  8. 2024The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scalecited by 1,104
  9. 2016The LAMBADA dataset: Word prediction requiring a broad discourse contextcited by 960
  10. 2023RWKV: Reinventing RNNs for the Transformer Eracited by 814
  11. 2016Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networkscited by 541
  12. 2023Gated Linear Attention Transformers with Hardware-Efficient Trainingcited by 490
30 more
  1. 2024Gated Delta Networks: Improving Mamba2 with Delta Rulecited by 483
  2. 2021Linear Transformers Are Secretly Fast Weight Programmerscited by 471
  3. 2024Parallelizing Linear Transformers with the Delta Rule over Sequence Lengthcited by 373
  4. 2014Memory Networkscited by 354
  5. 2024Learning to (Learn at Test Time): RNNs with Expressive Hidden Statescited by 329
  6. 2024Titans: Learning to Memorize at Test Timecited by 303
  7. 2025Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Freecited by 266
  8. 2023Landmark Attention: Random-Access Infinite Context Length for Transformerscited by 206
  9. 2024RedPajama: an Open Dataset for Training Large Language Modelscited by 178
  10. 2024Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modelingcited by 167
  11. 2016Using Fast Weights to Attend to the Recent Pastcited by 161
  12. 2025RWKV-7 "Goose" with Expressive Dynamic State Evolutioncited by 134
  13. 2019Large Memory Layers with Product Keyscited by 120
  14. 2025Nested Learning: The Illusion of Deep Learning Architecturescited by 69
  15. 2025Test-time regression: a unifying framework for designing sequence models with associative memorycited by 63
  16. 2025It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimizationcited by 60
  17. 2020Learning Associative Inference Using Fast Weight Memorycited by 51
  18. 2024Memory Layers at Scalecited by 39
  19. 2025Understanding Transformer from the Perspective of Associative Memorycited by 31
  20. 2016Scaling Memory-Augmented Neural Networks with Sparse Reads and Writescited by 28
  21. 2023Approximating Two-Layer Feedforward Networks for Efficient Transformerscited by 27
  22. 2025Key-value memory in the braincited by 22
  23. 2024Ultra-Sparse Memory Networkcited by 19
  24. 2024Memory Mosaicscited by 16
  25. 2025Blending Complementary Memory Systems in Hybrid Quadratic-Linear Transformerscited by 12
  26. 2025Artificial Hippocampus Networks for Efficient Long-Context Modelingcited by 9
  27. 2025TNT: Improving Chunkwise Training for Test-Time Memorizationcited by 7
  28. 2025Memory Mosaics at scalecited by 5
  29. 2025UltraMemV2: Memory Networks Scaling to 120B Parameters with Superior Long-Context Learningcited by 5
  30. 2023Inverse distance weighting attentioncited by 4

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 11.

11 papers.

  1. 2026RoboTTT: Context Scaling for Robot Policiescited by 13
  2. 2026IndexMem: Learned KV-Cache Eviction with Latent Memory for Long-Context LLM Inferencecited by 7
  3. 2026Online Vector Quantized Attentioncited by 1
  4. 2026MARCH: Scaling Recurrent Memory with Content-Routed State Anchorscited by 1
  5. 2026Proteus: Incremental Memory Activation for Long-Context Sequence Modelingcited by 1
  6. 2026Safin-1: Safety from Within through Memory-Native State Evolutioncited by 1
  7. 2026MemDLM: Memory-Enhanced DLM Trainingcited by 0
  8. 2026TokenMem: Faithful Knowledge Injection for Frozen LLMscited by 0
  9. 2026MindMemOS: A Portable and Self-Evolving Memory Operating Layer for AI Agentscited by 0
  10. 2026Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modelingcited by 0
  11. 2026Decision Titan: Test-Time Training for Long-Term Memory in Offline Reinforcement Learningcited by 0