# Scaling Memory-Augmented Neural Networks with Sparse Reads and Writes

###### Abstract

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 faster and with 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 s 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.

Scaling Memory-Augmented Neural Networks with Sparse Reads and Writes

Jack W Rae^{1}^{1}1These authors contributed equally.
jwrae
Jonathan J Hunt^{1}^{1}1These authors contributed equally.
jjhunt
Tim Harley
tharley
Ivo Danihelka
danihelka
Andrew Senior
andrewsenior
Greg Wayne
gregwayne
Alex Graves
gravesa
Timothy P Lillicrap
countzero
Google DeepMind
@google.com

noticebox[b]30th Conference on Neural Information Processing Systems (NIPS 2016), Barcelona, Spain.\end@float

## 1 Introduction

Recurrent neural networks, such as the Long Short-Term Memory (LSTM) hochreiter1997long (), have proven to be powerful sequence learning models graves2013speech (); sutskever2014 (). However, one limitation of the LSTM architecture is that the number of parameters grows proportionally to the square of the size of the memory, making them unsuitable for problems requiring large amounts of long-term memory. Recent approaches, such as Neural Turing Machines (NTMs) graves2014neural () and Memory Networks weston2014memory (), have addressed this issue by decoupling the memory capacity from the number of model parameters. We refer to this class of models as memory augmented neural networks (MANNs). External memory allows MANNs to learn algorithmic solutions to problems that have eluded the capabilities of traditional LSTMs, and to generalize to longer sequence lengths. Nonetheless, MANNs have had limited success in real world application.

A significant difficulty in training these models results from their smooth read and write operations, which incur linear computational overhead on the number of memories stored per time step of training. Even worse, they require duplication of the entire memory at each time step to perform backpropagation through time (BPTT). To deal with sufficiently complex problems, such as processing a book, or Wikipedia, this overhead becomes prohibitive. For example, to store memories, a straightforward implementation of the NTM trained over a sequence of length consumes physical memory; to store memories the overhead exceeds (see Figure 1).

In this paper, we present a MANN named SAM (sparse access memory). By thresholding memory modifications to a sparse subset, and using efficient data structures for content-based read operations, our model is optimal in space and time with respect to memory size, while retaining end-to-end gradient based optimization. To test whether the model is able to learn with this sparse approximation, we examined its performance on a selection of synthetic and natural tasks: algorithmic tasks from the NTM work graves2014neural (), Babi reasoning tasks used with Memory Networks sukhbaatar2015end () and Omniglot one-shot classification santoro2016 (); lake2015human (). We also tested several of these tasks scaled to longer sequences via curriculum learning. For large external memories we observed improvements in empirical run-time and memory overhead by up to three orders magnitude over vanilla NTMs, while maintaining near-identical data efficiency and performance.

Further, in Supplementary D we demonstrate the generality of our approach by describing how to construct a sparse version of the recently published Differentiable Neural Computer graves2016dnc (). This Sparse Differentiable Neural Computer (SDNC) is over faster than the canonical dense variant for a memory size of slots, and achieves the best reported result in the Babi tasks without supervising the memory access.

## 2 Background

### 2.1 Attention and content-based addressing

An external memory is a collection of real-valued vectors, or words, of fixed size . A soft read operation is defined to be a weighted average over memory words,

(1) |

where is a vector of weights with non-negative entries that sum to one. Attending to memory is formalized as the problem of computing . A content addressable memory, proposed in graves2014neural (); weston2014memory (); bahdanau2014neural (); sukhbaatar2015end (), is an external memory with an addressing scheme which selects based upon the similarity of memory words to a given query . Specifically, for the th read weight we define,

(2) |

where is a similarity measure, typically Euclidean distance or cosine similarity, and is a differentiable monotonic transformation, typically a softmax. We can think of this as an instance of kernel smoothing where the network learns to query relevant points . Because the read operation (1) and content-based addressing scheme (2) are smooth, we can place them within a neural network, and train the full model using backpropagation.

### 2.2 Memory Networks

