Flexible mapping of abstract domains by grid cells via self-supervised extraction and projection of generalized velocity signals
Authors: Abhiram Iyer
Sarthak Chandra
Sugandha Sharma
Ila R Fiete
Venue: NeurIPS 2024 poster (2024)
Keywords: grid cells
cognitive mapping
cognitive maps
self-supervised learning
entorhinal cortex
path integration
neuroscience
dimensionality reduction
Primary area: neuroscience_and_cognitive_science

Abstract: Grid cells in the medial entorhinal cortex create remarkable periodic maps of explored space during navigation. Recent studies show that they form similar maps of abstract cognitive spaces. Examples of such abstract environments include auditory tone sequences in which the pitch is continuously varied or images in which abstract features are continuously deformed (e.g., a cartoon bird whose legs stretch and shrink). Here, we hypothesize that the brain generalizes how it maps spatial domains to mapping abstract spaces. 
To sidestep the computational cost of learning representations for each high-dimensional sensory input, the brain extracts self-consistent, low-dimensional descriptions of displacements across abstract spaces, leveraging the spatial velocity integration of grid cells to efficiently build maps of different domains.
Our neural network model for abstract velocity extraction factorizes the content of these abstract domains from displacements within the domains to generate content-independent and self-consistent, low-dimensional velocity estimates. 
Crucially, it uses a self-supervised geometric consistency constraint that requires displacements along closed loop trajectories to sum to zero, an integration that is itself performed by the downstream grid cell circuit over learning. This process results in high fidelity estimates of velocities and allowed transitions in abstract domains, a crucial prerequisite for efficient map generation in these high-dimensional environments. We also show how our method outperforms traditional dimensionality reduction and deep-learning based motion extraction networks on the same set of tasks.
This is the first neural network model to explain how grid cells can flexibly represent different abstract spaces and makes the novel prediction that they should do so while maintaining their population correlation and manifold structure across domains. Fundamentally, our model sheds light on the mechanistic origins of cognitive flexibility and transfer of representations across vastly different domains in brains, providing a potential self-supervised learning (SSL) framework for leveraging similar ideas in transfer learning and data-efficient generalization in machine learning and robotics.

