Flexible mapping of abstract domains by grid cells via self-supervised extraction and projection of generalized velocity signals

We would like to thank the reviewer for their detailed feedback and for raising excellent questions. We will try to address all questions one-by-one.

Relative to CSCG, TEM and models like them, I think this work is simply answering a different part of the cognitive mapping question

Thank you for highlighting these points and for your careful review! As we note in the text, the crucial difference between our approach and those of others such as TEM and CSCG, is the following: previous approaches perform simultaneous learning of representations of external states and transitions between states; in contrast, we propose that only transitions need to be learned, which are then processed by a reusable grid attractor to enable integration and mapping. Thus, velocity-extraction can be done more simply relative to models that would involve both representational embedding and transition learning. However, we completely agree with the reviewer’s point that our model can augment models such as TEM and CSCG, by taking advantage of our model’s state-independent extracted velocities (or equivalently, actions). We will note this in our revised text and we will correspondingly edit the framing of our comparison with background work, which we did not intend to be confrontational.

Claim that the velocity signals are learned independently of the stimuli also seems misleading (figure 1g)

We claim that the extracted velocity signals are independent of the stimuli within each environment. We do not make any claims that a network trained on one environment should be out-of-the-box generalizable to other abstract domains. Fig. 1g shows the extraction of generalized velocity signals within an example environment, Stretchy Blob. What we are saying is that the same velocity is extracted at various points within a space, independent of the specific location in that space. We will clarify this additionally in the text.

Yet, in lines 78-82 and 257-265 the authors make exactly this series of claims and frame it as a novel contribution. This is certainly a prediction, but it is far from unique to this work!

Thank you for the opportunity to clarify!

Our prediction here is not simply that grid cells can encode physical and abstract space, or that grid cell-like tuning exists across different spaces. Our prediction is, rather, that cell-cell correlation is preserved across different spaces. The key nuance in our framework is the use of a single prestructured grid cell attractor network across different spatial and non-spatial environments through being able to extract generalized velocities in each task domain. Because we hypothesize that the brain generates velocity signals and pipes them into a rigid, prestructured set of attractor networks, we can predict that the cell-cell correlations within each grid network should be preserved across different modalities (i.e. if two grid cells are co-active for a spatial task, they should remain coactive for a non-spatial task). Since other models (such as TEM) do not use prestructured attractor networks, and learn the appropriate representations from scratch in different domains, it is unclear whether the cell-cell correlations will be preserved across domains/modalities. Although our conclusions may seem straightforward, other computational models that build cognitive maps (such as TEM, CSCG, SR) do not use fixed, prestructured attractors. We are unaware of these models being able to make such a prediction.

The requirement for the latent code to be a linear function of the true velocities

Since grid cells are doing path integration through linear addition of velocity vectors, constraining our latent space to also be a linear function of the true velocity signals was crucial. Indeed, as the reviewer correctly noted, linearity is built into our choice of losses. In fact, the loop-closure loss imposes a linearity constraint (as can also be seen through our ablation studies; we will add a brief explanation of why loop closure ensures linearity in the appendix), with the additional shortcut and isotropy losses refining the obtained linear representations of velocity estimates. Indeed, what the reviewer noted is accurate: next-state prediction leads to a nonlinear representations of velocities, loop closure makes the encoding linear, and additional terms then refine the representations. We thank the reviewer for this succinct description and will include it in the text.

I couldn't find the loss [calculated in the table] described in detail anywhere

We apologize for not including these details. We will include additional details in the text describing calculation of the error used in Table 1. We provide a brief description of the error metric here. As discussed above, we aim for the estimated velocities $\hat{\mathcal{V}}$ to be a linear function of the true velocities $\mathcal{V}$. Correspondingly, we first estimate the best-fit linear transformation $T$ from $\hat{\mathcal{V}}$ to $\mathcal{V}$ via a pseudoinverse, $T = \mathcal{V} \hat{\mathcal{V}}^\dagger$. Then, we compute our error metric as a normalized mean-squared error between the $N$ transformed points and the true distribution: $e = \frac{|| T\hat{\mathcal{V}} - \mathcal{V} ||_2}{N \text{var}( \mathcal{V} ) }$. The presence of outliers, particularly in some of the poorly performing baseline methods, can lead to particularly poor best-fits $T$. To alleviate this, we find the transformation $T$ after removing a small number of outliers in $\hat{\mathcal{V}}$ via the DBSCAN clustering algorithm (any outlier rejection tool will suffice); however, we report the error $e$ evaluated on the entire dataset including outliers.

