Relating transformers to models and neural representations of the hippocampal formation — part 2/4

Cleaned-text code points [15846, 31289). Bibliography omitted here; source.md includes it.

The linking is done by binding every element of 𝒈 {\bm{g}} with every element of 𝒙 {\bm{x}} , in other words an outer product that is flattened back into a vector; 𝒑 = f ​ l ​ a ​ t ​ t ​ e ​ n ​ ( 𝒙 T ​ 𝒈 ) {\bm{p}}=flatten({\bm{x}}^{T}{\bm{g}}) (4) These conjunctive 𝒑 {\bm{p}} representations are stored in ‘fast weights’ 1 1 1 We note such ‘fast weights’ have previously been thought of as an alternative to the LSTM ( Ba et al., 2016a ) . via Hebbian learning; 𝑴 t = ∑ τ = 1 t − 1 𝒑 τ T ​ 𝒑 τ {\bm{M}}_{t}=\sum^{t-1}_{\tau=1}{\bm{p}}_{\tau}^{T}{\bm{p}}_{\tau} (5) And they can later be retrieved using an attractor network (a continuous version of the Hopfield network). Here a query vector 𝒒 {\bm{q}} (details next paragraph) is inputted into the network and updated via; 𝒒 ← σ ⁡ ( 𝒒 ​ 𝑴 t ) {\bm{q}}\leftarrow\sigma({\bm{q}}{\bm{M}}_{t}) (6) where σ ⁡ ( ⋯ ) \sigma(\cdots) is a non-linear activation function; a ReLu in TEM. Crucially, because the memories are formed using both 𝒈 {\bm{g}} and 𝒙 {\bm{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 𝒑 {\bm{p}} , TEM imagines (path-integrates) the next location 𝒈 {\bm{g}} and provides this as input to the attractor network in the form 𝒒 = f ​ l ​ a ​ t ​ t ​ e ​ n ​ ( 𝟙 T ​ 𝒈 ) {\bm{q}}=flatten(\mathbbm{1}^{T}{\bm{g}}) . Equation 6 is then iterated until a memory is retrieved.

Finally, to make sensory predictions, the retrieved conjunctive memory ( 𝒑 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d {\bm{p}}^{retrieved}_{t} ) is ‘deconjunctified’ into sensory and location components. The sensory component is obtained by unflattening 𝒑 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d {\bm{p}}^{retrieved}_{t} and summing over the 𝒈 {\bm{g}} dimension (Figure 8 ); 𝒙 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d = s ​ u ​ m ​ ( u ​ n ​ f ​ l ​ a ​ t ​ t ​ e ​ n ​ ( 𝒑 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d ) , 1 ) {\bm{x}}_{t}^{retrieved}=sum(unflatten({\bm{p}}^{retrieved}_{t}),1) (7) Finally, to make the sensory prediction 𝒙 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d {\bm{x}}^{retrieved}_{t} is fed through a MLP 𝒛 t = f θ ​ ( 𝒙 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d ) {\bm{z}}_{t}=f_{\theta}({\bm{x}}^{retrieved}_{t}) ) to classify (predict) the upcoming sensory observation.

It is also possible, and often helpful, to project 𝒈 {\bm{g}} and 𝒙 {\bm{x}} via 𝑾 g {\bm{W}}_{g} and 𝑾 x {\bm{W}}_{x} ; 𝒈 ~ = 𝒈 ​ 𝑾 g {\bm{\tilde{g}}}={\bm{g}}{\bm{W}}_{g} and 𝒙 ~ = 𝒙 ​ 𝑾 x {\bm{\tilde{x}}}={\bm{x}}{\bm{W}}_{x} before they are combined conjunctively 2 2 2 In fact, in the TEM code online, 𝑾 g {\bm{W}}_{g} and 𝑾 x {\bm{W}}_{x} are set as fixed weight matrices, where 𝑾 g {\bm{W}}_{g} sub-samples 𝒈 {\bm{g}} and 𝑾 x {\bm{W}}_{x} transforms (compresses) 𝒙 {\bm{x}} from a one-hot to a two-hot representation. .

## 5 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, 𝒈 {\bm{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 𝒒 t ​ 𝑴 t = 𝒒 t ​ ∑ τ t 𝒑 τ T ​ 𝒑 τ {\bm{q}}_{t}{\bm{M}}_{t}={\bm{q}}_{t}\sum^{t}_{\tau}{\bm{p}}_{\tau}^{T}{\bm{p}}_{\tau} is simply equal to 𝒑 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d = ∑ τ t [ 𝒒 t ​ 𝒑 τ T ] ​ 𝒑 τ {\bm{p}}^{retrieved}_{t}=\sum^{t}_{\tau}[{\bm{q}}_{t}{\bm{p}}_{\tau}^{T}]{\bm{p}}_{\tau} (8) Since [ 𝒒 t ​ 𝒑 τ T ] [{\bm{q}}_{t}{\bm{p}}_{\tau}^{T}] is just a dot-product ( [ 𝒒 t ⋅ 𝒑 τ ] [{\bm{q}}_{t}\cdot{\bm{p}}_{\tau}] ), 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 𝒈 ~ {\bm{\tilde{g}}} and 𝒙 ~ {\bm{\tilde{x}}} ( 𝒑 τ = f ​ l ​ a ​ t ​ t ​ e ​ n ​ ( 𝒙 ~ τ T ​ 𝒈 ~ τ ) {\bm{p}}_{\tau}=flatten({\bm{\tilde{x}}}_{\tau}^{T}{\bm{\tilde{g}}}_{\tau}) ), and these memories can be retrieved by a query based on 𝒈 ~ {\bm{\tilde{g}}} or 𝒙 ~ {\bm{\tilde{x}}} alone (e.g. 𝒒 t = f ​ l ​ a ​ t ​ t ​ e ​ n ​ ( 𝟙 T ​ 𝒈 ~ t ) {\bm{q}}_{t}=flatten(\mathbbm{1}^{T}{\bm{\tilde{g}}}_{t}) ). Together, these properties mean we can reduce the above dot product even further; [ 𝒒 t ​ 𝒑 τ T ] = x ~ ¯ τ ​ [ 𝒈 ~ t ⋅ 𝒈 ~ τ ] → 𝒑 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d = 𝒈 ~ t ​ 𝑮 ~ T ​ 𝚲 x ​ 𝑷 [{\bm{q}}_{t}{\bm{p}}_{\tau}^{T}]=\bar{{{\tilde{x}}}}_{\tau}[{\bm{\tilde{g}}}_{t}\cdot{\bm{\tilde{g}}}_{\tau}]\qquad\rightarrow\qquad{\bm{p}}^{retrieved}_{t}={\bm{\tilde{g}}}_{t}{\bm{\tilde{G}}}^{T}{\bm{\Lambda}}_{x}{\bm{P}} (9)

Where x ~ ¯ = ∑ i ( 𝒙 ~ τ ) i \bar{{{\tilde{x}}}}=\sum_{i}({\bm{\tilde{x}}}_{\tau})_{i} and 𝚲 x {\bm{\Lambda}}_{x} is a diagonal matrix with elements x ~ ¯ τ \bar{{{\tilde{x}}}}_{\tau} (see Appendix for an alternative derivation using vector elements). Thus to retrieve a conjunctive 𝒑 {\bm{p}} memory, all that was necessary is weighting past 𝒑 {\bm{p}} representations via ‘self-attention’ of 𝒈 ~ t {\bm{\tilde{g}}}_{t} to past representations 𝑮 ~ {\bm{\tilde{G}}} .

To simplify this even further, we consider what happens when we ‘deconjunctify’ 𝒑 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d {\bm{p}}^{retrieved}_{t} to obtain the sensory component of the memory . Following the TEM procedure described above (Figure 8 ); 𝒙 ~ t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d = s ​ u ​ m ​ ( u ​ n ​ f ​ l ​ a ​ t ​ t ​ e ​ n ​ ( 𝒑 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d ) , 1 ) = ∑ τ t 𝒙 ~ τ ​ g ~ ¯ τ ​ x ~ ¯ τ ​ [ 𝒈 ~ t ⋅ 𝒈 ~ τ ] = 𝒈 ~ t ​ 𝑮 ~ T ​ 𝚲 g ​ 𝚲 x ​ 𝑿 ~ {\bm{\tilde{x}}}^{retrieved}_{t}=sum(unflatten({\bm{p}}^{retrieved}_{t}),1)=\sum^{t}_{\tau}{\bm{\tilde{x}}}_{\tau}\bar{{{\tilde{g}}}}_{\tau}\bar{{{\tilde{x}}}}_{\tau}[{\bm{\tilde{g}}}_{t}\cdot{\bm{\tilde{g}}}_{\tau}]={\bm{\tilde{g}}}_{t}{\bm{\tilde{G}}}^{T}{\bm{\Lambda}}_{g}{\bm{\Lambda}}_{x}{\bm{\tilde{X}}} (10)

Where 𝚲 g {\bm{\Lambda}}_{g} is a diagonal matrix with elements g ~ ¯ τ = ∑ i ( 𝒈 ~ τ ) i \bar{{{\tilde{g}}}}_{\tau}=\sum_{i}({\bm{\tilde{g}}}_{\tau})_{i} . Now all that is necessary to retrieve the sensory component of the memory is weighting past 𝒙 ~ {\bm{\tilde{x}}} representations with via ‘self-attention’ of 𝒈 ~ t {\bm{\tilde{g}}}_{t} to past representations 𝑮 ~ {\bm{\tilde{G}}} . This equation is now very similar to equation 1 except without the softmax and with additional weightings 𝚲 x {\bm{\Lambda}}_{x} and 𝚲 g {\bm{\Lambda}}_{g} . These weighting however are likely learned to be constant ( α \alpha ) because otherwise some memories will be preferentially retrieved. In this case TEM is retrieving memories using 𝒙 ~ t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d = ( α ​ 𝒈 ~ t ​ 𝑮 ~ T ) ​ 𝑿 ~ c ​ f . s ​ o ​ f ​ t ​ m ​ a ​ x ​ ( 𝒈 ~ t ​ 𝑮 ~ T d k ) ​ 𝑿 ~ {\bm{\tilde{x}}}^{retrieved}_{t}=(\alpha{\bm{\tilde{g}}}_{t}{\bm{\tilde{G}}}^{T}){\bm{\tilde{X}}}\qquad cf.\qquad softmax(\frac{{\bm{\tilde{g}}}_{t}{\bm{\tilde{G}}}^{T}}{\sqrt{d_{k}}}){\bm{\tilde{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 𝒈 {\bm{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.

## 6 Place 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 & Hopfield, 2020 ) . In particular, self-attention (equation 1 ) can be split into two steps which correspond to two pools of neurons (Figure 5 A); 1) calculate s ​ o ​ f ​ t ​ m ​ a ​ x ​ ( 𝒒 t ​ 𝑲 T d k ) softmax(\frac{{\bm{q}}_{t}{\bm{K}}^{T}}{\sqrt{d_{k}}}) . 2) multiply by 𝑽 {\bm{V}} . In this light, 𝑲 {\bm{K}} and 𝑽 {\bm{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 5 D-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 𝒈 {\bm{g}} and 𝒙 {\bm{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 {\bm{\tilde{g}}}={\bm{g}}{\bm{W}}_{g} and 𝒙 ~ = 𝒙 ​ 𝑾 x {\bm{\tilde{x}}}={\bm{x}}{\bm{W}}_{x} , but which can connect to the same memory neurons in hippocampus (Figure 5 B-C). These feature sub-populations are updated alternately rather than simultaneously, depending on the direction of retrieval; for example, when retrieving 𝒙 ~ {\bm{\tilde{x}}} the 𝒈 ~ {\bm{\tilde{g}}} feature neurons stay constant while the 𝒙 ~ {\bm{\tilde{x}}} neurons are updated (in turn updating 𝒙 {\bm{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 𝒑 {\bm{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 5 F) then there will be hippocampal spatial cells 𝒈 ~ {\bm{\tilde{g}}} that maintain their relationship to grid cells across different environment (as they are inherited from 𝒈 {\bm{g}} via a projection 𝒈 ~ = 𝒈 ​ 𝑾 g {\bm{\tilde{g}}}={\bm{g}}{\bm{W}}_{g} ). Thus across the population of hippocampal cells, there will be those that maintain their relationship to grid cells (e.g. 𝒈 ~ {\bm{\tilde{g}}} and those that don’t (e.g. memory neurons and 𝒙 ~ {\bm{\tilde{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, 𝒈 {\bm{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 2^{\frac{N}{2}} , where N 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.

## 7 Discussion

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 ( 𝒈 {\bm{g}} and 𝒙 {\bm{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 n_{c} ; the number of neurons from brain region 𝒄 ~ {\bm{\tilde{c}}} . TEM-t does not scale so badly. Instead it just requires an additional n c 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 5 G).