Transformer as a hippocampal memory consolidation model based on NMDAR-inspired nonlinearity
Authors: Dong-Kyum Kim
Jea Kwon
Meeyoung Cha
C. Justin Lee
Venue: NeurIPS 2023 poster (2023)
Keywords: Transformer
NMDA
long-term memory
reference memory
memory consolidation

Abstract: 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 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.

Transformer as a hippocampal memory consolidation model based on NMDAR-inspired nonlinearity
Dong-Kyum Kim
Jea Kwon Meeyoung Cha
C. Justin Lee
IBS Data Science Group
IBS Center for Cognition and Sociality
AKAIST School Computing
[kdkkyu,jeekkwnn macha cjl)@ibs.ree re..iibs.ee............ re.
Abstract
The hippocampus plays a critical role in learning, memory, and spatial representa- tion, 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 memory consolida- tion in the mammalian brain. We design navigation task assessing these two memory functions and show that manipulating the activation function (i.c., mim- icking the Mg of NMDAR) disrupts long-term memory processes. Our experiments suggest that place cell-like functions and reference memory reside the fred-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.
Introduction
Synaptic plasticity in the hippocampus is crucial for converting short-term memories into long-term memories during memory consolidation [133 The N-methyl-D-asparte acid receptor (NMDAR) is
essential for mediating this memory formation well as spatial representation. NMDAR serves as
switch for such plasticity and long-term memory formation [ Hippocampal place cells [7] and
entorhinal cortex grid cells [ are known to be central for spatial navigation in animals. NMDAR has been highlighted for its importance in place cell representations through hippocampal CAI neurons (9]
3. These discoveries have led to deeper understanding of hippocampal function, inspiring recent efforts to replicate such spatial representation mechanisms deep neural networks [1] 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 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 the whole-cell current voltage (I-V) relationship (Fig. [b) and are modulated by Mg ion blockade at the pore region of the channel. Previous research indicates that the -dee-dden nonlinear dynamics of NMDAR play a key role in synaptic plasticity and memory formation []]
Recently, transformer-based deep learning models have been reported have functions that resemble hippocampal formations [3 Transformers comprise two consecutive modules: self attentttntion
"qqual contribution. Cooresponding authors.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

(a)
(b)
X
Linear
Feod Forward Net
Nonlinear Unit
LayerNorm
Linear
My
Salf Balf-attention
NMDA,
RelU
++00"
GEUU
91U
LayerNorm
Synaptic
Plaslicit
AMIAAR
NUDAR
Input embeddings
Figure 1: (a) Schematic diagram of Mg -gated NMDAR modulating synaptic plasticity (left), current- voltage dynamics (I-V curve; top right) and an NMDAR -insiiede activation function, NMDA, , (x) (bottom right). (b) Transformer architecture and nonlinear activation functions in its feed- forward network (bottom left): ReI.U, Gaussian Error Linear Unit (GELUU, and Sigmoid Linear Unit (SiLU).
layer and a fordard network (FFN; see Fig. TD). The self-attention layer is closely related to 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 spatial navigation task has yet to be elucidated and compared to hippocampal characteristics.
This paper uncovers resemblance between the NMDAR nonlinearity and recently developed activation functions commonly used in FFNs of deep learning models (Fig. [T]. NMDAR functions operate activity-dependent repulsion of Mg ions [5 16 interesting because it supports self- gaffigatin of information flow (Ca2+ ion influx) the post-synaptic region. Similar to this NMDAR activity-dependee neural networks combine input with a self-gating function (i.c., a sigmoidal function that ranges between 0 to 1) that determines information flow. Our observation prompted the following inquiry: Can NMDAR R-like nonlinearity the feed-forward network layer transformers enhance the formation of long-term memory and spatial representation by place cells?
and this phenomenon is particularly
gating mechanism, activation functions in modern
Here we derive a novel NMDAR-like activation function using the NMDAR current-voltagee (I-V) curve. In parallel we design spatial navigation task in 2D grid environment that assesses two memory types well-defined by neuroscience research: working memory and reference memory [1 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- 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 Al models and for understanding similar processes that occur in the brain and AI models.
Methods
2.1 Designing 2D navigation task to test the role of working memory and reference memory
We designed 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. 2 [2]. The agent

---

Working memory Visited place prediction
Reference memory Unvisited piace prediction
(b)
transformer (causal)
T
L
Action 68)
query Figure 2: (a) Sensory observation prediction task in 2D grid, where dotted squares indicate the target position predict given 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 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.
Opsercatio A
ARAFAL
context
receives sequence of previous [Action (a), Observation (x)) pairs input, and its goal is 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 x 2D grids. A random sensory observation chosen from ten letters placed at each position 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 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 working memory error (Fig. (2 left). contrast, when the prediction based on unvisited nodes is incorrect, it counts as reference memory error (Fig. 2 right). See details of the navigation task and definitions Appendix
Fig. S33
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] [2] 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 3 that hypotheses about two separate memory systems are based on- 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 (x1,..., xT) of d-dimensional input embeddings, the self-attention layer calculates the interaction term between each embedding element within context window via the self-attention mechanism. More formally, each input embedding applies two linear layers (W& and W.) to the embeddings to form the key matrix K and value matrix V:
KT
[k__ kj-ct1 ... k, VT
[v -VT vi-ct1 ... v.].
where V; = W, (W, E (Rddds) and k; = x, We (W& E [dxds). Here, e denotes the context length. The key matrix K E (e(++)) is then used to compute the interaction score between an input embedding at step and all the vectors in K via dot products:
s; = KT where q. x.W, (W, E Rdddd).
The normalized values of s; (R(11), called attention values, are calculated via the softmax function; the final output of the self-attention mechanism is weighted sum of the value vectors in
3

---

Table 1: Comparison of common activation functions with NMDAa.
NMDA o,
Name
Equation
Reference
NMDA, 1MIDA
SiLU(()
zo [z)
20
NMDA 1,
GELU(X)
xo (o.70.02
21
NMDA.= N,DA0=1,,8=0012
ReLU(x)
max (0, 2)
22
NMDA
Swish(x)
xo(xx)
23
V E [R(++)xxde with the attention values:
':K' Vlg
sof tmax
V.
(3)
After this update, Yt E Rd is updated by another linear transformation W. E Rdxxd, 2e W The output 2, is added to the Xx: Ze + X₁ 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
FFN(X,) o(x.U)))
where U1 E Rddxx xd and U2 E Rd trainable weight matrices. Sukhbaatar et al. [19 showed that Eq. (3) and Eq. 4 have similar structures except for the following: 1) U, and U2 matrices are fixed over different imput sequences while and 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-gate gatge--gtedddd Mg repulsion at the NMDAR channel's 's pore [15 16 (Fig. a left). Previous work showed this nonlinear [-V relationship to be:
Juom = VPoo,a()
(5)
where V represents an input voltage and Pa,, (V) is voltage-dependent channel open probability that follows the ion blockade model [24]:
Pa.a (V)
+ Qe -V
where a and ( is temperature constant. For further details, see Appendix A.2
is a parameter determined by [Mg2+], K Km is dissociation constant,
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 new NMDAR-inspired inspireed activation function (see details in Appendix A.3) as follows:
NMDAa, = xP (x) + ae 9a
This NMDA (x) incorporates modern activation functions with varying values of temperature constant, 3 (Table [}. To investigate the Mg -gated nonlinear dynamics, we compared various activation functions with NMDA, (x) = rP (1))
2.4 Testing the NMDAR R-ins--nspired activation in navigation tasks
The transformer model is trained using softmax pross-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 closely related the most advanced neuroscience model of the hippocampus [2].

---

We generate the embedding vectors of the sensory observation sequence with word embedding layer, but the embedding vectors of the action sequence are generated using an RNN; 0x41 = tanh (e, W.). where 6, is a recurrent positional embedding at step 4, and Wa is the action-dependent trainable weight matrix. The input given by {[e] xi]; [e2, x2)...... (t; xz]}: where x denotes the embedding vector of sensory observation x; the initial recurrent positional embedding e1 is sampled from normal distribution, and we mask the last observation Z₁.
In our experiment, FFN in the transformer model consists of two linear layers (see Fig. Tp and Eq. 4 with the NMDAR inspired activation function NMDAa- We use TransformerXL [25] with an extended memory length of 32 and segment length of 32 so that the context length e is 64 and working memory error is measured when the node to predict its sensory observation is in the context
window (see Fig. (2); i.e., node that the agent has not been visited in the last 64 steps is treated
as an unvisited node. Note that our model 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 concatenated vector [e, x] of the word embedding x (dimension of 256), and the recurrent positional embedding 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 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) x 2,048 (total number of steps in a trial) x 200 (number of epochs) (see Fig. [S3 in Appendix A.44. The average number of unvisited nodes in a single trial was 561. The Py PyTorch code for reproducing our results is available at https://// thub. com/kdkyuu bransformer ormer www._morryccooooooeeoeee
3 Results
3.1
Working memory and reference memory errors
To measure the impact of nonlincarity in the FNNs, trained transformer models with different values of a in [0, 0.01, 0.05, 0.1, 0.5, 1, 5, 10] and evaluated the working memory and reference memory errors the train maps (i.e., familiar maps) and test maps (i.e., novel maps). The changing a values in a transformer model mimic the changes in NMDAR Mg²+-gating -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. 3h shows that the reference memory error on the training maps rapidly decreased over training trials when a is larger than zero, with a larger improvement for increasing a. The reference memory error on the novel maps. however, was nearly constant at a chance level of 0.9 (= 1//uumber of letters)) for all a (see Fig. 3 top right). Fig. 3 bottom right shows that working memory was active 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. 3 shows that increasing nonlinearity (i.e., a) helped to activate reference memory, and the trend of improvement was consistent, as shown for N = 32, 48, 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- activation function (a- = impairs reference memory formation (Fig. (b) but leaves working memory formation intact (Fig. 3 This means that even though the trainable parameters exist in the self-attention layer, 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 (ro (.700x)). ReLU (max (0, .r),, LeakyRel.U (max(0, x) + 0.01 min(0..))) sigmoid, and tanh. Fig. 3e shows that the newly proposed biologically-inspired function (NMDA,=10) yielded substantially lower reference memory errors relative to other activation functions. A comparison with the second-best berf-performing
5

---

GELU activation, for example, indicated a statistically sigmificant 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 hyperparameter space of the NMIDA-inspird activation function.
Expanding the evaluation, one may consider other memory types like recurrent positional embed- ding 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 eπror while increasing a of NMDA,, (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 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. [56lin 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 77 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 navigation task). Selective impairment of NMDAR in hippocampal CAI disrupts place cell emergence and memory consolidation of forg-term memory 23.99 9
We investigated whether place cells emerge during the formation of long- meng-temmm reference memory in a transformer model, building upon our previous results demonstrating selective impairment in long-term reference memory with NMDA,, modulation. We newly designed metric called the place cell score described in Eq. ST in Appendix [A. examined the spatial representation within the transformer architecture. This metric, ranging from 0 (i.e., homogeneous firing) to I (i.e., firing is specific to the agent's position in the grid) quantifies the firing specificity of each neuron for spatial
(a)
(b)
Reference memory erter rate
Train (familiar)
Test (novel)
Working memory eiror cale Reference memory arrer rate
0.8
N
0.8
32
0.0
0.6
0.6
0.01
64
0.4
04-
0.01
0.1
0.5
0.2
1.0
0.2
5.0
0.0
a= 10.0
0.0
0.0
0.01
0.05
0.1
0.5
1.0
5.0
10.0
(c)
Reference mentryy entor rate
0.8
0.10
0.10-
0.6
0.15
0.4
0.5-
0.10
305-
0.4
0.05
0.2
0.3
0.00
0.00
0.2
x10
0.0
0.1
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.0
No. of train brilss
x10
No. of train trials x10"
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 eтrors were evaluated on training maps over different values of in NMDA,, and N. (c) Reference memory error comparison between NMDA, = 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
vs other AFs
0.5
10.0
Place fwids lend foroeer
0.5-
0.3
denalty
0.5
0.-
1.0
tanh
0.-
10.2
0-
0.00
0.
05
63
04
0G
Plase coll soore
Avcrage plaze ool boord
0.0-
0.0
0.
Self-attention layer
score
(d)
4
Plere fedoa veluitttion
0.01
0.775
0.0
0.8-
0.5
0.5
Referance
0.44
1.0
10.
02-
10.0
0.2
10-
05
0.4
02
65
0.4
0.
Place scoe
Avarage place soors
(h)
Top-16 Place Fields
Feed forward layer (a 10)
Feed forward layer (a= 1)
Feed ferddffowwrrd layer (a
O
Figure 4: Reference memory-related place cells selectively emerge in FFN but not in the self-attention
layer over increasing a. (a. b) Example rate maps with place scores in FFN and self-attention layers at a = 10; from top left (high) to bottom right (low); color bar indicates the firing rate between Pmmx and Pmin- (c-d) Place cell score distributions with varying a in FFN (c) and self 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, RU, LeakyRel.U, tanh, and sigmoid. All results evaluated using training maps. (h) Rate maps of neurons with top- 16 place cell scores in the FFNs with varying values of a; = 10 (left), (middle), and (right).
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 accumulated the activation values for each
environment, which had an 11 x 11 grid structure, 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, shown in Fig/4h and F))
Fig. fa op show the rate maps of neurons with place cell scores the FFN and ceff-attention layers, respectively. For the self-attention layer, the total number of neurons in the softmax layer was 65 (context length + masked sensory observation) x 8 (number of heads) x 2 (number of layers). The total number of neurons in the FFN layer was set at 2,048 (number of neurons) x 2 (number of layers). As can be seen, our metric accurately represents location specificity. Fig. 4 and q show the distribution of place cell scores in the two layers with different values of a. When the value is increased, the place cell score distribution found in the FFN becomes positively shifted (see Fig. T rate map examples for = 0, 1.0, and 10.0), whereas the place cell score distribution in the self- attenttton layers remains.
Fig. 4e and 4 show a relationship between the average place cell score and the reference memory error for each a. The average place cell scores in the self-attention layers show no correlation with reference memory errors, while neurons in the FFN exhibit clear correlation. These results imply that reference memory formation and place cell emergence can be enhanced by NMDAR- like nonlinearity in the FFN.
7

---

Fig. compares the place cell representations of our NMDA (a O, 10) with the representations in FFNS with the activation functions used in Bi. indicating that the case of NMDA,,10 outperforms other activation functions, in both reference memory formation and place cell representation. Our finding that large a (the [Mg2+] component) enhances reference memory in line with the biological observation that increasing the [Mg in the brain enhances long-term memory formation [
We also explored the possibility of extending the hyper-parametee space by contrasting the NMDA- inspired activation function with an alternative activation function based on a. LeakyReLU was considered for this experiment because of its adjustable negative slope. which was re-defined as (max(0, (0, 2) + min x)). The results reported in Fig. S10 of Appendix A.7 show that LeakyReLU activation, for all tested a values, produced lower place cell scores than the best-performin inspired function (i.e.. NMDA,=10)) Our findings suggest that NMDA- nonlinearity, which is already known be critical in biological systems, could be used to enhance transformer models.
NMDA-
Building upon these results, our work connects memory consolidation theory from neuroscience to transformer-basee processing of short-term and long-term memory. Unlike previous studies that directly provide 2D geometric information to models [0 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 AA.I0 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 the brain.
4 Related works
The current study is inspired by recent studies that connect neuroscience and Al models. One such seminal work is by Whittington et al. [1], which showed a relationship between the self-attention layer and a recent hippocampal model called the et al. [4). Our work expands the literature by focusing on FFNs in the transformer and making a connection to the emergence of place cells.
Machine (TEM; Whittington
TEM is a neuroscience-base 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 xd, Every outer product of key and value vector k, v, at each step i is simply stored in M via the Hebbian update rule. M is initialized to zero matrix at the beginning of the task and adds every outer product at each time step:
M Σ k Vv. -ov.aKTT
(8)
where a is a weighting factor. In the memory retrieving phase with the query vector q. TEM uses an attractor network:
q = aqK V.
Whittington et al. 12 found that the memory retrieving process in TEM has 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 address high computational costs when context length e is very long [27 28). The subsequent model of TEM, called TEM- [[2]. replaces the attractor network (Eq. with self-attention mechanism (Eq. This study demonstrated that TEM-t learns significantly faster than TEM.
TEM-t and TEM do not have fixed context length c; therefore, these models store all information
before step t. i.c., c = t. The computational cost of the self-attention layer in TEM-t is O(P). and retaining all previous information is too expensive from both biological and computational standpoin |
For TEM, the Hebbian update rule has no quadratic computational cost and can add all previous information in fixed number of synapses dt: however, the memory capacity of the Hebbian matrix M is (dg)) and the speed of memory retrieval is substantially slower than the
'Due to the computational cost, TEM-t does not store all historical data but selectively chooses which data to store in K and V.
8

---

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 fixed context length and 2) context-independen fixed matrices U1 and U2 (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 s performance. 2) We proposed a new brain-inspired activation function, NMDAa. which relates to modern nonlinear activation functions and allows for the analysis of the effect of a 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 more detailed analysis of the model's 's performance.
5 Discussion
While extensive efforts have been directed toward finding the optimal nonlinear activation function for improving modern deep neural network models [21H23], 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 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 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 CAI regions of the hippocampus to be involved in the initial formation of new memories, specifically in pattern completion [33). while CA1 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- attentttntion layer, while the CAI region may function similarly to FFN. However, further research 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
32] In the hippocampus, the CA3 region is thought
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 transformer block functioning as a key-value memory 135 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 33..
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 Mg -gatin of NMDA receptors modulates the formation of long-term memories [5] d. 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 CAI NMDAR inhibition does not disrupt
working memory 41 but impairs the long-term memory formation (2). 2) changing NMDAR -gatig (changing a in this work) enhances or disrupts long-term memory formation 53 6). 3) NMDAR is required for long-tern stabilization of newly forming place fields (9 3. These similarities between hippocampal memory consolidation and our results suggest that the transformer an effective memory consolidation model.
Our research points to exciting future directions. The current study examined whet-where memory using sensory observation task in static environment. However, our real-world environment

---

changes constantly and provides new inputs over time. Modern deep learning systems are generally incapable of adapting dynamic environment or reordering sensory imputs. We intend to explore what-where-when memory, called episodic memory, in transformers and other deep models.
Acknowledgments and Disclosure of Funding
This work 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.
References
[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 sccepto--deennden synaptic plasticity in spatial memory. Cell, 87(7):1327-1338 1996.
[3] Clifford Kentros, Eric Hargreaves, Robert D Hawkins, Eric Kandel, Matthew Shapiro, and Robert V Muller. Abolition of long-term stability of new hippocampal place cell maps by NMDA receptor blockade. Science, 28066372)2121-22126, 1998.
[4] Tim VP Bliss and Graham L Collingridge. 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-Gat Gag Zhao, Min Zhuo, al. Enhancement of learning and memory by elevating brain magnesium. Neuron, 6552):165-177, 2010.
[6] Tomoyuki Miyashita, Yoshiaki Oda, Junjiro Horiuchi, Jerry CP Yin, Takako Morimoto, and Minoru Saitoe. Mg2+ block of drosophila NMDA receptors is required for long-term memory formation and CREB-dependens gene expression. Neuron, 74(5):887-88 2012.
[7] John O'eeefe and Jonathan Dostrovsky. The hippocampus as 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. Mi- crostructure of spatial map in the entorhinal cortex. Nature, 436(7052):801-888 8066, 2005.
[9] Thomas J McHugh, Kenneth I Blum, Joe Z Tsien, Susumu Tonegawa, and Matthew A Wilson. Impaired hippocampal representation space cal-specific nmdarl knockout mice, Cell, 87 (7):1339-134) 1349, 1996.
[10] Andrea Banino, Caswell Barry. Benigno Uria, Charles Blundell, Timothy Lillicrap, Piotr Mirowski, Alexander Pritzel, Martin 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::/openreriee..nt wttp://wwwww......oooooooddd//i.
[12] James C. R. Whittington, Joseph Warren, and Tim E.J. Behrens. Relating transformers models and neural representations the hippocampal formation. In International Conference on Learn- ing Representations, 2022. URL httpp::/ a
[13] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
10

---

[14] James CR Whittington, Timothy H Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy EJ Behrens. The tolman- cemaan-ichennaum machine: unifying space and relational memory through generalization in the hippocampal formation. Cell, 183(5):1249- 1263, 2020.
[15] Nowak, P Bregestovski, Ascher, A Herbet, and A Prochiantz. Magnesium gates glutamate- activated channels in mouse central neurones. Nature, 307(5950):462 465, 1984.
[16] Mark L. Mayer, Gary L. Westbrook, and Peter B Guthrie. Voltage-dependent block by mg2+ of nmda responses in spinal cord neurones. Nature, 309(5965):261-222 1984.
[17] David S Olton, Christine Collison, and Mary Ann Werz. Spatial memory and radial arm maze performance of rats. Learning and motivation, 8(3):289-314, 1977.
[18] David S Olton, James Becker, and Gail E Handelmann. Hippocampus, space, and memory. Behavioral and Brain sciences, 2(3):313-322, 1979.
[19] Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Herve Jegou, and Armand Joulin. Augmenting self-attention with persistent memory, 2019. [tt.httss::///at org/ abs/
1907.01477
[20] Stefan Elfwing, Eiji Uchibe, and Kenji Doya. function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018.
linear units for neural network
[21] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint a60iv:1606.0 08415, 2016.
[22] Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In leml, 2010.
[23] Prajit Ramachandran, Barret Zoph, and Quoe V Le. Searching for activation functions. arXiv preprint arXiv: arXiv: 1710.05911, 2017.
[24] Ann M Woodhull. Ionic blockage of sodium channels in nerve. The Journal of general physiology, 616666887- 1973.
[25] Zihang Dai, Zhilin Yang. Yiming Yang, Jaime G. Carbonell, Quoc Viet Le, and Ruslan Salakhut- dinov. Transformer-xl: Attentive language models beyond a fixed-length context. In Association for Computational Linguistics, pages 2978-2988, 2019.
[26] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, In Interna- tional Conference on Learning Representations, 2015. http://aaari org.abg/////ooooooooooo.oo.....t /abs/1112 1412 6980
[27] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5156-5165. PMLR, 13-18 Jul 2020. URL https::/prooceednng:
wttp://reeeeiiinnn..eeeeeeeeeeee mlr. wtt.a//pppp//////tattttoooooooop html
[28] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers secretly fast weight programmers. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th Inter- national Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9355-9366 PMLR, 18-24 Jul 2021. URL. httpss:/pproeedingg mlr
press/v139/schi ag2la. html
[29] Mete Demircigil, Judith Heusel, Matthias Löwe, Sven Upgang, and Franck Vermet. On model of associative memory with huge storage capacity. Journal of Statistical Physics, 168(2): 288-299, may 2017.
[30] Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gruber, Markus Holzleitner, Thomas Adler, David Kreil, Michael Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need. In International Conference on Learning Representations, 2021. ///p:ss///ppeeviiee iaaieerrrviivenee//ruunn net/forum?
id=t189Rnzlicd

---

[31] Dmitry Krotov and John J. Hopfield. Large associative memory problem in neurobiology and machine learning. International Conference on Learning Representations. 2021. URL htttps://openreviexx net/forummiddX4yy: 1000--X
[32] Can Dong, Antoine D Madar, and Mark EJ Sheffield. Distinct place cell dynamics in cal 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 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. Empirical Methods 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 SXXi::2104008666 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 Jaszezur, Aakanksha Chowdhery, Afroz Mohiuddin, Lukasz Kaiser, Wojeiech Gajew- ski, 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 EIShowk, 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- Doddd, 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///tenssommerr-ciiiiitbbbbp pww./0020sss/////eexee.moo.
[40] James L McGaugh. Memory-a century of consolidation. Science, 287(5451):248-22 2518 2000.
[41] Inah Lee and Raymond P Kesner. Differential contribution of nmda receptors in hippocampal subregions spatial working memory. Nature neuroscience, 5(2):162-16) 1622216661111 2002.
12

---

A Appendix
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 certain location [ST], and provide a sparse population code for sell-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 create and maintain spatial representations, which are essential for navigation and memory S4
Relevance to machine learning Accurately identifying place cells and understanding their prop- erties 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 navigation tasks with deep models [S5 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 's performance declined, indicating the role of grid cells in place cell formation and spatial memory. An interesting hypothesis neuroscience is that the interaction between grid cells, path integration, and place cells forms critical loop [S9] [ST0] 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
(a)
(b)
(c)
x#
#4
F
96
#4
U
O
R
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 G. Colors represent the firing rate of nodes. (b) Convert G into directed graph by comparing the firing rate o 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. [S]] is scalar value that quantifies each neuron's 's firing specificity for grid location. We record every neuron's 's activation value at every step during the random walk process. Since the map has an 11 x 11 grid structure, we accumulate the activation value for each grid point as the agent random walks, creating 2D array. This array called the rate map or the place field (i.e., 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 K x K 2D grid environment as graph G (V. Æ) and firing rate (i.e.. cumulative activation value at node divided by the length of evaluation trial) of node i E V as Pi. we defined a maximally firing node as fmax and its firing rate as Pmax. E are directed edges that connect high to low firing nodes in G. From G, we ran a depth-first-search from source node max to build a sub-graph G (V,E) which we call all connected components SI Given G and G, the place cell score is defined as follows:
Place cell score = Eiev Pi 'iev
(SI)
is a discount factor and V is V without node imax and leaf nodes. To measure place cell score, we recorded the firing rate P: of neurons over a random walking trajectory with 10 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 when it fires homogeneously across all nodes.
where 7 = -
Since 2D-grid spaces cannot be directly interpreted as Euclidean spaces, it is not appropriate to apply space-related functions, such Gaussian kernel density estimation, to the place fields. Place cell evaluation methods in previous research can be broadly categorized into four types [SII]:
1. Peak method: Classifies cells based on the average firing rate in one location being higher than 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 node with a maximal firing rate, i Starting from this node, we find all connected components as 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 imax. DFS can be considered an inverted watershed algorithm that fills single pool containing the global minimum. The advantage of our place cell metric 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 PlaceCellScorepppace field)
Input place field (K' X K 2D array)
Output: place score
G == 2D grid graph (K x K)
== empty directed graph (K x K) for edge (node; nodee) in G do firing rate P; > firing rate Pj then G add edge (node; node;) end end
Find node, of firing rate Pmaar for node, do
if node, is not descendant of node; found with DFS(k) then delete node, from G end end
conn. components sum of all nodes' firing rates in total components = sum of all nodes' firing rates in G conn. components
place score = total components
I is discount factor, determined by connected component size
return place score
Algorithm Pseudo code for calculating place cell score metric
The place field Algorithm|] 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 ag-
2. Every time the agent visits node again, value a; is added cumulatively to the recorded value; A;+ ag such that A; is the cumulative activation value at node i. We assume the initial value of A, is zero.
3. After the random walk done, A, divided by the length of the random walk trajectory the firing rate at node i of the neuron (place field E K xxXK
In our place cell evaluation experiment, the length of the random walk is 10 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 Mg and NMDAR following previous literature (S12- At the molecular level, one Mg ion binds to one NMDAR receptor when opening the NMDAR channel. Thus, the chemical equation of binding reaction between Mg2+ ion and NMDAR receptor, R, can be described as
Mg + R
2g² R.
(S2)
From this chemical equation, the equilibrium constant K is given by
Mg R
K
(S3)
Mg
[
Thus, dissociation constant Kp, which correspond to Mg dissociation from NMDAR, follows
Mg
2+
R
KD=K-1=
(S4)
Mg
R
in which [R] and [M are the free and 2- bound NMDARs respectively. The fraction of opened NMDAR channels (number of unbound NMDAR over a number of total NMDAR) at equilibrium follows,

---

R
R + Mg R 1
(S5)
1+ MMg2* |/K
Experimentally, the voltage-dependent dynamics of KD has been described following equation by Ascher and Nowak [S13]
KD K @v
(S6)
where, V is membrane voltage, 3 is temperature constant and KMg-- is dissociation constant at 0. If Eq. [S6 substituted into Eq. S5 voltage-dependent open fraction NMDAR can be expressed as follows:
P(V)
+
Mg2+] ev Ma
9V
(S7)
1 + ae 8V BV
in which a MMM+++, the parameter determined by the [Mg2+ Given the voltage-dependent open fraction of NMDAR, PVV), and NMDAR's maximal conductance, Jmax then voltage-dependent NMDAR conductance g (V) can be described as
g(V) - 9maxp (VV)
(S8)
Given g(V). and driving force, V- Vr, and current 1, they have relationship of / = (V (V), which V, 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, = 0 S16 S18 I (V).. Then, the normalized NMDAR current Iporm /mamax follows:
morm Vp(V)
(S9)
From Eq. S9 and previous clectrophysiologicall experimental results S19), reconstruct IV curve in (Fig. 1. right top).
A.3 NMDAR-inspired nonlinear activation function
Here, we propose an AMDAR-inspires nonlinear activation function from the nonlinear dynamics of the NMDAR-IV curve. If we consider the nonlinear IV curve of NMDAR (S9 as nonlinear mapping function, Φ, the membrane voltage, V, can be viewed as input, æ, and normalized NMDAR current, / Forms as an output, dxxx). Therefore, we can rewrite the nonlinear mapping function s follows
(xz) = xp(x).
(S10)
We define the NMDAR-inspired activation function as a nonlinear mapping function, (MDA(2)==(()). By substituting Eq. SS into Eq. S10 we show the generalized expression of (2DA(z) equation with and 3 parameters follows:
NMDA (x) =IP...(>))
x
(S11)
ae Ba
Given = and 3 = 1, p(x) is identical to the sigmoid function, (ι) 1/(1 + e-)). This
particular case and 3 leads to za wa(xx) Sigmoid Linear Unit (SiLU) activation function (S20] In the case of a 1 and [ = 1.702, ro (1. 702x)) corresponds to the GELU activation function, which popular transformer models [S21 S22 [22 Ramachandran al. [823] introduced the swish
activation function, 28xx), which is a generalized form of GELU and SiLU. They demonstrated

---

NMDAs=26
14
NMDKe-13
GELU
SwishtS
1.2
Swiskg8
RELU
RELU
1.0
Figure S2: Comparison of common activation functions (left) and their derivatives (right) with NMDAa.9.
that when B + oc, the activation function resembles RELU. We summarized these four activation functions by comparing them with our NMDAa, (2) ()) in Table and Fig. S2
In contrast the extensive research 3 in NMDA...( a, the Mg2+-gating component, is explored. For this reason, we focused on the parameter a over 8, and investigated NMDA, 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 NMDA(z) is inspired by the physical Mg²--gating mechanism that occurs at the real biological synapses. These shared mechanisms of self gelf-gating artificial models and biological observations raise the interesting possibility that NMDAR is 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 relatively recent experience can be defined as short-term working memory (STWM), and the memory of relatively old experience can be defined as long-lerm 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 the novel map will establish baseline of chance-level accuracy; above this baseline, the formation of long-term memory can be observed (Fig. S3).
5

---

Familiar maps
Novel maps
Agent
Trajectory
Within trial
Trial1
Training phase
Trial2
Trial3
Epoch1
Query (Batch size)
Across trial
Test phase
Unwisited
Unvisited
Unvisited
Visited
Unwisited
Unvisited
Unwisited
Unvisited
Unwisited
Unvisited
Visited
Unvisited
Epoch3 Epoch2
Training phase
Test phase
Training phase
Test phase
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 map from training maps (familiar maps), the initial position of the agent is random, and the agent does random walk.

---

A.5
Non-recurrent positional embeddings and prediction errors on the node visited for the first time
Train (familarr
Test (novel)
Reference memory error
0.0
0.01
0.8
0.05
(b)
0.1
0.4
1.0
0.8
5.0
Reference memory
0.2
Working memory etror
0.4
0.2
0.3
0.0
0.01
0.05
0.1
0.5
1.0
5.0
10.0
0.1
0.0
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
No. of train trall
No. of frain trials
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 32 where N is the number of training maps. (b) Reference memory errors evaluated on training maps over different values of in NMDAa for = 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 e, with the action embedding A(ax). where A is the embedding layer and ag is the action at step t. Compared to Fig. [S4p, the result demonstrates significant increase working memory error and reference memory error (Fig. 3 vs. Fig. (54).
Nonetheless, the model's behavior is comparable to the trend of decreasing reference memory error with increasing a of NMDAa (see Fig. S4D).
Moreover, we found that place cells do emerge the feedforward layers for this standard positional embedding method (see Fig. S5 It is worth noting that these place cell scores noticeably lower compared to those obtained through the recurrent positional embedding. Despite the hindered performance in working memory and reference memory, strong correlation persists between the place cell scores FFNs and the reference memory errors (see Fig. S5c and (Sdd).
Prediction errors 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 trial). As shown in Fig. (66) the prediction node do not differ from the reference memory error results.
results for the first visited
These results strongly support the conclusions that (1) while the path- intbernteerrted information from recurrent positional embedding important for learning the spatial structure of the map, this informa- tion is not used in predicting the unvisited node, and (2) the reference memory is used for predicting an unvisited node on familiar map.

---

14.
12-
Reference memory eroir
0.0
P=
Probablitty denaty
0.05
10-
r=0.51
0.1
10
06-
0.5
10.0
0.4
0.0
0.0
01
02
Place cell score
0.3
0.4
0.
06
0.10
0.15
0.20
0.25
0.30
0.35
Average place score
(c)
0.9
Reference memory erorr
0.8
12-
p< 0.001 r=-0.436
0.01
Probabiity danakt
0.7
2.0
0.1
0.01
0.5
0.05
5.0
06-
044
0.0
0.1
02
0.3
06
0.15
Place call score
Average place cell score
Figure S5: Place cell analysis for non-recurrent positional embeddings for N 32. Place cell score distributions with varying a in self-attention layers (a) and of average place cell scores and reference memory errors for self-attention layers (b) and feed-forward layers (d). r and denote Spearman's 's rank correlation coefficient and significance score, respectively.
layers (c). Scatter plot
Airst visited node prediction cton error
N
48
0.6
64
0.4
0.0
0.01
0.05
0.1
0.5
1.0
5.0
10.0
First visited node prediction error evaluated on training maps over different values of a in NMDA,, for N 32, 48, and 64. Error bars and shaded areas represent the standard deviation of emror from three independently trained models.
Figure S6:
A.6 Statistical significance between the NMDA and GELU activation functions & learning rate effect: it an under-trained effect?
In this section, we address concerns about the NMDA style activation with large 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 claim, conducted

---

statistical across various activation functions. Our findings revealed significant difference between NMDA,, and other activation functions, indicating that the NMDA style activation beneficial for the emergence of place cells and the improvement of learning speed (see Table S7
(4)
(b)
(6)
25
20
08
-10
100
00
(4)
N--22
(4)
Aghood
agheol
syroed
kmh
122
100000
PNLU
Hireed
PMU
Agrodd
Figure S7: Statistical test result across various activation functions with different numbers maps (N) related to Figure 3. (a-c) Mann-Whitne Whitney test across various of NMDA (d-f) Mann- Mann-Whitney test across various activation functions. The sample size (or trials) used for tests ranges from to 11. Colors represent p-values and asterisks indicates statistical significances: 0.01<p<0.1; ** 0.001<p<0.01; *** p<0.001. a = for NMDA and a 0.01 for LeakyReLU.
Gradient around r 0 and learning speed The performance gain large alpha values might be due to an increased gradient around æ 0. However, the gradient around x = 0 at =1 (green) is actually larger than not due to an increased effective learning speed around O. Instead, it may be attributed to the activation function's 's ability better capture nonlinear relationships.
(blue) in Fig. S2 Therefore, the performance gain at large alpha values is
Learning rate effeec 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. [88 larger learning rate at low alpha values is 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. (99 shows that higher place cell scores correspond to lower reference memory errors. Note that we employed linear decay learning rate schedule (from the starting learning rate to 0) and the starting learning rate used in manuscript is 0.0001.
In conclusion, our additional experiments and analyses demonstrate that the NMDA style activation with large plays critical role in the emergence of place cells and the improvement of reference memory, and it not trivial effect an under-trained effect.

---

Reference memory effor
0.8
0.1
1.0
0.6
10.0
0.4-
0.0001
0.0002
0.0003
0.0005
0.001
Learning rate
Figure S8: Familiar map reference memory error results with five different learning rates, where 1.01.1.0, 10.0). for total of 64 training maps (N = 64). The error bars represent the standard deviation of errors from three independently trained models.
Forward Layer
Forward Layer
Feod Forword Layer o
18
9.8
00001
00001
07-
08003
08885
Nelesnso mamor
00-
occes
00063
326
8.30
659
Aetrage place coll Store
Aetrage place cell score
Average place coll score
Feed-Formord Layer 2-0.1
Feed-formard
Feed-foroard
22-
20-
20-
2.-
borming rote
Rarning
haming cite
E001
81 02
0.4
0.2
0.5
0.4
05
0.1
0.4
05
Place acore
Figure S9: (top) Scatter plot between reference memory eтror and average place score; (bottom) Distribution of place cell scores with five different learning rates, where {0.11 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, investigated the consequence of changing nonlincarity with activation function other than NMDA,, We chose LeakyReLU (max(0, x)) + a min(O, (0, z)) activation function to compare with NMDA, Regarding Lea,yRel.U, function similar to a = 0 of NMDA,- Compared to NMDA,=10. LeakyRel.U exhibits a lower average place score in the allowed range of a, indicating that NMDA, performs better in terms of robust place cell emergence (see Fig. S10)
of Leaky also leads to a fully linear activation
10

---

LeakyReLU
Reference memory error
0.2
0.1
0.01
0.02
0.05
0.1
0.2
LRELU 0.05
L --1:1
Reforence monsory crorr
0.8-
LRELU a 0.2
LRELU a 0.1
LRELU a 0.05
0.3-
LRELU a 0.01
0.02
LRELU e 0.02
UHELU
Probodiitty
0.2
C4+
0.1-
0.2
80
0.0
0.4 Paice call scor
0.2
0.8
0.2
0.3
Average place oel soore
Figure S10: Evaluation of reference memory error in LeakyReLU (LRELU) while modulating a<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 NMDA,, activation function in the feed-forward layer of the transformer and role memory formation and place cell representation. We show that modulating a corresponds to change in extracellular IMg*+], J. by deriving the nonlinear activation function from the real NMDAR nonlinear IV curve. The reconstructed real nonlinear IV curve is shown in Fig. I (top right pancl).
The modulation of 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. 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 short-term memory systems via attention, and some are transferred long-ten term-termmm memory system while others quickly disappear.
In neuroscience, the transfer of short-term memory into long-term system is called memory consolidation [S28). Animal studies have demonstrated that the CAI region of the hippocampus essential for memory consolidation SS29 S30 In hippocampal CA1, the postsynaptic NMDA receptor mediates synaptic plasticity, and the selective perturbation of these receptors leads to impairment long-term memory formation [S3 S30 Later research revealed that Mg -gating of NMDA receptors modulates the formation of long-term memory [S32 S33 These observations imply that the nonlinear dynamics of NMDA receptors CA1 are critical for consolidating short-term memory into long-term memory [S1T
On the basis of a previous link between the hippocampus and the transformer model, we hypothesize that the latter is model for memory consolidation. Given the resemblance of the GELU nonlinear activation function and CAI NMDAR nonlinear IV curve, assumed that the GELU activation function serves as key component that links short-term working memory and long-term reference memory. Our experimental results indicate that that the formation of long-term reference memory is impaired when the activation function is completely linear (corresponding to no Mg²+). contrast, increasing a (which corresponds an increase in Mg level) revealed that for long-term reference memory, our model outperforms other activation functions (e.g., RELU, GELU, IRELU, Sigmoid, Tanh). Based on these similarities between hippocampal memory consolidation and our results, propose that transformer is effective memory consolidation model.
11

---

Transformer
Sensory
Self-attention layer
Feed-forward
NMDAa
layer
Action
Experience
Short-term memory
Consolidation (Norrincar gatin))
Long-term memory (Place call;
Brain
LEC
CA3
NMDAR
CA1
MEC
Figure $11: Schematic illustration of the transformer memory consolidation model of hippocam- pus via NMDAR nonlinearity. Top. transformer with NMDA, nonlinear activation Canction.Botton 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 NMDA,, found that modulating the a affects the emergences place cells the feed-forward layer and conclude that there significant correlation between place cell score and long-term reference memory formation. Our results align with previous biological findings that perturbation of CAI 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-tern term spatial memories. Altogether, our findings suggest the exciting possibility that the nonlinear IV curve of NMDAR in the hippocampal CAI 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 NMDA, 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 NMDA activation function and conducted experiments on the Open- WebText dataset [S37]. We employed the GPT2 training Py PoTcc code from carpathy/nanoooe
GitHub repository [S38 The application of NMDA, showed a slight increase in performance,
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 the CIFAR-100 and TinylmageNet datasets. Three trials conducted for each condition, and an increasing, though statistically insignificant, tendency in performance was observed. The t test accuracies for these datasets are provided in TT
Table S1: Top-1 test accuracies with ViT model on CIFAR-100 and Tiny datasets.
Dataset
GELU
NMDA,
NMDA,,
CIFAR100
69.92 0.34
70.27 0.47
49.91 4 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 range of hyperpa- rameters for either NMDA or the GELUU-basdd model. The optimal hyperparameters for NMDA might differ from those for the GELU-based model.
12

---

val/loss
ffn_act_ftn. gelu, alpha:
ffn_act_ftm nmda, alpha: 10
3.3
(2,856, 2.856) ffn: alpha:
295880:
(2.547, 2.847) rfn_act_fene neda, alpha:
itel
50k
100k
150k
200k
250k
300k
Figure $12: Validation loss curve over training iteration for GPT2 model on the OpenWebText dataset. The blue (red) line represents the GPT2 with GELU (NMDA,=10 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 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 NMDA, activation function, NMDA, (2) = x/(1 xe=
we can rewrite this function as follows:
NMDA = x (e- where a e".
(S12)
In the given expression, an increase in a results in shifting the sigmoid function towards the positive direction by increasing 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 might enhance the sparsity of the NMDA activation.
Empirical Results To validate the aforementioned mathematical intuition, conducted measure- ments 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 follows:
Ix Zz
G
(S13)
2K
where 24 is the i-th neuron's activation value and 5 values the feed-forward layer (K 2048 is the total number of neurons in a feed-forward layer).
Σ
3i/ K is the mean of the activation
The Gini index ranges from 0 to When only few neurons have very high activation values and others are small, the Gini index is close to 1; on the other hand, when most neurons have homogeneous activation values, this value is close to 0 (1 absolute sparsity, 0 = all activations equal).
According Fig. S13a when the value of a in NMDA,, (x) is increased, the Gini index also increases. This increase sigmihes that the population activities become more heterogeneous, leading to heavy- tailed distribution. Furthermore, the Gini index of NMDA with an a value of is greater than that of the GELU activation function (indicated by the dashed line). These findings imply that NMDA, (a) could potentially improve long-term memory formation by promoting sparsity in the activations of the feed-forward layer.
13

---

0.85
Ieyer
043
0.81
0.00
0.75
Myari
0.70
kayw?
patul
0.10
getll
0.65
0.25
0.1
10.0
1MD4u.t5
Swish
GEU
Rack
LeakyRoll
Sigmod
Activation Function
(a)
(b)
Figure S13: (a) Comparison of Gini index for different values of a in NMDA (b) Comparison of Gini index between different nonlinear activation functions.
Previous studies have explored the generalization performance and sparsity in overparameterized models [S40]. Although the mechanism behind the emergence of sparse representation large models is not fully understood, it is worth noting that overparameterized models converge to simpler models than those with dense representations.
The high sparsity properties of the NMDA function may contribute to the increased score of place cells. In the case of Transformers, enforcing neuron activation sparsity in MLPs has been found to improve the interpretability or selectivity of a higher percentage of neurons [S41]] This evidence could potentially explain the higher place cell score observed when the alpha value is increased.
We observed that NMDA with = 10 exhibits the highest Gini index among different activation functions. See Fig. ST3a for comparison of Gini index among activation functions.
References
[S1] John O'Keefe. Place units in the hippocampus of the freely moving rat. Experimental neurology, 51(1):78-109, 1976.
[S2] Matthew A Wilson and Bruce L McNaughton. Dynamics of the hippocampal ensemble code for space. Science, (5 (5)) 1055--0888 1993.
[S3] John O'keefe and Lynn Nadel. Précis of o o'keefe & nadel's the hippocampus as cognitive map. Behavioral and Brain Sciences, 2(4):487-494, 1979.
[S4] Howard Eichenbaum, Paul Dudchenko, Emma Wood, Matthew Shapiro, and Heikki Tanila. The hippocampus, memory, and place cells: is spatial memory or a memory space? Neuron, 23(2):209-22) 226, 1999.
[S5] Christopher Cueva and Xue-Xin Xin Wei. Emergence of grid-like representations by training recurrent neural networks to perform spatial localization. International Conference on Learning Representations, 2018. https: /opprrevie//erriieeee wwt.
[S6] 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): 433, 2018.
[S7] James CR Whittington, Timothy H Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy EJ Behrens. The tolman-eichenbuum machine: unifying space and relational memory through generalization in the hippocampal formation. Cell, 183(5): 1263, 2020.
[S8] James C. R. Whittington, Joseph Warren, and Tim E.J. Behrens. Relating transformers to models and neural representations of the hippocampal formation. In International Confer- ence on Learning Representations, 2022. URI. B8DVo9B1YE0
14

---

[S9] César Rennó-Costa and Adriano BL Tort. Place and grid cells in loop: implications for memory function and spatial coding. Journal of Neuroscience, 37(34):8062-8880 00008006666666 2017.
[S10] NM Van Strien, NLM Cappaert, and MP Witter. The anatomy of memory: an interactive overview of the peatippocamaal-hipocamaae network. Nature reviews neuroscience, 10(4); 272-282, 2009.
[S11] Dori M Grijseels, Kira Shaw, Caswell Barry, and Catherine N Hall. Choice of method of place cell classification determines the population of cells identified. PLoS Computational Biology, 17(7):e 1008835, 2021.
[S12] Ann M Woodhull. Ionic blockage of sodium channels in nerve. The Journal of general physiology, 61 6166)6887700, 1973.
[S13] Craig Jahr and Charles F Stevens. Voltage dependence of nmda-activated macroscopic conductances predicted by single-channel kinetics. Journal of Neuroscience, 10(9):3178-318) 1990.
[S14] MICHAEL Perouansky and YOEL Yaari. Kinetic properties of NMDA receptor-mediated synaptic currents in rat hippocampal pyramidal cells versus interneurones. The Journal of Physiology, 465 (36(11):22) 1):22)22244 1993.
[S15] Philippe Ascher and Linda Nowak. The role of divalent cations in the N-methyl-D-uspataen responses of mouse central neurones in culture. The Journal of physiology. 399(1):247-26) 1988.
[S16] ML Mayer and GL Westbrook. Permeation and block of N-methyl-D-aspariic acid receptor channels by divalent cations in mouse cultured central neurones. The Journal of physiology. 394(1):501-527, 1987.
[S17] Tomomi Ichinose, Shun Yu, Xue Qing Wang, and Shan Ping Yu. Ca2+-independent, but voltage-and activity-dependent regulation of the NMDA receptor outward k+ current in mouse cortical neurons. The Journal of physiology, 551(2):403-177, 2003.
[S18] G Kohr, Yves De Koninck, and Istvan Mody. Properties of nmda receptor channels in neurons acutely isolated from epileptic (kindled) rats. Journal of Neuroscience, 13(9):3612-3627, 1993.
[S19] Eilon D Kirson, Claudia Schirra, Arthur Konnerth, and Yoel Yaari. Early postnatal switch in magnesium sensitivity of nmda receptors in rat cal pyramidal cells. The Journal of Physiology, 5P1 1):99, 1999.
[S20] Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Siggoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018
[S21] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self celf-supervised learning of language representations. arXiv preprint arXiv: 11942.9..999........ 2019.
[S22] Yinhan Liu, Myle Ott, Naman Goyal. Jingfei Du. Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1692, 2019.
[S23] Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arXiv preprint arXivv1710.0550 2017.
[S24] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, (8):1735-1780, 1997.
[S25] David S Olton, Christine Collison, and Mary Ann Werz. Spatial memory and radial arm maze performance of rats. Learning and motivation, 8(3):289-311 1977.
[S26] David Olton, James T Becker, and Gail E Handelmann. Hippocampus, space, and memory. Behavioral and Brain sciences, 233)2313-32, 1979.
15

---

[S27] Richard C Atkinson and Richard M Shiffrin. Human memory: A proposed system and its control processes. In Psychology of learning and motivation, volume 2, pages 89-195, Elsevier, 1968.
[S28] James L McGaugh. Memory-a century of consolidation. Science, 287 2000.
[S29] Eiji Shimizu, Ya-Ping Tang, Claire Rampon, and Joe Z. Tsien. Nmda receptor-dependent synaptic reinforcement as a crucial process for memory consolidation. Science, 290(5494)) 1170-1174, 2000.
[S30] Miguel Remondes and Erin M Schuman. Role for a cortical input to hippocampal area cal in the consolidation of a long-term memory. Nature, 431((009):699--73,, 2004.
[S31] Joe Z Tsien, Patricio T Huerta, and Susumu Tonegawa. The essential role of hippocampal cal nmda receptor-dependent synaptic plasticity in spatial memory. Cell,
1996.
[S32] Inna Slutsky, Nashat Abumaria, Long-Jun Wu, Chao Huang, Ling Zhang, Bo Li, Xiang Zhao, Arvind Govindarajan, Ming-Gas Zhao, Min Zhuo, al. Enhancement of learning and memory by elevating brain magnesium. Neuron, 65(2):165-177 2010.
[S33] Tomoyuki Miyashita, Yoshiaki Oda, Junjiro Horiuchi, Jerry CP Yin, Takako Morimoto, and Minoru Saitoe. Mg2+ block of drosophila NMDA receptors is required for long-term memory formation and CREB-dependent gene expression. Neuron, 74(5):887-898, 2012.
[S34] Thomas J McHugh, Kenneth I Blum, Joe Z Tsien, Susumu Tonegawa, and Matthew A Wilson. Impaired hippocampal representation of space cal esl-speccfic nmdarl knockout mice. Cell, 87 1111339-1444 1996.
[S35] Clifford Kentros, Eric Hargreaves, Robert D) Hawkins, Eric R Kandel, Matthew Shapiro, and Robert V Muller. Abolition of long-term stability new hippocampal place cell maps by 1998.
NMDA receptor blockade. Science, 280(5372)
[S36] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenA/ blog, 1(8):9, 2019.
[S37] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027 2020.
[S38] Andrej Karpathy. The simplest, fastest repository for training/finetuning medium-sizer gpts. hhttps: gi thub. httm//aarpathyynnoooPP 2023.
[S39] Michelle Miller, Yue Chung, and Kenneth D. Miller. Divisive feature normalization improves image recognition performance alexnet. In International Conference on Learning Representations, 2022. URL httpp:///oppeerreiiir net/foorumidd-00333553
[S40] 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.
[S41] 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. Trans- former Circuits Thread, 2022. http://taaattommmecccciiiis//poo
html.
16


Reviews:
[{"review_id": "TTGd6z9w5d", "reviewer": "Reviewer_TzRp", "summary": "This paper investigates the resemblance between the NMDA receptor (NMDAR) in the hippocampus of the human brain and the activation functions used in the transformer architecture (e.g., ReLU, GELU). Then, this paper presents a new activation function that exhibits similarities to NMDAR. It demonstrates that by adjusting the hyperparameter associated with this activation function, the memory capabilities of transformers can be fine-tuned. A 2D grid navigation experiment with transformers is investigated to examine the working memory and the reference memory.", "questions": "\n1. The technical contributions of this work could be better elucidated.\n\n2. More practical empirical studies e.g., on real-world large-scale datasets or more general and complicated tasks will be more interesting and helpful.\n\nMy final score will largely depend on the rebuttal and the discussion with other reviewers. I am willing to increase my score if the concerns are adequately addressed.\n", "limitations": "I did not find potential negative societal impacts in this work. See the “Weaknesses” section for my concerns.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper explored the famous transformer model from a neuroscience perspective, by drawing connections with the hippocampus in the human brain.\n\n2. This paper first investigates the similarity between the NMDAR in the hippocampus and the activation function such as ReLU and GELU in transformer models.\n\n3. The paper in general is well-written and easy to follow.", "weaknesses": "\n1. Although the proposed idea is interesting and neuro-inspired, the technical contribution seems limited (for ML venues). Based on my understanding, Sec 2.2 is a theoretical review of existing work, and the derivation of the NMDAR activation function in Sec 2.3 (w/. A.3 in appendix) is in general straightforward given previous work.\n\n2. The empirical results on the 2D navigation task seem promising, but it may be worthwhile to explore more general tasks that transformer models are typically applied to, e.g., language tasks, to better validate the efficacy of the proposed activation function.\n\n3. In Figure 3 (a), the result shown in the upper right subfigure (test) demonstrates that the proposed method could not predict the unvisited nodes in the novel map. Are there more detailed explanations for the phenomenon?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your reviews. The key __contribution__ of our work was to explore the resemblance between the NMDA receptor in the human hippocampus and activation functions in transformers. We appreciate that the reviewer finds our work to be a novel approach introducing a new activation function that reflects the properties of NMDAR and highlights its potential to fine-tune memory capabilities in transformer models.\n\nHere, we would like to address the reviewer’s __comments__ regarding\n1) technical contribution to ML venues\n2) tests on more general and practical tasks, and \n3) more detailed explanations for specific results. \n\nThese are important comments, and we believe that addressing them will improve the quality of our work. Please see our responses below: \n\n---\n\n> __Response 1)__ The technical contribution seems limited for ML venues. The technical contributions of this work could be better elucidated.\n\nThe following items demonstrate our technical contributions in terms of improvements over previous work (items 1 & 2) and new experiments (items 3 & 4). \n\n1. We designed a spatial navigation task that allows the intuitive separation of short-term memory and long-term memory performance.\n2. We integrated standard activation functions (ReLU, GELU, Swish) into NMDA$_{\\alpha, \\beta}$ (Table 1) and proposed additional hyperparameter space, $\\alpha$, which can potentially be beneficial for ML tasks.\n3. (Newly added) We conducted additional experimental tasks for standard vision and language models (ViT and GPT2). See details in __Gloabl Response 1)__ and __Figure S12 in attached pdf for rebuttal__.\n4. (Newly added) We provide additional sparsity analysis for the NMDA$_\\alpha$ activation function. This includes the mathematical intuition and empirical results on how increasing $\\alpha$ corresponds to an increase in sparse activation in the FFN population. We believe that this analysis will provide a better understanding of our NMDA nonlinearity to ML venues. See details in __Global Response 2)__ and __Figure S13 in attached pdf for rebuttal__.\n\n---\n\n> __Response 2)__ It may be worthwhile to explore more general tasks that transformer models are typically applied to, e.g., language tasks, to better validate the efficacy of the proposed activation function. More practical empirical studies e.g., on real-world large-scale datasets or more general and complicated tasks will be more interesting and helpful.\nWe agree that exploring more general tasks with our proposed activation function would be more interesting and helpful. Therefore, we have conducted additional experiments on language and image classification tasks.\n\nThank you for the suggestion. Testing the proposed activation function over more complex tasks and datasets will strengthen our findings. We have conducted the following set of additional experiments on language and image classification tasks.\n\nFor the result, please refer to  __Gloabl Response 1)__ and __Figure S12 in attached pdf for rebuttal__.\n\n---\n\n> __Response 3)__ In Figure 3 (a), the result shown in the upper right subfigure (test) demonstrates that the proposed method could not predict the unvisited nodes in the novel map. Are there more detailed explanations for the phenomenon?\n\nThis is a good catch. It is due to the transformer model’s design constraint. The problem in our proposed model is caused by the fixed length of the context window (c=64), as depicted in Figure 2b. Any node that has not been visited by the agent in the previous 64 steps is classified as unvisited, which means it is outside the current context window.\n\nThe key difference between our model and the transformer model is the recurrent positional embedding, which encodes the previous action sequence information rather than the preceding sequence of observations. Because of the context window size, our model is unable to access the sensory observations of unvisited nodes via the self-attention mechanism.\n\nWe will try to explain this limitation better. \n"}, {"review_id": "nACBM1kMuw", "reviewer": "Reviewer_r1mY", "summary": "Many recent papers have shown a relationship between Transformers and biological structures, particularly the hippocampal formation. This work demonstrates that the types of non-linearities provided by NMDAR dynamics (which have known biological importance) can be beneficial in a Transformer architecture for a toy what-where memory/navigation task.", "questions": "1. Different activation functions have different computational footprints. How does the computational footprint of the proposed activation function compare to others and is the trade-off worth the performance increase?\n\n2. Line 318. Why is this surprising?\n\n3. While TEM has been popularised and cited widely, and is discussed here, how is it considered 'state-of-the-art' when it is essentially a very fancy Markov chain model? Or do the authors argue that that's what hippocampus is/does?", "limitations": "Limitations need more discussion.", "rating": 3, "confidence": 4, "soundness": 1, "presentation": 3, "contribution": 3, "strengths": "This paper deepens the connection between Transformers and biology, which helps both the neuroscience and machine learning communities in understanding these systems. It may also help improve Transformer models, which are already very popular and performant but may be further improved by inspiration from neuroscience.", "weaknesses": "The values of $N$ used in Figure 3 suggest that in this case there is a ceiling effect, i.e., once N is sufficiently large, the benefit of the NMDA-like non-linearity becomes negligible. Further, while error bars are included, I could not find the statistical testing to demonstrate differences between the groups. I think it is important to do such tests to measure the size and significance of any effects.\n\nAlthough the what-where task is clearly neuroscientifically-inspired and has been used in past similar studies, it is not clear whether anything contained here will generalise to other, more ecologically-valid tasks from neuroscience or more naturalistic real-world data from machine learning. The authors should test on more complex tasks/data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate that the reviewer finds our __contributions__ regarding 1) the connection between transformers and hippocampal formation, including the incorporation of NMDAR nonlinear dynamics, and, 2) its potential benefits for both neuroscience and machine learning communities.\n\nHere, we would like to address the reviewer’s  __comments__ regarding \n1) ceiling effect via increasing N and lack of statistical testing in Figure 3\n2) generalization to more complex or real-world tasks\n3) the computational footprint of the proposed activation function\n4) clarification on why is surprising in line 318\n5) TEM's characterization and relevance to the hippocampus.\n\nThese are great comments, and we believe that addressing them will improve the quality of our work. Please see our responses below:\n\n---\n\n> __Response 1)__ Values of N used in Fig 3 suggest there is a ceiling effect. If N is sufficiently large, the benefit of the NMDA-like non-linearity becomes negligible? It is important to do statistical testing to demonstrate differences between the groups to measure significances.\n\nThank you for the feedback. We agree that conducting statistical tests of significance will provide more convincing evidence. Regarding this comment, we ran additional experiments to the level where statistical analysis was possible. We used a nonparametric statistical test, i.e., Mann-Whitney U test, across all groups in Figure 3. This result is included in provided __(Figure S7 in attached pdf for rebuttal)__. We will include this result in our Appendix.\n\nAs shown in Figure S7d-f in attached pdf for rebuttal, the overall significance level of NMDA vs others increases (i.e., yellow color means lower $p$-values). We acknowledge we are unable to fully address the reviewer’s concern about the ceiling effect by increasing N. However, based on our current statistical test results, as suggested by the reviewer, NMDA-like non-linearity appears to be effective even in larger N (up to 64) conditions.\n\n---\n\n> __Response 2)__ It is not clear whether anything contained here will generalize to other. The authors should test on more complex tasks/data.\n\nThank you for the suggestion. Testing the proposed activation function over more complex tasks and datasets will strengthen our findings. We have conducted experiments on language modeling and image classification tasks. Please refer to __Global Response 1)__.\n\n---\n\n> __Response 3)__ How does the computational footprint of the proposed activation function compare to others and is the trade-off worth the performance increase?\n\nThe computational footprint of the NMDA function is comparable to that of Swish and requires less computation than GELU. Under the JIT compile feature in PyTorch, the computation speeds of Swish, NMDA, and GELU are essentially the same on an actual GPU. \n\nReLU offers low computational cost and high memory efficiency in large-language model training, but it is not widely used in language models due to its inefficiency in training stability and learning speed. GELU is widely used, despite its higher computational demands, because of its stable training process and rapid reduction in training loss. \n\nWhile we did not observe the stability issues associated with ReLU in our work, considering these trade-offs into account when evaluating the proposed activation function will be important. We appreciate the reviewer's insight and will consider these factors in future studies.\n\n---\n\n> __Response 4)__ Line 318. Why is this surprising?\n\nThe role of NMDAR in the CA1 region of the hippocampus is known to be vital for long-term memory formation in neuroscience, leading to discoveries such as place cells (i.e., a finding that won a Nobel Prize).\n\nWe were surprised to discover parallels between this process and transformers, a leading architecture in deep learning. We observed similarities such as the use of NMDAR-like nonlinearity, the emergence of place cells in feed-forward networks (FFNs), and the effect of linear activation functions on long-term memory.\n\nTo clearly convey these points, we have newly added a figure to illustrate these connections in __Figure S11 in attached pdf for rebuttal__; we will include this figure in Appendix.\n\n---\n\n> __Response 5)__ While TEM has been popularised and cited widely, and is discussed here, how is it considered 'state-of-the-art' when it is essentially a very fancy Markov chain model? Or do the authors argue that that's what hippocampus is/does?\n\nThank you for your insightful comment regarding the classification of the Tolman-Eichenbaum Machine (TEM) as a 'state-of-the-art' model.\n\nThe navigation problem we addressed is inherently non-Markovian. However, the TEM allows us to render the problem Markovian by including a memory M, an integral part of the model that encompasses location-sensory conjunctions.\n\nThe TEM is based on the hypothesis that hippocampal cells encode these conjunctions (p = flatten($x^{T} * g$)) and that memories are rapidly stored in weights M through Hebbian learning ($M = M + p^T * p$). This approach not only explains various neural representations in spatial tasks but also extends to non-spatial tasks.\n\nThe ability of TEM to represent abstract spatial relationships via sensory input (lateral entorhinal cells) and abstract locations (medial entorhinal cells) uniquely positions it as a model capable of explaining complex neural phenomena. Its ability to store memories via simple Hebbian learning also underscores its novelty.\n\nWe referred to TEM as the 'state-of-the-art' model, not only because of its popularity but also because of its comprehensive ability to mimic hippocampal synaptic potentiation, aligning closely with observed neural representations in both spatial and non-spatial tasks.\n\nWe hope this clarifies our rationale behind classifying the TEM as 'state-of-the-art' and we welcome any additional comments. (We will also consider giving it a more neutral name.)"}, {"review_id": "QSMIwCPCEO", "reviewer": "Reviewer_xzxf", "summary": "This paper builds on recent work connecting the hippocampus to the transformer architecture by introducing a new hippocampus-inspired activation function for the transformer's feedforward modules. In a toy navigation task, several empirical investigations show that the choice of this activation function has a large effect on the model's ability to store and recall information from longer-term memory, and that models using brain-like activation functions contain neurons with place cell-like firing patterns.", "questions": "* Regarding the above weakness: How might the choice of nonlinearity give rise to the effects (reference memory error rate, place cell score, etc) documented here?\n* Why is a recurrent embedding used for actions? What happens if a different embedding is used?\n* The definition of place cell score is crucial to the paper and should be in the main body, I think. On a related note, I find the definition given in the appendix difficult to understand. In particular, given how the auxiliary graph $\\mathcal G$ is constructed, it seems to me that every vertex will be a direct descendant of $node_k$ since $node_k$ has the highest firing rate. How this (a) measures sensitivity of a neuron to a specific location or (b) resembles the peak method used to identify place cells in the neuroscience literature should be explained more carefully. A figure might also be helpful here.\n* The discussion on lines 318-324 notes that the results here agree qualitatively with experimental results from neuroscience. It would be nice to have a figure summarizing these findings, and perhaps comparing against this paper's results, ", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* This paper continues an interesting research direction -- exploring the connections between the hippocampus and the transformer architecture. This is exciting work that is relevant to both machine learning and neuroscience.\n* The empirical analysis is presented well with many attractive and easy-to-interpret figures, seems quite thorough, and largely supports the paper's conclusions.", "weaknesses": "* The paper presents a convincing argument that the nonlinearity in the feedforward module of a transformer is important for forming and recalling certain types of memories, and that the brain-inspired nonlinearity is among the best. However, it provides essentially no understanding (even intuitively) of why this should be the case.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad that the reviewer recognizes our __contributions__ in 1) connecting the hippocampus to the transformer architecture 2) the convincing empirical analysis which supports the transformer’s longer-term memory formation, and 3) unveiling place cell-like firing patterns in a feed-forward layer, 4) with many attractive and easy-to-interpret figures. \n\nHere, we would like to address the reviewer’s  __comments__ regarding: \n1) Nonlinearity's importance in memory formation.\n2) Choice of nonlinearity and its effects.\n3) Recurrent embedding for actions and alternatives.\n4) Place cell score definition and explanation.\n5) Summary figure linking findings to neuroscience.\n\nWe believe addressing these comments will enhance our work. Please see our responses below:\n\n---\n\n> __Response 1)__  It provides essentially no understanding (even intuitively) of why this should be the case.\n\nBased on the reviewer’s comment, we find it important to give a more intuitive evidence for why the proposed nonlinearity is effective. \n\nTo provide a better understanding of our work, we present both analytical and empirical results indicating that increasing alpha corresponds to increasing sparsity. Please refer to __Global Response 2)__.\n\n---\n\n> __Response 2)__ How might the choice of nonlinearity give rise to the effects (reference memory error rate, place cell score, etc) documented here?\n\nThe high sparsity properties of NMDA function may contribute to the increased score of place cells (__Figure S13 in attached pdf for rebuttal__). In the case of Transformers, enforcing neuron activation sparsity in MLPs has been found to improve the interpretability or selectivity of a higher percentage of neurons [1]. This evidence could explain why the place cell score increases when 𝛼 is increased. Among the various activation functions, NMDA with 𝛼=10 has the highest Gini index. A comparison of the Gini index among activation functions is shown in the figure.\n\nWe will include this result in Appendix.\n\n[1] Elhage et al., (2022). Softmax linear units. Transformer Circuits Thread.\n\n---\n\n> __Response 3)__  Why is a recurrent embedding used for actions? What happens if a different embedding is used?\n\nA recurrent positional embedding for actions was used to capture the temporal dependencies in the agent's actions. The model can encode information about the agent's previous actions and incorporate it into the current position prediction by using recurrent positional embeddings. This can be particularly useful in tasks where the sequence of actions is important, such as navigation. If a different embedding method is used, such as a non-recurrent learnable positional encoding, the model may still be able to learn the task. \n\nIn Appendix A.5, we conducted an experiment that may be relevant to the reviewer's question. In this experiment, we disrupted the embedding layers to make them non-recurrent, effectively preventing them from retaining previous action information in embeddings. The results indicate that working memory error and reference memory error increased significantly (see Fig. 3a and Fig. S3a). However, the behavior observed in this experiment is similar to the trend seen when increasing 𝛼 of NMDA$_{\\alpha}$ (see Fig. S3b). \n\nThese findings imply that, while path-integrated information from recurrent positional embedding is useful for learning the spatial structure of the map, it is not required nor essential to predict the unvisited node.\n\nThis finding supports the idea that working memory is crucial for memory consolidation and that disrupting it can cause impairment in reference memory.\n\n---\n\n> __Response 4)__ The definition of place cell score is crucial to the paper and should be in the main body. … A figure might also be helpful.\n\nThank you for the in-depth examination of our place cell section. We agree that giving definitions for the place cell score is important. We had tried various writing styles, including having the definition in the main body. In the current version, we tried to go directly to main results. We would be happy to incorporate this feedback and rearrange our content.\n\nAs the reviewer suggested, we believe including a schematic figure to explain how an auxiliary graph is constructed would help readability. As shown in __Figure S1 in attached pdf for rebuttal__, not all nodes are direct descendants as we constructed from the 2d grid map. We are grateful for the reviewers' suggestions to improve our work. \n\n---\n\n> __Response 5)__ It would be nice to have a figure summarizing that the results here agree qualitatively with experimental results from neuroscience.\n\nWe will include a summary figure that connects the transformer (this work) and the brain (experimental findings). Please see our __Figure S11 in attached pdf for rebuttal__.\n\n"}, {"review_id": "oelC27hLgf", "reviewer": "Reviewer_FYwU", "summary": "This paper uses a transformer on neuroscience relevant task of spatial navigation, and shows when trained on both novel and familiar tasks simultaneously, place cell representations appear is different parts of the transformer depending on the current task – place cells in the feedforward net (post self attention) for familiar tasks, and in the self-attention layer for both novel and familiar tasks. The ability for the transformer to perform on familiar tasks is related to the activation function in the feedforward net.", "questions": "Presumably episodic memory is a better name than working memory, being that you’re saying it’s hippocampal related.\n\nSee weaknesses for other questions.\n", "limitations": "As described in the weaknesses, the main limitation is about the interpretation of NMDAR and the relevance of the activation function for ML. Otherwise the paper is good.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Interesting question of different representations for different task distributions. Gets at hippocampal consolidation which is under-explored with models. Convincing simulation results. ", "weaknesses": "1)\tThe NMDAR part of the paper is the least convincing, being that standard activation functions like ReLU seem to work just as well. Don’t think you can have sentences like ‘We find that NMDAR-like nonlinearity is essential for shifting short-term working memory long-term reference memory in transformers’. \n2)\tThe NMDAR impairment is essentially just making the activation function more linear, which of course will mean that it can’t effectively store long term memories. \n3)\tTo effectively make the NMDAR receptor point as a different activation function class, you’d need to test it on more standard ML tasks, rather that bespoke neuroscience tasks.\n4)\tYou have shown a nice potential neuro-ai link via NMDAR, but it’s currently just a potential relationship. It would need a lot more testing to make it concrete, e.g. showing that alpha=10 (the best performing model) is the regime that the brain operates in etc.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your reviews and encouragement. The key __contributions__ of our work was to explore the resemblance between the NMDA receptor in the human hippocampus and activation functions in transformers. We are happy to see that the reviewer finds our results convincing.  \n\nHere, we would like to address the reviewer’s  __comments__ regarding \n1) strong statement on the NMDAR-like nonlinearity\n2) NMDAR impairment experiments\n3) results on standard ML tasks\n4) potential evidence of the relationship with the brain\n5) terminology regarding working memory and episodic memory.\n\nThese are important comments, and we believe that addressing them will improve the quality of our work. Please see our responses below:\n\n---\n\n> __Response 1)__ The NMDAR part is the least convincing, being that standard activation functions work just as well. Don’t think you can have sentences like ‘We find that NMDAR-like nonlinearity is essential for shifting short-term working memory long-term reference memory in transformers’.\n\nAs the reviewer noted, the standard activation function ReLU works fine in the transformer model. Our intention was to mention that standard activation functions widely used in deep models such as ReLU, GELU, and Swish can be viewed as a subset of the NMDAR-like activation function (Table 1). We will consider ways to convey this message while toning down our expression, including the following in the abstract:  \n\n(Original) We find that NMDAR-like nonlinearity is essential for shifting short-term working memory to long-term reference memory in transformers  \n(Revised) We find that NMDAR-like nonlinearity has a beneficial role in shifting short-term working memory to long-term reference memory in transformers\n\n---\n\n> __Response 2)__ The NMDAR impairment is essentially just making the activation function more linear, which of course will mean that it can’t effectively store long term memories.\n\nYour understanding is correct that making the activation function more linear in feed-forward layers affects the ability of transformers to efficiently store long-term memory. This finding is consistent with neuroscience, which shows that removing the nonlinearity caused by Mg2+ gating affects long-term memory formation [1]. While using linear activation functions may demonstrate limitations in machine learning, this experiment is linked to neuroscientific discoveries. It bridges the gap between understanding the transformer as a memory consolidation model and biological memory processes,  extending beyond traditional machine learning paradigms.\n\n[1] Miyashita et al., (2012). Mg2+ block of Drosophila NMDA receptors is required for long-term memory formation and CREB-dependent gene expression. Neuron, 74(5), 887-898.\n\n\n---\n\n> __Response 3)__ To effectively make the NMDAR receptor point as a different activation function class, you’d need to test it on more standard ML tasks, rather that bespoke neuroscience tasks.\n\nThank you for the suggestion. We agree that testing our NMDA function on standard ML tasks would be effective to get NMDAR points as a new function class. We have conducted additional experiments on language modeling and image classification tasks. \nPlease refer to our __Global Response 1)__.\n\n---\n\n> __Response 4)__ You have shown a nice potential neuro-ai link via NMDAR, but it’s currently just a potential relationship. It would need a lot more testing to make it concrete, e.g. showing that α=10 (the best performing model) is the regime that the brain operates in etc.\n\nAs suggested by the reviewer, we calculated the corresponding α value in the physiological CA1 hippocampal neurons following real experimental values [1]. The calculated α was ranged between 0.01~0.2.\n\n|Term|Symbol|Typical Value|Units|\n|-|-|-|-|\n| Magnesium Ion Concentration | [ $\\text{Mg}^{2+} ]$ |1| mM|\n| Physiological Temperature|T| 37 (310) |°C (K)|\n|Dissociation Constant at $V=0$ (mV) | $K_{\\text{Mg}^{2+}}$ | 1 – 20 | mM |\n|Temperature Constant | $\\beta$ | 0.062 | mV$^{-1}$ | \n\nThese values are typical and can vary depending on the specific biological context, experimental conditions, and the model used. While previous research has shown that increasing the Mg2+ level in the brain with a specific compound such as MgT(magnesium-L-threonate) increases long-term memory formation [2], its effective concentration only increased about 15% from the baseline (expected maximal α in the brain to be ~0.23), possibly due to the physiological ion excretion process in humans. \n\nWhile the brain's effective α was lower than expected, nonetheless, we are grateful to the reviewer for this suggestion. \n\n[1] Kirson et al., (1999). Early postnatal switch in magnesium sensitivity of NMDA receptors in rat CA1 pyramidal cells. The Journal of Physiology, 521(Pt 1), 99.  \n[2] Slutsky et al., (2010). Enhancement of learning and memory by elevating brain magnesium. Neuron, 65(2), 165-177.\n\n---\n\n> __Response 5)__ Presumably episodic memory is a better name than working memory, being that you’re saying it’s hippocampal related.\n- Terminology: Episodic memory is concerned with the ability to recall specific events and the contexts in which they occurred (What-Where-When), whereas working memory retains information over short durations to perform cognitive tasks. Therefore we used 'working memory' to emphasize short-term retention and processing in the hippocampus, akin to the limited context in transformers.\n- Why Not Episodic Memory?: We avoided this term as our study lacked the time element (When) essential for episodic memory, focusing on What-Where components instead.\n- Future Work: We plan to investigate What-Where-When aspects of episodic memory in future research on transformer models.\n\nWe appreciate the reviewer's feedback and hope this response clarifies our terminology choice.\n"}]

Author rebuttal:
We are grateful to the reviewers for their insightful comments on our study. All reviewers recognized our contribution to investigating NMDAR-like nonlinearity in the transformer's feed-forward network, which will be beneficial for both communities of ML and neuroscience. Reviewers find our work well-written and easy to follow (__Reviewer TzRp__) many attractive and easy-to-interpret figures (__Reviewer xzxf__). Reviewers also think our empirical analysis results support our paper conclusion well (__Reviewer FYwU & xzxf__). 

During the rebuttal period, we made additional figures and conducted additional experiments for addressing the reviewers’ questions and concerns. Highlights include experiments on
1) standard machine learning tasks on ViT and GPT2 (Figure S12), 
2) sparsity analysis on activation functions (Figure S13),
3) improved illustrations of place cells (Figure S1), 
4) statistical significance tests (Figure S7),
5) and clear comparisons between the hippocampus and the transformer model (Figure S11).

