Large Memory Layers with Product Keys

By Guillaume Lample, Alexandre Sablayrolles, Marc'Aurelio Ranzato, Ludovic Denoyer, Herv'e J'egou. 2019

181 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Model Architectures

Citations in the corpus

Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks (2020, background): “There has been extensive previous work proposing architectures to enrich systems with non-parametric memory which are trained from scratch for speciﬁc tasks—e.g. in memory networks [58, 49], stack-augmented networks [21] and memory layers for transformers [26].”
REALM: Retrieval-Augmented Language Model Pre-Training (2020, background): “From this view, our work share motivations with works such as product key memory (Lample et al., 2019), which enables sub-linear memory access in a memory network (Weston et al.”
CTRL: A Conditional Transformer Language Model for Controllable Generation (2019, background): “Recent work on memory mechanisms (Dai et al., 2019; Lample et al., 2019) has improved perplexities on the most common benchmarks, and even without these memories, large Transformer architectures (Vaswani et al.”
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (2021, background): “Finally, Product Key networks (Lample et al., 2019) were proposed to scale up the capacity of neural networks by doing a lookup for learnable embeddings based on the incoming token representations to a given layer.”
Compressive Transformers for Long-Range Sequence Modelling (2019, background): “Several works have focused on reducing the computational cost of attention with sparse access mechanisms (Rae et al., 2016; Child et al., 2019; Sukhbaatar et al., 2019; Lample et al., 2019).”
Efficient Content-Based Sparse Attention with Routing Transformers (2020, methodology): “Another related work is (Lample et al., 2019) who use product quantization based key-value lookups to replace the feed forward network in the Transformer.”
Reformer: The Efficient Transformer (2020, methodology): “Using product-key attention to increase the key space has also been used to reduce memory requirements in the feed-forward layers with no loss in performance (Lample et al., 2019).”
Entities as Experts: Sparse Memory Access with Entity Supervision (2020, background): “Another approach to scale memory networks is given by Lample et al. (2019) who introduce product-key memories to ef-ﬁciently index a large store of values.”
The Efficiency Misnomer (2021, background): “Sparse Models and Mixture-of-Experts A deﬁning feature of sparse models (Lample et al., 2019; Fedus et al., 2021) is that they remain compute-matched while enabling scaling to a large number of parameters.”
Controlling Computation versus Quality for Neural Sequence Models (2020, positive): “This also corroborates findings in previous work exploring where high capacity layers should be used in Transformer-based language models (Lample et al., 2019).”
Do Transformer Modifications Transfer Across Implementations and Applications? (2021, background): “Product key memories (Lample et al., 2019) should only slightly increase FLOPs over the vanilla Trans-former, but the total number of operations is artiﬁcially extremely high due to an ineﬃcient implementation in Mesh Tensorﬂow.”
BASE Layers: Simplifying Training of Large, Sparse Models (2021, background): “For example, Lample et al. (2019) introduce a large memory layer that supports efﬁcient sparse queries.”
Efficient Transformers: A Survey (2020, methodology): “…2021) O ( kN ) M TokenLearner (Ryoo et al., 2021) O ( k 2 ) DS Adaptive Sparse Transformer (Correia et al., 2019) O ( N 2 ) Sparse Product Key Memory (Lample et al., 2019) O ( N 2 ) Sparse Switch Transformer (Fedus et al., 2021) O ( N 2 ) Sparse ST-MoE (Zoph et al., 2022) O ( N 2 ) Sparse…”
Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention (2020, methodology): “Lample et al. (2019) used product-key attention to increase the capacity of any layer with negligible computational overhead.”
Hash Layers For Large Sparse Models (2021, methodology): “[15] uses a different approach, where product keys enable nearest neighbor search to select parameters.”
Nyströmformer: A Nystöm-based Algorithm for Approximating Self-Attention (2021, background): “2019), product-key attention in (Lample et al. 2019), and the Transformer-XL model in (Dai et al.”
Augmenting Transformers with KNN-Based Composite Memory for Dialog (2020, background): “Lample et al. (2019) focus on learning memories of up to one million slots and how to efﬁciently access the slots using product keys.”
Learning Knowledge Bases with Parameters for Task-Oriented Dialogue Systems (2020, background): “Moreover, even with the most advanced transformer architecture (Kitaev et al., 2020; Lample et al., 2019; Child et al., 2019), end-to-end models struggle when the input becomes too large (Neelakantan et al.”