Flexible mapping of abstract domains by grid cells via self-supervised extraction and projection of generalized velocity signals
Abhiram Iyer 3,4, Sarthak Chandra Sugandha Sharma and Ila F,
Department of Electrical Engineering and Computer Science, MIT, Cambridge, MA Department Brain Cognitive Sciences, MIT; Cambridge, MA
"MGoovem Institute for Brain Research, MIT, Cambridge, MA *K. Lisa Yang Integrative Computational Neuroscience (ICoN), MIT, Cambridge, MA. {abiyer, sarthakc, susharma, sustrhaakcsuuaaaaaa wieteemmit..t.e.eee edu
Abstract
Grid cells in the medial entorhinal cortex create remarkable periodic maps of explored space during navigation. Recent studies show that they form similar maps of abstract cognitive spaces. Examples of such abstract environments include audi- tory tone sequences in which the pitch is continuously varied or images in which abstract features are continuously deformed (e.g., a cartoon bird whose legs stretch and shrink). Here, we hypothesize that the brain generalizes how it maps spatial domains to mapping abstract spaces. To sidestep the computational cost learning representations for each high-dimensional sensory input, the brain extracts self- consistent, low-dimensional descriptions of displacements across abstract spaces, leveraging the spatial velocity integration of grid cells to efficiently build maps different domains. Our neural network model for abstract velocity extraction factorizes the content of these abstract domains from displacements within the domains to generate content-independeet and self-consistent, low- low-dimensional velocity estimates. Crucially, it uses sulf-supervised geometric consistency straint that requires displacements along closed loop trajectories to sum to zero, integration that is itself performed by the downstream grid cell circuit over learning. This process results in high fidelity estimates of velocities and allowed transitions in abstract domains, crucial prerequisite for efficient map generation in these high-dimensional environments. We also show how our method outperforms traditional dimensionality reduction and deep-learning based motion extraction networks on the same set of tasks. This is the first neural network model to explain how grid cells can flexibly represent different abstract spaces and makes the novel prediction that they should do so while maintaining their population correlation and manifold structure across domains. Fundamentally, our model sheds light on the mechanistic origins of cognitive flexibility and transfer of representations across vastly different domains brains, providing potential self-supervised learning (SSL. framework for leveraging similar ideas in transfer learning and data-efficient generalization machine learning and robotics.
Introduction
Grid cells in the medial entorhinal cortex of paramount importance for navigating and representing spatial domains. Interestingly, a series of recent experiments have shown that the brain still uses the same cells to represent non-spatial environments that are continuously traversed. These include free
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

firing fields
b
spk cade (40)
parce
Lini Losp)
P
placess where cell fired
hejedbay
wit w- w-tR gold
movement direction
helpet
learned map consists of state-dependent transitions
with
velcaity
blob height
learned independent
dm3
biob with
dim
Figure 1: Conceptual understanding of learning velocities in abstract cognitive environments. a. Grid cells generate hexagonal tuning in spatial navigation experiments (illustration borrowed from (I). b. Similar grid-like tuning has been found in auditory frequency space in rodents 2]. (c.) in visual space in monkeys [3], and (d.) in an abstract cartoon bird space in human experiments. e. Example of an abstract, non-spatial domain called 'Stretchy Blob', 2D Gaussian that can either stretch or shrink along two axes. f. SR, CSCG, and TEM 4477 learn transition structures of this cognitive domain by traversing These models require encountered states as well as structures enabling transitions between them. g. Our approach learns self-consistent movement (velocity) signal that is independent of the states traversed and encodes the global transition structure of the enviromment in minimally low-dimensional representation. h. How can the brain initially extract a general notion of velocity within abstract, non-spatial domains? Solving this important prerequisite challenge can show how grid cells flexibly map and organize various abstract spaces. The actual movement signals this particular example are low-dimensional. These velocities are independent from the states they connect.
build of representations for these
viewing of naturalistic images, listening to and modifying a sound changing in pitch, navigating conceptual 'Stretchy Bird' space (images of bird that stretch or shrink via joystick input), or even traversing an 'odor' space among many others [814 [5]
(Fig. Ia-d).
How does the brain transfer its ability to represent space these non-spatial environments? More specifically, how does the brain infer a metric layout in these abstract domains? Answering this question involves understanding how the brain perceives and processes structure in different domains and modalities, and integrates them into coherent cognitive maps.
How the brain maps abstract cognitive environments has been question of intense interest, with several proposed approaches (7] These works propose that when the brain learns the transition structure of cognitive domain (such the 'Stretchy Blob' space in Fig. [e.) it simultancously builds a set of representations for these states as well as structures enabling transitions between them, Fig.
Our approach investigates how the brain exploits an existing scaffold of structured states, provided by grid cells, to represent new cognitive domains by projecting them onto the invariant grid coding space. We hypothesize that this projection involves anchoring a state in the cognitive domain to grid phase and then extracting self-consistent movement signal to measure displacements in the cognitive domain (Fig. Tg. This abstract velocity signal serves as the input to the grid cell network, updating grid phases and encoding transitions between high-dimensional states using the existing

---

low-dimensional transition mechanism within the grid cell circuit. Crucially, rather than building new set of continuous stable states for each explored cognitive domain - process that both difficult and slow - this mechanism allows for the efficient reuse of a canonical scaffold of cognitive states for the memory and coding of continuous variables.
In spatial domains, continuous attractor models of the grid cell circuit [16] efficiently reuse prestruc- tured grid cell states encode transitions between high-ddimensioal states. However, using such models to map non-spatial, abstract domains (and by extension, explaining how grid cells can be effi- ciently reused) requires the extraction of faithful representations of velocity in these domains. Thus, the fundamental, prerequisite challenge lies in extracting a general notion of minimal dimensional velocity signal from high-ding sigh-dimensional, Limevaryyig data from various abstract domains, Fig. L Crucially, the representations of this extracted velocity must be independent of the specific states, and must be self-consistent with a net-zero velocity corresponding to a net identity transformation in the abstract space. These constraints point towards a self-superviser learning paradigm, wherein neither the coordinate system of the space nor the dimensionality of the estimated velocity are known priori.
The following are the key contributions of the paper:
First neural network model for abstract velocity extraction. We present the first neural network model explaining how grid cells flexibly represent different abstract spaces through providing learning framework for velocity estimation arbitrary spaces. This reuse of grid cells across domains leads to cognitive flexibility in mapping spaces, providing insights towards transfer learning and data-efficient generalization in machine learning and robotics.
State-independent velocity extraction. Our framework for abstract velocity extraction gen- erates state-independent, self-consistent low-dimensional velocity estimates by separating the content abstract domains from their displacements (Fig. (gig.Tg.i).
Self-supervised geometric consistency constraint. We show how consistent metric rep- resentations of velocity can be learned through self-supervised geometric consistency constraint requiring displacements along closed loop trajectories to sum to zero an operation facilitated by the downstream grid cell circuit.
Superior dimensionality reduction performance. Our method surpasses traditional di- mensionality reduction and deep learning- based motion extraction networks. specifically in environments characterized by underlying low-dimensional transitions and motions between states.
Preservation of cell-cell relationships. The model predicts that cell-cell relationships between grid cells are preserved across different spatial and non-spatial domains. For example, if two grid cells are co-active in spatial task, they should remain co-active in non-spatial task. This prediction critically relies on the same continuous attractor- based grid module performing integration across domains realized through extraction of a state-independent velocity from abstract domains. This prediction aligns with the observed invariance of internal neural representations across different spatial environments and brain states [3] This forms testable hypothesis for future experimental studies neural representations in abstract cognitive domains.
capability that can only be
1.1
Related work
Our work on learning velocities and transition operators in abstract cognitive spaces can be compared to two key areas: neuroscience models of spatial mapping and dimensionality reduction models for high-dimensional data. We first discuss cognitive space mapping models and then relate them to dimensionality reduction work.
The
Machine (TEM) 7 learns map of 'cognitive domain' via a recurrent neural network, predicting sensory observations based on given actions. However, TEM does not infer transition structures or affordances, relying instead on predefined actions at each time step. When presented with new environment, TEM requires re-learning these affordances along with grid cell-like representations from scratch. Similarly, the successor representation (SR) future states as weighted sum of expected future occurrences and uses an underlying discrete action space. State-action variants of SR are dependent these discrete action inputs to construct
4 5 predicts
3

---

cognitive maps. Clone-structured cognitive graphs (CSCG) [6] learn hidden Markov models of sensory representations from observational imputs without prior domain structure assumptions but are limited to only discrete domains. By construction, these models are not guaranteed to preserve cell- cell correlations across domains and modalities and are unable to efficiently reuse the prestructured states and transition operators provided by grid cells.
More generally, these models posit that the brain simultaneously learns representations of the external states and the transition structures to form cognitive maps. Here, we propose that approaching the cognitive mapping problem from a selocity-extraatioon learning a www-dimesiional velocity that captures transitions between external inputs, mapping the environment reusable grid cell states can be performed simply by continuous attractor network, instead of something more complex like TEM. More generally, via our self- self-supervised learning framework that infers velocity solely from sensory inputs in continuous domains, we can entirely avoid the computationally expensive task of representation and transition learning through capturing the minimal low-dimensional structure of these abstract observations.
perspective offers key benefits. By
Our work thus bridges the tasks of building cognitive maps of abstract domains and dimensionality reduction. Previous dimensionality reduction methods 244 focus on preserving proximity in high-dimensional spaces on low-dimensional manifolds but do not take advantage of the structure imposed on the tangent spaces of these manifold (i.e., velocity spaces) through transitions between states. As result, while these methods learn oww-dimensional mappings, they often fail to reproduce the underlying metric space of abstract domains. Unlike these methods, our framework learns global velocity operator applicable to any input within the environment's state manifold. Our work also relates to the task of motion decomposition and next-frame prediction of temporally varying inputs 30 However, these works have not attempted to decipher minimally low dimensional description of velocity, making their learned representations unsuitable for grid cells to effectively map observed environments.
2 Self-supervision for velocity extraction
Learning action primitives purely through self-supervision, without access to true velocities in the space, requires careful consideration of the data generation process, the loss terms to extract meaningful neural representations, and the chosen neural network architecture. We will discuss each of these aspects in detail.
2.1
Task setup
The initial step in our methodology involves the construction of tasks designed to facilitate the model's inference of the lowest-dimensional representations of velocity within cognitive domains.
We revisit the 'Stretchy Blob' environment that we previously introduced. This environment can be procedurally generated where state transitions are characterized by changes in the blob's width and height. As seen in Fig. hh, given two images, i and i2, we aim to learn function f that infers a low-dimensional velocity é representing the transition from i i2 (Fig.) 2
To ensure that é a useful and faithful metric representation of the transition between two inputs, we demand that it satisfy two properties.
First, the estimated transition velocity i can used to transform given input, i.e., there exists a function y that uses é to transform i to predict ig. In practice, to ensure that f does not merely memorize image features of ig, we demand that g predict an unseen i3 obtained by transforming i2 using the same transition é. More generally, this ensures that the estimates of displacements, ú, in this space must be independent the route taken, ensuring both path and state independence. We refer to this requirement as "next-state prediction'.
Second, the estimated velocities in the abstract space, é, must be geometrically self-consistent, i.e.. the start and end point of a trajectory in the abstract space are identical if, and only if, the sum of estimated velocities zero (Fig. 2 (Note that this also ensures that the identity transformation should be represented by the pero-vector) We refer to this requirement as 'loop-closure'. This summation of velocities must be performed by a neural integrator, such as grid cells, which results in calibrated velocity input for grid cells to then themselves use map out the input abstract domain.

---

velocity estimation
next-step prediction
blob height
U1
((1,,
12)
i g
blob width
loop closure
d
12
14
Stretchy Blob (2D)
Moving Blobs (2D)
Stretchy Bird (2D)
image manifold
Stretchy Bird (3D)
N Ut 10
Frequency Modulation (1D)
IN
Figure 2: Task and fundamental problem setup. a. Two states in the 'Stretchy Blob' task that have a low-dimensional velocity representing the transition between them. b. two consecutive images in this space i and i2 are separated by velocity 01-2,, can we learn function f that estimates this velocity and a function g that 'applies' this quantity i to predict iz? c. The key self-consistency constraint we introduce called 'loop-closure': estimated velocities along a closed-loop trajectory must sum to zero. This computation performed by a neural integrator, such as grid cells. d. Our various procedurally generated abstract cognitive domains: Stretchy Blob (2D), Moving Blobs (2D), Stretchy Bird (2D), Stretchy Bird (3D), and Frequency Modulation (ID).
These self-consistency constraints on è suggests a self-supervised learning paradigm, obviating the need for a specific coordinate system within the explored abstract space.
To rigorously evaluate our system, we procedurally generate random trajectories across five different "abstract cognitive domains." These trajectories are generated by starting from an initial random state on the image manifold and then taking random velocities to determine the subsequent states on the same manifold. We assume for simplicity that states are unique at each point space within these domains. We visualize a randomly sampled trajectory of each of these domains in Fig. 2d:
1. A 2D Stretchy Blob environment (discussed above) where a blob in the center of the visual field stretches or shrinks height and/or width.
2. A 2D Stretchy Bird environment where bird's legs and neck stretch and shrink. This environment is specifically constructed to mimic its experimental counterpart (15).
3. A 3D variant of Stretchy Bird where the two bird legs can independently transform.
4. A 2D Moving Blobs environment where a set of Gaussian blobs uniformly translate, moving in and out of the visual field.
5. A ID Frequency Modulation task that emulates the experimental sound modulation task created by 2 with a sum of sine waves uniformly changing in frequency.
Detailed information regarding the generated data provided in Sec.
Encoder-Decoder architecture
We construct encoder-decoder architecture, wherein the encoder f and decoder y are both multi- layer perceptrons (MI.ss) (Fig. (3)). The encoder f processes two adjacent inputs in and i from any of our procedurally generated trajectories and maps them to low-dimensional velocity space. It
5

---

1++1
decoder
02++2
Integrator:
Next-State Prediction
Closure
min di d0222 6422
min || edr
imegge manifold
image menitoid
Isotropy dees)
Shortout Estimation
61-5 61-12 + 02-0
Pa
02-44 62-93 + Rg-ed
Are
Bre
Image manifoid
image manifoid
Figure 3: Self-supervised learning framework. a. Model diagram consisting of an encoder, decoder, and an integrator which acts the oow-dimensional velocity latent space. The model takes two consecutive imput frames and predicts an unseen frame with the learned velocity. various self- oelf-supervise loss terms. The two critical loss terms (top) are 'next-state prediction' and 'loop-closure.. The auxiliary loss terms (hottom) which further refine the solution space are 'shortcut estimation' and "isotropy".
Our
important to note that the low-dimensional velocity signal not known a priori. The decoder then upsamples this latent velocity representation and combines it with the last input to generate the subsequent 4,42 which remains unseen to the model.
single sample in batch of data for this model comprises a trajectory of states T 1 N} and velocities V = {V+-+411 < < N - 1} (individual velocities abbreviated as for ease). The model is trained through triplets of frames of the input trajectory, using i and 44+1 to predict it42 for 1 < t S N 1. As discussed above, we train our models on triplets of states solely to ensure that the encoder does not memorize features of the image be predicted. Training on pairs of images instead of triplets does not affect any of our results (SI Sec. B.2).
Detailed experimental procedures regarding the training of this model across various constructed domains are provided in SI Sec. C We note that the same architectural motif was employed for training all our experiments.
2.3 Loss functions
We formalize our two requirements the extracted low-dimensional velocity signal into two critical loss terms, next-state prediction loss, and loop-closure loss. These losses from the core of our self-supervised learning framework. To further refine the solution space, we also employ auxiliary losses in addition described in detail in SI Sec. C
primary constraints. Our loss terms are visualized in Figure 3p and
Next-State Prediction Loss. We quantify this based on the difference between the next- state prediction 4ez2 and the ground-trutt truth frame 4 This loss term operates on individual samples of the generated trajectory, and ensures decodability of the estimated velocity 6.
Loop-Closure Loss. We quantify this as a norm of the sum of velocities along closed
loop trajectory T. i.e., the model must produce velocity estimates such that Σner 0 = Ō. The error signal for this loss operates at the scale of the entire generated trajectory. For convenience, construct all trajectories of our training data random loops in the

---

considered abstract spaces, such that the start and end state of trajectories are identical. See SI Sec. B.2 for training data that solely loops.
Shortcut Estimation Loss. The first of our two auxiliary loss terms is a shortcut estimation loss, which further tests the generalization ability of our decoder From &t+1, we predict future states 44+33 or 4144 by directly modifying Specifically. if 02-13 inferred from 4-2 to 44+3 and 03-44 inferred from 4++3 4t+-4, then ie+4 should be 02-33 03->4 away from 4.42. This loss is important for further refining our velocity estimates and ensuring their validity during generalization.
Isotropy Loss. The loss terms considered so far do not ensure isotropy in the inferred velocity space, allowing differential scaling factors for transformations in different input space directions. To induce isotropy, we introduce an auxiliary isotropy loss term that acts on the norm of the velocities, independent of direction. Since we don't assume access to the global velocity distribution in the training data, the isotropy loss applied only near zero velocity. In particular, we minimize the variance of {116-+t+11) d(ée, i(+1) < 0), where d(ie, 1441 is a similarity metric the input image space, and is some small threshold. In practice, use 1- cosine similarity as our distance metric d.
There is no direct supervision signal regressing the model outputs onto known distribution of velocities; instead, the velocities are extracted and inferred automatically. Further, we do not priori assume knowledge of the dimensionality of the underlying transition structure of the training domains.
Regardless of the training environment, the relative weighting ratio of the two critical loss terms remains consistent, with loop-closure loss always weighted ten times higher than next-state prediction loss. To stabilize the training of models with three-dimensional latent spaces, we include small L regularization during training. All models are evaluated on an unseen testing dataset consisting of random walks within the domain. To prevent overfitting, all models are deliberately underparame- terized relative to the training dataset. conduct ablation studies on our loss terms, which can be found in SI Sec. B.
3 Experimental results
Stretchy Blob (2D)
Moving Blobs (2D)
Stretchy Bird (2D)
Stretchy Bird (3D)
Frequency Modulation (1D)
a
Latent dim manifold
Latent manifold dim
917%
100%
Figure 4: Model results. Our model produces low-dimensional velocity latents that similar to the ground truth (g.t.) distribution without knowing this distribution across a variety of cognitive environments. b. In cases where the model's 's latent dimensionality is higher than the intrinsic velocity dimensionality of the environment, our model still identifies the lowest-dimensional representations embedded in higher-dimensional space.
3.1 Single learning framework infers geometrically consistent representations of velocity across cognitive domains
For each abstract environment, we compare the model-inferred velocity representations (Fig. 4h) to the true distribution of velocities. Note that the self-supervised framework does not result in exact identity mapping between the ground truth velocities and the model outputs it suffices for the obtained output to be linear transformation of the ground-truth velocity space. Correspondingly,
7

---

Task
Our Model
MCNet
Autoencoder
PCA
Isomap
UMAP
Stretchy Blob (2D)
0.05 0.01
1.95 = 0.14
(1.59 + 3.40) x x10²
0.63
0.42
0.79
Stretchy Bird (2D)
0.07 = 0.03
2.01 0.01
20,90 = 38.04
0.21
0.36
0.46
Stretchy Bird (3D)
0.07 0.02
2.96 0.10
(2.66 + 5.08) x102
0.31
0.64
1.05
Moving Blobs (2D)
0.02 0.01
2.00 0.01
2.03 + 0.05
1.94
0.66
0.62
Freq. Modulation (1D)
0.02 + 0.02
2.01 = 0.01
2.00 + 0.01
1.97
2.00
2.00
Table Mean and standard deviation of errors for various tasks and models. Mean and standard deviation of eπrоrs computed across different rums for each experiment. Each run seeded to ensure reproducibility. The 6 seeds were picked randomly and are the same seeds used across different experiments where multiple runs were run.
construct an error metric on the inferred velocities as the mean squared error in mapping the predicted velocities to the ground-truth via a single linear transformation, after removing small number of outliers via the DBSCAN clustering algorithm 36.. More details about this error metric can be found in SI Sec.
In all cases, irrespective of the dimensionality of the input manifold space or the detailed statis- ties and structure of the environment states, we see that the inferred velocities are faithful metric representations of the ground-truth velocities, quantified in Table
While we primarily consider cases where the latent dimensionality of the encoder output matches the underlying dimensionality of transitions in the input space, this is not necessity. In Fig. 4p. we set up our framework to have latent dimensionalities larger than the true data manifold dimensionality. In all cases, the model outputs automatically occupy subspace of dimensionality that corresponds to the actual input manifold transition space, with a PCA of the inferred velocities capturing greater than 97% of the variance within the correct number of dimensions (ef. Fig. 4p). Thus, our model can effectively identify the appropriate low-dimensional structure within the high-dimensional embedding space of the inputs. Results from the other synthetic cognitive domains can be found in SI Sec. A
3.2 Comparison to existing dimensionality reduction methods
Our model
PCA
Isomap
UMAP
Autoencoder
MCNet
Moving Blobs
(2D)
0
b
Frequency Modulation (1D) g4.
Figure 5: Comparison to baselines. We show model comparisons various dimensionality reduction and motion-ppedictiin baselines in the a. 2D Moving Blobs and b. ID Frequency Modula- tion tasks. Existing baselines cannot identify the diw-dineniional velocity signals between arbitrary transitions in this space, even failing to do so a simple one- me-dimensional domain. Meanwhile, our model produces results that closely match the true, underlying velocity distribution.
Our model estimates low-dimensional velocities between successive high high-dimensione -iieensional states. These velocities can then be integrated to determine low-dimensional representation for each state. In this sense, it possible to view our work as dimensionality reduction method for continuously varying inputs. Traditional dimensionality reduction methods rely the statistics of distances between points across an ensemble of states. contrast, our approach finds a structured tangent manifold around each state that captures the low-dimensional transitions successive states.
We can compare standard dimensionality reduction techniques such as PCA, Isomap, UMAP, and deep autoencoders with method. To do these techniques to embed the inputs into the known

---

b
Our model
PCA
dim 256
ofvarisncs wistin df
dim
dim
dim 255
dim dim
dim
dim 1
included dims (x)
Figure 6: Dimensionality reduction through our model, versus PCA. a. Schematic of the raw input states from a random trajectory the 2D Moving Blob task, showing the states as points in a 16x16 dimensional space. b. We estimate 3-dimensional velocities between states, and integrate these estimated velocities obtain low-dimensional representation of the initial input states. A 2 2-dimensional plane is shown in gray for perspective, demonstrating our model-pree dimensional representations are approximately in a 2D subspace. c. Left: Computing PCA on the same dataset shows representations occupying volume in 3-dimensional space. Right: Around 24 dimensions are required for PCA to capture 95% of the variance the data, indicating that PCA unable to find low-dimensional space describing the dataset.
model-produced pproduced low-
latent dimensionality of the data manifold, then compute an estimated wiweemmmnnioaal velocity between states by taking the difference of their corresponding low-dimessional representations.
Fig. 5 compares our model's 's estimated velocities with those produced by these standard baselines. Since our self-supervised framework uses a next-step prediction component, we also compared our results to MCNet (30] flexible deep network designed specifically for future frame prediction) and constrained the model to use low-dimensional representations of the transitions between frames. In all cases, our model significantly outperformed other baseline models (cf. Table [ and produced velocities that were closely aligned with the true velocity distributions. Remarkably, for even one- dimensional manifolds embedded in high-dimensional spaces (as in our ID Frequency Modulation task), existing dimensionality reduction techniques struggle to find a coherent representation for velocity and produce errors that are two orders magnitude larger than our model.
We also examine dimensionality reduction through constructing representations of the original data integrating the model velocity estimates, Fig. 6 For a sample trajectory through the high- dimensional states of the 2D Moving Blob environment, our models representation collapses onto two-dimensional plane (consistent with the data, since the states can minimally described through transitions in R2). 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.
3.3 Model outputs allow reuse of grid cells mapping abstract domains
Stretchy Blob (2D)
Moving Blobs (2D)
Stretchy Bird (2D)
Frequency
Modulation (1D)
blob height
blob bloby-cood y-coords
fning rate
blob width
blob x-coords
neck length
frequencies
Figure 7: Model outputs can be used to generate grid-like firing fields. Our model's outputs can be used as input to a synthetic grid cell network across variety of cognitive domains. We predict clear hexagonal-like firing field when traversing these environments, illustrating how the grid cell circuit is crucial to building these cognitive maps.
We hypothesize that the brain can map abstract cognitive domains to manifold of grid cell states by learning low-dimensional velocity representations of the input space. Extracting this low-dimensional signal allows the brain to reuse continuous attractor dynamics and path integration functionalities to stably represent and traverse along the manifold of inputs in the mapped domain.

---

To verify that the model inferred velocities are faithful enough for accurate path integration, we construct the tuning curve random neuron from a synthetic grid cell module (details of which are briefly described in SI Sec.C) that is fed these inferred velocities (up to best-fit linear transformation; for simplicity we choose our transformation to map to aligned velocities across environments, leading to orientation-aligna curves) as the path integration inputs. As seen in Fig. 7 our model results in hexagonal tuning curves in the abstract cognitive domains, consistent with previous work grid representations [8] 14 3 15 2 (grid firing fields from baseline models explored in SI Sec. A.. Through extracting faithful www-iieesiional representations of velocities across abstract domains, the same continuous attractor-basee grid modules can be used across tasks. As a result, if two grid cells fire in an overlapping (or von-overlappingg way in one mapped domain, they continue to be overlapping (or non-overlapping) in all other domains. Thus, cell-cell relationships are preserved between grid cells across abstract domains. This forms testable hypothesis for future experimentation, that may be falsified the brain were to use distinct, independent grid cell modules to organize information from different modalities, or if grid cells were not prestructured networks that function independently of the nature of the inputs.
grid tuning curves this need not hold in experimentally observed tuning
We also note that while each grid cell module is a two-dimensional toroidal manifold, enabling integration of awo-dimensional velocities, the grid cell system consisting of multiple modules can integrate velocities in higher dimensions Thus, higher dimensional velocities extracted by our model (e.g.. 3D Stretchy Bird) do not pose a problem for integration by grid cells.
4
Discussion
Our work introduces the first neural network model that can infer velocities within abstract cognitive domains. This enables circuits like grid cells to encode transitions between high nigh-dimensional sensory states through low-dimensional path integration, mapping sensory inputs to prestructured states instead of learning independent representations for each state. Our velocity extraction mechanism itself requires path-integration (via the 'loop-closure' loss), necessitating a grid-cell-like neural model, highlighting how grid cells form the foundation of mapping abstract spaces.
Future Work. Our research offers new perspectives for neuroscientists on the flexible utilization of grid cells to organize and map non-spatial domains. Future neuroscience research may test our hypothesis on the conservation of cell-cell relationships across cognitive domains and identify brain regions that generate velocity signals, along with their experimental signatures. Future machine learning research directions include scaling our framework for naturalistic environments, learning non-Euclidean spaces like family trees, and possibly using our framework to augment existing cognitive space-mapping models like TEM or CSCG.
Limitations. While our core SSL loss terms (next-state prediction and loop-closure) are biologically plausible and may align with sensory prediction error and neural integration, the auxiliary losses (shortcut and isotropy) are less biologically supported. Additionally, we assume velocity vectors in our latent space commute, which prevents them from directly representing tangent vectors in Enclludlld spaces like sphere. However, non- -on-Euclidean spaces can be represented by embedding them in higher-dimensions Euclidean space where velocities commute (e.g., sphere embedded in dimenssinlii space) 39
Broader Impact to AI. Our novel SSL framework can be applied for invertible dimensionality reduction (by virtue of a generative decoder which generates high-dimensional states corresponding to points in a low-dimensional latent) and manifold learning tasks. Our model significantly outperforms non-invertible dimensionality reduction baselines on datasets that contain relatively lower-dimensional transitions (suggesting applications to video data, for example). Our method also naturally lends itself to manifold alignment-related challenges, which is particularly effective when the data exhibits small number of continuous modes of variability. Moreover, with small number of "gluing" points, our method allows for building me-to-one correspondences between different domains. Our work also shows how a fixed integrator circuit can leverage common velocity representations to navigate between abstract spaces efficiently. For example, in complex maze where learning action strategies is costly, our model maps transitions and actions to the grid-cell representational space, enabling strategies learned in a simpler, topologically similar space to be effectively applied to the complex domain relevant challenge in robotics.
10

---

Acknowledgements
We thank the McGovern Institute and the K. Lisa Yang Integrative Computational Neuroscience (ICoN) Center for supporting and funding this research. Al was supported by the K. Lisa Yang Integrative Computational Neuroscience (ICoN) Fellowship.
References
[1] Rylan Schaeffer, Mikail Khona, Tzuhsuan Ma. Cristobal Eyzaguirre, Sanmi Koyejo, and Ila Fiete. Self-supervised learning of representations for space generates multi-modular grid cells. Advances in Neural Information Processing Systems, 36, 2024.
[2] Dmitriy Aronov, Rhino Nevers, and David W Tank. Mapping of non-spatial dimension by the hippocampal-cntorhinal circuit. Nature, 543(7647):719-722, 2017.
[3] Nathaniel Killian, Michael J Jutras, and Elizabeth A Buffalo. A map of visual space in the primate entorhinal cortex. Nature, 49726))77777 November 2012.
[4] Peter Dayan. Improving generalization for temporal difference learning: The successor repre- sentation. Neural computation, 514)6613 1993.
[5] Kimberly Stachenfeld, Matthew M Botvinick, and Samuel Gershman, The hippocampus as a predictive map. Nature neuroscience, 20(11):1633-165. 2017.
[6] Dileep George, Rajeev V Rikhye, Nishad Gothoskar, J Swaroop Guntupalli, Antoine Dedieu, and Miguel Lázaro Gredill.. Clone-structured graph representations enable flexible learning and vicarious evaluation of cognitive maps. Nature communications, 12(1):2392, 2021.
[7] James CR Whittington, Timothy H Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy EJ Behrens. The tolman-eichenbaum machine: unifying space and relational memory through generalization in the hippocampal formation. Cell, 183(5): 1249 1263, 2020.
[8] Niklas Wilming, Peter König, Seth König, and Elizabeth A Buffalo. Entorhinal cortex receptive fields are modulated by spatial attention, even without movement. Elife, 7, 2018.
[9] Seongmin A Park, Douglas Miller, Hamed Nili, Charan Ranganath, and Erie D Boorman. Map making: Constructing, combining, and inferring on abstract cognitive maps. Neuron, July 2020.
[10] Seongmin A Park, Douglas S Miller, and Erie D Boorman. Inferences on a multidimensional social hierarchy use grid-like code. Nature neuroscience, 24(9): 1292-1312, 2021.
[11] Matthias Nau, Tobias Navarro Schröder, Jacob LS Bellmund, and Christian F Doeller. Hexadirec- tional coding of visual space in human entorhinal cortex. Nature neuroscience, 21(2):188-19) 2018.
[12] Xiaojun Bao, Eva Gjorgieva, Laura K Shanahan, James D Howard, Thorsten Kahnt, and Jay A Gottfried. Grid-like neural representations support olfactory navigation of a two-dimensional odor space. Neuron, 102(5): 10616006667777 2019.
[13] Joshua B Julian, Alexandra T Keinath, Giulia Frazzetta, and Russell A Epstein. Human entorhinal cortex represents visual space using a boundary-anchord grid. Nature neuroscience, 21(2):191-194, 2018.
[14] Simone Viganò, Valerio Rubino, Antonio Di Soccio, Marco Buiatti, and Manuela Piazza. Grid-like and distance codes for representing word meaning in the human brain. Neurolmage, 232: 11776, 2021.
[15] Alexandra O Constantinescu, Jill X O' O'Reilly, and Timothy E Behrens. Organizing conceptual knowledge humans with gridlike code. Science, 352(6292): 1464-1468, June 2016.
[16] Yoram Burak and lla R Fiete. Accurate path integration in continuous attractor network models of grid cells. PLoS computational biology, 1(2)xc000029, 2009.

---

[17] Mikail Khona and lla R Fiete. Attractor and integrator networks in the brain. Nat. Rev. Neurosci., 23(12):744-766, December 2022.
[18] Kijung Yoon, Michael A Buice, Caswell Barry, Robin Hayman, Neil Burgess, and lla R Ficte. Specific evidence low-dimensional continuous attractor dynamics in grid cells. Nat. Neurosci., 16(8):1077-1084, August 2013.
[19] K J Yoon, S Lewallen, A A Kinkhabwala, D W Tank, and I R Fiete. Grid cell responses in 1D environments assessed as slices through 2D lattice. Neuron, 8965):108611099 March 2016.
[20] Richard J Gardner, Li Lu, Tanja Wernle, May-Britt Moser, and Edvard I Moser. Correlation structure of grid cells is preserved during sleep. Nat. Neurosci., 22(4):598- 608, 2019.
[21] Sean G Trettel, John Trimper, Ernie Hwaun, Ila Fiete, and Laura Lee Colgin. Grid cell co-activity patterns during sleep reflect spatial overlap of grid fields during active behaviors. Nat. Neurosci, 22(4):609-617, April 2019.
[22] Richard J Gardner, Erik Hermansen, Marius Pachitariu, Yoram Burak, Nils A Baas, Benjamin A Dunn, May-Britt Moser, and Edvard Moser. Toroidal topology of population activity in grid cells. February 2021.
[23] Ling I. Dong and Ila R Fiete. Grid cells in cognition: Mechanisms and function. Annual Review of Neuroscience, 47, 2024.
[24] Leland Mclnnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv:1802.03422, 2018.
[25] Joshua B Tenenbaum, Vin de Silva, and John C Langford. A global geometric framework for nonlinear dimensionality reduction. science, (9065500):2319923332 2000.
[26] Diederik P Kingma and Max Welling, Auto-encoding variational bayes. arXiv preprint 1311v:1312.6114, 2013.
[27] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
[28] S Velliangiri, SJPCS Alagumuthukrrshnnan et al. review of dimensionality reduction tech- niques for efficient computation. Procedia Computer Science, 165:104-111,21
[29] Weikuan Jia, Meili Sun, Jian Lian, and Sujuan Hou. Feature dimensionality reduction: review. Complex & Intelligent Systems, 8(3):2663-200 2293, 2022.
[30] Ruben Villegas. Jimei Yang, Seunghoon Hong, Xunyu Lin, and Honglak Lee. Decomposing motion and content for natural video sequence prediction. arXiv preprint 1rX.v:1706.0803: 2017.
[31] Yufei Zhang, Jeffrey O Kephart, and Qiang Ji. Incorporating physics principles for precise human motion prediction. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 6164 044-61144 74,, 2024.
[32] Wei Mao, Miaomiao Liu, Mathieu Salzmann, and Hongdong Li. Learning trajectory dependen- cies for human motion prediction. In Proceedings of the IEEE/CVF international conference on computer vision, pages 9489 9497, 2019.
[33] Yunbo Wang, Zhifeng Gao, Mingsheng Long, Jianmin Wang, and S Yu Philip. Predmnt++: Towards a resolution of the deep-in-time dilemma in spatiotemporal predictive learning. In International conference machine learning, pages 5123-5132. PMLR, 2018.
[34] Judith Butepage, Michael J Black, Danica Kragic, and Hedvig Kjellstrom. Deep representation learning for human motion prediction and classification. Proceedings of the IEEE conference on computer vision and pattern recognition, pages 6158-6166, 2017.
12

---

[35] Sergiu Oprea, Pablo Martinez-Gonzalle Alberto Garcia Garcia, John Alejandro Castro- Vargas, Sergio Orts-Escolang Jose Rodriaezo and Antonis Argyros. A review on deep learn- ing techniques for video prediction. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(6):2806-2255 2020.
Garcia-Gareia,
[36] Martin Ester, Hans-Peter Kriegel, Jörg Sander, Xiaowei Xu, et al. A density-based algorithm for discovering clusters in large spatial databases with noise. In kdd, volume 96, pages 226-231, 1996.
[37] Mirko Klukas, Marcus Lewis, and lla Fiete. Efficient and flexible representation of higher- dimensional cognitive variables with grid cells. PLoS computational biology, 1007)9 2020.
[38] John Nash. Analyticity of the solutions of implicit function problems with analytic data. Annals of Mathematics, 84(3):345-335 1966.
[39] Robert E Greene and Howard Jacobowitz. Analytic isometric embeddings. Annals of Mathe- matics, 93(1):189-204, 1971.
13

---

b A Additional experiments
A.1
Additional results across domains
Fig. S includes further experiments of our framework across different domains and also compare them to baselines. Fig. 9 visualizes the grid firing fields using model outputs as input synthetic grid cell network and compares them to the same baselines. Our model produces the most faithful representations of velocity across various domains.
Our model
PCA
Isomap
UMAP
Autoencoder MCNet
Stretchy Bird (2D)
Stretchy Bird (3D)
c Stretchy Blob (2D)
Figure 8: Continued, comparison to baselines. a. Model inferred velocity space in the 2D Stretchy Bird environment compared with baselines. b. Model inferred velocity space in the 3D Stretchy Bird environment compared with baselines. c. Model inferred velocity space in the 2D Stretchy Blob environment compared with baselines.
Our model
PCA
Isomap
UMAP
Autoencoder
MCNet
Stretchy Bird (2D)
Stretchy Blob (2D)
Moving Blobs (2D)
Frequency Modulation
(1D)
Figure Can baseline models produce faithful representations of velocity? We pipe the outputs of our baselines into the same synthetic grid cell network to observe grid firing rates along trajectory. Qualitatively, our model produces the most hexagonal grid firing field comparison to other basclines.
14

---

A.2 Decoder implicitly learns boundaries of training data manifold
While the encoder is -ndepeendenn and infers generalized notion of velocity between any two high-din -igh-dmmensionl states within our abstract domains, the decoder is state dependent by construction as outputs predicted state given a state and velocity. We investigate how the decoder performs at boundaries in the 2D) Stretchy Bird environment, as illustrated in Fig. TO We find that the decoder implicitly understands the boundaries of the training data's s underlying manifold. That is, once velocity produces bird state that is unseen in the training distribution, i.e. bird that cannot further shrink 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.
a
b
neck length
training data
11
22
2/3
24
55
26
27
2 N
2N
domain
fixed velocity
i7... N
26
15
time
leg length
Figure Decoder respects boundaries inferred from the extent of the training data. a. Starting from random point in the 2D Stretchy Bird environment, we use the decoder to estimate the state after applying velocity v. We then iteratively apply this same velocity estimate from the previous step, generating series of states obtained by traversing the environment while following this fixed velocity. We find that once the neck and leg length have maximally shrunk to the extent observed in the training data, the decoder arrives at a fixed point. Thus, the decoder performs state-dependent transformations: after reaching an inferred boundary of the training data, further velocities along the same direction do not continue to transform the data along that dimension.
the generated state
B
Ablation studies
B.1
Loss Ablations
next-state prediction loss
next-state prediction loss loop closure loss
next-state prediction loss loop closure loss two auxiliary losses
g.t.
blob bloby y-coords
y-coords
blob y-coords
blob x-coords
blob X-coords
blob X-coards
Figure 11: Loss Ablation studies. Loss ablation studies in the 2D Moving Blobs environment show
that most of velocity estimation learning process comes from our two critical loss terms, 'next-state prediction' and 'loop-closure'. Further loss terms refine the solution space. Grid cell firing rates show that model generated velocities are faithful representations of the true, underlying velocity distribution.
We ablate our various loss terms order to identify which losses are critical for faithful velocity extraction. As seen in Fig. left, if we are interested in extracting some nonlinear low-dimensional
15

