The Tolman-Eichenbaum Machine: Unifying Space and Relational Memory through Generalization in the Hippocampal Formation

Abstract

The hippocampal-entorhinal system is important for spatial and relational memory tasks. We formally link these domains, provide a mechanistic understanding of the hippocampal role in generalization, and offer unifying principles underlying many entorhinal and hippocampal cell types. We propose medial entorhinal cells form a basis describing structural knowledge, and hippocampal cells link this basis with sensory representations. Adopting these principles, we introduce the Tolman-Eichenbaum machine (TEM). After learning, TEM entorhinal cells display diverse properties resembling apparently bespoke spatial responses, such as grid, band, border, and object-vector cells. TEM hippocampal cells include place and landmark cells that remap between environments. Crucially, TEM also aligns with empirically recorded representations in complex non-spatial tasks. TEM also generates predictions that hippocampal remapping is not random as previously believed; rather, structural knowledge is preserved across environments. We confirm this structural transfer over remapping in simultaneously recorded place and grid cells.

Keywords: entorhinal cortex; generalization; grid cells; hippocampus; neural networks; non-spatial reasoning; place cells; representation learning.

Graphical abstract

Figure 1

Spatial and Relational Inferences Cast as Structural Generalization (A–C) Structured relationships exist in many situations and can often be formalized on a connected graph, e.g., (A) social hierarchies, (B) transitive inference, and (C) spatial reasoning. Often the same relationships generalize across different sets of sensory objects (e.g., left/right in A). This transferable structure allows quick inference, e.g., seeing only the blue relationships allows you to infer the green ones. (D) Our task is predicting the next sensory observation in sequences derived from probabilistic transitions on a graph. Each node has an arbitrary sensory experience, e.g., a banana. An agent transitions on the graph observing only the immediate sensory stimuli and associated action taken, e.g., having seen motorbike $\to$ book $\to$ table $\to$ chair, it should predict the motorbike next if it understands the rules of the graph. (E) If you know the underlying structure of social hierarchies, observing a new node (in red) via a single relationship, e.g., Emily is Bob’s daughter, allows immediate inference about the new node’s (Emily’s) relationship to all other nodes (shown in black/gray). (F) Similarly for spatial graphs observing a new node on the left (solid red line) also tells us whether it is above or below (dashed red lines) other surrounding nodes. (G) Our agent performs this next step prediction task in many worlds sharing the same underlying structure (e.g., 6- or 4-connected graphs), but differing in size and arrangement of sensory stimuli. The aim is to learn the common structure in order to generalize and perform quick inferences. (H) Knowing the structure allows full graph understanding after only visiting all nodes, not all edges. Here, only 18 steps (red line) are required to infer all 42 links. (I) An agent that knows structure (node agent) will reach peak predictive performance after it has visited all nodes, quicker than one that has to see all transitions (edge agent). Icons from https://www.flaticon.com. See also Figure S1.

Figure S1

Task Schematics for TEM, Related to Figures 1 and 3 (A) Learning to predict the next sensory observation in environments that share the same structure but differ in their sensory observations. TEM only sees the sensory observations and associated action taken, it is not told about the underlying structure - this must be learned. (B) Transitive inference graph. When a new node (red) is seen to be one higher, all other (dotted) relations can be inferred i.e., 3 higher. (C) Example graph for a social hierarchy. (D) Example graph for 2D structure. (E) A complex task embedded in a spatial world. This is a schematic representation of the state space for the task in Sun et al. (2020). Each lap is of length 4 as the sensory objects (A, B, C, D) repeat every 4 nodes. There are 3 laps in total, and that defines the true state-space as a reward, r, is given every 3 laps.

Figure 2

