SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Conformal Isometry Loss Consider a velocity  $v_{t}$  that shifts the neural code from  $g_{t}$  to  $g_{t+1}$ . A neural code is conformally isometric if the distance between the neural codes remains constant if  $v_{t}$  changes direction (but not magnitude). Motivated by [22, 67], we propose a novel and simpler conformal isometry loss that minimizes the variance in ratios of changes in code space to changes in physical space for small velocities:
$$
\mathcal {L} _ {\text {C o n I s o}} \stackrel {\text {d e f}} {=} \mathbb {V} \left[ \left\{\frac {| | \boldsymbol {g} _ {t} - \boldsymbol {g} _ {t - 1} | |}{| | \boldsymbol {v} _ {t} | |} \right\} _ {t: 0 <   | | \boldsymbol {v} _ {t} | | <   \sigma_ {x}} \right], \tag {7}
$$
where  $\mathbb{V}$  denotes the variance. Intuitively,  $\mathcal{L}_{\text{ConIso}}$  asks that neural representations change by some constant amount for the same magnitude of velocity, regardless of the velocity's direction. Our overall self-supervised loss is then a weighted combination of these four losses:
$$
\mathcal {L} \stackrel {\text {d e f}} {=} \lambda_ {S e p} \mathcal {L} _ {S e p} + \lambda_ {I n v} \mathcal {L} _ {I n v} + \lambda_ {C a p} \mathcal {L} _ {C a p} + \lambda_ {C o n I s o} \mathcal {L} _ {C o n I s o} \tag {8}
$$
We term this framework Separation-Invariance-Capacity (SIC). This SIC framework is naturally suited to the setting where neural representations are generated by a deep recurrent network which receives velocity inputs. Unlike previous (supervised) approaches, our SIC framework requires no information about absolute spatial position, only relative spatial separation, and does not require any tuned supervised target, thus mitigating the shortcomings of previous works identified recently [48].
6.1 Emergence of multiple modules of periodic cells
Under our SIC SSL framework, we find that trained networks learn neurons with spatially multi-periodic responses. In each run, most neurons in the network exhibit periodic tuning. Moreover, the network clearly exhibits a few discrete periods. We used grid search over hyperparameters with roughly  $1/2$  an order of magnitude variation for key hyperparameters  $(\sigma_x, \sigma_g, \lambda_{Sep}, \lambda_{Inv}, \lambda_{Cap})$  and found that multiple modules robustly emerge across this range. Computational limitations prevented us from sweeping more broadly. Example units from a run are shown in Fig. 4. Below, we characterize different facets of the results.
6.2 Generalization to input trajectory statistics
A drawback in previous (supervised) deep learning grid cell works was poor generalization when the networks are evaluated in arenas larger than their training arenas. Grid cells in our networks
Figure 5: Grid tuning generalizes to different velocity statistics and to environments much larger than the training environment. (a) [left] A kernel density estimate plot of pairwise distances in neural space for randomly chosen pairs of points in real space showing rapid decorrelation over a length scale that is  $\approx \sigma_x$ . A cumulative CDF of the pairwise distance distribution over 1 training batch (black), showing that the network generalizes to environment sizes much larger than it has seen during training. [right] A plot of pairwise neural space distance as a function of both spatial separation and temporal separation. Red vertical line at 0 spatial separation denotes perfect path integration. (b) Training trajectory statistics (left) and evaluation trajectory statistics (right) (c-e) Ratemaps for trajectories in boxes of sizes  $2m$ ,  $3m$ ,  $4m$ .
generalize to arenas many times larger than the distances explored during training (Fig. 5c-e) (training trajectories extend  $< 1.2\mathrm{m}$  from the starting position (Fig. 5a)). Secondly, the training trajectories were drawn i.i.d. from a uniform velocity distribution, Fig. 5b (left), while evaluation trajectories were constructed to be smoother and avoid the walls of a confined arena (Fig. 5b, left). As found in experiments, the periodicities of the emergent grid cells remained the same, regardless of the size of the arena (Fig 5c-e).
6.3 State space analysis of trained RNNs
We turn to a more-detailed characterization of the grid cells in each module and their cell-cell relationships/population states. For a big enough set of cells for comparison with experimental results, we needed sufficiently many co-modular cells and so used the network corresponding to Fig. 7a, in which most of the units formed grids of the same period.
We found that cells with similar periods were tightly clustered to have essentially identical periods and orientations (Fig. 6a). Moreover, this population tiled the space of all phases (the Fourier PSD contains a hexagonal structure with wavevectors  $k_{1}, k_{2}, k_{3}$ ; we defined three non-orthogonal phases satisfying  $\sum_{a=1,2,3} \phi_{i}^{a} = 0 \mod 2\pi$  for each unit from its ratemap  $r(\pmb{x})$ :  $\phi_{i}^{a} = \text{phase}\left[\int d\pmb{x} e^{i\pmb{k}_{a} \cdot \pmb{x}} r_{i}(\pmb{x})\right]$ ), exhibiting a uniform phase distribution (Fig. 6b).
a mean fourier power
C toroidal state space
b phase coverage
d
Figure 6: Analysis of low-dimensional structure in one of the grid cell modules: (a) (left) Mean Fourier power spectral density of all 128 units in the module, showing hexagonally arranged peaks. (right) Histogram of cell periods and orientations. (b) Distribution of phases with phases  $\phi_1, \phi_2, \phi_3$  calculated by evaluating the Fourier transform of each ratemap along each lattice vector. (c) Cells from one module exhibit a toroidal manifold (visualization by non-linear dimensionality reduction (App D)). (d) Projection of the state space manifold onto the three principal directions defined by vectors  $k_1, k_2, k_3$  shows 3 rings.
In theoretical models and real data, grid cells exhibit low-dimensional dynamics and the states of all cells in a module lie on a 2D torus. We used spectral embedding [47] to find that the emergent grid cell population states also fell on a toroidal manifold (Fig. 6a). Following [55], we projected the activity manifold onto 3 axes defined by each set of phases calculated for each unit. This revealed the existence of 3 rings, Fig. 6d, confirming that the attractor manifold has a 2d twisted torus topology.
Ablation experiments We also ablated each of the three loss terms, the data augmentation, i.e., trajectory permutations and the hyperparameter  $\sigma_{g}$  to reveal when multiple modules of grid cells changes to a single module and then to a place-cell like code (Fig 7a-f).
7 Discussion & Future Directions
In this work, we extracted insights from previous first principles coding theoretic properties of grid cells and mechanistic models that produce grid-like firing. We then proposed a novel set of self-supervised loss functions that with the architectural constraint of a recurrent neural network can produce multiple modules of grid cells when trained on a spatial navigation task. In doing so, we have highlighted the distinction between separation and capacity for a neural representation and shown that optimizing for capacity leads to efficient representations that are better for generalization to larger amounts of test data.
a capacity loss ablation
b capacity loss ablation  $+\mathbb{O}_{\mathrm{g}} = 0.1$
C capacity loss ablation  $+O_{g} = 0.2$
d trajectory permutations ablation
e path invariance loss ablation
f separation loss ablation
Figure 7: Ablation experiments: (a) Removal of capacity loss leads to a single module of grid cells. (b-c) Reducing  $\sigma_{g}$  makes the representation place cell-like, with size of place fields decreasing with increasing  $\sigma_{g}$ . (d-f) Ablating other key ingredients leads to loss of spatial tuning.
Looking to the future, we see four broad directions, two in neuroscience and two in machine learning. Firstly, we believe significantly more work is necessary to understand these trained recurrent networks and connect them back to grid cells. Key questions include: What components of the data, augmentations, losses and architecture control the number of modules, the distribution of the module periods, the period ratios, the structure (square or hexagonal) of the lattices? What do the state spaces of these networks look like? Can we reverse-engineer the learnt mechanisms, using mechanistic analyses [50, 30, 37]? Another interesting question is what underlying connectivity leads to discrete modular responses [29]. How do the learnt representations change when exposed to experimental manipulations, e.g., environmental deformations or rewards [33, 34, 35, 28, 8, 9, 7]? What new experiments can these networks suggest in biological agents? Secondly, we see many opportunities in neuroscience broadly. The biological brain cares about high capacity representations, and grid cells have been implicated in many diverse cognitive and non-spatial tasks [15, 1, 6, 45, 63, 44]. How might our grid cell recurrent networks, and more generally, SSL principles be applied to drive computational neuroscience forward?
Thirdly, we are excited to push our SIC framework further towards machine learning. We envision empirically exploring whether our SIC framework can improve performance in other popular SSL domains such as vision and audition, as well as mathematically characterizing the relationship between SIC and other SSL frameworks. Fourthly, one puzzling aspect of our capacity loss is that it can be viewed as minimizing the entropy of the neural representations in representation space, which flies against common information theoretic intuition arguing for maximizing the entropy of the representations [66, 43, 62, 36, 53, 52]. In some sense, the capacity loss is akin to a "dual" minimum description length problem [25]: instead of maximizing compression by having a fixed amount of unique data and trying to minimize the code words' length, here we maximize compression by having fixed length code words and trying to maximize the amount of unique data. How can this apparent contradiction be resolved?
IRF is supported by the Simons Foundation through the Simons Collaboration on the Global Brain, the ONR, the Howard Hughes, Medical Institute through the Faculty Scholars Program and the K. Lisa Yang ICoN Center. SK acknowledges support from NSF grants No. 2046795, 1934986, 2205329, and NIH 1R01MH116226-01A, NIFA 2020-67021-32799, the Alfred P. Sloan Foundation, and Google Inc. RS is supported by a Stanford Data Science Scholarship. MK is supported by a MathWorks Science Fellowship and the MIT Department of Physics.
Architecture and training data + augmentations
Our code was implemented in PyTorch [46] and PyTorch Lightning [20]. Hyperparameters for our experiments are listed in Table 1. Our code will be made publicly available upon publication.
Hyperparameters	Values
Batch size	130
Trajectory length	60
Velocity sampling distribution	vt ∈ R2~i.i.d. Uniform(-0.15, 0.15) meters
RNN nonlinearity	Norm(ReLU(·))
Number of RNN units	128
Number of MLP layers	3
Spatial length scale σx	0.05 meters
Neural length scale σg	0.4
Separation loss coefficient λSep	1.0
Invariance loss coefficient λInv	0.1
Capacity loss coefficient λCap	0.5
Optimizer	AdamW [38]
Optimizer scheduler	Reduce Learning Rate on Plateau
Learning rate	2e-5
Gradient clip value	0.1
Weight decay	None
Accumulate gradient batches	2
Number of gradient descent steps	2e6
Table 1: Hyperparameters used for training the networks.
Figure 8: A trained recurrent neural network learns multiple modules of grid cells. (a) The architecture of the recurrent neural network used. Inputs are 2D Cartesian velocities  $\mathbf{v}_t \in \mathbb{R}^2$  and the non-linearity is Norm-ReLU. No positional readout exists. (b) For each batch, input trajectory velocities are drawn i.i.d. from a uniform distribution, then randomly permuted to create a batch; 3 trajectories shown.
b
Figure 9: All 128 ratemaps evaluated on trajectories inside a  $2\mathrm{m}$  box. (a) Ratemaps from the corresponding to Fig. 4 (b) Ratemaps corresponding to the run in Fig.6
B All Ratemaps
Figure 10: (a) 2 Example cells from 2 modules, with preferred phase 0 and  $\pi$ . (b) Visualizing the state space defined by  $\{\phi^1, \phi^2\}$  as a torus (left) and on a square with periodic boundary conditions (right), which an equivalent construction of a torus.
C Construction of the grid code
To explain the structure of the grid code, we consider idealized tuning curves in 1d.
Each cell  $i$  is defined by its periodicity  $\lambda^{\alpha} \in \mathcal{R}$  and preferred phase  $\phi_i \in S^1$ . All cells in the same module  $\alpha$  have the same periodicity and uniformly tile all allowed phases. For position  $x \in \mathcal{R}$ ,
$$
r _ {i} ^ {\alpha} (x) = R _ {\max } \operatorname {R e L U} \left[ \cos \left(\frac {2 \pi}{\lambda_ {\alpha}} x + \phi_ {i}\right) \right] \tag {9}
$$
The tuning curves corresponding to this module can be seen in Fig. 10a.
For this module, we can define  $\phi^{\alpha}(x) = \frac{2\pi}{\lambda_{\alpha}} x$  modulo  $2\pi$ . Here  $\phi^{\alpha} \in S^{1}$ .
So the firing rate can now be written as
$$
r _ {i} ^ {\alpha} (x) = R _ {\max } \operatorname {R e L U} \left[ \cos \left(\phi^ {\alpha} (x) + \phi_ {i}\right) \right] \tag {10}
$$
All information about the current state of the module is encoded in the single variable  $\phi^{\alpha}$ . Thus the set of phases  $\{\phi^{\alpha}\}_{\alpha} \in S^{1} \times \ldots \times S^{1}$  uniquely define the coding states of the set of grid modules.
For 2 modules, defined by  $\{\phi^1,\phi^2\}$ , these states can be visualized as being on a torus  $S^1\times S^1$ , Fig.10b.
D Nonlinear Dimensionality Reduction
For Fig. 6, we qualitatively followed the methodology used by seminal experimental papers examining the topology of neural representations [12, 24]: we used principal components analysis to 6 dimensions followed by a non-linear dimensionality (in our case, spectral embedding) reduction to 3 dimensions. Similar results are obtained if one uses Isomap [60].