Modelling Working Memory using Deep Recurrent Reinforcement Learning

Pravish Sainath, Pierre Bellec, Guillaume Lajoie

TL;DR: LSTMs can more effectively model the working memory if they are learned using reinforcement learning, much like the dopamine system that modulates the memory in the prefrontal cortex

Keywords: deep learning, working memory, recurrent neural networks, reinforcement learning, brain modelling

Abstract:

In cognitive systems, the role of a working memory is crucial for visual reasoning and decision making. Tremendous progress has been made in understanding the mechanisms of the human/animal working memory, as well as in formulating different frameworks of artificial neural networks. In the case of humans, the visual working memory (VWM) task is a standard one in which the subjects are presented with a sequence of images, each of which needs to be identified as to whether it was already seen or not.

Our work is a study of multiple ways to learn a working memory model using recurrent neural networks that learn to remember input images across timesteps. We train these neural networks to solve the working memory task by training them with a sequence of images in supervised and reinforcement learning settings. The supervised setting uses image sequences with their corresponding labels. The reinforcement learning setting is inspired by the popular view in neuroscience that the working memory in the prefrontal cortex is modulated by a dopaminergic mechanism. We consider the VWM task as an environment that rewards the agent when it remembers past information and penalizes it for forgetting.

We quantitatively estimate the performance of these models on sequences of images from a standard image dataset (CIFAR-100). Further, we evaluate their ability to remember and recall as they are increasingly trained over episodes. Based on our analysis, we establish that a gated recurrent neural network model with long short-term memory units trained using reinforcement learning is powerful and more efficient in temporally consolidating the input spatial information.

This work is an initial analysis as a part of our ultimate goal to use artificial neural networks to model the behavior and information processing of the working memory of the brain and to use brain imaging data captured from human subjects during the VWM cognitive task to understand various memory mechanisms of the brain.