RELATING TRANSFORMERS TO MODELS AND NEURAL REPRESENTATIONS OF THE HIPPOCAMPAL FORMATION
Year: 2021

Abstract:
Many deep neural network architectures loosely based on brain networks have recently been shown to replicate neural firing patterns observed in the brain. One of the most exciting and promising novel architectures, the Transformer neural network, was developed without the brain in mind. In this work, we show that transformers, when equipped with recurrent position encodings, replicate the precisely tuned spatial representations of the hippocampal formation; most notably place and grid cells. Furthermore, we show that this result is no surprise since it is closely related to current hippocampal models from neuroscience. We additionally show the transformer version offers dramatic performance gains over the neuroscience version. This work continues to bind computations of artificial and brain networks, offers a novel understanding of the hippocampal-cortical interaction, and suggests how wider cortical areas may perform complex tasks beyond current neuroscience models such as language comprehension.

heading: 1 INTRODUCTION
text: 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.

heading: 2 TRANSFORMERS
text: 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.

heading: 3 TRANSFORMERS LEARN ENTORHINAL REPRESENTATIONS
text: 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.

heading: 4 TEM
text: The Tolman-Eichenbaum Machine (TEM; Figure 2, further details in Appendix) is a neuroscience model that captures many known neural phenomena in hippocampus (HPC) and entorhinal cortex (medial/lateral; MEC/LEC). TEM is a sequence learner trained on tasks exactly like the one described in the previous section. TEM consists of two parts;
1) A module that aims to understand where it is in space, using a representation g to represent location. To update its location, TEM uses path-integration - the accumulation of self movement vectors a - enacted in a recurrent neural network:
gt+1 = σ(gtWa) (3) Where Wa is a learnable action dependent weight matrix and σ(· · · ) is a non-linear activation function. It is in this path-integrating representationg that TEM learns grid and other entorhinal cells for self-localisation (Figure 2b).
2) To make sensory predictions, location representations g alone are not enough; they must each link to a sensory observation x, corresponding to the stimulus at that position. Note that these links are specific to an environment, since each environment consists of a different arrangement of stimuli in space (i.e. different stimulus-position pairings).
The linking is done by binding every element of g with every element of x, in other words an outer product that is flattened back into a vector;
p = flatten(xTg) (4) These conjunctive p representations are stored in ‘fast weights’ 1via Hebbian learning;
Mt = t−1∑ τ=1 pTτ pτ (5)
And they can later be retrieved using an attractor network (a continuous version of the Hopfield network). Here a query vector q (details next paragraph) is inputted into the network and updated via; q ← σ(qMt) (6) where σ(· · · ) is a non-linear activation function; a ReLu in TEM. Crucially, because the memories are formed using both g and x, they can be retrieved (pattern-completed) using just one of those representations alone i.e. ‘what did I see the last time I was here’ or ‘where was I the last time I saw this’. To retrieve a memorised conjunction p, TEM imagines (path-integrates) the next location g and provides this as input to the attractor network in the form q = flatten(1Tg). Equation 6 is then iterated until a memory is retrieved.
Finally, to make sensory predictions, the retrieved conjunctive memory (pretrievedt ) is ‘deconjunctified’ into sensory and location components. The sensory component is obtained by unflattening pretrievedt and summing over the g dimension (Figure 8);
xretrievedt = sum(unflatten(p retrieved t ), 1) (7)
Finally, to make the sensory prediction xretrievedt is fed through a MLP zt = fθ(x retrieved t )) to classify (predict) the upcoming sensory observation.
It is also possible, and often helpful, to project g and x via Wg and Wx; g̃ = gWg and x̃ = xWx before they are combined conjunctively 2.

