Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This well-written and clear paper clarifies the relationship between transformers and a recent exciting model of the medial temporal lobe in neuroscience. The authors find that a transformer with recurrent positional encodings trained on a spatial navigation task ends up learning representations that resemble some classic findings from the medial temporal lobe in the neuroscience (e.g., grid and place cells). This finding alone is not that surprising -- as the authors note, a variety of groups have previously found such representations are learned in navigation tasks with other models (particularly place cells). The major contribution of this paper is that the authors then demonstrate the connection between transformers and a contemporary model of the hippocampal formation, the Tolman-Eichenbaum Machine (TEM). Specifically, the memory retrieval process in the TEM resembles self-attention and the path integration representations are comparable to a transformer's learned position encodings. I believe that by offering a novel perspective on transformers and models of biological memory systems, this work may lead to fruitful future work both in machine learning and neuroscience.

main_review: Major issues:
- Many papers have reported finding learned representations that resemble place cells and some have found grid cells as well. It would be quite helpful to contextualize the findings in this manuscript with others, as well as more rigorously evaluate the degree of evidence for these types of response properties, by: (1) defining a metric of place/grid cell-ness, and then quantifying the extent to which a given model unit resembles a place cell or grid cell, and then (2) quantifying how common these representations are across all model units, rather than simply showing one or two example units (and ideally comparing these numbers to estimates of the prevalence of such units in the relevant [biological] neural populations). And then (3) showing the top K units that score on this rather than one or two (ideally k=50 or 100 and then is just a big appendix figure). If the authors also wish to discuss band cells then they should do this for that type as well. Lastly, an obvious gap is that the authors don't mention boundary cells -- do they find any units that mimic those?

- I found myself wanting a full end-to-end diagram of both models, but particularly the TEM, as I am (and I imagine the ICLR audience more broadly is) less familiar with it. Obviously, space is tight, but including a full schematic of each model would be  particularly helpful. If need be, this could be in the appendix.

- Further clarity on what the task exactly is. The authors briefly discuss it briefly in a few paragraphs, but the manuscript would benefit from more details on the exact nature of the task -- e.g., example diagrams of the environments, and perhaps even examples of performance over training. Since space is tight, this could just be in the appendix.


Minor issues:
- What do the authors mean that TEM-t both has better sample efficiency and has reduced training time? Naively, one would imagine these things are largely equivalent statements, but I assume they mean something more here.
- Typos: 
  - Pg 5: "via via Hebbian learning" --> "via Hebbian learning"
  - Pg 6: "TEM can tackle much larger problems" --> "TEM-t can tackle..."
  - Pg 7: "relation ship" --> "relationship"

summary_of_the_review: By offering a novel perspective on transformers and models of biological memory systems, this work may lead to fruitful future work both in machine learning and neuroscience.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper demonstrates emergence of grid and place cells in transformer architectures with positional encodings. Transformers are related to TEM models and Hopfield networks in computational neuroscience. A neurobiological model of TEM-transformers is presented, which contains feature neurons and memory neurons. It is shown that memory neurons resemble place cells in the hippocampus. Remapping phenomenon is also discussed. 
main_review: This paper nicely connects several ideas discussed recently in the AI and comp. neuroscience communities, such as transformers, Hopfield Networks, memory systems, place cells in the hippocampus, grid cells, etc. I want to give a couple of pointers to the literature that might make these connections more transparent and mathematically more precise. 

