Large Memory Layers with Product Keys

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.

arXiv commentAdvances in Neural Information Processing Systems, 2019

Cited by in corpus (30)

• Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks
• CTRL: A Conditional Transformer Language Model for Controllable Generation
• REALM: Retrieval-Augmented Language Model Pre-Training
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• Reformer: The Efficient Transformer
• NAS-BERT: Task-Agnostic and Adaptive-Size BERT Compression with Neural Architecture Search
• Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention
• Generating Radiology Reports via Memory-driven Transformer
• Zero-shot Entity Linking with Efficient Long Range Sequence Modeling
• Augmenting Transformers with KNN-Based Composite Memory for Dialogue
• Towards Crowdsourced Training of Large Neural Networks using Decentralized Mixture-of-Experts
• Pruning Self-attentions into Convolutional Layers in Single Path
• Controlling Computation versus Quality for Neural Sequence Models
• Memory-Based Graph Networks
• Vector Quantized Contrastive Predictive Coding for Template-based Music Generation
• PIC: Permutation Invariant Convolution for Recognizing Long-range Activities
• Transformers for Limit Order Books
• Efficient Transformers in Reinforcement Learning using Actor-Learner Distillation
• Not All Memories are Created Equal: Learning to Forget by Expiring
• Multi-scale Transformer Language Models
• NWT: Towards natural audio-to-video generation with representation learning
• SMYRF: Efficient Attention using Asymmetric Clustering
• WaLDORf: Wasteless Language-model Distillation On Reading-comprehension
• Beyond Nyströmformer -- Approximation of self-attention by Spectral Shifting
• On-The-Fly Information Retrieval Augmentation for Language Models
• Knowledge Efficient Deep Learning for Natural Language Processing
• Large Product Key Memory for Pretrained Language Models
• The Efficiency Misnomer
• Differentiable Random Access Memory using Lattices
• CNN with large memory layers