Learning Working Memory Tasks by Reward Prediction in the Basal Ganglia

Transcription

1 Learning Working Memory Tasks by Reward Prediction in the Basal Ganglia Bryan Loughry Department of Computer Science University of Colorado Boulder 345 UCB Boulder, CO, Michael J. Frank Department of Psychology University of Colorado Boulder 345 UCB Boulder, CO Randall C. O Reilly Department of Psychology University of Colorado Boulder 345 UCB Boulder, CO oreilly@psych.colorado.edu Abstract We present a detailed computational model of working memory as subserved by the basal ganglia and prefrontal cortex. The model expands on our previous work by incorporating a learning mechanism based on discrepancies in predicted rewards, the implementation of which is informed via a detailed consideration of the basal ganglia anatomy and physiology. This mechanism enables the system to learn to selectively store and maintain relevant stimuli. The model can learn a demanding working memory task through appropriate shaping, demonstrating the sufficiency of the proposed mechanisms. 1 Introduction The basal ganglia (BG) have long been implicated in motor function and dysfunction. More recent evidence supports a broader view of BG function, including higher level cognitive function and executive planning (e.g., Brown & Marsden, 1990; Wise, Murray, & Gerfen, 1996). Furthermore, the motor role of the BG can be refined to that of planning or initiation of actions (Chevalier & Deniau, 1990; Passingham, 1993), though this is not universally accepted. We and others have argued that the BG is involved in updating working memory representations in the prefrontal cortex, and that one can understand this function using the same mechanisms that can also explain its role in motor initiation (Frank, Loughry, & O Reilly, in press). We have implemented a preliminary model that shows how the unique patterns of connectivity between the prefrontal cortex and basal ganglia can implement this updating/initation mechanism. Here, we extend our model to include a learning mechanism based on the widely acknowl-

2 Posterior Cortex Striatum Frontal Cortex reverberatory loops Thalamus GPi tonically active inhibition net disinhibition Figure 1: Biological schematic, with one stripe highlighted. GPi=internal segment of globus pallidus (substantia nigra, pars reticulata is functional equivalent as well). edged role of the basal ganglia in reward processing and expectation (e.g., Schultz, Apicella, & Ljungberg, 1993). This extension to the model builds on existing ideas that have related the biology of the basal ganglia with the computational framework of the temporal differences algorithm (TD) for reinforcement learning (Barto, 1995; Schultz, Romo, Ljungberg, Mirenowicz, Hollerman, & Dickinson, 1995; Houk, Adams, & Barto, 1995; Montague, Dayan, & Sejnowski, 1996). The model incorporates a somewhat novel biological mechanism for computing the temporal derivative that lies at the heart of the TD algorithm, and is unique in relating this learning mechanism to the role of the basal ganglia in working memory updating. 1.1 Key properties of working memory Our model addresses the subset of working memory function associated with maintained activation states in the prefrontal cortex (Miller & Cohen, 2001; O Reilly, Braver, & Cohen, 1999). The contribution of the basal ganglia in our model is to provide a dynamic gating signal that can rapidly and selectively update these prefrontal working memory representations, while also enabling them to be robustly maintained over delays and in the face of interference from ongoing processing (Frank et al., in press). We illustrate the characteristics and requirements of working memory by way of the 1,2-AX task, which our model can simulate as described later. The 1,2-AX task is a version of the standard CPT-AX task, where subjects respond to target sequences of letters or digits presented one at a time on the screen. There are two possible target sequences, either A-X or B-Y which one is the current target sequence is determined by whether a 1 or a 2 was seen most recently in the stimulus stream. If a 1, the target is A-X, 2 = B-Y. Thus, the 1 or 2 must be maintained over repeated trials while the subject searches for the target sequence in a series of letter presentations. We call this the outer loop. Meanwhile, detecting the target itself requires maintaining stimuli for one step of the sequence, which we call the inner loop. Thus, this task requires rapid selective updating of only the inner loop stimuli while robustly maintaining the outer loop information about the identity of the current target. 2 Biological architecture of the model Our model builds upon the architecture developed in Frank et al. (in press) (Figure 1), which is based on the following facts of the relevant biology. The frontal cortex projects to the striatum of the basal ganglia, which then sends projections via the globus pallidus (GP) and thalamus back to the frontal cortex. Both the frontal cortex and striatum receive sensory information from a variety of posterior cortical areas. Regions of frontal cortex