Why should I expect other dimensionality reduction techniques to learn the velocity?

PCA extracts a different, but also meaningful encoding. It is interesting that this is not the same subspace, and certainly it shows the proposed method is the best at this, but the phrasing made it sound like a key contribution is beating these other methods

This is an excellent question. By construction, the datasets that we have generated have a specified intrinsic dimensionality (1, 2 or 3, depending on the environment). For example, in the Stretchy Bird environment, there exists a minimal representation of the dataset that is captured by a two-dimensional coordinate, neck and leg lengths. Thus, within the high-dimensional space in which the raw inputs live, the dataset is explicitly constructed to lie on a (nonlinear) two-dimensional manifold. If a dimensionality reduction method were able to flatten this manifold and represent it in two dimensions, differences between points would produce a veridical representation of velocities. While PCA does provide an encoding of the data (that describes the projection of the data onto the principal components), it does not provide a meaningful representation, since it is unable to find even a low-dimensional representation of the data. We demonstrate our viewpoint further through the new Fig. R3 in the attached PDF. Here, we examine a dataset generated through a single trajectory within our 2D Moving Blobs environment. Through integration of our velocity estimates, we find that the high-dimensional states collapse onto a two-dimensional plane, and are thus completely captured by two coordinates. In contrast, PCA appears to require 24 dimensions to capture 95% of the variance within the data, and appears to occupy a volume when plotted in three dimensions. Since PCA is unable to find a low-dimensional representation of a dataset that is intrinsically two-dimensional, we believe that the projection learned by PCA is not strongly meaningful for this dataset. These new results continue to point towards dimensionality reduction as a key contribution and strength of our model. As a result, we felt that it was reasonable to use the other dimensionality reduction techniques as necessary baselines, which the other reviewers have appreciated. In summary, we find that integrating estimates of low-dimensional transitions between high-dimensional states can be an effective tool for dimensionality reduction.

Stated equivalent between velocities summing to 0 and inputs being the same assumes no aliasing

Thank you for pointing this out. We had implicitly assumed that within the abstract cognitive spaces we investigate, states are unique at each point in space. We will make this assumption explicit in the text.

You cannot then use this for spheres

Thank you for raising this subtlety in our model. We do in fact assume that the velocity vectors in the space of our latents commute. As a result, the estimated velocity vectors cannot directly represent tangent vectors in a non-Euclidean space, such as a sphere. We will add this note to the paper. However, this does not entirely preclude the representation of these non-Euclidean spaces through our method, since we can consider these spaces as embedded in a higher dimensional Euclidean space wherein velocity vectors will commute again. For example, in representing a point in an abstract domain with spherical geometry, our method will fail if we attempt to estimate two-dimensional velocity vectors representing transitions. However, if instead we estimate three-dimensional vectors, we expect our method to continue to work since the surface of a sphere can be embedded in three-dimensional space. For data arising from a general manifold, we will require a few extra dimensions (cf. Whitney embedding and Nash embedding theorems) but will always be able to consider a Euclidean space of sufficient dimensions that will embed the data generating manifolds. We thank the reviewer for leading us to consider this subtlety and will include this cost of requiring extra dimensions for non-flat spaces as a limitation in the discussion section of our paper.

Further more details of the exact form of the losses would be nice

Thank you for the comment. We have currently listed the description of each loss in L173-195 along with equations in Fig. 3 and are happy to explicitly write out our loss functions and other algorithmic details in the appendix, that we briefly mention here:

Given two states $i\_t$ and $i\_{t+1}$, the next-state prediction loss minimizes the distance between the predicted state and the true state: $\text{min} ~ || i\_{t+2} - \hat{i}\_{t+2}||\_2$.

The loop closure loss ensures that all the predicted velocities in a given trajectory sum to zero given that the trajectory is a loop: $\text{min} \sum\_{0 \leq t \leq T-1} \hat{v}\_{t \rightarrow t+1} = \text{min} ~ || \oint \hat{v} dt ||$.

The shortcut loss ensures that the decoder $g$ can generalize given a state and a velocity. For instance, $\text{min} ~~ || g(i\_{t+2}, \hat{v}\_{2 \rightarrow 3} + \hat{v}\_{3 \rightarrow 4}) - \hat{i}\_{t+4} ||$.

Finally, the isotropy loss induces an isotropy in the inferred velocity space: $\text{min} ~ \text{var} [ ||\hat{v}\_{t \rightarrow t+1} || ~ \mid ~ d(i\_t, i\_{t+1}) < \theta ]$, where $d$ is a similarity function in the input image space and $\theta$ is some small threshold.

