Relating transformers to models and neural representations of the hippocampal formation β part 1/4
Cleaned-text code points [0, 15846). Bibliography omitted here; source.md includes it.
# Relating transformers to models and neural representations of the hippocampal formation
###### 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.
## 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 mathematical 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; π = { π 1 , π 2 , π 3 , β― , π T } {\mathbb{X}}=\{{\bm{x}}_{1},{\bm{x}}_{2},{\bm{x}}_{3},\cdots,{\bm{x}}_{T}\} ( π t {\bm{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. π t > T {\bm{x}}_{t>T} , or could be a missing part of a sentence or image, i.e. { π 1 = the , π 2 = cat , π 3 = sat , π 4 = ? , π 5 = the , π 6 = mat } \{{\bm{x}}_{1}=\texttt{the},{\bm{x}}_{2}=\texttt{cat},{\bm{x}}_{3}=\texttt{sat},{\bm{x}}_{4}=\texttt{?},{\bm{x}}_{5}=\texttt{the},{\bm{x}}_{6}=\texttt{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 ( π t {\bm{x}}_{t} is a row vector) emits a query ( π t = π t β πΎ q {\bm{q}}_{t}={\bm{x}}_{t}{\bm{W}}_{q} ) and compares it to other elements keys ( π Ο = π t β πΎ k {\bm{k}}_{\tau}={\bm{x}}_{t}{\bm{W}}_{k} ). Each element is then updated using π t = β Ο ΞΊ β‘ ( π t , π Ο ) β π Ο {\bm{y}}_{t}=\sum_{\tau}\kappa({\bm{q}}_{t},{\bm{k}}_{\tau}){\bm{v}}_{\tau} , where ΞΊ β‘ ( π t , π Ο ) \kappa({\bm{q}}_{t},{\bm{k}}_{\tau}) is kernel describing the similarity of π t {\bm{q}}_{t} to π Ο {\bm{k}}_{\tau} and π Ο {\bm{v}}_{\tau} is the value computed by each element π Ο = π t β πΎ v {\bm{v}}_{\tau}={\bm{x}}_{t}{\bm{W}}_{v} . Intuitively, the similarity measure ΞΊ β‘ ( π t , π Ο ) \kappa({\bm{q}}_{t},{\bm{k}}_{\tau}) 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 ΞΊ β‘ ( π t , π Ο ) \kappa({\bm{q}}_{t},{\bm{k}}_{\tau}) , i.e. cat, sat, mat ) are βattendedβ to and therefore contribute more to the output y t y_{t} . Typically, the similarity measure is a softmax i.e. ΞΊ β‘ ( π t , π Ο ) = e Ξ² β π t β
π Ο β Ο β² e Ξ² β π t β
π Ο β² \kappa({\bm{q}}_{t},{\bm{k}}_{\tau})=\frac{e^{\beta{\bm{q}}_{t}\cdot{\bm{k}}_{\tau}}}{\sum_{\tau^{\prime}}e^{\beta{\bm{q}}_{t}\cdot{\bm{k}}_{\tau^{\prime}}}} . These equations can be succinctly expressed in matrix form, with all elements updated simultaneously: π t = s β o β f β t β m β a β x β ( π t β π² T d k ) β π½ β π = s β o β f β t β m β a β x β ( πΈ β π² T d k ) β π½ {\bm{y}}_{t}=softmax(\frac{{\bm{q}}_{t}{\bm{K}}^{T}}{\sqrt{d_{k}}}){\bm{V}}\qquad\rightarrow\qquad{\bm{Y}}=softmax(\frac{{\bm{Q}}{\bm{K}}^{T}}{\sqrt{d_{k}}}){\bm{V}}\\ (1)
Here πΈ {\bm{Q}} , π² {\bm{K}} , π½ {\bm{V}} are matrices with rows filled by π t {\bm{q}}_{t} , π t {\bm{k}}_{t} , π t {\bm{v}}_{t} respectively, and the softmax is taken independently for each row. After this update, each π t {\bm{y}}_{t} is then sent through a deep network ( f ΞΈ β ( β― ) f_{\theta}(\cdots) ) typically consisting of residual ( He et al., 2016 ) and layer-normalisation ( Ba et al., 2016b ) layers to produce π t = f ΞΈ β ( π t ) {\bm{z}}_{t}=f_{\theta}({\bm{y}}_{t}) . π {\bm{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 π {\bm{x}} is in the sequence. This is typically done by adding a βposition encodingβ that uniquely identifies each time-step ( π t {\bm{e}}_{t} - typically sines and cosines) to each input: π t β π t + π t {\bm{x}}_{t}\leftarrow{\bm{x}}_{t}+{\bm{e}}_{t} . Alternatively the position embedding can be appended i.e. π t = [ π t , π t ] {\bm{h}}_{t}=[{\bm{x}}_{t},{\bm{e}}_{t}] , with self attention then performed using π t {\bm{h}}_{t} 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 π t + 1 {\bm{x}}_{t+1} conditioned on taking an action π t {\bm{a}}_{t} while moving around spatial environments (Figure 1 a). For example, after seeing { ( π 1 = cat , π 1 = North ) , ( π 2 = dog , π 2 = East ) , ( π 3 = frog , π 3 = South ) , ( π 4 = pig , π 4 = West ) , ( π 5 = ? , π 5 = β― ) } \{({\bm{x}}_{1}=\texttt{cat},{\bm{a}}_{1}=\texttt{North}),({\bm{x}}_{2}=\texttt{dog},{\bm{a}}_{2}=\texttt{East}),({\bm{x}}_{3}=\texttt{frog},{\bm{a}}_{3}=\texttt{South}),({\bm{x}}_{4}=\texttt{pig},{\bm{a}}_{4}=\texttt{West}),({\bm{x}}_{5}=\texttt{?},{\bm{a}}_{5}=\cdots)\} , the aim is to predict π 5 = cat {\bm{x}}_{5}=\texttt{cat} . For simplicity, we treat sensory observations as one-hot vectors, thus the prediction problem is a classification problem.
When faced with an unseen stimulus-action pair (e.g. π 4 = pig , π 4 = West {\bm{x}}_{4}=\texttt{pig},{\bm{a}}_{4}=\texttt{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 π 5 = cat {\bm{x}}_{5}=\texttt{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 1 a) - 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 ; π t = s β o β f β t β m β a β x β ( π t β π² T d k ) β π½ {\bm{y}}_{t}=softmax(\frac{{\bm{q}}_{t}{\bm{K}}^{T}}{\sqrt{d_{k}}}){\bm{V}} , where πΈ = π― β πΎ q {\bm{Q}}={\bm{H}}{\bm{W}}_{q} , π² = π― β πΎ k {\bm{K}}={\bm{H}}{\bm{W}}_{k} , π½ = π― β πΎ v {\bm{V}}={\bm{H}}{\bm{W}}_{v} , and π― {\bm{H}} is a matrix of inputs and position encodings (i.e. its rows are π t = [ π t , π t ] {\bm{h}}_{t}=[{\bm{x}}_{t},{\bm{e}}_{t}] ). We restrict these weight matrices such that queries ( πΈ {\bm{Q}} ) and keys ( π² {\bm{K}} ) are the same; πΈ , π² = π¬ β πΎ e {\bm{Q}},{\bm{K}}={\bm{E}}{\bm{W}}_{e} . We refer to this matrix as π¬ ~ {\bm{\tilde{E}}} . Thus the keys and queries only focus on position encodings. Meanwhile, values are exclusively dependent on the stimulus component of π― {\bm{H}} i.e. π½ = πΏ β πΎ x {\bm{V}}={\bm{X}}{\bm{W}}_{x} . We refer to this matrix as πΏ ~ {\bm{\tilde{X}}} . π t = s β o β f β t β m β a β x β ( π t β π² T d k ) β π½ β π t = s β o β f β t β m β a β x β ( π ~ t β π¬ ~ T d k ) β πΏ ~ {\bm{y}}_{t}=softmax(\frac{{\bm{q}}_{t}{\bm{K}}^{T}}{\sqrt{d_{k}}}){\bm{V}}\qquad\rightarrow\qquad{\bm{y}}_{t}=softmax(\frac{{\bm{\tilde{e}}}_{t}{\bm{\tilde{E}}}^{T}}{\sqrt{d_{k}}}){\bm{\tilde{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. π < t {\bm{e}}_{<t} and π < t {\bm{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 π t {\bm{e}}_{t} .
3. The position encodings are recurrently generated (as in Wang et al. (2019) ; Liu et al. (2020) ); π t + 1 = Ο β‘ ( π t β πΎ a ) {\bm{e}}_{t+1}=\sigma({\bm{e}}_{t}{\bm{W}}_{a}) , where πΎ a {\bm{W}}_{a} is a learnable action-dependent weight matrix, and Ο β‘ ( β― ) \sigma(\cdots) 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 1 C; 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 self-attention, 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.
## 4 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 π {\bm{g}} to represent location. To update its location, TEM uses path-integration - the accumulation of self movement vectors π {\bm{a}} - enacted in a recurrent neural network: π t + 1 = Ο β‘ ( π t β πΎ a ) {\bm{g}}_{t+1}=\sigma({\bm{g}}_{t}{\bm{W}}_{a}) (3) Where πΎ a {\bm{W}}_{a} is a learnable action dependent weight matrix and Ο β‘ ( β― ) \sigma(\cdots) is a non-linear activation function. It is in this path-integrating representation π {\bm{g}} that TEM learns grid and other entorhinal cells for self-localisation (Figure 2 b).
2) To make sensory predictions, location representations π {\bm{g}} alone are not enough; they must each link to a sensory observation π {\bm{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).