Relational recurrent neural networks

Title

This paper introduces the Relational Memory Core (RMC), a novel memory module that enhances neural networks' ability to perform relational reasoning with stored information.

Abstract

The Relational Memory Core (RMC) improves relational reasoning in memory-based neural networks by using multi-head dot product attention for memory interaction, leading to state-of-the-art results in RL, program evaluation, and language modeling.

Introduction

This work proposes that explicitly modeling memory interactions, rather than just storage and retrieval, improves relational reasoning, and introduces the Relational Memory Core (RMC) to address deficits in current memory architectures.

Relational Reasoning

Relational reasoning involves understanding how entities are connected to achieve a goal, and this paper explores how neural network architectures can be biased to better perform this type of reasoning, especially in the temporal domain.

Relational Memory Core Architecture

The Relational Memory Core (RMC) combines elements of LSTMs, memory-augmented networks, and non-local networks, using multi-head dot product attention to allow fixed memory slots to interact and update based on attended information.

Memory Encoding

The RMC efficiently incorporates new information into its memory matrix using a memory-size preserving attention operation, and this update mechanism can be integrated into recurrent structures like LSTMs.

Tasks

The Relational Memory Core (RMC) was tested on supervised and reinforcement learning tasks, including the Nth Farthest task for relational reasoning and language modeling on large datasets, with a variety of tunable parameters.

Language Modeling

The RMC was evaluated on language modeling tasks using WikiText-103, Project Gutenberg, and GigaWord datasets, demonstrating improved performance over state-of-the-art LSTM models by achieving lower perplexity.

Results

The RMC significantly outperformed LSTM and DNC baselines on the Nth Farthest task and achieved competitive results on program evaluation, while attention analysis provided insights into its internal memory processing.

Performance Comparison

The RMC achieved higher scores in Mini Pacman and showed lower perplexity and better data efficiency in language modeling compared to LSTM baselines, particularly excelling with frequent words.

Discussion

Explicitly modeling memory interactions via multi-head attention in the RMC improves performance on tasks requiring relational reasoning, although the exact causal mechanisms remain an area for further investigation and potential improvements.

Acknowledgements

The authors express gratitude to colleagues for their feedback, discussions, and support in the development of this research.

Model Configurations

This section details model configurations including total units, number of heads, memory slots, blocks, and gate style.

Nth Farthest Task Details

The Nth Farthest task uses sequences of vectors and labels with specific input formats and optimizers, with architectural parameters tested for LSTM and DNC models.

Program Evaluation Setup

The Program Evaluation task uses the Learning to Execute dataset with specific training parameters and encoder-decoder configurations, comparing RMC against baselines like LSTM and DNC.

Program Evaluation Samples

Examples of Addition, Control, and Full Program tasks are presented, demonstrating RMC's superior performance and data efficiency compared to baseline models on these program evaluation tasks.

Viewport Box World Description

Viewport Box World is a reinforcement learning environment requiring relational reasoning, where the agent has limited perceptual access to a grid and must navigate keys and locks to reach a gem.

Mini Pacman Experiment Details

The Mini Pacman experiment configures a Recurrent Memory Core with specific memory and head settings and compares its performance against a ConvLSTM baseline using an actor-critic setup.

BoxWorld Results

An RMC-augmented agent achieved 98% accuracy on BoxWorld levels requiring up to 5 box openings, significantly outperforming a ConvLSTM-augmented agent.

Language Modeling Configuration

The RMC was trained for language modeling using Adam optimizer with gradient clipping and a truncated backpropagation window, with optimal architecture parameters selected based on validation error.

Language Modeling Validation

Validation perplexity on WikiText-103 shows RMC's strong performance compared to LSTM, with RMC exhibiting a smaller perplexity increase at shorter unroll lengths, indicating focus on recent words.