One recent architecture, Memory Networks, make use of a content addressable memory that is accessed via a series of read operations weston2014memory (); sukhbaatar2015end () and has been successfully applied to a number of question answering tasks weston2015towards (); hill2015goldilocks (). In these tasks, the memory is pre-loaded using a learned embedding of the provided context, such as a paragraph of text, and then the controller, given an embedding of the question, repeatedly queries the memory by content-based reads to determine an answer.

### 2.3 Neural Turing Machine

The Neural Turing Machine is a recurrent neural network equipped with a content-addressable memory, similar to Memory Networks, but with the additional capability to write to memory over time. The memory is accessed by a controller network, typically an LSTM, and the full model is differentiable — allowing it to be trained via BPTT.

A write to memory,

(3) |

consists of a copy of the memory from the previous time step decayed by the erase matrix indicating obsolete or inaccurate content, and an addition of new or updated information . The erase matrix is constructed as the outer product between a set of write weights and erase vector . The add matrix is the outer product between the write weights and a new write word , which the controller outputs.

## 3 Architecture

This paper introduces Sparse Access Memory (SAM), a new neural memory architecture with two innovations. Most importantly, all writes to and reads from external memory are constrained to a sparse subset of the memory words, providing similar functionality as the NTM, while allowing computational and memory efficient operation. Secondly, we introduce a sparse memory management scheme that tracks memory usage and finds unused blocks of memory for recording new information.

For a memory containing words, SAM executes a forward, backward step in time, initializes in space, and consumes space per time step. Under some reasonable assumptions, SAM is asymptotically optimal in time and space complexity (Supplementary A).

### 3.1 Read

The sparse read operation is defined to be a weighted average over a selection of words in memory:

(4) |

where contains number of non-zero entries with indices ; is a small constant, independent of , typically or . We will refer to sparse analogues of weight vectors as , and when discussing operations that are used in both the sparse and dense versions of our model use .

We wish to construct such that . For content-based reads where is defined by (2), an effective approach is to keep the largest non-zero entries and set the remaining entries to zero. We can compute naively in time by calculating and keeping the largest values. However, linear-time operation can be avoided. Since the largest values in correspond to the closest points to our query , we can use an approximate nearest neighbor data-structure, described in Section 3.5, to calculate in time.

Sparse read can be considered a special case of the matrix-vector product defined in (1), with two key distinctions. The first is that we pass gradients for only a constant number of rows of memory per time step, versus , which results in a negligible fraction of non-zero error gradient per timestep when the memory is large. The second distinction is in implementation: by using an efficient sparse matrix format such as Compressed Sparse Rows (CSR), we can compute (4) and its gradients in constant time and space (see Supplementary A).

### 3.2 Write

The write operation is SAM is an instance of (3) where the write weights are constrained to contain a constant number of non-zero entries. This is done by a simple scheme where the controller writes either to previously read locations, in order to update contextually relevant memories, or the least recently accessed location, in order to overwrite stale or unused memory slots with fresh content.

The introduction of sparsity could be achieved via other write schemes. For example, we could use a sparse content-based write scheme, where the controller chooses a query vector and applies writes to similar words in memory. This would allow for direct memory updates, but would create problems when the memory is empty (and shift further complexity to the controller). We decided upon the previously read / least recently accessed addressing scheme for simplicity and flexibility.

The write weights are defined as

(5) |

where the controller outputs the interpolation gate parameter and the write gate parameter . The write to the previously read locations is purely additive, while the least recently accessed word is set to zero before being written to. When the read operation is sparse ( has non-zero entries), it follows the write operation is also sparse.

We define to be an indicator over words in memory, with a value of when the word minimizes a usage measure

(6) |

If there are several words that minimize then we choose arbitrarily between them. We tried two definitions of . The first definition is a time-discounted sum of write weights where is the discount factor. This usage definition is incorporated within Dense Access Memory (DAM), a dense-approximation to SAM that is used for experimental comparison in Section 4.

The second usage definition, used by SAM, is simply the number of time-steps since a non-negligible memory access: . Here, is a tuning parameter that we typically choose to be . We maintain this usage statistic in constant time using a custom data-structure (described in Supplementary A). Finally we also use the least recently accessed word to calculate the erase matrix. is defined to be the expansion of this usage indicator where is a vector of ones. The total cost of the write is constant in time and space for both the forwards and backwards pass, which improves on the linear space and time dense write (see Supplementary A).