heading: 5 TEM AS A TRANSFORMER
text: Here we show that the above equations of TEM can be written so that: 1) the memory retrieval components looks like a transformer self-attention; 2) the path integration representation, g look like position encodings.
1) When considering the TEM memory retrieval process more closely (in this analysis, for direct comparison, we are only considering 1 attractor step in TEM with no non-linearity), we see that the attractor update qtMt = qt ∑t τ p T τ pτ is simply equal to
pretrievedt = t∑ τ [qtp T τ ]pτ (8)
1We note such ‘fast weights’ have previously been thought of as an alternative to the LSTM (Ba et al., 2016a).
2In fact, in the TEM code online, Wg and Wx are set as fixed weight matrices, where Wg sub-samples g and Wx transforms (compresses) x from a one-hot to a two-hot representation.
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)
TEM
g0 x0 g1 x1 g2 x2 g3 ?
x3
q0
k0 v0
q1
k1 v1
q2
k2 v2
q3
Input with position encoding
Query
Key, Value
Figure 3: Self-attention in (a) Transformers and (b) TEM.
Since [qtpTτ ] is just a dot-product ([qt · pτ ]), a single step of the attractor just retrieves memories weighted by their similarity (dot product) to the query. As noted by Ramsauer et al. (2020), this is exactly like a transformer but without the softmax scaling the dot-products. Thus the TEM memory retrieval process behaves like transformer self-attention.
2) We can however go further since TEM’s input to the transformer (i.e. the TEM memories) are special; they are learnable and built from an outer product between g̃ and x̃ (pτ = flatten(x̃Tτ g̃τ ) ), and these memories can be retrieved by a query based on g̃ or x̃ alone (e.g. qt = flatten(1T g̃t)). Together, these properties mean we can reduce the above dot product even further;
[qtp T τ ] = ¯̃xτ [g̃t · g̃τ ] → pretrievedt = g̃tG̃TΛxP (9)
Where ¯̃x = ∑ i(x̃τ )i and Λx is a diagonal matrix with elements ¯̃xτ (see Appendix for an alternative derivation using vector elements). Thus to retrieve a conjunctive p memory, all that was necessary is weighting past p representations via ‘self-attention’ of g̃t to past representations G̃.
To simplify this even further, we consider what happens when we ‘deconjunctify’ pretrievedt to obtain the sensory component of the memory. Following the TEM procedure described above (Figure 8);
x̃retrievedt = sum(unflatten(p retrieved t ), 1) = t∑ τ x̃τ ¯̃gτ ¯̃xτ [g̃t · g̃τ ] = g̃tG̃TΛgΛxX̃ (10)
Where Λg is a diagonal matrix with elements ¯̃gτ = ∑ i(g̃τ )i. Now all that is necessary to retrieve the sensory component of the memory is weighting past x̃ representations with via ‘self-attention’ of g̃t to past representations G̃. This equation is now very similar to equation 1 except without the softmax and with additional weightings Λx and Λg . These weighting however are likely learned to be constant (α) because otherwise some memories will be preferentially retrieved. In this case TEM is retrieving memories using
x̃retrievedt = (αg̃tG̃ T )X̃ cf. softmax(
g̃tG̃ T
√ dk )X̃ (11)
Which can be seen to be very closely related to the transformer equation (shown on the right), and diagrammatically shown in Figure 3. The model presented in this paper utilises the full transformer softmax rule.
The TEM-transformer. Thus the TEM-transformer (TEM-t; from Section 3) is this transformer that is directly analogous to TEM. Additional modelling details (analogous to modelling details in TEM) can be found in the Appendix. TEM-t offers dramatic performance improvements over the original TEM model (Figure 4; code will be released on publication). In particular, 1) Sample efficiency is increased; TEM-t requires many fewer data samples than TEM, and thus training time is reduced 2) TEM-t can tackle much larger problems, with the ability to store and retrieve many more memories (not shown here). Additionally to improved performance, TEM-t learns grid cells (Figure 1) and has potential implications for what place cells are (see next section).
Path integrating position encodings. This leads us to an interesting observation; we see that TEM’s representations for path integration g plays the role of position encodings in transformers. However the structure of these positional encodings are not hard-coded, but instead learned via path integration (the structure of space!), with the particular position encoding depending on the particular
sequence of actions taken. Other (non-spatial) structural representations could also be learned depending on the task at hand, i.e. grammar for language. This is a very different (and we think fruitful) re-understanding of position encodings; representing ‘location’ in a (learned) structure that can be inferred on the fly.