The Tolman-Eichenbaum Machine (A) Factorization and conjunction as a principle for generalization. Separating structural codes (the transition rules of the graph) from sensory codes allows generalization over environments sharing the same structure. The conjunctive code represents the current environment in the context of this learned structure. (B and C) The two key elements of TEM. (B) Representations for path integration $\left(\mathbf{g}\right)$ on arbitrary graphs and C) relational memories $\left(\mathbf{p}\right)$ that bind abstract locations to sensory observations. (B) TEM must learn structural codes $\left(\mathbf{g}\right)$ that (1) represent each state differently so that different memories can be stored and retrieved and (2) have the same code on returning to a state (from any direction) so the appropriate memory can be retrieved. (C) Relational memories conjunctively combine the factorized structural (in blue representing location C) and sensory (in red representing the person) codes, thus these memories know what was where. The memories are stored in Hebbian weights $\left(\mathbf{M}\right)$ between the neurons of $\mathbf{p}$. (D) Depiction of TEM at two time points, with each time point described at a different level of detail. Red shows predictions; green shows inference. Time point t shows network implementation and $t+1$ describes each computation in words. Circles depict neurons (blue is $\mathbf{g}$, red is $\mathbf{x}$, blue/red is $\mathbf{p}$); shaded boxes depict computation steps; arrows show learnable weights; looped arrows describe recurrent attractor. Black lines between neurons in the attractor describe Hebbian weights $\mathbf{M}$. Yellow arrows show errors that are minimized during training. Overall, TEM transitions through latent variables $\mathbf{g}$ and stores and retrieves memories $\mathbf{p}$ using Hebbian weights $\mathbf{M}$. We note that this is a didactic schematic; for completeness and a faithful interpretation of the Bayesian underpinnings, please see STAR Methods and Figures S2, S3, and S4.

Figure S2

Full TEM Model, Related to Figure 2 (A) Generative model. (B) Inference model. Circled/ boxed variables are stochastic/ deterministic. Dashed arrows/boxes are optional as explained in the text. (C) Schematic to show the model flow in the neural network. Depiction of TEM at three time-points, with each time-point described at a different level of detail. Green/ red show inference and generative networks. Time point $t-1$ shows the overall Bayesian logic, t shows network implementation, $t+1$ describes each computation in words. Circles depict neurons (blue is $\mathbf{g}$, red is $\mathbf{x}$, blue/red is $\mathbf{p}$); shaded boxes depict computation steps; arrows show learnable weights (green and red are weights in inference and generative networks); looped arrows describe recurrent attractor. Black lines between neurons in attractor describe Hebbian weights $\mathbf{M}$. ${\mathbf{W}}_a$ are learnable, action dependent, transition weights. ${\mathbf{W}}_{repeat}$ and ${\mathbf{W}}_{tile}$ are fixed weights that make the dimensions of the structural (blue) and sensory (red) inputs, respectively, to the attractor the same. Yellow arrows show training errors. We note we do not show the temporal filtering of sensory data $\mathbf{x}$ in this schematic.

Figure S3

Computations in TEM Generative Model, Related to Figure 2 This shows each computation, as described in Generative architecture, making clear the fixed ${\mathbf{W}}_{tile}$ and ${\mathbf{W}}_{repeat}$ matrices perform appropriate dimension changes, though we note that the matrices may not be the sole computation in each step. Attractor dynamics are described in Retrieval using an attractor network. Red/Blue boxes describe two different ‘streams’.

Figure S4

Computations in TEM Inference Model, Related to Figure 2 This shows each computation, as described in Inference architecture, making clear the fixed ${\mathbf{W}}_{tile}$ and ${\mathbf{W}}_{repeat}$ matrices perform appropriate dimension changes ensuring the entorhinal and sensory input to hippocampus have the same dimension. We note that the matrices may not be the sole computation in each step. Red/Blue boxes describe two different ‘streams’. In the bottom left, we show that multiplying together the representation after the ${\mathbf{W}}_{tile}$ and ${\mathbf{W}}_{repeat}$ operations is equivalent to an outer product.

Figure 3

TEM Learns and Generalizes Abstract Relational Knowledge (A–C) Learning to learn: when TEM has only seen a few environments (blue/green) it takes many visits to each node to remember it. This is because it (1) does not yet understand the structure of the graph and (2) has not learned how to use memories. After visiting more environments and learning the common structure (cyan/yellow), TEM correctly predicts a node on the second visit regardless of the edge taken—TEM now understands both the rules of the graph (path integration) and how to store and retrieve memories. (A) Transitive inference, (B) social hierarchies, and (C) 2D graphs. (D–F) On 2D graphs. (D) TEM is able to predict sensory observations when returning to a node for the first time via a new direction—this is only possible with learned structural knowledge. (E) TEM can store long-term memories. (F) TEM’s performance tracks nodes visited, not edges. These results all demonstrate that TEM has learned and generalized abstract structural knowledge. See also Figure S1.

Figure 4

