# Interpretable statistical representations of neural population dynamics and
geometry

Adam Gosztolai,1∗ Robert L. Peach,2 Alexis Arnaudon,3,4
Mauricio Barahona,4 Pierre Vandergheynst1

1Signal Processing Laboratory (LTS2), EPFL, Lausanne, 1016, Switzerland
2Department of Neurology, University Hospital Würzburg, Würzburg, 97070,
Germany
3 Blue Brain Project, EPFL, Campus Biotech, 1202, Geneva, Switzerland
4Department of Mathematics, Imperial College London, London, SW7 2AZ, United
Kingdom

∗To whom correspondence should be addressed; E-mail<EMAIL_ADDRESS>
###### Abstract

The dynamics of neuron populations during many behavioural tasks evolve on
low-dimensional manifolds. However, it remains challenging to discover latent
representations from neural recordings that are interpretable and consistently
decodable across individuals and conditions without explicitly relying on
behavioural information. Here, we introduce MARBLE, a fully unsupervised
geometric deep learning framework for the data-driven representation of non-
linear dynamics based on statistical distributions of local dynamical
features. Using both in silico examples from non-linear dynamical systems and
recurrent neural networks and in vivo recordings from primates and rodents, we
demonstrate that MARBLE can infer latent representations that are highly
interpretable in terms of global system variables such as decision-thresholds,
kinematics or internal states. We also show that MARBLE representations are
consistent across neural networks and animals, so that they can be used to
compare cognitive computations or train universal decoders. Through extensive
benchmarking, we show that unsupervised MARBLE provides best-in-class within-
and across-animal decoding accuracy, comparable to or significantly better
than current supervised approaches, yet without the need for behavioral
labels. Our results suggest that using the manifold structure in conjunction
with temporal information of neural dynamics provides a common framework to
develop better decoding algorithms and assimilate data across experiments.

## Introduction

It is increasingly recognised that computations in the brain and in artificial
neural networks can be understood as outputs of a high-dimensional dynamical
system conformed by the activity of large populations of neurons [1, 2, 3]. A
prominent line of research has observed that behaviourally-relevant neural
activity often takes place on low-dimensional smooth subspaces of the full
state space, usually called _neural manifolds_ [4, 5, 6, 7, 8, 9, 10, 11, 12].
From this perspective, several works have focused on how the geometry [4, 7,
13, 14] or topology [6, 8, 9] of neural manifolds relates to the underlying
task or computation. By contrast, others have suggested that dynamical flows
of neural population activity have an equally prominent role [15, 11, 16, 17]
and that the geometry of the manifold is merely an imprint that changes over
time or across individuals [4]. Although recent experimental techniques
provide means to simultaneously record the activity of large neuron
populations [18, 19, 20], revealing the structure of the underpinning latent
dynamical processes from data and interpreting their relevance in
computational tasks remains a fundamental challenge [16].

Overcoming this challenge requires computational frameworks that leverage the
manifold structure of neural dynamical flows for the unsupervised discovery of
its behavioural correlates. Yet current techniques fall short of this aim.
Indeed, unsupervised dimensionality reduction approaches, both linear methods
such as PCA or non-linear manifold learning methods such as t-SNE [21] or UMAP
[22], treat the neural activations as a point cloud, from which they infer a
manifold but cannot distinguish different dynamical processes that
distributionally yield the same manifold. Alternatively, dynamical systems
models can be used to learn temporal information of neural activity [23, 24,
25, 26, 27, 28]. However, these models depend on the particular embedding of
trajectories in state space and, in turn, the measured variables (i.e.,
neurons), hindering their ability to compare representations across
experiments. Topological data analysis methods can circumvent some of these
issues by characterising invariant structures in the data, e.g., loops [6] and
tori [9]. However, these methods are designed to capture qualitative
behaviours and changes (e.g., bifurcations) but fail to capture quantitative
changes in dynamics and geometry that can be crucial during representational
drift [4] or gain modulation [29]. Only recently, representation learning
methods such as pi-VAE [30] and CEBRA [31] have been introduced to infer
interpretable latent embeddings of neural trajectories and provide accurate
decoding thereof into behaviour. Yet, these methods are either supervised or
self-supervised, i.e., requiring user-defined ground-truth labels such as task
conditions, kinematics, or time labels. This is undesirable during discovery
as such labels are often unavailable or can introduce unintended
correspondence between experimental conditions or trials. For the same
reasons, supervised approaches cannot provide an unbiased distance metric to
compare neural dynamics across conditions or systems.

Here, we introduce a fully unsupervised representation-learning method called
MARBLE, or MAnifold Representation Basis LEarning, which obtains highly
interpretable and decodable representations from neural dynamics and provides
a well-defined similarity metric between dynamical processes, even in
different systems. MARBLE combines ideas from empirical dynamical modelling
[32] and the statistical descriptions of collective systems [33, 34] to
represent non-linear dynamics over manifolds. Specifically, stemming from
notions in differential geometry, MARBLE decomposes the dynamics into a set of
local flow fields around each sample point and uses them to encode global
information of the dynamics by building a similarity-preserving embedding
using geometric deep learning [35, 36, 37] and unsupervised contrastive
learning [38]. We show that MARBLE representations inferred from
electrophysiological recordings of the premotor cortex of macaques during a
reaching task and of the hippocampus of rats during a spatial navigation task
are substantially more interpretable and decodable than those obtained using
current representation learning frameworks [31]. Further, the local viewpoint
of MARBLE provides a robust similarity metric between the sparsely and
irregularly sampled dynamical systems typical in neural recordings obtained
over a limited number of trials. We show that this data-driven metric is
expressive enough to infer continuous and qualitative changes in the dynamical
landscape of recurrent neural networks during sensory gain modulation or in
decision-making tasks at the decision threshold, respectively. Finally, we
show that MARBLE facilitates comparing dynamics across entirely different
systems by discovering consistent embeddings across animals without auxiliary
transformations and offering a well-defined similarity metric irrespective of
the embedding of the dynamics.

Our results suggest that differential geometric notions can reveal as yet
unaccounted-for non-linear variations in neural data that can further our
understanding of neural dynamics underpinning neural computations and
behaviour.

## Results

Figure 1: Data-driven statistical representation of measured dynamics over
manifolds. a Trajectories measured in different trials (colours) evolving over
a latent manifold. b Vector field representation of trajectories. A
$k$-nearest neighbour graph approximates the manifold. The black circle marks
a sample point and its corresponding local flow field (LFF). c Geometric deep
learning model, MARBLE, for the node-wise embedding of LFFs. The vector field
is smoothed by an optional vector diffusion layer. The LFFs are approximated
to $p$-th order by gradient filters and optionally transformed into rotation-
invariant inner product features. Finally, a multilayer perceptron (MLP)
performs a node-wise similarity-preserving embedding of LFFs into a latent
space. d Top: four toy vector fields, two constant and two rotational fields.
Bottom: geometry-aware and geometry-agnostic MARBLE embeddings. Geometry-aware
embeddings distinguish rotational information in LFFs, whereas geometry-
agnostic embeddings learn only vector field expansion and contraction. e
Vector fields of the Van der Pol oscillator over a paraboloid of variable
curvature in the unstable ($\mu=-0.25$) and stable ($\mu=0.25$) regimes
sampled from randomly initialised trajectories. Insets show the limit cycle in
red and representative trajectories from a vertical projected view. f
Distribution distances across conditions. Clustering indicates an abrupt
dynamical change at $\mu=0$. g Two-dimensional MDS embedding of the
distribution distance matrix recovers the ordering of parameter $\mu$ over two
weakly connected one-dimensional manifolds.

### Unsupervised data-driven distributional representation of vector fields
over manifolds

To characterise neural computations during a task, e.g., decision-making or
arm-reaching, a typical experiment involves a set of trials under condition
$c$, e.g., a given stimulus or task condition, producing a corresponding set
of $d$-dimensional time-series $\\{\mathbf{x}(t)\\}$ of neural activity, e.g.,
firing rate or fluorescence intensity. Here, $d$ is the number of neurons or a
dimensionally-reduced set of variables. Frequently, one performs recordings
under diverse conditions and seeks to discover, without supervision, how the
structure of neural dynamics relates to task structure. In other setups, one
may desire to compare trajectories across conditions or even animals to reveal
alterations in neural mechanisms. Both applications are challenging as
trajectories initialised differently across trials produce uneven and possibly
sparse sampling of the relevant state space features so that similar dynamics
may unfold over manifolds with different embeddings across subjects or
conditions [4, 27].

To address these challenges, MARBLE departs from the current modelling
paradigm of learning the temporal evolution of the states in single isolated
trajectories [27, 28, 31]. Instead, it represents the dynamics statistically
over ensembles of trajectories using a latent distribution obtained from the
local dynamical context of each sample point. This is possible as the
trajectories of neural dynamics [15] or, more generally, of dissipative
dynamical systems [39] converge to a submanifold of the original
$d$-dimensional state space (Fig. 1a). We may equivalently treat the
trajectories from a set of trials under a given experimental condition $c$ as
a vector field $\mathbf{F}_{c}=(\mathbf{f}_{1},\dots,\mathbf{f}_{n})$ anchored
to a point cloud $\mathbf{X}_{c}=(\mathbf{x}_{1},\dots,\mathbf{x}_{n})$ (Fig.
1b), where $n$ is the number of sample points, e.g., recorded neural
activations across all trials under this condition. By approximating the
underlying manifold using a proximity graph to $\mathbf{X}_{c}$, we define the
dynamical context of a point $i$ as the local flow field (LFF) supported on
points that are at most a distance $p$ from $i$ measured over the graph. We
jointly embed individual LFFs into a $E$-dimensional shared latent space to
represent the dynamical system as an empirical distribution $P_{c}$ formed by
a collection of features $\\{\mathbf{z}_{i}\\}$. To define a metric between
dynamical systems under conditions $c$ and $c^{\prime}$, we use the optimal
transport distance (Eq. S18), $d(P_{c},P_{c^{\prime}})$, between their
distributional latent representations, which generally outperforms entropic
measures (e.g., KL divergence) when detecting complex interactions based on
overlapping distributions [34].

We note that the above approach is substantially different from classical
dimensionality reduction approaches such as PCA, t-SNE or UMAP, which map the
neural activity $\mathbf{x}_{i}$ directly into latent features and, therefore,
are oblivious to the dynamics that have generated the point cloud $X_{c}$. By
learning the variation of the vector field within the LFF around each point,
we also learn dynamical information, which substantially enhances the
representational capability of our method.

