Transformer needs NMDA receptor nonlinearity for long-term memory
Authors: Dong-Kyum Kim
Jea Kwon
Meeyoung Cha
C. Justin Lee
Venue: Submitted to ICLR 2023 (2023)
Keywords: NMDAR
hippocampus
transformer
memory

Abstract: The NMDA receptor (NMDAR) in the hippocampus is essential for learning and memory. We find an interesting resemblance between deep models' nonlinear activation function and the NMDAR's nonlinear dynamics. In light of a recent study that compared the transformer architecture to the formation of hippocampal memory, this paper presents new findings that NMDAR-like nonlinearity may be essential for consolidating short-term working memory into long-term reference memory. 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 formation. Our experimental data suggest that the concept of place cells and reference memory may reside in the feed-forward network layer of transformers and that nonlinearity plays a key role in these processes. Our findings propose that the transformer architecture and hippocampal spatial representation resemble by sharing the overlapping concept of NMDAR-like nonlinearity.

Under review as a conference paper at ICLR 2023
NMDA RECEPTOR NONLINEARITY ATTRIBUTES TO MEMORY CONSOLIDATION IN TRANSFORMERS
Anonymous authors Paper under double-blind review
ABSTRACT
The NMDA receptor (NMDAR) in the hippocampus is essential for learning and memory. We find an interesting resemblance between deep models' nonlinear activation function and the NMDAR's nonlinear dynamics. In light of a recent study that compared the transformer architecture to the formation of hippocampal memory, this paper presents new findings that NMDAR-like nonlincarity may be essential for consolidating short-term working memory into long-term reference memory. We design navigation task assessing these memory functions and show that manipulating the activation function (i.e., mimicking the Mga+-gating of NMDAR) disrupts long-term memory formation. Our experimental data sug- gest that the concept of place cells and reference memory may reside in the feed- forward network layer of transformers and that nonlinearity plays key role these processes. Our findings propose that the transformer architecture and hip- pocampal spatial representation resemble by sharing the overlapping concept of NMDAR-like nonlinearity.
INTRODUCTION
In the hippocampus, NMDAR is regarded as an essential component that mediates synaptic plas- ticity, memory formation, and spatial representation (Li & Tsien 2009 Tsien et al.] [996] [entros et al., 1998). NMDAR serves as switch for synaptic plasticity and long-term memory forma- tion (Bliss& & Collingridge, 1993; Slutsky al. 2010 Miyashita et al., 2012). In addition, NMDAR has been highlighted for its importance in place cell representations in hippocampal CA1 (McHugh (ct al. 1996; Kentros et al. 1998). Place cells the hippocampus Oe'Keefe & Dostrovsky/ 1971) and grid cells in the entorhinal cortex (Hafting et al. 2005) are thought to be crucial for spatial navigation in an animal. These discoveries have triggered recent efforts to replicate these spatial representations found in the brain by using deep neural networks (Banino et al. 2018: Cueva & Wei 2018 Whittington et al. 2022).
In NMDAR depicted in Fig. Ta the ion channels that reside in the post-synaptic region have unique characteristics that distinguish them from other ion channels in the brain. Their nonlinear dynamics are modulated by Mg ion blockade at the pore region. NMDAR requires activity-dependen re- pulsion of Mg ion (Nowak et al. 1984 Mayer et al. 1984) to be functional, and this phenomenon partly interesting because it serves as self-gating of ion influx in the post-synaptic region. In particular, the Mg gated nonlinear dynamics of NMDAR plays a key role in synaptic plasticity and memory formation (Slutsky et al. 2010; Miyashita et al. 2012).
Recently, the relationship between the transformer Vaswani et al. 2017) and hippocampal formation model has been reported (Whittington et al., 2022). The transformer the most advanced deep learning model, showing unprecedented results in various tasks such as language modeling (Devlin ct al. 2018 Brown et al. 2020), computer vision (Dosovitskiy et al. 2020 Radford et al. 2022). and art generation (Ramesh et al. 2022). This model has two consecutive modules, a self self-attention layer and a feed-forward network (see Fig. Пp). W Whittington et al. (2022) show the self-attention layer is closely related the state-of ofate-of-the-ar neuroscience model (Whittington softmax neurons in the self-attention layer behave like place cells in a navigation task. However, studies the role of neurons in feed-forward networks have been absent.
al. 2020) and claim that

---