These efforts demonstrate our commitment to addressing the reviewers' questions and enhancing the connection between neuroscience and machine learning.  ___Please find our attached PDF regarding our response.___ Thank you.

---
> __Global Response 1)__ standard machine learning tasks on ViT and GPT2

Testing our NMDA activation function over more complex tasks and datasets will strengthen our findings. We have conducted the following set of additional experiments on language and image classification tasks.

__1. Language modeling with GPT2__

We have tested the NMDA$_\alpha$ on the GPT2 model with the Open Web Text dataset. Our analysis indicates a slight performance improvement. Please find this result of loss curves from __Figure S12 in attached pdf for rebuttal__.

__2. Image classification tasks with ViT__

We also tested the NMDA$_\alpha$ on the ViT model (3 trials for each condition) and found an increasing tendency of performance (although statistically non-significant). The table below reports top-1 test accuracies for the CIFAR-100 and the TinyImageNet datasets. 
| **Dataset**     | GELU | NMDA$_{\alpha=10}$ | NMDA$_{\alpha=0}$ |
|--------------------|----------|-------------------------------|-----------------------------|
| **CIFAR100** | $69.92\pm0.34$ | $70.27\pm0.47$ | $49.91\pm0.58$ |
| **TinyImNet** | $54.90\pm0.45$ | $55.72\pm0.03$ | $40.36\pm0.52$ |

