Scaling Memory-Augmented Neural Networks with Sparse Reads and Writes

By Jack W. Rae, Jonathan J. Hunt, Ivo Danihelka, Tim Harley, Andrew W. Senior, Greg Wayne, Alex Graves, Timothy P. Lillicrap. 2016

171 citations recorded by the source corpus.

Neural networks augmented with external memory have the ability to learn algorithmic solutions to complex tasks. These models appear promising for applications such as language modeling and machine translation. However, they scale poorly in both space and time as the amount of memory grows — limiting their applicability to real-world domains. Here, we present an end-to-end differentiable memory access scheme, which we call Sparse Access Memory (SAM), that retains the representational power of the original approaches whilst training efficiently with very large memories. We show that SAM achieves asymptotic lower bounds in space and time complexity, and find that an implementation runs 1,000 x faster and with 3,000x less physical memory than non-sparse models. SAM learns with comparable data efficiency to existing models on a range of synthetic tasks and one-shot Omniglot character recognition, and can scale to tasks requiring 100,000s of time steps and memories. As well, we show how our approach can be adapted for models that maintain temporal associations between memories, as with the recently introduced Differentiable Neural Computer.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, 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 Computer Vision and Pattern Recognition

In Model Architectures

Citations in the corpus

Universal Transformers (2018, neutral): “3 (0/20) Sparse DNC (Rae et al., 2016) 2.”
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): “Memory augmented neural networks with sparse reads and writes have also been proposed in (Rae et al., 2016) as a way to scale Neural Turing Machines (Graves et al., 2014).”
Reformer: The Efficient Transformer (2020, background): “The requirement that the memory be ﬁxed before has been removed in Santoro et al. (2016) at the cost of memory size and later alleviated by Rae et al. (2016).”
A simple neural network module for relational reasoning (2017, neutral): “When training jointly on all tasks using 10K examples per task, Memory Networks pass 14/20, DNC 16/20, Sparse DNC 19/20, and EntNet 16/20 (the authors of EntNets report state-of-the-art at 20/20; however, unlike previously reported results this was not done with joint training on all tasks, where th”
Evaluating Theory of Mind in Question Answering (2018, background): “models fail at only one or two (out of 20) tasks (Rae et al., 2016; Santoro et al., 2017).”
Categorical Reparameterization with Gumbel-Softmax (2016, background): “Discrete representations are often more interpretable (Chen et al., 2016) and more computationally efﬁcient (Rae et al., 2016) than their continuous analogues.”
Fast Parametric Learning with Activation Memorization (2018, background): “Although the cache is of fixed size n, it can be defined to be very large with sparse attention and efficient data-structures (Rae et al., 2016; Kaiser et al., 2017; Grave et al., 2017).”
Recurrent Relational Networks (2017, neutral): “7 LSTM [Rae et al., 2016] 15 28 .”
Augmenting Transformers with KNN-Based Composite Memory for Dialog (2020, neutral): “For example, Chandar et al. (2016) propose a hierarchical memory network rather than a ﬂat one and Rae et al. (2016) learn sparse operations to read and write.”
Dynamic Neural Turing Machine with Soft and Hard Addressing Schemes (2016, methodology): “ Least Recently Used Addressing We introduce a memory addressing operation that can learn to put more emphasis on the least recently used (LRU) memory locations. As observed in (Santoro et al., 2016; Rae et al., 2016), we ﬁnd it easier to learn the write operations with the use of LRU addressing. To”
Hierarchical Memory Networks (2016, methodology): “While Rae et al. (2016) focus on architectures where the memory is written by the controller itself, we focus on handling memory access to large external knowledge bases.”
Practice on Long Sequential User Behavior Modeling for Click-Through Rate Prediction (2019, methodology): “Previous works [22, 23] in NLP area have proposed to use the LRU strategy to balance the utilization of each memory.”
Towards Neural Theorem Proving at Scale (2018, methodology): “Similarly to this work, Rae et al. (2016) also rely on approximate nearest neighbours to speed up Memory-Augmented neural networks.”
MEMO: A Deep Network for Flexible Combination of Episodic Memories (2020, methodology): “A more recent extension incorporated sparsity into the DNC(Rae et al., 2016), allowing the model to perform well at larger-scale tasks, including the bAbI task suite (Weston et al., 2015).”
\infty-former: Infinite Memory Transformer (2021, background): “…et al. (2014), Weston et al. (2014), Joulin and Mikolov (2015) and Grefenstette et al. (2015) proposed extending RNNs with an external memory, while Chandar et al. (2016) and Rae et al. (2016) proposed efﬁ-cient procedures to read and write from these memories, using hierarchies and sparsity.”
Recurrent Memory Transformer (2022, background): “NTMs followed by Differentiable Neural Computer (DNC) (Graves et al., 2016) and Sparse DNC (Rae et al., 2016) are implemented as recurrent neural networks able to write to memory storage over time.”
Learning to Reason with Third-Order Tensor Products (2018, methodology): “Task REN [51] DNC [49] SDNC [50] TPR-RNN (ours) Avg Error (%) 9.”