Under review as a conference paper at ICLR 2023
(a)
(b)
NMDAR
X
Linear
Feed Forward Net
MPP11
Nonlinear Unit
LayerNorm
Linear
(x)A,
GELU
Self-attention
1 00
LayerNorm
Synaptic
Plastictt
AMPAR
NMDWR
Input embeddings
Figure 1: (a) Schematic diagram of Mgtt-gated NMDAR modulating synaptic plasticity (left), its IV curve of current-voltage dynamics (right top) and NMDAR-inspired activation function, NMDAa (x) (right bottom). (b) Transformer architecture and its feed-forward network's s activa- tion function, Gaussian Error Linear Unit (GELU; left bottom).
We find an interesting resemblance of NMDAR nonlinearity with the Gaussian Error Linear Unit (GELU), nonlinear activation function popularly used in the transformer's 's feed-forward network (Fig. M Similar to NMDAR's aetivity-dependent gating mechanism of ion influx, the form of the GELU function has combination of input with the self-gating function. Biological experi- ments have shown the critical consequence of changing NMDAR's 's nonlinearity in synaptic plastic- ity and long-term memory formation (Slutsky al. 2010 Miyashita et al. 2012), while the role of NMDAR-like nonlinearity place cell representation remains unclear.
This work is inspired by the fascinating resemblance of NMDAR's nonlinearity dynamics with the GELU activation function and the recent model relating transformer's self-attention mechanism to hippocampal formation (Whittigton et al.) 2020; 2020). These findings motivated to ask ques- tion; Can NMDAR-like nonlinearity in the foed-forward network layer of transformers en- hance the formation of long-term memory and spatial place cell representation?
To address this question, we propose novel NMDAR-like activation function derived from NM- DAR IV curve and design spatial navigation task in a 2D grid environment that can different memory types well formulated in neuroscience experiments (Olton et al., 1977 1979): working memory and reference memory. Working memory controls the events from within- while reference memory controls across-trials from the unchanging environment. We evaluate the transformer model with the NMDAR-like activation function this task; the results show that 1) place cell representations emerge in feed-forward networks, 2) the reference memory can be con- trolled by the nonlinearity of the NMDAR-lik activation function, 3) the place cell in ferw-forward networks is strongly correlated with the reference memory, while the place cell in self-attention lay- ers has no correlation, 4) the proposed NMDAR-like activation shows best reference memory performance when compared to other widely used nonlinear activation functions.
two
Our experimental data suggest that NMDAR-like nonlinearity the cerd-forward network layer of the transformer can enhance the long-term memory formation and place cell representation.
2 TRANSFORMER
The transformer architecture VVawaani al. 2017) can be constructed by stacking multiple blocks of self-attention layers and forward networks (see Fig. [b).. Here we briefly review the self- attention mechanism and the feed-forward networks.
2

---

Under review as a conference paper at ICLR 2023
Self-attention mechanism Given sequence [Xi,..,;yy xT) } of d-dimensional input embeddings, the self-attention layer calculates the interaction term between each embedding element within a context window via the self-attention mechanism. More formally, each input embedding applies two linear layers (W& and W,) to the embeddings to form the key matrix K and value matrix V:
KJ = k K- ... k,1, where k; = x (W& E Rdxde); VT = Vi V Vt-e+1 v,Jl. where V, = x; (W, E Rdxd)).
(1)
Here, c denotes the context length. The key matrix K E (c(+))xdd is then used to compute the interaction score between input embedding step and all the vectors in K via dot products:
= q where qu - W (W₂ E dd
The normalized values of s, E (R(++), called attention values, are calculated via the softmax func- tion; the final output of the self-attention mechanism weighted sum of the value vectors in V E (g(++))xd with the attention values:
qqKK" Vdg
Y7 so somaaax
V.
After this update, Yı E R is updated by another linear transformation W. Rdxxd: z, W- The output z, is added to the x,; z, + x, is the final output of the self-attention layer, and this information sent through the following subsequent layer.
Feed-forward networks Another important component of transformer layer is the ford-forward network. This consists of two linear layers with point-wise nonlinear activation function d:
FFN( = (xx JU2,
where U1 E Rd and U2 Rd are trainable weight matrices. Sukhbaatar et al. (2019) pointed out that equation 4 equation havve similar structures except for the following two major differences: 1) UT and U2 matrices are fixed over different input sequences while K and V matrices are dynamically changed as input is and 2) operations in feed-forward networks are entirely 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 feed-forward networks store general knowledge about the task that does not depend on the situation.
METHODS
3.1 RELATING ACTIVATION FUNCTION IN TRANSFORMERS WITH NMDAR NONLINEARITIES
NMDAR's nonlinear dynamics arise from the voltage- gated Mg repulsion at the NMDAR chan- nel's pore (Nowak et al. 1984 Mayer al. 1984) (Fig. Dp left). Previously, 2g²+-gated -gated NMDAR open probability has been shown to follow ion blockade model of Woodhull (1973):
Pa (x) 1tae-Bz Bz
(5)
where I represent an input voltage, a = [M is a parameter determined by [Mg2+]. Kmga+ is dissociation constant, and 33 is temperature constant. Experimentally, increasing the Mg level in the brain can enhance long-term memory formation (Slutsky et al. 2010). We ob- served the NMDAR's 's nonlinear dynamics of IV curve (Fig. D right top; current-voltage relation- ship) in the synapse to closely resemble the form of GELU activation function. GELU a widely used activation function in transformers (Fig. Tb left bottom; GELU(x) z 20 20.(1.702. where o the sigmoid function) (Hendrycks Gimpel| 2016 Devlin et al. 2018 Brown et al. 2020). In- spired by this resemblance, we define new nonlineat activation function (Fig. right bottom) with a parameter which modulates dynamics as following (see details in Appendix
A.2
NMDA (I) = xP. (z)
(6)
+ ae
To investigate this NMDAR-like nonlinearity in transformer memory formation, we replaced the GELU(x)) activation function with NMDA, (zr) in standard transformer model.
3

---

Under review as a conference paper at ICLR 2023
(a) Visted place prediction
Unvisited place prediction
(b)
transformer (causal)
B
A
B
1 D
?
Action 02)
(4]
A
context
query
Figure 2: (a) Sensory observation prediction task in 2D grid, where dotted squares indicate the target position to predict given a sequence of past actions and observations. The unvisited (visited) places are represented in gray (black) letters. (b) A transformer model for predicting the next loca- tion's 's observation based on sequences of [action, observation] pairs. Using the 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.
3.2 TRANSFORMERS LEARN SPATIAL NAVIGATION TASKS
We train the transformer model to predict the subsequent sensory observations of agent that randomly walks a 2D grid environment (Whittington et al. 2022) (Fig. 2h. A sequence of previ- ous [Action (a), Observation (z:] pairs is input to the model, and the subsequent observation is masked for prediction (Fig. (2). Instead of using sinusoidal positional encoding (Vaswani et al., 2017) that commonly used in transformers, we employ the recurrent positional embedding which encoding the location of an input element by using the recurrent neural network (RNN) (Whit-
tington al. 2022
We generate the embedding vectors of the sensory observation sequence with word embed- ding layer, but the embedding vectors of the action sequence are generated by RNN; et+1 = tanh (e, where e; is recurrent positional embedding at step I, and Wa is the action-dependent trainable weight matrix. The input is given by {[]] [e2, embedding vector of sensory observation z; the initial recurrent positional embedding e is sampled from normal distribution and we mask the last observation øt. We generate N maps of 11 x 11 2D grids. A random sensory observation among ten letters is placed at each position each map. Agents can move 'up,' 'right,' 'down,' 'left,' or 'stay.' An agent starts at a random position and initiates random walk on the map, randomly selected map among N training maps, for 2,048 steps for each trial.
[e,, x,]}. where x denotes the
The model is trained with the softmax cross-entropy loss and predicts the subsequent sensory obser- vations (i.e., dotted squares). We evaluate two types of memory: working memory and reference memory. When the prediction on nodes that were previously visited during the random walking is incorrect, it will count as the working memory егrоr (see Fig. 2p left). On the other hand, when the prediction on unvisited nodes is incorrect, it will count as the reference memory error (see Fig. 2 right). Minimizing the reference memory error by memorizing input sequences is infeasible; the possible number sequence configurations exponential since the input sequence is randomly generated at each trial. To solve this task, the model should be able to 1) understand the abstract structure of 2D space, 2) infer which map it is on from input sequence data, and 3) memorize what sensory observation is placed at each position in that map. We describe more training, evaluation, and transformer model details in the following section. See details for task and definition regarding working and reference memory in Appendix A.3
RESULTS
4.1
IMPLEMENTATION DETAILS
In our experiment, the feed-forward network (FFN) in the transformer model consists of two linear layers (see Fig. D and equation 4 with the NMDAR-inspines activation function NMDA,, (Eq. 6
This method closely related to the most advanced neuroscience model of the hippocampus.
4

---

Under review as a conference paper at ICLR 2023
We use TransformerXL (Dai et al. 2019) with extended memory length of 32 and segment length of 32 so that the context length is 64 and working memory error is measured when the node to predict its sensory observation is in the context window (see Fig. 2b); i... node that the agent had never visited within recent 64 steps is treated as an unvisited node. 'Te input embedding is concatenated vector (,, x] of the word emhedding x (dimension of 256) and the recurrent positional embedding (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 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 ADAM (Kingma & Baj 2015) optimizer and learning rate schedule with a linear decay from 0.0001 (start) (end). We run 512 random walk simulations (trials) in parallel for collecting training trajectories. The total number of random walking steps is 2,048 for each simulation so the total number of steps for training model is 512 (batch size; number of trials per epoch) x 2,048 (total number of steps in a trial) x 200 (number of epochs) (see Fig, in Appendix A.3). All runs are performed on a single NVIDIA TITAN V GPU.
4.2 WORKING MEMORY ERROR & REFERENCE MEMORY ERROR
To measure the impact of nonlinearity a in the FNNs, we train the transformer models with differ- ent values of a in [0, 0.01, 0.05, 0.1, 0.5, 1, 5, 10] and evaluate the working memory and reference memory errors on the train maps (i.e., familiar maps) and test maps (i.e., novel maps). The average number of unvisited nodes in single trial is 561.
The top left plot in Fig- 3 shows that the reference memory error on the train maps is rapidly de- creased over train trials when a is larger than zero, with larger improvement for increasing a. The reference memory error on the novel maps, however, is nearly constant at the chance level of 0.9 (= 1/(nmbeer of letters)) for all (see Fig. 3 top right). Fig. 3a (bottom right) shows that working memory is active on novel maps that had not been shown during training. This find- ing suggests that the working memory formation is intact novel maps. Training the models different numbers of maps N, Fig. [3p shows that increasing nonlinearity (i.e., a) helps activate the reference memory, and the trend of improvement consistently shown for N 32, 48, and 64 cases. Training over more maps leads to bigger reference memory errors. This is because more
(a)
Reference memery eirar
Train (Tamiliar)
0
Test (novel)
(b)
Reference memory enrr
0.8
N
0.8
32
0.0
0.6
0.6
64
0.05
0.4-
a=
0.1
0.4
0.2
1.0
0.2
0.0
10.0
10.0
0.0
0.0
0.01
0.05
0.1
0.5
1.0
5.0
10.0
Working memory erter
0.6
Refarence memory erro
0.6
0.15
0.4
0.10
0.4
0.05
0.3
0.00
0.2
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 trials
x10
No. of train triats
x10"
LuakyRulU
Figure 3: (a) Reference and working memory errors over training trials for training (familiar) maps and testing (novel) maps for N 64 where N is the number of training maps. (b) Reference mem- ory errors evaluated on training maps over different values of a in NMDA, and N. (c) Reference memory errors comparison between NMDA, activation functions. Inset: zoom on the top 4 activation functions. Error bars and shaded areas represent the standard deviation errors from three independently trained models.
10, GELU, ReLU, LeakyReLU, sigmoid, and tanh
5

---

Under review as a conference paper at ICLR 2023
maps require the model to store more pains of "what''where' memory (i.e., each training contains unique "ehat'-wherrs information).
In addition, demonstrate other nonlinear activation functions which are widely used in the machine learning literature. We test GELU (1.702)) ReLU (max (0, x), LeakyReLU (max(0, 2) 0.01 min(0, sigmoid, and tanh in the FFNs. As can be seen in Fig. 3 NMDA shows the lowest reference memory errors on the training maps.
Other memory types, such as information in path integration (i.e., recurrent positional embedding), may be used instead of reference memory. To test this assumption, we used non-recurrent positional embeddings to train the models. The result shows that working memory and reference memory er- rors increase substantially. However, it exhibits similar behavior to the trend of decreasing reference memory error while increasing a of NMDAa (see Fig. S in Appendix A.4). We also 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 the previous 65 steps, the first visited node prediction error is prediction error on a node that the agent visits for the first time in trial. The results for the first visited node prediction error in training maps are identical to the results for the reference memory error (see Fig. 2 in Appendix A.4). These findings suggest that the reference memory is used in training maps to predict the unvisited node.
4.3 PLACE CELLS IN TEED-FORWARD NETWORKS
Place cell is a neuron in the hippocampus which fires at a particular place of the environment (eeKee & Dostrovsky 1971). Studies have shown that hippocampal place cells encode the spa- tial location through localized firing patterns. They have been considered a substrate for long-term memory of the location where specific events occurred (i.e., previously visited position in our navi- gation task). Selective impairment of NMDAR in hippocampal CAI disrupts place cell emergence and long-term memory formation TTien al. 1996. Kentros al. 1998, McHugh et al. 1996).
We investigate the role of neurons in the FFNs and self-attention layers by measuring the neuron's place specificity. Given a K 2D grid environment as graph G EV and firing rate
(a)
(e)
tods in food forward boyor
NHEI 10.0
0.0
0.5
0.1
0.01
0.00
1.0
Reforecc momory
D
000
0.06
Profoditty derely
0.4-
02-
tari
sig
ROLU
NMEAA 0.0
GELU
Locky
6.0
10.0
10.0
0.2-
0.1-
0.0
0.2
0.3
0.1
06
0.0
0.0
0.2
0.6
Place osl rcor
Averape plaze call soore
Place octl score
Plaza to sof cttertoon aper
10.0
Refbre op motroy orort
770
0.0
Uhetoabiity dennty
04
F. 0.055
D:01
0.06
ReLl
0.1
0.4
0.5
Reference
diyroii
50
0.2
0.0
0.2
0/4
05
0.4
Average call
Figure 4: Reference memory-related place cells selectively emerge in the forward layer but not in the self-attention layer along with increase (N 64). (a, b) Example rate maps with place scores in feed-forward layers and self-attention layers at a = 10; from top left (high) to bottom right (low); color bar indicates the firing rate between Pruxx and Pmin- (c-d) Place cell score distri- butions with varying a in feed-forward layers (c) and self-attention layers (d). (e-f) Scatter plot of average place cell scores and reference memory errors. and p denote Spearman's 's rank correla- tion coefficient and significance score, respectively. (g) place cell score distribution and relationship of average place cell scores and reference memory errors in common activation functions: GELU, ReLU, LeakyReLU, tanh, and sigmoid. All results are evaluated from training maps.

---

Under review as a conference paper at ICLR 2023
(cumulative activation value at node i divided by the length of evaluation trial) of node i V as Ps. we define maximally firing node as fmax and its firing rate as Pmax- Where E is directed edges, cepth-first-searcl from source node,
which connect high to low firing nodes in G. From G, Fmax to build a sub-graph g = (V,E) which we call all connected components. Given G and g, the place cell score is defined following
Place cell score = Eiey ΣiEV Pi
where 7 = 11 /V is a discount factor and V is without node imax and leaf nodes (see details in Appendix A.6 To measure place cell score, record the firing rate P; of neurons over a random walking trajectory with 10 steps in one of the training maps; then we measure the place cell scores of neurons in the FFN and self-attention layers. The place cell score is when the neuron is firing only at certain node; the score is 0 when the neuron is firing homogeneously across all nodes.
Fig. Dp and 4p show the rate maps of neurons with place cell scores in the FFN and self-attention layers, respectively. For the self-attention layer, the total number of neurons in the softmax layer is 65 (context length + masked sensory observation) X 8 (number of heads)
(number of layers). (number
The total number of neurons in the FFN layer is set as 2,048 (number of neurons) of layers). As can be seen, our metric well represents place specificity. Fig. 4 and 4i show the distribution of place cell scores in the two layers with different values of a. When the a value is increased, the place cell score distribution found in the FFN layer becomes positively shifted (see Fig. 5 rate map examples for a = 0, 1.0, and 10.0), whereas place cell score distribution in the self-atten layers remains.
Fig. 4p and A show a relationship between the average place cell score and the reference memory error for each While average place cell scores in self-attention layers show no correlation with reference memory errors whatsoever, neurons in the FFN exhibit a substantial correlation. These results imply that the reference memory formation and place cell emergence be enhanced by NMDAR-like nonlinearity in the FFN.
In Fig. 4l. compare the place cell representations of our NMDA (a 0, 10) with the repre- sentations in FFNs with the activation functions used in Fig. 3e. Our results show that the case of NMDA,=110 outperforms other activation functions, in both reference memory formation and place cell representation. Our finding that increasing ([Mg2+] component) enhances reference memory in line with the biological observation that increasing the [Mg2+] in the brain enhances long-term memory formation (Miyashita et al. 2012)
In addition, we investigate the consequence of changing nonlinearity with other than NMDA We choose LeakyReLU with controllable negative slope (max((0,2) 2) + min(0, z)) to compare with
NMDA,, Compared to NMDA,=100 LeakyRel.U exhibits lower average place score the al- lowed range of a, indicating that NMDA,, is better in place cell emergence (see Fig. 10 in Ap- pendix A.7
Feed-forward layer 10)
Feed-forward layer (a Feed-forward layer
8
Figure 5: Rate maps of neurons with top-64 place cell scores in the feed-forward network with varying values of a; a = 10 (left), = 1 (middle), and = 0 (right)..

---

Under review as a conference paper at ICLR 2023
5 RELATED WORKS
The current study is inspired by recent observations that connect neuroscience and Al models. One such seminal work is by [Whittington et al.] (022). where the authors showed the relation- ship between the self-attention layer and the state-of-the-ar hippocampal model called the Tolman- Eichenbaum Machine (TEM; Whittington et al. (020 ). The current work expands the literature by
focusing on the ferdard networks (FFNs) In the transformer and the emergence of place cells.
TEM is a neuroscience-basse 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 assumes to store memory at the Hebbian weight matrix M E Rd4 and every outer product of key and value vector k, v, at each step i are simply stored in M via 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 a Σ kTv ak
where a a weighting factor. the memory retrieving phase with the query vector q. TEM uses an attractor network
q.V = =aq VV.
Whittington et al. 2022) found that the memory retrieving process in TEM has close mathematical
structure with equation 3 when the softmax function is replaced with a linear function. Their sub-
sequent model, called TEM-t (Whittington et al. 2022), replaces the attractor network (equation 9
with self-attention mechanism (equation 3). They demonstrated that TEM-t learns significantly faster than TEM.
TEM-t and TEM do not have a fixed context length c; therefore, these models store all information before step 1, i.e., e = f. The computational cost of the self-attention layer TEM-t O(t²), and retaining all previous information is too expensive from both biological and computational stand- poinf For TEM, the Hebbian update rule has no quadratic computational cost and can add all pre- vious information in fixed number of synapse di; matrix M is (dg) and the speed of memory retrieval is substantially slower than the self-attention mechanism (Demircigil& al. 2017 Ramsauer et al. 2021 Krotov & Hopfield (02T). In contrast to TEM and TEM-t that rely on a single memory system, the transformer model employs two separate memory systems: 1) context-dependent matrices K and V in the self-attention layer with fixed context length and 2) context-independee fixed matrices U1 and U2 (in equation 4 in the FFNs.
i however, the memory capacity of the Hebbian
We also focus on the observation by Whittington al. (2022) that softmax neurons in the self- attention layer behave like place cells. Nonetheless, the role of neurons in FFNs of transformers has not been thoroughly investigated, which is our contribution. In our work, we newly 1) propose a method for assessing the reference memory; 2) compare the effects of various nonlinear activation functions in FFNs with our NMDA NMDA-inspired activation functions on reference memory performance; 3) demonstrate the emergence of place cells in FFNs. We note that TEM and TEM-t only evaluated working memory errors in test maps.
6 DISCUSSION AND CONCLUSION
Searching for biological substrate of nonlinear activation function Rigorous previous efforts in finding the optimal nonlinear activation function underlie the great success of modern deep neural network models (Nair & Hinton| 2010: Hendrycks & Gimpel 2016]] Raaacchnnrann et al. 2017). However, the neural substrates that mediate nonlinearity in the human brain and their role in intelli- gence have not been clearly understood. Our work is one of the first to put together the biologically inspired nonlinearity and its effect on long-term memory formation and the place cell representation in the previously described transformer model of the hippocampal formation. This idea was tested on a sensory observation task in the 2D grid environment and with the implementation of NMDAR- like nonlinearity. Our data indicated that NMDAR-like nonlinearity in the feed-forward network
NNote that this a simplified description TEM and it not exactly the same. DDue to this limitation, TEM-t does not store all historical data. Instead, the model selectively chooses which data to store in K and V.
8

---

Under review as a conference paper at ICLR 2023
layer of transformers can enhance the formation of long-tenn memory and spatial place cell repre- sentation. Furthermore, this design choice improves long-term memory more than other commonly used nonlinear functions.
Our results agree qualitatively with previous NMDAR impairment experiments from neuroscience: 1) hippocampal CA1 NMDAR perturbation does not impair working memory (Lee & Kesner 2002), 2) changing NMDAR Mg -gat-ga (changing a in this work) enhances or disrupts long-term mem- ory formation (Slutsky et al., 2010 Miyashita al. 2012), 3) NMDAR is required for long-term stabilization of newly forming place fields (McHugh et al. 1996, Kentros et al. 1998). Our contri- bution showing these patterns experimentally for the first time.
Short-term working memory and long-term reference memory Memories can be divided into short-term memory and long-term memory with respect to time (Atkinson & Shiffrin 1968). Later, these two different memory systems were denoted as working memory and reference memory, con- cerning functional aspect (Baddeley & Hitch 1974). In neuroscience, there is a consolidation theory that some short-term working memories are converted to a long-term reference memory system, and others fade out. There has been accumulating evidence that NMDAR is essential for memory con- solidation in the hippocampal CAI region (Kentros et al. 1998, Shimizu et al. 2000). Our work assessed the short-term working memory and long-term reference memory in the transformer's ' nav- igation task by measuring the visited error and unvisited error, respectively.
The modulation of selectively affects the formation of long-term reference memory (i.e., predic- tion 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 ford-forward layer. A similar idea has been proposed in psychology, which we illustrate in detail in Appendix A.5
In neuroscience, the transfer of short-term memory into long-term system is called memory con- solidation (McGaugh 2000). Various research has revealed that Mg2--gating of NMDA receptors modulates the formation of long-term memory (Slutsky al., 2010 Miyashita et al. 2012 These observations imply that the nonlinear dynamics of NMDA receptors in hippocampus CAI are critical for consolidating short-term memory into long-term memory.
Current research showed that the transformer could model memory consolidation. We assumed that the GELU activation function links short-term working memory and long-term reference mem- ory. Our experiments indicate that the formation of long-term reference memory is impaired when the activation function completely linear (corresponding to no Mg²+). In contrast, increasing a (which corresponds to an increase in Mg level) has resulted in superior performance in long-tew -eemtermm reference memory compared to other activation functions (c.g., ReLU, GELU, LeakyReLU, sig- moid, tanh). These similarities between hippocampal memory consolidation and our results suggest a transformer as an effective memory consolidation model.
We have investigated the role of NMDAR-like nonlinearity with novel consideration of activation function in a transformer model and have demonstrated the effects of altering the nonlinear dynamics of the activation function. Even though there are trainable parameters in the self-attention layer, the quantitative analysis of the place cell score indicates that most of the reference memory is stored in bred-forward networks. Surprisingly, our result that loss of nonlinearity in NMDAR-like activation function (a = 1) show impaired reference memory formation (Fig. (3)) but intact working memory formation (Fig. 3 seems to replicate the previous findings in neuroscience; selective inhibition of hippocampal CA1 NMDAR inhibition does not disrupt working memory (Lee & Kesner (002) but impairs the long-term memory formation (Tsien et al.) 1996). These similarities provide an exciting possibility that our brain selectively updates short-term working memory into long-term reference memory by activity-dependent Mg 2 2+ -gating of NMDAR.
Future directions Our research has exciting future directions. The current study only examined what-where memory using a sensory observation task in a static environment. However, our real- world environment is changing dynamically. Modern deep learning systems are generally incapable of adapting to a dynamic environment or reordering of sensory inputs. In future work, we intend to explore what-where- when memory, called episodic memory, in transformer and other deep models.

---

Under review as a conference paper at ICLR 2023
7 REPRODUCIBILITY STATEMENT
We included the Pytorch implementation code available along with the paper (see the Supplemental Material file). For reproducibility, we included thorough overview of the experimental setup, task design, and place cell score evaluation pseudo-code in the Appendix.
REFERENCES
Philippe Ascher and Linda Nowak. The role of divalent cations in the N-methyl-D-aspanee sponses of mouse central neurones in culture. The Journal of physiology. 399(1):247-266, 1988.
Richard C Atkinson and Richard M Shiffrin. Human memory: A proposed system and its control processes. In Psychology of learning and motivation, volume 2. pp. 89-195. Elsevier, 1968.
Alan D Baddeley and Graham Hitch. Working memory. In Psychology of learning and motivation, volume 8, pp. 47-89. Elsevier, 1974.
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(77055::2994333 2018.
Tim VP Bliss and Graham Collingridge. synaptic model memory: long-term potentiation in the hippocampus. Nature, 361(640)) 1993.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Ma- teusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models shot learners. In Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901, 2020. URL https: http:ss//pproceedngs..eeeee neuripss.......un.o.n.....m..... wttp://pppp...ppppp///////// /ctps/////perr/////lleee 1457c0d6bf Cb496 741 2ttp:///.....................22.
few-
Christopher Cueva and Xue-Xin Wei. Emergence of grid-like representations by training recur- rent neural networks to perform spatial localization. In International Conference on Learning Representations, 2018. URL /tt/://cte////eeeeeeeeeeeeeeetee het.:e/we.ettttrrmoeeee/111000
Zihang Dai, Zhilin Yang. Yiming Yang, Jaime G. Carbonell, Quoe Viet Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond fixed-length context. In Association for Computational Linguistics, pp. 2978-29888 2019.
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.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 18Xiv:1810.04880, 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthiss Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 2020.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Nerworks, 107:3-11, 2018.
Torkel Hafting, Marianne Fyhn, Sturla Molden, May-Britt Moser, and Edvard I Moser. Microstruc- ture of spatial map in the entorhinal cortex. Nature, 436(7052):8018866 2005.
Dan Hendrycks and Kevin Gimpel. arXiv: 0ee116660881151 2016.
Gaussian error linear units (gelus). arXiv preprint
10

---

Under review as a conference paper at ICLR 2023
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
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-417, 200.
Craig Jahr and Charles F Stevens. Voltage dependence of nmda-activated macroscopic conduc- tances predicted by single-channel kinetics. Journal of Neuroscience, 1009)3177-31188 1990.
Clifford Kentros, Eric Hargreaves, Robert D Hawkins, Eric R Kandel, Matthew Shapiro, and Robert V Muller, Abolition of long-term stability of new hippocampal place cell maps by NMDA 1998.
receptor blockade. Science, 280(5372)21111-2222
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. URL
6980
Eilon D Kirson, Claudia Schirra, Arthur Konnerth, and Yoel Yaari. Early postnatal switch in mag- nesium sensitivity of nmda receptors rat cal pyramidal cells. The Journal of Physiology, 521 (Pt 1):99, 1999.
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(8):3612-3627 1993.
Dmitry Krotov and John J. Hopfield. Large associative memory problem neurobiology and machine learning. In International Conference Learning Representations, 2021. URI.
htt.ps:://ppenev/
net/ Forum?id=xmy 10oxXx4__10xxhhh
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. Albert: A lite bert for self-supervised learning of language representations. arXiv preprint arXiv: 1942, 2019.
Inah Lee and Raymond Kesner. Differential contribution of nmda receptors in hippocampal sub- regions to spatial working memory. Nature neuroscience, 5(2):162-168, 2002.
Fei Li and Joe Z. Tsien. Memory and the nmda receptors. The New England journal of medicine, (3)33), 2009.
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.
Mark L Mayer, Gary L Westbrook, and Peter Guthrie. Voltage-dependent block by mg2+ of nmda responses in spinal cord neurones. Nature, 3096566)):611-66,, 1984.
ML Mayer and GL Westbrook. Permeation and block of -mnethy----aspartie acid receptor channels by divalent cations in mouse cultured central neurones. The Journal of physiology. 394(1):501- 527, 1987.
James I. McGaugh. Memory-a a century of consolidation. Science, (:755)5)):248--1 2000.
Thomas J McHugh, Kenneth I Blum, Joe Z Tsien, Susumu Tonegawa, and Matthew A Wilson. Impaired hippocampal representation of space in cal-specific umdarl knockout mice. Cell, 87 (7):1399-1399, 1996.
Tomoyuki Miyashita, Yoshiaki Oda, Junjiro Horiuchi, Jerry CP Yin, Takako Morimoto, and Minoru Saitoe. Mg2+ block of drosophila NMDA receptors required for long-term memory formation and CREB-dependent gene expression. Newron, 74(5):887-898, 2012.
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Jeml, 2010.
LPPAA Nowak, Bregestovski. P Ascher, A Herbet, and Aa Prochiantz. Magnesium gates dlutamate-activated channels in mouse central neurones. Nature, 307(5950):46244655 1984.

---

Under review as a conference paper at ICLR 2023
John O' Keefe and Jonathan Dostrovsky. The hippocampus spatial map: Preliminary evidence from unit activity in the frvely- rat. Brain research, 1971.
David S Olton, Christine Collison, and Mary Ann Werz. Spatial memory and radial arm maze performance of rats. Learning and motivation, 8(3):299311,, 1977.
David Olton, James T Becker, and Gail Handelmann. Hippocampus, space, and memory. Behavioral and Brain sciences, (3)):1133222, 1979.
MICHAEL Perouansky and YOEL Yaari. Kinetic properties of NMDA receptor-mediated synaptic currents in rat hippocampal pyramidal cells versus interncurones. The Journal of Physiology, 465 (1):223-244, 1993.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, pp. 8748-8763. PMLR, 2021.
Prajit Ramachandran, Barret Zoph, and Quoe V Le. Searching for activation functions. arXiv preprint arXiv: urXiv:/710.0594|, 2017.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 2022.
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gru- ber, Markus Holzleitner, Thomas Adler, David Kreil, Michael K Kopp, Günter Klambauer, Jo- hannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need. In International Conference on Learning Representations, 2021. UtLpttps: hptp///pppeeeeeeeeeeeeeeeeeeeeee /
net forrmm?r
id=t 1.8 9Rnz. Cc
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(7009):699-7000 2004.
Eiji Shimizu, Ya-Ping Tang, Claire Rampon, and Joe Z Tsien. Nmda receptor-dependerr synaptic reinforcement as crucial process for memory consolidation. Science, 290(5494) 2000.
Inna Slutsky, Nashat Abumaria, Long-Jun Wu, Chao Huang, Ling Zhang, Bo Li, Xiang Zhao, Arvind Govindarajan, Ming-Gao Zhao, Min Zhuo, et al. Enhancement of learning and mem- ory by elevating brain magnesium. Neuron, 65(2):165-177, 2010.
Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Herve Jegou, and Armand Joulin. Aug- menting self-attention with persistent memory, 2019. URL htttps:/arxxiv.org:
1907. 1907001470
Joe Z Tsien, Patricio T Huerta, and Susumu Tonegawa. The essential role of hippocampal cal nmda receptor-dependent synaptic plasticity in spatial memory. Cell, 87(7):1327-1338, 1996.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz. Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
James C. R. Whittington, Joseph Warren, and Tim E.J. Behrens. Relating transformers to models and neural representations of the hippocampal formation. In International Conference on Learning Representations, 2022. URL
net.
YEO
James CR Whittington, Timothy H Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy EJ Behrens. The colman-eichenbann machine: unifying space and relational mem- ory through generalization in the hippocampal formation. Cell, 183(5):
2020.
Ann M Woodhull. Ionic blockage of sodium channels in nerve. The Journal of general physiology, 61(6):687-708, 1973.
12