### 3.3 Controller

We use a one layer LSTM for the controller throughout. At each time step, the LSTM receives a concatenation of the external input, , the word, read in the previous time step. The LSTM then produces a vector, , of read and write parameters for memory access via a linear layer. The word read from memory for the current time step, , is then concatenated with the output of the LSTM, and this vector is fed through a linear layer to form the final output, . The full control flow is illustrated in Supplementary Figure 6.

### 3.4 Efficient backpropagation through time

We have already demonstrated how the forward operations in SAM can be efficiently computed in time. However, when considering space complexity of MANNs, there remains a dependence on for the computation of the derivatives at the corresponding time step. A naive implementation requires the state of the memory to be cached at each time step, incurring a space overhead of , which severely limits memory size and sequence length.

Fortunately, this can be remedied. Since there are only words that are written at each time step, we instead track the sparse modifications made to the memory at each timestep, apply them in-place to compute in time and space. During the backward pass, we can restore the state of from in time by reverting the sparse modifications applied at time step . As such the memory is actually rolled back to previous states during backpropagation (Supplementary Figure 5).

At the end of the backward pass, the memory ends rolled back to the start state. If required, such as when using truncating BPTT, the final memory state can be restored by making a copy of prior to calling backwards in time, or by re-applying the sparse updates in time.

### 3.5 Approximate nearest neighbors

When querying the memory, we can use an approximate nearest neighbor index (ANN) to search over the external memory for the nearest words. Where a linear KNN search inspects every element in memory (taking time), an ANN index maintains a structure over the dataset to allow for fast inspection of nearby points in time.

In our case, the memory is still a dense tensor that the network directly operates on; however the ANN is a structured view of its contents. Both the memory and the ANN index are passed through the network and kept in sync during writes. However there are no gradients with respect to the ANN as its function is fixed.

We considered two types of ANN indexes: FLANN’s randomized k-d tree implementation flann_pami_2014 () that arranges the datapoints in an ensemble of structured (randomized k-d) trees to search for nearby points via comparison-based search, and one that uses locality sensitive hash (LSH) functions that map points into buckets with distance-preserving guarantees. We used randomized k-d trees for small word sizes and LSHs for large word sizes. For both ANN implementations, there is an cost for insertion, deletion and query. We also rebuild the ANN from scratch every insertions to ensure it does not become imbalanced.

## 4 Results

### 4.1 Speed and memory benchmarks

We measured the forward and backward times of the SAM architecture versus the dense DAM variant and the original NTM (details of setup in Supplementary E). SAM is over times faster than the NTM when the memory contains one million words and an exact linear-index is used, and times faster with the k-d tree (Figure 1\subrefsf:speed). With an ANN the model runs in sublinear time with respect to the memory size. SAM’s memory usage per time step is independent of the number of memory words (Figure 1\subrefsf:memory), which empirically verifies the space claim from Supplementary A. For memory words SAM uses of physical memory to initialize the network and to run a 100 step forward and backward pass, compared with the NTM which consumes .

### 4.2 Learning with sparse memory access

We have established that SAM reaps a huge computational and memory advantage of previous models, but can we really learn with SAM’s sparse approximations? We investigated the learning cost of inducing sparsity, and the effect of placing an approximate nearest neighbor index within the network, by comparing SAM with its dense variant DAM and some established models, the NTM and the LSTM.

We trained each model on three of the original NTM tasks graves2014neural (). 1. Copy: copy a random input sequence of length 1–20, 2. Associative Recall: given 3-6 random (key, value) pairs, and subsequently a cue key, return the associated value. 3. Priority Sort: Given 20 random keys and priority values, return the top 16 keys in descending order of priority. We chose these tasks because the NTM is known to perform well on them.

Figure 2 shows that sparse models are able to learn with comparable efficiency to the dense models and, surprisingly, learn more effectively for some tasks — notably priority sort and associative recall. This shows that sparse reads and writes can actually benefit early-stage learning in some cases.

Full hyperparameter details are in Supplementary C.

### 4.3 Scaling with a curriculum