Unsupervised learning of the latent representation $P_{c}$ follows from the
continuity of LFFs over the manifold, which means that close LFFs are more
similar than randomly sampled ones. To capture global information using LFFs,
we develop a similarity-preserving embedding of LFFs using a geometric deep
learning architecture consisting of four components (Fig. 1c, see Sect. S1 in
Materials and Methods). Following an optional vector diffusion kernel with
learnable parameter $\tau$ (Eq. S3) which tunes the spatial extent of the LFFs
we apply: (i) $p$ gradient filter layers whose output together with
$\mathbf{f}_{i}$ gives the best $p$-th order approximation of the LFF around
$i$ (Figs. S1–S3, Eq. S13); (ii) inner product features involving learnable
linear transformations which make the gradient features rotation invariant
(Figs. S4–S5, Eq. S15); and (iii) a multilayer perceptron (MLP) with learnable
weights $\omega$ that produces the latent feature $\mathbf{z}_{i}$.

We perform unsupervised training via a soft constraint in the loss function
(Eq. S17) to learn to embed adjacent LFFs over the graph to lie close by
whereas randomly sampled LFFs are placed further apart. In contrast to
supervised representation-learning methods [30, 31], MARBLE embeds LFFs across
all manifolds (i.e., conditions) with shared weights, and thus assigns similar
latent features to similar LFFs across conditions, allowing their unsupervised
comparison.

### Comparing dynamics within and across systems: geometry-aware vs. geometry-
agnostic representations

The inner product features in MARBLE give rise to two operation modes that can
be exploited to provide insights about the inferred neuronal representations.
To illustrate this issue, we consider a toy model consisting of linear and
constant (rotational) vector fields over a two-dimensional plane ($d=2$,
trivial manifold) in Fig. 1d, top. As shown later, MARBLE can capture complex
non-linear dynamics over non-linear manifolds beyond this simple illustrative
example. In geometry-aware mode, the inner product features are disabled (Fig.
1c), allowing MARBLE to learn the orientation of the LFFs, ensuring maximal
expressivity and interpretability. In our toy example, constant fields (left
and right) are mapped into two distinct clusters, whereas rotational fields
are distributed based on their angular orientation (Fig. 1d, bottom). This
mode of operation in MARBLE is useful when representing dynamics across
conditions within a given animal or neural network, i.e. when the same neurons
are sampled (Fig. 1d-g, Fig. 3-4b,c), or when the geometry of the neural
manifold is essential to characterise the behaviour (Fig. 4e).

In geometry-agnostic mode, on the other hand, the inner product features are
present and make the learnt features rotation invariant. In our toy example,
this means that vector fields are no longer distinguishable based on LFF
orientation, but we still capture other differences, e.g., expansion and
contraction of the flows (Fig. 1d, bottom). This geometry-agnostic mode is
able to discount the arbitrary rotations in the LFFs induced by different
embeddings [40] (Figs. S4–S5), and is therefore useful when comparing across
systems, such as neural networks trained from different initialisations (Fig.
2h-j). To demonstrate this for a non-linear dynamical system on a non-linear
manifold, we analysed the Van der Pol oscillator mapped to a paraboloid while
varying the damping parameter $\mu$ and the manifold curvature (Fig. 1e, see
Sect. S3.1 in Materials and Methods). Using short simulated trajectories from
random initial conditions for $20$ different values of $\mu$, we trained a
geometry-agnostic MARBLE network to jointly embed the corresponding vector
fields into a shared latent space ($E=5$). Despite the sparse and irregular
sampling, we detected robust dynamical variation as $\mu$ was varied.
Specifically, the matrix of distribution distances
$D_{cc^{\prime}}=d(P_{c},P_{c}^{\prime})$ displays a two-partition structure,
indicating the presence of two different types of dynamic behaviours (stable
and unstable) separated by the Hopf bifurcation at $\mu=0$ (Fig. 1f). This
result is observed independently of manifold curvature (Fig. S6). Furthermore,
and despite the irregular and sparse sampling across conditions, embedding
$D_{cc^{\prime}}$ using multidimensional scaling (MDS) revealed a one-
dimensional manifold that captures the variation of $\mu$ and consequent
compression and expansion of flow (Fig. 1g). Notably, this continuous manifold
was not captured when training a geometry-aware network on the same data (Fig.
S6), thus confirming that geometry-agnostic MARBLE can extract both abrupt and
continuous dynamical variation while being invariant to manifold embedding at
a slight loss of expressivity compared to the geometry-aware mode.

### Discriminating computational mechanisms in recurrent neural networks

Figure 2: Statistical representations represent and discriminate computations
across recurrent neural networks. a Low-rank RNN takes two stimuli as input
and produces a decision variable as a read-out. b Representative stimulation
patterns and decision outcomes for the delayed-match-to-sample task (one
stimulus shown). Input amplitude equals the ’gain’ during stimulation epochs
(red) and zero otherwise (grey). c Neural dynamics of a trained rank-two RNN
evolves on a randomly oriented plane in neural state space. Phase portrait
shows mean field dynamics superimposed with a trajectory during a trial
(orange). d Input weight space to neurons in a trained RNN. Colours are
obtained by k-means clustering, indicating subpopulations specialised in
respective inputs in a. Ellipses represent $3$ std of fitted Gaussian
distributions. e Phase portraits for different stimulus gains (0, 0.32, 1.0).
f Distribution distances across gains obtained from geometry-aware embeddings.
Hierarchical clustering indicates two clusters in the non-zero gain region,
showing a qualitative change, i.e., bifurcation (red dashed line). g The
inferred bifurcation point predicts the drop in task performance. h Networks
trained from different initialisations can have different dynamics over
differently embedded manifolds (Solution I, Solution II). As a control, we
sampled networks from the weight distribution in d to obtain the same dynamics
over differently embedded manifolds (Solution I/1, Solution I/2). i
Distribution distances across the three networks and gains obtained from
geometry-agnostic embeddings compared with that obtained from CCA. j MDS
embedding of the distance matrix in i shows quantitative, continuous variation
and alignment of gain-modulated conditions across similar dynamics.

There has been significant recent interest in RNNs as surrogate models for the
computations performed by the brain [15, 41, 42, 43]. Previous approaches for
comparing RNN computations for a given task relied on representations of
population activity as point clouds, which were then compared using shape
analysis [44, 14, 41]. However, these methods cannot distinguish dynamics
producing similar point cloud distributions.

To demonstrate the advantages of MARBLE’s LFF-based representation, we trained
RNNs with a rank-two connectivity matrix (Fig. 2a, see Sect. S3.2 in Materials
and Methods) on the delayed-match-to-sample task [45]. Rank-two RNNs were
previously shown to be sufficiently expressive to learn this task [46]. This
common contextual decision-making task comprises two distinct stimuli with
variable gain and two stimulation epochs of variable duration interspersed by
a delay (Fig. 2b). Setting the gain to unity, the RNN was trained to converge
to an output of $1$ if either stimulus was presented during both epochs and
$-1$ otherwise (Fig. 2b). As expected [47], the neural dynamics of a trained
network during a trial evolves on a randomly oriented plane in state space
(Fig. 2c and Fig. S7a-c). Yet we found that differently initialised networks
produce two classes of solutions characterised by the clustering of their
input weight distributions
$\mathbf{w}^{\mathsf{in}}_{1},\mathbf{w}^{\mathsf{in}}_{2}$ (Eq. S22):
solution I consists of two subpopulations specialised in sensing the two
stimuli (Fig. 2d), whilst in solution II the nodes generalise across the two
stimuli (Fig. S7d). These two solutions exhibit substantially different neural
dynamics consisting of multiple fixed points, saddles and line attractors, as
shown by two representative examples, which we selected for further analysis
(Fig. 2e, Fig. S7e).

We first asked whether MARBLE could infer dynamical neural correlates of loss
of task performance as the stimulus gain is decreased. For $20$ different
input gains, we simulated $200$ randomly initialised trials. For each gain, we
extracted the part of the trajectories where the stimulus was present and used
the remainder of the trajectory as a negative control because we expected no
gain-related variation therein. We then trained a geometry-aware MARBLE
network to embed all groups into a common latent space. The resulting matrix
of pairwise distribution distances exhibits a block-diagonal structure (Fig.
2f). The top left block contains vanishing entries corresponding to the zero-
gain controls, implying that MARBLE is robust to sampling variability. The
other two subblocks of the bottom right block indicate a quantitative change
in dynamics, which remarkably corresponds to a sudden drop in task performance
(Fig. 2g) from $1$ to $0.5$ (random). This shows that MARBLE enables the
unsupervised detection of dynamical events that are interpretable in terms of
global decision variables.

Next, we asked if MARBLE can define a similarity metric for the dynamics
across distinct RNNs. To this end, we took network Solutions I and II, which
have different dynamics evolving over differently embedded manifolds (Fig.
2h). As a negative control, we randomly sampled two new networks from the
weight distribution of Solution I because low-rank RNNs sampled from the same
Gaussian weight distribution have the same mean field dynamics [46] albeit on
differently oriented manifolds. Due to the arbitrary embedding of the
manifolds across networks, we used geometry-agnostic MARBLE to produce feature
embedding for $200$ trials drawn from these three RNNs during stimulation
($E=5$). We found the rotation-invariant features of the geometric-agnostic
MARBLE is insentitive to the embedding of dynamics, i.e., plane orientation,
and can uniquely detect a block diagonal structure in the distribution
distance matrix across gain-modulated conditions, clearly discriminating
between the on-manifold dynamics of networks sampled from Solution I and II
while producing similar dynamics for the control (Fig. 2i, left). By contrast,
canonical correlation analysis (CCA), which compares datasets based on point
cloud-based representations, cannot detect dynamical changes because these
occur within planes spanned by the same principal axes (Fig. 2i, right). This
also suggests that previous methods that used CCA to detect changes across
RNNs [13] or animals [4, 48] could have detected the compound effect of
changing manifold curvature and different sampling of state space, without
concluding dynamical differences. In addition to correctly clustering
solutions, the MDS embedding of the distance matrix (Fig. 2j) also shows that
MARBLE reveals the underlying one-parameter variation due to gain modulation,
which is correctly ordered within a network and across sampled networks. These
findings confirm that our method can obtain a robust metric between dynamical
processes generated by different system architectures with possible
implications towards guiding their design for faster training or more accurate
predictions.

### Representing and decoding neural dynamics in a macaque reaching task