---

Under review as a conference paper at ICLR 2023
A
APPENDIX
A.1 DERIVATION OF NMDAR NONLINEARITY FROM THE MOLECULAR LEVEL CHEMICAL
INTERACTION
Here, describe the NMDAR nonlinear dynamics from chemical interaction between Mg and NMDAR following previous literature Woodhull y Jahr & Stevens 1990); 1990) Perouansky & Yaari 1993). At the molecular level, Mg ion binds to one NMDAR receptor when opening the NMDAR channel. Thus, the chemical equation of binding reaction between Mg 2+ ion and NMDAR receptor, R, be described
Mg
R
Mg 2 R.
(10)
From this chemical equation, the equilibrium constant is given by
MMg
R
K
(11)
Mg
[R]
Thus, dissociation constant KD, which correspond to Mg dissociation from NMDAR, follows
Mg
2+
R
KD = --
(12)
Mg
R
in which [R] and [Mg2+R] RJ are the free and Mg -bound NMDARs respectively. The fraction of opened NMDAR channels (number of unbound NMDAR over number of total NMDAR) at equilibrium follows,
Mg2+-bound
R
R
+ Mg 2+ R
(13)
1
1 Mg
1/KD
Experimentally, the voltage-dependent dynamics of Kp has been described following equation by Ascher & Nowak (1988)
KD KMg²1 @V
(14)
where, V is membrane voltage, / is temperature constant and M is dissociation constant at V = Eq is substituted into Eq [3] geltageddeendenn open fraction of NMDAR can be expressed as follows
P(V)
+
Mg Kxa++
9V
(15)
1
1 DV
in which Mg 21 1gg+*1KKgaar M/Kg the parameter determined by the [Mg2*]. Given the voltage- dependent open fraction of NMDAR, p(V), and NMDAR's maximal conductance, Imax, then voltage-dependent NMDAR conductance g(V) can be described as
(VV) (v)
(16)
Given (VV), and driving force, V- V,, and current I, they have relationship of [ (V-VV))
in which V, is reversal potential (the value membrane potential above which current inverts the direction). As experimental investigations the physiological reversal potential of NMDAR to be
13

