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

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

With multiple inputs to hippocampus [ 𝒙 ~ , 𝒈 ~ , 𝒄 ~ , … ] [{\bm{\tilde{x}}},{\bm{\tilde{g}}},{\bm{\tilde{c}}},\dots] , any subset of those brain areas can reinstate a memory in the other brain regions i.e. 𝒙 ~ {\bm{\tilde{x}}} and 𝒈 ~ {\bm{\tilde{g}}} can reinstate a 𝒄 ~ {\bm{\tilde{c}}} memory or 𝒈 ~ {\bm{\tilde{g}}} alone could reinstate 𝒙 ~ {\bm{\tilde{x}}} and 𝒄 ~ {\bm{\tilde{c}}} memories. As an analogy to the TEM triple conjunction, TEM-t proposes that 𝒄 ~ t {\bm{\tilde{c}}}_{t} is updated via 𝒄 ~ t ← s ​ o ​ f ​ t ​ m ​ a ​ x ​ ( ( 𝒈 ~ t ​ 𝑮 ~ T ) ⊙ ( 𝒙 ~ t ​ 𝑿 ~ T ) ) ​ 𝑪 ~ {\bm{\tilde{c}}}_{t}\leftarrow softmax(({\bm{\tilde{g}}}_{t}{\bm{\tilde{G}}}^{T})\odot({\bm{\tilde{x}}}_{t}{\bm{\tilde{X}}}^{T})){\bm{\tilde{C}}} , where ⊙ \odot is an element wise product. We note an alternate, and perhaps more intuitive, option could also be 𝒄 ~ t = s ​ o ​ f ​ t ​ m ​ a ​ x ​ ( 𝒈 ~ t ​ 𝑮 ~ T + 𝒙 ~ t ​ 𝑿 ~ T ) ​ 𝑪 ~ {\bm{\tilde{c}}}_{t}=softmax({\bm{\tilde{g}}}_{t}{\bm{\tilde{G}}}^{T}+{\bm{\tilde{x}}}_{t}{\bm{\tilde{X}}}^{T}){\bm{\tilde{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.

## 8 Conclusion

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 mathematical 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.

## Appendix A Appendix

### A.1 The maths using elements

Here we derive the main results using vector and matrix elements. Since each element in 𝒑 {\bm{p}} is multiplicative combination of every pair of elements in 𝒙 ~ {\bm{\tilde{x}}} and 𝒈 ~ {\bm{\tilde{g}}} , then p i ​ j = g ~ i ​ x ~ j {p}_{ij}={{\tilde{g}}}_{i}{{\tilde{x}}}_{j} (12)

Where we label the elements of vector 𝒑 {\bm{p}} with two indices even though it is a vector. Similarly, the memory matrix uses 4 indices; M i ​ j ​ k ​ l = ∑ τ g ~ i τ ​ x ~ j τ ​ g ~ k τ ​ x ~ l τ {M}_{ijkl}=\sum_{\tau}{{\tilde{g}}}^{\tau}_{i}{{\tilde{x}}}^{\tau}_{j}{{\tilde{g}}}^{\tau}_{k}{{\tilde{x}}}^{\tau}_{l} (13)

Retrieving a memory via path integration. Querying the network becomes q i ​ j = g ~ i P ​ I {q}_{ij}={{\tilde{g}}}^{PI}_{i} , which reduces to ∑ i ​ j ​ k q i ​ j ​ M i ​ j ​ k ​ l = ∑ τ ∑ i ​ j ​ k g ~ i P ​ I ​ g ~ i τ ​ x ~ j τ ​ g ~ k τ ​ x ~ l τ = ∑ τ x ~ l τ ​ ( ∑ k g ~ k τ ) ​ ( ∑ i g ~ i P ​ I ​ g ~ i τ ) ​ ( ∑ j x ~ j τ ) → 𝒈 ~ t P ​ I ​ 𝑿 ~ T ​ 𝚲 g ​ 𝚲 x ​ 𝑮 ~ \begin{split}\sum_{ijk}{q}_{ij}{M}_{ijkl}&=\sum_{\tau}\sum_{ijk}{{\tilde{g}}}^{PI}_{i}{{\tilde{g}}}^{\tau}_{i}{{\tilde{x}}}^{\tau}_{j}{{\tilde{g}}}^{\tau}_{k}{{\tilde{x}}}^{\tau}_{l}\\ &=\sum_{\tau}{{\tilde{x}}}^{\tau}_{l}(\sum_{k}{{\tilde{g}}}^{\tau}_{k})(\sum_{i}{{\tilde{g}}}^{PI}_{i}{{\tilde{g}}}^{\tau}_{i})(\sum_{j}{{\tilde{x}}}^{\tau}_{j})\\ &\rightarrow{\bm{\tilde{g}}}^{PI}_{t}{\bm{\tilde{X}}}^{T}{\bm{\Lambda}}_{g}{\bm{\Lambda}}_{x}{\bm{\tilde{G}}}\end{split} (14)

Retrieving a memory using a sensory observation alone. Similarly when the query is the sensory observation alone q i ​ j = x j {q}_{ij}={x}_{j}

∑ i ​ j ​ l q i ​ j ​ M i ​ j ​ k ​ l = ∑ i ​ j ​ l ​ τ x ~ j ​ g ~ i τ ​ x ~ j τ ​ g ~ k τ ​ x ~ l τ = ∑ τ g ~ k τ ​ ( ∑ l x ~ l τ ) ​ ( ∑ i g ~ i τ ) ​ ( ∑ j x ~ j ​ x ~ j τ ) → 𝒙 ~ t ​ 𝑿 ~ T ​ 𝚲 g ​ 𝚲 x ​ 𝑮 ~ \begin{split}\sum_{ijl}{q}_{ij}{M}_{ijkl}&=\sum_{ijl\tau}{{\tilde{x}}}_{j}{{\tilde{g}}}^{\tau}_{i}{{\tilde{x}}}^{\tau}_{j}{{\tilde{g}}}^{\tau}_{k}{{\tilde{x}}}^{\tau}_{l}\\ &=\sum_{\tau}{{\tilde{g}}}^{\tau}_{k}(\sum_{l}{{\tilde{x}}}^{\tau}_{l})(\sum_{i}{{\tilde{g}}}^{\tau}_{i})(\sum_{j}{{\tilde{x}}}_{j}{{\tilde{x}}}^{\tau}_{j})\\ &\rightarrow{\bm{\tilde{x}}}_{t}{\bm{\tilde{X}}}^{T}{\bm{\Lambda}}_{g}{\bm{\Lambda}}_{x}{\bm{\tilde{G}}}\end{split} (15)

Retrieving a memory using a sensory observation and path integration. When the query is q i ​ j = g i P ​ I ​ x j {q}_{ij}={g}^{PI}_{i}{x}_{j} , and we want to retrieve a position encoding memory

∑ i ​ j ​ l q i ​ j ​ M i ​ j ​ k ​ l = ∑ i ​ j ​ l ​ τ g ~ i P ​ I ​ x ~ j ​ g ~ i τ ​ x ~ j τ ​ g ~ k τ ​ x ~ l τ = ∑ τ g ~ k τ ​ ( ∑ l x ~ l τ ) ​ ( ∑ i g ~ i P ​ I ​ g ~ i τ ) ​ ( ∑ j x ~ j ​ x ~ j τ ) → ( 𝒙 ~ t ​ 𝑿 ~ T ⊙ 𝒈 ~ t P ​ I ​ 𝑮 ~ T ) ​ 𝚲 g ​ 𝑮 ~ \begin{split}\sum_{ijl}{q}_{ij}{M}_{ijkl}&=\sum_{ijl\tau}{{\tilde{g}}}^{PI}_{i}{{\tilde{x}}}_{j}{{\tilde{g}}}^{\tau}_{i}{{\tilde{x}}}^{\tau}_{j}{{\tilde{g}}}^{\tau}_{k}{{\tilde{x}}}^{\tau}_{l}\\ &=\sum_{\tau}{{\tilde{g}}}^{\tau}_{k}(\sum_{l}{{\tilde{x}}}^{\tau}_{l})(\sum_{i}{{\tilde{g}}}^{PI}_{i}{{\tilde{g}}}^{\tau}_{i})(\sum_{j}{{\tilde{x}}}_{j}{{\tilde{x}}}^{\tau}_{j})\\ &\rightarrow({\bm{\tilde{x}}}_{t}{\bm{\tilde{X}}}^{T}\odot{\bm{\tilde{g}}}^{PI}_{t}{\bm{\tilde{G}}}^{T}){\bm{\Lambda}}_{g}{\bm{\tilde{G}}}\end{split} (16)

### A.2 Details of implementation

We will release code on publication.

Scaling beta. Since the number of memories is not constant, we use an adaptive β \beta 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 l ​ o ​ g ​ ( n m ​ e ​ m ​ o ​ r ​ i ​ e ​ s ) log(n_{memories}) .

Losses. We use same set of losses as in TEM; • a one-step sensory prediction cross entropy loss, i.e. using 𝒈 t P ​ I {\bm{g}}^{PI}_{t} as input to the memory retrieval process

• a sensory prediction cross entropy loss using 𝒈 t {\bm{g}}_{t} as input to the memory retrieval process

• a squared error loss between 𝒈 t {\bm{g}}_{t} and 𝒈 t P ​ I {\bm{g}}^{PI}_{t}

• l2 weight regularisation

• l2 regularisation on 𝒈 t {\bm{g}}_{t}

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 𝒈 {\bm{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; 𝒒 t = f ​ l ​ a ​ t ​ t ​ e ​ n ​ ( 𝒙 ~ t T ​ 𝟙 ) {\bm{q}}_{t}=flatten({\bm{\tilde{x}}}_{t}^{T}\mathbbm{1}) , and memories of positional encodings are retrieved as; 𝒈 ~ t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d = s ​ u ​ m ​ ( u ​ n ​ f ​ l ​ a ​ t ​ t ​ e ​ n ​ ( 𝒒 t ​ 𝑴 t ) , 0 ) = ∑ τ t 𝒈 ~ τ ​ g ~ ¯ τ ​ x ~ ¯ τ ​ [ 𝒙 ~ t ⋅ 𝒙 ~ τ ] = 𝒙 ~ t ​ 𝑿 ~ T ​ 𝚲 g ​ 𝚲 x ​ 𝑮 ~ {\bm{\tilde{g}}}^{retrieved}_{t}=sum(unflatten({\bm{q}}_{t}{\bm{M}}_{t}),0)=\sum^{t}_{\tau}{\bm{\tilde{g}}}_{\tau}\bar{{{\tilde{g}}}}_{\tau}\bar{{{\tilde{x}}}}_{\tau}[{\bm{\tilde{x}}}_{t}\cdot{\bm{\tilde{x}}}_{\tau}]={\bm{\tilde{x}}}_{t}{\bm{\tilde{X}}}^{T}{\bm{\Lambda}}_{g}{\bm{\Lambda}}_{x}{\bm{\tilde{G}}} (17)

The final position encoding, 𝒈 {\bm{g}} is computed on a basis of path integration ( 𝒈 t P ​ I {\bm{g}}^{PI}_{t} ; equation 3 ) and stored landmark information ( 𝒈 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d {\bm{g}}^{retrieved}_{t} ; equation 17 ), i.e. 𝒈 t = 𝒈 t P ​ I + f θ ​ ( 𝒈 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d , 𝒈 t P ​ I ) ​ ( 𝒈 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d − 𝒈 t P ​ I ) {\bm{g}}_{t}={\bm{g}}^{PI}_{t}+f_{\theta}({\bm{g}}^{retrieved}_{t},{\bm{g}}^{PI}_{t})({\bm{g}}^{retrieved}_{t}-{\bm{g}}^{PI}_{t}) , where 𝒈 r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d = f θ ​ ( 𝒈 ~ r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d ) {\bm{g}}^{retrieved}=f_{\theta}({\bm{\tilde{g}}}^{retrieved}) and f θ ​ ( ⋯ ) f_{\theta}(\cdots) are different MLPs.

We note that a better query to the memory network would be 𝒒 t = f ​ l ​ a ​ t ​ t ​ e ​ n ​ ( 𝒙 ~ t T ​ 𝒈 t P ​ I ) {\bm{q}}_{t}=flatten({\bm{\tilde{x}}}_{t}^{T}{\bm{g}}^{PI}_{t}) since sensory observations may be aliased, including 𝒈 t P ​ I {\bm{g}}^{PI}_{t} can help disambiguate such aliasing. In this case, the retrieved memory is ( 𝒙 ~ t ​ 𝑿 T ⊙ 𝒈 t P ​ I ​ 𝑮 ) ​ 𝚲 x ​ 𝑮 ({\bm{\tilde{x}}}_{t}{\bm{X}}^{T}\odot{\bm{g}}^{PI}_{t}{\bm{G}}){\bm{\Lambda}}_{x}{\bm{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 𝒈 t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d = f θ ​ ( s ​ o ​ f ​ t ​ m ​ a ​ x ​ ( 𝒙 ~ t ​ 𝑿 T ⊙ 𝒈 t P ​ I ​ 𝑮 ) ​ 𝑮 ) {\bm{g}}^{retrieved}_{t}=f_{\theta}(softmax({\bm{\tilde{x}}}_{t}{\bm{X}}^{T}\odot{\bm{g}}^{PI}_{t}{\bm{G}}){\bm{G}}) , this can be thought of as another attention head i.e. from input [ 𝒈 , 𝒙 ] [{\bm{g}},{\bm{x}}] , the key and query ‘attends’ to 𝒙 {\bm{x}} while the value ‘attends’ to 𝒈 {\bm{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 𝒈 {\bm{g}} and 𝒙 {\bm{x}} (similarity determined by dot product) then no new memory is added 3 3 3 The TEM model actually does something very much like this - memories are only added if they did not already exist - 𝑴 = ∑ τ ( 𝒑 τ − 𝒑 ^ τ ) T ​ ( 𝒑 τ + 𝒑 ^ τ ) {\bm{M}}=\sum_{\tau}({\bm{p}}_{\tau}-\hat{{\bm{p}}}_{\tau})^{T}({\bm{p}}_{\tau}+\hat{{\bm{p}}}_{\tau}) , where 𝒑 ^ τ \hat{{\bm{p}}}_{\tau} is the memory that was retrieved at τ \tau , i.e. only the un-predicted components of the memory 𝒑 τ {\bm{p}}_{\tau} are added. .

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 𝒙 ~ t r ​ e ​ t ​ r ​ i ​ e ​ v ​ e ​ d {\bm{\tilde{x}}}^{retrieved}_{t} can be fed into the next iteration, along with positional encoding i.e. the path integrated 𝒈 {\bm{g}} . In this case (see Appendix for derivation), TEM suggests that 𝒙 ~ t {\bm{\tilde{x}}}_{t} is iteratively updated via 𝒙 ~ t ← s ​ o ​ f ​ t ​ m ​ a ​ x ​ ( 𝒙 ~ t ​ 𝑿 ~ T ⊙ 𝒈 ~ t ​ 𝑮 ~ ) ​ 𝑿 ~ {\bm{\tilde{x}}}_{t}\leftarrow softmax({\bm{\tilde{x}}}_{t}{\bm{\tilde{X}}}^{T}\odot{\bm{\tilde{g}}}_{t}{\bm{\tilde{G}}}){\bm{\tilde{X}}} , but with the first iteration via 𝒙 ~ t ← s ​ o ​ f ​ t ​ m ​ a ​ x ​ ( 𝒈 ~ t ​ 𝑮 ~ ) ​ 𝑿 ~ {\bm{\tilde{x}}}_{t}\leftarrow softmax({\bm{\tilde{g}}}_{t}{\bm{\tilde{G}}}){\bm{\tilde{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.

### A.3 Additional 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).

Should one passively move on a graph (e.g. Figure 6 ), where each node is associated with a non-unique 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 r ​ i ​ g ​ h ​ t → d ​ o ​ w ​ n → l ​ e ​ f ​ t → u ​ p right\to down\to left\to 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.