Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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. 
strength_and_weaknesses: Strengths: 

(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. 

(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. 

(3) The idea and method in this paper is simple and interesting. 

Weaknesses: 

(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. 

(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. 

clarity,_quality,_novelty_and_reproducibility: The key idea, method and findings of this paper are original. 

The quality of empirical study is good. 

But the clarity of the presentation can be improved. The main text should provide more details about the transformer architecture used in this paper. Some discussions on its biological plausibility may also help. 
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.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: 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.

strength_and_weaknesses: Strengths:
- 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.

Major Weaknesses:
- 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.
- 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? 
- 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.
- 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.
- Figure 3b: If I’m understanding correctly, training the model on *more* maps leads to *bigger* errors? Is there an explanation for this?
- 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.

Minor weaknesses:
There are some weaknesses inherited from TEM-t: 
- The model doesn’t account for some well-established hippocampal phenomenon, eg. replay; 
- 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.

Questions:
- 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?
- 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.

clarity,_quality,_novelty_and_reproducibility: This paper is very clearly written with original ideas built upon current hippocampal models. Since the code is provided, I assume the work is reproducible.
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.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: 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.
strength_and_weaknesses: **Strengths**

The 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.

**Weaknesses**

The two considered memory types are ill-defined, leading to void claims:
* 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.
* 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.

The 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.

clarity,_quality,_novelty_and_reproducibility: **Clarity**

Several aspects of the paper are unclear:
* What do the NMDAR IV curves in Figure 1 a) correspond to? The NMDAR open probability lacks a factor of $x$ if it is supposed to model the behavior in the figure.
* The paper does not justify why the positional encodings, which are particularly relevant in a 2D navigation setup, are omitted.
* The inset working memory plot in Figure 3 a) is quite confusing and should be presented as a standalone figure.
* Figure 3 b) does not specify whether the reference memory error is measured on train or test maps.
* An illustration of the place cell score computation would help the reader understand the metric.

**Quality**

I am not sure if ICLR is the appropriate venue for the paper, given both its topic and length (i.e., 6.5 pages).

**Novelty**

To the best of my knowledge, the proposed activation function and connection to the NMDAR have not been investigated before.

**Reproducibility**

Given the lack of clarity of the experimental setup (see above), I am not convinced that the results would be entirely reproducible.

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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
strength_and_weaknesses: ### Strengths

- The link between GELU and NMDAR dynamics modulated by the concentration $[Mg^{2+}]$ is highly interesting.
- Their model relies on state-of-the-art computational modelling of neurological functions

### Weaknesses

- Several parts of the paper are unclear (details in the next field).
- The main conclusion of the experiments and article is too extrapolated:
“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.”
The 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.
- 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.
- 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.
- 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.

More of a suggestion rather than a weakness:

- Giving the number of reference memory error vs. working memory error in addition to the rates (as given in Figure 3) would be informative.
clarity,_quality,_novelty_and_reproducibility: ### Clarity

Several points seems unclear:

- In the 2.2 and 3.1 subsections, the task is insufficiently detailed. How are the trials conditioned in batches (a figure would be highly appreciated)? What loss function is being used? How are the validation and test sets created and used?
- Page 3: “the initial positional embedding $e_1$ is sampled from a normal distribution”. Can you explain this choice? Why not using a special token?
- Page 4, subsection 3.1, you call $e$ a “positional embedding” whereas you introduced it as an action embedding.
- Page 4: “We run 512 random walk simulations in parallel for collecting training trajectories. The total number of random walking steps is 2,048 for each simulation so the total number of gradient steps for each run is 512 (batch size) × 2,048 (total number of steps in a trial) × 200 (number of trials)”. This part is quite unclear but a better description of the task (as already suggested) would solve it.
- Page 5, subsection 3.3: The explanations are given in reverse:
    - you explain the environment on which you measure the place cell score before defining the place cell score;
    - you define the place cell score before introducing the variables;
    - the notion of firing rate $\rho_i$ is undefined in the paper.
- Page 5, subsection 3.3: “by defining a $K \times K$ 2D grid environment”. It seems implicit but it is new, isn’t it?

### Quality

The figures are explanatory and more of them would be highly appreciated. Further work on the topic could give valuable understanding on the comparison between the brain and artificial neural networks.

### Novelty

The comparison between the GELU activation function and NMDAR dynamics is novel to the best of my knowledge and researches and this model could benefit to the computational neuroscience community.

That being said, other work in computational neuroscience might have tackled this modelling and did not appear when I researched for them. This paper would highly benefit from the review of an expert in neuroscience.

### Reproducibility

The code runs, which is already good sign of reproducibility. More time would be required to determine complete reproducibility. The code is well written and answers a few questions.
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.

EDIT:
After clarification, and changing the claim, the authors correctly addressed my concerns and correctly support their work. I hence increased the score.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: The paper makes interesting connections between the NMDA receptor in the Hippocampus and gelu nonlinearity in transformers. In response to reviewer's comments, the authors toned down the statement of "transformer needs ...". Indeed Relu nonlinearity works pretty well, just slightly worse than gelu (and NMDA with the right setting), consistent with researchers switching from the former to latter to improve the performance. 
Overall the reviewers find this paper interesting, but not good enough and not the right venue for iclr, but likely a good workshop paper or a more neuroscience oriented venue.
justification_for_why_not_higher_score: While the paper is interesting, it is not strong enough for iclr (and not the best venue for it).
Decision: Reject