R ELATING TRANSFORMERS TO MODELS AND NEURAL REPRESENTATIONS OF THE HIPPOCAMPAL FORMATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Many deep neural network architectures loosely based on brain networks have recently been shown to replicate neural ﬁring 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 artiﬁcial 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
I NTRODUCTION
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 ﬁnally it tells us something formal about position encodings in transformers.
T RANSFORMERS
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 = { x 1 , x 2 , x 3 , · · · , x T } ( x t 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. x t>T , or could be a missing part of a sentence or image, i.e.
{ x 1 = the , x 2 = cat , x 3 = sat , x 4 = ?
, x 5 = the , x 6 = 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 4 th element ( ? ) could attend to the 2 nd ( cat ), 3 rd ( sat ), and 6 th ( mat ) to understand it should be on .
Formally, to attend to another element each element ( x t is a row vector) emits a query ( q t
= x t W q ) and compares it to other elements keys ( k τ = x
t W k ).
Each element is then updated using y t = (cid:80) τ κ ( q t , k τ ) v τ , where κ ( q t , k τ ) is kernel describing the similarity of q t to k τ and v τ is the value computed by each element v τ = x t W v .
Intuitively, the similarity measure κ ( q t , 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 κ ( q t , k τ ) , i.e. cat, sat, mat ) are ‘attended’ to and therefore contribute more to the output y t .
Typically, the similarity measure is a softmax i.e. κ ( q t , k τ ) =
[Equation]
These equations can be succinctly expressed in matrix form, with all elements updated simultaneously:
[Equation]
Here Q , K , V are matrices with rows ﬁlled by q t , k t , v t respectively, and the softmax is taken independently for each row.
After this update, each y t 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 z t = f θ ( y t ) .
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 identiﬁes each time-step ( e t - typically sines and cosines) to each input: x t ← x t + e t .
Alternatively the position embedding can be appended i.e. h t =
[ x t , e t ] , with self attention then performed using h t as input.
T RANSFORMERS LEARN ENTORHINAL REPRESENTATIONS
Here we show that transformers (with a small modiﬁcation) recapitulate spatial representationsgrid 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 x t +1 conditioned on taking an action a t while moving around spatial environments (Figure 1a).
For example, after seeing { ( x 1 = cat , a 1 = North ) , ( x 2 = dog , a 2 = East ) , ( x 3 = frog , a 3 = South ) , ( x 4 = pig , a 4 = West ) , ( x 5 = ? , a 5 = · · · ) } , the aim is to predict x 5 = cat .
For simplicity, we treat sensory observations as one-hot vectors, thus the prediction problem is a classiﬁcation problem.
exexexe?
yqkvqkvqkvqiW v iW q iW k
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
[Figure]
Environment
Observations (x) Actions (a) A B A ?
C A A D ?
D
Recurrent connections W(a 0 )
W(a 1 )
W(a 2 )
[Figure] activations
With ReLu activations
[Figure]
Autocorrelations
Cell [Figure]
Autocorrelations
[Figure]
When faced with an unseen stimulus-action pair (e.g. x 4 = pig , a 4 = West above; an action you have never taken at that stimulus before), successful prediction requires more than just remembering speciﬁc sequences of stimulus-action pairs; knowledge of the rules of space must be known; i.e. North + East + South + West = 0 allows prediction of x 5 = 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 modiﬁcations to the transformer are:
Recall equation 1; y t = softmax ( q t K T √ d k ) V , where Q = HW q , K = HW k , V = HW v , and H is a matrix of inputs and position encodings (i e. its rows are h t =
[ x t , e t ] ).
We restrict these weight matrices such that queries ( Q ) and keys ( K ) are the same; Q , K = EW e .
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 = XW x .
We refer to this matrix as ˜ X .
[Equation]
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.
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 e t .
The position encodings are recurrently generated (as in Wang et al.
Liu et al.
; e t +1 =
σ ( e t W a ) , where W a 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 modiﬁcations are sufﬁcient 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.
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 Hopﬁeld network (Hopﬁeld, 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.
TEM
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:
[Equation]
Where W a is a learnable action dependent weight matrix and σ ( · · · ) is a non-linear activation function.
It is in this path-integrating representation g 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 speciﬁc 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 ﬂattened back into a vector;
[Equation]
These conjunctive p representations are stored in ‘fast weights’ 1 via Hebbian learning;
[Equation]
And they can later be retrieved using an attractor network (a continuous version of the Hopﬁeld network).
Here a query vector q (details next paragraph) is inputted into the network and updated via;
[Equation]
) 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 ( 1 T g ) .
Equation 6 is then iterated until a memory is retrieved.
Finally, to make sensory predictions, the retrieved conjunctive memory ( p retrievedt ) is ‘deconjunctiﬁed’ into sensory and location components.
The sensory component is obtained by unﬂattening p retrievedt and summing over the g dimension (Figure 8);
[Equation]
Finally, to make the sensory prediction x retrievedt is fed through a MLP z t = f θ ( x retrievedt ) ) to classify (predict) the upcoming sensory observation.
It is also possible, and often helpful, to project g and x via W g and W x ; ˜ g = gW g and ˜ x
= xW x before they are combined conjunctively 2 .
TEM AS A TRANSFORMER
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 q t M t =
q t (cid:80) tτ p
Tτ p τ is simply equal to
[Equation]
exexexe?
yqkvqkvqkvqiW v iW q iW k
Transformer (with causal attention mask)
TEM gxgxgxg?
xqkvqkvqkvq
Input with position encoding
Query
Key, Value
Since [ q t p Tτ ] is just a dot-product ( [ q t · 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. 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. q t = flatten ( 1 T ˜ g t ) ).
Together, these properties mean we can reduce the above dot product even further;
[Equation]
Where ¯˜ x = (cid:80) 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’ p retrievedt to obtain the sensory component of the memory .
Following the TEM procedure described above (Figure 8);
[Equation]
Where Λ g is a diagonal matrix with elements ¯˜ g τ = (cid:80)
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
[Equation]
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 efﬁciency 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 [Figure] 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 ﬂy.
P LACE CELLS IN TRANSFORMERS
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 & Hopﬁeld, 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 ( q t K T √ d k ) .
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 & Hopﬁeld (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 speciﬁcs 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 = gW g and ˜ x
= xW x , 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 ﬁring (as predicted by conjunctive memory cells p in TEM and veriﬁed experimentally in Whittington et al. However another mechanism for this [Table]
+v7kHY7lWP3GZQyzpTPGMM4ouNT3o/OoeThAWuGpVftMGLF4YAFWVk0SWE+QihXJWVdOo2hkd7o4e74ZeVFUbPBcFtIjFtsV/MsMdGGfqMWS8aeYHpC2ANV4eJ4b/b3JcTXv98ChcRbAtokb0SROX096fSapZIUEhE9Ta2yjMMS6pQc4EuMEKCzlld3QGt04qKsHG5cpCFRy6TBpk2rhPYbDKPqwoqbR2KRNHSopzu3lWJwfuX1fbWuBcDhK5cStmX+KSq7xAUGx9aVaIAHVQP9Eg5QYYiqUTlBnufAdsTg1l6B5yq1N6z3PbTLBYj9B1C4s217MtrodHkdNXJ/2z82Z1++Q9+UA+koickjPynVySEWHenjfwPnmf/VM/9lM/W6Nep6l5R1rh679LTkPx</latexit>
G [Figure] 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 = gW g ).
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 & Hopﬁeld (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 Hopﬁeld 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 & Hopﬁeld (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 & Hopﬁeld (2020) architectures does not solve the scaling problem of conventional Hopﬁeld networks; the number of memories that original Hopﬁeld 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 & Hopﬁeld (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.
D ISCUSSION
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 n c ; the number of neurons from brain region ˜ c .
TEM-t does not scale so badly.
Instead it just requires an additional n c 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 t ˜ G T ) (cid:12) ( ˜ x t ˜ X T ))
˜ C , where (cid:12) is an element wise product.
We note an alternate, and perhaps more intuitive, option could also be ˜ c t = softmax ( ˜ g t ˜ G T + ˜ x t ˜ X 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 ﬁne with major hippocampal deﬁcits (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.
; Lewis (2021) but speciﬁcally 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 reﬂect 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 ﬂy 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 sufﬁce.
C ONCLUSION
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.
A A PPENDIX
A. T MATHS USING ELEMENTS
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
[Equation]
Where we label the elements of vector p with two indices even though it is a vector.
Similarly, the memory matrix uses 4 indices; [Equation]
[Equation]
Retrieving a memory via path integration.
Querying the network becomes q ij = ˜
g PIi , which reduces to
[Equation]
Retrieving a memory using a sensory observation alone.
Similarly when the query is the sensory observation alone q ij = x j
[Equation]
Retrieving a memory using a sensory observation and path integration.
When the query is q ij = g PIi x j , and we want to retrieve a position encoding memory
[Equation]
A. D ETAILS OF IMPLEMENTATION
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 ( n memories ) .
Losses.
We use same set of losses as in TEM;
• a one-step sensory prediction cross entropy loss, i.e. using g PIt as input to the memory retrieval process • a sensory prediction cross entropy loss using g t as input to the memory retrieval process • a squared error loss between g t and g PIt • l2 weight regularisation • l2 regularisation on g t
Normalisation.
We ﬁnd that using layernorm (Ba et al., 2016b) on the positional encodings (not in the RNN, but on the input to transformer) to be beneﬁcial, 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 ﬁxed 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; q t = flatten ( ˜ x
Tt 1 ) , and memories of positional encodings are retrieved as;
[Equation]
The ﬁnal position encoding, g is computed on a basis of path integration ( g PIt ; equation 3) and stored landmark information ( g retrievedt ; equation 17), i.e. g t =
g PIt + f θ ( g retrievedt , g PIt )( g retrievedt − g PIt ) , where g retrieved =
f θ ( ˜ g retrieved ) and f θ ( · · · ) are different MLPs.
We note that a better query to the memory network would be q t
= flatten ( ˜ x
Tt g PIt ) since sensory observations may be aliased, including g PIt can help disambiguate such aliasing.
In this case, the retrieved memory is ( ˜ x t X T (cid:12) g PIt G ) Λ x G .
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 g retrievedt =
f θ ( softmax ( ˜ x t X T (cid:12) g PIt 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 Hopﬁeld network (as in Ramsauer et al.
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 t ˜ X T (cid:12) ˜
g t ˜ G ) ˜ X , but with the ﬁrst iteration via ˜ x t ← softmax ( ˜ g t ˜ G ) ˜ 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 & Hopﬁeld, 2016; Demircigil et al., 2017) than traditional
Hopﬁeld nets with linear activation.
A. A DDITIONAL DETAILS OF THE SPATIAL TASK
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).
[Figure]
[Figure]
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 ﬁrst presentation.
Formally, given data of the form D = { ( x k ≤ T , a k ≤ 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.
A. A DDITIONAL DETAILS OF TEM
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 = W g g .
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 ‘deconjunctiﬁed’ into a sensory representation (Figure 8C).
Model ﬂow.
TEM transitions through time and infers g t and p t at each time-step.
g t is inferred before forming each new memory p t .
In other words variables g and p are inferred in the following order g t , p t , g t +1 , p t +1 , g t +2 · · · .
This ﬂow of information is shown in a schematic in Figure 7.
+ action a t+
Form memory p from g and x, via outer product, then change weights M
Query memory network, and retrieve memory p via attractor dynamics
Path integrate new g
Sensory data
Predict sensory from p
Step
[Equation]
Step
[Equation]
g
[Figure]
[Figure]
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 g t the one
I would have predicted from g t − 1 . 2)
Is the inferred p t the one I would have predicted from g t . 3)
Is x t what I would have predicted from p t .
This leads to errors (at each timestep) between inferred and predicted variables g t and p t , and between sensory data x t 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 inﬂuence those at later timesteps.
Thus when learning (using back-propagation through time - BPTT), gradient information ﬂows 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.
A. R ELATED WORK
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).
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).
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 ﬁnding 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.
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 conﬁguration 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.
A. F URTHER LEARNED CELL REPRESENTATIONS
[Figure]
[Figure]