SlowGuess/NeurIPS_2023_OCR huggingface dataset:
locations, similar to the approach described in the TEM-t model [12]. We recorded the activation values of each neuron at every step during a random walk process. As the agent traversed the environment, which had an  $11 \times 11$  grid structure, we accumulated the activation values for each grid point, resulting in a 2D array known as the rate map or place field (i.e., the spatial distribution of single neuron activation during a random walk, as shown in Fig.4a and Fig.4b).
Fig. 4a and 4b show the rate maps of neurons with place cell scores in the FFN and self-attention layers, respectively. For the self-attention layer, the total number of neurons in the softmax layer was 65 (context length + masked sensory observation) × 8 (number of heads) × 2 (number of layers). The total number of neurons in the FFN layer was set at 2,048 (number of neurons) × 2 (number of layers). As can be seen, our metric accurately represents location specificity. Fig. 4c and 4d show the distribution of place cell scores in the two layers with different values of  $\alpha$ . When the  $\alpha$  value is increased, the place cell score distribution found in the FFN becomes positively shifted (see Fig. 4h rate map examples for  $\alpha = 0, 1.0, \text{and} 10.0$ ), whereas the place cell score distribution in the self-attention layers remains.
Fig. 4e and 4f show a relationship between the average place cell score and the reference memory error for each  $\alpha$ . The average place cell scores in the self-attention layers show no correlation with reference memory errors, while neurons in the FFN exhibit a clear correlation. These results imply that reference memory formation and place cell emergence can be enhanced by NMDAR-like nonlinearity in the FFN.
Fig. 4g compares the place cell representations of our NMDA  $(\alpha = 0,10)$  with the representations in FFNs with the activation functions used in Fig. 3c, indicating that the case of  $\mathrm{NMDA}_{\alpha = 10}$  outperforms other activation functions, in both reference memory formation and place cell representation. Our finding that large  $\alpha$  (the  $[\mathrm{Mg}^{2+}]$  component) enhances reference memory is in line with the biological observation that increasing the  $[\mathrm{Mg}^{2+}]$  in the brain enhances long-term memory formation [6].
We also explored the possibility of extending the hyper-parameter space by contrasting the NMDA-inspired activation function with an alternative activation function based on  $\alpha$ . LeakyReLU was considered for this experiment because of its adjustable negative slope, which was re-defined as  $(\max(0, x) + \alpha \min(0, x))$ . The results reported in Fig. S10 of Appendix A.7 show that LeakyReLU activation, for all tested  $\alpha$  values, produced lower place cell scores than the best-performing NMDA-inspired function (i.e.,  $\mathrm{NMDA}_{\alpha=10}$ ). Our findings suggest that NMDA-like nonlinearity, which is already known to be critical in biological systems, could be used to enhance transformer models.
Building upon these results, our work connects memory consolidation theory from neuroscience to transformer-based processing of short-term and long-term memory. Unlike previous studies that directly provide 2D geometric information to models [11, 10], our model only takes sensory-action pairs as an input, which was sufficient to observe the emergence of place cells in the transformer model. Specifically, we showed that nonlinearity in the transformer's FFNs plays a crucial role in transforming short-term working memory from the self-attention layer into long-term memory, and the place cell emergence. The observed place cells may correspond to sparse interpretable activation in FFN. See Appendix A.10 for further analysis on the sparsity and place cells. These findings provide new insight into the inner workings of transformers and their connection to neural processes in the brain.
The current study is inspired by recent studies that connect neuroscience and AI models. One such seminal work is by Whittington et al. [12], which showed a relationship between the self-attention layer and a recent hippocampal model called the Tolman-Eichenbaum Machine (TEM; Whittington et al. [14]). Our work expands the literature by focusing on FFNs in the transformer and making a connection to the emergence of place cells.
TEM is a neuroscience-based model that reproduces neural representations in the hippocampus and entorhinal cortex. Instead of storing memory in the key matrix  $K$  and value matrix  $V$ , it instead stores memory using a Hebbian weight matrix  $M \in \mathbb{R}^{d_k \times d_k}$ . Every outer product of key and value vector  $\mathbf{k}_i^\top \mathbf{v}_i$  at each step  $i$  is simply stored in  $M$  via the Hebbian update rule.  $M$  is initialized to a zero matrix at the beginning of the task and adds every outer product at each time step:
$$
M = a \sum_ {i = 1} ^ {t} \mathbf {k} _ {i} ^ {\top} \mathbf {v} _ {i} = a K ^ {\top} V, \tag {8}
$$
where  $a$  is a weighting factor. In the memory retrieving phase with the query vector  $\mathbf{q}$ , TEM uses an attractor network:
$$
\mathbf {q} M = a \mathbf {q} K ^ {\top} V. \tag {9}
$$
Whittington et al. [12] found that the memory retrieving process in TEM has a close mathematical structure to Eq. (3) when the softmax function is replaced with a linear function. We note that methods for linearizing softmax in self-attention layers have been studied to address high computational costs when context length  $c$  is very long [27, 28]. The subsequent model of TEM, called TEM-t [12], replaces the attractor network (Eq. (9)) with a self-attention mechanism (Eq. (3)). This study demonstrated that TEM-t learns significantly faster than TEM.
TEM-t and TEM do not have a fixed context length  $c$ ; therefore, these models store all information before step  $t$ , i.e.,  $c = t$ . The computational cost of the self-attention layer in TEM-t is  $O(t^{2})$  and retaining all previous information is too expensive from both a biological and computational standpoint. For TEM, the Hebbian update rule has no quadratic computational cost and can add all previous information in a fixed number of synapses  $d_{k}^{2}$ ; however, the memory capacity of the Hebbian matrix  $M$  is  $O(d_{k})$  and the speed of memory retrieval is substantially slower than the
<sup>3</sup>Due to the computational cost, TEM-t does not store all historical data but selectively chooses which data to store in  $K$  and  $V$ .
self-attention mechanism [29-31]. In contrast to TEM and TEM-t that rely on a single memory system, we investigated two separate memory systems: 1) context-dependent matrices  $K$  and  $V$  in the self-attention layer with a fixed context length  $c$  and 2) context-independent fixed matrices  $U_{1}$  and  $U_{2}$  (in Eq. (4)) in the FFNs.
Our research differs from previous studies in the following ways: 1) We designed a navigation test that assesses working memory and reference memory separately, providing a more comprehensive evaluation of the model's performance. 2) We proposed a new brain-inspired activation function,  $\mathrm{NMDA}_{\alpha}$ , which relates to modern nonlinear activation functions and allows for the analysis of the effect of  $\alpha$  on long-term reference memory formation. 3) We demonstrated that place cell-like neurons emerge in FFNs in conjunction with reference memory formation, which is a novel finding that has not been addressed in the TEM or TEM-t models. 4) TEM and TEM-t focus on working memory errors and do not cover reference memory errors. In contrast, our work evaluates both types of memory errors, providing a more detailed analysis of the model's performance.
While extensive efforts have been directed toward finding the optimal nonlinear activation function for improving modern deep neural network models [21-23], their relationship to neural substrates that mediate nonlinearity in the human brain remains obscure. Furthermore, the role of nonlinearity in intelligent functions remain unclear. Our research attempts to fill this gap by proposing and testing how a biologically inspired nonlinearity functions in a transformer model that was previously related to the hippocampal formation. We examined functions in terms of long-term reference memory formation and place cell representation. This idea was evaluated in a carefully designed 2D grid environment and by implementing an activation function derived from NMDAR-like nonlinearity.
Our research reveals that place cell-like neurons that are critical to spatial navigation can be found in both the self-attention layers and FFNs of the transformer model, similar to the presence of place cells in the CA3 and CA1 regions of the hippocampus [32]. In the hippocampus, the CA3 region is thought to be involved in the initial formation of new memories, specifically in pattern completion [33], while CA1 is thought to be important for the long-term consolidation of memories. As such, we suggest that the CA3 region may serve a similar function to the self-attention layer, while the CA1 region may function similarly to FFN. However, further research is needed to fully understand the similarities and differences between the properties of place cells in the transformer model and those found in the CA3 and CA1 regions of the hippocampus. For more discussions regarding the biological plausibility of this proposal, see Appendix A.8.
Recent machine learning research has tested whether the transformer architecture is analogous to different types of biological memory. It has been suggested that (1) transformer FFN modules resemble associative memory [34], (2) the FFN in a transformer block functioning as a key-value memory [35, 36], (3) activation sparsity in the transformer FFN enhances robustness to noisy input [37, 38], and (4) sparse activity in FFNs increases the percentage of neurons that selectively activate to human interpretable input features [39].
Our data instead suggests the transformer architecture resembles memory consolidation by the animal brain, which refers to the transfer process of a short-term memory into a long-term memory system in neuroscience research [40]. Previous research has revealed that  $\mathrm{Mg}^{2+}$ -gating of NMDA receptors modulates the formation of long-term memories [5, 6]. These observations imply that the nonlinear dynamics of NMDA receptors in hippocampus CA1 are critical for consolidating short-term memory into long-term memory.
To our surprise, our results agree qualitatively with previous NMDAR impairment experiments from neuroscience: 1) selective inhibition of hippocampal CA1 NMDAR inhibition does not disrupt working memory [41] but impairs the long-term memory formation [2], 2) changing NMDAR  $\mathrm{Mg}^{2+}$ -gating (changing  $\alpha$  in this work) enhances or disrupts long-term memory formation [5, 6], 3) NMDAR is required for long-term stabilization of newly forming place fields [9, 3]. These similarities between hippocampal memory consolidation and our results suggest that the transformer is an effective memory consolidation model.
Our research points to exciting future directions. The current study examined what-where memory using a sensory observation task in a static environment. However, our real-world environment
changes constantly and provides new inputs over time. Modern deep learning systems are generally incapable of adapting to a dynamic environment or reordering of sensory inputs. We intend to explore what-where-when memory, called episodic memory, in transformers and other deep models.
This work was supported by the Institute for Basic Science in Korea (IBS-R029-C2, IBS-R001-D2). We sincerely thank Dr. Sungho Hong of the Okinawa Institute of Science and Technology for his thoughtful feedback on computational neuron modeling.
[1] Fei Li and Joe Z Tsien. Memory and the nmda receptors. The New England journal of medicine, 361(3):302, 2009.
[2] Joe Z Tsien, Patricio T Huerta, and Susumu Tonegawa. The essential role of hippocampal cal nmda receptor-dependent synaptic plasticity in spatial memory. Cell, 87(7):1327-1338, 1996.
[3] Clifford Kentros, Eric Hargreaves, Robert D Hawkins, Eric R Kandel, Matthew Shapiro, and Robert V Muller. Abolition of long-term stability of new hippocampal place cell maps by NMDA receptor blockade. Science, 280(5372):2121-2126, 1998.
[4] Tim VP Bliss and Graham L Collingridge. A synaptic model of memory: long-term potentiation in the hippocampus. Nature, 361(6407):31-39, 1993.
[5] Inna Slutsky, Nashat Abumaria, Long-Jun Wu, Chao Huang, Ling Zhang, Bo Li, Xiang Zhao, Arvind Govindarajan, Ming-Gao Zhao, Min Zhuo, et al. Enhancement of learning and memory by elevating brain magnesium. Neuron, 65(2):165-177, 2010.
[6] Tomoyuki Miyashita, Yoshiaki Oda, Junjiro Horiuchi, Jerry CP Yin, Takako Morimoto, and Minoru Saitoe.  $\mathrm{Mg2 + }$  block of drosophila NMDA receptors is required for long-term memory formation and CREB-dependent gene expression. Neuron, 74(5):887-898, 2012.
[7] John O'Keefe and Jonathan Dostrovsky. The hippocampus as a spatial map: Preliminary evidence from unit activity in the freely-moving rat. *Brain research*, 1971.
[8] Torkel Hafting, Marianne Fyhn, Sturla Molden, May-Britt Moser, and Edvard I Moser. Microstructure of a spatial map in the entorhinal cortex. Nature, 436(7052):801-806, 2005.
[9] Thomas J McHugh, Kenneth I Blum, Joe Z Tsien, Susumu Tonegawa, and Matthew A Wilson. Impaired hippocampal representation of space in cal-specific nmdar1 knockout mice. Cell, 87 (7):1339-1349, 1996.
[10] Andrea Banino, Caswell Barry, Benigno Uria, Charles Blundell, Timothy Lillicrap, Piotr Mirowski, Alexander Pritzel, Martin J Chadwick, Thomas Degris, Joseph Modayil, et al. Vector-based navigation using grid-like representations in artificial agents. Nature, 557(7705):429-433, 2018.
[11] Christopher J. Cueva and Xue-Xin Wei. Emergence of grid-like representations by training recurrent neural networks to perform spatial localization. In International Conference on Learning Representations, 2018. URL https://openreview.net/forum?id=B17JT0e0-.
[12] James C. R. Whittington, Joseph Warren, and Tim E.J. Behrens. Relating transformers to models and neural representations of the hippocampal formation. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=B8DVo9B1YE0.
[13] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[31] Dmitry Krotov and John J. Hopfield. Large associative memory problem in neurobiology and machine learning. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=X4y_100X-hX.
[32] Can Dong, Antoine D Madar, and Mark EJ Sheffield. Distinct place cell dynamics in ca1 and ca3 encode experience in new environments. Nature communications, 12(1):2977, 2021.
[33] Kazu Nakazawa, Michael C Quirk, Raymond A Chitwood, Masahiko Watanabe, Mark F Yeckel, Linus D Sun, Akira Kato, Candice A Carr, Daniel Johnston, Matthew A Wilson, et al. Requirement for hippocampal ca3 nmda receptors in associative memory recall. science, 297 (5579):211-218, 2002.
[34] Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in GPT. Advances in Neural Information Processing Systems, 36, 2022.
[35] Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In Empirical Methods in Natural Language Processing (EMNLP), 2021.
[36] Damai Dai, Li Dong, Yaru Hao, Zhifang Sui, Baobao Chang, and Furu Wei. Knowledge neurons in pretrained transformers. arXiv preprint arXiv:2104.08696, 2021.
[37] Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank J Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, et al. The lazy neuron phenomenon: On emergence of activation sparsity in transformers. In The Eleventh International Conference on Learning Representations.
[38] Sebastian Jasczur, Aakanksha Chowdhery, Afroz Mohiuddin, Lukasz Kaiser, Wojciech Gajewski, Henryk Michalewski, and Jonni Kanerva. Sparse is enough in scaling transformers. Advances in Neural Information Processing Systems, 34:9895-9907, 2021.
[39] Nelson Elhage, Tristan Hume, Catherine Olsson, Neel Nanda, Tom Henighan, Scott Johnston, Sheer ElShowk, Nicholas Joseph, Nova DasSarma, Ben Mann, Danny Hernandez, Amanda Askell, Kamal Ndousse, Andy Jones, Dawn Drain, Anna Chen, Yuntao Bai, Deep Ganguli, Liane Lovitt, Zac Hatfield-Dodds, Jackson Kernion, Tom Conerly, Shauna Kravec, Stanislav Fort, Saurav Kadavath, Josh Jacobson, Eli Tran-Johnson, Jared Kaplan, Jack Clark, Tom Brown, Sam McCandlish, Dario Amodei, and Christopher Olah. Softmax linear units. Transformer Circuits Thread, 2022. https://transformer-circuits.pub/2022/solu/index.html.
[40] James L McGaugh. Memory-a century of consolidation. Science, 287(5451):248-251, 2000.
[41] Inah Lee and Raymond P Kesner. Differential contribution of nmda receptors in hippocampal subregions to spatial working memory. Nature neuroscience, 5(2):162-168, 2002.
A.1 Place cell
A.1.1 Place cell overview
Place cells are a distinct class of neurons in the hippocampus that fire selectively in a certain location [S1], and provide a sparse population code for self-location. Place cells play a key role in understanding how the hippocampus represents spatial locations [S2], which is characterized by their place fields, which are spatially stable regions where the cell preferentially fires. As the subject moves through the environment, different place cells fire in response to its position, collectively forming a cognitive map [S3]. This dynamic encoding allows the brain to create and maintain spatial representations, which are essential for navigation and memory [S4].
Relevance to machine learning Accurately identifying place cells and understanding their properties is crucial for machine learning practitioners working on spatial representation, navigation, and memory problems. Recent studies investigated the role of place cells and grid cells in spatial representation in navigation tasks with deep models [S5-S8]. The concept of place cells can provide insights that support the development of novel algorithms and techniques for potential machine learning applications such as robotics, autonomous vehicles, virtual reality, and other areas where spatial understanding is essential.
Grid cells and path integration While our work was mainly on the emergence of place cells, we could also observe grid cells based on the experiment, similar to TEM and TEM-t models. When we removed the recurrent connection in the action embedding layer, grid cells did not emerge, and the model's performance declined, indicating the role of grid cells in place cell formation and spatial memory. An interesting hypothesis in neuroscience is that the interaction between grid cells, path integration, and place cells forms a critical loop [S9, S10] in the neural circuitry underlying spatial memory formation, which is similar to our findings from removing recurrent connections. This finding awaits future investigation on how grid cells and path integration affect place cell emergence and spatial memory formation.
A.1.2 Place cell score calculation
(c)
Figure S1: Schematic diagram for creating auxiliary graph for place cell score calculation (a) Consider the 2D grid environment of the place field as an undirected graph  $\mathcal{G}$ . Colors represent the firing rate of nodes. (b) Convert  $\mathcal{G}$  into a directed graph by comparing the firing rate  $\rho$  between nodes (direction of higher to lower firing rates) (c) Operate Depth First Search algorithm from peak node to construct all connected components \} (yellow). Dashed circles are leaf nodes.
Our place cell score metric is inspired by the metric described in the TEM-t model [S8]. The place cell score in Eq. S1 is a scalar value that quantifies each neuron's firing specificity for a grid location. We record every neuron's activation value at every step during the random walk process. Since the map has an  $11 \times 11$  grid structure, we accumulate the activation value for each grid point as the agent random walks, creating a 2D array. This array is called the rate map or the place field (i.e., the spatial distribution of single neuron activation during a random walk).
We investigated the role of neurons in the FFNs and self-attention layers by measuring their place specificity. Given a  $K \times K$  2D grid environment as graph  $G = (V, E)$  and a firing rate (i.e., cumulative activation value at node  $i$  divided by the length of evaluation trial) of node  $i \in V$  as a  $\rho_{i}$ , we defined a maximally firing node as  $i_{\max}$  and its firing rate as  $\rho_{\max}$ .  $E$  are directed edges that connect high to low firing nodes in  $G$ . From  $G$ , we ran a depth-first-search from source node  $i_{\max}$  to build a sub-graph  $\mathcal{G} = (\mathcal{V}, \mathcal{E})$  which we call all connected components S1. Given  $G$  and  $\mathcal{G}$ , the place cell score is defined as follows:
$$
\text {P l a c e c e l l s c o r e} = \gamma \frac {\sum_ {i \in V} \rho_ {i}}{\sum_ {i \in V} \rho_ {i}}, \tag {S1}
$$
where  $\gamma = 1 - |\mathcal{V}^{*}| / |V|$  is a discount factor and  $\mathcal{V}^*$  is  $\mathcal{V}$  without node  $i_{\mathrm{max}}$  and leaf nodes. To measure place cell score, we recorded the firing rate  $\rho_{i}$  of neurons over a random walking trajectory with  $10^{5}$  steps in one of the training maps. We then measured the place cell scores of neurons in the FFN and self-attention layers. The place cell score is 1 when the neuron is firing only at a specific node; the score is 0 when it fires homogeneously across all nodes.
Since 2D-grid spaces cannot be directly interpreted as Euclidean spaces, it is not appropriate to apply space-related functions, such as Gaussian kernel density estimation, to the place fields. Place cell evaluation methods in previous research can be broadly categorized into four types [S11]:
1. Peak method: Classifies cells based on the average firing rate in one location being higher than in the rest of the environment.
2. Stability method: Classifies cells with stable firing patterns across locations over time.
3. Information method: Classifies cells based on the increased amount of spatial information they hold about the animal's location.
4. Combination method: A combined approach that considers cell's place fields, including their size, peak, and activity.
We chose the Peak method approach because it easily applies to 2D-grid structures.
To implement the Peak method in a 2D grid environment, we represent the 2D grid environment as a graph structure where neuron activities at each node are tracked while the agent explores the environment. Averaging neuronal activities per node yields the neuron's place fields, as illustrated in Figures 4 and 5. In these place fields, we find a node with a maximal firing rate,  $i_{\mathrm{max}}$ . Starting from this node, we find all connected components as a subgraph by connecting neighboring nodes within the nearest minimum. We create directed edges connecting nodes from high to low firing rates and run depth first search (DFS) from the node  $i_{\mathrm{max}}$ . DFS can be considered an inverted watershed algorithm that fills a single pool containing the global minimum. The advantage of our place cell metric is its ability to generalize to various graph structures, such as 2D hexagonal grids or 3D grid spaces.
A.1.3 Pseudo code for calculating place cell score metric
function PlaceCellScore(place field)
Input : place field  $(K\times K$  2D array)
Output: place score
$G\coloneqq 2\mathrm{D}$  grid graph  $(K\times K)$
$\mathcal{G} \coloneqq$  empty directed graph  $(K \times K)$
for edge  $(node_{i}\rightarrow node_{j})$  in  $\mathbf{G}$  do
if firing rate  $\rho_{i} >$  firing rate  $\rho_{j}$  then  $\mathcal{G}$  add edge (node  $\rightarrow$  node  $j$  end
Find node  $k$  of firing rate  $\rho_{max}$
for node  $v$  in  $\mathcal{G}$  do
if node  $v$  is not descendant of node  $k$  found with DFS  $(k)$  then delete node  $v$  from  $\mathcal{G}$  end