All losses have a weight / prefactor, listed in Table 2 of our paper. As mentioned in L200, regardless of the training environment, the relative weighting of the two critical loss terms remains consistent (the loop-closure loss weighted ten times higher than the next state prediction loss). We also plan to discuss more about how each synthetic task was generated, as also recommended by $**sK54**$.

We hope that you agree that our paper has improved given your detailed feedback and suggestions! Given that we have addressed the primary concerns raised in the review, we kindly ask you to adjust your score while taking our rebuttal into account.

We would like to thank the reviewer for their detailed feedback and for raising excellent questions. We will try to address all questions one-by-one.

State-based SR doesn't assume action inputs

Thank you for this note. We will clarify this in our text and distinguish between state-based SRs and state-action SRs. We additionally note that the original SR paper (Dayan 1993) and recent extensions (Stachenfeld 2017) use a discrete action space whether SR is state-based (i.e. expectation over actions) or state-action based.

Grammatical changes in text

Thank you for suggesting these revisions! We will change them in text.

Somewhere in the supplement, the authors should completely spell out their loss function and other algorithmic details

Thank you for the comment. We have currently listed the description of each loss in L173-195 along with equations in Fig. 3 but are happy to explicitly write out our loss functions and other algorithmic details in the appendix, that we briefly mention here:

Given two states $i\_t$ and $i\_{t+1}$, the next-state prediction loss minimizes the distance between the predicted state and the true state: $\text{min} || i\_{t+2} - \hat{i}\_{t+2}||\_2$.

The loop closure loss ensures that all the predicted velocities in a given trajectory sum to zero given that the trajectory is a loop: $\text{min} \sum\_{0 \leq t \leq T-1} \hat{v}\_{t \rightarrow t+1} = \text{min} || \oint \hat{v} dt ||$.

The shortcut loss ensures that the decoder $g$ can generalize given a state and a velocity. For instance, $\text{min} || g(i\_{t+2}, \hat{v}\_{2 \rightarrow 3} + \hat{v}\_{3 \rightarrow 4}) - \hat{i}\_{t+4} ||$.

Finally, the isotropy loss induces an isotropy in the inferred velocity space: $\text{min} ~ \text{var} \left[ ||\hat{v}\_{t \rightarrow t+1} || \mid d(i\_t, i\_{t+1}) < \theta \right]$, where $d$ is a similarity function in the input image space and $\theta$ is some small threshold.

All losses have a weight / prefactor, listed in Table 2 of our paper. As mentioned in L200, regardless of the training environment, the relative weighting of the two critical loss terms remains consistent (the loop-closure loss weighted ten times higher than the next state prediction loss). We also plan to discuss more about how each synthetic task was generated, as also recommended by $**sK54**$.

Boundary/obstacle effects

This is an excellent question. We clarify that while the decoder is state dependent by construction (outputs a predicted state given a state and velocity), the encoder is state-independent since it infers a generalized notion of velocity between any two high-dimensional states within our abstract domains. We investigate how the decoder performs at boundaries (Fig. R4 in the attached PDF) within the 2D Stretchy Bird environment. We find that the decoder implicitly understands the boundaries of the training data’s underlying manifold. That is, once a velocity produces a bird state that is unseen in the training distribution, i.e. a bird that cannot further shrink its legs or extend its neck, further transitions in the same direction do not produce any changes in the predicted state. Thus, the decoder understands the boundaries observed in the training data.

Model assumes that state variables live in R^D

Thank you for raising this point! We do indeed assume that the input state variables occupy a region of $\mathbb{R}^D$ (where $D$ is the number of pixels in the input, for example); we also further assume that the estimated velocities live in $\mathbb{R}^d$ for some smaller $d$. This in and of itself is not particularly constraining: while a single module of grid states live on a torus, they can integrate velocities that are obtained generally from $\mathbb{R}^2$ without bounds. While grid states wrap around, the velocities do not necessarily need to wrap around boundaries. Further, if we include additional grid modules, grid cells can integrate velocities in the higher dimensional Euclidean space $\mathbb{R}^d$, as has been shown by recent theoretical work (Klukas et al. 2020, PLOS Computational Biology 16(4): e1007796). We are happy to add a discussion of this point in text.

Concerns on loops

