Abstract
This paper introduces a structured memory which can be easily integrated intoa neural network. The memory is very large by design and thereforesignificantly increases the capacity of the architecture, by up to a billionparameters with a negligible computational overhead. Its design and accesspattern is based on product keys, which enable fast and exact nearest neighborsearch. The ability to increase the number of parameters while keeping the samecomputational budget lets the overall system strike a better trade-off betweenprediction 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 weplug our memory layer in a state-of-the-art transformer-based architecture. Inparticular, we found that a memory augmented model with only 12 layersoutperforms a baseline transformer model with 24 layers, while being twicefaster at inference time. We release our code for reproducibility purposes.