Global Distortions from Local Rewards: Neural Coding Strategies in Path-Integrating Neural Systems

Thank you for taking the time to read and review our paper. Your feedback was invaluable, as it highlighted that we had not properly emphasized several crucial aspects of our work. In particular, we had not sufficiently explained the generality of our theoretical framework and the fact that our experiments were made possible only when freezing the place cell read-out weights. We address your comments and questions in detail below.

Theoretical contributions

We thank you for this important remark. Motivated by your feedback, we provide additional elements to our theory, as discussed in the global answer to all reviewers.

We now explicitly model the emergence of the four deformations: diffused units, band units, ring units, and attracted units. Interestingly, this only required a simple modification of the theory. We look forward to hearing your thoughts about it.

Experimental contributions

You noted that we did not sufficiently study how particular changes in the loss function correspond to particular changes in the rate map distortions. Indeed, the earlier experiments were simple in the sense that differences in reward location, magnitude, or spread were not well explored. We now performed these experiments, please see our global response for details.

As you mentioned, the most relevant modification in our work is the differential weighing of the loss function to incorporate the reward location. Here, we wish to emphasize another of our unique contributions that was missing from your summary. When we first performed the saliency training, we observed that the topology of the neural manifold was not preserved (Fig. 7B). This was a surprising result, one we had not anticipated. Only after we froze the pre-trained output weights to the place cells, we were able to observe both the global distortions and the preservation of the topology (Fig. 7C).

So, when referring to saliency training, we do not mean a simple modification of the loss function, but also this non-trivial aspect of weight freezing. We believe setting up this framework to study the effects of local rewards is an important technical contribution of our work, which the future work can build on.

Signal periodicity

In Fig 3, showing results on synthetic neural responses, the diffused units have the same spatial periodicity as the original units – as the diffused units were created by applying a Gaussian convolution to the original units. That is, the locations of the signal peaks are the same.

Wording on diffused and band units

This is a very good point, thank you for bringing it up. As you correctly point out, our theory does not explain why some distortions appear and some others don’t. Instead, we provide a mathematical formulation of their distortions, and explain what is the effect of this distortion on the neural manifold. We apologize for the confusing wording, which we propose to correct as follows: "we observed the emergence of diffused, band and ring units, which fall into the category of distortion via isotropic and anisotropic deformations from our theoretical framework".

Attracted units

You mentioned that the lack of attracted units in the trained piRNN was surprising. We also expected to see some local distortions of the grid cell firing patterns, given the standard interpretation of neuroscience experiments like Boccara et al. However, to our surprise, we did not observe this kind of local “magnification” or deformations in any of our experiments. Instead, we observed more global changes in the rate maps – which inspired the title of the work, “Global Distortions from Local Rewards”.

We believe this is a strength of our work, not a weakness. Our results challenge the conventional interpretation that reward modulation results in local magnification of the grid cell lattice, which may be a simplistic explanation. Indeed, in hindsight, the original findings might have been inconsistent with the conventional wisdom that grid cells provide a global code for space. Thus, our results provide the first evidence that reward modulation may result in global deformations of spatial responses in MEC, and the local distortions may arise due to interactions with other experimental variables not controlled in Boccara et al., 2019, or a reasonable but incorrect interpretation of the data.

It may also be true that future changes to the saliency training may recapitulate the experimental findings. Our work is just the first step (which sets up the framework) to study this problem theoretically and in detail.

Rephrasing introduction.

We agree with the suggestion of rephrasing the introduction so that it no longer suggests that “reward will be incorporated into the loss function”. Indeed, as you mention, the actual method is simpler. We plan to rephrase as follows: “we modify the loss function to differentially weigh position estimation error based on a saliency map of space.” Thank you for the suggestion!

Quantifying distortion

This is an excellent suggestion! To address this, we performed statistical tests with structural similarity index and Wasserstein distances. Specifically, we computed these distances between images before and after distortion for each class of distortions, and now defined the distortion units as a result of this statistical analysis.

Overall, we believe your feedback was very helpful for increasing our impact. To address your concerns, we have performed additional analyses (leading to new results, see our general response!) and expanded on our theory. If you believe we have adequately addressed your concerns, would you consider increasing your score to support the acceptance of our work?

Thank you very much for your thorough review. We appreciate your attention and address your comments and questions below.

Main concern: link between theory, and firing patterns observed in piRNN experiments versus real experiments

Thank you for raising the point that distortions of firing patterns in real rodent experiments differ from those observed in piRNN experiments. We think that this is an important point, and is actually one of our most interesting findings, which is answered below in response to your question #2.

On piRNN experiments

We agree that more thorough experiments would be beneficial to this work. This concern was shared across several reviewers. Consequently, we have run additional experiments. These experiments include a wider range of saliency functions as per your suggestion. Please see our global answer where we present and discuss their results.

Answer to question 1. on continual learning

