Source-linked AI summary
Scaling Memory-Augmented Neural Networks with Sparse Reads and Writes
Jack W Rae, Jonathan J Hunt, Tim Harley, Ivo Danihelka, Andrew Senior, Greg Wayne, Alex Graves, Timothy P Lillicrap
TL;DR
Large external memories make memory-augmented neural networks expensive in space and time, limiting their practical use. The paper introduces Sparse Access Memory, which restricts reads and writes to sparse subsets while retaining differentiable training, and reports substantial efficiency gains with comparable learning performance. It also applies the approach to a sparse Differentiable Neural Computer.
Problem
Memory-augmented neural networks scale poorly in space and time as memory grows, limiting their applicability to real-world domains.
Method
SAM thresholds memory modifications to sparse subsets and uses efficient data structures for content-based reads while retaining end-to-end gradient optimization.
Results
Across synthetic and natural tasks, SAM improves runtime and memory overhead by up to three orders of magnitude while maintaining near-identical data efficiency and performance.
Takeaways & Limitations
Sparse access enables training with large memories and extends to a sparse Differentiable Neural Computer, including tasks with longer sequences and memories.
Takeaways & Limitations
The paper notes that the approximate-nearest-neighbor structures used for sublinear forward passes may become imbalanced during online training and require further empirical assessment.
Abstract
from arXiv · showhide
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\times$ faster and with $3,\!000\times$ 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,\!000$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.
1 Introduction
Memory-augmented neural networks separate memory capacity from model parameters, enabling algorithmic solutions beyond traditional LSTMs but introducing prohibitive scaling costs. SAM addresses these costs with sparse memory access while retaining differentiable training and comparable task performance.
- Motivation: MANNs decouple external memory capacity from model parameters, allowing algorithmic solutions and longer-sequence generalization beyond traditional LSTMs.LSTM parameters grow with the square of memory size, whereas MANNs use external memory.
- Scaling challenge: Smooth reads and writes impose linear computational overhead in stored memories and require duplicating the entire memory during BPTT.These costs become prohibitive for complex inputs such as books or Wikipedia.
- Scaling challenge: 64 memories consume ≈30 MiB physical memory, while 64,000 memories exceed 29 GiB in a straightforward NTM implementation.The comparison uses sequences of length 100.
- Approach: SAM thresholds memory modifications to a sparse subset and uses efficient content-based read structures, targeting optimal space and time while preserving end-to-end gradient optimization.The evaluation covers synthetic algorithmic tasks, bAbI reasoning, Omniglot one-shot classification, and longer sequences through curriculum learning.
- Results: SAM improves empirical runtime and memory overhead by up to three orders of magnitude over vanilla NTMs while maintaining near-identical data efficiency and performance.The paper also describes a sparse Differentiable Neural Computer that is over 400× faster at 2,000 memory slots.
2 Background
Content-addressable memories read by weighting stored vectors according to their similarity to a query, with differentiable operations enabling backpropagation. Neural Turing Machines extend this framework by allowing a recurrent controller to write to memory over time.
- Memory and attention: An external memory is a collection of N real-valued vectors, or words, of fixed size M, and a soft read computes a weighted average over them.The read weights are non-negative and sum to one.
- Memory and attention: Content-based addressing computes read weights from the similarity between each memory word and a query q.The similarity measure is typically Euclidean distance or cosine similarity, followed by a differentiable monotonic transformation such as softmax.
- Differentiability: Because soft reads and content-based addressing are smooth, the complete memory system can be trained with backpropagation.The differentiable formulation places these operations inside a neural network.
- Memory Networks: Memory Networks repeatedly query a pre-loaded content-addressable memory to answer questions from an embedded context.The controller receives an embedding of the question and performs successive content-based reads.
- Neural Turing Machines: Neural Turing Machines combine a recurrent controller with content-addressable memory and add the ability to write to memory over time.The full model is differentiable and can be trained through BPTT.
- Neural Turing Machines: NTM memory updates retain the previous memory after erasing obsolete content and then add new or updated information.The erase and add operations are constructed from write weights combined with erase or write vectors.
3 Architecture
SAM makes external-memory access sparse, using constant-size reads and writes with efficient indexing and memory-state rollback. This yields logarithmic time and constant per-step space while retaining differentiable optimization.
- Sparse access: SAM constrains reads and writes to sparse memory subsets and adds memory management to locate unused blocks for new information.The controller updates previously read locations or the least recently accessed location.
- Complexity: Θ(log N) forward and backward time, Θ(N) initialization space, and Θ(1) space per timestep make SAM asymptotically optimal under stated assumptions.An approximate-nearest-neighbor index provides the logarithmic access cost, while sparse updates support the space bound.
- Sparse read: K=4 or K=8 non-zero read weights are retained independently of memory size, and approximate nearest-neighbor search finds them in O(log N) time.The K closest points to the query correspond to the largest read weights, avoiding an O(N) scan.
- Sparse read: Sparse reads pass gradients through only K memory rows and use sparse matrix formats to compute reads and gradients in constant time and space.This produces a negligible fraction of non-zero error gradients per timestep when memory is large.
- Sparse write: SAM's write operation has constant forward and backward time and space by maintaining usage statistics and writing to previously read or least recently accessed locations.The least recently accessed word is erased before receiving fresh content, while writes to previously read locations are additive.
- Backpropagation: O(T) space replaces naive O(NT) memory-state caching by tracking sparse modifications and reverting them during backpropagation.Applying sparse updates in place computes each memory state in O(1) time and O(T) space.
4 Results
SAM preserves learning performance while making sparse memory access practical at much larger scales. It trains efficiently on synthetic reasoning tasks, long curricula, Babi, and one-shot Omniglot recognition.
- Speed and memory benchmarks: 1600× faster and 3700× less memory: SAM reduces the cost of training with 1M or 64,000 memory words.At 1M memories, NTM takes 12 s versus 7 ms for SAM; at 64,000 words, NTM uses 29 GiB versus SAM’s 7.8 MiB.
- Learning with sparse memory access: SAM learns comparably to dense models on Copy and reaches asymptotic error faster on Associative Recall and Priority Sort.Sparse reads and writes do not impair learning and can improve early-stage learning on some tasks.
- Scaling with a curriculum: SAM advances through curricula requiring thousands of time steps and thousands of stored words, including Associative Recall sequences longer than 4000.This training used full backpropagation through time rather than truncated backpropagation.
- Scaling with a curriculum: SAM trained on sequences up to 10,000 and generalized Associative Recall to sequences of length 200,000.The result tests whether SAM learns algorithmic solutions beyond the sequence lengths observed during training.
- Question answering on the Babi tasks: The SDNC solves all but one Babi task, while the other MANNs achieve results comparable to prior best results and fail at only two tasks.The SDNC result is reported as the best known Babi result in the supplied passages.
5 Discussion
The paper argues that sparse differentiable memory access can scale neural memory systems while remaining applicable beyond SAM. It identifies approximate-nearest-neighbor indexing as an area for further improvement.
- Scalable memory architectures: Sparse reads and writes with efficient data structures enable large-memory neural networks and can generalize to architectures such as Memory Networks.The discussion presents this as a route to significant training speedups rather than a restriction to SAM alone.
- Limitations and future work: The paper does not establish which scalable memory architecture will become standard, while arguing that alternatives should not surpass SAM’s space and time efficiency.The comparison is framed as an unresolved direction rather than a settled architectural conclusion.
- Limitations and future work: The tested randomized k-d trees and LSH reduce forward-pass cost sublinearly, but their online behavior may still be improved.The paper specifically notes that k-d trees can become imbalanced during training and suggests Mondrian forests or LSH forests as alternatives.
Supplementary Information
SAM achieves asymptotically optimal time and space complexity for a reasonable class of content-addressable memories by restricting reads and writes to sparse operations. Its memory-efficient backpropagation stores and reverts sparse modifications instead of duplicating the full memory.
- A Time and space complexity: SAM is optimal in time and space complexity under a reasonable class of content-addressable memory architectures.
- A Time and space complexity: Content-based reads use an approximate-nearest-neighbor index, requiring O(log N) time, while subsequent sparse reads, writes, and gradient calculations take O(1) time.
- A.2 Read: Sparse reads have only K non-zero gradient terms, allowing constant-time computation and O(1) space per time step.
- A.3 Write: Sparse writes update previously read locations or the least recently accessed location, with sparse add and erase matrices applied in O(1) time.
- A.3 Write: During training, SAM consumes O(1) space per time step by storing only sparse memory modifications and reverting them during backpropagation.
- B Control flow: The controller’s LSTM output is projected into memory read and write parameters; the read result contributes to the output and the next controller input.
C Training details
The experiments used a fixed training setup for Figure 2, selecting the learning rate based on DAM rather than SAM and testing several sparsity levels.
- C Training details: The learning rate was selected for DAM to avoid bias, with 10−5 producing the best DAM training performance.
- C Training details: Testing K values of 4, 8, and 16 found no significant performance differences.
- C Training details: Training used 100 hidden units, minibatches of 8, eight asynchronous workers, and RMSProp optimization.
D Sparse Differentiable Neural Computer
The sparse access approach extends to the Differentiable Neural Computer, producing an SDNC that preserves sparsity’s computational advantages while retaining similar data efficiency.
- D Sparse Differentiable Neural Computer: The Sparse Differentiable Neural Computer adapts SAM’s sparse methods to the Differentiable Neural Computer.
- D Sparse Differentiable Neural Computer: SDNC learns with similar data efficiency while retaining the computational advantages of sparsity.
D.1 Architecture
The SDNC approximates dense temporal memory linkage with sparse structures while preserving scalable memory access. Its implementation omits gradient propagation through temporal linkage matrices for simplicity.
- D.1 Architecture: DNC temporal linkage associates memory locations written in temporal order through a linkage matrix and supports forward or backward linked reads.
- D.1 Architecture: Naive temporal linkage requires O(N^2) memory and computation, while the DNC offers an O(N log N)-time, O(N)-memory alternative.
- D.1 Architecture: SDNC avoids dependence on N by maintaining sparse matrices Nt and Pt with KL columns that approximate the temporal linkage matrix and its transpose.
- D.1 Architecture: The sparse temporal-linkage implementation constrains relevant matrix rows to at most KL non-zero entries and computes updates using sparse operations.
- D.1 Architecture: KL is typically set to 8 because the computation takes 100µs on one CPU thread and larger values provide no observed learning benefit.
- D.1 Architecture: For implementation simplicity, SDNC does not pass gradients through the temporal linkage matrices.
D.2 Results
SDNC substantially reduces runtime and physical memory relative to a naive DNC implementation, while remaining competitive on the Babi task.
- The sparse temporal transition matrices in SDNC consume much less memory than the corresponding dense transition matrix in DNC.
- SDNC benchmarks compared wall-clock time for one forward and backward pass and total memory usage during training over 10 time steps.
- SDNC learned competitively on the Babi task and achieved the best reported result.
E Benchmarking details
The benchmarking setup used controlled model and hardware configurations, while generalization tests evaluated SAM on sequences far longer than those seen during training.
- Each benchmark model used an LSTM controller with 100 hidden units, N memory slots, 32-dimensional words, and 4 access heads.Speed tests used minibatches of 8, whereas memory tests used minibatches of 1.
- Implementations used Torch7 and Eigen v3 sparse tensor operations on a Linux desktop with 32 GiB of RAM and an Intel Xeon E5-1650 processor.
- SAM models trained to sequence length 10,000 performed much better than chance on associative-recall sequences of length 200,000.The evaluation included SAM variants with and without approximate nearest-neighbor search.
G Babi results
On the jointly trained Babi benchmark, SAM and DAM solved all but two tasks, while SDNC solved all but one and achieved the best reported result with unsupervised memory access.
- The Babi experiments used straightforward 1-hot word encodings and 10,000 training examples per task version.A small subset was used for validation and selecting the best run and hyperparameters.
- SDNC achieved the best reported Babi result with unsupervised memory access, solving all but 1 task.
- SAM and DAM passed all but 2 Babi tasks without supervision of their memory accesses.
- Models were trained and tested jointly across all Babi tasks, with approximately equal training resources for each task.
- Mean and variance test errors were reported across 15 runs using the best hyperparameter set selected on validation data.