Gating Mechanisms Underlying Sequence-to-Sequence Working Memory

1 INTRODUCTION . Recurrent neural networks ( RNNs ) transform stimuli across multiple time-points to produce nonlinear working memory representations that can be used to solve complex tasks ( Elman , 1990 ; Hochreiter & Schmidhuber , 1997 ; Mante et al. , 2013 ) . Memorization and manipulation of discrete sequences of elements are a common low level requirement to many broad families of such problems ( Hochreiter & Schmidhuber , 1997 ; Jordan et al. , 2021 ; Yang et al. , 2019 ) . However , a full understanding of how the underlying dynamics learned by these networks accurately bring rise to the necessary computations remains an open area of research ( Sussillo & Barak , 2013 ) . That is , from a dynamical system ’ s point of view ( Guckenheimer & Holmes , 1983 ; Jordan et al. , 2021 ) , how does the network ’ s internal phase-flow and attractor structures , brought forth by training , play a part in the found solution to the desired task ? Furthermore , many of the known properties of trained RNNs ’ learned dynamical mechanisms are specific to individual problems ( Henaff et al. , 2016 ; Jarne , 2020 ; Ichikawa & Kaneko , 2021 ) . Therefore , due to their narrow formalism , such attributes are often difficult to extrapolate to alternative tasks of the same family . It has been demonstrated # 82 # 83 # 84 Time N e u ro n I n d e x Encoding Variable Length Delay Decoding Figure 1 : Superimposed trajectories for neurons 82−84 across one thousand trials of VDCM , from a perfectly trained GRU . Neurons 82 and 83 demonstrate slow manifold dynamics during the delay period of each trial . These three neurons represent the qualitative behavior across all 250 neurons . that the underlying attractor structures of RNNs successfully trained on an individual task are often similar , if not topologically equivalent , across networks , regardless of architecture , initialization , and hyper-parameters ( Maheswaranathan et al. , 2019 ) . However , one would also expect commonalities between the underlying structures of RNNs trained on similar , but none the less , different tasks ( Flesch et al. , 2021 ; Yang et al. , 2019 ) – a research feat not yet well explored . To better train and interpret the solutions found by RNN models we require a finer grain analysis , where results can be broadened to very general classes of problems . ( Karpathy et al. , 2016 ) demonstrated a more in depth empirical exploration of GRU and LSTM architectures computing with sequential data , and identified different failure cases that can arise when training these models . However , this work studied the existence of underlying dynamical mechanisms indirectly from their effects on network output and single-neuron behavior , leaving out details on the functionality of each mechanism at the population level . If done , found mechanisms can be further studied and synthetically recreated to be made more understandable . Moreover , the synthetic realization of dynamical mechanisms , inspired by those obtained through gradient based optimization , can be combined and extrapolated from to form synthetic RNN solutions to related tasks . We surgically analyze a single RNN trained on a discrete sequence working memory task . The network performs with no mistakes on a sufficiently sized validation set of trials . Inspired by the behavior discovered in the network , we design and experimentally validate a synthetic solution to a simplified version of the same task , realizable in relatively low dimensions . We then discuss how such findings apply to networks trained on different but mechanistically similar tasks , including a sequence-to-sequence translation task . 2 VARIABLE DELAY COPY MEMORY TASK AND RNN MODEL . In choosing an appropriate task , we look at copy memory ; a standard benchmark to evaluate a neural network ’ s ability to accurately recall information seen many time-steps in the past ( Hochreiter & Schmidhuber , 1997 ; Henaff et al. , 2016 ; Arjovsky et al. , 2016 ) . Let A = { ai } Ki=1 be a set of K symbols . We then pick S , T ∈ N. The input is a vector of categories , length T+2S , where each entry is one-hot encoded . Each trial of the task consists of three phases , of length S , T , and S respectively . During the first phase ( encoding ) , the network is presented with with S entries uniformly sampled from { ai } Ki=1 , to be remembered sequentially . During the second phase ( delay period ) , the network is fed T − 1 inputs of aK+1 , a blank category , indicating no important information is entering the network . At the final time-step of the second phase S+T , a delimiter aK+2 is input to the network , indicating that the RNN should output the original S entries of input in the same order by which they appeared in the first phase of the trial beginning at the next time-step ( third phase – decoding ) . During these last S time-steps , the inputs are all set to the blank category aK+1 . During the first T + S time-steps ( encoding and delay period ) , the network should output aK+1 . The task is to minimize the average cross-entropy of the outputs at every time-step . As such , the networks should remember a sequence of S elements for T time-steps . Henaff , Szlam , and LeCun developed and experimentally verified a synthetic solution to this task ( Henaff et al. , 2016 ) . However , if we allow T to vary from trial to trial , the underlying dynamical mechanisms allowing the RNN to properly enact the computation remains elusive . We will refer to this task as variable delay copy memory ( VDCM ) , as coined by ( Henaff et al. , 2016 ) . We successfully trained a GRU network ( Cho et al. , 2014 ) , by novel means ( Anonymous ) , with a linear readout on VDCM , such that it performs perfectly on all the test trials . We were unable to train other network architectures on VDCM , including LSTM . For each trial we set K = 8 , S = 10 , and T ∼ U ( 100 , 101 , ... , 120 ) . The model used is represented as follows : zt = σ ( Wzxt + biz +Uzht−1 + bhz ) ( 1 ) rt = σ ( Wrxt + bir +Urht−1 + bhr ) ( 2 ) ht = ( 1− zt ) tanh ( Whxt + bih + rt ( Uhht−1 + bhh ) ) + zt ht−1 ( 3 ) yt = Voutht + bout ( 4 ) ychoice = argmax ( [ yt ] m ) , m ∈ [ 0 , 1 , ... , K ] ⊂ N ∪ { 0 } ( 5 ) where ht ∈ Rd , d = 250 , xt ∈ RK+2 , Wz , Wr , Wh ∈ Rd×S and Uz , Ur , Uh ∈ Rd×d are the parameter matrices , biz , bir , bih , bhz , bhr , bhh ∈ Rd are bias vectors , represents element-wise multiplication , and σ ( z ) = 1/ ( 1 + e−z ) is the element-wise logistic sigmoid function . For the linear readout yt ∈ RK+1 , bout ∈ RS−1 and Vout ∈ R ( S−1 ) ×d . ychoice is taken as the largest element of yt at each time-step , and represents the chosen class readout . 3 ENCODING ON SLOW MANIFOLDS . The computations required for VDCM fall into two main parts . The first is the memory structure . How does the GRU retain information about the elements presented to the network at each of the ten encoding time-steps ? The second is memory recall . How does the GRU pull stored information from memory in the correct order ? In this section we will focus on the former , and explain the structure by which our trained RNN uses to properly encode inputted information . Fig . 1 demonstrates the behavior of hidden-state neurons 82 − 84 ( i.e . [ ht ] j for index j ∈ [ 82 , 83 , 84 ] ) of the trained network , while performing VDCM . Let neuron refer to hidden-state neuron unless otherwise specified . The trajectory of these selected neurons across one thousand trials are superimposed , clearly indicating the three segments of the task . If we look at the delay period , beginning at t = 11 , we notice that the neural activity appears to be near constant across each trial . The selection of neurons shown in Fig . 1 are representative of the behavior of across most neurons in the network , with the exception of several oscillatory modes that are rare and inconsistent across trials . A complete collection of neurons across trials can be found in the appendix E. What is of primary interest are the neurons that significantly vary trial to trail during the delay period , such as neurons 82 and 83 . These varying neurons indicate the existence of a slow manifold , an observation in line with recent research ( Ghazizadeh & Ching , 2021 ) . A slow manifold is similar to an attractor ( i.e . a fixed point , an attracting line , an attracting ring , etc . ) , where if the state of the system ht lies on the attractor it will not change unless perturbed . However , the manifold is not entirely made up of fixed points , rather the speed of the phase flow in these regions is arbitrarily slow in a subset of directions . In the GRU architecture , such behavior results from either a pseudo-line attractor ( Jordan et al. , 2021 ) , or from the influence of the update-gate zt ( Cho et al. , 2014 ) . If we look at the update-gate for dimensions with analogous behavior to 82 and 83 , what we find is that , during the delay period , most demonstrate a value close to 1 for zt . Such activity simplifies equation 3 to the following approximate form : [ ht ] i ≈ [ ht−1 ] i ( 6 ) where i indexes the neurons with a high update-gate during the delay period . As such , the network retains near perfect memory of the past in these directions . We assume that the neurons ’ updategates are the primary mechanism to enact slow-manifolds in our trained network . In the case of pseudo-line attractors , such slow flow is the result of the nullclines of the underlying continuoustime dynamical system , by which the network can be interpreted as a numerical approximation of , existing sufficiently close together in the hidden state-space ( Jordan et al. , 2021 ) . However , the nullclines of this system can not be oriented such that they form a pseudo-line attractor in any canonical direction ( in the direction of a single neuron ) . We will see in the next section that the means by which our network is decoding are canonical in nature , and so we disregard this mechanism . Given that this is the most likely mechanism used to encode information , how exactly this computation enacted ? We trained this network to be able to encode 810 possible sequences , K = 8 elements to choose from at each time-step , across S = 10 time-steps . If such a computation is implemented by carefully placing each trajectory on a slow manifold , the manifold can be segmented into regions where each individual readout element is outputted . Due to the implementation of the argmax function for VCDM , the largest element in the readout vector will be chosen as the class . A cartoon representation of such a regime is depicted in Fig . 2 . To empirically show that a slow manifold is the dynamical feature used and to determine how it is organized , we implement a perturbation based experiment . The low level details can be found in appendix A . We can assume an encoding time-step q ∈ [ 1 , ... , 10 ] , and an element p ∈ [ 0 , ... 7 ] . If across two trials , we input a sequence where at time-step q the element to be encoded is p , and an identical sequence where at time-step q the element is not p , we can determine how the neuron representation differs after the encoding phase of each trial ( t = 11 ) . By comparing this difference over many sessions , we can determine which neurons are used to encode for this q − p pair , how often each neuron is used , and approximate the expected value each neuron takes . We can then test the accuracy of our q − p representation by inputting many sequences through the network , where the elements at t = q are not p , but we set , at t = 11 , the n most frequently used neurons in the representation to their approximate expected values for the representation . If the network output remains unaltered , but at decoding time-step q the readout is incorrectly changed to p , we consider the trial a success . The results of this method are displayed in Fig . 3 . Only 73 of the 250 neurons in the network may be used to encode memory . We can reorder these 73 neurons by their center of mass with respect to the 80 possible combinations of q and p , where we sweep through all values p can take on for q = 1 , then q = 2 , etc . Fig . 4 ( Left ) depicts such a reordering , revealing a block-like structure . We see that a nonempty set of neurons account for every element at each specific time-step q . The plot is colored by the expected value that neuron ( row ) takes when presented with a specific element at a given time-step ( column ) . To analyze the finer details of the manifold ’ s structure , let ’ s consider the neurons primarily tuned to a single time-step ; those tuned to encode information at time-step 9 , for example . We can use principal component analysis ( PCA ) to visualize the activity of these selected neurons in low dimensions ( Bishop , 2006 ) . We project down the state of the set of neurons tuned to time-step 9 at t = 11 , across the set of test trials as demonstrated in Fig . 4 ( Right ) . Data points are colored by which element K was presented to the network at t = 9 , and form eight separable clusters , one for each element . The points in each cluster vary in a single direction , indicated by the blue arrow . Since VDCM is a deterministic task , the only source for this variability is the influence of the input to the network at all time-steps preceding t = 9 . This suggests that the neurons primarily tuned to encode information at t = 9 are not fully decoupled from the neurons tuned to other time-steps . However , this observation brings about an important point regarding the neurons in our trained network that are primarily tuned to input presented at t = 9 . While PCA brings about the dimensions with the highest variability across trials , it does not indicate which dimensions are most crucial for enacting the computation . While insightful in understanding how training brought forth various sub-mechanisms that make up the finer memory structures in the network , this direction allocated to previous inputs may not be used to indicate class during the ninth step of decoding . In the following section , we will dive into the second major computation required for VDCM , decoding . It will be shown that analysis of decoding will enable us to better understand the mechanism used for encoding . For all other encoding time-steps , plots analogous to Fig . 4 ( right ) can be found in appendix B .

Summary: The paper analyzes GRU's underlying mechanisms that store and retrieve information in the delay copy task of a sequence of K symbols. It proposes a perturbation-based method to determine which neurons are responsible for encoding a step-element pair. The paper then shows that at each step of the decoding phase, certain neurons are reset by GRU, generally mapping to those tuned to hold information of step-element pair.  Finally, the paper provides a synthetic solution to the delay copy task for the case K=2.