Memory Layers at Scale: Trainable Key-Value Lookup for Facts

This episode examines "Memory Layers at Scale" from Meta FAIR, which asks whether a language model can gain huge numbers of fact-storing parameters without adding any arithmetic per token. It explains how a memory layer works: a token's query retrieves the top-k rows from a large learned key-value table, and product keys keep the search cheap. It also covers where the idea came from, including the 2019 product-key work, which this paper scales up to billions of parameters with custom GPU kernels, and the related million-expert approach. The discussion examines what "matched FLOPs" leaves out, namely stored parameters and memory-bandwidth and communication costs. It also looks at how closed-book question answering on NaturalQuestions and TriviaQA is used to test the factual-recall claims against dense and mixture-of-experts baselines. Listeners interested in how models store facts, and in the gap between FLOP counts and real hardware cost, will find it useful.

Interactive Visualization: Memory Layers at Scale: Trainable Key-Value Lookup for Facts

Sources:
1. Memory Layers at Scale — Vincent-Pierre Berges, Barlas Oğuz, Daniel Haziza, Wen-tau Yih, Luke Zettlemoyer, Gargi Ghosh, 2024
http://arxiv.org/abs/2412.09764
2. Mixture of A Million Experts (PEER) — Xu Owen He, 2024
https://scholar.google.com/scholar?q=Mixture+of+A+Million+Experts+%28PEER%29
3. Large Memory Layers with Product Keys — Guillaume Lample, Alexandre Sablayrolles, Marc'Aurelio Ranzato, Ludovic Denoyer, Hervé Jégou, 2019
https://scholar.google.com/scholar?q=Large+Memory+Layers+with+Product+Keys
4. DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models — Damai Dai et al., 2024
https://scholar.google.com/scholar?q=DeepSeekMoE%3A+Towards+Ultimate+Expert+Specialization+in+Mixture-of-Experts+Language+Models
5. Scaling Laws for Fine-Grained Mixture of Experts — Jan Ludziejewski et al., 2024
https://scholar.google.com/scholar?q=Scaling+Laws+for+Fine-Grained+Mixture+of+Experts
6. Physics of Language Models: Part 3.3, Knowledge Capacity Scaling Laws — Zeyuan Allen-Zhu, Yuanzhi Li, 2024
https://scholar.google.com/scholar?q=Physics+of+Language+Models%3A+Part+3.3%2C+Knowledge+Capacity+Scaling+Laws
7. Generalization through Memorization: Nearest Neighbor Language Models (kNN-LM) — Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, Mike Lewis, 2020
https://scholar.google.com/scholar?q=Generalization+through+Memorization%3A+Nearest+Neighbor+Language+Models+%28kNN-LM%29
8. Improving language models by retrieving from trillions of tokens (RETRO) — Sebastian Borgeaud et al., 2022
https://scholar.google.com/scholar?q=Improving+language+models+by+retrieving+from+trillions+of+tokens+%28RETRO%29
9. Product Key Memory with Mixture of Experts / Mixture of Memory-style hybrids; and Titans: Learning to Memorize at Test Time — Ali Behrouz, Peilin Zhong, Vahab Mirrokni, 2024
https://scholar.google.com/scholar?q=Product+Key+Memory+with+Mixture+of+Experts+%2F+Mixture+of+Memory-style+hybrids%3B+and+Titans%3A+Learning+to+Memorize+at+Test+Time
10. Locating and Editing Factual Associations in GPT (ROME) and Transformer Feed-Forward Layers Are Key-Value Memories — Kevin Meng et al.; Mor Geva et al., 2022 / 2021
https://scholar.google.com/scholar?q=Locating+and+Editing+Factual+Associations+in+GPT+%28ROME%29+and+Transformer+Feed-Forward+Layers+Are+Key-Value+Memories
11. Memory Mosaics / Mixture of Lookup Experts / Engram-style conditional-memory work (2025) — Various, 2025
https://scholar.google.com/scholar?q=Memory+Mosaics+%2F+Mixture+of+Lookup+Experts+%2F+Engram-style+conditional-memory+work+%282025%29
Interactive Visualization: Memory Layers at Scale: Trainable Key-Value Lookup for Facts