Figure 3: Intepretable representation and decoding of neural activity in the
prefrontal cortex in macaques during a centre-out reaching task. a Ground
truth hand trajectories in seven reaching directions. b Single-trial spike-
trains in premotor cortex during a single trial for three respective reach
directions ($24$ recording channels). The shaded area shows the analysed
traces after the GO cue. c Ten representative firing rate trajectories for a
given reach condition PCA-embedded in 3D for visualisation. d Vector field
obtained from firing rate trajectories. e MARBLE embedding of neural data
alone across conditions in a single session, compared with CEBRA (using reach
conditions as labels) and LFADS. MARBLE embedding of neural dynamics for a
given session geometrically corresponds to the reach configuration in physical
space. f Linear decoding of hand trajectories from latent representations. g
Decoding accuracy measured by $R^{2}$ between ground truth and decoded
trajectories across all sessions for the final position (left) and
instantaneous velocity (right). Wilcoxon tests (paired samples), **:
$p<1\times 10^{-2}$, ***: $p<1\times 10^{-3}$, ****: $p<1\times 10^{-4}$.
Horizontal and vertical bars show mean and one std, respectively.

State-of-the-art representation-learning of neural dynamics uses a joint
embedding of neural and behavioural signals [31]. However, true biological
discovery should be based on post hoc interpretation of unsupervised
representations rather than supervised learning based on pre-determined, user-
defined condition labels. To exemplify this issue, we reanalysed
electrophysiological recordings of a macaque performing a delayed centre-out
hand-reaching task [27] (see Sect. S3.3 in the Material and Methods). During
the task, a trained monkey was instructed to move a handle towards seven
distinct targets placed at radial locations from the start position. This
dataset comprises simultaneous recordings of hand kinematics (Fig. 3a) and
neural activity via a $24$-channel probe inserted into the premotor cortex
(PMd) over $44$ recording sessions (Fig. 3b shows one session).

Since a subset of neurons in PMd is directionally tuned [49], we argued that
neural representations must be sensitive to the orientation of the neural
manifold. For maximal expressivity, in this case, we trained a geometry-aware
MARBLE network by constructing separate vector fields from the firing rate
traces under individual reach conditions across which we expected the dynamics
to be consistent (Fig. 3c-d). For comparison, we produced embeddings using two
other prominent approaches: CEBRA [31], a supervised model using reach
condition as labels (CEBRA-behaviour), and LFADS [27], an unsupervised model
that uses generative recurrent neural networks. Our results show that the
unsupervised MARBLE feature embeddings in 3D strongly reflect the spatial
geometry of the physical reaches (Fig. 3e and Fig. S9a). By comparison,
although the supervised CEBRA could partly unfold the arrangement of the
reaches, the temporal evolution of trajectories was lost, and the latent
factors from unsupervised LFADS lacked interpretability (Fig. 3e). The
embeddings from point-wise methods such as PCA, t-SNE or UMAP showed similar
loss of information (Fig. S8). This geometric configuration is also confirmed
by the diagonal and periodic structure of the condition-averaged distance
matrix between feature distributions across reach directions (Fig. S9b) and
its associated MDS embedding (Fig. S9c). Hence, MARBLE can unfold global
geometric information in the neural code that mirrors kinematics in physical
space.

This interpretability based on a geometric correspondence between neural and
behavioural representations suggests a potent decoding algorithm. To show
this, we fitted an optimal linear estimator between the latent representations
and their corresponding hand positions. Remarkably, the decoded kinematics
showed excellent visual correspondence to ground truth kinematics both in
directionality, comparable to supervised CEBRA and substantially better than
unsupervised LFADS (Fig. 3f), quantified by a 10-fold cross-validated
classification of final reach direction (Fig. 3g). Further, we find that while
a three-dimensional ($E=3$) MARBLE embedding accurately encodes spatial
positions, the delay between the GO cue and the beginning of the movement
could only be accounted for with higher embedding dimensions ($E=20$, Fig.
3g). In this regard, MARBLE significantly outperformed both supervised CEBRA
and LFADS models, which could not sufficiently unfold the latency and temporal
variation of reaches to achieve linear decodability. These results highlight
that MARBLE can generate representations of neural dynamics that are
simultaneously interpretable and decodable into behavioural variables.

### Consistent embeddings of neural dynamics across animals in the rat
hippocampus

Figure 4: Intepretable representation and cross-animal decoding of neural
activity in rat hippocampus during linear maze navigation. a Top: Experimental
setup. Bottom: raster plot showing spiking activity in 120 neurons in a single
session. b Comparison of latent representation ($E=3$) of unsupervised MARBLE
against self-supervised (time-only labels) and supervised (time, position and
direction labels) CEBRA. c Time traces of linearly decoded animal position
within Rat 1 ($E=32$, default settings from CEBRA decoding notebook examples).
d Decoding accuracy within the same animal. Wilcoxon tests (paired samples),
****: $p<1\times 10^{-4}$. Horizontal and vertical bars show mean and one std,
respectively. e Cross-animal consistency as measured by $R^{2}$ of linear fit
between the optimally aligned 3D embeddings of a source animal to a target
animal.

Recent evidence suggests there exists a strong similarity between neural
representations across animals in a given task [48], with profound
implications for, e.g., brain-machine interfacing. However, as shown above
(Fig. 2i, Fig. S8), point-wise methods based on CCA [13] and similar shape
metrics [44, 14, 13] do not capture dynamical variation that otherwise
preserves the geometry of the neural manifold. While it is possible to
‘stitch’ multiple datasets by fitting auxiliary transformations [27] or
aligning the neural dynamics to behaviour [31], this relies on the assumption
that the respective neurons encode the same computation.

Given that MARBLE can produce embeddings that are comparable across RNNs (Fig.
2h-i) and are interpretable within a given RNN (Fig. 2e-g) or animal (Fig.
3h-i), we finally asked whether it can produce consistently decodable
embeddings across animals. To this end, we re-analysed electrophysiological
recordings from the rat hippocampus during navigation of a linear track [50]
(Fig. 4a, Sect. S3.4 in Materials and Methods). We found that from the neural
data alone, MARBLE can infer interpretable representations consisting of a
one-dimensional manifold in neural state space representing the animal’s
position and walking direction (Fig. 4b). Remarkably, unsupervised MARBLE
representations were significantly more interpretable than those obtained with
CEBRA supervised by time labels alone (CEBRA-time) and comparable to providing
both positional and directional information as labels (CEBRA-behaviour, Fig.
4b). This finding was corroborated by significantly higher decoding accuracy
using a K-means decoder (Fig. 4c,d).

Finally, to demonstrate the consistency of MARBLE embeddings across animals,
we aligned them post hoc using a linear transformation between animals. We
found that the consistency, measured using the $R^{2}$ fit from a linear model
trained using one animal as the target and one animal as the source, was
significantly higher for unsupervised MARBLE than for self-supervised CEBRA-
time, although not as good as supervised CEBRA-behaviour (Fig. 4e). This is
remarkable given that MARBLE does not rely on external behavioural data yet
finds consistent representations despite experimental and neurophysiological
differences across animals. These findings underscore MARBLE’s potential for
data-driven discovery and applications such as brain-computer interfaces.

## Discussion

A hallmark of large collective systems such as the brain is the existence of
many system realisations that lead to equivalent computations defined by
population-level dynamical processes [51, 52]. The growing recognition that
dynamics in biological and artificial neural networks evolve over low-
dimensional manifolds [5, 6, 8, 9] offers an opportunity to reconcile the
variability of dynamics across system realisations with the observed
invariance of computations by using manifold geometry as an inductive bias for
designing data-driven models. We have shown that non-linear dynamical systems
can be represented as statistical distributions built from a point-wise,
similarity-preserving embedding of local flow fields. Due to the continuity of
the dynamics over the manifold, latent features can be learnt in a fully
unsupervised manner using a geometric deep learning architecture. Further,
features can be made rotation-invariant to achieve robust representations of
perturbations in manifold geometry and orientation. These properties open the
door to independently testing the relevance of manifold geometry and dynamics
and comparing dynamics across system realisations.

Our formalism can be framed as a statistical generalisation of the convergent
cross mapping (CCM) framework by Sugihara et al. [32], which tests the
topological equivalence of two dynamical attractors through a one-to-one map
between their local neighbourhoods. CCM tests the causality between two long,
concurrent time series without learning. In contrast, our method uses
geometric deep learning to extract interpretable representations and provides
a similarity metric between any collection of dynamical systems based on
ensembles of variable-length, sparsely sampled time series. In addition, due
to the locality of representations, our approach diverges from typical
geometric deep learning models that learn vector fields globally [37, 53] and
are thus unable to consider the manifold geometry and the dynamics separately.
Locality also implies that our method can generalise to assimilate different
datasets without additional trainable parameters to increase the statistical
power of the model even when individual datasets are poorly sampled. Although
our method does not explicitly learn temporal dynamics over manifolds [26,
54], we showed that temporal ordering naturally emerges from our similarity-
preserving embeddings of local vector fields. We demonstrated that this
property enables kinematic decoding from embeddings of neural data. Beyond
their use in interpreting and decoding neural dynamics, we expect MARBLE
embeddings to provide powerful representations for general machine-learning
tasks, which is highly relevant to research in brain-machine interfaces.

Even though there is an increasing trend to jointly embed neural activity and
behaviour to infer interpretable and consistent embeddings that do not suffer
from ’nuisance variables’, we have shown that MARBLE can uncover such
embeddings from neural data alone. Hence, the LFF-based representations of
neural dynamics obtained by MARBLE contain more behaviourally relevant
information than neural data alone without the need for supervised learning
using additional behavioural information to obtain comparable embeddings. This
suggests a hypothesis that the neural readout into behaviour in biological
brains might rely on the context of neural trajectories in a broader dynamical
landscape.

## Code availability

The code to carry out the simulations and analysis can be found at
github.com/agosztolai/MARBLE.

## Data availability

The data generated during the simulations is available with DOI:
10.7910/DVN/KTE4PC.

## Acknowledgements

We thank Nicolas Aspert for the much-needed computing support. We also thank
Adrian Valente, Nikolas Karalias and Matteo Vinao-Carl for the interesting
discussions. AG acknowledges support from an HFSP Cross-disciplinary
Postdoctoral Fellowship (LT000669/2020-C). MB acknowledges funding through
EPSRC awards EP/N014529/1 (Centre for Mathematics of Precision Healthcare) and
EP/W024020/1 (Statistical Physics of Cognition). RP acknowledges the Deutsche
Forschungsgemeinschaft (DFG, German Research Foundation) Project-ID
424778381-TRR 295. A.A. was supported by funding to the Blue Brain Project, a
research centre of the École polytechnique fédérale de Lausanne (EPFL), from
the Swiss government’s ETH Board of the Swiss Federal Institutes of
Technology.

## Author contributions

Conceptualisation: A.G., R.L.P., A.A., M.B., P.V., Methodology: A.G., R.L.P.,
A.A., P.V., Software: A.G., R.L.P., A.A., Analysis: A.G., R.L.P., A.A.,
Writing - Original Draft: A.G., R.L.P., A.A., Writing - Review & Editing:
M.B., P.V.