The computational efficiency of SAM opens up the possibility of training on tasks that require storing a large amount of information over long sequences. Here we show this is possible in practice, by scaling tasks to a large scale via an exponentially increasing curriculum.

We parametrized three of the tasks described in Section 4.2: associative recall, copy, and priority sort, with a progressively increasing difficulty level which characterises the length of the sequence and number of entries to store in memory. For example, level specifies the input sequence length for the copy task. We exponentially increased the maximum level when the network begins to learn the fundamental algorithm. Since the time taken for a forward and backward pass scales with the sequence length , following a standard linearly increasing curriculum could potentially take , if the same amount of training was required at each step of the curriculum. Specifically, was doubled whenever the average training loss dropped below a threshold for a number of episodes. The level was sampled for each minibatch from the uniform distribution over integers .

We compared the dense models, NTM and DAM, with both SAM with an exact nearest neighbor index (SAM linear) and with locality sensitive hashing (SAM ANN). The dense models contained 64 memory words, while the sparse models had words. These sizes were chosen to ensure all models use approximately the same amount of physical memory when trained over 100 steps.

For all tasks, SAM was able to advance further than the other models, and in the associative recall task, SAM was able to advance through the curriculum to sequences greater than (Figure 3). Note that we did not use truncated backpropagation, so this involved BPTT for over steps with a memory size in the millions of words.

To investigate whether SAM was able to learn algorithmic solutions to tasks, we investigated its ability to generalize to sequences that far exceeded those observed during training. Namely we trained SAM on the associative recall task up to sequences of length , and found it was then able to generalize to sequences of length (Supplementary Figure 8).

### 4.4 Question answering on the Babi tasks

weston2015towards () introduced toy tasks they considered a prerequisite to agents which can reason and understand natural language. They are synthetically generated language tasks with a vocab of about 150 words that test various aspects of simple reasoning such as deduction, induction and coreferencing.

We tested the models (including the Sparse Differentiable Neural Computer described in Supplementary D) on this task. The full results and training details are described in Supplementary G.

The MANNs, except the NTM, are able to learn solutions comparable to the previous best results, failing at only 2 of the tasks. The SDNC manages to solve all but 1 of the tasks, the best reported result on Babi that we are aware of.

Notably the best prior results have been obtained by using supervising the memory retrieval (during training the model is provided annotations which indicate which memories should be used to answer a query). More directly comparable previous work with end-to-end memory networks, which did not use supervision sukhbaatar2015end (), fails at 6 of the tasks.

Both the sparse and dense perform comparably at this task, again indicating the sparse approximations do not impair learning. We believe the NTM may perform poorly since it lacks a mechanism which allows it to allocate memory effectively.

### 4.5 Learning on real world data

Finally, we demonstrate that the model is capable of learning in a non-synthetic dataset. Omniglot lake2015human () is a dataset of 1623 characters taken from 50 different alphabets, with 20 examples of each character. This dataset is used to test rapid, or one-shot learning, since there are few examples of each character but many different character classes. Following santoro2016 (), we generate episodes where a subset of characters are randomly selected from the dataset, rotated and stretched, and assigned a randomly chosen label. At each time step an example of one of the characters is presented, along with the correct label of the proceeding character. Each character is presented 10 times in an episode (but each presentation may be any one of the 20 examples of the character). In order to succeed at the task the model must learn to rapidly associate a novel character with the correct label, such that it can correctly classify subsequent examples of the same character class.

Again, we used an exponential curriculum, doubling the number of additional characters provided to the model whenever the cost was reduced under a threshold. After training all MANNs for the same length of time, a validation task with characters was used to select the best run, and this was then tested on a test set, containing all novel characters for different sequence lengths (Figure 4). All of the MANNs were able to perform much better than chance, even on sequences longer than seen during training. SAM outperformed other models, presumably due to its much larger memory capacity. Previous results on the Omniglot curriculum santoro2016 () task are not identical, since we used 1-hot labels throughout and the training curriculum scaled to longer sequences, but our results with the dense models are comparable ( errors with characters), while the SAM is significantly better ( errors with characters).

## 5 Discussion