---

velocity between states in any abstract domain, the next-step prediction loss is enough (representation error of 0.134). However, since grid cells are doing path integration through linear addition of velocity vectors, constraining our latent space to also be linear function of the true velocity signals essential. Thus, to demand that the estimated velocities be linear functions of the ground truth velocities, we require the loop-closure loss, the addition of which is visualized in Fig. [ middle (representation error of 0.044). If we require further refinement and guaranteed isotropic representations of velocities, we include our isotropy and shorteut losses, the addition of which visualized in Fig. TT right (representation error of 0.02). Examining the grid cell tuning curve within this domain reveals that using only the next-step prediction loss results in non non-hexagonal firing fields. In contrast, incorporating the loop-closure loss results in hexagonal tuning curves, with further refinement achieved by adding auxiliary losses.
For completeness, we prove that the loop-closure loss exacts a strong linearity constraint on its inputs.
Proof: Let the function e 14-+++1 + 0t-t+1 represent the mapping from the true velocity to the model estimated one. The loop-closure loss, applied on trajectories that form loops, applies the following constraint:
Σ
e(vx-1411)= =0
OS1ST-1
0SIST-1
We wish to show that this constraint implies that e must be a linear function over the reals. To show this, first prove homogeneity, i.e., e(ov)) ae(v) for all a R. Then we will prove additivity, i.c. e(U1 + 02) = e(vi) + e(v2).
1. Homogeneity:
Consider a trajectory where T 2, such that up-+1 -21-22 v, and hence uo-11 + 21-2 Loop-closure implies:
e(v) + e(-u) =
(1)
-e(u) =(--))
(2)
This shows that e is an odd function, so e(0) =0
Now, consider trajectory where T 1 for any N:
v for 0 S t S n -
12-1+11 =
mw for 1 = n.
Loop-closure implies:
ne(v) + e(-nv) = = -ne(v) = e(-nv).
Using Eq. we thus obtain
e(nc) ne(u),
Vn Z
(3)
Thus, e is homogeneous over the integers.
To show that e is homogenous over the rationals, consider the loop with T given by u0-,1 = my for E Z, and v1-22 -uw for E - with w mu The loop-closure condition implies to:
e(me) + e(-nw) =
me(v)
ne
C V
(5)
m ((v)
m
(6)
Thus e(av) = ae(()) for a E Q Assuming that is continuous function, and since the rationals are dense the reals, e(arr) = (e(w) for a E R, i.e., is homogeneous over the reals.
2. Additivity:
Assume T' = 3 with Up = V1-2 w. and U2-3 - (v + w). Loop-closure implies:
e(v) + e(w) + + ( =
(7)
e(v) e(w)
= =-e(-(u+op (v+w))
(8)
<(v) e(w)
= =e(u + ww)
(9)
16

---

for all velocity vectors e and w.
Conclusion: Since is both homogeneous and additive, is linear over the reals. Thus, the loop- closure loss constrains the estimated velocities to be linear function of the true velocities, thereby inducing global metric structure in this estimated space.
B.2 Data Ablations
a
b
By
Dy
g.t.
Ux
V/y
Up
Ux
Figure 12: Data Ablation studies. a. Data ablation studies in the 2D Stretchy Bird environment show that training data does not need to consist only of loops. We created a dataset with 50% of trajectories as loops and 50% as independent random walks. The loop-closure loss applies only on loops while other loss terms apply on all trajectories. We report transformation eπrоr of 0.035, comparable to models trained on only closed-loop trajectories. The training paradigm we previously employed asked the model to predict 4t+2 given 0 and 4++1, assuming that Pt+1-0422 0t-+++1- This was done by construction to ensure that image features of 442 were not memorized in the velocity latent space. This paradigm is not necessary. We retrain our network on the 2D Moving Blobs task with trajectories that vary in velocity randomly at each time-step. We predict an estimated 6 from i and 4++ and then predict 4 from 0z-1+11 and iz. We report a transformation error of 0.02, comparable to models trained under our earlier training paradigm.
We conduct two ablation studies our data generation process.
First, we note that while our models trained with data consisting of closed-loop trajectories, this training paradigm is for convenience and not necessity of our framework. If arbitrary random walks were selected instead for training samples, self-intersections would automatically lead to loops within subsequences of the walks. These loops could then be used for the loop-closure loss, with all other losses applied on all loop and non-loop trajectories. As simple proof of concept, we generated dataset of 2D) Stretchy Bird trajectories with 50% of the trajectories as loops and the remainder as independent random walks. While the next-state prediction, isotropy, and shorteut losses applied to all trajectories, the loop-closure loss only applied trajectories that formed loops. Training this modified dataset achieves a transformation error of 0.035 (cf. Fig. 12,,, comparable to the error reported in Table T (corresponding generality, we note that the loop closure loss can also work on trajectories that are "almost loops". ie.,, loss whose coefficient scales with how close trajectory forms closed loop.
models trained only closed-loop trajectories). In greater
Second, ensure that our model velocity latent space did not memorize image features, we estimate velocity 0t-s++1 from i, to it+1, and predict unseen it+2 that is Ut- = 0t---1 away from 1++1 Fig. 12p shows that this training paradigm is not necessity. We train on trajectories from the Moving Blob environment whose velocities vary at each timestep. Training our model to estimate a velocity Ut-++11 from i, to 441 and predicting 4442 from Ut-++11 from i, also results in similar performance. We obtain a transformation error of 0.02, again comparable to the error reported in Table (corresonnng to models trained under our earlier paradigm).
C Experimental Details
Code All experiments were run on single NVIDIA Titan RTX GPU. Each experiment took anywhere from 1-5 hours to train. Code and all experimental runs can be found here: httttt:::/tttiiie com/ http:///www.aa...ee////otttttttr
17

---

Dataset construction. Details about our dataset construction can be found in Table 2 Each sample in the dataset is described as RT which state consisting of two parts where the trajectory length:
I Xframe size is random walk.
X2. which is a negative permutation of x1 such that the velocities between states satisfy So dl) O.
Loss function details. The loss terms, whose specific coefficients are also described in Table 2 are explicitly written out here:
Next-state prediction loss. Given two states i, and Re+1, the next-state prediction loss minimizes the distance between the predicted state and the true state: min | 2212
Loop-Closure Loss. The loop-closure loss ensures that all the predicted velocities in given trajectory sum zero given that the trajectory is loop: min Losest- 0c+++1
min || Edt|l.
Shorteut Loss. The shortcut loss ensures that the decoder g can generalize given state and velocity. For instance, min (1g( f2-13 + 03-+4) -++41l-
Isotropy loss. Finally, the isotropy loss induces isotropy in the inferred velocity space: min var ||ي-- d(i, 4441) < 0 where is a similarity function in the input image space and 0 is some small threshold.
Best-fit linear transform as an error metric. We aim for the estimated velocities & to be linear function of the true velocities V. Correspondingly. we first estimate the best-fit linear transformation T from V' to V via pseudoinverse, T = VQt. Then, we compute our error metric as normalized mean-squared error between the transformed points and the true distribution: 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 '/ via the DBSCAN clustering algorithm (any outlier rejection tool will suffice); however, we report the error evaluated on the entire dataset including outliers.
122vv12 The Vvar(V
Grid cell model. To visualize grid firing fields given inputs from our model, we use an approximation of a continuous attractor model for module of grid cells: we simulate patterned activity on lattice neurons as the sum of three plane waves, resulting in 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.
18

---

Hyperparams
Stretchy Blob (2D)
Stretchy Bird (2D)
Stretchy Bird (3D)
Moving Blobs (2D)
Frequency Modulation
(1D)
Frame size
16 16 81
32 12
32 12
16 16
x 100
Trajectory
81
81
81
length
Velocity dis-
U (0.05, 0.6)
U² (-55,11))))) (-55,1.
20)
U(0.1, 10)
tribution
Max velocity
0.08
1.5
1.5
1.0
0.05
step
Optimizer
Adam
Adam 5e-4
Adam 5e-4
Adam 5e-4
Adam
Learning
5e-4
5c-4
Rate
Epochs
800
800
800
800
1200
Trained
Batch size
256
192
192
256
256
Learnable Pa-
536e3
622e3
622e3
536e3
544e3
rameters
State Predic-
lel
lel
tion Weight
Loop-
1e1
le2
1e2
1e1
1e1
Closure
Weight
Shortcut
lel
lel
1
Estimation
Weight
Isotropy
1e2
le2
1e2
1e2
le2
Weight
Isotropy
le-4
6e-3
6e-3
le-2
le-4
Threshold
Training
Set
800e3
800e3
800e3
800e3
800e3
Size
Testing
Set
200e3
200c3
200e3
200c3
200c3
Size
Table 2: Hyperparameters used for generating the datasets and training the networks.
19

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No] or [NA]. .
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] " "provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, the scope of what our paper tries solve is clearly explained in several sections. We are clear about what challenges our paper solves and what future challenges can be solved result of this work.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Yes, limitations to this paper, including proposals to address them in future directions of research, have been explained in the discussion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
20

