SlowGuess/NeurIPS_2023_OCR huggingface dataset:
conn. components = sum of all nodes' firing rates in  $\mathcal{G}$
total components  $=$  sum of all nodes' firing rates in  $G$
place score =  $\gamma \frac{\text{conn. components}}{\text{total components}}$
$\dagger$ $\gamma$  is discount factor, determined by connected component size
return place score
Algorithm 1: Pseudo code for calculating place cell score metric
The place field in Algorithm 1 is measured as following procedure:
1. During a random walk simulation, the activation value of a neuron at node  $i$ , where the agent is located, is measured every 65 steps. Let's say this value is  $a_{i}$ .
2. Every time the agent visits node  $i$  again, value  $a_i$  is added cumulatively to the recorded value;  $A_i + = a_i$  such that  $A_i$  is the cumulative activation value at node  $i$ . We assume the initial value of  $A_i$  is zero.
3. After the random walk is done,  $A_{i}$  divided by the length of the random walk trajectory is the firing rate  $\rho_{i}$  at node  $i$  of the neuron (place field  $\in \mathbb{R}^{K\times K}$ ).
In our place cell evaluation experiment, the length of the random walk is  $10^{5}$  and  $K = 11$ ; the evaluate map is one of the training maps.
A.2 Derivation of NMDAR nonlinearity from the molecular level chemical interaction
Here, we describe the NMDAR nonlinear dynamics from chemical interaction between  $\mathrm{Mg}^{2+}$  and NMDAR following previous literature [S12-S14]. At the molecular level, one  $\mathrm{Mg}^{2+}$  ion binds to one NMDAR receptor when opening the NMDAR channel. Thus, the chemical equation of binding reaction between  $\mathrm{Mg}^{2+}$  ion and NMDAR receptor, R, can be described as
$$
\mathrm {M g} ^ {2 +} + \mathrm {R} \rightleftharpoons \mathrm {M g} ^ {2 +} \mathrm {R}. \tag {S2}
$$
From this chemical equation, the equilibrium constant  $K$  is given by
$$
K = \frac {\left[ \mathrm {M g} ^ {2 +} \mathrm {R} \right]}{\left[ \mathrm {M g} ^ {2 +} \right] \left[ \mathrm {R} \right]}. \tag {S3}
$$
Thus, dissociation constant  $K_{D}$ , which correspond to  $\mathrm{Mg}^{2+}$  dissociation from NMDAR, follows
$$
K _ {D} = K ^ {- 1} = \frac {\left[ \mathrm {M g} ^ {2 +} \right] [ \mathrm {R} ]}{\left[ \mathrm {M g} ^ {2 +} \mathrm {R} \right]}, \tag {S4}
$$
in which [R] and  $\mathrm{[Mg^{2 + }R]}$  are the free and  $\mathrm{Mg}^{2 + }$  -bound NMDARs respectively. The fraction of opened NMDAR channels (number of unbound NMDAR over a number of total NMDAR) at equilibrium follows,
$$
\begin{array}{l} \mathbf {p} = \frac {[ \mathrm {R} ]}{[ \mathrm {R} ] + [ \mathrm {M g} ^ {2 +} \mathrm {R} ]} \tag {S5} \\ = \frac {1}{1 + [ \mathrm {M g} ^ {2 +} ] / K _ {D}} \\ \end{array}
$$
Experimentally, the voltage-dependent dynamics of  $K_{D}$  has been described as following equation by Ascher and Nowak [S15]
$$
K _ {D} = K _ {\mathrm {M g} ^ {2 +}} e ^ {\beta V}, \tag {S6}
$$
where,  $V$  is membrane voltage,  $\beta$  is a temperature constant and  $K_{\mathrm{Mg}^{2+}}$  is a dissociation constant at  $V = 0$ . If Eq. S6 is substituted into Eq. S5, voltage-dependent open fraction of NMDAR can be expressed as follows:
$$
\begin{array}{l} \mathbf {p} (V) = \frac {1}{1 + \frac {\left[ \mathrm {M g} ^ {2 +} \right]}{K _ {\mathrm {M g} ^ {2 +}}} e ^ {- \beta V}}. \tag {S7} \\ = \frac {1}{1 + \alpha e ^ {- \beta V}}. \\ \end{array}
$$
in which  $\alpha = [\mathrm{Mg}^{2+}] / K_{\mathrm{Mg}^{2+}}$ , the parameter determined by the  $[\mathrm{Mg}^{2+}]$ . Given the voltage-dependent open fraction of NMDAR,  $\mathbf{p}(V)$ , and NMDAR's maximal conductance,  $g_{\mathrm{max}}$ , then voltage-dependent NMDAR conductance  $g(V)$  can be described as
$$
g (V) = g _ {\max } \mathbf {p} (V) \tag {S8}
$$
Given  $g(V)$ , and driving force,  $V - V_{r}$ , and current  $I$ , they have a relationship of  $I = (V - V_{r})g(V)$  in which  $V_{r}$  is reversal potential (the value of membrane potential above which current inverts the direction). As experimental investigations on the physiological reversal potential of NMDAR to be  $V_{r} = 0$  [S16-S18],  $I = Vg(V)$ . Then, the normalized NMDAR current  $I_{\mathrm{norm}} = I / g_{\mathrm{max}}$  follows:
$$
I _ {\text {n o r m}} = V \mathbf {p} (V) \tag {S9}
$$
From Eq. S9 and previous electrophysiological experimental results [S19], we reconstruct IV curve in (Fig. 1a, right top).
A.3 NMDAR-inspired nonlinear activation function
Here, we propose an NMDAR-inspired nonlinear activation function from the nonlinear dynamics of the NMDAR-IV curve. If we consider the nonlinear IV curve of NMDAR (Eq. S9) as a nonlinear mapping function,  $\phi$ , the membrane voltage,  $V$ , can be viewed as an input,  $x$ , and normalized NMDAR current,  $I_{\mathrm{norm}}$ , as an output,  $\phi(x)$ . Therefore, we can rewrite the nonlinear mapping function  $\phi$  as follows
$$
\phi (x) = x \mathbf {p} (x). \tag {S10}
$$
We define the NMDAR-inspired activation function as a nonlinear mapping function,  $\mathrm{NMDA}(x) \coloneqq \phi(x)$ . By substituting Eq. S7 into Eq. S10, we show the generalized expression of  $\mathrm{NMDA}(x)$  equation with  $\alpha$  and  $\beta$  parameters as follows:
$$
\begin{array}{l} \mathrm {N M D A} _ {\alpha , \beta} (x) = x \mathbf {p} _ {\alpha , \beta} (x) \\ = \frac {x}{1 + \alpha e ^ {- \beta x}}. \tag {S11} \\ \end{array}
$$
Given  $\alpha = 1$  and  $\beta = 1$ ,  $\mathbf{p}(x)$  is identical to the sigmoid function,  $\sigma(x) = 1/(1 + e^{-x})$ . This particular case of  $\alpha$  and  $\beta$  leads to  $x\sigma(x)$ , Sigmoid Linear Unit (SiLU) activation function [S20]. In the case of  $\alpha = 1$  and  $\beta = 1.702$ ,  $x\sigma(1.702x)$  corresponds to the GELU activation function, which is popular in transformer models [S21, S22, S21]. Ramachandran et al. [S23] introduced the swish activation function,  $x\sigma(\beta x)$ , which is a generalized form of GELU and SiLU. They demonstrated
Figure S2: Comparison of common activation functions (left) and their derivatives (right) with  $\mathrm{NMDA}_{\alpha,\beta}$ .
that when  $\beta \to \infty$ , the activation function resembles RELU. We summarized these four activation functions by comparing them with our  $\mathrm{NMDA}_{\alpha,\beta}(x)$  in Table 1 and Fig. S2.
In contrast to the extensive research on  $\beta$  in  $\mathrm{NMDA}_{\alpha,\beta}(x)$ ,  $\alpha$ , the  $\mathrm{Mg}^{2+}$ -gating component, is not explored. For this reason, we focused on the parameter  $\alpha$  over  $\beta$ , and investigated  $\mathrm{NMDA}_{\alpha}(\mathrm{X})$ . It is interesting to note that the Swish function was originally proposed as a self-gating function, inspired by the use of the sigmoid function as a gating of information flow in the long short-term memory (LSTM) network [S24]. In contrast, our activation function  $\mathrm{NMDA}(x)$  is inspired by the physical  $\mathrm{Mg}^{2+}$ -gating mechanism that occurs at the real biological synapses. These shared mechanisms of self-gating in artificial models and biological observations raise the interesting possibility that NMDAR is a neural substrate of nonlinear activation function in the brain.
A.4 Detailed description of task design and definition of short-term working memory and long-term reference memory
Our task is based on a widely employed neuroscience experiment for two separate memory systems: working memory and reference memory [S25, S26]. Errors in working memory are measured by within-trial error, whereas errors in reference memory are measured by across-trial error. The training phase and the test phase alternate at each trial. In the test phase, the unvisited place prediction error and visited place prediction error for the familiar map and the novel map, respectively, are measured. The memory of a relatively recent experience can be defined as short-term working memory (STWM), and the memory of relatively old experience can be defined as long-term reference memory (LTRM). Within trial visited place prediction measures relatively short-term experience for our task. On the other hand, the across-trial unvisited place prediction task in the familiar map measures the relatively long-term experience. Measuring unvisited place prediction error in the novel map will establish a baseline of chance-level accuracy; above this baseline, the formation of long-term memory can be observed (Fig. S3).
Figure S3: Detailed task design of working and reference memory evaluation. At each random walk step, a batch is created (which is then used in the backpropagation step). The batch size is 512 since there are 512 parallel random walkers in use. Note that at each trial the agent randomly selects a map from training maps (familiar maps), the initial position of the agent is random, and the agent does a random walk.
A.5 Non-recurrent positional embeddings and prediction errors on the node visited for the first time
Figure S4: Experiment with non-recurrent positional embeddings. (a) Reference and working memory errors over training trials for training (familiar) maps and testing (novel) maps for  $N = 32$  where  $N$  is the number of training maps. (b) Reference memory errors evaluated on training maps over different values of  $\alpha$  in  $\mathrm{NMDA}_{\alpha}$  for  $N = 32$ . Error bars and shaded areas represent the standard deviation of errors from three independently trained models.
Non-recurrent positional embeddings We tested the non-recurrent positional embedding by substituting the recurrent positional embedding  $\mathbf{e}_t$  with the action embedding  $A(a_{t})$ , where  $A$  is the embedding layer and  $a_{t}$  is the action at step  $t$ . Compared to Fig. S4a, the result demonstrates a significant increase in working memory error and reference memory error (Fig. 3 vs. Fig. S4). Nonetheless, the model's behavior is comparable to the trend of decreasing reference memory error with increasing  $\alpha$  of  $\mathrm{NMDA}_{\alpha}$  (see Fig. S4b).
Moreover, we found that place cells do emerge in the feedforward layers for this standard positional embedding method (see Fig. S5). It is worth noting that these place cell scores are noticeably lower compared to those obtained through the recurrent positional embedding. Despite the hindered performance in working memory and reference memory, a strong correlation persists between the place cell scores in FFNs and the reference memory errors (see Fig. S5c and S5d).
Prediction errors on the node visited for the first time We compared unvisited node prediction error (unvisited within context window, in this case, 64 steps) versus first visited node prediction error (unvisited for within a trial). As shown in Fig. S6, the prediction error results for the first visited node do not differ from the reference memory error results.
These results strongly support the conclusions that (1) while the path-integrated information from recurrent positional embedding is important for learning the spatial structure of the map, this information is not used in predicting the unvisited node, and (2) the reference memory is used for predicting an unvisited node on a familiar map.
Figure S5: Place cell score analysis for non-recurrent positional embeddings for  $N = 32$ . Place cell score distributions with varying  $\alpha$  in self-attention layers (a) and feed-forward layers (c). Scatter plot of average place cell scores and reference memory errors for self-attention layers (b) and feed-forward layers (d).  $r$  and  $p$  denote Spearman's rank correlation coefficient and significance score, respectively.
Figure S6: First visited node prediction error evaluated on training maps over different values of  $\alpha$  in  $\mathrm{NMDA}_{\alpha}$  for  $N = 32,48$ , and 64. Error bars and shaded areas represent the standard deviation of errors from three independently trained models.
A.6 Statistical significance between the NMDA and GELU activation functions & learning rate effect: is it an under-trained effect?
In this section, we address concerns about the NMDA style activation with large  $\alpha$  being a trivial effect for the emergence of place cells and the improvement of learning speed. We provide additional evidence to support our claims and demonstrate that the observed effects are not trivial or due to under-training.
Statistical test across activation functions Our primary claim regarding the effectiveness of NMDA-style activation is based on the results shown in Fig. 3. To support our claim, we conducted a
statistical across various activation functions. Our findings revealed a significant difference between  $\mathrm{NMDA}_{\alpha}$  and other activation functions, indicating that the NMDA style activation is beneficial for the emergence of place cells and the improvement of learning speed (see Table S7).
Figure S7: Statistical test result across various activation functions with different numbers of maps (N) related to Figure 3. (a-c) Mann-Whitney test across various  $\alpha$  of  $\mathrm{NMDA}_{\alpha}$ . (d-f) Mann-Whitney test across various activation functions. The sample size (or trials) used for tests ranges from 3 to 11. Colors represent  $p$ -values and asterisks indicate statistical significances: * 0.01 < p < 0.1; ** 0.001 < p < 0.01; *** p < 0.001.  $\alpha = 10$  for NMDA and  $\alpha = 0.01$  for LeakyReLU.
Gradient around  $x = 0$  and learning speed The performance gain at large alpha values might be due to an increased gradient around  $x = 0$ . However, the gradient around  $x = 0$  at  $\alpha = 1$  (green) is actually larger than  $\alpha = 10$  (blue) in Fig. S2. Therefore, the performance gain at large alpha values is not due to an increased effective learning speed around  $x = 0$ . Instead, it may be attributed to the activation function's ability to better capture nonlinear relationships.
Learning rate effect on reference memory error To clarify that our result is not simply due to the under-trained model, we performed experiments with varying learning rates: 0.001, 0.0005, 0.0003, 0.0002, and 0.0001. As can be seen in Fig. S8, a larger learning rate at low alpha values is not helpful for reducing reference memory error. It is worth noting that all models' working memory error is almost zero. Larger learning rates lead to faster decrease in working memory error, but at high learning rates, reference memory error stays around the chance level of 0.9, i.e., no reference memory emerges. Additionally, Fig. S9 shows that higher place cell scores correspond to lower reference memory errors. Note that we employed a linear decay learning rate schedule (from the starting learning rate to 0) and the starting learning rate used in our manuscript is 0.0001.
In conclusion, our additional experiments and analyses demonstrate that the NMDA style activation with large  $\alpha$  plays a critical role in the emergence of place cells and the improvement of reference memory, and it is not a trivial effect or an under-trained effect.
Figure S8: Familiar map reference memory error results with five different learning rates, where  $\alpha \in \{0.1, 1.0, 10.0\}$ , for a total of 64 training maps ( $N = 64$ ). The error bars represent the standard deviation of errors from three independently trained models.
Figure S9: (top) Scatter plot between reference memory error and average place score; (bottom) Distribution of place cell scores with five different learning rates, where  $\alpha \in \{0.1, 1.0, 10.0\}$  (shown at the top of each figure), for a total of 64 training maps ( $N = 64$ ).
A.7 Consequent of changing nonlinear dynamics in Leaky ReLU activation function
Here, we investigated the consequence of changing nonlinearity with an activation function other than  $\mathrm{NMDA}_{\alpha}$ . We chose LeakyReLU  $(\max(0,x) + \alpha \min(0,x))$  activation function to compare with  $\mathrm{NMDA}_{\alpha}$ . Regarding LeakyReLU,  $\alpha = 1$  of LeakyReLU also leads to a fully linear activation function similar to  $\alpha = 0$  of  $\mathrm{NMDA}_{\alpha}$ . Compared to  $\mathrm{NMDA}_{\alpha=10}$ , LeakyReLU exhibits a lower average place score in the allowed range of  $\alpha$ , indicating that  $\mathrm{NMDA}_{\alpha}$  performs better in terms of robust place cell emergence (see Fig. S10)
Figure S10: Evaluation of reference memory error in LeakyReLU (LRELU) while modulating  $\alpha$  (top) and relationship of average place cell score and reference memory error (bottom).
A.8 Transformer as a memory consolidation model and its biological plausibility
Next, we examined the biologically inspired  $\mathrm{NMDA}_{\alpha}$  activation function in the feed-forward layer of the transformer and its role in memory formation and place cell representation. We show that modulating  $\alpha$  corresponds to a change in extracellular  $[\mathrm{Mg}^{2+}]$ , by deriving the nonlinear activation function from the real NMDAR nonlinear IV curve. The reconstructed real nonlinear IV curve is shown in Fig. 1a (top right panel).
The modulation of  $\alpha$  selectively affects the formation of long-term reference memory (i.e., prediction of unvisited places across trials) while leaving the formation of short-term working memory (i.e., prediction of unvisited places within trials) intact. 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 feed-forward layer. In psychology, a multi-store model for the animal brain that comprises short-term memory and long-term memory was historically suggested in [S27]. In this biological model, sensory inputs are stored in short-term memory systems via attention, and some are transferred to a long-term memory system while others quickly disappear.
In neuroscience, the transfer of short-term memory into a long-term system is called memory consolidation [S28]. Animal studies have demonstrated that the CA1 region of the hippocampus is essential for memory consolidation [S29, S30]. In hippocampal CA1, the postsynaptic NMDA receptor mediates synaptic plasticity, and the selective perturbation of these receptors leads to impairment in long-term memory formation [S31, S30]. Later research revealed that  $\mathrm{Mg}^{2+}$ -gating of NMDA receptors modulates the formation of long-term memory [S32, S33]. These observations imply that the nonlinear dynamics of NMDA receptors in CA1 are critical for consolidating short-term memory into long-term memory S11.
On the basis of a previous link between the hippocampus and the transformer model, we hypothesize that the latter is a model for memory consolidation. Given the resemblance of the GELU nonlinear activation function and CA1 NMDAR nonlinear IV curve, we assumed that the GELU activation function serves as a key component that links short-term working memory and long-term reference memory. Our experimental results indicate that the formation of long-term reference memory is impaired when the activation function is completely linear (corresponding to no  $\mathrm{Mg}^{2+}$ ). In contrast, increasing  $\alpha$  (which corresponds to an increase in  $\mathrm{Mg}^{2+}$  level) revealed that for long-term reference memory, our model outperforms other activation functions (e.g., RELU, GELU, LRELU, Sigmoid, Tanh). Based on these similarities between hippocampal memory consolidation and our results, we propose that transformer is an effective memory consolidation model.
Figure S11: Schematic illustration of the transformer as a memory consolidation model of hippocampus via NMDAR nonlinearity. Top, transformer with  $\mathrm{NMDA}_{\alpha}$  nonlinear activation function. Bottom, summarized brain model with NMDA receptor mediated long-term memory formation. Colors represent the corresponding stage of the memory consolidation process.
In addition to the performance gain in long-term memory formation with  $\mathrm{NMDA}_{\alpha}$ , we found that modulating the  $\alpha$  affects the emergencies of place cells in the feed-forward layer and conclude that there is a significant correlation between place cell score and long-term reference memory formation. Our results align with previous biological findings that perturbation of CA1 NMDARs leads to impairment in both place cell representation and long-term memory formation [S31, S34, S35, S29]. These similarities support the idea that place cells are the neural correlates of long-term spatial memories. Altogether, our findings suggest the exciting possibility that the nonlinear IV curve of NMDAR in the hippocampal CA1 is a neural substrate of nonlinear activation function in the brain.
A.9 GPT2 and Vision Transformer with NMDA activation functions
In this section, we present the detailed experimental setup and results for applying the proposed  $\mathrm{NMDA}_{\alpha}$  activation function to the GPT2 [S36] and Vision Transformer (ViT) models.
A.9.1 Language Modeling with GPT2
We replaced the GELU activation function in the feed-forward networks (FFNs) of the GPT2 model (model size: 124M) with the  $\mathrm{NMDA}_{\alpha = 10}$  activation function and conducted experiments on the Open-WebText dataset [S37]. We employed the GPT2 training PyTorch code from karpathy/nanoGPT GitHub repository [S38]. The application of  $\mathrm{NMDA}_{\alpha = 10}$  showed a slight increase in performance, as illustrated in Fig. S12.
A.9.2 Image Classification Tasks with ViT
We conducted experiments on image classification tasks using the ViT model, and the results were examined on the CIFAR-100 and TinyImageNet datasets. Three trials were conducted for each condition, and an increasing, though statistically insignificant, tendency in performance was observed. The top-1 test accuracies for these datasets are provided in Table S1.
Table S1: Top-1 test accuracies with ViT model on CIFAR-100 and TinyImageNet datasets.
Dataset	GELU	NMDAα=10	NMDAα=0
CIFAR100	69.92 ± 0.34	70.27 ± 0.47	49.91 ± 0.58
TinyImNet	54.90 ± 0.45	55.72 ± 0.03	40.36 ± 0.52
Hyperparameter Consideration It is important to note that we did not scan a range of hyperparameters for either NMDA or the GELU-based model. The optimal hyperparameters for NMDA might differ from those for the GELU-based model.
Figure S12: Validation loss curve over training iteration for GPT2 model on the OpenWebText dataset. The blue (red) line represents the GPT2 with GELU  $(\mathrm{NMDA}_{\alpha = 10})$  function.
Future Works Although these results are promising, they are not conclusive in stating that NMDA generally outperforms GELU for various real-world tasks. Future works will include evaluating our NMDA function on different model sizes of GPT or other language models, as well as evaluating the pre-trained models on downstream NLP tasks such as reading comprehension, question answering, common sense, MMLU, and BIG-bench.
A.10 Sparsity of Activation
In this section, we present more intuitive evidence to support the effectiveness of our proposed nonlinearity, based on the sparsity of activation in FFNs.
Mathematical Intuition Given the  $\mathrm{NMDA}_{\alpha}$  activation function,  $\mathrm{NMDA}_{\alpha}(x) = x / (1 + \alpha e^{-x})$  we can rewrite this function as follows:
$$
\mathrm {N M D A} _ {\alpha (x)} = x \cdot \frac {1}{1 + e ^ {- (x - c)}}, \quad \text {w h e r e} \alpha = e ^ {c}. \tag {S12}
$$
In the given expression, an increase in  $\alpha$  results in shifting the sigmoid function towards the positive direction by increasing  $c$ . This shift can be understood as an increase in the threshold of information gating, which is mediated by the sigmoid function. Consequently, an increase in  $\alpha$  might enhance the sparsity of the NMDA activation.
Empirical Results To validate the aforementioned mathematical intuition, we conducted measurements on the sparsity of the activities in the feed-forward layer by calculating the Gini index [S39] for each input sequence. The Gini index is defined as follows:
$$
G = \frac {\sum_ {i = 1} ^ {K} \sum_ {j = 1} ^ {K} \left| x _ {i} - x _ {j} \right|}{2 K ^ {2} \bar {x}}, \tag {S13}
$$
where  $x_{i}$  is the i-th neuron's activation value and  $\bar{x} = \sum_{i=0}^{K} x_{i} / K$  is the mean of the activation values in the feed-forward layer ( $K = 2048$  is the total number of neurons in a feed-forward layer).