NeurIPS 2019 · Lample, Sablayrolles, Ranzato, Denoyer, Jégou · Facebook AI Research / Sorbonne Université

Product Keys Give Transformers Million-Slot Memory Layers

Companion to the episode on “Large Memory Layers with Product Keys”.
arXiv 1907.05242 NeurIPS PDF Interactive page

One token through the memory layer

Click a stage (or use the buttons) to light it up.

Where the memory sits

Example: 16-layer model, memory layers 6 and 12 (as described in the episode). Memory replaces the FFN in those layers only.

Toy product-key searchIllustrative values

3 sub-keys per codebook → 9 implied keys. Sub-key vectors are made up for the demo; the procedure is the paper’s. Rotate each query half.
left half angle
right half angle
keep k per half

Why top-k of each half is enough

A key’s score = left score + right score. If a true winner used a left sub-key outside the left top-k, then k better left sub-keys paired with the same right partner would all beat it — contradiction.
Scope: exact over product-structured (Cartesian) keys only, not a million freely chosen vectors.

Search cost: flat vs product keys

Cost in full-query-dimension comparison units, from the formula in the episode: flat = N; product ≈ √N + k² (two codebooks of n half-size sub-keys ≈ n full-size, plus k² pairs). N = n².
sub-keys per codebook n = 2^
k

Implied keys vs stored sub-keys

Hover a cell: key (i, j) = left sub-key i ⊕ right sub-key j. First 16×16 of the grid shown.

Reported numbers

Values as stated in the episode. Single runs, no variance reported. Lower perplexity is better.

Main setup & training recipe

What the comparison leaves out

Hover or tap a row. Statuses summarize the episode’s critique.

Back-of-envelope parameter billAssumption

The paper reports no parameter count. Value dimension is not confirmed in the text; slide to see what 262,144 slots × dim would imply.
value dim

References

  1. Large Memory Layers with Product Keys — Lample et al., 2019 (arXiv 1907.05242)
  2. Product Quantization for Nearest Neighbor Search — Jégou, Douze, Schmid, 2011
  3. Scaling Memory-Augmented Neural Networks with Sparse Reads and Writes — Rae et al., 2016
  4. Memory Layers at Scale — Berges et al., 2024
  5. Outrageously Large Neural Networks: Sparsely-Gated MoE — Shazeer et al., 2017
  6. Generalization through Memorization: kNN-LM — Khandelwal et al., 2020
  7. Transformer Feed-Forward Layers Are Key-Value Memories — Geva et al., 2021
  8. Switch Transformers — Fedus, Zoph, Shazeer, 2021
  9. Improving language models by retrieving from trillions of tokens (RETRO) — Borgeaud et al., 2022
  10. Mixture of A Million Experts (PEER) — He, 2024
  11. Scaling Laws for Fine-Grained Mixture of Experts — Ludziejewski et al., 2024
  12. Large Product Key Memory for Pretrained Language Models — Kim, Jung, 2020
  13. Adaptive Input Representations for Neural Language Modeling — Baevski, Auli, 2019