Relational recurrent neural networks

TL;DR

The paper introduces the Relational Memory Core (RMC), a memory-augmented recurrent architecture that enables explicit interactions between memory slots through multi-head attention. By integrating recurrence and attention, the RMC improves relational reasoning over time, addressing limitations of traditional memory systems. Empirical results span a diverse set of tasks, including Nth Farthest, LTE program evaluation, Mini Pacman with partial observability, and language modeling on WikiText-103, Gutenberg, and GigaWord, where the RMC achieves notable gains and in some cases state-of-the-art perplexities. The work suggests that explicit memory-to-memory interactions can substantially enhance relational reasoning in sequential domains and points to future directions involving hybrid memory schemes and online, scalable processing. Overall, the RMC provides a flexible, tunable backbone for relational reasoning in sequences with broad applicability across RL, symbolic reasoning, and language modeling.

Abstract

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.

Figures (12)

• Figure 1: Relational Memory Core. (a) The RMC receives a previous memory matrix and input vector as inputs, which are passed to the MHDPA module labeled with an "A". (b). Linear projections are computed for each memory slot, and input vector, using row-wise shared weights $W^q$ for the queries, $W^k$ for the keys, and $W^v$ for the values. (c) The queries, keys, and values are then compiled into matrices and $\text{softmax}(QK^T)V$ is computed. The output of this computation is a new memory where information is blended across memories based on their attention weights. An MLP is applied row-wise to the output of the MHDPA module (a), and the resultant memory matrix is gated, and passed on as the core output or next memory state.
• Figure 2: Tasks. We tested the RMC on a suite of supervised and reinforcement learning tasks. Notable are the $N^{th}$ Farthest toy task and language modeling. In the former, the solution requires explicit relational reasoning since the model must sort distance relations between vectors, and not the vectors themselves. The latter tests the model on a large quantity of natural data and allows us to compare performance to well-tuned models.
• Figure 3: Model analysis. Each row depicts the attention matrix at each timestep of a particular sequence. The text beneath spells out the particular task for the sequence, which was encoded and provided to the model as an input. We mark in red the vector that is referenced in the task: e.g., if the model is to choose the $2^{nd}$ farthest vector from vector $7$, then the time point at which vector $7$ was input to the model is depicted in red. A single attention matrix shows the attention weights from one particular memory slot (y-axis) to another memory slot (columns), or the input (offset column), with the numbers denoting the memory slot and "input" denoting the input embedding.
• Figure 4: $N^{th}$ Farthest hyperparameter analysis. Timestamp refers to hours of training. There is a clear effect with the number of memories, with $8$ or $16$ memories being better than $1$. Interestingly, when the model had $1$ memory we observed an effect with the number of heads, with more heads ($8$ or $16$) being better than one, possibly indicating that the RMC can learn to compartmentalise and relate information across heads in addition to across memories.
• Figure 5: LSTM and DNC training curves for the $N^{th}$ Farthest task.