Thank you for raising these excellent points! While our models were trained with data consisting of loops, this training paradigm is simply for convenience (L179-181) and is not a necessity of the model. If arbitrary random walks were selected instead for training samples, self-intersections would automatically lead to loops within sub-sequences of the walks (as detected by $k<m$ with $i_k=i_m$) — these loops could then be used for the loop-closure loss, with all other losses applied on all loop and non-loop trajectories. As a simple proof of concept, we generated a dataset of arbitrary trajectories with 50% of the trajectories as loops and the remainder as non-loops. We applied loop closure loss on trajectories that formed loops and all other losses on the entire dataset of trajectories. Fig. R1 in the attached PDF demonstrates this paradigm for the 2D Stretchy Bird environment, where we continue to get consistent velocity representations with low transformation error to the true velocity distribution.

In greater generality, we also note that the loop closure loss can also work on trajectories that are “almost loops” – i.e., a loss whose prefactor term scales with how close a trajectory is to forming a closed loop. However, running this analysis and tuning these parameters is beyond the scope of the current rebuttal period, and we will examine these results for the camera ready version of our paper.

Lack of limitations

We apologize for not including sufficient discussion of limitations. We will add a discussion of the limitations to our model in text, as also mentioned to Reviewer $**sK54**$.

Compelling case for AI applications

We believe our work has potential application to machine learning beyond neuroscience, primarily through the context of dimensionality reduction.

A crucial step in analysis of large datasets of high-dimensional data is often the generation of dimensionality reduced representations of the data. Our method provides a novel SSL framework for dimensionality reduction and manifold learning, which appears to significantly outperform baselines on datasets that contain relatively lower-dimensional transitions. Video data, for example, naturally contains transitions between frames that are typically lower dimensional than the frames themselves, and thus may be particularly amenable to our techniques.

Moreover, typical nonlinear dimensionality reduction techniques, such as Isomap and UMAP are non-invertible. In contrast, similar to an autoencoder, we can use the decoder in our model to generate the high-dimensional states that correspond to points within the low-dimensional space. In this sense our model also acts like a generative model, by being able to generate inputs that correspond to different regions of the low-dimensional space.

Furthermore, our method naturally lends itself to manifold alignment, which is particularly effective when the data exhibits a small number of continuous modes of variability. Through mapping two spaces onto the same low-dimensional space of transitions, our technique can perform equivalent transformations on datasets from different modalities. Moreover, with a small number of “gluing” points, our method allows for building one-to-one correspondences between different domains.

We will be glad to include these discussion points in the camera ready version of the text.

We would like to thank the reviewer for their detailed feedback and for raising excellent questions. We will try to address all questions one-by-one.

There is little discussion about grid cells, and the description of how the continuous attractor network of grid cells works is insufficient.

Thank you for highlighting this. We briefly describe why extracting velocity signals from non-spatial, abstract domains is an important prerequisite challenge for grid cells to path integrate between states in these domains (L35-38, L95-98, and Fig. 1 caption). Here we provide an additional explanation: continuous attractor models for grid cells have gained much favor through repeated confirmation via experimental results. However, relating these continuous attractor models to mapping of abstract domains requires the extraction of a faithful representation of velocity in these domains. This is in contrast to other models for spatial mapping, such as TEM, which learn grid cell-like representations from scratch while mapping new abstract spaces. Our work shows how it can be sufficient to use a prestructured attractor network and the experimental predictions that can arise from this modeling choice. We will elaborate on this discussion and make the connection back to grid cells clearer in the text. We apologize for not including a sufficiently detailed description of the continuous attractor network that we used. We provide here a brief description of the grid cell model used and will include a more detailed description in the appendix of the revised paper. We use an approximation of a continuous attractor model for a module of grid cells: we simulate patterned activity on a lattice of neurons as the sum of three plane waves, resulting in a hexagonal pattern of activity. Input velocities are used to update the state of the activity on the lattice of neurons through updating the phases of the plane waves, leading to accurate integration of the input velocities.

Why this model is superior to the baselines

Thank you for highlighting this point as well. While we have briefly discussed why our model outperforms other dimensionality reduction techniques in L102-105, we elaborate on this further here and will include this in the paper:

Typical dimensionality reduction techniques are agnostic to the transition structure between high-dimensional input states when drawn from a trajectory. They usually rely on the statistics of distances between points across the entire ensemble of inputs. In contrast, our method explicitly looks for a structured tangent manifold around each point that captures the low-dimensional transitions between successive states along a trajectory (see Fig. R3). Through using this additional sequential information, with a prior of assuming the existence of a tangent manifold along the manifold of input points, we obtain a superior performance. We have also previously included comparison to another technique that uses sequential information (MCNet). However, there too we find that our method is significantly better through assuming the existence of a low-dimensional description of transitions in a state-independent fashion, rather than a high-dimensional maximal description of the state-dependent transitions.