## Competing interests

The authors declare no competing interests.

## References

  * [1] J J Hopfield  “Neural networks and physical systems with emergent collective computational abilities.”  In _Proc. Natl. Acad. Sci. U. S. A._ 79.8, 1982, pp. 2554–2558
  * [2] Mark M Churchland et al.  “Neural population dynamics during reaching”  In _Nature_ 487.7405, 2012, pp. 51–56  DOI: 10.1038/nature11129
  * [3] Saurabh Vyas et al.  “Computation through Neural Population Dynamics”  In _Annu. Rev. Neurosci._ 43, 2020, pp. 249–275  DOI: 10.1146/annurev-neuro-092619-094115
  * [4] Juan A Gallego et al.  “Neural manifolds for the control of movement.”  In _Neuron_ 94.5, 2017, pp. 978–984  DOI: 10.1016/j.neuron.2017.05.025
  * [5] SueYeon Chung, Daniel D Lee and Haim Sompolinsky  “Classification and geometry of general perceptual manifolds”  In _Phys. Rev. X_ 8.3, 2018, pp. 31003  DOI: 10.1103/physrevx.8.031003
  * [6] Rishidev Chaudhuri et al.  “The intrinsic attractor manifold and population dynamics of a canonical cognitive circuit across waking and sleep”  In _Nat. Neurosci._ 22.9, 2019, pp. 1512–1520  DOI: 10.1038/s41593-019-0460-x
  * [7] Nikolaus Kriegeskorte and Xue-Xin Wei  “Neural tuning and representational geometry”  In _Nat. Rev. Neurosci._ 22.11, 2021, pp. 703–718  DOI: 10.1038/s41583-021-00502-3
  * [8] Mikail Khona and Ila R. Fiete  “Attractor and integrator networks in the brain”  In _Nat. Rev. Neurosci._ 23.12, 2022, pp. 744–766  DOI: 10.1038/s41583-022-00642-0
  * [9] Richard J. Gardner et al.  “Toroidal topology of population activity in grid cells”  In _Nature_ 602.7895, 2022, pp. 123–128  DOI: 10.1038/s41586-021-04268-7
  * [10] Ramon Nogueira et al.  “The geometry of cortical representations of touch in rodents”  In _Nat. Neurosci._ 26.2, 2023, pp. 239–250  DOI: 10.1038/s41593-022-01237-9
  * [11] Manuel Beiran et al.  “Parametric control of flexible timing through low-dimensional neural manifolds”  In _Neuron_ 111.5, 2023, pp. 739–753.e8  DOI: 10.1016/j.neuron.2022.12.016
  * [12] Christopher Langdon, Mikhail Genkin and Tatiana A Engel  “A unifying perspective on neural manifolds and circuits for cognition”  In _Nat. Rev. Neurosci._ Springer US, 2023  DOI: 10.1038/s41583-023-00693-x
  * [13] Niru Maheswaranathan et al.  “Universality and individuality in neural dynamics across large populations of recurrent networks”  In _Adv. Neural Inf. Process. Syst._ 32.NeurIPS, 2019
  * [14] Alex H. Williams et al.  “Generalized shape metrics on neural representations”  In _Adv. Neural Inf. Process. Syst._ 6.NeurIPS, 2021, pp. 4738–4750
  * [15] David Sussillo and Omri Barak  “Opening the black box: low-dimensional dynamics in high-dimensional recurrent neural networks”  In _Neural Comput._ 25.3, 2013, pp. 626–649  DOI: 10.1162/NECO˙a˙00409
  * [16] Lea Duncker and Maneesh Sahani  “Dynamics on the manifold: Identifying computational dynamical activity from neural population recordings”  In _Curr. Opin. Neurobiol._ 70, 2021, pp. 163–170  DOI: 10.1016/j.conb.2021.10.014
  * [17] Robert L. Peach et al.  “Implicit Gaussian process representation of vector fields over arbitrary latent manifolds”, 2023  arXiv:2309.16746 [cs.LG]
  * [18] Phan Tan Toi et al.  “In vivo direct imaging of neuronal activity at high temporospatial resolution”  In _Science_ 378.6616, 2022, pp. 160–168  DOI: 10.1126/science.abh4340
  * [19] Nicholas A. Steinmetz et al.  “Neuropixels 2.0: A miniaturized high-density probe for stable, long-term brain recordings”  In _Science_ 372.6539, 2021, pp. eabf4588  DOI: 10.1126/science.abf4588
  * [20] Vincent Villette et al.  “Ultrafast Two-Photon Imaging of a High-Gain Voltage Indicator in Awake Behaving Mice”  In _Cell_ 179.7, 2019, pp. 1590–1608.e23  DOI: https://doi.org/10.1016/j.cell.2019.11.004
  * [21] Laurens Maaten and Geoffrey Hinton  “Visualizing Data using t-SNE”  In _Journal of Machine Learning Research_ 9.86, 2008, pp. 2579–2605  URL: http://jmlr.org/papers/v9/vandermaaten08a.html
  * [22] Leland McInnes et al.  “UMAP: Uniform Manifold Approximation and Projection”  In _The Journal of Open Source Software_ 3.29, 2018, pp. 861
  * [23] Steven L. Brunton et al.  “Chaos as an intermittently forced linear system”  In _Nat. Commun._ 8.1, 2017, pp. 1–8  DOI: 10.1038/s41467-017-00030-8
  * [24] Bethany Lusch, J Nathan Kutz and Steven L Brunton  “Deep learning for universal linear embeddings of nonlinear dynamics”  In _Nat. Commun._ 9.1, 2018, pp. 851–10  DOI: 10.1038/s41467-018-07210-0
  * [25] Thomas Kipf et al.  “Neural Relational Inference for Interacting Systems”  In _arXiv preprint arXiv:1802.04687_ , 2018
  * [26] Ryan J. Low et al.  “Probing variability in a cognitive map using manifold inference from neural dynamics”  In _bioRxiv_ Cold Spring Harbor Laboratory, 2018  DOI: 10.1101/418939
  * [27] Chethan Pandarinath et al.  “Inferring single-trial neural population dynamics using sequential auto-encoders”  In _Nat. Methods_ , 2018, pp. 1–21  DOI: 10.1038/s41592-018-0109-9
  * [28] Feng Zhu et al.  “A deep learning framework for inference of single-trial neural population dynamics from calcium imaging with subframe temporal resolution”  In _Nat. Neurosci._ 25.12, 2022, pp. 1724–1734  DOI: 10.1038/s41593-022-01189-0
  * [29] Jake P. Stroud et al.  “Motor primitives in space and time via targeted gain modulation in cortical networks”  In _Nat. Neurosci._ 21.12 Springer US, 2018, pp. 1774–1783  DOI: 10.1038/s41593-018-0276-0
  * [30] Ding Zhou and Xue-Xin Wei  “Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-VAE”, 2020  arXiv:2011.04798 [stat.ML]
  * [31] Steffen Schneider, Jin Hwa Lee and Mackenzie Weygandt Mathis  “Learnable latent embeddings for joint behavioural and neural analysis”  In _Nature_ 617.7960, 2023, pp. 360–368  DOI: 10.1038/s41586-023-06031-6
  * [32] George Sugihara et al.  “Detecting causality in complex ecosystems”  In _Science (80-. )._ 338.6106, 2012, pp. 496–500  DOI: 10.1126/science.1227079
  * [33] Adam Gosztolai and Alexis Arnaudon  “Unfolding the multiscale structure of networks with dynamical Ollivier-Ricci curvature”  In _Nat. Commun._ 12.1, 2021, pp. 4561  DOI: 10.1038/s41467-021-24884-1
  * [34] Dominic J. Skinner et al.  “Topological metric detects hidden order in disordered media”  In _Phys. Rev. Lett._ 126.4, 2021, pp. 48101  DOI: 10.1103/PhysRevLett.126.048101
  * [35] Nicholas Sharp et al.  “DiffusionNet : discretization agnostic learning on surfaces”  In _ACM Trans. Graph._ 99.99 Association for Computing Machinery, 2020  DOI: 10.1145/9999999.9999999
  * [36] Dominique Beaini et al.  “Directional graph networks”  In _International Conference on Machine Learning_ , 2021, pp. 748–758  PMLR
  * [37] Cristian Bodnar et al.  “Neural Sheaf Diffusion: A topological perspective on heterophily and oversmoothing in GNNs”  In _Advances in Neural Information Processing Systems_ , 2022
  * [38] William L. Hamilton, Rex Ying and Jure Leskovec  “Inductive representation learning on large graphs”  In _Adv. Neural Inf. Process. Syst._ 2017-Decem.Nips, 2017, pp. 1025–1035
  * [39] Charles Fefferman, Sanjoy Mitter and Hariharan Narayanan  “Testing the manifold hypothesis”  In _Journal of the American Mathematical Society_ 29.4, 2016, pp. 983–1049
  * [40] George Sugihara and Robert M May  “Nonlinear forecasting as a way of distinguishing chaos from measurement error in time series”  In _Nature_ 344.6268, 1990, pp. 734–741  DOI: 10.1038/344734a0
  * [41] Timo Flesch et al.  “Orthogonal representations for robust context-dependent task performance in brains and neural networks”  In _Neuron_ 110.7, 2022, pp. 1258–1270.e11  DOI: 10.1016/j.neuron.2022.01.005
  * [42] Rishi Rajalingham, Aída Piccato and Mehrdad Jazayeri  “Recurrent neural networks with explicit representation of dynamic latent variables can mimic behavioral patterns in a physical inference task”  In _Nat. Commun._ 13.1 Springer US, 2022  DOI: 10.1038/s41467-022-33581-6
  * [43] Aniruddh R. Galgali, Maneesh Sahani and Valerio Mante  “Residual dynamics resolves recurrent contributions to neural computation”  In _Nat. Neurosci._ , 2023  DOI: 10.1038/s41593-022-01230-2
  * [44] Simon Kornblith et al.  “Similarity of neural network representations revisited”  In _36th Int. Conf. Mach. Learn. ICML 2019_ 2019-June, 2019, pp. 6156–6175
  * [45] Yasushi Miyashita  “Neuronal correlate of visual associative long-term memory in the primate temporal cortex”  In _Nature_ 335.6193, 1988, pp. 817–820  DOI: 10.1038/335817a0
  * [46] Alexis Dubreuil et al.  “The role of population structure in computations through neural dynamics”  In _Nat. Neurosci._ 25.6, 2022, pp. 783–794  DOI: 10.1038/s41593-022-01088-4
  * [47] Francesca Mastrogiuseppe and Srdjan Ostojic  “Linking Connectivity, Dynamics, and Computations in Low-Rank Recurrent Neural Networks”  In _Neuron_ 99.3, 2018, pp. 609–623.e29  DOI: 10.1016/j.neuron.2018.07.003
  * [48] Mostafa Safaie et al.  “Preserved neural dynamics across animals performing similar behaviour”  In _Nature_ 623.September 2022 Springer US, 2023  DOI: 10.1038/s41586-023-06714-0
  * [49] A. Riehle and J. Requin  “Monkey primary motor and premotor cortex: single-cell activity related to prior information about direction and extent of an intended movement”  In _Journal of Neurophysiology_ 61.3, 1989, pp. 534–549  DOI: 10.1152/jn.1989.61.3.534
  * [50] Andres D Grosmark and György Buzsáki  “Diversity in neural firing dynamics supports both rigid and learned hippocampal sequences”  In _Science (80-. )._ 351.6280 American Association for the Advancement of Science, 2016, pp. 1440–1443  DOI: 10.1126/science.aad1935
  * [51] Eve Marder and Adam L. Taylor  “Multiple models to capture the variability in biological neurons and networks”  In _Nat. Neurosci._ 14.2, 2011, pp. 133–138  DOI: 10.1038/nn.2735
  * [52] Adam Gosztolai and Pavan Ramdya  “Connecting the dots in ethology: applying network theory to understand neural and animal collectives”  In _Curr. Opin. Neurobiol._ 73, 2022, pp. 102532  DOI: 10.1016/j.conb.2022.102532
  * [53] Daniele Grattarola and Pierre Vandergheynst  “Generalised implicit neural representations”  In _Advances in Neural Information Processing Systems_ , 2022
  * [54] Daniel Floryan and Michael D Graham  “Data-driven discovery of intrinsic dynamics”  In _Nat. Mach. Intell._ 4.12, 2022, pp. 1113–1120  DOI: 10.1038/s42256-022-00575-4
  * [55] David M Mount and Sunil Arya  “Finding the nearest neighbors for points in d-dimensional Euclidean space”  In _ACM Transactions on Mathematical Software (TOMS)_ 24.1 ACM, 1998, pp. 96–103
  * [56] Tyrus Berry and Timothy Sauer  “Consistent manifold representation for topological data analysis”  In _Foundations of Data Science_ 1.1, 2019, pp. 1–38  DOI: 10.3934/fods.2019001
  * [57] Max Budninskiy et al.  “Parallel Transport Unfolding: A Connection-Based Manifold Learning Approach”  In _SIAM Journal on Applied Algebra and Geometry_ 3.2, 2019, pp. 266–291  DOI: 10.1137/18M1196133
  * [58] W. Kabsch  “A solution for the best rotation to relate two sets of vectors”  In _Acta Crystallographica Section A_ 32.5, 1976, pp. 922–923  DOI: 10.1107/S0567739476001873
  * [59] Vergne M. Berline N.  “Heat kernels and Dirac operators”  Springer, 1996
  * [60] A. Singer and H.. Wu  “Vector diffusion maps and the connection Laplacian”  In _Commun. Pure Appl. Math._ 65.8, 2012, pp. 1067–1144  DOI: 10.1002/cpa.21395
  * [61] Kaiming He et al.  “Deep Residual Learning for Image Recognition”  In _2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2016, pp. 770–778  DOI: 10.1109/CVPR.2016.90
  * [62] Matthias Fey and Jan E. Lenssen  “Fast Graph Representation Learning with PyTorch Geometric”  In _ICLR Workshop on Representation Learning on Graphs and Manifolds_ , 2019
  * [63] Diederik P. Kingma and Jimmy Ba  “Adam: A Method for Stochastic Optimization”  arXiv, 2014  DOI: 10.48550/ARXIV.1412.6980
  * [64] M. Denker, A. Yegenoglu and S. Grün  “Collaborative HPC-enabled workflows on the HBP Collaboratory using the Elephant framework”  In _Neuroinformatics 2018_ , 2018, pp. P19  DOI: 10.12751/incf.ni2018.0019