We would like to mention that although these results are promising, we are not conclusive in stating that the NMDA function outperforms GELU on all real-world tasks. Future work will include testing our NMDA function on different model sizes of GPT and other language models, as well as testing the pre-trained models on downstream NLP tasks such as reading comprehension, question answering, common sense, MMLU, and BIG-bench. 

---

> __Global Response 2)__ sparsity analysis on activation functions

To provide a better understanding of our work, we present both analytical and empirical results indicating that increasing alpha corresponds to increasing the sparsity of the feed-forward layer neuronal population, which supports our last paragraph of results section 3.2 (lines 246-248), and third paragraph of the discussion section (lines 309-311).

__Mathematical intuition__:  
Given $ \text{NMDA}_{\alpha} = x / (1+\alpha e^{-x}) $, we can rewrite this function as following:

$
\text{NMDA}_\alpha=x \cdot \frac{1}{1+e^{-(x-c)}}, \\
\text{where,} \alpha=\exp{c}.
$

From the above expression, increasing $\alpha$ corresponds to shifting the sigmoid function towards a positive direction by increasing c, which can be interpreted as increasing the threshold of information gating mediated by the sigmoid function. As a result, an increase in $\alpha$ may cause the sparsity of the downstream population activities.

__Empirical result__:  
To confirm the above mathematical intuition, we measured the Gini index [1] to determine the sparsity of the feed-forward layer activities. For each input sequence, we calculate the Gini index $ G = \frac{{\sum_{i=1}^{K} \sum_{j=1}^{K} |x_i - x_j|}}{{2K^2 \bar{x}}} $ 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). 