You asked whether this scheme would be capable of training a salience-trained piRNN on a new changed reward function. This is a very interesting question! In the same way that we pretrained a network with uniform saliency (no reward) and then fine tuned the network with non-uniform saliency s(x), a good starting point would be to introduce additional phases of fine-tuning with different saliency maps, and analyze the resulting deformations after each phase. We expect that the neural responses will adapt by deforming in the ways we have described (developing a mixture of band units, diffused units, and ring units), while phases of training without rewards would see a convergence back to the typical hexagonal grid cells.

Answer to question 2. on global distortions

You noted that previous works reporting experimental observations of firing rate distortions interpret them as local changes as opposed to global distortions. This is what we also expected to see a priori, given the claims from neuroscience experiments in Boccara et al. However, to our initial surprise, we did not observe this kind of local “magnification” or deformations in any of our experiments. Instead, we observed more global changes in the rate maps – which inspired the title of the work, “Global Distortions from Local Rewards”. Thus we challenge the conventional interpretation that reward modulation results in local magnification of the grid cell lattice as too simplistic an explanation. Indeed, because grid cells provide a global code for space, we propose that reward modulation should result in global deformations of spatial responses in MEC, as observed in our piRNN experiments.

Answer to question 3. on link between loss function and diffeomorphism theory

You raise a thought-provoking question about how a grid cell deformation, modeled as a diffeomorphism of the 2D environment, depends on the chosen loss function. Dorrell & Latham et al., 2023 showed that hexagonal grid cell representations result from minimizing a particular loss function subject to biological and representational constraints. Their loss function encourages the separation of neural representations of different positions in space, and includes a term that weighs the importance of separating any pair of positions x and x’. This term is directly analogous to the saliency map we introduce in the loss of our piRNN. Because Dorrell & Latham et al. do not study grid cell deformations, we believe that the framework we introduce in our paper for studying grid cell deformations opens a new route to make progress on this question, complementing existing work in the literature. We note that this question is of utmost importance to neuroscience and AI: given some task and objective function, can we predict and explain the representations learned by neural systems?

Conclusion

We thank you for your thorough review and specifically for prompting us to discuss the disconnect between the results from the piRNN experiments and the real experiments, as well as its link with our theoretical framework. We hope that we have addressed your concerns and we look forward to discussing these topics further.

We thank you for the thoughtful review, and address your comment and questions below. We look forward to hearing from you in case you have any additional comments!

Answer to weakness on Originality/Novelty: We agree that the originality and novelty of our approach could have been made clearer in the original paper, which did not emphasize these points enough. In the final version, we propose to add a subsection 2.3 in Section 2 on Background & Related Works to explain the following:

In this work, we use the supervised piRNN approach as a framework to study the deformations in grid cell firing patterns reported in the experimental neuroscience literature. We do so by modifying the piRNN’s supervised loss in an interpretable way to account for varying saliency of space. Particularly, our saliency training allows for systematic study of two highly entangled behavioral variables: i) where the rewards are located and ii) which locations the agents have the tendency to visit the most. To our knowledge, this approach is completely novel in the literature.

In fact, in our new Fig. N4, we show that the grid cells contributing to the coding of the origin (the most visited location), but not the saliency map (the reward location), are robust to global distortions during the saliency training. In other words, as a secondary objective is introduced, grid cells with projections to the place fields at the edges are more likely to become distorted. To our knowledge, we are the first to make this prediction, which can be tested experimentally, for instance, by introducing both reward and landmark locations. Furthermore, we introduce a mathematical, geometric framework to explain and quantify the firing patterns’ deformations observed in our experiments; this is also novel and original. Finally, we show “that global distortion (geometry) can be added while maintaining the toroidal topology for path integration” (citing reviewer yMbD), which is true once the read-out weights to place fields are frozen, but not otherwise!

One related work is Nayebi, et. al. “Explaining heterogeneity in medial entorhinal cortex with task-driven neural networks” (NeurIPS 2021). The authors were also motivated by results from experimental neuroscience on reward-modulated neural responses by grid cells. While not the central question of their work, Nayebi, et al. propose a modification of training piRNNs to account for reward modulation: they modify the way training paths are generated in the presence of reward. Specifically, the authors modified the training paths ge to be either (1) direct movement toward a reward location, called “pure exploitation”, (2) random walk, as is standard in the piRNN literature, called “pure exploration”, or (3) an intermediate policy consisting of a mix of movement toward reward and random walk exploration. Their results are consistent with the existence of nontrivial reward-modulated response changes in piRNNs. However, they did not seek to characterize the specific deformations to the neural responses, as we do in our work.

Answer to question on Equations (1) and (3): Thank you for pointing this out. On the one hand, there is $\phi(x_t)$, which is the vector of place cell activity corresponding to position $x_t$. This is an “encoding” of 2D position into a place cell code.

On the other hand we have $\hat{\phi}_t = Qr_t$, which is a linear readout of the recurrent grid cell activity $r_t$. We apologize for the typo in Equation (1), which thus should have been: $\hat\phi = Qr$ and $\hat x = \text{argmax} \hat \phi.$