Supplementary Material for:
Interpretable statistical representations of neural population dynamics and
geometry

Adam Gosztolai,1∗ Robert L. Peach,2,3 Alexis Arnaudon,4
Mauricio Barahona,4 Pierre Vandergheynst1

1Signal Processing Laboratory (LTS2), EPFL, Lausanne, 1016, Switzerland
2Department of Neurology, University Hospital Würzburg, Würzburg, 97070,
Germany
3Department of Brain Sciences, Imperial College London, London, W12 0NN,
United Kingdom,
4Department of Mathematics, Imperial College London, London, SW7 2AZ, United
Kingdom

∗To whom correspondence should be addressed; E-mail<EMAIL_ADDRESS>
##  S1 The MARBLE method

The MARBLE method aims to find a similarity-preserving map from the local flow
fields (LFFs) to a shared latent space. Since the method processes vector
fields point-by-point, it suffices to describe it applied on a single manifold
and associated dynamical flows obtained from an ensemble of trials. The
generalisation of the method to the joint embedding of multiple vector fields
is immediate.

Given a smooth, compact $m$-dimensional submanifold $\mathcal{M}$ of
$\mathbb{R}^{d}$ described by a point cloud
$\mathbf{X}=(\mathbf{x}_{1},\dots,\mathbf{x}_{n})$, MARBLE aims at
representing the discrete vector field
$\mathbf{F}=\\{\mathbf{f}_{1},\dots,\mathbf{f}_{n}\\}$ supported over the
point cloud as an empirical distribution
$P(\mathbf{Z})=\sum_{i=0}^{n}\delta(\mathbf{z}_{i})$ of latent features
$\mathbf{Z}=(\mathbf{z}_{1},\dots,\mathbf{z}_{n})$. The vectors
$\\{\mathbf{f}_{i}\\}$ can be obtained from time-series data, for example, by
taking first-order finite differences
$\mathbf{f}_{i}:=\mathbf{x}(t_{i}+1)-\mathbf{x}(t_{i})$. We now introduce an
unsupervised geometric deep learning architecture for obtaining the latent
features $\\{\mathbf{z}_{i}\\}$ based on the vector field restricted to a
neighbourhood around each point $\\{\mathbf{x}_{i}\\}$.

###  S1.1 Approximating the manifold by a proximity graph

To define the locality on $\mathcal{M}$, we describe it with a proximity
graph. We must ensure that the graph is as homogeneous as possible to obtain a
faithful representation of the tangent space of $\mathcal{M}$ via the vector
field $\mathbf{F}$. In particular, consecutive temporal points over dynamical
trajectories should not be overrepresented in neighbourhoods of any point
relative to points in close trajectories. For this reason, we generate the
proximity graph from a subsample of points obtained by farthest point sampling
[55], with a criterion $\delta\in[0,1]$ that controls the spacing of the
points relative to the diameter of the manifold
$\max_{ij}(||\mathbf{x}_{i}-\mathbf{x}_{j}||_{2})<\delta\,\text{diam}(\mathcal{M})$.

Several proximity graph algorithms can be used, such as the $k$-NN algorithm
and the $\epsilon$-ball algorithm, but we found that the continuous
$k$-nearest neighbour ($ck$-NN) algorithm [56] creates more representative
neighbourhoods. Indeed, contrary to the classical $k$-NN graph algorithm, it
can be interpreted as a local kernel density estimate and accounts for
sampling density variations over $\mathcal{M}$. The $ck$-NN algorithm extends
the classical $k$-NN algorithm by accounting for density variations by
connecting $i$ and $j$ whenever
$||\mathbf{x}_{i}-\mathbf{x}_{j}||_{2}^{2}<\delta||\mathbf{x}_{i}-\mathbf{x}_{u}||_{2}\,||\mathbf{x}_{j}-\mathbf{x}_{v}||_{2}$,
where $\delta>0$ is a scaling parameter, $u,v$ are the $k$-th nearest
neighbours of $i,j$, respectively, and $||\cdot||_{2}$ is the Euclidean norm.
This proximity graph endows $\mathcal{M}$ with a geodesic structure, i.e., for
any $i,j\in\mathcal{M}$, there is a shortest path with distance $d(i,j)$. We
can then define the local vector fields as the $p$-hop geodesic neighbourhood
$\mathcal{N}(i,p)$ around each point $i$.

###  S1.2 Parametrising the tangent spaces

The tangent spaces of $\mathcal{M}$ do not have a preferred coordinate system.
However, being isomorphic to $\mathbb{R}^{m}$, $\mathcal{T}_{i}\mathcal{M}$
can be parametrised by $m$ orthogonal vectors in the ambient space
$\mathbb{R}^{d}$ to form a local frame or gauge. Specifically, we assume that
the tangent space at a node $i$, $\mathcal{T}_{i}\mathcal{M}$, is spanned by
the edge vectors $\mathbf{e}_{ij}\in\mathbb{R}^{d}$ pointing from $i$ to $K$
nodes $j$ in its neighbourhood on the proximity graph. The $m$ largest
singular values $\mathbf{t}^{(\cdot)}_{i}\in\mathbb{R}^{d}$ of the matrix
formed by column-stacking $\mathbf{e}_{ij}$ yield the orthonormal coordinate
frame $\mathbb{T}_{i}\in\mathbb{R}^{d\times
m}=(\mathbf{t}_{i}^{(1)},\dots\mathbf{t}_{i}^{(m)})$ spanning
$\mathcal{T}_{i}\mathcal{M}$. In practice, we pick $K>deg(i)$ closest nodes to
$i$ on the proximity graph where $K$ is a hyperparameter. Larger $K$ increases
the overlaps between the nearby tangent spaces, and we find that
$K=2|\mathcal{N}(i,1)|$ is often a good compromise between locality and
robustness to noise of the tangent space approximation. Note that because
$\mathbb{T}_{i}$ defines an orthogonal basis,
$\mathbb{T}^{T}_{i}\mathbf{f}_{i}$ acts as a projection of the signal to the
tangent space in the $\ell_{2}$ sense. We perform these computations using a
modified Parallel Transport Unfolding package [57]. We illustrate the computed
frames on a spherical manifold (Fig. S1).

###  S1.3 Connections between tangent spaces

Having the local frames, we next define the parallel transport map
$\mathcal{P}_{j\to i}$ aligning the local frame at $j$ to that at $i$, which
is necessary to define convolution operations in a common space. While
parallel transport is generally path dependent, we assume that adjacent nodes
$i,j$ are close enough to consider the unique smallest rotation, known as the
Lévy-Civita connection. Thus, for adjacent edges, $\mathcal{P}_{j\to i}$ can
be computed as the matrix $\mathbf{O}_{ji}$ corresponding to
$\mathcal{P}_{j\to i}$, as the orthogonal transformation (rotation and
reflection)

