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: The paper postulates a task that requires the ML model to capture the spatial nature of the task in order to perform well. The chosen ML method is a Transformer with positional encoder, where the encoding is not fixed, but is learnable, making it possible for the transformer to "choose" what kind of representation would it "prefer" for the task at hand. The Transformer is modified to act comparably to a neuroscientific model of hippocampal function (TEM). After the learning process the authors discover, that, when visualized, the representation that has emerged resembles the grid-cell pattern similar to the one that is being encoded by hippocampal place and grid cells. From there the works suggests the existence of a special relationship between the Transformer architecture and the mechanism by which hippocampus encodes spatial information.