3 Posterior Cortex Frontal Cortex... Orbital PFC Ext Rew Striatum S M S M Limbic Striatum SN & VTA modulatory dopamine Figure 2: Biological areas involved in the control of dopamine, which regulates learning. S=striosome, M=matrisome, SN=substantia nigra, pars compacta, Ext Rew=external reward. project to corresponding regions of the striatum, which in turn project preferentially back to the same frontal cortex areas (Alexander, DeLong, & Strick, 1986). We assume that this loop-like connectivity holds at a smaller scale of stripes, where an individual stripe is akin to a hypercolumn in the prefrontal cortex (Levitt, Lewis, Yoshioka, & Lund, 1993). It seems clear that stripes in the basal ganglia can provide specific control signals to corresponding stripes in the frontal cortex. Our earlier model shows that the overall disinhibitory effect of striatal firing (arising because the GP neurons are tonically active and inhibiting the thalamus, and the striatal neurons inhibit these GP neurons when they fire) can produce a gating-like effect on working memory representations in prefrontal cortex. This gating signal is specific to individual stripes, allowing fine-grained control (i.e., selectivity). The present model includes additional areas of the basal ganglia that are specifically important for the learning mechanisms (Figure 2). First, the striatum can be divided into two cell types that show different connectivity and reactivity: striosomes (patches) and matrisomes (the matrix) (Graybiel, Ragsdale, & Mood Edley, 1979). The basic disinhibitory gating function just described is associated with the matrisomes. In contrast, the striosomes project to the substantia nigra pars compacta (SN) and the ventral tegmental area (VTA). These are the midbrain nuclei that project dopamine to the rest of the brain (with greatest concentration in the basal ganglia and frontal cortex), so that dopamine can modulate learning. Thus, the striosomes are important for regulating learning, not for updating working memory. Similarly, the limbic striatum (nucleus accumbens) contributes a dominant projection to these dopamine nuclei, and is likewise important for learning. The SN and VTA receive tonic excitation from the subthalamic nucleus (STN) we capture this by ensuring that these neurons have a base tonic firing level. Finally, the orbital region of the prefrontal cortex provides the dominant activation of the limbic striatum, and is included in our model. 3 Learning via reward prediction The model is based on the temporal differences learning mechanism (Sutton, 1988), which divides the learning task into two components, an actor and a critic. The actor s job is to select actions appropriate for the current state in our model, this amounts to the working memory update actions initiated by firing of the matrisomes. The critic s job is to predict how likely it is that the actor s actions will lead to reward differences in these predictions and actual outcomes can be used to drive learning. Specifically, if unexpected reward is obtained, learning should reinforce the actions that led to this reward, and viceversa if expected rewards are not obtained. As the critic learns to perfectly predict rewards, learning stops. Thus, the centerpiece of the critic s computation is a temporal derivative of