$\mathbf{O}_{ji}=\arg\min_{\mathbf{O}\in
O(m)}||\mathbb{T}_{i}-\mathbb{T}_{j}\mathbf{O}||_{F},$ (S1)

where $||\cdot||_{F}$ is the Frobenius norm. The unique solution to this
problem is found by the Kabsch algorithm [58].

Note that $\mathbb{T}_{i}$ is defined only up to an orthogonal transformation
(rotation and reflection) within the tangent space of $\mathcal{M}$ because
the $m$-dimensional $\mathcal{T}_{i}\mathcal{M}$ only constrains $m$
coordinates of the frame. However, when the signal is projected to the local
frame, the tangent frame alignment by $\mathcal{P}_{j\to i}=\mathbf{O}_{ij}$
removes this ambiguity. Indeed, suppose that each node carries the same signal
$\mathbf{f}$, then, parallel transport alignment of the projected signal from
$j$ to $i$ yields

$\mathbb{T}_{i}^{T}\mathbf{f}=\mathbf{O}_{ij}\mathbb{T}_{j}^{T}\mathbf{f}=(\mathbb{T}_{j}\mathbf{O}_{ji})^{T}\mathbf{f}=\mathbb{T}_{i}^{T}\mathbf{f}\,,$
(S2)

where the first equality used the definition of the parallel transport, the
second equality used the transpose operation, while the third equality used
Eq. S1. Note that the same result does not hold when parallel transporting
signals in the ambient space (without projection) because, in that case, the
ambiguity in the frame orientation introduces ambiguity in the signal

###  S1.4 Vector diffusion

Before constructing the vector field features, we use a vector diffusion layer
to smooth out noisy vector field samples. Vector diffusion is a generalisation
of the scalar (heat) diffusion and can be expressed as a kernel associated
with the vector diffusion equation [59] to produce a smoothed vector field

$\text{vec}(\mathbf{F}(\tau))=e^{-\tau\mathcal{L}}\text{vec}(\mathbf{F})\,,$
(S3)

where $\text{vec}(\mathbf{F})\in\mathbb{R}^{nm\times 1}$ the row-wise
concatenation of vector-valued signals, $\tau$ is a learnable parameter that
controls the scale of the local vector fields and $\mathcal{L}$ is the random-
walk normalised connection Laplacian defined as a block matrix whose nonzero
blocks are given by

$\mathcal{L}(i,j)=\begin{cases}\mathbf{I}_{m\times m}\;&\text{for}\;i=j\\\
-\text{deg}(i)^{-1}\mathbf{O}_{ij}&\text{for}\;j\in\mathcal{N}(i,1)\,.\end{cases}$
(S4)

See [60] for further details. Intuitively, rather than diffusion of the
vectors component-wise, vector diffusion smoothens vectors based on
differences between a vector at a given node and other vectors parallel
transported to this node.

###  S1.5 Gradient filters

Before we can learn the vector field features, we define gradient filters,
whose role is to approximate the variation of the vector field around a point
in the local tangent frame. To numerically compute the gradient at a point
$i\in\mathcal{M}$ we construct an anisotropic filter by extending the
directional derivative filter of [36] using local coordinates. To do so,
consider the local frame $\mathbb{T}_{i}$ and construct the directional
derivative filter in the direction of the $q$-th unit vector
$\mathbf{t}_{i}^{(q)}$. We follow [36] and decompose
$\mathbf{t}_{i}^{(q)}\in\mathbb{R}^{d\times 1}$ by projecting it to the set of
edge vectors $\mathbf{e}_{ij}$ to obtain a vector
$\widehat{\mathbf{t}}_{i}^{(q)}\in\mathbb{R}^{n\times 1}$ at node $i$ pointing
in the basis direction $q$ with components $j$

$\widehat{t}^{(q)}_{i}(j)=\begin{cases}\langle\mathbf{t}_{i}^{(q)},\mathbf{e}_{ij}\rangle/deg(i)\,&\text{if}\,j\in\mathcal{N}(i,1)\\\
0&\text{otherwise}\,.\end{cases}$ (S5)

Collating for all nodes, we can write a matrix for the $q$-th coordinate of
the local frame projected onto the edge vectors
$\widehat{\mathbb{T}}_{q}=(\widehat{\mathbf{t}}_{1}^{(q)},\dots,\widehat{\mathbf{t}}_{n}^{(q)})\in\mathbb{R}^{n\times
n}$. Using this decomposition, the directional derivative of the scalar field
$s_{i}$ at $i$ in the direction $\mathbf{t}_{i}^{q}$ is given by

$\mathcal{K}^{(i,q)}s_{i}:=\langle\nabla
s_{i},\widehat{\mathbf{t}}_{i}^{(q)}\rangle=\sum_{j\in\mathcal{N}(i,1)}\left(s_{j}-s_{i}\right)\widehat{t}_{i}^{(q)}(j)\,,$
(S6)

or, in matrix form, by

$\mathbf{K}^{(q)}\mathbf{s}=(\widehat{\mathbb{T}}_{q}-\text{diag}(\widehat{\mathbb{T}}_{q}\,\mathbf{1}_{n}))\mathbf{s}\,,$
(S7)

where $\mathbf{1}_{n}$ is the $n\times 1$ vector of ones. As a result, the
gradient of a scalar field can be obtained by column-wise concatenating (as
new channels) the derivatives against all directions in the basis set

$\nabla\mathbf{s}=(\mathcal{D}^{(1)}\mathbf{s},\dots,\mathcal{D}^{(d)}\mathbf{s})\,.$
(S8)

Figs. S2a,b show the output of the first and second-order filters applied to a
linear and a quadratic scalar field.

To compute the component-wise directional derivative for a vector field
$\mathbf{F}\in\mathbb{R}^{n\times m}$, one must first parallel transport the
local frames at the neighbours $j$ to $i$ before applying the anisotropic
filters in Eq. S7. Let $\mathbf{O}$ denote the $nm\times nm$ block matrix of
$m\times m$ blocks given by the connection matrices $\mathbf{O}_{ij}$. Then,
we may express Eq. S12 in matrix form as

$\mathcal{D}^{(q)}\mathbf{F}=((\mathbf{K}^{(q)}\otimes\mathbf{1}_{m}^{T}\mathbf{1}_{m})\odot\mathbf{O})\mathbf{F}\,.$
(S9)

Here the kronecker product in the inner brackets expands $\mathbf{K}^{(q)}$ to
the $nm\times nm$ block matrix where the $(i,j)$ $m\times m$ block is filled
with entries $K^{(q)}_{ij}$.

###  S1.6 Approximating local vector fields

We now define convolution kernels on $\mathcal{M}$ that act on the vector
field to represent the vector field around a given point. We first project the
vector signal to the manifold
$\mathbf{f}^{\prime}_{i}=\mathbb{T}^{T}_{i}\mathbf{f}_{i}$. This reduces the
dimension of $\mathbf{f}_{i}$ from $d$ to $m$ without losing information since
$\mathbf{f}_{i}$ was already in the tangent space. We drop the bar in the
sequel to understand that all vectors are expressed in local coordinates. In
this local frame, the best polynomial approximation of the vector field around
$i$ is given by the Taylor-series expansion of each component $f_{i,l}$ of
$\mathbf{f}_{i}$

$f_{j,l}\approx f_{i,l}+\nabla
f_{i,l}(\mathbf{x}_{j}-\mathbf{x}_{i})+\frac{1}{2}(\mathbf{x}_{j}-\mathbf{x}_{i})^{T}\nabla^{2}f_{i,l}(\mathbf{x}_{j}-\mathbf{x}_{i})+\dots\,.$
(S10)

Motivated by the Taylor approximation, we construct gradient operators of
increasing order and implement them as a set of $m$ anisotropic filters
$\\{\mathcal{D}^{(q)}\\}$ acting along $\\{\mathbf{t}_{i}^{(q)}\\}$,

$\nabla
f_{i,l}\approx\left(\mathcal{D}^{(1)}(f_{i,l}),\dots,\mathcal{D}^{(m)}(f_{i,l})\right)^{T}\,.$
(S11)

Here, $\mathcal{D}^{(q)}(f_{i,l})$ is the $l$-th component of

$\mathcal{D}^{(q)}(\mathbf{f}_{i})=\sum_{j=1}^{n}\mathcal{K}^{(i,q)}_{j}\mathcal{P}_{j\to
i}(\mathbf{f}_{j})\,,$ (S12)

where $\mathcal{P}_{j\to i}=\mathbf{O}_{ij}$ is the parallel transport
operator that takes the vector $\mathbf{f}_{j}$ from the adjacent frame $j$ to
a common frame at $i$. $\mathcal{K}^{(i,q)}\in\mathbb{R}^{n\times n}$ is a
directional derivative filter [36] (Eq. S7) expressed in local coordinates at
$i$ and acting along $\mathbf{t}_{i}^{(q)}$. As a result of the parallel
transport, the value of Eq. S12 is independent of the local curvature of the
manifold. The $p$-th order gradients can be defined by the iterated
application of the filters, which aggregates information in the $p$-hop
neighbourhood of points. Although we find that increasing the order of the
differential operators increases the expressiveness of the network (Fig. S3),
second-order filters ($p=2$) were sufficient for the application considered in
this paper.

The expansion in Eq. S10 suggests augmenting the vectors $\mathbf{f}_{i}$ by
the derivatives (Eq. S11), to obtain a matrix

$\mathbf{f}_{i}\mapsto\mathbf{f}^{\mathcal{D}}_{i}=\left(\mathbf{f}_{i},\nabla
f_{i,1},\dots,\nabla f_{i,m},\nabla(\nabla f_{i,1})_{1},\dots,\nabla(\nabla
f_{i,m})_{m}\right)\,,$ (S13)

of dimensions $m\times c$ whose columns are gradients of signal components up
to order $p$ to give a total of $c=(1-m^{p+1})/(m(1-m))$ vectorial channels.

###  S1.7 Inner product for geometry invariance

In the optional geometry-agnostic mode, deformations on the manifold have the
effect of introducing rotations into the local vector fields. Thus, we can
achieve invariance to these deformations by making the learnt features
rotation invariant. We do so by first transforming the $m\times c$ matrix
$\mathbf{f}^{\mathcal{D}}_{i}$ to a $1\times c$ vector as

$\mathbf{f}^{\mathcal{D}}_{i}\mapsto\mathbf{f}^{\mathsf{ip}}_{i}=\left(\mathcal{E}^{(1)}(\mathbf{f}^{\mathcal{D}}_{i}),\dots,\mathcal{E}^{(c)}(\mathbf{f}^{\mathcal{D}}_{i})\right)\,.$
(S14)

