Source-linked AI summary
Relational recurrent neural networks
Adam Santoro, Ryan Faulkner, David Raposo, Jack Rae, Mike Chrzanowski, Theophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, Timothy Lillicrap
TL;DR
Memory-based neural networks can retain information over long periods, but their ability to perform complex relational reasoning with that information is unclear. The paper introduces a Relational Memory Core that enables memory interactions through multi-head dot product attention and evaluates it across sequential reasoning tasks, reporting improved performance across reinforcement learning, program evaluation, comparative reasoning, and language modeling.
Problem
The paper addresses whether memory-based neural networks can perform complex relational reasoning with remembered information, since standard architectures may struggle to model how memories are connected.
Method
The paper introduces a Relational Memory Core that uses multi-head dot product attention to let memory vectors interact and applies it across sequential reasoning tasks.
Results
Explicit modeling of memory interactions improves performance across reinforcement learning, program evaluation, comparative reasoning, and language modeling.
Takeaways & Limitations
The results demonstrate the value of instilling a capacity for relational reasoning in recurrent neural networks.
Takeaways & Limitations
The authors cannot make concrete claims about the causal influence of their design choices on relational reasoning capacity or about the computations performed within the model.
Abstract
from arXiv · showhide
Memory-based neural networks model temporal data by leveraging an ability to remember information for long periods. It is unclear, however, whether they also have an ability to perform complex relational reasoning with the information they remember. Here, we first confirm our intuitions that standard memory architectures may struggle at tasks that heavily involve an understanding of the ways in which entities are connected -- i.e., tasks involving relational reasoning. We then improve upon these deficits by using a new memory module -- a \textit{Relational Memory Core} (RMC) -- which employs multi-head dot product attention to allow memories to interact. Finally, we test the RMC on a suite of tasks that may profit from more capable relational reasoning across sequential information, and show large gains in RL domains (e.g. Mini PacMan), program evaluation, and language modeling, achieving state-of-the-art results on the WikiText-103, Project Gutenberg, and GigaWord datasets.
1 Introduction
The paper argues that memory-based sequence models should explicitly model interactions among memories to improve relational reasoning over time. It introduces the Relational Memory Core and applies it to reinforcement learning, program evaluation, and language modeling.
- Existing memory architectures store and retrieve information across time but are not explicitly biased toward modeling interactions among memories.The paper motivates considering memory interactions alongside storage and retrieval.
- The Relational Memory Core uses multi-head dot product attention to allow memories to interact with one another.
- The RMC is evaluated on partially observed reinforcement learning, program evaluation, and language modeling tasks.The language-modeling datasets include WikiText-103, Project Gutenberg, and GigaWord.
2 Relational reasoning
Relational reasoning involves understanding how entities are connected and using those relations toward a higher-order goal. The paper examines whether neural architectures’ inductive biases limit their ability to compare and relate information across time.
- Relational reasoning uses connections among entities to accomplish a higher-order goal rather than reasoning about each entity in isolation.The paper illustrates this with comparing tree distances to a park bench.
- Relational reasoning can be framed across a spectrum of neural-network inductive biases, including convolutional kernels and message-passing networks.
- In temporal settings, relational reasoning can compare and contrast information seen at different points in time.Attention can relate previous hidden states using content rather than proximity.
- Memory-augmented networks may compartmentalize memories without enabling them to interact, while LSTMs may make compartmentalization and relational reasoning harder by storing information in one hidden vector.
3 Model
The RMC combines compartmentalized memory slots with attention-based interactions, then embeds these updates into recurrent processing. Its design supports configurable memory capacity and multiple attention mechanisms.
- 3.1 Allowing memories to interact using multi-head dot product attention: The RMC uses attention to let fixed memory slots interact rather than only storing and retrieving compartmentalized information.Each memory attends over other memories and updates its content using attended information.
- 3.1 Allowing memories to interact using multi-head dot product attention: Queries, keys, and values are linearly projected from each memory, and scaled dot-product attention produces weighted information transfers between memories.The resulting matrix has the same dimensionality as the original memory and acts as a proposed update.
- 3.1 Allowing memories to interact using multi-head dot product attention: Multiple attention heads independently compute memory interactions using distinct projection parameters before combining their outputs.For an N × F memory with two heads, each head produces an N × F/2 output.
- 3.2 Encoding new memories: The same attention operation can incorporate a new input while preserving the memory matrix size.Keys and values use the concatenated memory and input, while queries use only the existing memory.
- 3.3 Introducing recurrence and embedding into an LSTM: The memory updates are embedded into an LSTM-like recurrent structure, with row-wise shared parameters allowing the number and size of memories to be tuned separately.The authors report that some tasks need more memories, whereas language modeling needs fewer, larger memories.
4 Experiments
The experiments evaluate the RMC on supervised, reinforcement-learning, and language-modeling tasks requiring relational reasoning across sequential information. The suite includes synthetic distance reasoning, program execution, partially observed navigation, and large text datasets.
- Task suite: The task suite includes supervised and reinforcement-learning evaluations, with language modeling used to compare performance on substantial natural-data benchmarks.The experiments also reference BoxWorld as another reinforcement-learning task requiring relational reasoning in memory space.
- 4.1 Illustrative supervised tasks: The Nth Farthest task requires computing pairwise distances to a reference vector and sorting those relations across a sequence.The model must identify the nth farthest vector, even when the reference may be stored in memory or arrive later.
- 4.1 Illustrative supervised tasks: Learning to Execute evaluates whether models can transform sequences of pseudo-code characters into the numeric character sequences produced by program execution.The dataset covers addition, control, and full-program categories.
- 4.2 Reinforcement learning tasks: Mini Pacman uses a 5 × 5 viewport, requiring the agent to remember ghost dynamics and previously collected food while navigating the maze.The viewport makes the task partially observable.
- 4.3 Language modeling: Language modeling is tested on WikiText-103, Project Gutenberg, and GigaWord, which provide approximately 100M, 180M, and 4B training tokens respectively.These datasets contain contiguous sentences and differ in scale and writing style.
5 Results
Across relational reasoning, program evaluation, reinforcement learning, and language modeling, the RMC generally outperformed or matched recurrent memory baselines. Its strongest gains appeared on the Nth Farthest task, Mini Pacman, and language-modeling perplexity, though the harder Nth Farthest setting was less robust.
- 5.1 N th Farthest: 91% end-of-training accuracy was achieved by the RMC on the 16-dimensional Nth Farthest task, versus below 30% for LSTM and DNC baselines.The RMC reached similar performance with 32-dimensional vectors, but success was limited to fewer seeds and model configurations.
- 5.1 N th Farthest: The RMC changed its attention after observing the reference vector, with memory slots focusing on memories associated with that vector.Before the reference appeared, one or two slots mainly attended to the input.
- 5.2 Program Evaluation: The RMC performed at least as well as all baselines on each Learning to Execute task and surpassed equivalent teacher-forced results from prior work.It was marginally surpassed on double memorization, where both models effectively solved the task.
- 5.3 Mini-Pacman: 677 vs. 550 points were achieved by the RMC versus LSTM with viewport observations, while full observations produced 1159 vs. 598.The comparison concerns Mini Pacman performance under partial and full visual observations.
- 5.4 Language Modeling: 1.4–5.4 lower perplexity was observed across the three language-modeling tasks than the best published results, a 5–12% relative improvement.The improvement was consistent across datasets with varying size and style.
- 5.4 Language Modeling: The RMC showed slightly better data efficiency than an LSTM and performed well with relatively few context words.Its overall perplexity improvement came from better modeling of frequent words.
6 Discussion
The discussion argues that explicitly modeling interactions among memories improves relational reasoning across several sequential tasks, while cautioning that the underlying causal mechanisms remain uncertain.
- Growing buffers can capture long-distance interactions but face scaling and implementation challenges when sequence length is large or unbounded.
- The results provide evidence of improved function, but do not establish the causal influence of specific design choices on relational reasoning.
- The model’s design uses memory slots and attention to support both information compartmentalization and interactions between compartmentalized memories.
- The number of memories and individual memory size involve a task-specific trade-off rather than a universally best configuration.Program evaluation favored an intermediate four-slot, four-head configuration across subtasks, although some subtasks performed best with one or eight memories.
- Explicit memory interactions improved performance in reinforcement learning, program evaluation, comparative reasoning, and language modeling.
A Further task details, analyses, and model configurations
This section defines the terminology used to describe the model’s memory and attention configuration.
- “total units” denotes the total number of elements in memory matrix M.It equals memory size multiplied by the number of memories.
- “num heads” denotes the number of attention heads producing distinct query, key, and value sets for memories.
- “memory slots” or “number of memories” denotes the number of rows in matrix M.
- “num blocks” denotes the number of attention iterations performed at each time-step.
- “gate style” denotes whether gating is applied per unit or per memory slot.
A.1 N th Farthest
The N th Farthest experiments compare RMC, LSTM, and DNC configurations under controlled training and memory settings.
- Inputs were sequences of eight randomly sampled 16-dimensional vectors with randomly assigned one-hot labels and task specifications.
- All models used Adam, batch size 1600, tuned learning rates from 1e−5 to 1e−3, and softmax cross-entropy loss.The final learning rate was 1e−4 because learning rate did not appear to influence performance.
- LSTM and DNC architecture parameters did not appear to affect performance across the configurations tested.LSTM hidden sizes ranged from 64 to 4096 units, while DNC settings varied memories, memory sizes, and read-write counts.
- RMC experiments varied 1, 8, or 16 memories and 1, 8, or 16 attention heads while keeping 2048 total memory units.The sweep used one attention block and both unit and memory gating methods.
A.2 Program Evaluation
The program-evaluation experiments use recurrent encoder-decoder models on algorithmic pseudo-code tasks, with comparisons across memory architectures and training regimes.
- A.2 Program Evaluation: The Learning to Execute dataset contains pseudo-code snippets whose character-sequence inputs must be mapped to execution-output character sequences.Its addition, control, and full-program categories involve symbolic manipulation of variables and were selected to test relational reasoning.
- A.2 Program Evaluation: Training samples varied literal length and nesting depth, sampled uniformly up to task-specific maxima under a Mix curriculum.
- A.2 Program Evaluation: The encoder processes the input sequence, initializes the decoder state, and the decoder generates the target program output.
- A.2 Program Evaluation: Using model predictions rather than ground-truth outputs during decoder training tended to produce the strongest results.
- A.2 Program Evaluation: RMC training reached high performance quickly relative to baselines, while DNC and EntNet faltered on specific tasks.The strongest LSTM models also contained many more parameters than the strongest RMC models.
A.3 Viewport BoxWorld
Viewport BoxWorld tests whether an agent can use partial visual observations and memory to plan through key-lock relations toward a gem while avoiding distractors. Levels vary in path structure and difficulty, and the viewport is processed before the RMC.
- Environment: The agent sees only a 5 × 5 RGB viewport plus the color of the key currently held, requiring memory across observations.The full level is a 14 × 14 pixel grid, but perceptual access is restricted to the local viewport.
- Task structure: Each level contains a unique key-lock sequence leading to a gem, alongside distractor branches that can lead to dead ends.Keys disappear after use, so the agent must look ahead and avoid choosing an unrecoverable branch.
- Task structure: Randomized key and lock locations remove spatial shortcuts and emphasize reasoning about abstract key-lock relations in memory.The authors note that this setting may challenge CNN-based approaches whose biases favor proximal relations.
- Difficulty: Levels vary by solution-path length, distractor count, and distractor-path length, with training paths spanning 1 to 5 locks and up to 5 distractor branches.The sampling pool contains 20 key-lock colors, and training distractor branches have length 1.
- Architecture: The viewport is processed by two convolutional layers and then passed to an RMC with four memories and four attention heads.The convolutional layers use 12 and 24 kernels, while spatial position is added as two feature maps.
A.3.1 Results
In viewport BoxWorld, the RMC substantially outperforms a ConvLSTM on procedurally generated levels requiring up to five boxes and distractor branches.
- Results: 98% of BoxWorld levels were solved by the RMC agent after 1e9 steps, compared with 73% for the ConvLSTM.Training levels required opening at least 1 and up to 5 boxes, with 0 to 5 randomly sampled distractor branches.
A.4 Language Modeling
The language-modeling experiments train the RMC with selected architecture settings and evaluate perplexity across WikiText-103, GigaWord, and Project Gutenberg. The supplied passages also describe validation-perplexity and test-unroll-length comparisons, including a caveat about smoothed curves.
- Training setup: The RMC language-modeling setup uses 512-dimensional word embeddings, tied input-output embedding parameters, and truncated backpropagation through 100 steps.Training used Adam with learning rate 0.001, gradient clipping at L2 norm 0.1, and six Nvidia Tesla P100 GPUs.
- Architecture selection: The architecture sweep varies total memory units, attention heads, number of memories, MLP layers, and attention blocks.The listed ranges include 1000–3000 memory units, 1–5 attention heads, 1–2 memories, 1–5 MLP layers, and 1–4 attention blocks.
- Architecture selection: Validation error on WikiText-103 selects 2500 units, 4 heads, 1 memory, a 5-layer MLP, and 1 attention block, reused on GigaWord and Project Gutenberg without further sweeps.The reuse reflects the expense of additional training sweeps.
- Evaluation: Figure 11 presents WikiText-103 validation perplexity with an LSTM comparison, but visual curves may differ from Table 2 because of smoothing.This passage supports using the table for exact numerical comparisons rather than reading precise values from the plotted curves.
- Evaluation: Figure 12 compares perplexity across test unroll lengths against continuous full-test-set unrolling and reports that both models use little information beyond 500 words.The RMC shows a smaller perplexity increase when evaluated with shorter unrolls and no state transfer.
- Evaluation: Table 3 reports GigaWord test perplexity split by word-frequency buckets, including a >10K category for the most frequent words.The buckets are defined by the number of occurrences in the training set.