Training process relies on specially constructed datasets (loops)

Thank you for raising this excellent point! While our models were trained with data consisting of loops, this training paradigm is simply for convenience (L179-181) and is not a necessity of the model. If arbitrary random walks were selected instead for training samples, self-intersections would automatically lead to loops within sub-sequences of the walks (as detected by $k<m$ with $i_k=i_m$) — these loops could then be used for the loop-closure loss, with all other losses applied on all loop and non-loop trajectories. As a simple proof of concept, we generated a dataset of arbitrary trajectories with 50% of the trajectories as loops and the remainder as non-loops. We applied loop closure loss on trajectories that formed loops and all other losses on the entire dataset of trajectories. Fig. R1 in the attached PDF demonstrates this paradigm for the 2D Stretchy Bird environment, where we continue to get consistent velocity representations with low transformation error to the true velocity distribution.

In greater generality, we also note that the loop closure loss can also work on trajectories that are “almost loops” – i.e., a loss whose prefactor term scales with how close a trajectory is to forming a closed loop. However, running this analysis and tuning these parameters is beyond the scope of the current rebuttal period, and we will examine these results for the camera-ready version of our paper.

Situation when the dimension of velocity is higher than grid cells (2D)

Thank you for raising this point. We have included an example of an environment where the dimension of the velocity was higher than two-dimensional, cf. the 3D Stretchy Bird environment. This, however, does not pose a problem for integration by grid cells. While a single grid cell module can only integrate two-dimensional velocities, previous theoretical work (Klukas et al. 2020, PLOS Computational Biology 16(4): e1007796) has shown that multiple grid cell modules can collectively represent large volumes of high-dimensional Euclidean spaces through integration of velocities in higher dimensions. We will add a brief discussion of this point in the text.

We hope that you agree that our paper has improved given your detailed feedback and suggestions! Please let us know if there is anything else that we can clarify. We kindly ask you to consider revising your score if you agree that the primary concerns raised in the review were addressed.

We would like to thank the reviewer for their detailed feedback and for raising excellent questions. We will try to address all questions one-by-one.

Figure 4 var/axes and Figure 3 isotropy equation

Thank you for the suggestions! We will make the necessary edits to the figure.

Dimensionality reduction

Thank you for raising this point. Yes, our method is indeed a dimensionality reduction method. Our model estimates low-dimensional velocities between high-dimensional states, which can then be integrated to determine the low-dimensional representation of a given state. Dimensionality reduction methods typically apply over an ensemble of states, with the low-dimensional representation of a single state examined in relation to the low-dimensional representation of other states. Our method directly constructs these relative relationships by extracting low-dimensional velocities. As an example, we have generated a new figure (Fig. R3 in the attached PDF), showing states generated from a trajectory in the Moving Blobs environment. Through integration of our model-outputted velocities, the high-dimensional states within this trajectory fall on a plane (consistent with the data, since it can be minimally described through transitions in $\mathbb{R}^2$). PCA of the same set of states fails to capture this low-dimensional description, with ~24 principal components necessary to capture >95% of the variance.

Preservation of cell-cell relationships

The key prediction of our model is that the cell-cell relationship between two cells is preserved across task modalities and across abstract cognitive domains. To simplify our explanation, we consider a grid cell model with a single grid module. Preserving cell-cell relationships does not require the model to simultaneously receive inputs into the same grid module from multiple tasks. The only requirement is that the estimated velocities from each task are fed into the module, but not at the same time. The question of cell-cell relationships is whether, if two cells fire in an overlapping way in one mapped space/domain, they continue to do so in other domains; and if they fire in a non-overlapping way in one domain, whether they continue to be non-overlapping in all other domains (cf. Yoon et al. 2013 Nat Neuroscience). Our hypothesis is that the brain generates velocity signals from these different domains and feeds these signals into a fixed, prestructured grid cell circuit. Following the structure of the continuous attractor model assumed for the grid cells, we observe that if two grid cells are co-active for one task, they remain co-active for all other tasks. While this claim relies on the underlying grid cell model, it is not presupposed by it: preservation of cell-cell relationships follows from our hypothesis that low-dimensional velocities can be extracted from abstract domains, and thus the same continuous-attractor-based grid module can perform integration across domains.