heading: 6 PLACE CELLS IN TRANSFORMERS
text: Here we discuss, and demonstrate, how TEM-t offers a new interpretation of place representations. To do so we utilise a recent suggestion of how the transformer update can be performed in biological hardware (Krotov & Hopfield, 2020). In particular, self-attention (equation 1) can be split into two steps which correspond to two pools of neurons (Figure 5A); 1) calculate softmax(qtK T
√ dk ). 2) multiply by V . In this light, K and V can simply be seen as weight matrices between feature neuron (representing the query) and memory neurons (computing the softmax).
Since memory neurons are sparsely activated due to the softmax, they appear to have a spatial tuning for each environment resembling hippocampal place cells (Figure 5D-E; note Krotov & Hopfield (2020) stated memory neurons may correspond to place cells but without simulation). Similarly to experimentally recorded place cells, these neurons remap randomly between environments i.e. place cells being neighbours in one environment is not predictive of them being neighbours in another (unlike grid cells which maintain their phase neighbours across environments).
We can curate this architecture for the specifics of TEM-t. TEM-t explicitly considers factorised g and x representations (e.g. MEC and LEC), which project to feature neurons in hippocampus (or still in cortex). Thus the feature neurons consist of two separate sub-populations, g̃ = gWg and x̃ = xWx, but which can connect to the same memory neurons in hippocampus (Figure 5B-C). These feature sub-populations are updated alternately rather than simultaneously, depending on the direction of retrieval; for example, when retrieving x̃ the g̃ feature neurons stay constant while the x̃ neurons are updated (in turn updating x in LEC). In this vein, hippocampal memories link together cortical representations in potentially disparate brain areas. Thus TEM-t instantiates hippocampal indexing theory (Teyler & Rudy, 2007), which states that hippocampus provides an index that binds together cortical patterns across different brain regions.
The randomly remapping place cells described one paragraph ago cannot be the full picture since we know that place cell remapping is not random; instead individual place cells preferentially remap to locations consistent particular grid cell firing (as predicted by conjunctive memory cells p in TEM and verified experimentally in Whittington et al. (2020)). However another mechanism for this
phenomena born from TEM-t could be as follows. Should the feature neurons exist in hippocampus (Figure 5F) then there will be hippocampal spatial cells g̃ that maintain their relationship to grid cells across different environment (as they are inherited from g via a projection g̃ = gWg). Thus across the population of hippocampal cells, there will be those that maintain their relationship to grid cells (e.g. g̃ and those that don’t (e.g. memory neurons and x̃), but the population effect will exist, just like what is experimentally observed.
As a note, the particular relationship of our model to the model of Krotov & Hopfield (2020), is what they refer to as a ‘type B’ model. These are models with contrastive normalisation on the memory neurons (via a softmax in our case), as opposed to ‘type B’ models which have a power activation function on the memory neurons. TEM (left hand side of Equation 11) corresponds to a linear activation function on the memory neurons, and is directly analogous to the original Hopfield energy. Secondly, the notion that there are two types of feature neurons that can be bound together in the same memory, was explored in Krotov & Hopfield (2016) where pixel intensities were associated with labels of those images. In TEM-t, one of the feature vectors, g, is learned via a RNN and structures itself according to the underlying task structure.
As an additional aside, we note that Krotov & Hopfield (2020) architectures does not solve the scaling problem of conventional Hopfield networks; the number of memories that original Hopfield networks could store scaled linearly with the dimensionality of the recurrent attractor network (Amit et al., 1985). While recent analytical work has shown with exponential power activation functions, the number of memories that can be stored to scale as 2 N 2 , where {N} is the dimensionality of the feature neurons (Demircigil et al., 2017). This is a considerably more favourable scaling. However, unfortunately the architecture from Krotov & Hopfield (2020) instead requires a growing number of memory neurons (one for each memory), so the number of memories is still linear with the number of neurons! We note that mathematically derived scaling law was for an exponential activation function, not with a softmax as we use here.

heading: 7 DISCUSSION
text: We have shown that TEM, a current model of the hippocampal formation, is closely related to a transformer with recurrent position encodings. We now consider some wider implications for neuroscience.
Multiple cortical inputs to hippocampus. TEM considers hippocampal conjunctions between two cortical regions (g and x). It is, however, possible to consider conjunctions of more than two brain
regions. Indeed hippocampal neurons often respond to more than two task variables (McKenzie et al., 2014). In TEM, the naive approach of a ‘triple’ (or higher) conjunction would increase the number of hippocampal neurons would increase by a factor of nc; the number of neurons from brain region c̃. TEM-t does not scale so badly. Instead it just requires an additional nc feature neurons, and the number of memory neurons can stay the same since the each hippocampal memory neuron can simply index a memory across three (or more), rather than two, brain regions (Figure 5G).
With multiple inputs to hippocampus [x̃, g̃, c̃, . . . ], any subset of those brain areas can reinstate a memory in the other brain regions i.e. x̃ and g̃ can reinstate a c̃ memory or g̃ alone could reinstate x̃ and c̃ memories. As an analogy to the TEM triple conjunction, TEM-t proposes that c̃t is updated via c̃t ← softmax((g̃tG̃T ) (x̃tX̃T ))C̃, where is an element wise product. We note an alternate, and perhaps more intuitive, option could also be c̃t = softmax(g̃tG̃T + x̃tX̃T )C̃.
Beyond hippocampus: Cortex as a Transformer. We have considered transformers as a model of hippocampus and its connections. We know, however, that transformer representations predict language areas (Schrimpf et al., 2020), and that patients can talk and comprehend just fine with major hippocampal deficits (Elward & Vargha-Khadem, 2018). This indicates that the transformer, and TEM-like models, may also model other brain regions, such as language areas, that are seemingly independent from hippocampus (related ideas discussed in Hawkins et al. (2019); Lewis (2021) but specifically for grid cells in neocortex). This raises two questions. Firstly what is the analogue of spatial positional encodings for higher order tasks such as language, and secondly what takes the role of the memory neurons if not hippocampus. We offer some thoughts in the following two paragraphs.
In spatial tasks, TEM and TEM-t learn positional encodings that mirror the structure of space. The implication is that positional encoding should reflect the abstract underlying properties of the task at hand. In language for example, this structure is grammar. This contrasts to the typical positional encodings in Transformers - sines and cosines - which represent a linear structure. It is our contention that positional encodings that are inferred on the fly and consist of previously learned structures (like the spatial case we have considered) would offer an interesting and potentially fruitful research direction in problems of language, maths, and logic.
If the transformer were solely instantiated in cortex, then what about the memory neurons? It is possible that the memory neuron equivalent exists in cortex too, but for these tasks, since it is not necessary to store long term memories or bind knowledge across multiple brain areas hippocampus is not required; so short term cortical memory neurons suffice.

heading: 8 CONCLUSION
text: We have shown that transformers with recurrent positional encodings reproduce neural representations found in rodent entorhinal cortex and hippocampus. We then showed these transformers are close cousins to models of hippocampus that neuroscientists have developed over the last few years. We hope this work brings neuroscience and machine learning closer together, and offers understanding for both sides; for neuroscientists a road map to understanding cortical areas beyond the hippocampal formation; for machine learners a greater understanding of positional encodings in transformers.

heading: A APPENDIX

heading: A.1 THE MATHS USING ELEMENTS
text: Here we derive the main results using vector and matrix elements. Since each element in p is multiplicative combination of every pair of elements in x̃ and g̃, then
pij = g̃ix̃j (12)
Where we label the elements of vector p with two indices even though it is a vector. Similarly, the memory matrix uses 4 indices; Mijkl = ∑ τ g̃τi x̃ τ j g̃ τ k x̃ τ l (13)
Retrieving a memory via path integration. Querying the network becomes qij = g̃PIi , which reduces to ∑
ijk qijMijkl = ∑ τ ∑ ijk g̃PIi g̃ τ i x̃ τ j g̃ τ k x̃ τ l
= ∑ τ x̃τl ( ∑ k g̃τk)( ∑ i g̃PIi g̃ τ i )( ∑ j x̃τj )
→ g̃PIt X̃TΛgΛxG̃
(14)
Retrieving a memory using a sensory observation alone. Similarly when the query is the sensory observation alone qij = xj
∑ ijl qijMijkl = ∑ ijlτ x̃j g̃ τ i x̃ τ j g̃ τ k x̃ τ l
= ∑ τ g̃τk( ∑ l x̃τl )( ∑ i g̃τi )( ∑ j x̃j x̃ τ j )
→ x̃tX̃TΛgΛxG̃
(15)
Retrieving a memory using a sensory observation and path integration. When the query is qij = g PI i xj , and we want to retrieve a position encoding memory
∑ ijl qijMijkl = ∑ ijlτ g̃PIi x̃j g̃ τ i x̃ τ j g̃ τ k x̃ τ l
= ∑ τ g̃τk( ∑ l x̃τl )( ∑ i g̃PIi g̃ τ i )( ∑ j x̃j x̃ τ j )
→ (x̃tX̃T g̃PIt G̃T )ΛgG̃
(16)

heading: A.2 DETAILS OF IMPLEMENTATION
text: We will release code on publication.
Scaling beta. Since the number of memories is not constant, we use an adaptive β parameter in the softmax. This is because normalisation term in the softmax sums the number of memories, and so more memories down-weights probabilities. We want a self-attention value to not be affected by the number of elements in the set. In particular, we weight the softmax by log(nmemories).
Losses. We use same set of losses as in TEM;
• a one-step sensory prediction cross entropy loss, i.e. using gPIt as input to the memory retrieval process
• a sensory prediction cross entropy loss using gt as input to the memory retrieval process • a squared error loss between gt and gPIt • l2 weight regularisation • l2 regularisation on gt
Normalisation. We find that using layernorm (Ba et al., 2016b) on the positional encodings (not in the RNN, but on the input to transformer) to be beneficial, since the memory retrieval process can then be standardised - no one memory is up-weighted relative to others. Using layernorm before
self-attention in transformers is common practice (Vaswani et al., 2017). For simplicity, we use fixed weights on the layer norm, i.e. is is just a z score of g, Since we use a one-hot encodings of our sensory representations, they are already normalised.
Stabilising position representations. Recurrently generated positional encodings accumulate noise and drift. While bespoke path integration networks from neuroscience mitigate noise by enforcing their neural representations to stay close to a neural manifold (Burak & Fiete, 2009), this can not be guaranteed in learned recurrent networks. One method of stabilisation is via sensory landmarks - i.e. ‘what positional encoding did I have the last time I saw this landmark’. In this vein TEM uses the following query to the memory network; qt = flatten(x̃Tt 1), and memories of positional encodings are retrieved as;
g̃retrievedt = sum(unflatten(qtMt), 0) = t∑ τ g̃τ ¯̃gτ ¯̃xτ [x̃t · x̃τ ] = x̃tX̃TΛgΛxG̃ (17)
The final position encoding, g is computed on a basis of path integration (gPIt ; equation 3) and stored landmark information (gretrievedt ; equation 17), i.e. gt = g PI t + fθ(g retrieved t , g PI t )(g retrieved t − gPIt ), where g retrieved = fθ(g̃
retrieved) and fθ(· · · ) are different MLPs. We note that a better query to the memory network would be qt = flatten(x̃Tt g PI t ) since sensory observations may be aliased, including gPIt can help disambiguate such aliasing. In this case, the retrieved memory is (x̃tXT gPIt G)ΛxG. This is perhaps, different to what would be anticipated; using an element wise product rather than an addition. Translating this other TEM memory retrieval process into transformers we get gretrievedt = fθ(softmax(x̃tX
T gPIt G)G), this can be thought of as another attention head i.e. from input [g,x], the key and query ‘attends’ to x while the value ‘attends’ to g.
When to add memories. Memories should not be added at every step, otherwise the self-attention mechanism would be biased towards memories (locations) that have been visited more frequently. This is not a problem typically addressed in transformer, but here to mitigate this issue memories are only added when existing memories for that particular conjunction do not already exist i.e. if there is a memory already with the present combination of g and x (similarity determined by dot product) then no new memory is added 3.
Multiple iterations: Though in our simulations we only used a single iteration of the transformer block, it is possible to use multiple iterations - just like a Hopfield network (as in Ramsauer et al. (2020)). Here the retrieved sensory memory x̃retrievedt can be fed into the next iteration, along with positional encoding i.e. the path integrated g. In this case (see Appendix for derivation), TEM suggests that x̃t is iteratively updated via x̃t ← softmax(x̃tX̃T g̃tG̃)X̃ , but with the first iteration via x̃t ← softmax(g̃tG̃)X̃ . The multiplicative term in the softmax is perhaps unexpected - it may be thought that it should be additive instead.
Place-like representations without a softmax: Here we chose a softmax activation function on the memory neurons to make the relationship between TEM and transformers exact. However was this choice necessary to observe place cells in the memory neurons? While we have not examined this experimentally, we suspect place-like representations would emerge for a variety of activations, since memories must still be sparsely activated for correct prediction. We leave this for future work to verify, but note again that power and exponential activations have been shown to have a much greater memory capacity (Krotov & Hopfield, 2016; Demircigil et al., 2017) than traditional Hopfield nets with linear activation.

heading: A.3 ADDITIONAL DETAILS OF THE SPATIAL TASK
text: We formalise a task type that not only relates to known hippocampal function, but also tests the learning and generalising of abstract structural knowledge. We formalise this via relational understanding on graph structures (a graph is a set of nodes that relate to each other).
3The TEM model actually does something very much like this - memories are only added if they did not already exist - M = ∑ τ (pτ − p̂τ )
T (pτ + p̂τ ), where p̂τ is the memory that was retrieved at τ , i.e. only the un-predicted components of the memory pτ are added.
Should one passively move on a graph (e.g. Figure 6), where each node is associated with a nonunique sensory observation (e.g. an image of a banana), then predicting the subsequent sensory observation tests whether you understand the graph structure you are in. For example, if you return to a previously visited node (Figure 6 pink) by a new direction - it is only possible to predict correctly if you know that a right→ down→ left→ up means you’re back in the same place. Knowledge of such loop closures is equivalent to understanding the structure of the graph.
We thus train our models on these graphs with it trying to predict the next sensory observation. Our models are trained on many environments sharing the same structure, e.g. 2D graphs (Figure 6), however the stimulus distribution is different (each vertex is randomly assigned a stimulus). Should it be able to learn and generalise this structural knowledge, then it should be able to enter new environments (structurally similar but with different stimulus distributions) and perform feats of loop closure on first presentation.
Formally, given data of the form D = {(xk≤T ,ak≤T )} with k ∈ {1, · · · , N} (which environment it is in), where x≤T and a≤T are a sequence of sensory observations and associated actions/relations (Figure 6), N is the number of environments in the dataset, and T is the duration of time-steps in each environment, our model should maximise its probability of observing the sensory observations for each environment.
The sensory stimuli are chosen randomly, with replacement, at each node. We understand that this is not like the real world, where adjacent locations have sensory correlations - most notable in space (though names in a family tree will be less correlated). Sensory correlations help with sensory predictions, thus if we use environments with sensory correlations, we would not know what was causing the learned representations, sensory correlations, or transition structure. To answer this question cleanly, and to know that transition structure is the sole cause, we do not use environments with sensory correlations.

heading: A.4 ADDITIONAL DETAILS OF TEM
text: We present a more detailed model schematic of TEM in Figure 7. We see there are two components to TEM - a RNN for understanding position (g, in green top of Figure 7) that also indexes memories via ‘queries’ q = Wgg. A memory network that binds together x and g, via an outer product (middle green in 7, with more detail in Figure 8A). When the memory network is queried (red in Figure 7), it undergoes attractor dynamics to retrieve the full memory. To make a sensory prediction, the retrieved memory is ‘deconjunctified’ into a sensory representation (Figure 8C).
Model flow. TEM transitions through time and infers gt and pt at each time-step. gt is inferred before forming each new memory pt. In other words variables g and p are inferred in the following order gt, pt, gt+1, pt+1, gt+2 · · ·. This flow of information is shown in a schematic in Figure 7.
Independently, at each time-step, the model model asks ‘are the inferred variables what I would have predicted given my current understanding of the world (weights)’. I.e. 1) Is the inferred gt the one I would have predicted from gt−1. 2) Is the inferred pt the one I would have predicted from gt. 3) Is xt what I would have predicted from pt. This leads to errors (at each timestep) between inferred and predicted variables gt and pt, and between sensory data xt and its prediction.
At then end of a sequence, these errors are accumulated, with model parameters updated along the gradient (from back-propagation through time) that matches each others variables and also matches the data.
Since the model runs along uninterrupted, it’s activity at one time-step influence those at later timesteps. Thus when learning (using back-propagation through time - BPTT), gradient information flows backwards in time. This is important as, should a bad memory be formed at one-time step, it will have consequences for later predictions - thus BPTT allows us to learn how to form memories and latent representations such that they will be useful many steps into the future.