Then, by taking for each channel the inner product against all other channels,
weighted by a dense learnable matrix $\mathbf{A}^{(r)}\in\mathbb{R}^{m\times
m}$ and summing, we obtain

$\mathcal{E}^{(r)}(\mathbf{f}^{\mathcal{D}}_{i})=\mathcal{E}^{(r)}(\mathbf{f}^{\mathcal{D}}_{i};\,\mathbf{A}^{(r)}):=\sum_{s=1}^{c}\left\langle\mathbf{f}^{\mathcal{D}}_{i}(\cdot,r),\mathbf{A}^{(r)}\mathbf{f}^{\mathcal{D}}_{i}(\cdot,s)\right\rangle\,,$
(S15)

for $r=1,\dots,c$ (Fig. 1f). Taking inner products is valid because the
columns of $\mathbf{f}^{\mathcal{D}}_{i}$ all live in the tangent space at
$i$. Intuitively, Eq. S15 achieves coordinate independence by learning
rotation and scaling relationships between pairs of channels.

###  S1.8 Embedding with a multilayer perceptron (MLP)

To embed each local feature, $\mathbf{f}^{\mathsf{ip}}_{i}$ or
$\mathbf{f}^{\mathcal{D}}_{i}$, depending on if inner product features are
used, (Eq. S14) we use a multilayer perception (MLP) (Fig. 1g)

$\mathbf{z}_{i}=\text{MLP}(\mathbf{f}^{\mathsf{ip}}_{i};\,\omega)\,,$ (S16)

where $\omega$ are trainable weights. The MLP is composed of $L$ linear
(fully-connected) layers interspersed by ReLU non-linearities. We used $L=2$
with a sufficiently high output dimension to encode the variables of interest.
The parameters were initialised using the Kaiming method [61].

###  S1.9 Loss function

Unsupervised training of the network is possible due to the continuity in the
vector field over $\mathcal{M}$, which causes nearby local vector fields to be
more similar than distant ones. We implement this via negative sampling [38],
which uses random walks sampled at each node to embed neighbouring points on
the manifold close together while pushing points sampled uniformly at random
far away. We use the following unsupervised loss function [38]

$\mathcal{J}(\mathbf{Z})=-\log(\sigma(\mathbf{z}^{T}_{i}\mathbf{z}_{j}))-Q\mathbb{E}_{k\sim
U(n)}\log(\sigma(-\mathbf{z}_{i}^{T}\mathbf{z}_{k}))\,,$ (S17)

where $\sigma(x)=(1+e^{-x})^{-1}$ is the sigmoid function and $U(n)$ is the
uniform distribution over the $n$ nodes. To compute this function, we sample
one-step random walks from every node $i$ to obtain ’positive’ node samples
for which we expect similar local vector fields to that of node $i$. The first
term in Eq. S17 seeks to embed these nodes close together. At the same time,
we also sample nodes uniformly at random to obtain ’negative’ node samples
with likely different local vector fields from that of node $i$. The second
term in Eq. S17 seeks to embed these nodes far away. We also choose $Q=1$.

We optimise the loss Eq. S17 by stochastic gradient descent. For training, the
nodes from all manifolds were randomly split into training ($80\%$),
validation ($10\%$) and test ($10\%$) sets. The optimiser was run until
convergence of the validation set and the final results were tested on the
test set with the optimised parameters.

###  S1.10 Pseudo code of MARBLE algorithm

We implemented MARBLE architecture with Pytorch Geometric [62]. The general
algorithm is as follows.

Algorithm 1 MARBLE

$d$-dimensional vector field samples
$\mathbf{F}=(\mathbf{f}_{1},\dots,\mathbf{f}_{n})$ connection Laplacian
$\mathcal{L}$ derivative filters $\mathcal{D}_{i}^{(q)}$ for
$i\in\\{1,\dots,n\\}$ and $q\in\\{1,\dots,m\\}$ derivative order $p$

Embedding $\mathbf{z}_{i}$ for all $i\in\\{1,\dots,n\\}$ $\mathbf{F}\leftarrow
e^{\tau\mathcal{L}}\mathbf{F}$ $\triangleright$ Apply diffusion layer
(optional) $\mathbf{h}^{(0)}\leftarrow\mathbf{f}_{i}$

for $1\leq l\leq p$ do $\triangleright$ Loop over filter orders $\nabla
h_{q}^{(l)}=\left(\mathcal{D}^{(1)}(h_{q}^{(l)}),\dots,\mathcal{D}^{(m)}(h_{q}^{(l)})\right)^{T}$
$\triangleright$ Compute filters
$\mathbf{h}^{(l)}\leftarrow\text{concat}\left(\mathbf{h}^{(l-1)},\nabla
h_{1}^{(l)},\dots,\nabla h_{m}^{(l)}\right)$ $\triangleright$ Concatenate
derivatives

end for
$\mathbf{h}^{(l)}\leftarrow\left(\mathcal{E}_{1}(\mathbf{h}^{(l)};\mathbf{A}_{1}),\dots,\mathcal{E}_{c}(\mathbf{h}^{(l)};\mathbf{A}_{c})\right)$
$\triangleright$ Inner product features (optional)
$\mathbf{z_{i}}\leftarrow\text{MLP}(\mathbf{h}^{(l)};\,\omega)$
$\triangleright$ Pass through MLP

##  S2 Distributional distance between latent representations

To test whether shifts in the statistical representation of the dynamical
system can predict global phenomena in the dynamics we define a similarity
metric between pairs of vector fields $\mathbf{F}_{1},\mathbf{F}_{2}$ with
respect to their corresponding embeddings
$\mathbf{Z}_{1}=(\mathbf{z}_{1,1},\dots,\mathbf{z}_{n_{1},1})$ and
$\mathbf{Z}_{2}=(\mathbf{z}_{1,1},\dots,\mathbf{z}_{n_{2},1})$. We use the
optimal transport distance between the empirical distributions
$P_{1}=\sum_{i}^{n_{1}}\delta(\mathbf{z}_{i,1}),P_{2}=\sum_{i}^{n_{2}}\delta(\mathbf{z}_{i,2})$

$d(P_{1},P_{2})=\min_{\gamma}\sum_{uv}\gamma_{uv}||\mathbf{z}_{u,1}-\mathbf{z}_{v,2}||^{2}_{2}\,,$
(S18)

where $\gamma$ is the transport plan, a joint probability distribution subject
to marginality constraints that $\sum_{u}\gamma_{uv}=P_{2}$,
$\sum_{v}\gamma_{uv}=P_{1}$ and $||\cdot||_{2}$ is the Euclidean distance.

##  S3 Examples

We give here more details on the examples of the main text.

###  S3.1 Van der Pol oscillator on paraboloid

We used the following equations to simulate the Van der Pol system:

$\displaystyle\dot{x}=y$ (S19) $\displaystyle\dot{y}=\mu(1-x^{2})y-x\,,$

parametrised by $\mu$. If $\mu=0$, the system reduces to the harmonic
oscillator, if $\mu<0$, the system is unstable and if $\mu>0$, the system is
stable and converges to a limit cycle. In addition, we map this two-
dimensional system to a paraboloid as with the map

$\displaystyle x,y$ $\displaystyle\mapsto x,y,z=\mathrm{parab}(x,y)$
$\displaystyle\dot{x},\dot{y}$
$\displaystyle\mapsto\dot{x},\dot{y},\dot{z}=\mathrm{parab}(x+\dot{x},y+\dot{y})-\mathrm{parab}(x,y)\,,$

where $\mathrm{parab}(x,y)=-(\alpha x)^{2}-(\alpha y)^{2}$.

We sought to test whether we can distinguish on-manifold dynamical variation
due to variation in the damping parameter $\mu$ while being invariant to the
shape of the manifold. As conditions, we take different values of the damping
parameter $\mu$, which for negative values produces unstable oscillations and
stabilises as $\mu$ increases above zero in a Hopf bifurcation (Fig. 1d). In
addition, increasing $|\mu|$ causes continuous deformations of the limit cycle
from circular to asymmetric, corresponding to slow-fast dynamics. In Fig. S6,
we illustrate other examples of the application of MARBLE on this system, with
different regimes of $\mu$, and randomness in the paraboloid curvature across
$\mu$, demonstrating the difference between geometry-aware and geometry-
agnostic modes of MARBLE. We repeated training five times and confirmed that
the results were reproducible.

###  S3.2 Low-rank RNNs

We consider low-rank RNNs composed of $N=500$ rate units in which the
activation of the $i$-th unit is given by

$\tau\frac{dx_{i}}{dt}=-x_{i}+\sum_{j=1}^{N}J_{ij}\phi(x_{j})+\tilde{u}_{i}(t)+\eta_{i}(t),\quad
x_{i}(0)=0\,,$ (S20)

where $\tau=100$ ms is a time constant, $\phi(x_{i})=\tanh(x_{i})$ is the
firing rate, $J_{ij}$ is the rank-R connectivity matrix, $u_{i}(t)$ is an
input stimulus and $\eta_{i}(t)$ is a white noise process with zero mean and
std $3\times 10^{-2}$. The connectivity matrix can be expressed as

$\mathbf{J}=\frac{1}{N}\sum_{r=1}^{R}\mathbf{m}_{r}\mathbf{n}_{r}^{T}\,,$
(S21)

for vector pairs $(\mathbf{m}_{r},\mathbf{n}_{r})$. For the delayed-match-to-
sample task, the input is of the form

$\tilde{u}_{i}(t)=w^{\mathsf{in}}_{1i}u_{1}(t)+w^{\mathsf{in}}_{2i}u_{2}(t)\,,$
(S22)

where $w_{1i},w_{2i}$ are coefficients controlling the weight of inputs
$u_{1},u_{2}$ into node $i$. Finally, the network firing rates are read out to
the output as

$o(t)=\sum_{i=0}^{N}w^{\mathsf{out}}_{i}\phi(x_{i})\,.$ (S23)

To train the networks, we followed the procedure in [46]. The experiments
consisted of $5$ epochs; a fixation period of length between $100-500$ ms
chosen uniformly at random, a $500$ ms stimulation period, a delay period of
length between $500-3000$ ms, a $500$ ms stimulation period and a $1000$ ms
decision period. During training, the networks were subjected to the two
inputs, whose value varied discontinuously between zero and a non-zero gain
during stimulation. The networks were trained against a loss function

$\mathcal{L}=|o(T)-\hat{o}(T)|\,,$ (S24)

