Source-linked AI summary
Large Memory Layers with Product Keys
Guillaume Lample, Alexandre Sablayrolles, Marc'Aurelio Ranzato, Ludovic Denoyer, Hervé Jégou
TL;DR
Large neural networks improve modeling but become increasingly expensive to train and run. This paper adds a product-key key-value memory layer with exact search and sparse access, achieving strong language-modeling performance: a 12-layer model matches or outperforms a 24-layer baseline while running twice as fast at inference.
Problem
Increasing neural-network capacity improves modeling and generalization but raises computational complexity, motivating architectures that preserve capacity under limited computation.
Method
The paper integrates a trainable key-value memory layer using product keys, exact nearest-neighbor search, and sparse memory reads and updates.
Results
A 12-layer memory-augmented model outperforms a 24-layer transformer while being twice faster at inference time.
Takeaways & Limitations
The memory layer provides large capacity gains with negligible computational overhead and improves the capacity–efficiency trade-off in large-scale language modeling.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper addresses the cost of scaling neural-network capacity by introducing a large key-value memory layer with exact, fast product-key search and sparse access. Integrated into transformers, the layer improves the capacity–computation trade-off while achieving strong language-modeling results.
- Increasing neural-network capacity improves modeling and generalization but substantially raises training and inference costs.
- Keys are formed by concatenating two sub-keys, creating a very large implicit key set whose values provide most of the parameters.
- Sparse selection and updates make training and inference efficient despite the large memory size.
- The proposed memory layer adds large capacity with only slight computational overhead during training and testing.
- Product keys provide exact nearest-neighbor search without an indexing structure that must be repeatedly relearned during training.
- With 1 memory and 12 layers, the method outperforms a 24-layer transformer while running twice faster at inference time.
2 Related work
Related work increases neural-network capacity through conditional computation, memory augmentation, discretization, and sparsity. The paper distinguishes its approach by combining product-key structure with exact search and trainable sparse memory access.
- Conditional-computation models route inputs through subsets of large networks using mixture-of-experts, gating, or reinforcement-learning methods.
- Memory-augmented networks support variable-length inputs and feature-space operations, but many scale linearly with memory size.
- Discretization and hashing methods reduce computation, but approximate indexes create optimization challenges during training.
- The proposed method borrows product-quantization ideas to represent many key vectors compactly while updating them through regular back-propagation.
- Transformer memory components resemble self-attention query, key, and value networks, but use free embeddings and a much larger value set instead of input-token representations.
3 Learnable product key memories
The memory layer uses product keys to provide large capacity with efficient exact nearest-neighbor retrieval. It integrates query generation, structured key selection, sparse value access, and optional multi-head computation into neural networks.
- Memory structure: The memory maps inputs to queries, selects top-k product keys, and combines their associated values into an output.Its components are a query network, a key-selection module with two sub-key sets, and a value lookup table.
- Query generation: The query network maps inputs into a lower-dimensional space, typically reducing them to dq = 512, and batch normalization improves key coverage during training.Keys are randomly initialized and occupy the space relatively uniformly.
- Complexity: Product-key retrieval avoids exhaustive comparison over all flat keys by searching the two sub-key sets and then ranking the resulting k^2 candidates.Flat-key search requires O(|K| × dq) operations, whereas product-key search uses sub-key comparisons plus candidate ranking.
- Memory structure: Product keys are formed from two sub-key codebooks, creating |K| = |C|×|C′| implicit keys without explicitly storing every key.The Cartesian-product construction associates each combined key with a value memory slot.
- Key selection: Splitting the query into two sub-queries and retrieving k nearest sub-keys from each codebook guarantees that the k most similar product keys lie in the resulting candidate set.The candidate set contains k^2 combined keys, after which the best product keys are selected.
- Transformer integration: The memory can replace selected transformer FFN layers and can use multiple independent query heads whose outputs are summed.Each head has its own query network and sub-keys while sharing value vectors; different heads typically select different keys.
- Complexity: For a memory of size |K| = 1024^2 and small k, product-key retrieval requires about 10^3 fewer operations than exhaustive search.The study also reports better performance than a corresponding flat-key memory in ablation experiments.
4 Experiments
The experiments evaluate product-key memory layers integrated into transformers on large-scale language modeling, including memory usage, placement, size, and retrieval settings. Memory improves perplexity and capacity efficiency, with 12-layer memory models outperforming 24-layer memoryless models while inference remains faster or largely unchanged.
- Experimental setup: The evaluation integrates memory layers into transformer models for large-scale language modeling and measures test perplexity alongside memory usage.The study uses a 28-billion-word Common Crawl news corpus with separate validation and test sets.
- Overall results: A 12-layer model with one memory outperforms a memoryless 24-layer model of the same hidden dimension, while adding two or three memory layers further improves performance.This pattern holds for hidden dimensions of 1024 and 1600.
- Speed–quality trade-off: A 12-layer 1024-dimensional model with memory achieves better perplexity than a 24-layer model and is almost twice faster at inference.For 1600-dimensional models, adding memory barely increases inference time.
- Memory size: Increasing memory size from 16k to 1M slots lowers perplexity from 22.8 to 18.0 without changing inference time.Inference time is governed by the number of accessed values, determined by memory heads and k, rather than total memory size.
- Query normalization: Batch normalization raises 1M-memory usage from 25.8% to 80.3% and lowers perplexity from 19.8 to 18.0.For small memories, usage is already close to 100% without batch normalization.
- Ablations: Memory performs best when replacing the feed-forward block at transformer layers 4 or 5, while 4 heads and 32 nearest neighbors provide a reported speed–performance trade-off.Increasing heads or nearest neighbors improves performance and memory usage, but more heads also increase computation time.
5 Conclusion
The paper presents a memory layer that substantially increases neural-network capacity with negligible computational overhead. Integrated into language models, it enables a 12-layer system to match a 24-layer BERT-large model with half the running time.
- The memory layer increases neural-network capacity with negligible computational overhead.
- A 12-layer model reaches the performance of a 24-layer BERT-large model while using half the running time.
- Product-key factorization and sparse read/write accesses provide the layer’s efficiency.