The Gini index ranges from 0 to 1. When only a few neurons have high activation values and others have low values, the Gini index is close to 1. On the other hand, when most neurons have homogeneous activation values, the Gini index is close to 0 (1 = absolute sparsity, 0 = all activations equal).

As shown in __Figure S13 left in attached pdf for rebuttal__, increasing $ \alpha $ in NMDA causes the Gini index to rise. This increase suggests that population activities become more heterogeneous, resulting in a heavy-tailed distribution. Furthermore, the Gini index of NMDA with an alpha value of 10 is greater than that of the GELU activation function (represented by the dashed line). These results imply that $\text{NMDA}_{\alpha}$ may improve long-term memory formation by increasing the sparsity of activations in the feed-forward layer.

Prior studies [2] investigated the generalization performance and sparsity in overparameterized models. Although the mechanism underlying the emergence of sparse representation in large models is not fully understood, it is worth noting that overparameterized models with sparse representations are simpler models than those with dense representations. 

The discussion above will be included in the Appendix.

[1] Miller et al., (ICLR 2021). Divisive Feature Normalization Improves Image Recognition Performance in AlexNet.  
[2] Li et al., (ICLR 2022). The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers. 


Meta-review: Recent studies found parallels between transformers and the hippocampus. This study explores the relationship between NMDA receptor (NMDAR) dynamics in the hippocampus and activation functions used in transformer architectures. A new activation function inspired by NMDAR is introduced, impacting the memory capabilities of transformers. A navigation experiment demonstrates that this function affects working and reference memory processes in transformers, and the models that use this function show brain-like spatial representations.

Most reviewer criticism was addressed during the rebuttal period.
Decision: Accept (poster)