heading: A.5 RELATED WORK
text: Here we discuss and compare other models that recapitulate spatial representations found in the brain. These models fall into several categories. 1) Auto-encoder like models trained on spatial representations (Dordek et al., 2016). 2) Graph representation models (Gustafson & Daw, 2011; Stachenfeld et al., 2017). 3) Latent state inference models (George et al., 2021). 4) Path integrating
models with RNNs trained on spatial representations (Cueva & Wei, 2018; Banino et al., 2018). 5) Models mixing RNNs and memory networks that are trained on sequences of sensory observations.
We note that models (1-4) are learned using curated spatial representation, i.e. the modeller has already worked how how space works, and the model is finding an alternate representation of that space. This is not how you or I learn. Instead, we learn by extracting regularities from the sensory world, and transfer/generalise this knowledge from domain to domain. Model category (5) does unsupervised learning by predicting sequences of sensory observations alone. From just sensory observations, it can slowly build up a picture of how space is structured, and then transfer this knowledge to situations that share the same underlying spatial structure. Our model TEM-t fall in the model (5) category, along with TEM (Whittington et al., 2020) and other similar models Uria et al. (2020).
For model to make sensory predictions as fast as possible, it is necessary to have two components. a) A RNN that understand position, and b) a memory network that can remember what you see and where you see them. Now the model can be asked what ’what will you see when heading East’ and correctly respond Cat even it it has never gone East at that location before. This is because it can simulate going East via its RNN and then retrieve the memory it had stored at that location (it must have visited the Cat location before thought!). The other models, e.g. an autoencoder, could not solve this task since it learns slow statistical structures between sensory observations over many many training examples. Thus it would be clueless to the entirely new configuration of sensory observations in the new environment. In sum, it is necessary to have an understanding of position, and the ability to make and retrieve memories to successfully make sensory predictions as fast as possible.
We describe the relationship between TEM (and thus the hippocampal-entorhinal system) as close not because it produces the same representations, instead we are saying the relationship is close because we have shown a mathematical relationship between the two models. This could not happen for most ML models as most ML systems do not have both RNN structured representations and a memory system and so is not mathematically relatable to current models of the hippocampal
formation. However, equipped with the mathematical relationship, we can say that the memory part of the transformer is related to the hippocampal memory system in TEM. And the positional encodings are related to the entorhinal grid RNN system in TEM.

heading: A.6 FURTHER LEARNED CELL REPRESENTATIONS

Review summaries:
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.

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. 

This well-written and clear paper clarifies the relationship between transformers and a recent exciting model of the medial temporal lobe in neuroscience. The authors find that a transformer with recurrent positional encodings trained on a spatial navigation task ends up learning representations that resemble some classic findings from the medial temporal lobe in the neuroscience (e.g., grid and place cells). This finding alone is not that surprising -- as the authors note, a variety of groups have previously found such representations are learned in navigation tasks with other models (particularly place cells). The major contribution of this paper is that the authors then demonstrate the connection between transformers and a contemporary model of the hippocampal formation, the Tolman-Eichenbaum Machine (TEM). Specifically, the memory retrieval process in the TEM resembles self-attention and the path integration representations are comparable to a transformer's learned position encodings. I believe that by offering a novel perspective on transformers and models of biological memory systems, this work may lead to fruitful future work both in machine learning and neuroscience.