Successful path integration here means that $\hat \phi_t \simeq \phi(x_t)$, captured in the $L_{error}$ term. To further highlight this point, we can rewrite Equation (3) as

$L_{error} = \sum_{t=1}^T ||\phi(x + …) - \hat \phi (x + …)||^2 = \sum_{t=1}^T ||\phi(x + …) -Qr (x + …)||^2.$

We thank you for catching this.

Conclusion

We once again thank you for your time and attention. We hope that we were able to resolve your concern. If not, we look forward to discussing these topics further! If you believe we were able to adequately address your concerns, we hope that you will consider supporting our work with a strong acceptance!

I thank the authors for addressing my concerns, and for clearing up my confusion regarding Equation (3). I have accordingly increased my score to a Strong Accept.

We thank you very much for the thorough review and for your time. We address your comments and questions below.

On real neuroscience experiments:

You make the very important point that linking the theory with “real” data from rodent brains is important to increase the confidence of this work. Thank you for bringing it up. We think that this point deserves additional attention and we discuss it below.

The existing experimental neuroscience data on such reward-modulated deformations of grid cells firing patterns is limited in number and in the strength of the conclusions. We are aware of the results from [Boccara et al. 2019, Butler et al. 2019, Krupic et al. 2018, and Wang et al. 2021]. Boccara et al. 2019 in particular motivated our approach and we explained and discussed its results with our theoretical framework. To summarize: Boccara et al. observed an attraction of grid cells firing fields towards the reward – a phenomenon called “attracted units” in our paper. Furthermore, band units are well described in the literature – see Krupic et al. 2012. Are you aware of any other recent public datasets that may be relevant? If so, we would be happy to discuss them as well.

Due to the current limited experimental data, we chose to study in silico ANN-based experiments inspired by the existing, real neuroscience experiments. Therefore, the scope of our work is: we investigate how the representations learned by artificial path-integrating RNNs (piRNNs) change to account for “saliency” in the environment, in a way that is (1) interpretable, and (2) inspired by neuroscience experiments. We observe that after training with saliency conditions, our piRNN learns a more diverse set of neural representations (ring neurons, diffused neurons, band neurons) compared to observed in Xu et al. and Gao et al. earlier – while still maintaining toroidal topology. We then build a theoretical framework to explain how different deformations in the individual rate maps affect the geometry of the neural representation manifold, which is novel in the literature.

Unfortunately, being a theory lab, we cannot perform additional experiments on real rodents ourselves. Yet, we believe that our geometric study of piRNNs already provides interesting results for the study of RNNs, which is an important field of research in AI/ML. Additionally, we hope that our work will inspire experimental neuroscientists to perform corresponding experiments on rodents. In that perspective, it might even be preferable that theory and experiments come from two separate research teams, as its final validation on real data would even be stronger.

On additional piRNN experiments: Thank you for proposing this more comprehensive set of experiments. We have run them; please see the global response on experiments. We believe that these additional results strengthen our paper and we appreciate your suggestion.

Answer to question 1.): why such global distortions?

You asked why the distortions of grid cells firing patterns appear uniformly, while the reward (saliency) is local and localized at the center of the space. This is a great question! We also expected to see some local distortions of the grid cell firing patterns, given the claims from neuroscience experiments in Boccara et al. However, to our initial surprise, we did not observe this kind of local “magnification” or deformations in any of our experiments. Instead, we observed more global changes in the rate maps – which inspired the title of the work, “Global Distortions from Local Rewards”. Thus we challenge the conventional interpretation that reward modulation results in local magnification of the grid cell lattice as too simplistic an explanation. Indeed, because grid cells provide a global code for space, we propose that reward modulation should result in global deformations of spatial responses in MEC, as observed in our piRNN experiments.

Answer to question 2. on piRNN experiments that vary the saliency map

Thank you for proposing additional experiments with variations to the saliency map s(x)! Please see our results in the experiments section of the global response.

Answer to question 3. on the link between distortion and number of grid cells/modules

While we have not run experiments varying the number of grid cells, this is an interesting modification we will investigate in the near future. We would like to emphasize that the number of modules, as we have defined them, is an emergent property of the trained network and thus we do not have direct control over it. How the number of modules that emerge depend on various experimental parameters is a very interesting question that needs to be further explored, but is out of scope here.

Answer to question 4. on why different distortion patterns appear?

Thank you for the great thought-provoking question. We had not considered why different distortion patterns indeed jointly appear, i.e., why we observe both diffused and ring units after introduction of the reward. Your question also prompts the follow-up question: what role can diffused units, band units, and ring units play? To bring some insights to this question, in our new Fig. N4, we show that the grid cells contributing to the coding of the origin (the most visited location), and not the center of the saliency map (the reward location), are robust to global distortions during the saliency training. In other words, as a secondary objective is introduced, grid cells with projections to the place fields at the edges are more likely to become distorted.

Conclusion

We hope that we were able to address your comments and resolve your questions. These discussions have been beneficial for us, and we will gladly incorporate them into the final version of the paper. If you have any remaining questions, we'd be happy to discuss further!