Modelling Working Memory using Deep Recurrent Reinforcement Learning
Abstract
In cognitive systems, the role of a working memory is crucial for visual reasoningand decision making.
Tremendous progress has been made in understanding themechanisms of the human/animal working memory, as well as in formulatingdifferent frameworks of artiﬁcial neural networks.
In the case of humans, thevisual working memory (VWM) task [ 1 ] is a standard one in which the subjects arepresented with a sequence of images, each of which needs to be identiﬁed as towhether it was already seen or not.
Our work is a study of multiple ways to learn aworking memory model using recurrent neural networks that learn to rememberinput images across timesteps in supervised and reinforcement learning settings.
The reinforcement learning setting is inspired by the popular view in Neurosciencethat the working memory in the prefrontal cortex is modulated by a dopaminergicmechanism.
We consider the VWM task as an environment that rewards theagent when it remembers past information and penalizes it for forgetting.
Wequantitatively estimate the performance of these models on sequences of imagesfrom a standard image dataset (CIFAR-100 [ 2 ] ) and their ability to rememberand recall.
Based on our analysis, we establish that a gated recurrent neuralnetwork model with long short-term memory units trained using reinforcementlearning is powerful and more efﬁcient in temporally consolidating the input spatialinformation.
This work is an initial analysis as a part of our ultimate goal to modelthe behavior and information processing of the working memory of the brain and touse brain imaging data captured from human subjects during the VWM cognitivetask to understand various memory mechanisms of the brain.
Introduction
Memory is an essential component for solving many tasks intelligently.
Most sequential tasksinvolve the need for a mechanism to maintain a context .
In the brain, working memory serves asa work space to encode and maintain the most relevant information over a short period of time, inorder to use it to guide behavior for cognitive tasks.
Several cognitive tasks have been proposedin the Neuropsychology literature to study and understand the working memory in animals.
The
Visual Working Memory Task (VWM task)
[1] or the classic N-back task is one of the most simpleand popular ones.
It involves showing sequences of images to subjects and record their responsesindicating whether they have seen the image already.
On the other hand, with artiﬁcial intelligence systems, there has been very good progress in modelsthat learn from sequences of inputs using artiﬁcial neural networks as memory for all types of learning (supervised, unsupervised and reinforcement).
We intend to use these developments as an idealopportunity for synergy to computationally model the working memory system of the brain.
As memory is an important aspect of both artiﬁcial intelligence and neuroscience, there are somegood studies that helped choose our models as discussed in Section 2.  Models
[Figure]
Gated recurrent-units as cortical circuits [3] highlights that Long short-term memories (LSTMs[4]), the most commonly used gated recurrent units with memory cells controlled by input, outputand forget gates and these can be considered (at a high-level) as an abstraction of a cortical circuitunit composed of layers of pyramidal cells with gating as shown in Figure 1a.
Dopaminergic modulation of working memory As detailed in [5] and [6] and shown in Figure 1b,the stimulus receive a reward from the dopaminergic nuclei based on which it is stored in the prefrontalcortex (PFC) as persistent activations.
Thus, biologically there is a reward-processing mechanism(temporal difference, TD) that is operating outside the cortical memory.
Deep Recurrent Reinforcement Learning As an extension of Deep Q-Networks (DQN)[7] withrecurrent units were proposed as DRQNs by [8] for partially observable environments.
This modelnaturally suits the non-Markov decision process involved in learning.
Supervised Learning models We consider 2 different models for training recurrent networks tooutput the probability of past occurrence of each image in the sequence using supervised learning:
LSTM and LSTM-Attention (LSTM-A).
The attention augmented version ([12]) helps learn thecontext better.
These are used in conjunction with a CNN and trained using backpropagation throughtime (without truncation), as shown in Figure 2a.
Reinforcement Learning models We consider an environment where each image in the sequenceis a state that the DRQN uses as input to approximate the Q-values of the two different action : 0for not seen and 1 for seen .
An (cid:15) -greedy policy is learned using these Q-values to select the actionand receive a reward from the environment (positive value for correct action and negative value forwrong action) as shown in Figure 2b.
These are trained using common TD control algorithms such as
Q-learning and SARSA, indicated as Q+DQRN and SARSA+DQRN respectively.
Experiments and Results
For all experiments in both supervised and reinforcement learning settings, 100 images were drawnfrom CIFAR-100 ([2]) dataset for each sequence.
A CNN that was pre-trained on the same dataset [Figure] [Figure] for the classiﬁcation task was used for intermediate representation before LSTM layers (as shown in
Figure 2)
The problem solved by all the models is a binary classiﬁcation problem, predicting unseen(0)/seen(1) .
The performance of all the models in the experiments were measured using the accuracy metriccalculated based on the number of correct predictions for the 100 images in a sequence (as a %).
Thisevaluation was repeated for 10 independent trials as a part of ablation studies.
Figure 3a shows the test accuracy of both the reinforcement learning models after ﬁxed trainingepisodes to track the progress.
Figure 3b depicts the performance of the different models using meanﬁnal accuracy scores for two different conditions of repetitions of the stimuli images: N = 1 and N = 2 .
Figure 4 indicates the variation of test accuracy of each of the trained models over the courseof a testing sequence.
It can be seen that the models trained in the reinforcement learning settingoutperform both the LSTM and LSTM-A in the supervised setting.
Both on-policy and off-policymethods help in solving the task and learning good parametric functions in the LSTMs.
Conclusions and Future work
Popular literature in deep reinforcement learning ([8], [11]) outline many architectures and methodsto solve speciﬁc (mostly motor) decision making tasks in an environment.
In our work, we have takenan alternative view where we try to solve a strongly memory-oriented task using similar models usedin reinforcement learning to emulate the reward processing happening to store memory.
A good working memory model understands what to remember and what to forget.
From ourstudy, we conclude that modelling a working memory using gated recurrent neural networks (suchas LSTMs) to train using a reward-based learning approach offers signiﬁcant advantages giving areasonably superior performance.
From the models studied,it can be seen that in addition to thecapacity of the recurrent network, a training in the reinforcement learning setting offers a bettergeneralization with its power for acting as a good computational model for the biological workingmemory for explaining the VWM task.
This seems to indicate that the dopaminergic control ofmemory in the PFC is a high-level principle that is common to both artiﬁcial and biological neuralsystems.
Also, this observation could be attributed to the fact that deep reinforcement learning is abetter framework for learning in this case, given the non-stationary and non-Markov nature of thetask involved.
As a next step, we plan to model other cognitive tasks for memory as a reinforcement learningproblem and compare the performance of different algorithms and gating mechanisms in the networks.
Further, we would like to extend our work by using the brain responses of humans solving the taskfrom their fMRI data and identify the neural correlates of the visuo-temporal streams of information.
This would throw more light on the functional similarities of the biological and artiﬁcial neuralmodels explaining how the memory system functions computationally.