---

The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Theoretical motivations for of our SSL loss terms ('loop-closure) has been presented in the Supplementary Section. Assumptions for empirical experiments have been provided in various sections of the paper.
Guidelines:
The answer NA means that the paper does include theoretical results. the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Yes, all necessary information (code, experimental runs, etc.) has been provided reproduce the main results. Please see the Supplementary Material section for hyperparameters related to model training and dataset generation.
Guidelines:
The answer NA means that the paper does not include experiments.
21

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Link the code, Wandb runs, and hyperparameters have been provided in the Supplementary Section.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips.cc/public/eee.c http://ddddeeooeeuubiieeeeee ssi ssionyolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
22

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Hyperparams for running the experiments and generating the data we have used are the Supplementary Materials section. Experiments are fully reproducible.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Results include a table with mean and std of errors across a variety of experi- ments. These experiments, as mentioned in the paper, were across 6 random runs that are seed controlled (and thus fully reproducible).
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma еrror bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Discussion about the computer resources used to run these experiments are discussed in the Supplementary Section.
Guidelines:
23

---

The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https::/neurips.. httpp//rwr...cpoob/c/////aaaaaaa
Answer: [Yes]
Justification: Yes, have taken into account these ethics considerations. Our work follows the NeurIPS Code of Ethics fully.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e-g.. there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work does have any immediate societal impact. We build theoretical model of how grid cells can be used to map abstract cognitive spaces. Our model does not interact with any users nor does collect real-world data.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
24

---

Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: Our models pose no such risks. We do not train large models like LLMs that do pose risks to its users, nor do we collect real data for training our models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes)
Justification: Yes, we implemented baselines described by other papers and have cited them in work properly.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide the code for generating the synthetic tasks/datasets that we trained our models on. This generation process is well documented.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/codeemoodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
25