4 expected reward: the difference between the previous and current expectations of reward. Dopamine firing has many characteristics of this derivative-based TD learning signal (Montague et al., 1996). Specifically, it fires for unexpected rewards, but not for expected ones. Furthermore, it fires for unexpected stimuli that reliably predict rewards. Although TD does not explain all of dopamine firing (e.g., it also fires for novelty), it may be possible to incorporate these other firing properties within the overall TD framework (Kakade & Dayan, 2001). Our model shows how the temporal derivative can be computed in the projections to the midbrain dopamine areas. Specifically, we incorporated a suggestion from Rick Granger (personal communication, May, 2000) that differences in the time constants of GABA-A and GABA-B can produce a derivative computation. According to this idea, the striosomes and limbic striatum neurons (which are inhibitory) connect directly to dopamine neurons at synapses having GABA-B receptors, which have a long time constant for activation. In addition, these neurons also connect to inhibitory interneurons within the SN and VTA at synapses having GABA-A channels, and these interneurons also connect to the dopamine neurons with GABA-A. GABA-A has a fast time constant. This connectivity is indeed supported by some biological data (Charara, Heilman, Levey, & Smith, 1999). The net effect is that striatal firing causes an initial disinhibition of the dopamine neurons (via GABA-A), allowing them to fire at a faster than basal rate, but the slower GABA-B current then catches up and supplies inhibition that counteracts the disinhibition, restoring a steady-state balance. Thus, the SN/VTA firing effectively reflects changes (i.e., temporal derivatives) in striatal neuron activity. All we need assume is that these striosomes and limbic striatal neurons reflect expected levels of future reward, and the system will behave according to the TD algorithm. The orbital prefrontal cortex likely provides an important source of reward expectation information (e.g., Tremblay & Schultz, 2000). In our model, this orbital PFC area is specifically important for the active maintenance of expectations of future rewards, which then provide sustained activation of striosomes and neurons in the limbic striatum. Our model also includes an interesting division of labor between the striosomes and the limbic striatum. The limbic striatum represents a global prediction of reward, while the striosomes represent more selective reward predictions associated with the actions initiated within the associated stripe for each striosome. Thus, limbic striatal firing suggests that something good will happen, and striosomal firing refines this to say more specifically what action should be rewarded. 4 Computational simulation The model (Figure 3) incorporates the anatomical structures as shown in previous figures. The PFCm layer corresponds to the maintenance units within the PFC, while PFCg are the layer 4 gating units as described in Frank et al. (in press). There are three stripes throughout the PFC and BG layers, with different groups of input stimuli assumed to have already been associated with different stripes (future modeling work will be focused on relaxing this kind of prestructuring). Connectivity within stripes however is initially random. Projections from the striosomes and limbic striatum to the VTA and SN directly implement the GABA-based derivative computation described above. These dopamine neurons modulate leak currents in the basal ganglia and frontal cortex, such that high dopamine levels reduce leak and thus facilitate activation, and vice-versa for low dopamine levels. Thus, dopamine modulates activation and, thereby, learning. The matrisomal projection simulates disinhibitory gating of PFCg through the GP and thalamus by a similar decrease in leak current mechanism, and consequent firing of the PFCg units, if convergent with stimulus input, will trigger intracellular ionic maintenance currents in the PFCm units to maintain their current states (Frank et al., in press). Nevertheless, PFCm will maintain by default (if not already maintaining something else) based on recurrent excitatory connections.

5 3 C 2 B 1 A Input X Z Y R L Output PFCm PFCg OrbitalFC via GP & Thalamus Matrisomes Striosomes LimbicStr Disinhibition Derivitive Dopamine SubstantiaNigra VTA Reward Neg_Reward Figure 3: Working memory model. Activity shown during X stimuli of 1,A,X sequence. 4.1 Basic stimulus-reward association learning We initially tested the model by training it with a stimulus, X, that reliably predicts a subsequent reward. Prior to learning, the model responds to all stimuli as if they were nonpredictive of reward. Thus, the X activates a random PFCm unit within its corresponding stripe, and no other neurons are activated. As emphasized in Frank et al. (in press), it is important that prefrontal neurons always reflect current stimuli if they are not otherwise maintaining something this can happen if gating occurs via control of intracellular ionic conductances instead of modulation of connections into the prefrontal cortex as in other models (e.g., Braver & Cohen, 2000). Next, the reward is delivered by activating the Reward unit, which directly activates the limbic striatum, which in turn activates the SN and VTA because these units respond to changes in limbic striatum activation. The SN sends DA to the striosomes, matrisomes and back to the limbic striatum. The consequent reduced leak current enables the striosome in the X stripe, which had been receiving sub-threshold activation from the X input and the active PFCm unit, to actually get activated. This increases the dopamine level in the corresponding matrisome layer, enabling one of the matrisomes (at random) to become active. This matrisomal firing then triggers updating of PFCm working memory, activating the intracellular ion channels for the currently active unit. Simultaneously, the VTA releases dopamine to the OrbitalFC and PFCm layers, which allows the OrbitalFC to become active and strengthens the PFCm response. The DA activity causes all the above associations to be strengthened, which is how the model learns that an X is predictive of reward. When the weights have increased through repeated trials, the X stimulus will directly trigger OrbitalFC, limbic striatum, striosomal, and matrisomal firing, causing the X stimulus to be robustly encoded into working memory immediately. When the reward is subsequently delivered, the striosome and limbic striatal units are already activated (i.e., predicting the reward), so there is no change in activation and thus no dopamine firing (which is triggered only by changes via the derivative computation). This simulates the basic patterns of dopamine firing observed in neural recording studies (e.g., Schultz et al., 1993). Importantly, the default PFCm maintenance allows the network to learn over delays in stimulus-reward presentations.