---

Under review as a conference paper at ICLR 2023
V, = 0 Mayer & Westbrook (1987; Ichinose et al. )2003) Kohr et al. 1993
Vg(V). Then,
the normalized NMDAR current nerm 1/ 9max follows
norm = VV())
(17)
From Eq. 17 and previous electrophysiological experimental results Kirson et al. (1999), we recon- struct IV curve in (Fig. Th right top).
A.2 NMDAR-INSPIREH NONLINEAR ACTIVATION FUNCTION
Here, propose AMDAR-inspiree nonlinear activation function from the nonlinear dynamics of the NMDAR-IV curve. If we consider the nonlinear IV curve of NMDAR (Eq. [7) as a nonlinear mapping function, o, the membrane voltage, V, can be viewed NMDAR current, Jorms as an output, ((x). Then we can rewrite the nonlinear mapping function ø follows
input, T, and normalized
(xx) = rp(2).
(18)
We define the NMDAR-inspired activation function as a nonlinear mapping function, NMDA())==(()). By substituting Eq. 15 into Eq. 18 we show the generalized expression of (MDA(2) equation with a and 3 parameters as following
NMDA (z) =xPape())
(19)
1 + ae- fxx
Given a = and [ p(x) identical to sigmoid function,
--) This particular
case of and 3 leads to zor (z)) Sigmoid Linear Unit (SiLU) activation function Elfwing et al. 2018 In the case of a
and 3 = 1.702, xo (02x))) correspond to the GELU activation function, which is popular in transformer models Lan al. (2019); Liu al. (2019 Lan et al. 2019). RRamachandaan et al. (017) introduced the swish activation function, (Brr), which is generalized form of GELU and SiLU. They demonstrated that when
oo, the activation function resembles RELU. We summarized these four activation functions by comparing them with our NMDA (z) in table and 6
In contrast to the extensive research on 39 in N a, the Mg component, is not explored. For this reason, we focused on the parameter a over 3, and investigated NMDA, is interesting to note that the Swish function was originally proposed as self 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 Hochreiter & Schmidhuber 1997]. In contrast, our activation function NMDA(z) is inspired by the physical Mg=T-gating mechanism that occurs in the real biological synapses. These shared mechanisms of self-gating from the artificial models and biological observations raise an interesting possibility that NMDAR is a neural substrate of nonlinear activation function in the brain.
A.3 DETAILED DESCRIPTION OF TASK DESIGN AND DEFINITION OF SHORT-TERM WORKING MEMORY AND LONG-TERM REFERENCE MEMORY
Our task is based on widely employed neuroscience experiment for spatial working memory and reference memory Olton al. (1977 1979). Errors in working memory measured by within-trial error, whereas errors in reference memory are measured by across- caross-trial-o error. The training phase and
Table Comparison of common activation functions with NMDA
NMDA
Name
Equation
Reference
NMDA, 1,8 (x)
SiLU(x)
(x))
Elfwing et al. (2018)
NMDA, 1.8= 1.702
GELU(X)
xo #1.70022
Hendrycks & Gimpel (0016)
RELU(x)
max Φ)
Nair & Hinton (2010)
NMDA,=1,0(1
Swish(x)
No (3x)
Ramachandran et al. 2017)
14