Scaling memory systems is a pressing research direction due to potential for compelling applications with large amounts of memory. We have demonstrated that you can train neural networks with large memories via a sparse read and write scheme that makes use of efficient data structures within the network, and obtain significant speedups during training. Although we have focused on a specific MANN (SAM), which is closely related to the NTM, the approach taken here is general and can be applied to many differentiable memory architectures, such as Memory Networks weston2014memory ().

It should be noted that there are multiple possible routes toward scalable memory architectures. For example, prior work aimed at scaling Neural Turing Machines zaremba2015reinforcement () used reinforcement learning to train a discrete addressing policy. This approach also touches only a sparse set of memories at each time step, but relies on higher variance estimates of the gradient during optimization. Though we can only guess at what class of memory models will become staple in machine learning systems of the future, we argue in Supplementary A that they will be no more efficient than SAM in space and time complexity if they address memories based on content.

We have experimented with randomized k-d trees and LSH within the network to reduce the forward pass of training to sublinear time, but there may be room for improvement here. K-d trees were not designed specifically for fully online scenarios, and can become imbalanced during training. Recent work in tree ensemble models, such as Mondrian forests lakshminarayanan2014mondrian (), show promising results in maintaining balanced hierarchical set coverage in the online setting. An alternative approach which may be well-suited is LSH forests bawa2005lsh (), which adaptively modifies the number of hashes used. It would be an interesting empirical investigation to more fully assess different ANN approaches in the challenging context of training a neural network.

Humans are able to retain a large, task-dependent set of memories obtained in one pass with a surprising amount of fidelity brady2008visual (). Here we have demonstrated architectures that may one day compete with humans at these kinds of tasks.

### Acknowledgements

We thank Vyacheslav Egorov, Edward Grefenstette, Malcolm Reynolds, Fumin Wang and Yori Zwols for their assistance, and the Google DeepMind family for helpful discussions and encouragement.

## References

- (1) Sunil Arya, David M. Mount, Nathan S. Netanyahu, Ruth Silverman, and Angela Y. Wu. An optimal algorithm for approximate nearest neighbor searching fixed dimensions. J. ACM, 45(6):891–923, November 1998.
- (2) Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473, 2014.
- (3) Mayank Bawa, Tyson Condie, and Prasanna Ganesan. Lsh forest: self-tuning indexes for similarity search. In Proceedings of the 14th international conference on World Wide Web, pages 651–660. ACM, 2005.
- (4) Timothy F Brady, Talia Konkle, George A Alvarez, and Aude Oliva. Visual long-term memory has a massive storage capacity for object details. Proceedings of the National Academy of Sciences, 105(38):14325–14329, 2008.
- (5) Ronan Collobert, Koray Kavukcuoglu, and Clément Farabet. Torch7: A matlab-like environment for machine learning. In BigLearn, NIPS Workshop, number EPFL-CONF-192376, 2011.
- (6) Alex Graves, Abdel-rahman Mohamed, and Geoffrey Hinton. Speech recognition with deep recurrent neural networks. In Acoustics, Speech and Signal Processing (ICASSP), 2013 IEEE International Conference on, pages 6645–6649. IEEE, 2013.
- (7) Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXiv:1410.5401, 2014.
- (8) Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska-Barwińska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using a neural network with dynamic external memory. Nature, 2016.
- (9) Gaël Guennebaud, Benoıt Jacob, Philip Avery, Abraham Bachrach, Sebastien Barthelemy, et al. Eigen v3, 2010.
- (10) Felix Hill, Antoine Bordes, Sumit Chopra, and Jason Weston. The goldilocks principle: Reading children’s books with explicit memory representations. arXiv preprint arXiv:1511.02301, 2015.
- (11) Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):1735–1780, 1997.
- (12) Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 350(6266):1332–1338, 2015.
- (13) Balaji Lakshminarayanan, Daniel M Roy, and Yee Whye Teh. Mondrian forests: Efficient online random forests. In Advances in Neural Information Processing Systems, pages 3140–3148, 2014.
- (14) Rajeev Motwani, Assaf Naor, and Rina Panigrahy. Lower bounds on locality sensitive hashing. SIAM Journal on Discrete Mathematics, 21(4):930–935, 2007.
- (15) Marius Muja and David G. Lowe. Scalable nearest neighbor algorithms for high dimensional data. Pattern Analysis and Machine Intelligence, IEEE Transactions on, 36, 2014.
- (16) Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and T Lillicrap. Meta-learning with memory-augmented neural networks. In International conference on machine learning, 2016.
- (17) Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, et al. End-to-end memory networks. In Advances in Neural Information Processing Systems, pages 2431–2439, 2015.
- (18) Ilya Sutskever, Oriol Vinyals, and Quoc V Le. Sequence to sequence learning with neural networks. In Advances in Neural Information Processing Systems 27, pages 3104–3112. Curran Associates, Inc., 2014.
- (19) Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning, 4:2, 2012.
- (20) Jason Weston, Antoine Bordes, Sumit Chopra, Alexander M Rush, Bart van Merriënboer, Armand Joulin, and Tomas Mikolov. Towards ai-complete question answering: A set of prerequisite toy tasks. arXiv preprint arXiv:1502.05698, 2015.
- (21) Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks. arXiv preprint arXiv:1410.3916, 2014.
- (22) Wojciech Zaremba and Ilya Sutskever. Reinforcement learning neural turing machines. arXiv preprint arXiv:1505.00521, 2015.

## Supplementary Information

## Appendix A Time and space complexity

Under a reasonable class of content addressable memory architectures , SAM is optimal in time and space complexity.

###### Definition 1.

Let be a collection of real vectors of fixed dimension . Let be the set of all content addressable memory data structures that store and can return at least one word such that for a given norm , query vector , and ; provided such a memory exists with .

Existing lower bounds [14, 1] assert that for any data structure , requires time and space to perform a read operation. The SAM memory architecture proposed in this paper is contained within as it computes the approximate nearest neighbors problem in fixed dimensions [15]. As we will show, SAM requires time to query and maintain the ANN, to perform all subsequent sparse read, write, and error gradient calculations. It requires space to initialize the memory and to store intermediate sparse tensors. We thus conclude it is optimal in asymptotic time and space complexity.

### a.1 Initialization

Upon initialization, SAM consumes space and time to instantiate the memory and the memory Jacobian. Furthermore, it requires time and space to initialize auxiliary data structures which index the memory, such as the approximate nearest neighbor which provides a content-structured view of the memory, and the least accessed ring, which maintains the temporal ordering in which memory words are accessed. These initializations represent an unavoidable one-off cost that does not recur per step of training, and ultimately has little effect on training speed. For the remainder of the analysis we will concentrate on the space and time cost per training step.

### a.2 Read

Recall the sparse read operation,

(7) |

As is chosen to be a fixed constant, it is clear we can compute (7) in time. During the backward pass, we see the gradients are sparse with only non-zero terms,

and

where is a vector of zeros. Thus they can both be computed in constant time by skipping the computation of zeros. Furthermore by using an efficient sparse matrix format to store these matrices and vectors, such as the CSR, we can represent them using at most values. Since the read word and its respective error gradient is the size of a single word in memory ( elements), the overall space complexity is per time step for the read.

### a.3 Write

Recall the write operation,

(8) |

where is the add matrix, is the erase matrix, and is defined to be the erase weight matrix. We chose the write weights to be an interpolation between the least recently accessed location and the previously read locations,

(9) |

For sparse reads where is a sparse vector with non-zeros, the write weights is also sparse with non-zeros: for the least recently accessed location and for the previously read locations. Thus the sparse-dense outer product can be performed in time as is a fixed constant.

Since can be represented as a sparse matrix with one single non-zero, the erase matrix can also. As and are sparse matrices we can then add them component-wise to the dense in time. By analogous arguments the backward pass can be computed in time and each sparse matrix can be represented in space.

We avoid caching the modified memory, and thus duplicating it, by applying the write directly to the memory. To restore its prior state during the backward pass, which is crucial to gradient calculations at earlier time steps, we roll the memory it back by reverting the sparse modifications with an additional time overhead (Supplementary Figure 5).