1. The left equation in (Eq 11) seems to be closer related to the class of models that have power activation function (see for example equation 10 in Dense Associative Memory for Pattern Recognition by Krotov & Hopfield 2016 https://arxiv.org/abs/1606.01164, and assume that f is linear, or more generally a power function) than to models with f=softmax. In the classification scheme of Krotov & Hopfield 2020, the former is a class of models that they call models A, while the latter is the class of models that they call models B. The same applies to (Eq 8). 

2. The network motif with two kinds of feature neurons, similar to the one shown in Fig.5 B is commonly used in models of associative memory in the hetero-associative setting. When there are two kinds of features that need to be associated with each other in one memory. For instance, this idea was used in Krotov & Hopfield 2016 for associating pixel intensities of images with the labels of those images, see for example Figure 1A. The pixel intensities are presented to the network and hold clamped, while the label part of the feature neurons are updated to compute the desired output. 

3. The sentence “It has been shown that the use of a softmax corresponds to a variant of the Hopfield energy (Demircigil et al., 2017) where the number of memories is untethered from dimensionality of the attractor and therefore are potentially unbounded.” on page 8 is not entirely correct. First, the model by Demircigil and coauthors has a bounded capacity, which is at most $2^{(N/2)}$, see their theorem 1.3 in https://arxiv.org/abs/1702.01929. Second, the activation function in that model is f(x) = exp(x), rather than the softmax. I understand that the confusion comes from the derivation given in Ramsauer et al., who start with the model with f(x)=exp(x), then take a logarithm of the energy function for that model, and add an extra term to enforce the lower bound on the energy. The problem with this logic is that the addition of the quadratic term changes the mathematical properties of the model. For this reason, the model with the softmax activation, and the model with f=exp(x) are two mathematically different models.   

I also wish the paper had some kind of quantitative metrics of how closely the learned representations resemble place cells. Right now the similarity is only qualitative, e.g. Figure 5D. If the evaluation was somewhat more quantitative the authors might be able to ask interesting questions. For example how does the choice of the activation function (softmax, power, exponent) change the similarity of the learned representations to biological place cells? Maybe some activation functions model the biological system better than others?   

summary_of_the_review: In general, I like the scope of the paper, and I am inclined to vote accept provided that the authors can address the comments above.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper postulates a task that requires the ML model to capture the spatial nature of the task in order to perform well. The chosen ML method is a Transformer with positional encoder, where the encoding is not fixed, but is learnable, making it possible for the transformer to "choose" what kind of representation would it "prefer" for the task at hand. The Transformer is modified to act comparably to a neuroscientific model of hippocampal function (TEM). After the learning process the authors discover, that, when visualized, the representation that has emerged resembles the grid-cell pattern similar to the one that is being encoded by hippocampal place and grid cells. From there the works suggests the existence of a special relationship between the Transformer architecture and the mechanism by which hippocampus encodes spatial information.
main_review: In general I am very curious and supportive of the effort to find similarities between the methods and representations that biological and artificial learning systems discover during their learning process. However we must be very strict with ourselves when we declare that such a similarity has been found:
  * (a) Is it sufficiently similar to be non-trivial?
  * (b) Can the emergence of this similarity be explained by other factors, such as experimental setup?
  * (c) Does the similarity occur only when we are using one particular ML algorithm or would it also emerge when attempted using other learning methods / algorithms / architectures that are appropriately configured?

The emergent patterns resembling grid cell structure that are shown to emerge in this paper closely resemble the patterns that emerged in an RL agent in this work: https://deepmind.com/blog/article/grid-cells. In that work Transformers were not used, but the pattern has emerged nevertheless. This begs the question whether Transformers are indeed *the* architecture that encodes spatial representation closest to hippocampus, or whether any ML model, properly incentivized, will form such representations?

Authors mention in the abstract that "this result is no surprise since it is closely related to current hippocampal models from neuroscience" about the claim "transformers, when equipped with recurrent position encodings, replicate the precisely tuned spatial representations of the hippocampal formation". Could we have written the same paper by just replacing Transformers with something else that could mimic neuroscientific model of hippocampal grid cells? If yes, then Transformers have no special significance here and should not enter the discussion.

I am looking forward to the discussion and am very to open to changing my rating upon further exploration of these matters.


Major concerns / criticisms / discussion
----------------------------------------

* Seeing how Transformers are very flexible models with high capacity, wouldn't it be the case that (with a little twist) one could fit them to become good representations of almost any phenomena? The theory that is put forward in this paper is that Transformers model representations that are close to those used by hippocampal formation. What kind of result would invalidate such a theory? Let's say that instead of Transformers we would use a black box method "X" that would tune its parameters in such a way, that they correlate well with hippocampal responses. Would that mean that "X" is good model of representations that are being used by hippocampus? Note that the more flexible "X" is and the more time and computational resources we have, the closer we can bring "X"'s activations to correlate with hippocampal responses. What would be a hypothetical empirical indication from your experiments that would lead you to conclude that "nah, Transformers are nothing like hippocampus"?

* Our brain did not learn in an environment with a 4-connected Eucledian structure. Wouldn't having this structure most assuredly lead to a representation that reflects the pattern of the structure? Let's say you would conduct the same experiment in an environment with a some kind of 10-connected star-shaped Eucledian structure -- wouldn't the patterns of representations formed by the Transformer would try to capture that star-shape?


Questions
---------

* page 3: "It now becomes interesting to see what representations are learned" -- taking into account the properties of the task and its spatial structure, what would we expect the representation that captures that structure to look like? Does it have to be a Transformer to capture that representation or could be captured, for example, by an autoencoder trained to find an efficient representation of the samples constituting the task?

* page 4: "Hopfield network (Hopfield, 1982) which have recently been shown to be closely related to transformers" -- Could you please elaborate what is meant by being related and how this relationship pertains to the claim that "a transformer with recurrent positional encodings is closely related to current neuroscience models of the hippocampus and surrounding cortex"?



Minor remarks
-------------
page 2: typo "Crucially, such rules generalises"
summary_of_the_review: In my current evaluation (prior to discussion with the authors) it seems to me that the claim that Transformers have a special kind of relation to hippocampal grid cells is not supported by the evidence. The evidence does show that the condition is "sufficient" (Transformers do form grid-like representations), but it does not show that it is also "necessary" (meaning it does not have to be a Transformer to form a grid-like representation). If my understanding of the evidence is correct, then it does not support the claim of the paper and I would recommend rejecting this result.In my current evaluation (prior to discussion with the authors) it seems to me that the claim that Transformers have a special kind of relation to hippocampal grid cells is not supported by the evidence. The evidence does show that the condition is "sufficient" (Transformers do form grid-like representations), but it does not show that it is also "necessary" (meaning it does not have to be a Transformer to form a grid-like representation). If my understanding of the evidence is correct, then it does not support the claim of the paper and I would recommend rejecting this result.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The authors show that the Transformer architecture, with a specific choice of parameters, is analogous to the Tolman-Eichenbaum Machine (TEM), a neuroscientific model successful at explaining phenomena related to grid cells in the entorhinal cortex, place cells in the hippocampus, and their coordinated remapping in new environments. The authors show that, compared to TEM, its Transformer counterpart is more sample-efficient and should better generalize to multisensory inputs, thus offering a computationally efficient basis for models of the hippocampus and entorhinal cortex.
main_review: *Strengths*

The authors show the full analogy between the Transformer architecture and the TEMs – unlike the existing literature which has only identified some similarities between the two (e.g. the similarity between the attractor update rule in TEM and the matrix product in self-attention).

The Transformer architecture proposed in this paper inherits the biological plausibility of TEMs and, at the same time, is considerably more sample-efficient, allowing for building computationally efficient models of the hippocampus and entorhinal cortex.

The proposed architecture is expected to scale well with an increase in the number of cortical inputs to the model thus potentially enabling the studies of multisensory integration.

The paper has a nice introduction to the Transformer architecture.

*Weaknesses*

Whereas the authors, due to their results, expect the Transformer architecture to inform the future models of the hippocampus – and the other way around – such work has not been done in the scope of the current paper.

I think it would be nice to include a slightly more detailed description of the TEM model in the paper. I like the equation-based description, offered here, better than the text-based one in the original TEM paper; however, this approach would benefit from defining the variables and, perhaps, briefly commenting on their role before using them (e.g. the q-variable)

summary_of_the_review: The paper proves a parallel between the Transformers and the Tolman-Eichenbaum Machines, thus offering a computationally efficient and scalable implementation for this biologically-grounded model of the hippocampus. Even though no new biological insights were obtained in the scope of this paper using that knowledge, the parallel between the two frameworks deems useful for models of the hippocampus in follow-up works and highly relevant to the community. I, therefore, recommend accepting the paper.
review_title: null

Meta-review: This paper received 4 unanimous accept (including 1 marginal accept). This well-written and clear paper clarifies the relationship between transformers and a recent exciting model of the medial temporal lobe in neuroscience. There was some clarifications requested by the reviewers that were addressed during the revision. This paper will make a great computational neuroscience contribution to this year ICLR!
Decision: Accept (Poster)