[Paper Review] Large Memory Layers with Product Keys

Guillaume Lample, Alexandre Sablayrolles+3 more|HAL (Le Centre pour la Communication Scientifique Directe)|Jul 10, 2019

Topic Modeling50 citations

TL;DR

Introduces a scalable, memory-augmented layer based on product keys that enables very large memory capacity with exact nearest-neighbor search, integrated into a transformer for large-scale language modeling and achieving better accuracy with faster inference than deeper baselines.

ABSTRACT

This paper introduces a structured memory which can be easily integrated into a neural network. The memory is very large by design and significantly increases the capacity of the architecture, by up to a billion parameters with a negligible computational overhead. Its design and access pattern is based on product keys, which enable fast and exact nearest neighbor search. The ability to increase the number of parameters while keeping the same computational budget lets the overall system strike a better trade-off between prediction accuracy and computation efficiency both at training and test time. This memory layer allows us to tackle very large scale language modeling tasks. In our experiments we consider a dataset with up to 30 billion words, and we plug our memory layer in a state-of-the-art transformer-based architecture. In particular, we found that a memory augmented model with only 12 layers outperforms a baseline transformer model with 24 layers, while being twice faster at inference time. We release our code for reproducibility purposes.

Motivation & Objective

• Increase model capacity without proportional compute by introducing a large, trainable memory layer.
• Provide exact nearest-neighbor search within a scalable key-value memory using product keys.
• Demonstrate the memory layer’s effectiveness within a transformer on large-scale language modeling tasks.
• Show that memory-augmented models can outperform deeper baselines with lower or comparable compute.
• Analyze how memory size, placement in the network, and normalization affect performance and usage.

Proposed method

• Define a memory with a query network, a two-sub-key product key space, and a value lookup table.
• Use product keys formed by the outer product of two codebooks to implicitly define a large key set.
• Select top-k keys via inner-product similarity with a query, then compute a weighted sum of corresponding values (sparse reads).
• Employ multi-head memory attention where each head has its own query network and sub-keys but shares values.
• Train with sparse memory updates and batch normalization on the query network to improve key coverage and usage.
• Replace the FFN in selected transformer layers with the memory layer, allowing additive residual connections (x → x + PKM(x)).

Experimental results

Research questions

• RQ1 Does a product-key memory (PKM) provide scalable capacity with negligible overhead while maintaining exact search?
• RQ2 How do memory size, number of heads, k-NN, and memory placement influence perplexity, memory usage, and speed in large transformers?
• RQ3 How does product-key memory compare to flat-key memories in terms of performance, usage, and efficiency?
• RQ4 What training configurations (BatchNorm, sparse updates) optimize memory usage and model performance on large-scale language modeling datasets?

Key findings

• A memory with 12 layers and product keys outperforms a 24-layer transformer of the same dimension on CC-News, with roughly half the inference time.
• Increasing memory size consistently lowers perplexity, with memory size 1M slots achieving notable gains when batch normalization improves key usage.
• Product keys provide higher memory usage, lower perplexity, and faster inference than flat keys at comparable sizes.
• Memory placement matters; inserting memory at intermediate transformer layers (around 4–5) yields the best performance.
• 4 memory heads with 32 keys per head and 512^2 total slots yield favorable trade-offs between speed and perplexity; product keys reduce parameter count vs. flat keys and scale more efficiently at inference.
• Batch normalization in the query network significantly improves key usage and perplexity for large memories.