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  fiete
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 18:51:17

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.

Main paper: 
 introduction: 
 Grid cells in the medial entorhinal cortex are 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 a. Grid cells generate hexagonal tuning in spatial navigation experiments (illustration borrowed from [1] ). 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', a 2D Gaussian that can either stretch or shrink along two axes. f. SR, CSCG, and TEM [4] [5] [6] [7] learn transition structures of this cognitive domain by traversing it. These models require one to build a set of representations for these encountered states as well as structures enabling transitions between them. g. Our approach learns a self-consistent movement (velocity) signal that is independent of the states traversed and encodes the global transition structure of the environment in a 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. i. The actual movement signals in this particular example are low-dimensional. These velocities are independent from the states they connect. viewing of naturalistic images, listening to and modifying a sound changing in pitch, navigating a conceptual 'Stretchy Bird' space (images of a bird that stretch or shrink via joystick input), or even traversing an 'odor' space among many others [8-14, 3, 15, 2] (Fig. 1a-d ).How does the brain transfer its ability to represent space to 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 a question of intense interest, with several proposed approaches [4, 5, 7, 6] . These works propose that when the brain learns the transition structure of a cognitive domain (such as the 'Stretchy Blob' space in Fig. 1e ), it simultaneously builds a set of representations for these states as well as structures enabling transitions between them, Fig. 1f .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 a grid phase and then extracting a self-consistent movement signal to measure displacements in the cognitive domain (Fig. 1g ). 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 a new set of continuous stable states for each explored cognitive domain -a process that is 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 prestructured grid cell states to encode transitions between high-dimensional states. However, using such models to map non-spatial, abstract domains (and by extension, explaining how grid cells can be efficiently reused) requires the extraction of faithful representations of velocity in these domains. Thus, the fundamental, prerequisite challenge lies in extracting a general notion of a minimal dimensional velocity signal from high-dimensional, time-varying data from various abstract domains, Fig. 1h . 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-supervised learning paradigm, wherein neither the coordinate system of the space nor the dimensionality of the estimated velocity are known a 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 a learning framework for velocity estimation in 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 generates state-independent, self-consistent low-dimensional velocity estimates by separating the content of abstract domains from their displacements (Fig. 1g, i ).• Self-supervised geometric consistency constraint. We show how consistent metric representations of velocity can be learned through a 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 dimensionality 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 a spatial task, they should remain co-active in a non-spatial task. This prediction critically relies on the same continuous attractorbased grid module performing integration across domains -a capability that can only be 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 [17] [18] [19] [20] [21] [22] [23] . This forms a testable hypothesis for future experimental studies on neural representations in abstract cognitive domains.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 Tolman-Eichenbaum Machine (TEM) [7] learns a map of a '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 a new environment, TEM requires re-learning these affordances along with grid cell-like representations from scratch. Similarly, the successor representation (SR) [4, 5] predicts future states as a weighted sum of expected future occurrences and uses an underlying discrete action space. State-action variants of SR are dependent on these discrete action inputs to construct cognitive maps. Clone-structured cognitive graphs (CSCG) [6] learn hidden Markov models of sensory representations from observational inputs without prior domain structure assumptions but are limited to only discrete domains. By construction, these models are not guaranteed to preserve cellcell 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 velocity-extraction-first perspective offers key benefits. By learning a low-dimensional velocity that captures transitions between external inputs, mapping the environment to reusable grid cell states can be performed simply by a continuous attractor network, instead of something more complex like TEM. More generally, via our 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.Our work thus bridges the tasks of building cognitive maps of abstract domains and dimensionality reduction. Previous dimensionality reduction methods [24] [25] [26] [27] [28] [29] 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 a result, while these methods learn low-dimensional mappings, they often fail to reproduce the underlying metric space of abstract domains. Unlike these methods, our framework learns a 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] [31] [32] [33] [34] [35] . However, these works have not attempted to decipher a minimally low dimensional description of velocity, making their learned representations unsuitable for grid cells to effectively map observed environments.Consider a trajectory whereEQUATIONThis shows that e is an odd function, so e(0) = 0. Now, consider a trajectory where T = n + 1 for any n ∈ N:v t→t+1 = v for 0 ≤ t ≤ n -1, -nv for t = n.Loop-closure implies: ne(v) + e(-nv) = 0 =⇒ -ne(v) = e(-nv).Using Eq. 2, we thus obtain e(nv) = ne(v), ∀n ∈ Z (3) Thus, e is homogeneous over the integers.To show that e is homogenous over the rationals, consider the loop with T = 2 given by v 0→1 = mv for m ∈ Z, and v 1→2 = -nw for n ∈ Z + with w = mv n . The loop-closure condition implies to: e(mv) + e(-nw) = 0 (4)EQUATIONThus e(αv) = αe(v) for α ∈ Q. Assuming that e is a continuous function, and since the rationals are dense in the reals, e(αv) = αe(v) for α ∈ R, i.e., e is homogeneous over the reals.  
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.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. 2a , given two images, i 1 and i 2 , we aim to learn a function f that infers a low-dimensional velocity v representing the transition from i 1 to i 2 (Fig. 2b ).To ensure that v is a useful and faithful metric representation of the transition between two inputs, we demand that it satisfy two properties.First, the estimated transition velocity v can be used to transform a given input, i.e., there exists a function g that uses v to transform i 1 to predict i 2 . In practice, to ensure that f does not merely memorize image features of i 2 , we demand that g predict an unseen i 3 obtained by transforming i 2 using the same transition v. More generally, this ensures that the estimates of displacements, v, in this space must be independent of 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, v, 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 is zero (Fig. 2c ). (Note that this also ensures that the identity transformation should be represented by the zero-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 a calibrated velocity input for grid cells to then themselves use to map out the input abstract domain. These self-consistency constraints on v 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 in 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 in height and/or width.Bird environment where a bird's legs and neck stretch and shrink. This environment is specifically constructed to mimic its experimental counterpart [15] .We construct an encoder-decoder architecture, wherein the encoder f and decoder g are both multilayer perceptrons (MLPs) (Fig. 3a ). The encoder f processes two adjacent inputs i t and i t+1 from any of our procedurally generated trajectories and maps them to a low-dimensional velocity space. It is important to note that the low-dimensional velocity signal is not known a priori. The decoder g then upsamples this latent velocity representation and combines it with the last input to generate the subsequent ît+2 which remains unseen to the model.A single sample in a batch of data for this model comprises a trajectory of statesT = {i t | 1 ≤ t ≤ N } and velocities V = {v t→t+1 | 1 ≤ t ≤ N -1} (individualvelocities abbreviated as v t for ease). The model is trained through triplets of frames of the input trajectory, using i t and i t+1 to predict i t+2 for 1 ≤ t ≤ 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 to 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 in all our experiments.We formalize our two requirements of the extracted low-dimensional velocity signal into two critical loss terms, a next-state prediction loss, and a 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 to our primary constraints. Our loss terms are visualized in Figure 3b and are described in detail in SI Sec. C.• Next-State Prediction Loss. We quantify this based on the difference between the nextstate prediction ît+2 and the ground-truth frame i t+2 . This loss term operates on individual samples of the generated trajectory, and ensures decodability of the estimated velocity v. • Loop-Closure Loss. We quantify this as a norm of the sum of velocities along a closed loop trajectory T , i.e., the model must produce velocity estimates such that v∈ V v = ⃗ 0.The error signal for this loss operates at the scale of the entire generated trajectory. For convenience, we construct all trajectories of our training data as 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 is not 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 g. From i t+1 , we predict future states i t+3 or i t+4 by directly modifying v. Specifically, if v2→3 is inferred from i t+2 to i t+3 and v3→4 is inferred from i t+3 to i t+4 , then ît+4 should be v2→3 + v3→4 away from i t+2 . 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 is applied only near zero velocity. In particular, we minimize the variance of{∥v t→t+1 ∥ | d(i t , i t+1 ) < θ}, where d(i t , i t+1) is a similarity metric in the input image space, and θ is some small threshold. In practice, we use 1cosine similarity as our distance metric d.There is no direct supervision signal regressing the model outputs onto a known distribution of velocities; instead, the velocities are extracted and inferred automatically. Further, we do not a 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 a small L 1 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 underparameterized relative to the training dataset. We conduct ablation studies on our loss terms, which can be found in SI Sec. B.1. In cases where the model'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.Assume T = 3 with v 0→1 = v, v 1→2 =for all velocity vectors v and w.Conclusion: Since e is both homogeneous and additive, e is linear over the reals. Thus, the loopclosure loss constrains the estimated velocities to be a linear function of the true velocities, thereby inducing a global metric structure in this estimated space. 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 a transformation error of 0.035, comparable to models trained on only closed-loop trajectories. b. The training paradigm we previously employed asked the model to predict ît+2 given vt→t+1 and i t+1 , assuming that v t+1→t+2 = v t→t+1 . This was done by construction to ensure that image features of i t+2 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 vt→t+1 from i t and i t+1 , and then predict ît+1 from vt→t+1 and i t . We report a transformation error of 0.02, comparable to models trained under our earlier training paradigm.We conduct two ablation studies on our data generation process.First, we note that while our models are trained with data consisting of closed-loop trajectories, this training paradigm is for convenience and not a 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 a simple proof of concept, we generated a 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 shortcut losses applied to all trajectories, the loop-closure loss only applied to trajectories that formed loops. Training on this modified dataset achieves a transformation error of 0.035 (cf. Fig. 12a ), comparable to the error reported in Table 1 (corresponding to models trained on only closed-loop trajectories). In greater generality, we note that the loop closure loss can also work on trajectories that are "almost loops"i.e., a loss whose coefficient scales with how close a trajectory forms a closed loop.Second, to ensure that our model velocity latent space did not memorize image features, we estimate a velocity v t→t+1 from i t to i t+1 , and predict an unseen i t+2 that is v t→t+2 = v t→t+1 away from i t+1 . Fig. 12b shows that this training paradigm is not a necessity. We train on trajectories from the Moving Blob environment whose velocities vary at each timestep. Training our model to estimate a velocity v t→t+1 from i t to i t+1 and predicting i t+2 from v t→t+1 from i t also results in similar performance. We obtain a transformation error of 0.02, again comparable to the error reported in Table 1 (corresponding to models trained under our earlier paradigm).Code All experiments were run on a single NVIDIA Titan RTX GPU. Each experiment took anywhere from 1-5 hours to train. Code and all experimental runs can be found here: https://github.com/abhi-iyer/velocity_extraction.Dataset construction. Details about our dataset construction can be found in Table 2 . Each sample in the dataset is described as x ∈ R T×frame size , which is a state consisting of two parts where T is the trajectory length:• x 1 ∈ R T 2×frame size , is a random walk. • x 2 , which is a negative permutation of x 1 such that the velocities between states satisfy T 0 v(t) dt = 0. 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 t and i t+1 , the next-state prediction loss minimizes the distance between the predicted state and the true state: min ||i t+2 -ît+2 || 2 . • Loop-Closure Loss. The loop-closure loss ensures that all the predicted velocities in a given trajectory sum to zero given that the trajectory is a loop: min 0≤t≤T -1 vt→t+1 = min || vdt||. • Shortcut Loss. The shortcut loss ensures that the decoder g can generalize given a state and a velocity. For instance, min ||g(i t+2 , v2→3 + v3→4 ) -ît+4 ||. • Isotropy loss. Finally, the isotropy loss induces an isotropy in the inferred velocity space:min var[ ||v t→t+1 || | d(i t , i t+1 ) < θ ],where d is a similarity function in the input image space and θ is some small threshold.Best-fit linear transform as an error metric. We aim for the estimated velocities V to be a linear function of the true velocities V. Correspondingly, we first estimate the best-fit linear transformation T from V to V via a pseudoinverse, T = V V † . Then, we compute our error metric as a normalized mean-squared error between the transformed points and the true distribution: e = ||T V-V||2 N var{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 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.Grid cell model. To visualize grid firing fields given inputs from our model, we use an approximation of a continuous attractor model for a module of grid cells: we simulate patterned activity on a lattice of neurons as the sum of three plane waves, resulting in a hexagonal pattern of activity. Input velocities are used to update the state of the activity on the lattice of neurons through updating the phases of the plane waves, leading to accurate integration of the input velocities. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are 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, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
a 3d variant of stretchy: 
 Bird where the two bird legs can independently transform.For each abstract environment, we compare the model-inferred velocity representations (Fig. 4a ) to the true distribution of velocities. Note that the self-supervised framework does not result in an exact identity mapping between the ground truth velocities and the model outputs -it suffices for the obtained output to be a linear transformation of the ground-truth velocity space. Correspondingly, we 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 a small number of outliers via the DBSCAN clustering algorithm [36] . More details about this error metric can be found in SI Sec. C.In all cases, irrespective of the dimensionality of the input manifold space or the detailed statistics and structure of the environment states, we see that the inferred velocities are faithful metric representations of the ground-truth velocities, quantified in Table 1 .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 a necessity. In Fig. 4b , we set up our framework to have latent dimensionalities larger than the true data manifold dimensionality. In all cases, the model outputs automatically occupy a 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 (cf. Fig. 4b ). 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. Our model estimates low-dimensional velocities between successive high-dimensional states. These velocities can then be integrated to determine a low-dimensional representation for each state. In this sense, it is possible to view our work as a dimensionality reduction method for continuously varying inputs. Traditional dimensionality reduction methods rely on the statistics of distances between points across an ensemble of states. In contrast, our approach finds a structured tangent manifold around each state that captures the low-dimensional transitions to successive states.We can compare standard dimensionality reduction techniques such as PCA, Isomap, UMAP, and deep autoencoders with our method. To do so, we use these techniques to embed the inputs into the known latent dimensionality of the data manifold, then compute an estimated low-dimensional velocity between states by taking the difference of their corresponding low-dimensional representations. Fig. 5 compares our model'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] (a 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 1 ) and produced velocities that were closely aligned with the true velocity distributions. Remarkably, for even onedimensional manifolds embedded in high-dimensional spaces (as in our 1D Frequency Modulation task), existing dimensionality reduction techniques struggle to find a coherent representation for velocity and produce errors that are two orders of magnitude larger than our model.We also examine dimensionality reduction through constructing representations of the original data via integrating the model velocity estimates, Fig. 6 . For a sample trajectory through the highdimensional states of the 2D Moving Blob environment, our models representation collapses onto a two-dimensional plane (consistent with the data, since the states can be minimally described through transitions in 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. 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 a variety of cognitive domains. We predict a 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 the 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 of a 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 a best-fit linear transformation; for simplicity we choose our transformation to map to aligned velocities across environments, leading to orientation-aligned grid tuning curves -this need not hold in experimentally observed tuning 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 on grid representations [8-14, 3, 15, 2] (grid firing fields from baseline models explored in SI Sec. A). Through extracting faithful low-dimensional representations of velocities across abstract domains, the same continuous attractor-based grid modules can be used across tasks. As a result, if two grid cells fire in an overlapping (or non-overlapping) 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 a testable hypothesis for future experimentation, that may be falsified if 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.We also note that while each grid cell module is a two-dimensional toroidal manifold, enabling integration of two-dimensional velocities, the grid cell system consisting of multiple modules can integrate velocities in higher dimensions [37] . Thus, higher dimensional velocities extracted by our model (e.g., 3D Stretchy Bird) do not pose a problem for integration by grid cells.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: Theoretical motivations for one 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 not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • 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 if they appear in the supplemental material, the authors are encouraged to provide a 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. 
a 2d: 
 Moving Blobs environment where a set of Gaussian blobs uniformly translate, moving in and out of the visual field.5. A 1D 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 is provided in SI Sec. C.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-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 non-Euclidean spaces like a sphere. However, non-Euclidean spaces can be represented by embedding them in a higher-dimensional Euclidean space where velocities commute (e.g., a sphere embedded in three-dimensional space) [38, 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 a small number of continuous modes of variability. Moreover, with a small number of "gluing" points, our method allows for building one-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 a 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 -a relevant challenge in robotics. While the encoder is state-independent and infers a generalized notion of velocity between any two high-dimensional states within our abstract domains, the decoder is state dependent by construction as it outputs a 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. 10 . 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. We ablate our various loss terms in order to identify which losses are critical for faithful velocity extraction. As seen in Fig. 11 left, if we are interested in extracting some nonlinear low-dimensional 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 a linear function of the true velocity signals is 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. 11 middle (representation error of 0.044). If we require further refinement and guaranteed isotropic representations of velocities, we include our isotropy and shortcut losses, the addition of which is visualized in Fig. 11 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-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 : v t→t+1 → vt→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:0≤t≤T -1 v t→t+1 = 0 =⇒ 0≤t≤T -1 e(v t→t+1 ) = 0.We wish to show that this constraint implies that e must be a linear function over the reals. To show this, we first prove homogeneity, i.e., e(αv) = αe(v) for all α ∈ R. Then we will prove additivity, i.e., e(v 1 + v 2 ) = e(v 1 ) + e(v 2 ).Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it 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 to reproduce the main results. Please see the Supplementary Material section for hyperparameters related to model training and dataset generation.• The answer NA means that the paper does not include experiments.• 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 a 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 a specific model and empirical evaluation, it may be necessary to either make it 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
open access to data and code: 
 Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: Link to 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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 in the Supplementary Materials section. Experiments are fully reproducible.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a 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, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: Results include a table with mean and std of errors across a variety of experiments. 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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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 or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: Discussion about the computer resources used to run these experiments are discussed in the Supplementary Section.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or 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 in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
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 not have any immediate societal impact. We build a theoretical model of how grid cells can be used to map abstract cognitive spaces. Our model does not interact with any users nor does it collect real-world data.Guidelines:• The answer NA means that there is 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 or 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS 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. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.

To 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].

The 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.
strength_weakness: strength_weakness: **Technical Quality**

1. **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.
2. **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.
3. **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.
4. **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:
    (a) The biological plausibility of the proposed method, especially if this is meant to describe brain mechanisms for generating abstract velocities.
    (b) How naturalistic is the data collection process that requires closed-loops of random walks? 

**Presentation Clarity**

5. **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.
6. **Missing data generation details:** Line 153, It will be important to have these details on how the synthetic data was generated.
7. **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.
8. **Missing section number:** Line 207, Experimental Results should be differently numbered than Methods.
questions: questions: See above.
limitations: limitations: I didn't find limitations in the Discussion section as indicated in the Checklist. See above for particular questions.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2.	The 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.
3.	The 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.
4.  The authors didn’t discuss the situation when the dimension of velocity is higher than grid cells (2D).
questions: questions: Please refer to weaknesses.
limitations: limitations: NA
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.

- I didn't feel that the authors made a really compelling case for AI applications, though I understand the potential utility in principle.

Minor:

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

- p. 3: "learns simultaneously learns" -> "simultaneously learns"

- p. 8: "These inferred velocities" -> "With these inferred velocities"

- p. 8: "allow reuse grid cells" -> "allow reuse of grid cells"

- Somewhere in the supplement, the authors should completely spell out their loss function and other algorithmic details.
questions: 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.

- 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?

- 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: limitations: The authors dicuss future directions at the end, but don't directly address limitations. There are no potentially negative societal impacts.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 

Related 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?

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

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

Finally, 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.

Two smaller points:

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

I 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.
questions: 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: Limitations were not very clearly discussed.

Comments: title: Final Thoughts
Response, Part III
Response, Part II
Response, Part I
Response
score adjustment
Rebuttal, Part 3
Good points
Rebuttal, Part 2
Rebuttal, Part 2
Rebuttal, Part 2
comment: Thank you for your comments. We summarize our discussion here for convenience, please let us know if you disagree. You raised the following key comments/criticisms of our work:


1. **Of course piping velocities into a pre-existing CAN would work.** This criticism is analogous to “if you solve the problem and get good estimates (i.e., extracting these geometrically consistent velocities), using the estimates for computation will work (i.e., pipe them into a CAN and integrate within the abstract domain)”. This criticism skips over the crucial prerequisite challenge of extracting these velocities in the first place (i.e., solving the problem), which we crucially demonstrate in our work.
2. **TEM-like models are potentially better because they predict neural behavior.** Yet, the key prediction that you suggest is trivial to achieve (correlations being preserved across domains) cannot be replicated in these models.
3. **Dimensionality reduction methods should not be able to extract velocity.** However, we showed that since the data lies on a low-dimensional manifold, a dimensionality reduction method should lead to velocities if the low-dimensional structure was identified. Our method identifies this structure and is thus able to extract a consistent representation of velocity.
4. **Our method works specifically on datasets with dynamics.** Yes, we agree and believe this is a strength of our work that sets it apart from other methods. The synthetic tasks we use make it highly relevant to big data applications today where timeseries and video data are increasingly gaining importance.
5. **We require too many losses.** We specifically demonstrated that this is not the case, and that one can add more losses depending on what qualities are desired from the velocity representations. 
6. **We cannot work with non-Euclidean domains.** We show how our model can still work in nonlinear, non-flat manifolds as well. 


In summary, you conceded that our work is well-explained, is thorough, correct, and contains nice science. Further, you stated that the choice of problem was relevant and interesting, and that our exposition was clear and compelling. While we recognize that you appear to remain unconvinced by our arguments, we hope that this summarization can lead to a reconsideration of your position on our work.
I thank the authors for their further extended commentary. In summary, I'm afraid I still remain unconvinced on both counts, and will keep my score.

On the PCA point: there are problems where velocity extraction is key, and your method does well. There are others where there is no dynamical structure, and top PCs do well. There are probably also those with dynamical structure but on which PCA does better. Different problems demand different algorithms. Nonetheless, it's cool that yours does well.

On the preserved cell-cell correlation count: I realise that no other model predicts this. My point was more that I don't think we needed a model to tell us this would work. I agree probably no-one published this, but at the risk of sounding too harsh, I think that's likely because no-one thought they needed to - of course piping velocities into a pre-existing CAN would work like that? Further, there's a separation between a model of learning, and a model of flexible re-use of past learnt things, both CSCG and TEM seem like models of the learning of structure rather than statements that 'of course this structure never gets re-used'. 

Finally, the loss points you make are very interesting.
> **I do find it interesting so many losses are needed to make it perform well, I would have guessed that just next-state prediction and your architectural choices might have done it**

The number of losses needed for our model to “perform well” effectively depends on what criteria are being imposed in defining “well”. 

If we are interested in extracting some nonlinear low-dimensional velocity between states in any abstract domain, then just the next-step prediction loss is sufficient, as seen in Fig. 9. If we additionally demand that the estimated velocities be linear functions of the ground truth velocities, then the loop-closure loss needs to be added. If we then additionally demand isotropic representations of velocities, we need to include our isotropy loss. And, if we require further refinement, with very precise encoding of velocities, we include our shortcut loss. As indicated in our earlier rebuttal, we will certainly include an updated description of our model results to clarify this point. 

As such, we argued earlier that grid cells would require linear representations of velocities and thus posited that next-step prediction and loop-closure were essential losses to consider. We included isotropy and shortcut losses as additional auxiliary losses to help refine the obtained solution, and will make it clearer in the text that these are not essential to any of our key results. We are happy to include an elaboration of this discussion in text, and will also include a proof to show that loop-closure loss exacts a strict linearity constraint on its inputs.


> **But I'm afraid given my above thoughts, despite the paper being well-explained, thorough, correct, and generally nice science, I will keep my score.**

We hope our responses have further established the key interesting aspects of our results and would appreciate a re-evaluation of our score in this context.
> **Prediction of preserved cell-cell correlations across domains is what everyone and their mother (I checked with my mother) would have thought post-stretchy birds experiment…But could it ever have not been done? If you are capable of extracting velocities from pairs of images correctly then of course you can pipe them into a reusable grid cell system in which such cell-cell correlations are preserved. Did we need a model to tell us that?**

Prediction of cell-cell correlations across domains of different modalities may appear as a straight-forward prediction, however, it critically depends on two key ideas: (1) that low-dimensional velocities can indeed be extracted from abstract domains, and (2) that the _same_ continuous-attractor-based grid module can perform integration across domains. The bulk of our work establishes (1), through constructing an SSL framework for this velocity extraction that performs better than (both deep and non-deep) baselines. 

Note that simply the assumption of a single continuous-attractor-based grid module would not be sufficient, since typical continuous attractor models require a low-dimensional velocity input, and one would thus need to first establish that velocities can be extracted across inputs from different modalities. To summarize, in making predictions for grid cell correlations, we assume a continuous attractor network based grid model and provide a computational basis for linking such an attractor to inputs from different modalities through our framework for velocity extraction.

To establish that across domain cell-cell preservation is a nontrivial prediction, and do not follow directly from pre-existing literature, we point towards the other related models of cognitive mapping mentioned by the reviewer: (1) TEM learns grid cell representations from scratch in any given environment. If presented with a new environment, TEM requires re-learning of the affordances and statistics of the new environment, and as a result cells that were correlated within an environment of one modality are not guaranteed to be correlated in a different modality. (2) SR does not produce preserved cell-cell correlation structure across environments within the same modality itself, and (3) CSCG does not provide grid responses to be able to compare cell-cell correlations. 

In the context of these above models, we hope it is clearer that cell-cell correlations being preserved across modalities does not trivially follow from previous modeling work. We are also unaware of such a conclusion being drawn in the literature in the light of previous experimental work, such as the stretchy bird experiments, sound modulation experiments, or any others. We would be grateful if the reviewer could point us towards such a reference that has already made the predictions that we have stated in our work.

> **I think I learnt more from the algorithmic details needed than I did here and much more importantly (b) they predict neural behaviour!**

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.
Approaching the cognitive mapping problem from a velocity-extraction-first perspective, means that the mapping of the environment to grid states can be performed simply by a prestructured continuous attractor network, instead of something more complex like TEM. Thus, in terms of mapping spaces, our model can be thought of as making predictions for neural behavior corresponding to those obtained from the continuous attractor model, which can be distinct from the predictions made by the grid cells in other hippocampal models. If one wishes to examine neural predictions for hippocampal cells rather than grid cells, our method points towards the usage of a hippocampal complex model that can use prestructured attractor dynamics, such as Chandra et al. 2023 bioRxiv 2023.11.28.568960v2.
Thank you for your response to our rebuttal, and for recognizing our work as being well-explained, thorough, correct and containing generally nice science. 


> **I continue to disagree about the importance of the dimensionality reduction comparison to e.g. PCA. Sure, you generate datasets in which differences are a better way to extract out the underlying space than PCA. Your method, aimed at extracting velocities, therefore does better. That is wonderful, but I believe my original comments still hold.**

We do not believe our improved performance in comparison with PCA is simply because of our choice of datasets. Through the idea of velocity extraction as a means to perform dimensionality reduction, we can leverage local transition structures in the original data manifold that sequence agnostic methods (e.g., PCA, UMAP, etc.) cannot use. Thus, our method can utilize the crucial trajectory information in the data to perform more efficient dimensionality reduction. In this context, we do acknowledge that our method assumes the presence of trajectories characterized by continuous variation in inputs; however, we believe this is not a restrictive assumption, as many datasets, including video data, typically fall into this category. An examination of larger, more realistic datasets lies within the scope of future work that we plan to pursue. 

You mention that you believe that your original comments in this context still hold. Our understanding of your original comments were that you had three primary questions/concerns: (1) One may not expect other dimensionality reduction methods to extract a velocity. (2) PCA may still be extracting a useful embedding that happens to not correspond to the velocities that we extract. (3) Our loss terms specifically looked for a specific type of dimensionality reduction that other methods did not. 
In our earlier response, we attempted to respond to each of these concerns: (1) We argued that any dimensionality reduction method should be expected to extract a velocity, since the data lies on a very low-dimensional manifold embedded in the high-dimensional space. (2) Since PCA was unable to find a low-dimensional representation, we argued that it was not a meaningful dimensionality reduction. (3) _Yes_, our loss terms do look for dimensionality reduction based on transitions which other methods are unable to take into account; this allows us to be able to perform more effective dimensionality reduction as evidenced by extraction of a two-dimensional embedding of the data in a dataset where other methods failed to extract low-dimensional representations, even though the original data lived on a two-dimensional manifold.

We hope this response has helped to clarify the salience of our approach and address the concerns you raised in your original comments. If there are any unresolved questions about the importance of our method as a dimensionality reduction technique, we would appreciate further discussion to understand your perspective.
Thank you again for your detailed feedback and suggestions. We hope you agree that our paper has improved as a result of this discussion. Please let us know if there are any additional questions that we may be able to address to consider a re-evaluated score.
Thank you for your comments. I believe my concerns are properly addressed and I would like to maintain my score.
Thank you for your extensive response, I appreciate the assiduousness and engagement.

Broadly, the comments answer my questions, especially about aliasing and spheres.

I continue to disagree about the importance of the dimensionality reduction comparison to e.g. PCA. Sure, you generate datasets in which differences are a better way to extract out the underlying space than PCA. Your method, aimed at extracting velocities, therefore does better. That is wonderful, but I believe my original comments still hold.

Lastly, you got me musing about what it is we do when we build such models of the brain (cue authors' groan). I continue to think that the prediction of preserved cell-cell correlations across domains is what everyone and their mother (I checked with my mother) would have thought post-stretchy birds experiment, perhaps I am too new to the field and missed some 00/10s era debate on this. This model shows that, if you can design a scheme for extracting velocities, such an idea works. 

But, could it ever not have done? If you are capable of extracting velocities from pairs of images correctly then of course you can pipe them into a reusable grid cell system in which such cell-cell correlations are preserved. Did we need a model to tell us that?

The novelty of your work seems to be in the design of a system to extract velocities, the losses etc. that can make it work, and using it to do dimensionality reduction on appropriate datasets. Hypothetically, you could make the same attack at models like TEM/CSCG/SR, that they take a computation, and simply show a model of it working, but (a) I think I learnt more from the algorithmic details needed than I did here and much more importantly (b) they predict neural behaviour!

So it seems the potential neuroscience value I am drawing from this model is about the losses that are required to make it work. I do find it interesting so many losses are needed to make it perform well, I would have guessed that just next-state prediction and your architectural choices might have done it (and who knows, with another setup it still might). But I'm afraid given my above thoughts, despite the paper being well-explained, thorough, correct, and generally nice science, I will keep my score.

If you think there is something that has still not made it through my thick skull please do let me know.
Thank you for the comments. We will certainly improve the discussion on grid cells and our predictions in the paper. We will ensure that the novelty of our work and predictions in the context of grid cells will be made more apparent to people outside the grid cell community. 

Discussion of how our predictions can be falsified is an excellent idea, and we will be sure to include this. To restate from earlier, we predict that if two grid cells are co-active for one task, they remain co-active for all other tasks. Further, if they fire in a non-overlapping way in one domain, they will continue to be non-overlapping in all other domains. This is critically dependent on the following being true: (1) that low-dimensional velocities can indeed be extracted from abstract domains, which we demonstrate through our learning framework, and (2) that the _same_ continuous-attractor-based grid module can perform integration across domains. 

This prediction may thus be falsified if the brain were to use _distinct_, independent grid cell modules to organize information from different modalities. It would also be falsified if the grid cells are not modeled by a pre-structured continuous attractor network which would function independently of the nature of the inputs. We will include and elaborate on the falsifiability of our predictions in the revised text. 

We are certainly interested in constructing more biologically plausible architectures and learning rules that can perform the computation described in our work, but as correctly noted, this would fall outside the scope of our current work. 

Thank you for your score and for supporting our paper's acceptance. We hope that you agree that our paper has improved, given your detailed feedback and suggestions. Please let us know if there are any additional questions that we may be able to address to consider a re-evaluated score.
Thank you for your comments and thank you for your appreciation / advocacy of our paper!
My score was already quite high (8). While I really like this paper, I don't think it makes sense to raise it to 9, which would imply that this is a groundbreaking paper. I feel that is a valuable contribution and should definitely be accepted (I will advocate for it), but it's not going to revolutionize either neuroscience or machine learning. I agree that there are some potentially powerful contributions to AI, but these were not demonstrated in this paper.
We hope that you agree that our paper has improved given your detailed feedback and suggestions! Given that we have addressed the primary concerns raised in the review, we kindly ask you to adjust your score while taking our rebuttal into account.
We hope that you agree that our paper has improved given your detailed feedback and suggestions! Given that we have addressed the primary concerns raised in the review, we kindly ask you to adjust your score while taking our rebuttal into account.
We hope that you agree that our paper has improved given your detailed feedback and suggestions! Please let us know if there is anything else that we can clarify. We kindly ask you to consider revising your score if you agree that the primary concerns raised in the review were addressed.
Thank you for your rebuttal comments, additional results, and manuscript improvements! 

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

Overall, your rebuttal has addressed my feedback, which did not raise major concerns. I am happy to maintain my original high score and support this paper's acceptance.
> **Further more details of the exact form of the losses would be nice**

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

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

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

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

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

All losses have a weight / prefactor, listed in Table 2 of our paper. As mentioned in L200, regardless of the training environment, the relative weighting of the two critical loss terms remains consistent (the loop-closure loss weighted ten times higher than the next state prediction loss). We also plan to discuss more about how each synthetic task was generated, as also recommended by $\textbf{sK54}$.
Thanks for clarifying this! I think these are reasonable points to add to the discussion.
> **Why should I expect other dimensionality reduction techniques to learn the velocity?** 

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


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




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

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


> **You cannot then use this for spheres**

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

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

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

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

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

We will be glad to include these discussion points in the camera ready version of the text.
> **How naturalistic is the data collection process that requires closed-loops of random walks?**

Thank you for raising this question.

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

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


> **Missing grid cell model details**

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

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

> **Missing data generation details**

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

> **Network vs. circuit model; Missing section number**

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