Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Fast-weight Product Key Memory

Tianyu Zhao, Llion Jones

TL;DR

This work tackles the storage-computation trade-off in sequence modeling by transforming Product Key Memory (PKM) into Fast-weight Product Key Memory (FwPKM), a dynamic memory that updates during both training and inference via chunk-level gradient descent on a local memorization objective. By updating keys and values as fast weights and incorporating lookahead targets, marginal-entropy addressing, and gating, FwPKM acts as a high-fidelity episodic memory that complements semantic memory from slow-weight layers. Empirical results show significant perplexity reductions on long-context datasets and strong generalization to 128K-token contexts in Needle in a Haystack, with iterative readouts further boosting retrieval fidelity. The approach opens a path toward hybrid memory systems that combine large, sparse storage with rapid, context-specific updates, albeit with notable engineering and kernel-performance challenges to achieve broader scalability.

Abstract

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 costs, linear variants provide efficiency but suffer from limited, fixed-size storage. We propose Fast-weight Product Key Memory (FwPKM), a novel architecture that resolves this tension by transforming the sparse Product Key Memory (PKM) from a static module into a dynamic, "fast-weight" episodic memory. Unlike PKM, FwPKM updates its parameters dynamically at both training and inference time via local chunk-level gradient descent, allowing the model to rapidly memorize and retrieve new key-value pairs from input sequences. Experiments reveal 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.

Fast-weight Product Key Memory

TL;DR

This work tackles the storage-computation trade-off in sequence modeling by transforming Product Key Memory (PKM) into Fast-weight Product Key Memory (FwPKM), a dynamic memory that updates during both training and inference via chunk-level gradient descent on a local memorization objective. By updating keys and values as fast weights and incorporating lookahead targets, marginal-entropy addressing, and gating, FwPKM acts as a high-fidelity episodic memory that complements semantic memory from slow-weight layers. Empirical results show significant perplexity reductions on long-context datasets and strong generalization to 128K-token contexts in Needle in a Haystack, with iterative readouts further boosting retrieval fidelity. The approach opens a path toward hybrid memory systems that combine large, sparse storage with rapid, context-specific updates, albeit with notable engineering and kernel-performance challenges to achieve broader scalability.

Abstract

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 costs, linear variants provide efficiency but suffer from limited, fixed-size storage. We propose Fast-weight Product Key Memory (FwPKM), a novel architecture that resolves this tension by transforming the sparse Product Key Memory (PKM) from a static module into a dynamic, "fast-weight" episodic memory. Unlike PKM, FwPKM updates its parameters dynamically at both training and inference time via local chunk-level gradient descent, allowing the model to rapidly memorize and retrieve new key-value pairs from input sequences. Experiments reveal 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.
Paper Structure (50 sections, 19 equations, 10 figures, 3 tables)

This paper contains 50 sections, 19 equations, 10 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Product Key Memory
  3. Top-$k$ Key-value Memory
  4. Product Key Memory
  5. Fast-weight Product Key Memory
  6. Fast Weights
  7. Memorization Optimization
  8. MSE loss
  9. Loss aggregation and gradient shaping
  10. Reducing MSE losses to a scalar by summation
  11. Weighting gradient by row contribution
  12. Weighting MSE loss by token importance
  13. No gradient clipping
  14. Addressing Optimization
  15. Marginal entropy loss
  16. ...and 35 more sections

Figures (10)

  • Figure 1: Architecture of FwPKM.
  • Figure 2: Perplexity on Fineweb-Edu, LC64, and LAMBADA
  • Figure 3: FwPKM gating value distribution on Fineweb-Edu, LC64, and LAMBADA test sets. Each row represents one FwPKM layer.
  • Figure 4: Stacked bar plots for NIAH accuracy results on 4K-/8K-/32K-/128K-length test sets. Each stacked bar shows the accuracies of $\{1,2,3,4\}$-iter NIAH evaluations.
  • Figure 5: An example of FwPKM slot access of GDN+PKM@6+FwPKM@2,10 during generating an NIAH-4K answer. The model memorizes the haystack for 3 extra iterations, i.e.$4$-iter NIAH.
  • ...and 5 more figures