SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Rylan Schaeffer*
Mikail Khona*
mikail@mit.edu
Tzuhsuan Ma
Janelia Research Campus
mat@janelia.hhmi.org
Cristóbal Eyzaguirre
ceyzaguirre@stanford.edu
Ila Rani Fiete
fiete@mit.edu
To solve the spatial problems of mapping, localization and navigation, the mammalian lineage has developed striking spatial representations. One important spatial representation is the Nobel-prize winning grid cells: neurons that represent self-location, a local and aperiodic quantity, with seemingly bizarre non-local and spatially periodic activity patterns of a few discrete periods. Why has the mammalian lineage learnt this peculiar grid representation? Mathematical analysis suggests that this multi-periodic representation has excellent properties as an algebraic code with high capacity and intrinsic error-correction, but to date, synthesis of multi-modular grid cells in deep recurrent neural networks remains absent. In this work, we begin by identifying key insights from four families of approaches to answering the grid cell question: dynamical systems, coding theory, function optimization and supervised deep learning. We then leverage our insights to propose a new approach that elegantly combines the strengths of all four approaches. Our approach is a self-supervised learning (SSL) framework - including data, data augmentations, loss functions and a network architecture - motivated from a normative perspective, with no access to supervised position information. Without making assumptions about internal or readout representations, we show that multiple grid cell modules can emerge in networks trained on our SSL framework and that the networks generalize significantly beyond their training distribution. This work contains insights for neuroscientists interested in the origins of grid cells as well as machine learning researchers interested in novel SSL frameworks.
Spatial information is fundamental to animal survival, particularly in animals that return to a home base. Mammalian brains exhibit a menagerie of specialized representations for remembering and
*Denotes equal contribution and co-first authorship.
Figure 1: Grid Cells in the Mammalian Medial Entorhinal Cortex. (a) A grid cell fires at the vertices of a regular triangular lattice in space, regardless of direction or speed. (b) Examples of two grid cells' firing fields (top) and spatial autocorrelations (bottom). (c) The grid code is multimodular: across the population, grid cells cluster into discrete modules such that all cells within a module share the same spatial periodicity and orientation, and only a small number of spatial periods are expressed. (d) Grid cells within one grid module uniformly cover all phases. Figs b-d from [26, 59].
reasoning about space. Of these, grid cells (whose discovery earned a Nobel prize) are a striking and important example. Each grid cell fires in a spatially periodic pattern, at every vertex of a regular triangular lattice that tiles all two-dimensional environments [26] (Fig. 1ab). This neural code for space appears bizarre since position is a local and aperiodic variable, but the neural representations are non-local and periodic. Why has the mammalian brain learnt this peculiar representation?
Broadly, there have been two directions to answering this question. The first direction uses theoretical analysis to deduce the mathematical properties that the grid code possesses [21, 56, 64, 40, 58]. This direction yields many insights about the coding properties of grid cells, but has not validated these insights by proposing an optimization problem under which grid cells emerge in deep recurrent neural networks. The second direction trains deep recurrent neural networks (RNNs) [16, 5, 54] with the goal of identifying the optimization problem that grid cells solve. However, recent work [48, 51] showed that these models require highly specific, hand-engineered and non-biological readout representations to produce grid cells; removing these representations results in loss of grid cells even though the optimization problem is equally well-solved.
In this paper, we make three contributions. First, we extract critical insights from prior models of the grid code spanning diverse fields including coding theory, dynamical systems, function optimization and supervised deep learning. Second, we use those insights to propose a new self-supervised learning (SSL) framework for spatial representations, including data, data augmentations, loss functions and network architecture. Third, we show that multi-modular grid cells can emerge as optimal solutions to our SSL framework, with no hand-engineered inputs, internals or outputs. We also perform ablations to show under what constraints grid cells do and do not emerge.
By creating a minimal, theoretically-motivated and developmentally plausible learning problem for path-integrating deep recurrent networks through SSL, we close the loop from the discovery of biological grid cells, to the study of their theoretical properties, to their artificial genesis.
Grid cells in the mammalian medial entorhinal cortex Grid cells [26] are neurons found in the mammalian medial entorhinal cortex tuned to represent the animal's spatial location as it traverses 2D space. Each grid cell fires if and only if the animal is spatially located at any vertex of a regular triangular lattice that tiles the explored space, regardless of the speed and direction of movement through space (Fig. 1ab). As a population, grid cells exhibit several striking properties that provide support for a specialized and modular circuit: Grid cells form discrete modules (clusters) such that all cells within a module share the same spatial periodicity and orientation, while different modules express a small number of discretely different spatial periods (Fig. 1c) [59]. Grid cells in the same module (i.e., sharing the same lattice periodicity) exist for all phase shifts [68] (Fig. 1d).
Self-Supervised Learning (SSL) and its applications in Neuroscience SSL is a subset of unsupervised learning in which data is used to construct its own learning target [2, 13, 27, 61, 10, 11, 14, 69, 19]. Broadly, SSL methods belong to one of three families: deep metric learning, self-distillation or canonical correlations analysis; we refer the reader to a recent review [4]. SSL is increasingly gaining popularity as a normative approach in neuroscience since SSL mitigates the need for large-scaled supervised data that biological agents do not possess. SSL has been used in vision to explain the neural representations of ventral visual stream [32, 70], dorsal visual stream [41] and an explanation for the emergence of both [3]. In this paper, we provide the first model of SSL applied to the mammalian brain's spatial representations with an emphasis on sequential data and recurrent networks that to the best of our knowledge is not well explored in either the machine learning or neuroscience literature.
3 Notation and Terminology
We consider temporal sequences of three variables: spatial positions  $\boldsymbol{x}_t \in \mathbb{R}^2$ , spatial velocities  $\boldsymbol{v}_t \in \mathbb{R}^2$  and neural representations on the surface of the positive orthant of the  $N$ -dimensional sphere  $\boldsymbol{g}_t \in \mathbb{S}_{\geq 0}^{N-1} \stackrel{\mathrm{def}}{=} \mathbb{G}$ . Note that any neural representation  $\boldsymbol{g} \in \mathbb{G}$  has unit norm  $||\boldsymbol{g}||_2 = 1$ . We interchangeably refer to representations as states, codes or embeddings. We denote sequences of any of these variables from time  $t = 1$  to  $T$  as  $(\cdot_1, \cdot_2, \dots, \cdot_T)$ , e.g.,  $(\boldsymbol{v}_1, \boldsymbol{v}_2, \dots, \boldsymbol{v}_T)$ .
4 Insights from Grid Cell Modeling Approaches
The question of why the mammalian brain uses a high-dimensional, multiperiodic code to encode a low-dimensional, aperiodic variable has led researchers to theoretically analyze the grid code's properties from multiple mathematical perspectives: coding theory, dynamical systems, function optimization, and supervised deep learning. In this section, we extract key insights that suggest a convergence towards a SSL approach.
4.1 Insights from Coding Theory
Many interrelated coding-theoretic properties of grid cells have been pointed out before: Exponential capacity: The coding states form a space-filling curve (1D) (Fig. 2a) or space-filling manifold (2D) with well-separated coding states in the full coding volume defined by the number of cells. Linearly many distinct grid modules and therefore linearly many neurons generate exponentially many unique coding states for generic choices of periods [21, 56]. The grid code enables high dynamic-range representation (i.e., large value of range divided by resolution) of space using individual parameters that do not vary over large ranges [21] (Fig. 2d-e). Separation: Excluding the closest spatial locations, the coding manifold for spatial positions is interleaved in the coding volume, such that relatively nearby states are mapped to faraway locations (Fig. 2b) and relatively nearby spatial locations are mapped to faraway states (Fig. 2c); this gives the grid code strong error-correcting properties [56]. Decorrelation: Grid cells maximally decorrelate their representations of space beyond a small scale that is roughly half the period of the smallest module, so that all correlations are roughly equally small beyond that scale [21, 56] (Fig. 2f). Equinorm coding states: The grid coding states are equi-norm and lie on the surface of a hypersphere. Ease of update: The grid code permits carry-free updating arithmetic so that individual modules update their states without communicating with each other, meaning updates can be computed in parallel rather than sequentially [21]. Algebraic: The grid code is a faithful algebraic code [21, 56], in that the grid code for composition of displacements is equal to the composition of grid codes for each of the displacements. This property enables the grid code to self-consistently represent Euclidean displacements (Fig. 2g). Whitened representation: All modules (coding registers or phases) update by similar amounts for a given displacement in real space, as a consequence of module periods being similarly sized (in contrast with fixed-base number system representations, where the "periods" for each register are hierarchically organized in big steps as powers of the base, and the representation is not whitened).
These properties are all intrinsic to the grid code, without regard to other types of spatial neurons, e.g. place cells. The lesson is that the grid code likely originates from coding-theoretic properties, not from predicting place cells (c.f., [5, 55]). The challenge before us is identifying which subset of these partially inter-related mathematical coding-theoretic properties of the grid code - all useful for the function of grid cells - forms the minimally sufficient subset for grid cell emergence.
Figure 2: Mammalian grid cells and their coding theoretic properties: Schematics in (a)-(e) depict the coding of 1D rather than 2D space for simplicity, with phases of two grid modules. (a) Continuous locations translate to a continuous coding line that wraps around the torus defined by the two grid module phases. Right: torus cut open to a sheet with boundary conditions. (b) Near-to-far mapping from phases to locations. (c) Near-to-far mapping from locations to phases. (d-e) High capacity and separability: For generic choices of grid module periods, the code is a space-filling curve on the phase torus (d), with a capacity (before curves overlap for finite-thickness curves) that grows exponentially with the number of modules. The coding lines tend to be well-separated for any range of real-space being encoded (2 shown) because of the interleaving property of the code (e). (f) The distance between coding representations  $g(0)$  and  $g(x)$  very quickly increases (around  $\lambda / 2$  where  $\lambda$  is the scale of the smallest module.) (g) Path-independent encoding of displacements.
4.2 Insights from Dynamical Systems
Consider the following general framework for defining a neural code that is dynamically updated on the basis of a velocity signal: An agent starts with some neural representation  $\pmb{g}_0 \in \mathbb{G}$ , then moves along a spatial trajectory given by the sequence of velocities  $(v_1, \dots, v_T)$ . The sequence of neural representations  $(\pmb{g}_1, \dots, \pmb{g}_T)$  is updated based on velocity through the (arbitrary, non-linear and possibly time-dependent) dynamical translational operation  $f: \mathbb{G} \times \mathbb{R}^2 \to \mathbb{G}$ :  $\pmb{g}_t \stackrel{\mathrm{def}}{=} f(\pmb{g}_{t-1}, \pmb{v}_t)$ . In order for the neural representations to be good codes for Euclidean space, what properties should  $f$  and  $\mathbb{G}$  possess? The first property is that the representations should be path invariant:  $\pmb{g}_T$  should be the same regardless of whether the agent took path  $\pmb{v}_1, \dots, \pmb{v}_T$  or path  $\pmb{v}_{\pi(1)}, \dots, \pmb{v}_{\pi(T)}$  for any permutation  $\pi: [T] \to [T]$ . Specifically, we desire that  $\forall \pmb{g}, \pmb{v}_i, \pmb{v}_j, f(f(\pmb{g}, \pmb{v}_i), \pmb{v}_j) = f(f(\pmb{g}, \pmb{v}_j), \pmb{v}_i)$ . Note that path independence automatically implies that the coding states form a continuous attractor [21], because when the velocity  $\pmb{v} = 0$ , the state should remain constant:  $\forall \pmb{g}, f(\pmb{g}, \pmb{v} = 0) = \pmb{g}$ . Further, the grid cell representation is translationally invariant [39, 31]: representations at different positions are related by a shift operator that only depends on the relative separation between the positions, i.e., the correlational structure  $\langle g(\pmb{x}) g(\pmb{x} + \pmb{v}) \rangle_{\pmb{x}} = f(\pmb{v})$ , where  $f$  is a function only of  $\pmb{v}$ . The lessons are thus twofold: the data and loss function(s) should together train the network to learn path invariant representations and continuous attractor dynamics, and the dynamical translation operation of the recurrent neural networks should architecturally depend on the velocity.
4.3 Insights from Spatial Basis Function Optimization
A third approach to explaining grid cells is via optimization of spatial basis functions. The spatial basis function optimization approach is a non-dynamical one that posits a spatial neural code  $\pmb{g}(\pmb{x})$  of
a particular parametric form, e.g.,  $\pmb{g}(\pmb{x}) \stackrel{\mathrm{def}}{=} \pmb{a}_0 + \sum_d \pmb{a}_d \sin(\pmb{k}_d \cdot \pmb{x}) + \sum_d \pmb{b}_d \cos(\pmb{k}_d \cdot \pmb{x})$ . One can then define a loss function and optimize parameters to discover which loss(es) results in grid-like tuning, possibly under additional constraints [39, 18]. This functional form is used because it reflects the constraints of path-integration: if one path-integrates with action-dependent weight matrices, then one must use that functional form [18]. While sensible, these insights have not yet crossed to deep recurrent neural networks. The challenges are that recurrent networks given velocity inputs are not explicit functions of space, and their representations need not be comprised of Fourier basis functions. We seek to bring such insights, which requires discovering how to train recurrent networks that simultaneously learn a consistent representation of space (i.e., learn a continuous attractor), while at the same time, shape the learnt representations to have desirable coding properties.
4.4 Insights from Supervised Deep Recurrent Neural Networks
Deep recurrent neural networks sometimes learn grid-like representations when trained in a supervised manner on spatial target functions [5, 54, 23, 65, 22, 67], and studies regarding the successor representation in reinforcement learning [57, 42] show similar findings. These results share a common origin based on the eigendecomposition of the spatial supervised targets (putatively: place cells) which results in periodic vectors and was earlier noted in shallow nonlinear autoencoders [17]. This understanding was used by [54] to hand-design better supervised targets.
While high profile, this line of work has fundamental shortcomings. Firstly, because the networks learn to regurgitate whatever the researcher encodes in the spatial target functions, the learned representations are set by the researcher, not by the task or fundamental constraints [49]. Secondly, by assuming the agent already has access to its spatial position (via the supervised targets), this approach bypasses the central question of how agents learn how to track their own spatial position. Thirdly, previous supervised path-integrating recurrent neural networks either do not learn multiple grid modules or have them inserted via implementation choices, oftentimes due to statistics in the supervised targets [54, 49]. Fourthly, supervised networks fare poorly once the spatial trajectories extend outside the training arena box (i.e., where supervised targets have not been defined).
The lesson we extract from this is to remove supervised targets altogether, and instead directly optimize the neural representations via self-supervised learning. Supervised networks pick up whatever statistics exist in the supervised targets, and struggle to generalize beyond their supervised targets, meaning that if we wish to avoid putting grid cells in by hand or if we wish to obtain networks that generalize, we must remove the supervised targets altogether. Previous work lightly explored this direction by introducing unsupervised loss terms, but kept the supervised targets [22, 67]. In this work, we fully commit to this lesson and take a purely SSL approach.
5 Self-Supervised Learning of Representations for Space
We deduce that learning useful representations of space is a SSL problem requiring an appropriate SSL framework (comprising training data, data augmentations, loss functions and architecture). We hypothesize that three loss functions, extracted from the coding, path integration and optimization perspectives synthesized above, when combined with particular data and recurrent network architecture, will generate multiple modules of grid cells. See App. A for experimental details.
5.1 Data Distribution and Data Augmentations
For each gradient step, we sample a sequence of  $T$  velocities  $(\pmb{v}_1,\pmb{v}_2,\dots,\pmb{v}_T)$ , with  $\pmb{v}_t\sim_{i.i.d.}p(\pmb{v})$ , then construct a batch by applying  $B$  randomly sampled permutations  $\{\pi_b\}_{b = 1}^B$ ,  $\pi_b:[T]\to [T]$  to the sequence of velocities to obtain  $B$  permuted velocity trajectories; doing so ensures many intersections between the trajectories exist in each batch (SI Fig. 8b). We feed each trajectory through a recurrent network (Sec. 5.2) with shared initial state  $\pmb{g}_0$ , producing a sequence of neural representations  $(\pmb{g}_{\pi_b(1)},\pmb{g}_{\pi_b(2)},\dots,\pmb{g}_{\pi_b(T)})$ . This yields a batch of training data for a single gradient step.
$$
\mathcal {D} _ {\text {g r a d i e n t s t e p}} = \left\{\left(\boldsymbol {v} _ {\pi_ {b} (1)}, \boldsymbol {v} _ {\pi_ {b} (2)}, \dots , \boldsymbol {v} _ {\pi_ {b} (T)}\right), \left(\boldsymbol {g} _ {\pi_ {b} (1)}, \boldsymbol {g} _ {\pi_ {b} (2)}, \dots , \boldsymbol {g} _ {\pi_ {b} (T)}\right) \right\} _ {b = 1} ^ {B} \tag {1}
$$
Figure 3: A self-supervised framework for learning representations of space. (a) Separation incentivizes neural representations corresponding to different spatial locations to be well-separated. (b) (Path) Invariance incentivizes neural representations corresponding to the same spatial location to match. (c) The architecture of the recurrent neural network used. Inputs are 2D Cartesian velocities  $\boldsymbol{v}_t \in \mathbb{R}^2$  and the non-linearity is Norm-ReLU. No positional readout exists. (d) For each batch, input trajectory velocities are drawn i.i.d. from a uniform distribution, then randomly permuted to create a batch; three trajectories shown.
5.2 Recurrent Deep Neural Network Architecture
Motivated by our insights from dynamical systems (Sec. 4.2), we use an RNN (SI Fig. 8a) with interaction matrix given by:
$$
W \left(\boldsymbol {v} _ {t}\right) \stackrel {\text {d e f}} {=} M L P \left(\boldsymbol {v} _ {t}\right) \tag {2}
$$
where  $\pmb{v}_t \in \mathbb{R}^2$  is the velocity input to the recurrent network and  $W: \mathbb{R}^2 \to \mathbb{R}^{N \times N}$  for  $N$  hidden RNN units. The dynamics of the recurrent network are defined as:
$$
\boldsymbol {g} _ {t} \stackrel {\text {d e f}} {=} \sigma \left(W \left(\boldsymbol {v} _ {t}\right) \boldsymbol {g} _ {t - 1}\right) \tag {3}
$$
The neural representations  $g_{t}$  are enforced to be non-negative and unit-norm via the non-linearity  $\sigma(\cdot) \stackrel{\mathrm{def}}{=} \text{Norm}(ReLU(\cdot)) = ReLU(\cdot) / ||ReLU(\cdot)||$ . Constraining representations to be constant norm is a common technique in SSL to avoid representation collapse [4] and also was used in previous grid cell modelling work [67]. By choosing this particular architecture, we ensure that representations at different points in space are related only by a function of velocity, thereby ensuring that the transformation applied to the state is the same so long as the velocity is the same.
5.3 Loss Functions
Separation Loss Neural representations corresponding to different spatial locations, regardless of trajectory, should be well-separated (Fig. 3a). Here, "different spatial locations" is with respect to spatial length scale  $\sigma_{x} > 0$  and "well-separated" is with respect to neural length scale  $\sigma_{g} > 0$ .
$$
\mathcal {L} _ {S e p} \stackrel {\text {d e f}} {=} \sum_ {\substack {\forall \pi_ {b}, \pi_ {b ^ {\prime}}, t, t ^ {\prime}: \\ \left\| \boldsymbol {x} _ {\pi_ {b ^ {\prime}}} (t) - \boldsymbol {x} _ {\pi_ {b} \left(t ^ {\prime}\right)} \right\| _ {2} > \sigma_ {x}}} \exp \left(- \frac {\left\| \boldsymbol {g} _ {\pi_ {b ^ {\prime}}} (t) - \boldsymbol {g} _ {\pi_ {b}} \left(t ^ {\prime}\right) \right\| _ {2} ^ {2}}{2 \sigma_ {g} ^ {2}}\right) \tag{4}
$$
**Path Invariance Loss** Neural representations corresponding to the same spatial location, regardless of trajectory, should be identical (Fig 3b), i.e., neural representations should be invariant to the particular path the agent takes. If two trajectories are permuted versions of one another such that they begin at the same position and end at the same position (measured with respect to a given spatial length scale  $\sigma_x > 0$ ), the neural representations after both trajectories should be attracted together.
$$
\mathcal {L} _ {I n v} \stackrel {\text {d e f}} {=} \sum_ {\substack {\forall \pi_ {b}, \pi_ {b ^ {\prime}}, t, t ^ {\prime}: \\ \left\| \boldsymbol {x} _ {\pi_ {b ^ {\prime}} (t)} - \boldsymbol {x} _ {\pi_ {b} \left(t ^ {\prime}\right)} \right\| _ {2} <   \sigma_ {x}}} \left| \left| \boldsymbol {g} _ {\pi_ {b} (t)} - \boldsymbol {g} _ {\pi_ {b ^ {\prime}} \left(t ^ {\prime}\right)} \right| \right| _ {2} ^ {2} \tag{5}
$$
Capacity Loss Neural representations should be as high capacity as possible. That means that given a finite neural representation space, the number of unique spatial positions that the neural representations can uniquely encode should be maximized. In our particular setting, because our representations lie on the hypersphere, we can incentivize high capacity in the following manner:
$$
\mathcal {L} _ {C a p} \stackrel {\text {d e f}} {=} - \left| \left| \frac {1}{B T} \sum_ {\pi_ {b}, t} \boldsymbol {g} _ {\pi_ {b} (t)} \right| \right| _ {2} ^ {2} \tag {6}
$$
a
Figure 4: Multiple modules of grid cells in trained RNN: (a) The RNN's hidden units learn three discrete modules of spatially periodic representations. Four cells from each module are displayed.
This capacity loss pushes together all neural embeddings, regardless of path traveled or point in time, encouraging as little of the embedding space to be allocated for the training data as possible.