The location of the least recently accessed memory can be maintained in time by constructing a circular linked list that tracks the indices of words in memory, and preserves a strict ordering of relative temporal access. The first element in the ring is the least recently accessed word in memory, and the last element in the ring is the most recently modified. We keep a “head” pointer to the first element in the ring. When a memory word is randomly accessed, we can push its respective index to the back of the ring in time by redirecting a small number of pointers. When we wish to pop the least recently accessed memory (and write to it) we move the head to the next element in the ring in time.

### a.4 Content-based addressing

As discussed in Section 3.5 we can calculate the content-based attention, or read weights , in time using an approximate nearest neighbor index that views the memory. We keep the ANN index synchronized with the memory by passing it through the network as a non-differentiable member of the network’s state (so we do not pass gradients for it), and we update the index upon each write or erase to memory in time. Maintaining and querying the ANN index represents the most expensive part of the network, which is reasonable as content-based addressing is inherently expensive [14, 1].

For the backward pass computation, specifically calculating and with respect to , we can once again compute these using sparse matrix operations in time. This is because the non-zero locations have been determined during the forward pass.

Thus to conclude, SAM consumes in total space for both the forward and backward step during training, time per forward step, and per backward step.

## Appendix B Control flow

## Appendix C Training details

Here we provide additional details on the training regime used for our experiments used in Figure 2.

To avoid bias in our results, we chose the learning rate that worked best for DAM (and not SAM). We tried learning rates and found that DAM trained best with . We also tried values of and found no significant difference in performance across the values. We used 100 hidden units for the LSTM (including the controller LSTMs), a minibatch of , asynchronous workers to speed up training, and RMSProp [19] to optimize the controller. We used memory access heads and configured SAM to read from only locations per head.

## Appendix D Sparse Differentiable Neural Computer

Recently [8] proposed a novel MANN the Differentiable Neural Computer (DNC). The two innovations proposed by this model are a new approach to tracking memory freeness (dynamic memory allocation) and a mechanism for associating memories together (temporal memory linkage). We demonstrate here that the approaches enumerated in the paper can be adapted to new models by outlining a sparse version of this model, the Sparse Differentiable Neural Computer (SDNC), which learns with similar data efficiency while retaining the computational advantages of sparsity.

### d.1 Architecture

For brevity, we will only explain the sparse implementations of these two items, for the full model details refer to the original paper. The mechanism for sparse memory reads and writes was implemented identically to SAM.

It is possible to implement a scalable version of the dynamic memory allocation system of the DNC avoiding any operations by using a heap. However, because it is practical to run the SDNC with many more memory words, reusing memory is less crucial so we did not implement this and used the same usage tracking as in SAM.

The temporal memory linkage in the DNC is a system for associating and recalling memory locations which were written in a temporal order, for exampling storing and retrieving a list. In the DNC this is done by maintaining a temporal linkage matrix . represents the degree to which location was written to after location . This matrix is updated by tracking the precedence weighting , where represents the degree to which location was written to.

(10) | ||||

(11) |

The memory linkage is updated according to the following recurrence

(12) | ||||

(13) |

The temporal linkage can be used to compute read weights following the temporal links either forward

(15) |

or backward

(16) |

The read head then uses a 3-way softmax to select between a content-based read or following the forward or backward weighting.

Naively, the link matrix requires memory and computation although [8] proposes a method to reduce the computational cost to and memory cost.

In order to maintain the scaling properties of the SAM, we wish to avoid any computational dependence on . We do this by maintaining two sparse matrices that approximate and respectively. We store these matrices in Compressed Sparse Row format. They are defined by the following updates:

(17) | ||||

(18) | ||||

(19) | ||||

(20) |

Additionally, is, as with the other weight vectors maintained as a sparse vector with at most non-zero entries. This means that the outer product of has at most non-zero entries. In addition to the updates specified above, we also constrain each row of the matrices and to have at most non-zero entries — this constraint can be applied in because at most rows change in the matrix.

Once these matrices are applied the read weights following the temporal links can be computed similar to before:

(21) | ||||

(22) |

Note, the number of locations we read from, , does not have to equal the number of outward and inward links we preserve, . We typically choose as this is still very fast to compute ( in total to calculate on a single CPU thread) and we see no learning benefit with larger . In order to compute the gradients, and need to be stored. This could be done by maintaining a sparse record of the updates applied and reversing them, similar to that performed with the memory as described in Section 3.4. However, for implementation simplicity we did not pass gradients through the temporal linkage matrices.