6 4.2 Shaping the 1,2-AX task To train a monkey on a complex working memory task like the 1,2-AX task described earlier, one would need to train each step of the task in succession (i.e., shaping). This is what we did with the model. First, we trained the network that X was predictive of reward, as above. Then, we presented an A followed by an X followed by reward. Due to the well-known abilities of the TD algorithm to support higher-order conditioning, this resulted in A being predictive of reward as well. Specifically, when the X was activated after the A, it acted much like a primary reward itself, triggering dopamine firing and thus driving learning of the associations of A with reward. Once the network learned about A, the X no longer produced a change in expected reward, and therefore did not fire dopamine. The same process was then carried out with the 1-A-X sequence. Importantly, this chaining of reward leads to the robust maintenance of all three stimuli in working memory thus is the synergy between reward learning and working memory in our model. Next, the 2-B-Y stimulus was trained in the same way as 1-A-X. Because of the preassignment of stimuli to different stripes in the model, the activation of a 1 after a 2 will displace the memory of the previous 2, and likewise for all other stimuli within their respective functional class (e.g., A displaces B, X displaces Y). As a result, the working memory function of the model is exactly as required for actually performing the 1,2-AX task. All that is needed is to associate the memory states within PFCm with appropriate response outputs, which can be done with basic associative learning mechanisms. 5 Discussion To summarize, we have shown that we can simulate the learning of a relatively complex working memory task using biologically-based reinforcement learning mechanisms. Although this result is encouraging, the model still faces a number of important limitations, even as it represents advances over other approaches. One of the most salient limitations in the model is in how it integrates simple Pavlovian stimulus-reward learning with the kind of operant learning that is required for performing tasks. As the model is currently trained, rewards are not contingent on correct task performance. Instead, this task performance emerges out of associative learning on the working memory representations trained by non-contingent reward-based learning. Simulating performance-contingent learning requires much more trial-and-error, and therefore raises various issues about how to explore the space of possible actions, and exploration-exploitation tradeoffs, which we have yet to incorporate into the model. Nevertheless, we do not foresee any significant barriers to doing so. Also, as already mentioned, we plan to allow the model to develop its own assignments of stimuli to different working memory stripes, instead of using the current prestructuring. One of the most important advantages of our model is in how it integrates working memory function within the reinforcement learning paradigm. Other models have resorted to awkward representations such as the complete serial compound (e.g., Montague et al., 1996), where different units represent a given stimulus at each point in time. In contrast, relevant stimuli are specifically maintained in prefrontal cortex in our model, and these prefrontal representations can bridge temporal gaps between stimuli and subsequent outcomes associated with them. We plan to study the functional properties of our network using more abstract formalisms to gain greater analytical insight into the nature of this mechanism. Acknowledgments This work was supported by ONR grant N and NSF grant IBN