TEM Structural Neurons $\mathbf{g}$ Learn to Be Grid Cells that Generalize and TEM Conjunctive Memory Neurons $\mathbf{p}$ Learn to Be Place Cells that Remap We use 2D graphs with the number of nodes sampled from {61, 91, 127} or {64, 81, 100, 121} for hexagonal or square environments respectively. A cell’s rate map is obtained by allowing the agent to explore the environment then calculating its average firing rate at each point (graph node) in the environment. (A and B) TEM learned structural representations for random walks on 2D graphs. (A) Hexagonal worlds. Left to right: environments 1, 2, autocorrelation, real data (Krupic et al., 2012; Stensola et al., 2012), top to bottom: different cells. TEM learns grid-like cells, of different frequencies (top versus middle), and of different phases (middle versus bottom). (B) Square worlds. Two TEM learned structural cells—left/right; rate map/autocorrelation. (C) Raw unsmoothed rate maps. Left/right: bottom two cells from (A) both cells from (B). (D) TEM also learn band-like cells. Importantly, all TEM structural representations (A)–(D) generalize across environments. (E) Learned memory representations resemble place cells (left/right: environments 1/2; top 2 simulated, bottom 2 real cells) and have different field sizes. These cells remap between environments, i.e., do not generalize. (F) Grid scores of TEM grid-like cells correlate across environments. (G and H) To examine whether relationships between cells are preserved between environments, we correlated the spatial correlation coefficients of pairs of grid or place fields from each environment, using data from TEM or Barry et al. (2012) and Chen et al. (2018). (G) The spatial correlation coefficients of pairs of TEM structural cells and real data grid cells correlate strongly. (H) TEM hippocampal and real data place cells preserved their relationship to a lesser extent. This suggests that TEM structural cells, along with real grid cells, encode generalizable relationships to a greater extent than TEM hippocampal and real place cells. See also Figure S5.

Figure S5

Further TEM Cell Representations, Related to Figures 4 and 7 (A/B) Raw/smoothed structural cells, $\mathbf{g}$, learned by TEM in diffusive behavior. C/D) Raw/smoothed learned TEM entorhinal cells, $\mathbf{g}$, when trained on a square graph environment. D) E) Hippocampal cells, $\mathbf{p}$, learned by TEM during diffusive behavior. (F) Random sample of TEM hippocampal cells when trained on 4-lap task of Sun et al. (2020). (G) Random sample of TEM entorhinal cells when trained on 4-lap task of Sun et al. (2020).

Figure 5

TEM Learned Representations Reflect Transition Statistics When the agent’s transition statistics mimic different behaviors, TEM learns new representations (left to right: different cells; top to bottom: environments 1, 2, real data). (A) When biased to move toward objects (white dots) TEM learns structural cells with a vector relationship to the objects—object vector cells (Høydal et al., 2019). These cells generalize to all objects. (B) TEM hippocampal cells reflect this behavioral transition change with similar cells, though they do not generalize to all objects—landmark cells (Deshmukh and Knierim, 2013). (C) When biased toward boundaries, TEM learns border cell-like representations (Solstad et al., 2008).

Figure 6

Structural Knowledge Is Preserved over Apparently Random Hippocampal Remapping (A) TEM predicts place cells remap to locations consistent with a grid code, i.e., a place cell co-active with a grid cell will be more likely to remap to locations where that grid cell is also active. (B and C) Data from open-field remapping experiments with simultaneously recorded place and grid cells (Barry et al., 2012; Chen et al., 2018). We compute the grid cell firing rate at the location of place cell peak for every grid cell, place cell pair in each of the two environments and then correlate this measure across environments (left). We compare this correlation coefficient to those computed equivalently but with randomly permuted place cell peaks (right). This is done for two independent datasets (B) (Barry et al., 2012) and (C) (Chen et al., 2018). The true observed correlation coefficients lies off the null distribution (p < 0.05), demonstrating place cell remapping is not random but instead tied to structural constraints of grid cells. (D and E) The same analysis on TEM learned representations shows qualitatively similar results to Barry et al. (2012) (D) and Chen et al. (2018) (E). See also Figures S6 and S7 and Tables S3 and S4.

Figure S6