---

Under review as a conference paper at ICLR 2023
NMDD
GELU
GELU
SULU Swish(B-11
1.2
SILU, Swishg 11
RELU
RELU
1.0
0.5
3
0.6
2
0
-0
Figure 6: Comparison of common activation functions (left) and their derivatives (right) with NMDA
the test phase alternate 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-term reference memory (LTRM). Within trial visited place prediction measures relatively short-term experience for our task. On the other hand, the across-trial unvisited place prediction task in the familiar map measures the relatively long-term experience. Measuring unvisited place prediction error the novel map will establish baseline of chance-level accuracy; above this baseline, the formation of long-term memory can be observed (Fig. 7
15

---

Under review as a conference paper at ICLR 2023
Familiar maps
Novel maps
Agent
Trajectory
Within trial
Trial1
Training phase
Tria12
Trial3
Epoch
Query (Batch size)
Across trial
Test phase
Unvisited
Unvisited
Unvished
Unvisited
Visited
Unvisited
Epoch3 Epoch2
Training phase
Test phase
Training phase
Test phase
Figure 7: 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 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 a random walk.
16

---

Under review as a conference paper at ICLR 2023
A.4 NON-RECURRENT POSITIONAL EMBEDDINGS AND PREDICTION ERRORS ON THE NODE VISITED FOR THE FIRST TIME
We test the non-recurreh positional embedding by substituting the recurrent positional embedding e, with the action embedding A(a;), where A is the embedding layer and a, is the action at step Compared to Fig. Ra, the result sdemonstrates a significant increase in working memory error and reference memory error (Fig. 3 vs. Fig. { Nonetheless, its behavior is comparable to the trend of decreasing reference memory error while increasing a of NMDA, (see Fig. (D).
In addition, compare unvisited node prediction error (unvisited within context window, in this case, 64 steps) versus first visited node prediction error (unvisited for within a trial). As shown in results for the first visited node do not differ from the reference memory
Fig. 2 the prediction eπrо results.
These results strongly support that (I) while the path-integrated information from recurrent posi- tional embedding is important for learning the spatial structure of the map, this information is not used in predicting the unvisited node, and (2) the reference memory is used for predicting the unvis- ited node in a familiar map.
(a)
(familiar)
(novel)
Reference memory error
0.0
0.01
0.05
(b)
0.4
1.0
2= 10.0
10.0
Reference memory
0.6
0.2
0.5
Working memory error
0.4
0.2
0.3
0.0
0.2
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
train
x10
train
x10
Figure 8: Experiment with non-recurrent positional embeddings. (a) Reference and working mem- ory eιrоrs over training trials for training (familiar) maps and testing (novel) maps for N = 32 where N the number of training maps. (b) Reference memory errors evaluated on training maps over different values of in NMDA, for N 32. Error bars and shaded areas represent the standard deviation of errors from three independently trained models.
First visited node cradnction arror
0.8
48
0.6
0.4
0.2
0.0
0.0
0.01
0.05
0.1
0.5
1.0
5.0
10.0
Figure 9: 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 errors from three independently trained models.
17

---

Under review as a conference paper at ICLR 2023
A.5 TRANSFORMER AS MEMORY CONSOLIDATION MODEL AND ITS BIOLOGICAL PLAUSIBILITY
Next, examined the biologically inspired NMDA,, activation function in the feed-forward layer of the transformer and its role in memory formation and place cell representation. We show that modulating a corresponds
change in extracellular [Mg"* by deriving the nonlinear activation function from the real NMDAR nonlinear IV curve. We show the reconstructed real nonlinear IV curve in Fig. I (right top).
The modulation of selectively affects the formation of long-term reference memory (i.e., predic- tion of unvisited places across trials) while leaving the formation of short-term working memory (i.e., prediction of unvisited places within trials) intact. This result suggests that short-term working memory and long-term reference memory are physically stored in separate structures: the self- attention layer and the feed-forward layer. In psychology, the idea of multi-store model regarding short-term memory and long-term memory was historically suggested in Atkinsoo & Shiffrin 1968). In their model, sensory inputs are stored in short-term memory systems via attention, and some are transferred to a long-term memory system while others quickly disappear.
In neuroscience, the transfer of short-term memory into long-tem term system is called memory consol- idation McGaugh (2000). Animal studies have demonstrated that the CA1 region of the hippocam- pus is essential for memory consolidation (Shimizu et al.] (2000); [Remondes & Schuman (2004). In hippocampal CA1, 1, the postsynaptic NMDA receptor mediates synaptic plasticity, and the selec- tive perturbation of these receptors leads impairment long-term memory formation [sien al.
[196]; |Remondes & Schuman (2004). Later research revealed that Mg2+-gating of NMDA re- ceptors modulates the formation of long-term memor [Slutsky et al. (2010); Miyashita et al. (2012). These observations imply that the nonlinear dynamics of NMDA receptors in CAI are critical for consolidating short-term memory into long-term memory.
On the basis of a previous link between the hippocampus and the transformer, the current research hypothesized that the transformer could serve as model for memory consolidation. Given the resemblance of the GELU nonlinear activation function and CA1 NMDAR nonlinear IV curve, we assumed that the GELU activation function serves as a key component that links short- term working memory and long-term reference memory. Our experimental results indicate that that the formation of long-term reference memory is impaired when the activation function is completely linear (corresponding
no Mg ). In contrast, increasing (which corresponds to an increase in Mg level) has resulted in our model's superior performance in long-term reference memory com- pared to other activation functions (e.g., RELU, GELU, LRELU, Sigmoid, Tanh). Based on these similarities between hippocampal memory consolidation and our results, we propose a transformer as an effective memory consolidation model.
In addition to the performance gain in long-term memory formation with NMDA.. we found that modulating the a affects the emergences of place cells in the feed-forward layer and conclude significant correlation between place cell score and long-term reference memory formation. Our results align with previous biological findings that perturbation of CA1 NMDARS lead impairment in both place cell representation and long-term memory formation Tsien et al. (1996); McHHugh et al.) (1996); Kentros et al.] [998]; Shimizu et al.) (2000p. These similarities support the idea that place cells are the neural correlates of long-term spatial memories. Altogether, our findings suggest an exciting possibility that the nonlinear IV curve of NMDAR in the hippocampal CAI is neural substrate of nonlinear activation function in the brain.
18

---

Under review as a conference paper at ICLR 2023
A.6
PSEUDO CODE FOR CALCULATING PLACE CELL SCORE METRIC
Algorithm 1:
Pscudo code for calculating place cell score metric
function PlaceCelIScore(place 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; node;) in G do if firing rate p; > firing rate Pj then G add edge (node; node;) end
end
Find node; of firing rate Pmor 10 for node,, in G do
if node, is not descendant of node: found with DFS(k) then delete node, from G end end 15 conn. components = sum of all nodes' firing rates in g total components
12
of all nodes' firing rates in G
comn. components
1 return place score place score = total components
18 is discount factor, determined by connected component size
The place field in Algorithm | is measured following procedure: 1) During a random walk sim- ulation, the activation value of neuron node i, where the agent located, is measured every 65 steps. Let's say this value is a; 2) Every time the agent visits node again, value a; is added cumulatively to the recorded value; A,+ = a, such that A; is the cumulative activation value at node i. We assume the initial value of A; is zero. After the random walk is done, A, divided by the length of the random walk trajectory is the firing rate at node of the neuron (place field E K K ). In place cell evaluation experiment, the length of the random walk is 10 and K 11; the evaluate map is one the training maps.
A.7 CONSEQUENT OF CHANGING NONLINEAR DYNAMICS IN LEAKY RELU ACTIVATION FUNCTION
Here, we investigated the consequence of changing nonlinearity with other than NMDA,- Here, we choose Lea (max(0, .)) + a min(0, x)) activation function to compare with NMDA., 1 of LeakyReLU also leads to fully linear activation function similar
Regarding LeakyReLU, to & = 0 of NMDA,- Compared to NMDA.= LeakyReLU exhibits lower average place score in the allowed range of a, indicating that NMDA, is better in place cell emergence (see 10)
19

---

Under review as a conference paper at ICLR 2023
LeakyReLU
Reference memory error
0.2
0.1
0.01
0.02
0.05
0.1
0.2
a
LRELU
LRELU a 0.1
Referanec maneory croIt
0.8
LRELU a 0.2
LRELU = 0.05
LRELU o 0.1
0.3-
LRELU a 0.01
LRELU a 0.02
URELU
Probabilty
0.2
0.1-
0.0
0.2
CA
0.8
0.2
0.3
Payce cell spore
Averoge place ocl soore
Figure 10: Evaluation of reference memory error in LeakyReLU (LRELU) while modulating a(top) and relationship of average place cell score and reference memory error (bottom).
20


Reviews:
[{"review_id": "v9CfEHz0DC0", "reviewer": "Reviewer_o25e", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes to model the influence of the concentration of magnesium ions on the $Mg^{2+}$-gated NMDA (a neurotransmitter) receptors nonlinear dynamics - involved in several functions, especially place cells representations,  which are considered important for spatial navigation. The NMDAR nonlinearity is modeled with a GELU-like function. This paper experimentally shows on a 2D grid exploration task how this nonlinearity allows for the appearance of place cells-like structures in the feed forward network of a transformer model.", "summary_of_the_review": "While this work present highly interesting ideas for the computational neuroscience community, several points are unclear which hinders the understanding of the experiments. Also, the main claim is not sufficiently supported by the experiments. I would be more than glad to increase the score for this paper once the clarifications are made and the main claim is firmly supported (running the experiments with ReLU/tanh/sigmoid/leaky ReLU). Or defer my judgement to an additional reviewer from the neuroscience field.\n\nEDIT:\nAfter clarification, and changing the claim, the authors correctly addressed my concerns and correctly support their work. I hence increased the score.", "strength_and_weaknesses": "### Strengths\n\n- The link between GELU and NMDAR dynamics modulated by the concentration $[Mg^{2+}]$ is highly interesting.\n- Their model relies on state-of-the-art computational modelling of neurological functions\n\n### Weaknesses\n\n- Several parts of the paper are unclear (details in the next field).\n- The main conclusion of the experiments and article is too extrapolated:\n“Our data indicated that **NMDAR-like nonlinearity** in the feed-forward network layer of transformers **is necessary** for long-term memory and spatial place cell representation.”\nThe paper explores what they defined as the $NMDA_{\\alpha}$-family of nonlinear function only. It appears that the conclusion drawn is based on the comparison when $\\alpha=0$ vs. $\\alpha > 0$. Since the former case corresponds to a linear activation function, the conclusion to be drawn from the experiments is that such activation function prevents from place cells-like structures and long-term memory to appear. To sustain their claim, the authors should have compared different nonlinear activation functions with their $NMDA_{\\alpha}$. As such, the main claim of the article does not seem sufficiently supported.\n- The choice to consider “a node that the agent had never visited within recent 64 steps is treated as an unvisited node.” (page 4) seems lacking. A node visited during a trial on which the model has done a gradient step can be considered as depending on the reference memory, but what if the node is visited for the first time? The model would predict at random level, so no memory could be considered involved.\n- Page 4, on the Figure 3: “This finding suggests that the reference memory is non-active for predicting the visited nodes on novel maps.” I find this claim to be confusing. If the test maps are novel and **really** used as a test set and no optimisation step is performed, there is no way the model can integrate information on that map, hence no way for the model to develop a reference memory of the test map. The model can only perform at chance level when confronted to a new environment.\n- Page 6, section 4: “On a related note, Whittington et al. (2022) showed that softmax neurons in the self-attention layer behave like place cells and demonstrated that changing the softmax function to linear slows the learning process in the working memory.” But your experiment shows no particular appearance of place cells-like structure in the self-attention layers, right? On what condition do they appear in the feed forward layers rather than in self-attention layers? Without clarification, it feels like place cells can appear anywhere.\n\nMore of a suggestion rather than a weakness:\n\n- Giving the number of reference memory error vs. working memory error in addition to the rates (as given in Figure 3) would be informative.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "psLtKnKwwp", "reviewer": "Reviewer_oq3h", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this paper, the authors showed that place cells emerge in the feedforward layer of the Transformer that uses 1) NMDA-$\\alpha$ nonlinearity and 2) recurrent positional encoding when trained on the sensory observation prediction task. Moreover, they showed that bigger $\\alpha$ in the nonlinearity is simultaneously correlated with better reference memory and better place cell scores.\n", "summary_of_the_review": "It is a well-written paper with interesting findings and very clear communication, and can benefit the neuro-AI community. However, I believe the scope of this paper is more suited for a workshop or smaller conference. Thus I'd recommend rejecting the paper for ICLR.", "strength_and_weaknesses": "Strengths:\n- The paper presents a well-defined problem with clear communication and writing and easily understandable figures. It continues the line of research on the connections between Transformers and models of the hippocampus and takes inspirations from neurobiology to develop a new model for the hippocampus.\n\nMajor Weaknesses:\n- This paper only used one nonlinearity, which is NMDA-$\\alpha$. It would be helpful to include some comparison to other types of nonlinearity commonly used in hippocampal-entorhinal models.\n- This paper claims that NMDAR nonlinearity is needed for *long-term memory* when  it only shows the effect of $\\alpha$ on reference memory. What about other types of long-term memory, such as episodic memory? \n- What are the mathematical implications of replacing GELU with NMDA-alpha? It would strengthen the paper to see a comparison between GELU-based Transformer and NMDA-$\\alpha$ based Transformer in terms of behavioral performance and place cell representations.\n- Figure 3a shows that the model doesn’t perform well on predicting unvisited node in Novel Map, which, in my understanding, means that the model cannot learn latent structure of the map and do flexible binding, which is what the the original Tolman-Eichenbaum Machine (Whittington et al. 2020) is *supposed* to do. This should be discussed more in details in the discussion section.\n- Figure 3b: If I’m understanding correctly, training the model on *more* maps leads to *bigger* errors? Is there an explanation for this?\n- There might exist alternative explanations for why the model can predict nodes that haven’t been visited in the past 65 steps in a familiar map, other than reference memory. For example, such behavior could be explained by path integration, which can be attributed to recurrent positional encoding? As such, it might be interesting to see whether non-recurrent positional encoding gives similar results.\n\nMinor weaknesses:\nThere are some weaknesses inherited from TEM-t: \n- The model doesn’t account for some well-established hippocampal phenomenon, eg. replay; \n- The sensory prediction task is a toy task with simplified, pre-digested inputs. It’s hard to know whether the model is compatible with other types (eg. image or video) of inputs, or whether it can scale up to work with tasks more suited for current machine learning climate, such as Atari.\n\nQuestions:\n- TEM-t (Whittington et al. 2022) introduced 3 modifications to the original transformers. If I'm understanding correctly, you didn’t use TEM-t, but Transformer with recurrent positional encoding (that, and NMDA-$\\alpha$ instead of GELU). Then how do the other two modifications affect pace cell properties?\n- What are the biological implications of changing alpha? The biological parallel of NMDA-$\\alpha$ to Mg2+ gated ion channels seemed a little far-fetched to me.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "6E48gzMfv8", "reviewer": "Reviewer_91tN", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper applies the transformer model to spatial navigation problem in a grid world with labeled grid positions. The task is to predict the label of the next position that is either visited or unvisited. The paper connects this task to working memory and reference memory, as well as place cells. The finding is that NMDAR-like nonlinearity in the feedforward block of the transformer model is important for reference memory and neurons behave like place cells. ", "summary_of_the_review": "The paper compares transformer and hippocampus in navigation tasks, and the findings are interesting. However, the paper is too empirical with no theoretical investigation. ", "strength_and_weaknesses": "Strengths: \n\n(1) The analogy between the transformer model for prediction and the hippocampus in the brain is interesting, although this analogy has been explored in a recent paper. \n\n(2) The focus on the feedforward block of the transformer model seems novel, and the connection between the GeLU non-linearity and the NMDA receptor (NMDAR) is novel. \n\n(3) The idea and method in this paper is simple and interesting. \n\nWeaknesses: \n\n(1) The paper is entirely empirical. There is no theoretical or mathematical analysis. The empirical similarities between transformer and hippocampus are noteworthy, but some theoretical understanding can greatly improve the paper. \n\n(2) The focus on transformer architecture is understandable given its empirical successes and its popularity, but it may be worthwhile to explore simpler models that may have similar behaviors. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "18l8mawktl", "reviewer": "Reviewer_7x8r", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper draws connections between the NMDA receptor (NMDAR) in the hippocampus and the GELU activation function, which has been employed in the transformer architecture. The paper then proposes a novel activation function, which more closely resembles the behavior of the NMDAR, and shows that the transformers memory capabilities can be tuned with the hyperparameter of this activation function. In particular, the paper investigates the working memory (i.e., in-context memory of states observed during the current trajectory) and the reference memory (i.e., out-of-context memory of states observed during previous trajectories) of the transformer on a 2D grid navigation experiment. Finally, the paper also proposes a place cell score, which measures the sparsity of a neuron’s connections, and shows that it can also be tuned with the hyperparameter of the proposed activation function.", "summary_of_the_review": "Given the ill-defined memory metrics and tenuous connection to place cells, I do not recommend the paper for acceptance at this time.", "strength_and_weaknesses": "**Strengths**\n\nThe paper draws on insight from neuroscience to explain the behavior of the transformer architecture, which is interesting. Moreover, the experimental evaluation showcases an interesting connection between the sparsity of connections in the feed-forward network of the transformer and the proposed activation function.\n\n**Weaknesses**\n\nThe two considered memory types are ill-defined, leading to void claims:\n* In the case of working memory, any nonlinearity should be capable of decreasing the working memory error on test maps (as evidenced by the fact that the error is largely independent of the proposed activation function’s hyperparameter) due to the global attention mechanism of transformers (i.e., as long as the observation is within the context window). Thus, the connection to NDMARs is tenuous.\n* In the case of reference memory, evaluating the model on unseen maps is an ill-defined problem, since the unvisited places are inherently unpredictable. Thus, claiming that `reference memory is non-active for predicting the visited nodes on novel maps` or that `reference memory formation requires NMDAR-like nonlinearity` does not make sense.\n\nThe place cell analogy is tenuous, since the proposed score only measures sparsity and not the location in the 2D grid environment, unlike place cells. The results on the sparsity are still interesting, but the claim relating them to place cells has to be revised.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject