Bohte, S.M. (Sander)

Roelfsema, P.R. (Pieter)

Rombouts, J.O. (Jaldert)

English

Rombouts, Jaldert

Bohte, Sander

Roelfsema, Pieter

NARCIS 

Learning resets of neural working memory

CWI's Institutional Repository

Learning Resets of Neural Working MemoryJ.O. Rombouts1, P.R. Roelfsema2 and S.M. Bohte1 ∗1- Centrum Wiskunde & Informatica - Life SciencesScience Park 123 - The Netherlands2- Netherlands Institute for Neuroscience - Vision and CognitionMeibergdreef 47 - The NetherlandsAbstract. Working memory is a key component of intelligence thatthe brain implements as persistent neural activations. How do persistentneurons learn to store information, and how can they be made to forgetthis information once it is no longer relevant? When animals learn episodictasks, neurons in prefrontal cortex learn to represent task ends. We showthat a biologically plausible neural network model equipped with persistentmemory and a ‘reset’ action can learn to store and forget informationat task ends by reinforcement learning. The new model has competitiveperformance compared to a variety of (biologically implausible) models.1 IntroductionAnimals can learn very complex tasks based on simple reward and punishmentschemes. Such tasks may require working memory (e.g. [1, 2]) where the correctsequence of actions depends on past information. It could be argued that most,if not all, tasks require some form of working memory, making the study of thesetypes of tasks particularly relevant.Successful models for learning working memory tasks employ a form of persis-tent memory, e.g. [3, 4, 5, 6, 7]. From neuroscience, there is ample experimentalevidence that brains contain neurons that have persistent activations in workingmemory tasks [1, 2]. Persistent memory was also found to be powerful for super-vised training of Recurrent Neural Networks (RNNs) [8]. However, models thatemploy persistent memory also need some mechanism for forgetting [8]. Thereason for this is intuitive: Imagine performing a sequence of tasks where theinformation from a previous task may interfere with the execution of the currenttask (known as proactive interference). Clearing memory representations at taskboundaries eliminates this problem.The mathematical framework for modeling the learning of optimal sequencesof actions from rewards and punishments is Reinforcement Learning (RL, [9]).A large body of work has modeled (animal) learning in the RL framework. Mostwork has focussed on tasks that can be modeled as Markov Decision Problems,or MDPs. Working memory tasks fall in the class of Partially Observable MDPs,or POMDP tasks [6] as they require information presented at some previous timeto make an optimal decision at a later point in time.In the RL literature, learning episodic tasks involves a transition to a specialterminal state and a subsequent reset of all dynamic parameters such as eligibility∗JOR is funded by NWO grant 612.066.826111ESANN 2014 proceedings, European Symposium on Artificial Neural Networks, Computational  Intelligence and Machine Learning.  Bruges (Belgium), 23-25 April 2014, i6doc.com publ., ISBN 978-287419095-7. Available from http://www.i6doc.com/fr/livre/?GCOI=28001100432440.traces before starting a new trial [9]. However, autonomously learning agentsdo not have access to such ‘trial-end’ information. Learning task-ends shouldtherefore be part of learning the task, and this may in fact not be trivial. Wheninteracting with a new task, the animal has to learn that actions performedin the current task do not influence rewards in the next task, i.e. that theyare independent. If tasks are independent, then it is beneficial to also clearworking memory. Experimental neuroscience shows that animals indeed learnto recognize the beginnings and endings of tasks (e.g. [10]).Here, we introduce re-AuGMEnT, a biologically plausible neural networkmodel that can learn whole tasks, including trial ends, by reinforcement learning.The model is based on AuGMEnT [7], which can learn challenging workingmemory tasks but requires supervised ‘reset’-signals to reset working memoryrepresentations. We propose to include an ‘internal’ action that implements the‘reset’. This action can be included naturally in the AuGMEnT framework, andappropriately timed reset signals can then indeed be learned. As far as we areaware, this is the first paper that suggests a possible role for ‘end’ signals foundin the brain [10] within the RL framework. We show that the model can learnthe tasks in [7], and we show that AuGMEnT and re-AuGMEnT outperform arange of other RL working memory models on a T-Maze task.2 ModelWe describe re-AuGMEnT adopting the same notation as [7] and indicatingwhere the models differ. We suppress time indices t when expressions do notinclude information from previous time-steps. re-AuGMEnT is a three layerneural network that computes Q-values for alternative actions [9], and wherethe hidden layer includes integrating units that learn to store task relevant in-formation. (Fig 1a). The top layer contains two types of units: instantaneousand transient on(+)/off(-) units. Instantaneous units xi(t) encode sensory vari-ables si(t) at time step t, and on/off units code truncated temporal derivativesof these same sensory variables as:x+i (t) = [si(t)− si(t− 1)]+ ; x−i (t) = [si(t− 1)− si(t)]+ ,where [x]+ = x for x > 0 and 0 otherwise. The hidden layer also contains twokinds of units; regular units (superscripted with R) and memory units (super-scripted with M). Regular units j receive inputs from all instantaneous units inthe input layer via weights vRij (with vR0j a bias weight). Their activation yRj isdetermined by a standard sigmoidal transformation of their inputs aRj :yRj = σ(aRj , θ) = 1/(1 + exp (θ − aRj ))with aRj =∑ivRijxi .where θ shifts the squashing function. Memory units integrate input from theon/off units x′l = {x+l , x−l } through connections vMlm:aMm (t) = aMm (t− 1) +∑lvMlmx′l(t),where activations yMm are computed as yMm = σ(aMm , θ). Memory units and regu-lar hidden units project to output layer units k through, respectively, connections112ESANN 2014 proceedings, European Symposium on Artificial Neural Networks, Computational  Intelligence and Machine Learning.  Bruges (Belgium), 23-25 April 2014, i6doc.com publ., ISBN 978-287419095-7. Available from http://www.i6doc.com/fr/livre/?GCOI=28001100432440.wMmk and wRjk (with wR0k a bias weight). The Q-value for action k in state s, qs,k(where s is implicitly defined by the hidden layer activations) is:qs,k = qk =∑jyRj wRjk +∑myMmwMmk .Given the Q-values computed by the network, actions are selected using a max-Boltzmann Winner-Takes-All (WTA) mechanism [11] with exploration rate ,situated in an external controller. The network learns by minimizing SARSATemporal Difference errors δ(t) by stochastic gradient descent [7]:E(t) =12δ(t)2 =12(r(t) + γqs′,K′(t)− qs,K(t− 1))2 ,where r is the reward received after executing the action K in state s at time(t− 1), qs′,K′ is the predicted value of the winning action K ′ for the next states′, and γ is the discount rate. Learning is implemented by an interaction offeedforward and feedback signals and a globally available neuromodulator likedopamine which represents δ(t). The activation of the winning output unit is setto 1, and the activations of the other output units are set to 0: zk = δkK , whereδkK is the Kronecker delta function. The winning unit sends feedback throughfeedback connections w′ (Fig 1a, dashed connections). The feedback interactswith the feedforward activations to form synaptic tags Tagxy on each connectioncommensurate to the degree that this connection influenced the action selection.Connection strengths wxy are modified as Δwxy = βδTagxy, with β the learningrate. Synaptic tags are then updated as:ΔTagRjk = (λγ − 1)TagRjk + yRj zk ,ΔTagMjk = (λγ − 1)TagMmk + yMm zk ,ΔTagRij = (λγ − 1)TagRij + w′RKjyRj (1− yRj )sTraceRij ,ΔTagMlm = (λγ − 1)TagMlm + w′MKmyMm (1− yMm )sTraceMlm ,with λ a decay parameter [9] and w′RKj and w′MKj feedback weights from theoutput layer to the hidden layer. sTraces are intermediate variables:sTraceRij = xi(t) ; sTraceMlm =∑tt′=0x′l(t′) .The sTrace variables are new to re-AuGMEnT. A Tag effectively encodes aneligibility trace [9], and an sTrace encodes the history of activations that weretransmitted through the associated synapse. Having separate Tags and sTraceshas two advantages: all Tags are now treated identically (in contrast to [7] whereλ = 0 for memory Tags) and sTraces can be reset independently from Tags.When the ‘reset’ action is selected, parameters are updated as normal. Then,the values of memory sTraces and the activations of memory units are set to zero.After this reset, new feedforward activations are computed given the currentobservation, and an action is selected with the restriction that the reset actioncannot be selected twice in a row. Weights, sTraces and Tags are then againupdated as defined above. In addition to these modifications to AuGMEnT [7],we extended the observation layer with information about the reward and withthe previous action of the network (including the new reset action). This ishelpful as these signals can contain information about trial-ends.113ESANN 2014 proceedings, European Symposium on Artificial Neural Networks, Computational  Intelligence and Machine Learning.  Bruges (Belgium), 23-25 April 2014, i6doc.com publ., ISBN 978-287419095-7. Available from http://www.i6doc.com/fr/livre/?GCOI=28001100432440.+?0.90.70.50.3–0.3–0.5–0.7–0.9–?Assigned weightsFavouring greenFavouring redShapesa ProFixCueDelayGoAntibS1S4..c[2 , 0 , 1]GdFig. 1: (a) AuGMEnT (black lines) and extended re-AuGMEnT (grey lines). Dashedlines: feedback connections, Diamonds: memory units (b). Saccade/Antisaccade (SA)task after [1]. (c) Probabilistic Classification (YS) task after [2]. (d) T-Maze task after[4]. Corridor length N = 3. Dashed lines show agent’s observation space (coded as ingrey numbers). Circle indicates location of reward at end of maze (G - invisible)3 ExperimentsWe tested our model on the same set of (non-linear) tasks as were used for testingAuGMEnT [7] (Fig. 2b,c) and also to the T-Maze task from [4] (Fig. 1d). ForAuGMEnT we used networks with three regular units and four memory units inthe hidden layer, as in [7]. As the re-AuGMEnT model expands the input layer,we used ten regular and ten memory units in the hidden layer. The numberof output units (possible actions) depended on the task. For AuGMEnT weused the same parameters for the network and tasks as [7]. For re-AuGMEnTnetworks, we set β = 0.05 and λ = 0.10. Unless otherwise noted, results arebased on runs with 100 randomly initialized networks.Saccade/Antisaccade. The Saccade/Antisaccade (SA) task is based on [1](Fig. 1b). We implemented the task as in [7]. Briefly, the model is presentedwith a fixation mark (filled or empty square), then a cue (circle) appears on theleft or right. After a delay the fixation mark turns off and the model has to select‘left’ or ‘right’. For the filled square the model has to select the cue directionand the opposite direction for the empty square. The trial ends without rewardif the model breaks fixation before the fixation mark turns off. A correct trialyielded a reward of 1.5. We gave both models at most 1 × 105 trials to learnthe task. After networks made 90% optimal choices for each of the four possiblesubtasks over the last 50 examples of each, we checked for correct performanceby running 100 validation trials with β and  set to 0.Probabilistic Classification. The probabilistic classification task is based on [2](Fig. 1c). Trial structure is similar to that of the SA task discussed above,but now four shapes are presented to the model. After a delay period withonly the fixation mark visible, the model has to select ‘left’ or ‘right’. Theoptimal choice depends on the four shapes s1–s4 that were shown; each shapehas an associated weight (inset Fig. 1c) that determines the conditional rewarddistribution P (Red|W ) = 1/(1+10−W ), with W =∑4i=1 w(si); P (Green|W ) =114ESANN 2014 proceedings, European Symposium on Artificial Neural Networks, Computational  Intelligence and Machine Learning.  Bruges (Belgium), 23-25 April 2014, i6doc.com publ., ISBN 978-287419095-7. Available from http://www.i6doc.com/fr/livre/?GCOI=28001100432440.(log)(log)Mediana b c dMBitsRNN RL-LSTMre-AuGMEnTAuGMEnTFig. 2: Success rates for (a) AuGMEnT () and re-AuGMEnT () learning SA andYS task. (b) T-Maze task (N=100 for our models; N=10 for results from [4]). Seelabels inset in d. Bars: 95% confidence intervals (by R method prop.test). (c) Mediannumber of trials for learning the SA and YS tasks. Bars show 1st and 3rd quartile ofdistribution. (d) Mean number of time-steps (not trials) for learning the T-Maze.1− P (R|W ). The colored targets are randomly assigned to the left or right oneach trial. The model received a reward of 1.5 for its choice for the Red or Greentarget according to the conditional distributions. We gave both models at most5 × 105 trials to learn the task. Learning was complete when the model made85% optimal choices over the last 2× 104 trials.T-Maze. The T-Maze task is based on [4] (Fig. 1d). Information presentedat the beginning of the maze is required to make optimal decisions at the end.The agent has actions N, E, S, W to move in all compass directions; task difficultyis scaled by increasing the corridor length N . When the agent remains in thesame place (e.g. by moving into a wall), it receives a negative reward (−.1).The correct decision at the end of the maze is worth 4, and the wrong decision−.1. We added a time-out condition to the task: after 1.2N + 2 time-steps weautomatically stopped the trial, and started a new one. For the simulations wegave each network at most 5 × 105 trials to learn the task. Convergence wasdetermined as for the SA task, but checked at 80% optimal choices as in [4].The results in Fig. 2a show that re-AuGMEnT is able to learn the same tasksas AuGMEnT with similar success rates. The learning process does take sig-nificantly longer, and the within-model variance of learning speed is also larger(Fig. 2c). Given the increased complexity of the task, such increased learningtime is to be expected. Figures 2b,d compare the performance of AuGMEnT,re-AuGMEnT, and the methods tested in [4]: RL-LSTM, Memory Bits [3] andElman Simple Recurrent Networks; note that none of these methods are biolog-ically plausible. It is clear that AuGMEnT outperforms all other algorithms.At N = 70, this method still learns the task perfectly while convergence forRL-LSTM, the second best algorithm, drops to 30%. re-AuGMEnT does fairlywell: it outperforms all models except RL-LSTM and AuGMEnT. Learningin re-AuGMEnT is significantly slower than the learning in AuGMEnT. Thisdemonstrates that supervised reset signals contain a significant amount of in-formation about the task. Importantly, both models automatically generalize115ESANN 2014 proceedings, European Symposium on Artificial Neural Networks, Computational  Intelligence and Machine Learning.  Bruges (Belgium), 23-25 April 2014, i6doc.com publ., ISBN 978-287419095-7. Available from http://www.i6doc.com/fr/livre/?GCOI=28001100432440.over different delays (SA and YS tasks) and corridor lengths due to the on/offunits; e.g. a re-AuGMEnT network that learned the N = 5 maze also solves theN = 70 maze (results not shown). This is not guaranteed for the other models.4 DiscussionThe connection between memory and forgetting has been noted earlier: neuralalgorithms that include persistent memory are known to fail when this memoryis not reset at appropriate times [8]. This was resolved for LSTM by couplingmemory elements to special forget-gates that can reset the memory and theassociated gradients in [8]. These gates are sigmoidal units which set the ‘leak’of an associated memory cell in a multiplicative fashion. The LSTM algorithmwith forget gates is one of the most powerful supervised learning algorithms fortraining RNNs. While AuGMEnT shares a key idea of LSTM, we implementedresets by an internal reset action rather than forget-gates. A reset provides abinary on/off signal that is hard to achieve with sigmoidal gates, since gradientsvanish at the extremes–solutions thus tend to be fit to typical timescales of thetask, and do not automatically generalize to changes in delay length, unlikere-AuGMEnT. Some ideas in the RL literature are related to the current work.One notion is that of a memory cell that can be reset or overwritten by ‘internal’actions as in [3, 6], but neither of these models is biologically plausible. ThePBWM model [5], a biologically based learning scheme for learning workingmemory tasks is also closely related. Although the model is based in RL, itrequires a teaching signal that provides the correct actions on each time-step andthe architecture and learning rules are elaborate. In summary, we have shownthat ‘reset’ actions allow a neural network to learn sequences of difficult workingmemory tasks, including when to forget, purely by trial-and-error learning. Wehypothesize that such reset actions might explain the presence of task-end signalsfound in the brain [10].References[1] Gottlieb and Goldberg. Activity of neurons in the lateral intraparietal area of the monkeyduring an antisaccade task. Nat. Neurosci., 1999.[2] Yang and Shadlen. Probabilistic reasoning by neurons. Nature, 2007.[3] Peshkin, Meuleau, and Kaelbling. Learning Policies with External Memory. ICML, 1999.[4] Bakker. Reinforcement learning with long short-term memory. In NIPS, 2002.[5] O’Reilly and Frank. Making working memory work: a computational model of learningin the prefrontal cortex and basal ganglia. Neural Comp., 2006.[6] Todd, Niv, and Cohen. Learning to use working memory in partially observable environ-ments through dopaminergic reinforcement. NIPS, 2009.[7] Rombouts, Bohte, and Roelfsema. Neurally Plausible Reinforcement Learning of WorkingMemory Tasks. NIPS, 2012.[8] Gers, Schmidhuber, and Cummins. Learning to forget: Continual prediction with LSTM.Neural Comp., 2000.[9] Sutton and Barto. Reinforcement Learning: an introduction. MIT Press, 1998.[10] Fujii and Graybiel. Representation of action sequence boundaries by macaque prefrontalcortical neurons. Science, 2003.[11] Wiering and Schmidhuber. HQ-learning. Adaptive Behavior, 1997.116ESANN 2014 proceedings, European Symposium on Artificial Neural Networks, Computational  Intelligence and Machine Learning.  Bruges (Belgium), 23-25 April 2014, i6doc.com publ., ISBN 978-287419095-7. Available from http://www.i6doc.com/fr/livre/?GCOI=28001100432440.

Learning resets of neural working memory

Abstract

Similar works

Full text

Available Versions

NARCIS

CWI's Institutional Repository