7 6 References Alexander, G. E., DeLong, M. R., & Strick, P. L. (1986). Parallel organization of functionally segregated circuits linking basal ganglia and cortex. Annual Review of Neuroscience, 9, Barto, A. G. (1995). Adaptive critics and the basal ganglia. In J. C. Houk, J. L. Davis, & D. G. Beiser (Eds.), Models of information processing in the basal ganglia (pp ). Cambridge, MA: MIT Press. Braver, T. S., & Cohen, J. D. (2000). On the control of control: The role of dopamine in regulating prefrontal function and working memory. In S. Monsell, & J. Driver (Eds.), Control of cognitive processes: Attention and performance XVIII (pp ). Cambridge, MA: MIT Press. Brown, R. G., & Marsden, C. D. (1990). Cognitive function in parkinson s disease: From description to theory. Trends in Neurosciences, 13, Charara, A., Heilman, C., Levey, A., & Smith, Y. (1999). Pre-and postsynaptic localization of GABA- B receptors in the basal ganglia in monkeys. Neuroscience, 95, Chevalier, G., & Deniau, J. M. (1990). Disinhibition as a basic process in the expression of striatal functions. Trends in Neurosciences, 13, Frank, M. J., Loughry, B., & O Reilly, R. C. (in press). Interactions between the frontal cortex and basal ganglia in working memory: A computational model. Cognitive, Affective, and Behavioral Neuroscience. Graybiel, A. M., Ragsdale, C. W., & Mood Edley, S. (1979). Compartments in the striatum of the cat observed by retrograde cell labeling. Experimental Brain Research, 34, Houk, J. C., Adams, J. L., & Barto, A. G. (1995). A model of how the basal ganglia generate and use neural signals that predict reinforcement. In J. C. Houk, J. L. Davis, & D. G. Beiser (Eds.), Models of information processing in the basal ganglia (pp ). Cambridge, MA: MIT Press. Kakade, S., & Dayan, P. (2001). Dopamine bonuses. In T. Leen, & T. Dietterich (Eds.), Advances In Neural Information Processing Systems, 13. Cambridge, MA: MIT Press. Levitt, J. B., Lewis, D. A., Yoshioka, T., & Lund, J. S. (1993). Topography of pyramidal neuron intrinsic connections in macaque monkey prefrontal cortex (areas 9 & 46). Journal of Comparative Neurology, 338, Miller, E. K., & Cohen, J. D. (2001). An integrative theory of prefrontal cortex function. Annual Review of Neuroscience, 24, Montague, P. R., Dayan, P., & Sejnowski, T. J. (1996). A framework for mesencephalic dopamine systems based on predictive Hebbian learning. Journal of Neuroscience, 16, O Reilly, R. C., Braver, T. S., & Cohen, J. D. (1999). A biologically based computational model of working memory. In A. Miyake, & P. Shah (Eds.), Models of working memory: Mechanisms of active maintenance and executive control. (pp ). New York: Cambridge University Press. Passingham, R. E. (1993). The frontal lobes and voluntary action. Oxford: Oxford University Press. Schultz, W., Apicella, P., & Ljungberg, T. (1993). Responses of monkey dopamine neurons to reward and conditioned stimuli during successive steps of learning a delayed response task. Journal of Neuroscience, 13, Schultz, W., Romo, R., Ljungberg, T., Mirenowicz, J., Hollerman, J. R., & Dickinson, A. (1995). Reward-related signals carried by dopamine neurons. In J. C. Houk, J. L. Davis, & D. G. Beiser (Eds.), Models of information processing in the basal ganglia (pp ). Cambridge, MA: MIT Press. Sutton, R. S. (1988). Learning to predict by the method of temporal diferences. Machine Learning, 3, Tremblay, L., & Schultz, W. (2000). Reward-related neuronal activity during go-nogo task performance in primate orbitofrontal cortex. Journal of Neurophysiology, 83, Wise, S. P., Murray, E. A., & Gerfen, C. R. (1996). The frontal cortex-basal ganglia system in primates. Critical Reviews in Neurobiology, 10,