where $T$ is the length of the trial and $\hat{o}(T)=1$ when both stimuli were
present and $-1$ otherwise. Coefficient vectors were initially drawn from
zero-mean and unit std Gaussian distributions and then optimised. For
training, we used the ADAM optimiser [63] with moment decay rates $0.9-0.999$
and learning rates $10^{-3}-10^{-2}$. We trained MARBLE in geometry-agnostic
mode. See Table 1 for the parameters. We repeated training five times and
confirmed that the results were reproducible.

###  S3.3 Macaque hand reaching data

We used publically available [27] electrophysiology data during a centre-out
instructed-delay reaching task. The neural activity was recorded using linear
multielectrode arrays (V-Probe, 24-channel linear probes) from rhesus macaque
motor (M1) and dorsal premotor (PMd) cortices. See [27] for further details on
experimental procedures. Each trial began with the monkey’s hand at the centre
position, after which one (or more) radial targets at $10$cm from the centre
position were marked visually (target onset). After a variable delay period,
one of the radial targets was highlighted, indicating the go-cue for reaching.
We analysed the $700$ms period after the go-cue consisting of a delay followed
by the reach. A total of $44$ consecutive experimental sessions with a
variable number of trials were considered.

For each of the $24$ channels of each trial, we extracted the spike trains
using the neo package in Python (http://neuralensemble.org/neo/) and converted
them into instantaneous rates using Gaussian kernel convolution with a
standard deviation of $100$ms. We then binned the rates at $20$ms intervals
using the elephant Python package [64] to match the sampling frequency in the
decoded kinematics in [27]. Finally, we fitted a PCA across all trials in a
single session to reduce the dimension of the $24$-channel rates to five
principal components. We trained MARBLE in geometry-aware mode (without inner
product features) separately on each session, treating each of the seven
movement conditions as individual manifolds. See Table 1 for parameters. We
repeated training five times and confirmed that the results were reproducible.

For benchmarking, we took the trained LFADS models directly from the authors.
To train CEBRA, we used the seven reach directions as labels and trained CEBRA
supervised until convergence. We obtained the best results with an initial
learning rate of $0.01$, Euclidean norm as metric, number of iterations
$10,000$ and fixed temperature $1$.

We followed the same procedure as in [27] to decode the hand kinematics from
the neural trajectories. Specifically, we used Optimal Linear Estimation (OLE)
to decode the $x$ and $y$ reaching coordinates and velocities from the neural
embeddings. For each session, using $10$-fold cross-validation, we fitted an
OLE to map from the MARBLE embeddings to the kinematics. To assess the
accuracy of decoded movements, we computed the goodness of fit ($R^{2}$)
between the decoded and real velocities for both $x$ and $y$ before taking the
mean across them. We also trained a support vector machine classifier
(regularisation of $1.0$ with a radial basis function) on the real kinematic
movements against the known condition labels. We then evaluated the classifier
on the decoded kinematics and reported the accuracy for each session.

###  S3.4 Rat linear maze task

We used publically available [50] electrophysiology data during a 1.6m linear
track traversal both ’leftwards’ and ’rightwards’. The neural activity was
recorded with two 8- or 6-shank silicon probes in well-isolated CA1 pyramidal
single units in the hippocampus across four rats. See [50] for further details
on experimental procedures. Each rat had a single recording involving multiple
’leftward’ and ’rightward’ movements along the linear maze. Each recording had
a variable number of measured neurons, ranging from 48 neurons to 120 neurons.

For each of the animals, we extracted the spike trains using the neo package
in Python (http://neuralensemble.org/neo/) and converted them into
instantaneous rates using Gaussian kernel convolution with a standard
deviation of $10$ms. We then fitted a PCA across all data for a given animal
to reduce the variable dimension size to 10 principle components. We trained
MARBLE in geometry-aware mode (without inner product features) separately on
each animal, including the original position. See Table 1 for parameters.

For benchmarking, we took the CEBRA models unchanged from the publicly
available notebooks in the online documentation associated with their
manuscript. For decoding position, we used a 32-dim output and 10000
iterations as per their demo notebook
(https://cebra.ai/docs/demo_notebooks/Demo_decoding.html), whilst for
consistency between animals we used a 3-dim output and 15000 iterations also
in line with their demo notebook
(https://cebra.ai/docs/demo_notebooks/Demo_consistency.html). The same output
dimensions were used for MARBLE on each task, respectively.

We followed the procedure to decode the rat position along the linear maze
from the neural trajectories as in the CEBRA [31]. Specifically, we fit a KNN
Decoder with 36 neighbours and a cosine distance metric using the neural
embedding to predict the $x$ position along the track. The data was split such
that the initial 80% of the data was used for training, and the remaining 20%
was used for testing. To assess the accuracy of decoding, we computed the mean
absolute error between the predicted and true rat positions at each time
point.

To assess the consistency of the embeddings across animals, we used the
inbuilt functions from CEBRA for fair comparison. The embeddings are first
aligned optimally using Procrustes analysis. Secondly, a linear model is
fitted using the aligned embeddings of any two animals: one animal as the
source (independent variables) and another as the target (dependent
variables). The $R^{2}$ from the fitted model describes the amount of variance
in the embedding of one animal that can be explained by another animal, i.e.,
a measure of consistency between their embeddings.

Table 1: Parameters used for MARBLE embedding in the different experiments | Dataset
---|---
Parameters | van der Pol | RNN | Macaque | Rat
Inner product features | True | False/True | False | False
$k$ | 20 | 15 | 30 | 15
Diffusion | False | False | True | False
Feature order, $p$ | 2 | 2 | 2 | 1
MLP hidden channels | 32 | 32 | 100 | 64
MLP output channels | 5 | 3 | 3/20 | 3/32
Number of parameters | 731 | 2886/1339 | 15803 | 7940/9825

##  S4 Supplementary figures

Figure S1: Illustration of local coordinate frames. a,b Local coordinate
frames were fitted to eight neighbours at each point on the rectangular grid
over a sphere (manifold of dimension two) embedded into $\mathbb{R}^{3}$. Unit
vectors representing within-manifold basis. c Unit vector represents normal to
the manifold. Note that the orientation of the normal vectors is not
necessarily consistent. Figure S2: Output of gradient filters. a Scalar valued
linear function. b Scalar valued parabolic function. The first column shows
the output of the scalar signal convolved with the gradient filter with
respect to a principal spatial coordinate, representing the directional
(partial) derivatives along this direction. The second and third columns show
second-order mixed partial derivatives obtained by a second application of the
gradient filter to the derivative signal. In each case, we used a uniform
rectangular grid with eight neighbours. Figure S3: Effect of filter order. a
Scalar functions sampled ($n=512$) uniformly at random and fitted with a
continuous $k$-nearest neighbour graph ($k=15$). From left to right: constant,
linear, parabola, saddle. b Joint embedding of local scalar fields from all
functions based on first-order ($1$-hop) directional derivative filters. Dots
represent nodes drawn from a with nodes close together, signifying similar
signal distributions in their neighbourhoods. Black lines show the convex
polygon obtained from $k$-means clustering ($10$ clusters). c As in b, but
with first ($1$-hop) and second-order ($2$-hop) filters. Including second-
order filters increases the clustering of features. d Histogram representation
of the clustered neighbourhood types. The latter is shown at the bottom in
circular insets. e As in d but with first and second-order filters. Second-
order features better discriminate the parabola and saddle but show little
difference for constant and linear fields. Figure S4: Effect of rotational
invariant filters. a Vector fields sampled ($n=512$) uniformly at random in
the interval $[-1,1]^{2}$ and fitted with a continuous $k$-nearest neighbour
graph. b Joint embedding of local vector fields based on first-order ($1$-hop)
directional derivative filters. Dots represent nodes drawn from a with nodes
close together signifying similar signal distributions in their
neighbourhoods. Features from the linear vector fields cluster together
(clusters $2$ and $9$) while those drawn from the vortex fields fall on
separate halves of a one-dimensional circular manifold corresponding to the
one-parameter (angle) variation between them. c Same as b but with rotation
invariant features. Features from linear fields can no longer be distinguished
(cluster $15$) because the filter does not learn the orientation. Features
from vortex fields fall on a linear one-dimensional manifold parametrised by
the distance from the centre. d The distribution of orientation-preserving
derivative features can distinguish all fields. e The distribution of
rotation-invariant features can discriminate linear fields from vortex fields
but not the orientation.  Figure S5: Effect of rotational invariant filters on
local vector field types. Types are local vector field types with circular
insets representing a local vector field drawn from each of the four fields in
Fig S4. Top left: linear left; top right: linear right; bottom left vortex
left; bottom right; vortex right. White insets indicate that the given type is
not present in the vector field. a Without using rotation invariant filters
the network tends to classify local vector fields into separate types. b With
rotation invariant filters, local vector fields cluster more irrespective of
the orientation. The separate types are instead distinguished more based on
the expansion, and rotation properties of the features.  Figure S6: Additional
Van der Pol examples We illustrate the MARBLE embedding of the Van der Pol
example with $40$ values of $\mu$ in the following cases. a Using a larger
parameter range $\mu\in[-1,1]$ increases the definition of the clusters. b
Using a smaller parameter range $\mu\in[-0.1,0.1]$, corresponding to the red
square in a shows lower definition. c In geometry-agnostic mode, varying the
curvature of the paraboloid $\beta(x^{2}+y^{2})$ by drawing $\beta$ uniformly
at random from $[-0.2,0.2]$ $-0.2$ to $0.2$) does not alter the MARBLE
embedding. d In geometry-aware mode, the same variation as in c destroys the
cluster structure.  Figure S7: Neural activities in a rank-2 RNN during the
delayed-match-to-sample task. a Schematic of a low-rank RNN, taking as input
two stimuli and producing a decision variable as the output. Arrow endings
represent inhibitory and excitatory connections. b Two example input patterns
for one of the stimuli, and corresponding output patterns. Red and grey-shaded
bands show stimulated and unstimulated periods. c Mean-field dynamics (heatmap
and stream plot) superimposed with a sampled trajectory (orange) during one
trial.  Figure S8: Latent representations obtained using point-wise methods on
the macaque reaching task, such as PCA, t-SNE and UMAP. These methods are
unsuitable to reveal the geometric structure of latent dynamics.  Figure S9:
Comparison of MARBLE and LFADS for learning neural dynamics for individual
sessions of macaque reaching task. a Feature embeddings obtained from MARBLE
for various example sessions. The 3D-embedded points better reflect the
arrangement of reaches in physical space when compared to LFADS. b The matrix
of distribution distances between pairwise conditions for separate sessions
shows a stronger periodic structure compared to LFADS. c MDS embedding of the
distance matrix consistently recovers the spatial arrangement of reaches
across sessions, when compared to LFADS. d Hand trajectories linearly decoded
from MARBLE embeddings showed much stronger spatial correspondence to ground-
truth kinematics than LFADS.