Fitting Ideal Grid Maps and Analysis of Real Data Showing Grid Cells Realign and Place Cells Remap, Related to Figure 6 (A-C) Ideal grid. We fit an idealized grid rate map using the formula from Stemmler et al. (2015) to the original grid cell rate maps to remove any possible confounds and to ensure that we obtain accurate grid cell peaks. (A) An example original grid cell rate map. (B) An idealized rate map fit to that in (A). (C) Accurate finding of grid cell peaks (white crosses) on the idealized grid rate map, which also allows peaks that extend outside the box to be used (red crosses). D-E) Grid realignment and place cell remapping across environments in dataset 1. (D) Histograms showing the distributions of spatial correlations for place and grid cells both within and across environments. (E) Bar plots showing the mean ($\pm$ SEM) of these distributions. F-G) Grid realignment and place cell remapping across environments in dataset 2. (F) and (G) are same analyses as (D) and (E) but with dataset 2. They demonstrate distributions of spatial correlations near 0 for dataset 2. (G) has its axis locked to that of (E) for visualization.

Figure 7

TEM Represents Non-spatial Reinforcement Learning Tasks and Predicts Non-spatial Remapping (A) In Sun et al. (2020), rodents perform laps of a track, only “rewarded” every 4 laps. Different hippocampal cell types are found: spatial place-like cells (top), those that preferentially fire on a given lap (middle), and those that count laps (bottom). (B) TEM learns similar representations when only “rewarded” every 4 laps. (C) TEM medial entorhinal cells learn both spatially periodic cells (top) and cells that represent the non-spatial task structure of “every 4 laps” (bottom). The latter cells are yet experimentally observed but are predicted by TEM. (D and E) TEM offers a mechanistic understanding of remapping in both spatial and non-spatial tasks. (D) Top/middle/bottom: schematic of entorhinal/hippocampal/sensory cells. Left/right: environment 1/2. Only two laps shown for clarity. TEM says spatial hippocampal cells are active when they receive input from both sensory (LEC; cell codes for A in this example) and MEC input. Place cells, thus, can only remap to other peaks (or within a broad MEC cell field) provided it also receives sensory input there. (E) TEM says, however, cells will retain their lap specificity despite spatially remapping (i.e., a lap 2 cell stays a lap 2 cell); since sensory observations repeat each lap, lap specificity is driven by MEC input. (F and G) Analysis to show TEM (F) and real (Sun et al., 2020) (G) hippocampal cells retain lap specificity after remapping. Left: distribution of lap-specificity correlations are significantly higher than shuffles. Top-right: distribution of spatial correlations (spatial) after remapping is compared to the distribution of lap-specific correlations (ESR). Bottom right: for cells of high lap-specific correlation after remapping (defined by the blue box in the left panel). (H) TEM prescribes which hippocampal cells will be active in each environment. The proportion of active cells that are ESR cells in environment 1, 2, or both (brown, purple, gray) imply approximate independence of cells recruited. (I) Data from Sun et al.(2020) showing the same effect1. Icons are from https://www.flaticon.com. See also Figure S5. .

Figure S7

Schematic of Analysis Showing Preserved Grid-Place Relationships after Remapping, with Corresponding Results, Related to Figure 6 A) Schematic explaining the gridAtPlace analysis. Specifically how the scatterplot is generated. Note that in this figure original grid cell rate maps are shown, rather than ideal grid cell rate maps (Figures S6A–S6C) that were used to generate the main text figures. B-C) The grid cell correlation structure is preserved across environments in dataset 1. B) Dataset 1. Scatterplot shows the correlation across environments of the spatial correlations of grid cell-grid cell pairs (i.e., the correlation of the upper triangle of two grid cell by grid cell correlation matrices: one from environment 1 and one from environment 2). The histogram shows this correlation coefficient was significant relative to a null distribution of correlation coefficients obtained by permuting grid cell-grid cell pairs. (C) Same as A for place cells. D-E) Replication of preserved grid cell correlation structure across environments in dataset 2. D and E are the same format as (B) and (C). F-G) Preserved relationship between place and grid cells across environments in dataset 1. The scatterplots show the correlation of a given measure across trials, where each point is a place cell-grid cell pair. The histogram plots show where this correlation (gray line) lies relative to the null distribution of correlation coefficients. The p value is the proportion of the null distribution that is greater than the unshuffled correlation. (F) gridAtPlace (top) and minDist (bottom) measures are strongly significantly correlated over two trials within the same environment, as expected given the same place and grid code should be present. (G) These measures are also significantly correlated across the two different environments, providing evidence that place and grid cells retain their relationship across environments. (H) Replication of the preserved relationship between place and grid cells across environments in dataset 2. The gridAtPlace measure is significantly correlated at p$<$0.05 across real and virtual worlds and the minDist measure is trending very close to significance, replicating the preserved relationship between grid and place cells across environments.