SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dong-Kyum  $\mathbf{Kim}^{1*}$  Jea Kwon $^{2*}$  Meeyoung Cha $^{1,3\dagger}$  C. Justin Lee $^{2\dagger}$
$^{1}$ IBS Data Science Group
$^{2}$ IBS Center for Cognition and Sociality
$^{3}$ KAIST School of Computing
{kdkyum,jeakwon,mcha,cjl}@ibs.re.kr
The hippocampus plays a critical role in learning, memory, and spatial representation, processes that depend on the NMDA receptor (NMDAR). Inspired by recent findings that compare deep learning models to the hippocampus, we propose a new nonlinear activation function that mimics NMDAR dynamics. NMDAR-like nonlinearity shifts short-term working memory into long-term reference memory in transformers, thus enhancing a process that is similar to memory consolidation in the mammalian brain. We design a navigation task assessing these two memory functions and show that manipulating the activation function (i.e., mimicking the  $\mathrm{Mg}^{2+}$ -gating of NMDAR) disrupts long-term memory processes. Our experiments suggest that place cell-like functions and reference memory reside in the feed-forward network layer of transformers and that nonlinearity drives these processes. We discuss the role of NMDAR-like nonlinearity in establishing this striking resemblance between transformer architecture and hippocampal spatial representation.
Synaptic plasticity in the hippocampus is crucial for converting short-term memories into long-term memories during memory consolidation [1-3]. The N-methyl-D-aspartic acid receptor (NMDAR) is essential for mediating this memory formation as well as spatial representation. NMDAR serves as a switch for such plasticity and long-term memory formation [4-6]. Hippocampal place cells [7] and entorhinal cortex grid cells [8] are known to be central for spatial navigation in animals. NMDAR has been highlighted for its importance in place cell representations through hippocampal CA1 neurons [9, 3]. These discoveries have led to a deeper understanding of hippocampal function, inspiring recent efforts to replicate such spatial representation mechanisms in deep neural networks [10-12]. However, whether non-linear dynamics resembling hippocampal functions can be developed and used to support spatial representation in deep learning models remains unclear.
NMDAR is a post-synaptic ion channel that is characterized by nonlinear dynamics that distinguish it from other ion channels in the brain. These nonlinear dynamics are evident in the whole-cell current voltage (I-V) relationship (Fig. 1a) and are modulated by  $\mathrm{Mg}^{2+}$  ion blockade at the pore region of the channel. Previous research indicates that the  $\mathrm{Mg}^{2+}$ -dependent nonlinear dynamics of NMDAR play a key role in synaptic plasticity and memory formation [5, 6].
Recently, transformer-based deep learning models have been reported to have functions that resemble hippocampal formations [13]. Transformers comprise two consecutive modules: a self-attention
Figure 1: (a) Schematic diagram of  $\mathrm{Mg}^{2+}$ -gated NMDAR modulating synaptic plasticity (left), current-voltage dynamics (I-V curve; top right) and an NMDAR-inspired activation function,  $\mathrm{NMDA}_{\alpha}(x)$  (bottom right). (b) Transformer architecture and nonlinear activation functions in its feed-forward network (bottom left): ReLU, Gaussian Error Linear Unit (GELU), and Sigmoid Linear Unit (SiLU).
layer and a feed-forward network (FFN; see Fig. 1b). The self-attention layer is closely related to a recent neuroscience model that bridges with transformer [14] and it has been argued that softmax neurons in this layer behave like place cells in a navigation task [12]. However, the role of neurons in FFNs involved in a spatial navigation task has yet to be elucidated and compared to hippocampal characteristics.
This paper uncovers a resemblance between the NMDAR nonlinearity and recently developed activation functions commonly used in FFNs of deep learning models (Fig. 1). NMDAR functions operate on activity-dependent repulsion of  $\mathrm{Mg}^{2+}$  ions [15, 16] and this phenomenon is particularly interesting because it supports self-gating of information flow ( $\mathrm{Ca}^{2+}$  ion influx) in the post-synaptic region. Similar to this NMDAR activity-dependent gating mechanism, activation functions in modern neural networks combine input with a self-gating function (i.e., a sigmoidal function that ranges between 0 to 1) that determines information flow. Our observation prompted the following inquiry: Can NMDAR-like nonlinearity in the feed-forward network layer of transformers enhance the formation of long-term memory and spatial representation by place cells?
Here we derive a novel NMDAR-like activation function using the NMDAR current-voltage (I-V) curve. In parallel we design a spatial navigation task in a 2D grid environment that assesses two memory types well-defined by neuroscience research: working memory and reference memory [17, 18]. The former assesses events from within-trial, while the latter controls across trials in a stable (unchanging) environment. We then test the transformer model with the NMDAR-like activation function and find that 1) place cell representations emerge in the FFN layer rather than in the self-attention layer, 2) the nonlinearity of the NMDAR-like activation function can regulate reference memory, 3) place cell-like neurons in FFNs are strongly correlated with reference memory, whereas this correlation is not observed in the self-attention layer; and 4) NMDAR-like activation shows the best reference memory performance when compared to other available activation functions.
Collectively, these findings suggest that adopting NMDAR-like nonlinearity in FFN of transformer models can enhance the formation of long-term memory and spatial place cell representation, similar to previous observations in the animal brain. We believe these findings have important implications for developing new brain-inspired AI models and for understanding similar processes that occur in the brain and AI models.
2.1 Designing a 2D navigation task to test the role of working memory and reference memory
We designed a sensory observation prediction task in which an agent randomly walks in a 2D grid environment and is trained to predict subsequent sensory observations (see Fig. 2a) [12]. The agent
Figure 2: (a) Sensory observation prediction task in a 2D grid, where dotted squares indicate the target position to predict given a sequence of past actions and observations. The unvisited and visited places are represented in gray and black blocks, respectively. (b) A transformer model for predicting the location of an upcoming sensory observation based on sequences of [Action  $(a)$ , Observation  $(x)$ ] pairs. Using a sequence of pairs in the context, the model is trained to predict the masked observation (i.e., the subsequent observation) corresponding to the final query action.
receives a sequence of previous [Action  $(a)$ , Observation  $(x)]$  pairs as input, and its goal is to predict the next observation, which is masked (i.e., dotted squares in the sequence of Observation  $(x)$  in the figure). We use the transformer architecture as our model.
We generated  $N$  maps of  $11 \times 11$  2D grids. A random sensory observation chosen from ten letters is placed at each position on each map. Agents can move 'up,' 'right,' 'down,' 'left,' or 'stay.' An agent starts at a random position and initiates a random walk on the map, which is randomly selected from  $N$  training maps, for 2,048 steps for each trial.
Our design novelty is the consideration of two memory types: short-term working memory and long-term reference memory. When the prediction based on nodes that were previously visited during the random walk is incorrect, it counts as a working memory error (Fig. 2a left). In contrast, when the prediction based on unvisited nodes is incorrect, it counts as a reference memory error (Fig. 2a right). See details of the navigation task and definitions in Appendix A.4 and Fig. S3.
It is important to note that minimizing the reference memory error by memorizing input sequences is infeasible; the possible number of sequence configurations is exponential since the input sequence is randomly generated at each trial. To solve this task, the model needs to: 1) understand the abstract structure of 2D space, 2) infer which map it is on from the input sequence data, and 3) memorize the sensory observations and their positions on that map. Compared to the previous works [14, 12] that focused on working memory error, the current paradigm evaluates two distinct types of memory error and gives a more holistic view of model performance.
2.2 Transformer: two separate memory systems
We here review the self-attention mechanism and FFNs in the transformer architecture [13] that hypotheses about two separate memory systems are based on—working memory (formed within-trial) and reference memory (formed across-trials)—and why they are inferred to reside in self-attention layers and FFNs, respectively.
Self-attention mechanism Given a sequence  $\{\mathbf{x}_1,\dots,\mathbf{x}_T\}$  of  $d$  -dimensional input embeddings, the self-attention layer calculates the interaction term between each embedding element within a context window via the self-attention mechanism. More formally, each input embedding applies two linear layers  $(W_{k}$  and  $W_{v})$  to the embeddings to form the key matrix  $K$  and value matrix  $V$  ..
$$
K ^ {\top} = \left[ \mathbf {k} _ {t - c} ^ {\top} \mathbf {k} _ {t - c + 1} ^ {\top} \dots \mathbf {k} _ {t} ^ {\top} \right], \quad V ^ {\top} = \left[ \mathbf {v} _ {t - c} ^ {\top} \mathbf {v} _ {t - c + 1} ^ {\top} \dots \mathbf {v} _ {t} ^ {\top} \right], \tag {1}
$$
where  $\mathbf{v}_i = \mathbf{x}_i W_v$  ( $W_v \in \mathbb{R}^{d \times d_k}$ ) and  $\mathbf{k}_i = \mathbf{x}_i W_k$  ( $W_k \in \mathbb{R}^{d \times d_k}$ ). Here,  $c$  denotes the context length. The key matrix  $K \in \mathbb{R}^{(c + 1) \times d_k}$  is then used to compute the interaction score between an input embedding at step  $t$  and all the vectors in  $K$  via dot products:
$$
\mathbf {s} _ {t} = \mathbf {q} _ {t} K ^ {\top}, \quad \text {w h e r e} \mathbf {q} _ {t} = \mathbf {x} _ {t} W _ {q} \left(W _ {q} \in \mathbb {R} ^ {d \times d _ {k}}\right). \tag {2}
$$
The normalized values of  $\mathbf{s}_t\in \mathbb{R}^{(c + 1)}$ , called attention values, are calculated via the softmax function; the final output of the self-attention mechanism is a weighted sum of the value vectors in
Table 1: Comparison of common activation functions with  ${\mathrm{{NMDA}}}_{\alpha ,\beta }$
NMDAα,β	Name	Equation	Reference
NMDAα=1,β=1(x)	SiLU(x)	xσ(x)	[20]
NMDAα=1,β=1.702(x)	GELU(x)	xσ(1.702x)	[21]
NMDAα=1,β=∞(x)	ReLU(x)	max(0,x)	[22]
NMDAα=1,β(x)	Swish(x)	xσ(βx)	[23]
$V\in \mathbb{R}^{(c + 1)\times d_k}$  with the attention values:
$$
\mathbf {y} _ {t} = \operatorname {s o f t m a x} \left(\frac {\mathbf {q} _ {t} K ^ {\top}}{\sqrt {d _ {k}}}\right) V. \tag {3}
$$
After this update,  $\mathbf{y}_t\in \mathbb{R}^{d_k}$  is updated by another linear transformation  $W_{o}\in \mathbb{R}^{d_{k}\times d}\colon \mathbf{z}_{t} = \mathbf{y}_{t}W_{o}$ . The output  $\mathbf{z}_t$  is added to the  $\mathbf{x}_t;\mathbf{z}_t + \mathbf{x}_t$  providing the final output of the self-attention layer, and this information is sent through to the subsequent layer.
Feed-forward networks (FFNs) This component consists of two linear layers with a point-wise nonlinear activation function  $\phi$ :
$$
\operatorname {F F N} \left(\mathbf {x} _ {t}\right) = \phi \left(\mathbf {x} _ {t} U _ {1} ^ {\top}\right) U _ {2}, \tag {4}
$$
where  $U_{1} \in \mathbb{R}^{d_{f} \times d}$  and  $U_{2} \in \mathbb{R}^{d_{f} \times d}$  are trainable weight matrices. Sukhbaatar et al. [19] showed that Eq. (3) and Eq. (4) have similar structures except for the following: 1)  $U_{1}$  and  $U_{2}$  matrices are fixed over different input sequences while  $K$  and  $V$  matrices are dynamically changed with input and 2) operations in FFNs are point-wise or local while the self-attention layer has non-local operations, e.g., the softmax function and dot products between different elements. This observation suggests that the FFNs store "general" knowledge about the task that does not depend on the situation.
2.3 Resemblance of NMDA receptor nonlinear dynamics with modern activation functions
NMDAR's nonlinear dynamics mostly arise from the voltage-gated  $\mathrm{Mg}^{2+}$  repulsion at the NMDAR channel's pore [15, 16] (Fig. 1a left). Previous work showed this nonlinear  $I-V$  relationship to be:
$$
I _ {\text {n o r m}} = V \mathbf {p} _ {\alpha , \beta} (V) \tag {5}
$$
where  $V$  represents an input voltage and  $\mathbf{p}_{\alpha,\beta}(V)$  is a voltage-dependent channel open probability that follows the ion blockade model [24]:
$$
\mathbf {p} _ {\alpha , \beta} (V) = \frac {1}{1 + \alpha e ^ {- \beta V}} \tag {6}
$$
where  $\alpha = [\mathrm{Mg}^{2+}] / K_{\mathrm{Mg}^{2+}}$  is a parameter determined by  $[\mathrm{Mg}^{2+}]$ ,  $K_{\mathrm{Mg}^{2+}}$  is a dissociation constant, and  $\beta$  is a temperature constant. For further details, see Appendix A.2.
The dynamics of the NMDA receptor closely resemble those of new activation functions such as ReLU or GELU (see Fig. S2 in Appendix A.3). This visual resemblance motivated us to define a new NMDAR-inspired activation function (see details in Appendix A.3) as follows:
$$
\mathrm {N M D A} _ {\alpha , \beta} (x) = x \mathbf {p} _ {\alpha , \beta} (x) = \frac {x}{1 + \alpha e ^ {- \beta x}}. \tag {7}
$$
This  $\mathrm{NMDA}_{\alpha, \beta}(x)$  incorporates modern activation functions with varying values of temperature constant,  $\beta$  (Table 1). To investigate the  $\mathrm{Mg}^{2+}$ -gated nonlinear dynamics,  $\alpha$ , we compared various activation functions with  $\mathrm{NMDA}_{\alpha}(x) = x \mathbf{p}_{\alpha, \beta=1}(x)$ .
2.4 Testing the NMDAR-inspired activation in navigation tasks
The transformer model is trained using softmax cross-entropy loss to predict subsequent sensory observations (i.e., dotted squares in Fig. 2). Instead of using sinusoidal positional encoding [13], we employ recurrent positional embedding which encodes the location of an input element by using the recurrent neural network (RNN); this method is closely related to the most advanced neuroscience model of the hippocampus [12].
We generate the embedding vectors of the sensory observation sequence with a word embedding layer, but the embedding vectors of the action sequence are generated using an RNN;  $\mathbf{e}_{t + 1} = \tanh (\mathbf{e}_tW_a)$  where  $\mathbf{e}_t$  is a recurrent positional embedding at step  $t$ , and  $W_{a}$  is the action-dependent trainable weight matrix. The input is given by  $\{[\mathbf{e}_1,\mathbf{x}_1],[\mathbf{e}_2,\mathbf{x}_2],\ldots ,[\mathbf{e}_t,\mathbf{x}_t]\}$ , where  $\mathbf{x}$  denotes the embedding vector of sensory observation  $x$ ; the initial recurrent positional embedding  $\mathbf{e}_1$  is sampled from a normal distribution, and we mask the last observation  $x_{t}$ .
In our experiment, FFN in the transformer model consists of two linear layers (see Fig. 1b and Eq. (4)) with the NMDAR-inspired activation function  $\mathrm{NMDA}_{\alpha}$ . We use TransformerXL [25] with an extended memory length of 32 and segment length of 32 so that the context length  $c$  is 64 and working memory error is measured when the node to predict its sensory observation is in the context window (see Fig. 2b); i.e., a node that the agent has not been visited in the last 64 steps is treated as an unvisited node. Note that our model is unable to access the sensory observations of unvisited nodes via the self-attention mechanism due to the fixed context window size.
The input embedding is a concatenated vector  $[\mathbf{e},\mathbf{x}]$  of the word embedding  $\mathbf{x}$  (dimension of 256), and the recurrent positional embedding  $\mathbf{e}$  (dimension of 256) so that the total input embedding dimension is 512. The number of heads in the self-attention layer is 8, and the number of neurons in the FFN is 2,048. The dropout rate is set to 0.1, and the maximum clip norm of the gradient is set to 0.25. We employ the ADAM optimizer [26] and a learning rate schedule with a linear decay from 0.0001 (start) to 0 (end). We run 512 random walk simulations (trials) in parallel to collect training trajectories. The total number of random walking steps is 2,048 for each simulation, and the total number of steps for training a model is 512 (batch size; the number of trials per epoch)  $\times$  2,048 (total number of steps in a trial)  $\times$  200 (number of epochs) (see Fig. S3 in Appendix A.4). The average number of unvisited nodes in a single trial was 561. The PyTorch code for reproducing our results is available at https://github.com/kdkyum/transformer_memory_consolidation.
3.1 Working memory and reference memory errors
To measure the impact of nonlinearity  $\alpha$  in the FNNs, we trained transformer models with different values of  $\alpha$  in [0, 0.01, 0.05, 0.1, 0.5, 1, 5, 10] and evaluated the working memory and reference memory errors on the train maps (i.e., familiar maps) and test maps (i.e., novel maps). The changing  $\alpha$  values in a transformer model mimic the changes in NMDAR  $\mathrm{Mg}^{2+}$ -gating in the brain, as inspired by neuroscience findings that selective inhibition of hippocampal CA1 NMDAR can disrupt the consolidation of short-term working memory to long-term reference memory.
The top left plot in Fig. 3a shows that the reference memory error on the training maps rapidly decreased over training trials when  $\alpha$  is larger than zero, with a larger improvement for increasing  $\alpha$ . The reference memory error on the novel maps, however, was nearly constant at a chance level of  $0.9 (= 1 - 1 / (\text{number of letters}))$  for all  $\alpha$  (see Fig. 3a top right). Fig. 3a bottom right shows that working memory was active on novel maps that had not previously been shown during training. This finding suggests that working memory formation is intact on novel maps. Training the models on different numbers of maps,  $N$ , Fig. 3b shows that increasing nonlinearity (i.e.,  $\alpha$ ) helped to activate reference memory, and the trend of improvement was consistent, as shown for  $N = 32, 48, \text{and} 64$ . Training more maps led to larger reference memory errors, because increasing  $N$  means the model needs to memorize what-where pairs (i.e., each training contains unique what-where information).
In addition, we find that removing nonlinearity in the NMDAR-inspired activation function  $(\alpha = 0)$  impairs reference memory formation (Fig. 3b) but leaves working memory formation intact (Fig. 3a). This means that even though the trainable parameters exist in the self-attention layer, a lack of nonlinearity in the feed-forward layer significantly impairs the formation of long-term reference memory in familiar maps. This result suggests that short-term working memory and long-term reference memory are physically stored in separate structures (the self-attention layer and the FFN) of transformers and possibly gated by the nonlinear activation functions residing in the FFN.
We next tested other nonlinear activation functions, including GELU  $(x\sigma(1.702x))$ , ReLU  $(\max(0, x))$ , LeakyReLU  $(\max(0, x) + 0.01\min(0, x))$ , sigmoid, and tanh. Fig. 3c shows that the newly proposed biologically-inspired function  $(\mathrm{NMDA}_{\alpha=10})$  yielded substantially lower reference memory errors relative to other activation functions. A comparison with the second-best-performing
GELU activation, for example, indicated a statistically significant difference with  $p$ -value less than 0.05 (see Appendix A.6). These results open up the possibility of finding a better activation function through a new  $\alpha$  hyperparameter space of the NMDA-inspired activation function.
Expanding on the evaluation, one may consider other memory types like recurrent positional embedding instead of reference memory. We used non-recurrent positional embeddings to train the models (see Appendix A.5) and confirmed that working memory and reference memory errors increased more substantially. This finding supports the idea that working memory is crucial for memory consolidation and that disrupting it impairs reference memory. However, we also saw a similar trend of decreasing reference memory error while increasing  $\alpha$  of  $\mathrm{NMDA}_{\alpha}$  (see Fig. S4).
We further assessed the prediction error of the first visited node. While the reference memory error is defined as a prediction error on a node that the agent has not visited in the previous 65 steps, the first visited node prediction error is a prediction error on a node that the agent visits for the first time in a trial. The results for the first visited node prediction error were identical to the results for the reference memory error (see Fig. S6 in Appendix A.5). These findings suggest that reference memory is used in training maps to predict the unvisited node.
3.2 Place cells in feed-forward networks
One of the most striking characteristics of the hippocampus is the presence of the place cells [7], which are neurons in the brain that fire at a particular place in the environment [7]. Studies have shown that hippocampal place cells encode spatial location through localized firing patterns. They have been considered a substrate for long-term memory of the location where specific events occurred (e.g., a previously visited position in a navigation task). Selective impairment of NMDAR in hippocampal CA1 disrupts place cell emergence and memory consolidation of long-term memory [2, 3, 9].
We investigated whether place cells emerge during the formation of long-term reference memory in a transformer model, building upon our previous results demonstrating selective impairment in long-term reference memory with  $\mathrm{NMDA}_{\alpha}$  modulation. We newly designed a metric called the place cell score described in Eq. S1 in Appendix A.1 and examined the spatial representation within the transformer architecture. This metric, ranging from 0 (i.e., homogeneous firing) to 1 (i.e., firing is specific to the agent's position in the grid) quantifies the firing specificity of each neuron for spatial
Figure 3: (a) Reference and working memory error rates over training trials for training maps and testing maps for  $N = 32$  training maps. Insets in the bottom figures show working memory error rates during the initial training phase. (b) Reference memory errors were evaluated on training maps over different values of  $\alpha$  in  $\mathrm{NMDA}_{\alpha}$  and  $N$ . (c) Reference memory error comparison between  $\mathrm{NMDA}_{\alpha} = 10$ , GELU, ReLU, LeakyReLU, sigmoid, and tanh activation functions. Inset: magnified view of the top 4 activation functions. Error bars and shaded areas represent the standard deviation of errors from three independently trained models.
Feed-forward layer
Self-attention layer
vs other AFs
(h)
Top-16 Place Fields
Figure 4: Reference memory-related place cells selectively emerge in FFN but not in the self-attention layer over increasing  $\alpha$ . (a, b) Example rate maps with place scores in FFN and self-attention layers at  $\alpha = 10$ ; from top left (high) to bottom right (low); color bar indicates the firing rate between  $\rho_{\mathrm{max}}$  and  $\rho_{\mathrm{min}}$ . (c-d) Place cell score distributions with varying  $\alpha$  in FFN (c) and self-attention layers (d). (e-f) Scatter plot of average place cell scores and reference memory errors.  $r$  and  $p$  denote Spearman's rank correlation coefficient and significance score, respectively. (g) place cell score distribution and relationship of average place cell scores and reference memory errors in common activation functions: GELU, ReLU, LeakyReLU, tanh, and sigmoid. All results are evaluated using training maps. (h) Rate maps of neurons with top-16 place cell scores in the FFNs with varying values of  $\alpha$ ;  $\alpha = 10$  (left),  $\alpha = 1$  (middle), and  $\alpha = 0$  (right).