Relating transformers to models and neural representations of the hippocampal formation

1 INTRODUCTION . The last ten years have seen dramatic developments using deep neural networks , from computer vision ( Krizhevsky et al. , 2012 ) to natural language processing and beyond ( Vaswani et al. , 2017 ) . During the same time , neuroscientists have used these tools to build models of the brain that explain neural recordings at a precision not seen before ( Yamins et al. , 2014 ; Banino et al. , 2018 ; Whittington et al. , 2020 ) . For example , representations from convolutional neural networks ( Lecun et al. , 1998 ) predict neurons in visual and inferior temporal cortex ( Yamins et al. , 2014 ; Khaligh-Razavi & Kriegeskorte , 2014 ) , representations from transformer neural networks ( Vaswani et al. , 2017 ) predict brain representations in language areas ( Schrimpf et al. , 2020 ) , and lastly recurrent neural networks ( Cueva & Wei , 2018 ; Banino et al. , 2018 ; Sorscher et al. , 2019 ) have been shown to recapitulate grid cells ( Hafting et al. , 2005 ) from medial entorhinal cortex . Being able to use models from machine learning to predict brain representations provides a deeper understanding into the mechanistic computations of the respective brain areas , and offers deeper insight into the nature of the models . As well as using off-the-shelf machine learning models , neuroscience has developed bespoke deep learning models ( mixing together recurrent networks with memory networks ) that learn neural representations that mimic the exquisite spatial representations found in hippocampus and entorhinal cortex ( Whittington et al. , 2020 ; Uria et al. , 2020 ) , including grid cells ( Hafting et al. , 2005 ) , band cells ( Krupic et al. , 2012 ) , and place cells ( O ’ Keefe & Dostrovsky , 1971 ) . However , since these models are bespoke , it is not clear whether they , and by implication the hippocampal architecture , are capable of the general purpose computations of the kind studied in machine learning . In this work we 1 ) show that transformers ( with a little twist ) recapitulate spatial representations found in the brain ; 2 ) show a close mathematical relationship of this transformer to current hippocampal models from neuroscience ( with a focus on Whittington et al . ( 2020 ) though the same is true for Uria et al . ( 2020 ) ) ; 3 ) offer a novel take on the computational role of the hippocampus , and an instantiation of hippocampal indexing theory ( Teyler & Rudy , 2007 ) ; 4 ) offer novel insights on the role of positional encodings in transformers . 5 ) discuss whether similar computational principles might apply to broader cognitive domains , such as language , either in the hippocampal formation or in neocortical circuits . Note , we are not saying the brain is closely related to transformers because it learns the same neural representations , instead we are saying the relationship is close because we have shown a mathe- matical relationship between transformers and carefully formulated neuroscience models of the hippocampal formation . This relationship helps us get a better understanding of hippocampal models , it also suggests a new mechanism for place cells that would not be possible without this mathematical relationship , and finally it tells us something formal about position encodings in transformers . 2 TRANSFORMERS . Transformer Neural Networks ( Vaswani et al. , 2017 ) are highly successful machine learning algorithms . Originally developed for language , transformers perform well on other tasks that can be posed sequentially , such as mathematical understanding , logic problems ( Brown et al. , 2020 ) , and image processing ( Dosovitskiy et al. , 2020 ) . Transformers accept a set of observations ; X = { x1 , x2 , x3 , · · · , xT } ( xt could be a word embedding or image patch etc ) , and aim to predict missing elements of that set . The missing elements could be in the future , i.e . xt > T , or could be a missing part of a sentence or image , i.e . { x1 = the , x2 = cat , x3 = sat , x4 = ? , x5 = the , x6 = mat } . Self-attention . The core mechanism of transformers is self-attention . Self-attention allows each element to ‘ attend ’ to all other elements , and update itself accordingly . In the example data-set above , the 4th element ( ? ) could attend to the 2nd ( cat ) , 3rd ( sat ) , and 6th ( mat ) to understand it should be on . Formally , to attend to another element each element ( xt is a row vector ) emits a query ( qt = xtWq ) and compares it to other elements keys ( kτ = xtWk ) . Each element is then updated using yt = ∑ τ κ ( qt , kτ ) vτ , where κ ( qt , kτ ) is kernel describing the similarity of qt to kτ and vτ is the value computed by each element vτ = xtWv . Intuitively , the similarity measure κ ( qt , kτ ) places more emphasis on the elements that are relevant for prediction ; in this example , the keys may contain information about whether the word is a noun , verb or adjective , while the query may ‘ ask ’ for any elements that are nouns or verbs - elements that match this criteria ( large κ ( qt , kτ ) , i.e . cat , sat , mat ) are ‘ attended ’ to and therefore contribute more to the output yt . Typically , the similarity measure is a softmax i.e . κ ( qt , kτ ) = e βqt·kτ∑ τ′ e βqt·kτ′ . These equations can be succinctly expressed in matrix form , with all elements updated simultaneously : yt = softmax ( qtK T √ dk ) V → Y = softmax ( QK T √ dk ) V ( 1 ) Here Q , K , V are matrices with rows filled by qt , kt , vt respectively , and the softmax is taken independently for each row . After this update , each yt is then sent through a deep network ( fθ ( · · · ) ) typically consisting of residual ( He et al. , 2016 ) and layer-normalisation ( Ba et al. , 2016b ) layers to produce zt = fθ ( yt ) . Z is the output of the transformer which can then be used for prediction , or sent through subsequent transformer blocks . Position encodings . Self-attention is permutation invariant and so tells you nothing about order of the inputs . Should the data be sequential ( i.e . meaning depends on the order of elements , such as in language , or navigation as we will see later ! ) , it is necessary to additionally encode the position/ where x is in the sequence . This is typically done by adding a ‘ position encoding ’ that uniquely identifies each time-step ( et - typically sines and cosines ) to each input : xt ← xt+et . Alternatively the position embedding can be appended i.e . ht = [ xt , et ] , with self attention then performed using ht as input . 3 TRANSFORMERS LEARN ENTORHINAL REPRESENTATIONS . Here we show that transformers ( with a small modification ) recapitulate spatial representations - grid and band cells - when trained on tasks that require abstract spatial knowledge . Spatial understanding task . The task ( more detail in Appendix ) is to predict upcoming sensory observations xt+1 conditioned on taking an action at while moving around spatial environments ( Figure 1a ) . For example , after seeing { ( x1 = cat , a1 = North ) , ( x2 = dog , a2 = East ) , ( x3 = frog , a3 = South ) , ( x4 = pig , a4 = West ) , ( x5 = ? , a5 = · · · ) } , the aim is to predict x5 = cat . For simplicity , we treat sensory observations as one-hot vectors , thus the prediction problem is a classification problem . e0 x0 e1 x1 e2 x2 e3 ? y3 q0 k0 v0 q1 k1 v1 q2 k2 v2 q3 iWv iWq iWk Transformer ( with causal attention mask ) Input with position encoding Query Key , Value A B D B C A A D B B C A D A D C D C Environment 1 Environment 2 Observations ( x ) Actions ( a ) A B A ? C AAD ? D Recurrent connections W ( a0 ) W ( a1 ) W ( a2 ) With linear activations With ReLu activationsReal grid cell rate maps Cell rate maps Autocorrelations Cell rate maps Autocorrelations a b dc e f Band cells Figure 1 : ( a ) Sequence prediction in spatial navigation tasks test abstract spatial understanding since some sensory predictions can only be done by knowing ( generalising ) certain rules e.g . North + East + South + West = 0 or Parent + Sibling + Niece = 0 . Note , we use sequences drawn from much larger graphs . ( b ) Transformer with recurrent position encodings . ( c ) Real grid cell rate-maps ( Hafting et al. , 2005 ) . ( d-f ) Learned position embedding rate-maps ( i.e . average activity at each spatial location ; plots are spatially smoothed ) . ( d-e ) Resembling grid cells with ( e ) linear activation or ( e ) ReLu activation post transition . ( f ) Resembling band cells ( Krupic et al. , 2012 ) . When faced with an unseen stimulus-action pair ( e.g . x4 = pig , a4 = West above ; an action you have never taken at that stimulus before ) , successful prediction requires more than just remembering specific sequences of stimulus-action pairs ; knowledge of the rules of space must be known ; i.e . North + East + South + West = 0 allows prediction of x5 = cat . Crucially , such rules generalise to any 2D spaces and may therefore be transferred to aid prediction in entirely novel 2D environments . This is powerful , since unobserved relations between observed stimuli can be inferred in a zero-shot manner . However , these relational rules are not ‘ known ’ a priori and therefore must be learnt . We therefore train across multiple different spatial environments which share the same underlying 4-connected Euclidean structure ( Figure 1a ) - this means the model must learn and generalise the abstract structure of space to use for prediction in new environments . To perform on these tasks , the three modifications to the transformer are : 1 . Recall equation 1 ; yt = softmax ( qtK T √ dk ) V , where Q = HWq , K = HWk , V = HWv , and H is a matrix of inputs and position encodings ( i.e . its rows are ht = [ xt , et ] ) . We restrict these weight matrices such that queries ( Q ) and keys ( K ) are the same ; Q , K = EWe . We refer to this matrix as Ẽ . Thus the keys and queries only focus on position encodings . Meanwhile , values are exclusively dependent on the stimulus component of H i.e . V = XWx . We refer to this matrix as X̃ . yt = softmax ( qtK T √ dk ) V → yt = softmax ( ẽtẼ T √ dk ) X̃ ( 2 ) This is an extreme version of the realisation that , in transformers , best performance is when position encodings are used to compute keys and queries , but not values . 2 . We use causal transformers ; the key and value matrices contain the projected position encodings and sensory stimuli respectively at all previous time-steps ( i.e . e < t and x < t ) . This is equivalent to causal ‘ unmasking ’ as the agent wanders the environment accumulating new experiences ( not-yet-experienced stimulus-position pairs are inaccessible to the agent ) . Meanwhile the query at each time-point is the present positional encoding et . 3 . The position encodings are recurrently generated ( as in Wang et al . ( 2019 ) ; Liu et al . ( 2020 ) ) ; et+1 = σ ( etWa ) , where Wa is a learnable action-dependent weight matrix , and σ ( · · · ) is a non-linear activation function . This means that unlike traditional transformers , position encodings can be optimised and not the same for every sequence . It now becomes interesting to see what representations are learned . These modifications are sufficient to learn spatial representations , in the position encodings , that mimic representations observed in the brain ( Figure 1C ; see Appendix for model and training details ) . The rest of this paper now explains why this is not a surprising result ; namely we show that a transformer with recurrent positional encodings is closely related to current neuroscience models of the hippocampus and surrounding cortex ( Whittington et al. , 2020 ; Uria et al. , 2020 ) . Here we focus on the Tolman-Eichenbaum Machine ( TEM ) ( Whittington et al. , 2020 ) , though the same principles apply for Uria et al . ( 2020 ) . The critical points are : 1 ) the memory component of TEM can be viewed as a transformer selfattention , since the TEM memory network is analogous to a Hopfield network ( Hopfield , 1982 ) which have recently been shown to be closely related to transformers ( Ramsauer et al. , 2020 ) ; 2 ) TEM path integration ( see below ) can be viewed as a way to learn a position encoding .

Summary: This paper demonstrates emergence of grid and place cells in transformer architectures with positional encodings. Transformers are related to TEM models and Hopfield networks in computational neuroscience. A neurobiological model of TEM-transformers is presented, which contains feature neurons and memory neurons. It is shown that memory neurons resemble place cells in the hippocampus. Remapping phenomenon is also discussed.