Orthogonality of velocities in triplet

As discussed briefly in L127-129, we assume that $v_{23} = v_{12}$ and ask the model to predict an $i_3$ solely to demonstrate that the model does not memorize the image features of $i_2$ through some encoding in the predicted velocity / latent space. We will clarify this in our paper. We demonstrate that this assumption is not necessary for model training in Fig. R2 in the attached PDF. Here, we retrain our network on the Moving Blobs task with trajectories that vary in velocity randomly at each time-step. Here we only predict an estimated $\hat{v}$ from $i_1$ and $i_2$, and then predict $i_2$ from $\hat{v}$ and $i_1$. We show that qualitatively and quantitatively, the obtained results remain very similar to the training paradigm we previously employed.

Missing limitations: the biological plausibility of the proposed method

This is a great question. We believe our core loss framework is biologically plausible. The two critical losses we propose are the next-state prediction loss and the loop-closure loss: (1) The next-state prediction loss may occur through sensory systems computing an error that describes the difference between true and predicted sensory input; (2) The loop-closure loss can be computed through a neural integrator circuit such as the grid cell system. Experimental results show that grid cells process these abstract domains, so it is plausible that they provide some error / feedback towards the accurate mapping / organization of these spaces. We include further discussion related to this plausibility in text.

In addition to the critical losses, we also used two auxiliary losses. The biological plausibility of these losses is not as clear as the critical loss terms; however, we note that they primarily aided in refinement of the obtained representations and were not crucial to our results (as demonstrated in our ablation results). We also fully agree that the optimization process we use (stochastic gradient descent and backpropagation) is not a process supported by biology. We will add these limitations to our discussion section.

How naturalistic is the data collection process that requires closed-loops of random walks?

Thank you for raising this question.

While our models were trained with data consisting of loops, this training paradigm was for convenience (L179-181) and is not a necessity of the model. If arbitrary random walks were selected instead for training samples, self-intersections would automatically lead to loops within sub-sequences of the walks (as detected by $k<m$ with $i_k=i_m$) — these loops could then be used for the loop-closure loss, with all other losses applied on all loop and non-loop trajectories. As a simple proof of concept, we generated a dataset of arbitrary trajectories with 50% of the trajectories as loops and the remainder as non-loops. We applied loop closure loss on trajectories that formed loops and all other losses on the entire dataset of trajectories. Fig. R1 in the attached PDF demonstrates this paradigm for the 2D Stretchy Bird environment, where we continue to get consistent velocity representations with low transformation error to the true velocity distribution.

In greater generality, we also note that the loop closure loss can also work on trajectories that are “almost loops” – i.e., a loss whose prefactor term scales with how close a trajectory is to forming a closed loop. However, running this analysis and tuning these parameters is beyond the scope of the current rebuttal period, and we will examine these results for the camera-ready version of our paper.

Missing grid cell model details

Thank you for raising our attention to this. We provide here a brief description of the grid cell model used and will include a more detailed description in the appendix of the revised paper.

We use an approximation of a continuous attractor model for a module of grid cells: we simulate patterned activity on a lattice of neurons as the sum of three plane waves, resulting in a hexagonal pattern of activity. Input velocities are used to update the state of the activity on the lattice of neurons through updating the phases of the plane waves, leading to accurate integration of the input velocities.

Missing data generation details

Thank you for raising our attention to this as well! We will provide additional details about the synthetic data generation in the appendix section. We briefly touched upon the training / testing set sizes and velocity distributions from which each synthetic environment was generated in the appendix, but we will add more details. We plan to include relevant details about the construction of the training / test set (e.g., how we generated random loop trajectories) and statistics about the environment states (e.g., size of image frames, maximum velocity allowed between consecutive states, etc.).

Network vs. circuit model; Missing section number

Thank you for the excellent suggestions. We will change the terminology used to describe our model and fix the section numbering in our paper.

Thank you for your rebuttal comments, additional results, and manuscript improvements!

I urge you to clarify the "cell-cell relationships" discussion more in the paper, since the language can be a bit opaque, particularly for the non-grid-cell community. Maybe it would be good to discuss how this prediction could have been falsified (e.g. is it only if you weren't able to learn a good velocity signal to feed into the assumed grid cell model?). Also, in future work, it would be valuable to address the biological plausibility issues to make this a better model of the brain.

Overall, your rebuttal has addressed my feedback, which did not raise major concerns. I am happy to maintain my original high score and support this paper's acceptance.