Relating transformers to models and neural representations of the hippocampal formation

TL;DR

The paper demonstrates that Transformers equipped with recurrent position encodings can replicate spatial representations observed in the hippocampal-entorhinal system, notably grid and place cells. It establishes a close mathematical relationship between Transformers and the Tolman-Eichenbaum Machine (TEM), mapping TEM memory retrieval to self-attention and interpreting path integration as learned position encodings. Introducing TEM-t, a Transformer analogue of TEM, the authors report improved sample efficiency and memory capacity along with emergent grid-like representations, supported by a biologically plausible account of place-cell remapping. The work bridges neuroscience and machine learning, offering insights into hippocampal indexing, the role of position encodings, and potential extensions to language and broader cognitive tasks. Overall, it presents a unified view where hippocampal computations can be viewed through the lens of transformer-like architectures, with implications for understanding cortical processing beyond the hippocampus.

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.

Figures (11)

• 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 Hafting2005. (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 Krupic2012.
• Figure 2: (a) The TEM model, with a path integration component (equation \ref{['eq:path_int']}) and a memory network component (equation \ref{['eq:tem_mem_storage']} and \ref{['eq:attractor']}). TEM path integrates ${\bm{g}}$ and makes sensory predictions ${\bm{x}}$ via its memory network (dashed lines are additional connections for inference). (b) TEM recapitulates a host of empirically described cell representations Whittington2020. Top/bottom row: example TEM MEC/Hippocampal representations (plots are spatially smoothed). Figures adapted from Whittington2020. (c) Schematic of TEM (adapted from Sanders2020), showing that the same cortical representations (LEC and MEC) are reused in different environments allowing for generalisation, facilitated by different hippocampal combinations. (d) The TEM hippocampal conjunction is an outer product - cells receive input from particular MEC and LEC cells.
• Figure 3: Self-attention in (a) Transformers and (b) TEM.
• Figure 4: TEM-t is a more efficient learner than TEM, both in (a) sample efficiency and (b) time per gradient step. Zero-shot accuracy is prediction accuracy when taking links it has never taken before, but to a state it has visited before. Successful accuracy here is only possible with learned and generalised spatial knowledge. We have used the code from TEM from the TEM authors original code https://github.com/djcrw/generalising-structural-knowledge, and so have not optimised it for speed of learning etc, so we cannot claim this to be a fair comparison, nevertheless the difference is stark. We note that in the TEM paper, the authors say it takes up to 50,000 gradient updates for full training, whereas we stopped at 20,000.
• Figure 5: TEM-Transformer neural architecture. (a)Krotov2020 describe a neurally plausible architectural instantiation the 'Hopfield networks is all you need' with a separation between 'feature' neurons (i.e. ${\bm{h}}$) and memory neurons (i.e. softmax(${\bm{q}}_t {\bm{K}}^T$). (b-c) This can be extended for TEM-t, but now the feature neurons are not all updated simultaneously, but only those across brain regions. (d) Memory neurons resemble hippocampal place cells and (e) remap randomly across environments. (f) A possible architecture where cortical neurons project to feature neurons in hippocampus which in turn project to memory neurons in hippocampus. (g) Additional brain regions can be included easily in this architecture with minimal increase in hippocampal neuron number.