### d.2 Results

We benchmarked the speed and memory performance of the SDNC versus a naive DNC implementation (details of setup in Supplementary E). The results are displayed in Figure 7. Here, the computational benefits of sparsity are more pronounced due to the expensive (quadratic time and space) temporal transition table operations in the DNC. We were only able to run comparative benchmarks up to , as the DNC quickly exceeded the machine’s physical memory for larger values; however even at this modest memory size we see a speed increase of and physical memory reduction of . Note, unlike the SAM memory benchmark in Section 4 we plot the total memory consumption, i.e. the memory overhead of the initial start state plus the memory overhead of unrolling the core over a sequence. This is because the SDNC and DNC do not have identical start states. The sparse temporal transition matrices consume much less memory than the corresponding in the DNC.

In order to compare the models on an interesting task we ran the DNC and SDNC on the Babi task (this task is described more fully in the main text). The results are described in Supplementary G and demonstrate the SDNC is capable of learning competitively. In particular, it achieves the best report result on the Babi task.

## Appendix E Benchmarking details

Each model contained an LSTM controller with 100 hidden units, an external memory containing slots of memory, with word size and access heads. For speed benchmarks, a minibatch size of was used to ensure fair comparison - as many dense operations (e.g. matrix multiplication) can be batched efficiently. For memory benchmarks, the minibatch size was set to .

We used Torch7 [5] to implement SAM, DAM, NTM, DNC and SDNC. Eigen v3 [9] was used for the fast sparse tensor operations, using the provided CSC and CSR formats. All benchmarks were run on a Linux desktop running Ubuntu 14.04.1 with 32GiB of RAM and an Intel Xeon E5-1650 3.20GHz processor with power scaling disabled.

## Appendix F Generalization on associative recall

## Appendix G Babi results

See the main text for a description of the Babi task and its relevance. Here we report the best and mean results for all of the models on this task.

The task was encoded using straightforward 1-hot word encodings for both the input and output. We trained a single model on all of the tasks, and used the 10,000 examples per task version of the training set (a small subset of which we used as a validation set for selecting the best run and hyperparameters). Previous work has reported best results (Supplementary table 1), which with only 15 runs is a noisy comparison, so we additionally report the mean and variance for all runs with the best selected hyperparameters (Supplementary table 2).

LSTM | DNC | SDNC | DAM | SAM | NTM | MN S | MN U | |

1: 1 supporting fact | ||||||||

2: 2 supporting facts | ||||||||

3: 3 supporting facts | ||||||||

4: 2 argument relations | ||||||||

5: 3 argument relations | ||||||||

6: yes/no questions | ||||||||

7: counting | ||||||||

8: lists/sets | ||||||||

9: simple negation | ||||||||

10: indefinite knowledge | ||||||||

11: basic coreference | ||||||||

12: conjunction | ||||||||

13: compound coreference | ||||||||

14: time reasoning | ||||||||

15: basic deduction | ||||||||

16: basic induction | ||||||||

17: positional reasoning | ||||||||

18: size reasoning | ||||||||

19: path finding | ||||||||

20: agent’s motivations | ||||||||

Mean Error (%) | ||||||||

Failed tasks (err. > 5%) |

LSTM | DNC | SDNC | DAM | SAM | NTM | |

1: 1 supporting fact | ||||||

2: 2 supporting facts | ||||||

3: 3 supporting facts | ||||||

4: 2 argument relations | ||||||

5: 3 argument relations | ||||||

6: yes/no questions | ||||||

7: counting | ||||||

8: lists/sets | ||||||

9: simple negation | ||||||

10: indefinite knowledge | ||||||

11: basic coreference | ||||||

12: conjunction | ||||||

13: compound coreference | ||||||

14: time reasoning | ||||||

15: basic deduction | ||||||

16: basic induction | ||||||

17: positional reasoning | ||||||

18: size reasoning | ||||||

19: path finding | ||||||

20: agent’s motivations | ||||||

Mean Error (%) | ||||||

Failed tasks (err. > 5%) |