---

The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper did not crowdsource experiments nor collect nor train on real- world or human data.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper did not crowdsource experiments nor collect or train on real-world or human data.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
26


Reviews:
[{"review_id": "J9tq3idkdT", "reviewer": "Reviewer_sK54", "summary": "This work explores how the brain could theoretically generalize its representations of velocity signals used to map spatial domains into abstract velocity signals that map abstract cognitive domains.\n\nTo do so, proposes a self-supervised ANN algorithm and architecture to learn a low-dimensional latent space that describes abstract velocities (transformation parameters) between sensory inputs resulting from a parameterized low-dimensional generative process (where parameters are points in abstract space). For example, in a Stretchy Blob task, an image of a 2D Gaussian blob is the product of a generative process parameterized by[width, height], and the difference between two images/blobs can be succinctly captured as a 2D velocity corresponding to their difference in [width, height].\n\nThe self-supervised framework makes use of various loss terms that ensure the latent space is a well-formed / geometrically consistent metric space. The experimental evaluation shows that these methods successfully captures relationships across a variety of image and audio domains, better than standard dimensionality reduction baselines. When fed into a synthetic grid cell model, the learned upstream velocity signals result in much more grid-like firing patterns.", "questions": "See above.", "limitations": "I didn't find limitations in the Discussion section as indicated in the Checklist. See above for particular questions.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "**Significance:** The work addresses an important problem in reconciling how grid cells could be involved in both spatial and abstract tasks. This is relevant for both the neuroscience community interested in grid cells / integrators, and potentially the machine learning community interested in transfer / generalization.\n\n**Novelty:** The algorithm builds on ideas from SSL regarding prediction and promoting locally linear / geometrically consistent latent spaces. The perspective of abstract velocities is interesting and differently motivated than typical SSL work.\n\n**Technical Quality:** The qualitative results are convincing, with the learned velocities accurately map their abstract domains. The model is also elegantly simple, yet can be applied to a variety of modalities.\n\n**Presentation Clarity:** The writing is excellent, as is the figure design. The font for annotations in the figures (e.g. Figure 4 var/axes, Figure 3 isotropy equation) could be larger, but no major revisions suggested from me.", "weaknesses": "**Technical Quality**\n\n1. **Dimensionality reduction:** Line 74. Is this a dimensionality reduction method? Can a single image be compressed into a low dimensional repr like the baselines that were compared to (e.g. autoencoder, PCA)? It seems to me that only a _pair_ of images could be compressed into a low-dimensional repr of the velocity. If trying to pass a single image twice into the encoder, a well-trained velocity would be 0.\n2. **Preservation of cell-cell relationships:** Line 78, Section 2.6. I didn't follow the logic of this claim of what the model is predicting. Wouldn't showing that cell-cell relationships are preserved require showing that the velocity signals from multiple tasks can be fed to the same grid cell model simultaneously and use the same grid? Only one task is shown. Or is this claim presupposed by the grid cell model? In which case it wouldn't be a prediction.\n3. **Orthogonality of velocities in triplet:** Line 163. The model computes velocity v12 from (i1, i2), then regresses g(i2, v12) to i3. But the actual velocity v23 between i2 and i3 could be quite different from v12, and in the worst case it is orthogonal. So would this introduce a systematic error? Or maybe its an assumption that velocity doesn't change too quickly along a path, in which case this assumption should probably be made explicit.\n4. **Missing limitations:** I didn't find limitations in the Discussion section as indicated in the Checklist. Some limitations in particular that I'd like to see addressed are:\n    (a) The biological plausibility of the proposed method, especially if this is meant to describe brain mechanisms for generating abstract velocities.\n    (b) How naturalistic is the data collection process that requires closed-loops of random walks? \n\n**Presentation Clarity**\n\n5. **Missing grid cell model details:** Line 254, Is the grid cell model described in the paper/appendix? I couldn't find it or a reference to another work.\n6. **Missing data generation details:** Line 153, It will be important to have these details on how the synthetic data was generated.\n7. **Network vs circuit model:** Line 62 and elsewhere, I would suggest you rename this a \"neural network\" model, not a \"neural circuit\" model. The proposed SSL architecture is a deep learning style model. Typically, I think it's best practice to use \"circuit\" only if you are willing to provide specific commitments between parts of your model and the underlying neural circuit mechanisms in the nervous system. Meanwhile, \"network\" is less strict.\n8. **Missing section number:** Line 207, Experimental Results should be differently numbered than Methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n> **Figure 4 var/axes and Figure 3 isotropy equation**\n\nThank you for the suggestions! We will make the necessary edits to the figure. \n\n> **Dimensionality reduction**\n\nThank 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.\n\n> **Preservation of cell-cell relationships**\n\nThe 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. \nThe 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. \n\n\n> **Orthogonality of velocities in triplet**\n\nAs 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.\n\n\n> **Missing limitations: the biological plausibility of the proposed method**\n\nThis 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. \n\nIn 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.", "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": "", "extra_scores": {}}, {"review_id": "33HMHjlTDg", "reviewer": "Reviewer_SMXS", "summary": "This work proposes a novel neural circuit model that explains how grid cells in the medial entorhinal cortex can map both spatial and abstract cognitive spaces. The model extracts low-dimensional velocity signals from high-dimensional abstract domains using self-supervised learning, leveraging grid cells’ spatial velocity integration capabilities. This allows the creation of efficient cognitive maps across different domains.", "questions": "Please refer to weaknesses.", "limitations": "NA", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1.\tAlthough the idea of using grid cells as a scaffold for representing cognitive maps is not new, the ability to apply continuous actions (‘velocity’ in this paper) and learn from the data is indeed a highlight.\n2.\tThe paper conducts comprehensive ablation studies to verify the rationality of its objective function.\n3.\tThe comparison with baselines is also quite thorough.", "weaknesses": "1.  The paper’s starting point is that grid cells can serve as a scaffold for different cognitive domains, as reflected in the title, but the main content focuses on the extraction of velocity. There is little discussion about grid cells, and the description of how the continuous attractor network of grid cells works is insufficient.\n2.\tThe paper compares the performance of various baselines. The experimental results show that the presented model significantly outperforms the baselines in dimensionality reduction, but there is no clear explanation for why this model is superior to the baselines.\n3.\tThe training process relies on specially constructed datasets (loops). It’s obvious that organisms do not need to return to the original point to achieve good representations.\n4.  The authors didn’t discuss the situation when the dimension of velocity is higher than grid cells (2D).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> **There is little discussion about grid cells, and the description of how the continuous attractor network of grid cells works is insufficient.**\n\nThank 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.\nWe 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. \nWe 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.\n\n\n>  **Why this model is superior to the baselines**\n\nThank 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: \n\nTypical 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.\n\n> **Training process relies on specially constructed datasets (loops)**\n\nThank you for raising this excellent point! \nWhile 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.\n\nIn 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. \n\n> **Situation when the dimension of velocity is higher than grid cells (2D)**\n\nThank 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.", "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": "", "extra_scores": {}}, {"review_id": "GclUipA9Ij", "reviewer": "Reviewer_SEuj", "summary": "This paper develops a new dimensionality reduction approach based on velocity in latent space. It is inspired by, and tries to emulate, aspects of entorhinal grid cell activity. The critical ingredient is a \"loop-closure\" constraint that drives the model to build a metric map of the latent space. The results demonstrate advantages of this approach over traditional dimensionality reduction algorithms.", "questions": "- How does the model deal with boundary/obstacle effects? This seems like a situation where the effect of the velocity operator is not state-independent.\n\n- Related to the previous point, if I've understood correctly, the model assumes that state variables live in R^D (it would be helpful if this was made explicit), so there are no boundary conditions. How do the authors think about this in relation to the toroidal topology that appears to constrain grid cell population activity?\n\n- If the claim is that this model emulates how the entorhinal cortex learns velocity-based representations, I wonder how the authors think about the fact that often the training data animals receive will not have that many loops. For example, how many loops were there in the experiments on stretchy birds? What happens if the model doesn't have loops in its training set? What about trajectories that are *almost* loops?", "limitations": "The authors dicuss future directions at the end, but don't directly address limitations. There are no potentially negative societal impacts.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper proposes a novel solution to an interesting and important problem.\n\n- Writing and visualizations are very clear (some minor comments on this below).\n\n- The experimental results are impressive relative to baselines.\n\n- The work will be potentially impactful within neuroscience. I'm less sure about AI (see Weaknesses).\n\n- I appreciated that the authors laid out the predictions of their framework. It would be great to see these carefully tested.", "weaknesses": "- With respect to the brain, I think the model makes some assumptions with questionable plausible (see Questions below), but I would be interested to hear if the authors disagree.\n\n- I didn't feel that the authors made a really compelling case for AI applications, though I understand the potential utility in principle.\n\nMinor:\n\n- p. 3: Technically the state-based SR doesn't assume action inputs, since it's implicitly taken an expectation over actions. It's true that the state-action SR does take actions as input.\n\n- p. 3: \"learns simultaneously learns\" -> \"simultaneously learns\"\n\n- p. 8: \"These inferred velocities\" -> \"With these inferred velocities\"\n\n- p. 8: \"allow reuse grid cells\" -> \"allow reuse of grid cells\"\n\n- Somewhere in the supplement, the authors should completely spell out their loss function and other algorithmic details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> **State-based SR doesn't assume action inputs**\n\nThank 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.\n\n> **Grammatical changes in text**\n\nThank you for suggesting these revisions! We will change them in text.\n\n> **Somewhere in the supplement, the authors should completely spell out their loss function and other algorithmic details**\n\nThank 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:\n\nGiven 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$.\n\nThe 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 ||$.\n\nThe 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} ||$. \n\nFinally, 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. \n\nAll 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 $\\textbf{sK54}$.\n\n\n> **Boundary/obstacle effects**\n\nThis 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. \n\n> **Model assumes that state variables live in R^D** \n\nThank 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.\n\n> **Concerns on loops**\n\nThank you for raising these excellent points! \nWhile 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.\n\nIn 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. \n\n> **Lack of limitations**\n\nWe 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 $\\textbf{sK54}$.", "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": "", "extra_scores": {}}, {"review_id": "0PWmZ4ISkt", "reviewer": "Reviewer_AXig", "summary": "This work concerns stimuli that come from an underlying low-dimensional latent space. It studies a path-integrating problem on this space, in which a network is shown two images, has to infer a displacement signal between these images, and use it to traverse the space in a few different ways specified by a series of losses. It uses this as a model of velocity extraction for abstract cognitive maps in neuroscience.", "questions": "My questions are largely included in the weaknesses section. Further more details of the exact form of the losses would be nice.", "limitations": "Limitations were not very clearly discussed.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "I agree with the choice of problem. The brain must be extracting abstract velocity signals when it is traversing structured spaces. This is therefore an interesting and relevant problem.\n\nFurther, the stated network and losses do, indeed, nicely solve the problem.\n\nFinally, the exposition is, for the most part, clear and compelling. The figures were very nice.", "weaknesses": "First, I think the framing with respect to previous work is needlessly confrontational, leading to some silly claims. Currently it is framed as if the proposed approach is a fundamentally different way of constructing a cognitive map, I think this is wrong, rather it seems like a useful addition that could coherently be included in existing models, making them more powerful. Relative to CSCG, TEM and models like them, I think this work is simply answering a different part of the cognitive mapping question. As the authors correctly point out, these models require a velocity signal (though that signal has no semantics). They then learn the meaning of the velocity signal (for example that north is the opposite of south) and use that to create a representation that is claimed to match neural data. In contrast, the author's model is a model of velocity extraction, of turning pairs of points into a velocity signal. You could well imagine a combined model that extracted velocities according to the author's scheme, and used them to drive a CSCG/TEM like model (i.e. the function $g(i_1, v)$ would be CSCG/TEM). This would ideally have the cross-environment generalisation of CSCG/TEM (which the author's model doesn't have), and the cross-velocity-modality generalisation of the current model. I think SR is not really a model of cognitive mapping, so I'm not too worried about the comparisons there. \n\nRelated to this, the claim that the velocity signals are learned independently of the stimuli also seems misleading (figure 1g). The same network, f, cannot transfer across sets of stimuli, even if they share the same underlying structure, unlike CSCG/TEM. Have I misinterpreted, should I take this to mean the same velocity is extracted independent of where in the space that velocity begins?\n\nThe silly claims that I think this leads to are the claims of novel predictions. Everyone and their mother would have predicted, from the moment the stretchy birds experiment worked, that the cellular basis of the grid signal would be modules of grid cells behaving in the same way they do in physical space. 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! Every model of grid cells, of which there are likely hundreds, is ambiguous to whether the variable being encoded is physical or abstract space!\n\nA further confusion for me was the requirement for the latent code to be a linear function of the true velocities. Why is it so important? (As is implicitly assumed in the loss calculated in the table, though I couldn't find the loss described in detail anywhere) I see that the shortcut estimation loss and the isotropy loss build this linearity assumption into the velocities (though it seems from the ablation studies that it is actually not necessary to make the representation mostly linear), but I wonder whether it is not more impressive the fewer losses you use? It adds a lot of precision and clarity to be able to say that just the next-state prediction loss is enough to get a nonlinear encoding of the velocity, and from the ablation study, then adding loop closure gets you a linear representation.\n\nFinally, I thought both the comparison to other techniques, and the grid cell results were slightly gimmicky. That's fine, but I don't give them a huge amount of weight as a result. Of course, if you have a perfect velocity code and feed it to a hexagonal grid cell model it will look like a hexagon. Further, why should I expect other dimensionality reduction techniques to learn the velocity? It seems like that is an assumption about the dataset (that all pairs of points separated by the same velocity signal correspond to the most meaningful axes on which to describe differences between). 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 (lines 74-77). In reality, it seems like your loss is looking for a specific type of dimensionality reduction that your model is built for and the others aren't. Its still a non-trivial result, and don't get me wrong, using this technique for dimensionality reduction sounds interesting, but I think the framing is currently way off.\n\nTwo smaller points:\n\nThe stated equivalent between velocities summing to 0 and inputs being the same assumes no aliasing (i.e. no repeated stimuli at different positions) (bottom fig 2c), in reality the two are not equivalent (velocities summing to 0 implies $i_1 = i_N$, but not the other way around). It should either be relaxed to a one directional implication, or the assumption of no aliasing should be stated.\n\nI think the velocities in the latent space were literally added (like vector addition) to compute the shortcut loss, is that true? If so, it should be highlighted that this limits the spaces in which this can be used to ones in which velocities commute. For example, unless I've misunderstood, you cannot then use this for spheres, this should be mentioned as a limitation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n\n> **Relative to CSCG, TEM and models like them, I think this work is simply answering a different part of the cognitive mapping question**\n\nThank you for highlighting these points and for your careful review! \nAs 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.\n\n\n> **Claim that the velocity signals are learned independently of the stimuli also seems misleading (figure 1g)**\n\nWe 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.\n\n\n> **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!**\n\nThank you for the opportunity to clarify! \n\nOur 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.\n\n\n> **The requirement for the latent code to be a linear function of the true velocities**\n\nSince 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.\n\n\n> **I couldn't find the loss [calculated in the table] described in detail anywhere**\n\nWe 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.\nAs 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.", "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": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their careful and detailed review of our paper and their overwhelmingly positive feedback. In this general response, we address some common themes across the reviews. We will additionally provide detailed answers to each reviewer in subsequent comments. To go with this rebuttal, we provide a 1-page PDF with additional figures (labeled Fig. R1-R4).

All reviewers highlighted the importance and relevance of the problem, as well as the novelty of the proposed method. Reviewers also noted that our work was potentially impactful to both neuroscience and machine learning communities, found the model design to be simple yet elegant, the experiments and ablations to be thorough, and the writing and figure design to be excellent.

Two common critiques emerged across reviews: whether our training data consisting only of loops was justified, and whether our model was truly a/should be compared with dimensionality reduction method(s). 

To address these issues and other questions raised by reviewers, we run new analyses in this rebuttal that we include in the attached PDF. Specifically:

1. To show that training purely on loops is not necessary, we find relatively unchanged results from our model when trained on data consisting of both random walks (non-loops) and random loops ($\textbf{Fig. R1}$). We provide more details in our response to $\textbf{sK54, SMXS, SEuj}$.
2. To further demonstrate that our model can truly be considered a dimensionality reduction method, we show the low dimensional representations of an ensemble of states from our 2D Moving Blobs environments. The low-dimensional representations from our model occupy a 2D subspace, reflecting the true data structure, unlike a PCA embedding that we show for comparison ($\textbf{Fig. R3}$). We provide more details in our response to $\textbf{sK54, SMXS, AXig}$. 
3. To show how our model deals with boundary effects, we analyzed a trained model on a 2D Stretchy Bird task. Interestingly, the model’s encoder can continue to generate state-independent velocities, but the model’s decoder is state dependent and does not produce state predictions that lie beyond the implicit boundaries imposed by the training data ($\textbf{Fig. R4}$). We provide more details in our response to $\textbf{SEuj}$.

We sincerely thank the reviewers for their extensive feedback and believe that our manuscript has significantly improved as a result. We hope you agree that our proposed revisions and new analyses have enhanced our work and that our research can lead to a positive impact on both the neuroscience and machine learning communities.

Meta-review: The present study considers a self-supervised learning paradigm to extend grid cells from the spatial domain to abstract domains. This is an important extension because we know grid cells need to represent abstract knowledge, while the spatial location is only a convenient metric that can be easily measured in lab experiments. Overall, the paper is conceptually novel and technically solid, and thus I recommend to be accepted. Please integrate reviewers' suggestions in the revised manuscript.
Decision: Accept (poster)