<EMAIL_ADDRESS>
# Inferring stochastic low-rank recurrent neural networks from neural data

Matthijs Pals1,2 &A Erdem Sağtekin1,2,3
Felix Pei1,2
&Manuel Gloeckler1,2
&Jakob H Macke1,2,4

1Machine Learning in Science, Excellence Cluster Machine Learning, University
of Tübingen, Germany
2Tübingen AI Center, Tübingen, Germany
3Graduate Training Centre of Neuroscience, University of Tübingen, Germany
4Department Empirical Inference, Max Planck Institute for Intelligent Systems,
Tübingen, Germany

###### Abstract

A central aim in computational neuroscience is to relate the activity of large
populations of neurons to an underlying dynamical system. Models of these
neural dynamics should ideally be both interpretable and fit the observed data
well. Low-rank recurrent neural networks (RNNs) exhibit such interpretability
by having tractable dynamics. However, it is unclear how to best fit low-rank
RNNs to data consisting of noisy observations of an underlying stochastic
system. Here, we propose to fit stochastic low-rank RNNs with variational
sequential Monte Carlo methods. We validate our method on several datasets
consisting of both continuous and spiking neural data, where we obtain lower
dimensional latent dynamics than current state of the art methods.
Additionally, for low-rank models with piecewise linear nonlinearities, we
show how to efficiently identify all fixed points in polynomial rather than
exponential cost in the number of units, making analysis of the inferred
dynamics tractable for large RNNs. Our method both elucidates the dynamical
systems underlying experimental recordings and provides a generative model
whose trajectories match observed trial-to-trial variability.

##  1 Introduction

A common goal of many scientific fields is to extract the dynamical systems
underlying noisy experimental observations. In particular, in neuroscience,
much work is devoted to understanding the coordinated firing of neurons as
being implemented through underlying dynamical systems [1, 2, 3, 4, 5].
Recurrent neural networks (RNNs) constitute a common model-class of neural
dynamics [6, 7, 8, 9, 10, 11, 12, 13] which can be reverse-engineered to form
hypotheses about neural computations [14, 15].

As a result, several recent research directions have centered on interpretable
or analytically tractable RNN architectures. In particular, RNNs with low-rank
structure [16, 17, 18, 19, 20, 21, 22] admit a direct mapping between high-
dimensional population activity and an underlying low-dimensional dynamical
system. RNNs with piecewise-linear activations [9, 23, 24, 8, 25, 26] have
fixed points and cycles that can be accessed analytically.

To serve as useful models of brain activity, it is important that models also
capture the observed brain activity, including trial-to-trial variability.
Many methods that fit RNNs to data are restricted to RNNs with deterministic
transitions [6, 7, 8, 12, 11, 10]. It is unlikely that, in general, all
variability in the data can be explained by variability in the RNNs initial
state. Thus, adopting stochastic transitions is imperative. While
probabilistic sequence models are used effectively in neuroscience [27], they
have so far largely consisted of state space models without an obvious
mechanistic interpretation [28, 29, 30, 31, 32].

Here, we demonstrate that we can fit large stochastic RNNs to noisy high-
dimensional data. First, we show that, by combining variational sequential
Monte Carlo methods [33, 34, 35] with low-rank RNNs, we can efficiently fit
stochastic RNNs with many units by learning the underlying low-dimensional
dynamical system. Second, we show that, for low-rank networks with piecewise-
linear activation functions, the resulting dynamics can be efficiently
analyzed: In particular, we show how _all_ fix points can be found with a
polynomial cost in the number of units – dramatically more efficient than the
exponential cost in the general case.

We first validate our method using several teacher-student setups and show
that we recover both the ground truth dynamics and stochasticity. We then fit
our model to several real-world datasets, spanning both spiking and continuous
data, where we obtain a generative model which needs lower dimensional latent
dynamics than current state of the art methods. We also demonstrate how in our
low-rank RNNs fixed points can be efficiently inferred — potentially at a
lower cost than approximate methods [25], while additionally coming with the
guarantee that _all_ fixed points are found.

Figure 1: Our goal is to obtain generative models from which we can sample
realistic neural data while having a tractable underlying dynamical system. We
achieve this by fitting stochastic low-rank RNNs with variational sequential
Monte Carlo.

##  2 Theory and methods

###  2.1 Low-rank RNNs

####  2.1.1 Access to the low-dimensional dynamics underlying large networks

Our goal is to infer recurrent neural network models of the form

$\displaystyle\tau\frac{d\mathbf{x}}{dt}=-\mathbf{x}(t)+\mathbf{J}\phi(\mathbf{x}(t))+\Gamma_{\mathbf{x}}\xi(t),$
(1)

with neuron activity $\mathbf{x}(t)\in\mathbb{R}^{N}$, time-constant
$\tau\in\mathbb{R}_{>0}$, recurrent weights $\mathbf{J}\in\mathbb{R}^{N\times
N}$, element-wise nonlinearity $\phi$, an $R$ dimensional white noise process
$\xi(t)$ and $\Gamma_{\mathbf{x}}\in\mathbb{R}^{N\times R}$. In particular, we
are interested in the case where the weight matrix $\mathbf{J}$ has rank
$R\leq N$, i.e. it can be written as $\mathbf{J}=\mathbf{MN}^{\mathsf{T}}$,
with $\mathbf{M},\mathbf{N}\in\mathbb{R}^{N\times R}$ ([18, 19, 20, 21]).
Assuming that $\mathbf{x}(0)$ lies in the subspace spanned by the columns of
$\mathbf{M}$ and $\Gamma_{\mathbf{x}}=\mathbf{M}\Gamma_{\mathbf{z}}$, with
$\Gamma_{\mathbf{z}}\in\mathbb{R}^{R\times R}$ , we can rewrite Eq. 1 as an
equivalent $R$ dimensional system,

$\displaystyle\tau\frac{d\mathbf{z}}{dt}=-\mathbf{z}(t)+\mathbf{N}^{\mathsf{T}}\phi(\mathbf{M}\mathbf{z}(t))+\Gamma_{\mathbf{z}}\xi(t),$
(2)

where we can switch between Eq. 1 and Eq. 2 by means of linear projection,
$\mathbf{z}(t)=(\mathbf{M^{\mathsf{T}}M})^{-1}\mathbf{M}^{\mathsf{T}}\mathbf{x}(t)$
and $\mathbf{x}(t)=\mathbf{Mz}(t)$. Note that we can directly extend these
equations to include input, representing, e.g., experimental stimuli or
context (see Supplement C.2).

####  2.1.2 Low-rank RNNs as state space models

We consider nonlinear latent dynamical systems with observations
$\mathbf{y}_{t}$:

$\displaystyle p(\mathbf{z}_{1:T},\mathbf{y}_{1:T})$
$\displaystyle=p(\mathbf{z}_{1})\prod_{t=2}^{T}p(\mathbf{z}_{t}|\mathbf{z}_{t-1})\prod_{t=1}^{T}p(\mathbf{y}_{t}|\mathbf{z}_{t}),$
$\displaystyle p(\mathbf{z}_{t}|\mathbf{z}_{t-1})$
$\displaystyle=\mathcal{N}(F(\mathbf{z}_{t-1}),\Sigma_{\mathbf{z}}),\;p(\mathbf{z}_{1})=\mathcal{N}(\mu_{\mathbf{z}_{1}},\Sigma_{\mathbf{z}_{1}}),$
$\displaystyle p(\mathbf{y}_{t}|\mathbf{z}_{t-1})$
$\displaystyle=G(\mathbf{z}_{t}),$

where the transition distribution is parameterised by discretising a low-rank
RNN with timestep $\Delta_{t}$ (see Supplement C.1), we have mean
$F(\mathbf{z}_{t})=a\mathbf{z}_{t}+\tilde{\mathbf{N}}^{\mathsf{T}}\phi(\mathbf{M}\mathbf{z}_{t})$,
with $a=1-\frac{\Delta_{t}}{\tau}$ and
$\tilde{\mathbf{N}}=\frac{\Delta_{t}}{\tau}\mathbf{N}$, and covariance
$\Sigma_{\mathbf{z}}$. The specific form of the observation function $G$,
depends on the data-modality, e.g., here we use a Poisson distribution for
count observations. This formulation allows one to keep the one-to-one
correspondence between RNN units (or a subset of those) and recorded data
neurons, (as was desired in e.g., [11, 12, 10, 36]). For example, assuming
Gaussian observation noise, we can simply use that
$\mathbf{x}_{t}=\mathbf{Mz}_{t}$ and define
$G=\mathcal{N}(\mathbf{M}\mathbf{z}_{t},\Sigma_{\mathbf{y}})$.

Once we learn $p(\mathbf{z}_{1:T},\mathbf{y}_{1:T})$, we can use the obtained
RNN as generative model to sample trajectories, and reverse engineer the
underlying dynamics to gain insight in the data generation process. Given the
sequential structure of the RNN, we can do model learning by using variational
sequential Monte Carlo (also called Particle Filtering) methods [33, 34, 35].

###  2.2 Model learning with variational sequential Monte Carlo

####  2.2.1 Sequential Monte Carlo

Sequential Monte Carlo (SMC) can be used to approximate sequences of
distributions, such as those generated by our RNN, with a set of $K$
trajectories of latents $\mathbf{z}_{1:T}$ (commonly called particles) [37]. A
crucial choice when doing SMC is picking the right proposal distribution $r$,
from which we can sample latents at a given timestep, conditioned on the
previous latent $\mathbf{z}_{t-1}$ and observed data $\mathbf{y}_{1:t}$, or a
subset of those. Given initial samples $\mathbf{z}_{1}^{1:K}\sim r$ and
corresponding importance weights $\overline{w}_{1}^{1:K}$ (as defined below)
SMC progresses by repeatedly executing the following steps:

$\displaystyle resample$ $\displaystyle
a^{k}_{t-1}\sim\mathsf{Discrete}(a^{k}_{t-1}|\overline{w}_{t-1}^{k}),$
$\displaystyle propose$ $\displaystyle\mathbf{z}^{k}_{t}\sim
r(\mathbf{z}^{k}_{t}|\mathbf{y}_{t},\mathbf{z}_{t-1}^{a^{k}_{t-1}}),$
$\displaystyle reweight$ $\displaystyle
w_{t}^{k}=\frac{p(\mathbf{y}_{t},\mathbf{z}_{t}^{k}|\mathbf{z}^{a_{t-1}^{k}}_{t-1})}{r(\mathbf{z}_{t}^{k}|\mathbf{y}_{t},\mathbf{z}^{a_{t-1}^{k}}_{t-1})},$

with $\overline{w}^{k}_{t}=\frac{w_{t}^{k}}{\sum_{j=1}^{K}w_{t}^{j}}$. We
obtain, at time $t$, a filtering approximation to the posterior,

$\displaystyle
q_{\mathsf{filt}}(\mathbf{z}_{1:t}|\mathbf{y}_{1:t})=\sum_{k=1}^{K}\overline{w}_{t}^{k}\delta(\mathbf{z}_{1:t}^{k}).$
(3)

The unnormalised weights give an unbiased estimate to the marginal likelihood,

$\displaystyle\hat{p}(\mathbf{y}_{1:T})=\prod_{t=1}^{T}\frac{1}{K}\sum_{k=1}^{K}w_{t}^{k}.$
(4)

We now detail how we pick the proposal distribution $r$. For linear Gaussian
observations $G=\mathcal{N}(\mathbf{Bz}_{t},\Sigma_{\mathbf{y}})$, we set
$r(\mathbf{z}_{t}|\mathbf{y}_{t},\mathbf{z}_{t-1})=p(\mathbf{z}_{t}|\mathbf{y}_{t},\mathbf{z}_{t-1})$,
as this is available in closed form and is optimal (in the sense that it
minimises the variance of the importance weights [37])

$\displaystyle
r(\mathbf{z}_{t}|\mathbf{y}_{t},\mathbf{z}_{t-1})=\mathcal{N}((\mathbf{I}-\mathbf{K}\mathbf{B})F(\mathbf{z}_{t-1})+\mathbf{Ky}_{t},(\mathbf{I}-\mathbf{K}\mathbf{B})\Sigma_{\mathbf{z}}),$
(5)

with $\mathbf{K}$ the Kalman Gain:
$\mathbf{K}=\Sigma_{\mathbf{z}}\mathbf{B}^{\mathsf{T}}(\mathbf{B}\Sigma_{\mathbf{z}}\mathbf{B}^{\mathsf{T}}+\Sigma_{\mathbf{y}})^{-1}$.
For non-linear observations, we can not invert the observation process in
closed form, so we instead jointly optimize a parameterized ‘encoding’
distribution $e(\mathbf{z}_{t}|\mathbf{y}_{t-t^{\prime}:t}$) (as in a
variational autoencoder [38]). In particular, we assume $e$ to be a
multivariate normal with diagonal covariance, which we parameterize by a
causal convolutional neural network, such that each latent is conditioned on
the $t^{\prime}$ latest observations. We then use the following proposal:

$\displaystyle
r(\mathbf{z}_{t}|\mathbf{z}_{t-1},\mathbf{y}_{t-t^{\prime}:t})\propto
e(\mathbf{z}_{t}|\mathbf{y}_{t-t^{\prime}:t})p(\mathbf{z}_{t}|\mathbf{z}_{t-1}),$
(6)

where we now also assume $p(\mathbf{z}_{t}|\mathbf{z}_{t-1})$ has a diagonal
covariance matrix.

####  2.2.2 Relationship to Generalised Teacher Forcing

In our approach, the mean of the proposal distribution at time $t$ is a linear
interpolation between the RNN predicted state $F(\mathbf{z}_{t-1})$ and a
data-inferred state $\hat{\mathbf{z}}_{t}$. A recent study obtained state-of-
the art results for reconstructing dynamical systems by fitting deterministic
RNNs with a method called Generalised Teacher Forcing (GTF), which also
linearly interpolates between a data-inferred and an RNN predicted state at
every time-step [8]; the model propagates forward in time as
$\mathbf{z}_{t}=(1-\alpha)F(\mathbf{z}_{t-1})+\alpha\hat{\mathbf{z}}_{t}$. [8]
showed that by choosing the appropriate $\alpha$, one can completely avoid
exploding gradients, while still allowing backpropagation through time, and
thus obtaining long-term stable solutions [39]. The optimal $\alpha$ can be
picked based on the maximum Lyaponuv exponent of the system (a measure of how
fast trajectories diverge in a chaotic system) [8].

By including the RNN in the proposal distribution, we similarly to GTF allow
backpropagation through time through the sampled trajectories. The
interpolation is given by
$\alpha=\Sigma_{\mathbf{z}}\mathbf{B}^{\mathsf{T}}(\mathbf{B}\Sigma_{\mathbf{z}}\mathbf{B}^{\mathsf{T}}+\Sigma_{\mathbf{y}})^{-1}\mathbf{B}$
in Eq. 5, and in Eq. 6 by
$\alpha=\Sigma_{\mathbf{z}}(\Sigma_{\mathbf{z}}+\Sigma_{\hat{\mathbf{z}}_{t}})^{-1}$,
where $\Sigma_{\hat{\mathbf{z}}_{t}}$ is the predicted variance of the
encoding network. Thus, instead of interpolating based on an estimate of how
chaotic the system is, our approach interpolates adaptively (every time step,
if Eq. 6 is used) based on how relatively noisy the transition distribution is
with respect to the data-inferred states at time $t$, analogous to, e.g., the
gain of a Kalman filter.

####  2.2.3 Variational objective

We can fit our RNNs to data by using SMC to specify a variational objective
[33, 34, 35]. In variational inference, we specify a family of parameterized
distributions $Q$, and optimize those parameters such that a divergence
(usually the $\mathsf{KL}$ divergence) between the variational distribution
$q(\mathbf{z}_{1:T})\in Q$ and the true posterior
$p(\mathbf{z}_{1:T}|\mathbf{y}_{1:T})$ is minimized. We do this by maximising
a lower bound ($\mathsf{ELBO}$) to the log likelihood $p(\mathbf{y}_{1:T})$.
In particular, we can use Eq. 4 to specify the $\mathsf{ELBO}$ objective[33,
34, 35]

$\displaystyle\mathcal{L}=\mathbb{E}_{q_{\mathsf{smc}}(\mathbf{z}^{1:K}_{1:T},a^{1:K}_{1:T-1}|\mathbf{y}_{1:T})}[\log\hat{p}(\mathbf{y}_{1:T})],$
(7)

with $q_{\mathsf{smc}}$ the sampling distribution:
$q_{\mathsf{smc}}(\mathbf{z}^{1:K}_{1:T},a^{1:K}_{1:T-1}|\mathbf{y}_{1:T})=\prod_{k=1}^{K}r(\mathbf{z}^{k}_{1}|\mathbf{y}_{1})\prod_{k=1}^{K}\prod_{t=2}^{T}r(\mathbf{z}^{k}_{t}|\mathbf{z}^{a^{k}_{t-1}}_{t-1}\mathbf{y}_{t})\mathsf{Discrete}(a^{k}_{t-1}|\overline{w}_{t-1}^{k}).$
We approximate this objective with Monte Carlo samples during training. As
suggested in [33, 34, 35, 40], we use biased gradients during optimization by
dropping high-variance terms arising from the resampling.

###  2.3 Finding fixed points in piecewise-linear low-rank RNNs

After having learned our model, we can gain insight into the mechanisms
underlying the data generation process by reverse engineering the learned
dynamics [15], e.g. by calculating their fixed points. Here, we show that the
fixed points can be found analytically and efficiently for low-rank networks
with piecewise-linear activation functions. This class of activation functions
$\phi(\mathbf{x}_{i})=\sum_{d}^{D}\mathbf{b}^{(i,d)}\mathsf{max}(\mathbf{x}_{i}-\mathbf{h}_{i}^{(d)},0)$
includes, e.g., the standard ReLU
($\phi(\mathbf{x}_{i})=\mathsf{max}(\mathbf{x}_{i}-\mathbf{h}_{i},0)$) or the
‘clipped’ variant
($\phi(\mathbf{x}_{i})=\mathsf{max}(\mathbf{x}_{i}+\mathbf{h}_{i},0)-\mathsf{max}(\mathbf{x}_{i},0)$)
[8] which we used in all experiments with real-world data here.

Naively, the cost of finding all fixed points piecewise-linear networks scales
_exponentially_ with the number of units in the networks: we would have to
solve $(D+1)^{N}$ systems of $N$ equations [9, 24]. If networks are low rank,
it is straightforward to show that we can reduce this cost to solving
$(D+1)^{N}$ systems of $R$ equations (See Supplement A.1). In addition,
however, we show that the computational cost can be greatly reduced further:
One can find _all_ fixed points in a cost that is _polynomial_ instead of
_exponential_ in the number of units:

Proposition 1. Assume Eq. 1, with $\mathbf{J}$ of rank $R$ and piecewise-
linear activations $\phi$. For fixed rank $R$ and fixed number of basis
functions $D$, we can find all fixed points in the absence of noise, that is
all $\mathbf{x}$ for which $\frac{d\mathbf{x}}{dt}=0$, by solving at most
$\mathcal{O}(N^{R})$ linear systems of $R$ equations.

Figure 2: Proof sketch.

Proof. See Supplement A.1.

Sketch. Assuming $D=1$, activations
$\phi=\max(0,\mathbf{x}_{i}-\mathbf{h}_{i})$; $N$ units will partition the
full phase space into $2^{N}$ regions in which the dynamics are linear (2
units, 4 regions in Fig. 2). We can thus, in principle, solve for all fixed
points by solving all corresponding linear systems of equations [9, 24]. If
dynamics are confined to the $R$-dimensional subspace spanned by the columns
of $\mathbf{M}$, only a subset of the linear regions (3 in Fig. 2) can be
reached. Each unit partitions the space spanned by the columns of $\mathbf{M}$
with a hyperplane (pink points in Fig. 2). The amount of linear regions in
$\mathbf{M}$, becomes equivalent to ‘how many regions can we create in
$R$-dimensional space with $N$ hyperplanes?’ Using Zaslavsky’s theorem [41],
we can show that this at most $\sum^{R}_{r=0}{N\choose
r}\in\mathcal{O}(N^{R}).$

##  3 Empirical Results

###  3.1 RNNs recover ground truth dynamics in student-teacher setups

Figure 3: RNNs recover dynamics in teacher-student setups. a) Example ground
truth latent trajectory and phase plane of low-rank RNN trained to oscillate
(top left) and noisy observations of neuron activity (top right; 6/20 shown).
A second low-rank RNN trained on the activity of the first recovers ground
truth dynamics. b) Same set-up, but with Poisson observations. c) The teacher
network was trained on a task where it has to provide an output corresponding
to 8 different angles depending on an input cue. The student network, when
given the same input during fitting, recovers the approximate ring attractor
with 8 stable fixed points. d) Mean ($\pm 1$SD) autocorrelation of the latents
of the models from panel a, show the oscillation frequency is captured, as
well as the decorrelation due to recurrent noise. The scale of the observed
rates also agrees between student and teacher. e) Mean rates and ISI between
student and teacher units of panel b match. f) Example rate distribution of
one unit of the teacher and student RNN (of panel c), after onset of the 8
different stimuli.

We validated our method using several student-teacher setups (Fig. 3;
additional statistics in Fig. S4). We first trained a ‘teacher’ RNN, with the
weight matrix constrained to rank 2, to oscillate. We then simulated multiple
trajectories with a high level of stochasticity in the latent dynamics (Fig.
3a, top left) and additional additive Gaussian observation noise (Fig. 3b, top
left) on the observed neuron activity
($\mathbf{y}_{i}\sim\mathcal{N}(\mathbf{x}_{i},\sigma_{y}$), with
$\mathbf{x}=\mathbf{Mz}$). A second ‘student’ RNN was then fit to the data
drawn from the teacher, and both recovered the true latent dynamical system,
as well as the right level of stochasticity (Fig. 3a, bottom).

Given that neurons emit action potentials, which are commonly approximated as
discrete events, we repeated the experiment with Poisson observations
generated according to
$\mathbf{y}_{i}\sim\mathsf{Pois}(\mathsf{softplus}(\mathbf{w}_{i}\mathbf{x}_{i}-\mathbf{b}_{i}))$.
The student RNN again recovers the oscillatory latent dynamics. Note that
because of the affine transformation in the observation model, the inferred
dynamics can be scaled and translated with respect to the teacher model. To
verify that samples from our inferred model follow the same distribution as
samples from the teacher model, we computed several statistics, which all show
a close match (Fig. 3e; Fig. S4).

In our final teacher-student setup, we verified the ability to recover
dynamics when there are known stimuli or contexts. In particular, we trained a
rank-2 RNN on a task where, at each trial, it receives a transient pulse input
corresponding to a particular angle $\theta$ (given as
$\sin(\theta),\cos(\theta)$), and is asked to provide output matching the
input after stimulus offset. The teacher RNN learns to perform the task by
using an approximate ring attractor, with a stable fixed point for each of the
8 angles - which the student RNN accurately infers (Fig. 3c). Here, we
inferred all fixed points by making use of Preposition 1.

###  3.2 Stochasticity allows recovering low-dimensional latents underlying
EEG data

Figure 4: Example ground truth EEG [42, 43] and (unconditionally) generated
traces by our model. Shown are 5/64 EEG channels.

After validating our model on a toy example, we went on to several challenging
real-world datasets. We first used an EEG dataset [42, 43] with 64 channels
containing one minute of continuous data sampled at 160 Hz (Fig. 4). This
dataset was recently used in a study where generalized teacher forcing (GTF)
was used to fit deterministic RNNs with low-rank structure [8]. The GTF method
obtains state-of-the-art results on several dynamical systems reconstruction
tasks. It outperformed SINDy [44], neural differential equations [45], Long-
Expressive-Memory [46], and other methods, while using a smaller latent
dynamical system.

Here we show that using a stochastic RNN with SMC instead of a deterministic
RNN with GTF, we can decrease the latent dimensionality even further, from 16
to just 3 latents, while matching the original reconstruction accuracy (Table
1). We hypothesize this is because the data can be well explained by
stochastic transitions with simple underlying dynamics as opposed to complex
deterministic chaos.

Table 1: Lower dimensional latent dynamics than SOTA at same sample quality. We report median $\pm$ median absolute deviation over $20$ independent training runs, ‘dim’ refers to the dimensionality of the model’s underlying dynamics and $\lvert\theta\rvert$ denotes the total number of trainable parameters. Values for GTF taken from [8]. Dataset | Method |  $D_{\textrm{stsp}}$ $\downarrow$ |  $D_{H}$ $\downarrow$ | dim | $\lvert\theta\rvert$
---|---|---|---|---|---
EEG (64d) | GTF [8] | $2.1\pm 0.2$ | $0.11\pm 0.01$ | $16$ | $17952$
adaptive GTF [8] | $2.4\pm 0.2$ | $0.13\pm 0.01$ | $16$ | $17952$
SMC (ours) | $2.2\pm 0.2$ | $0.11\pm 0.01$ | $3$ | $3920$

We evaluated samples from our RNN with two measures which were used in [8],
one KL divergence-based measure between the states ($D_{\mathsf{stsp}}$), and
one measure over time, based on the power spectra of generated and inferred
dynamics ($D_{H}$; see Supplement D.3.3). Unlike [8], who applied smoothing,
we optimized our models directly on the raw EEG data.

###  3.3 Interpretable latent dynamics underlying spikes recorded from rat
hippocampus

Figure 5: RNNs reproduce the stationary distribution of spiking data. a) We
fit a rank-3 RNN to spike data recorded from rat hippocampus [47, 48] (left),
and generate new samples from the RNN (right). b) Single neuron statistics.
Mean rates and means of interspike interval (ISI) distributions of a long
trajectory of data generated by the RNN (gen) match those of a held-out set of
data (test). As a reference we additionally computed the same statistics
between the train and test set. c) Population level statistics. We plot the
pairwise correlations between all neurons for generated data against the
pairwise correlations in the test data. d) The corresponding latents generated
by running the RNN look visually similar to the local field potential (LFP).
e) The peak in the power spectrum matches between latents and LFP. f) The
posterior latents show coherence with the LFP. As a reference, we compute the
coherence between the LFP and the latents generated by the RNN.

We next investigated how well our model can capture the distribution of non-
continuous time series. In particular, we used publicly available
electrophysiological recordings from the hippocampus of rats running to drops
of water or pieces of food [47, 48]. We binned the spiking data into 10ms bins
and fit a rank-3 RNN to $\sim$850 s of data. Samples generated by running the
fit RNN autonomously closely matched the statistics of the recordings (Fig.
5a-c).

Figure 6: Posterior latents of our model (fit solely spikes) can be used to
predict rat position.

Previous investigations into this dataset have examined the relationship
between spikes and theta (5-10 Hz) oscillations in the local field potential
([47]), and found that units were locked to the LFP rhythm, with the relative
phase depending on the subregions from which the units were recorded. The
latents generated by the RNN are visually similar to the average local field
potential (Fig. 5d) and match its power spectrum (Fig. 5e). While the model
was solely trained on the spikes, the posterior latents (Eq. 3) have a clear
phase relationship with the LFP, as evidenced by a high coherence between the
posterior latents and LFP. In contrast, and as expected, latents from running
the RNN are not correlated with the LFP (Fig. 5f).

Units in rat hippocampus have been shown to code for position, i.e. through
place cells [49], which tend to fire if the animal is at a specific location.
To further investigate how well we can model recordings from the hippocampus,
we fit a rank-4 RNN to an additional set of recordings of rats running on a
linear track [50, 51, 52]. As in [53], we focus only on the spikes recorded
while the rat is moving, which we bin into 25 ms bins. The RNN again
accurately reconstructs the distribution of spikes and again has latent
oscillations with frequency matching the LFP (Fig. S5). While solely trained
on spikes, the posterior latents also allowed us to predict the position of
the rats with reasonable accuracy ($R^{2}=0.62\pm 0.073$ mean $\pm$ SD, N=4
RNNs; Fig. 6).

###  3.4 Extracting stimulus-conditioned dynamics in monkey reaching task

Figure 7: Inferred and generated dynamics from the model fit to macaque
spiking activity during reaching task. a) Latent states inferred from the
macaque spiking data prior to movement initiation (‘pre-movement’) and during
movement execution (‘movement’), colored by the intended reach target. b)
Reach trajectories decoded from model-inferred neural activity. c)
Dissimilarity matrices computed across the seven conditions (i.e., the seven
colors in a-b) for per-neuron mean firing rate and ISI. We generate neural
activity from the model by providing the same conditioning stimuli as in the
real data. Then, for each statistic, we compute and show the correlation
distance between conditions in the real data (left) and model-generated data
(right). d,e) Same as a-b, but with latent activity and behavioral predictions
generated from the model with conditioning inputs including directions not
seen in the real data (e.g., lime green). For clarity, we show only a subset
of conditions in the decoded reaches.

We further investigated how well we can recover stimulus-conditioned dynamics.
We applied our method to spiking activity recorded from the motor and premotor
cortices of a macaque performing a delayed reaching task. This type of data
has been popular for investigating neural dynamics underlying the control of
movement [2, 3] and evaluating neuroscientific latent variable models [54, 7,
55]. We first validated the ability of our method to obtain a sensible
posterior by evaluating it on the Neural Latents Benchmark [55] (Supplement
B.3, Table S 2).

We then went on to a set-up where we explicitly conditioned our model on
external context. For simplicity, we constrained our experiment to trials with
straight reach trajectories in the data. We fit a rank-5 model to these data
while conditioning its dynamics on the target position by providing the target
position as input to the RNN. Our model was able to infer single-trial latent
dynamics and neuron firing rates that predict reach velocity with high
accuracy at lower latent dimensionalities than models without inputs (Fig. 7b,
$R^{2}=0.90$ for this model, see Table S3 for additional statistics).

We examined the posterior latents inferred by the model and found that our
model recovers structured and interpretable latent dynamics. Before movement
onset, latent states corresponded to the intended reach targets, which were
near the edges of a rectangular screen (Fig. 7a, left), in line with [54].
During the movement period, the latents followed parallel curved trajectories
that preserve target information (Fig. 7a, right) and can be decoded to
predict monkey reach behavior (Fig. 7b).

We then generated neural data from the RNN conditioned on stimulus input.
Again, the distribution of spikes is well-captured (Fig. S6). We additionally
evaluated whether the model faithfully captures differences in spiking
statistics across the seven reach directions, finding reasonable
correspondence in dissimilarities between conditions in the generated and the
real data (Fig. 7c). Finally, we simulated our trained RNN with conditioning
inputs, including reach directions not present in the data, and found that the
structured latent space recovered by the model enables realistic
generalization to unseen reach conditions (Fig. 7d,e, lime green condition).

###  3.5 Searching for fixed points

Figure 8: Comparison of our analytic method (star) and the approximate method
proposed in [25] (blue) for finding the fixed points of the teacher RNN in
Fig. 3c. We can also use Proposition 1 to constrain the search space of the
approximate method (orange). We here show the number of fixed points found as
a function of the number of matrix inverses computed, with errorbars denoting
the minimum and maximum amount of fixed points found over 20 independent runs
of the algorithm.

In Proposition 1, we derived a bound on the number of systems of equations one
has to solve in order to find _all_ fixed points in piece-wise linear low-rank
RNNs. Recently, an approximate algorithm for finding fixed points in piece-
wise linear networks was proposed [25]. Here, we perform an exploration into
how this compares to our analytic method by searching for fixed points of the
RNN in Fig. 3c (top). For the same number of matrix inverses computed by our
analytic method, the approximate method generally does not find all 17 fixed
points (Fig. 8). We note, however, that (unlike ours) the convergence of the
approximate method depends on the dynamics of the RNN, and as a result, there
are theoretical scenarios where the approximate method can be shown to be
faster. Yet we empirically also found scenarios where the approximate methods
failed to converge within the time-frame of our experiments (Fig. S7).

Our analytic method relies on the insight that only a subset of all linear
subregions formed by the piece-wise linear activations can be reached in low-
rank networks.For networks with moderate rank, the cost of searching through
all of the subregions might still be too high. We can, however, hugely reduce
the search space of the approximate method [25] (from $(D+1)^{N}$ to
$\sum^{R}_{r=0}D^{r}{N\choose r}$), at an upfront cost (Supplement B.5; orange
line in Fig. 8).

##  4 Discussion

Here we proposed to fit low-rank RNNs to neural data using variational
sequential Monte Carlo. The resulting RNNs are generative models with
tractable underlying dynamics, from which we can sample long, stable
trajectories of realistic data. We validated our method on several teacher-
student setups and demonstrated the effectiveness of our method on multiple
challenging real-world examples, where we generally needed a latent dynamical
system with very few dimensions to accurately model the data. Besides our
empirical results, we obtained a theoretical bound on the cost of finding
fixed points for RNNs with piecewise-linear activation functions when they are
also low-rank.

Adding stochastic transitions to low-rank RNNs can potentially hugely reduce
the rank required to accurately model observed data, as demonstrated here with
a network fit to EEG data where we could reduce the dimensionality from 16 to
just 3. While many methods that fit RNNs to neural data (e.g, [6, 7, 8, 12,
11, 10]) assume deterministic transitions, there is a rich literature
concentrating on probabilistic sequence models in neuroscience (e.g., [28, 29,
30, 31, 32]). In particular, a recent work termed FINDR [31] uses variational
inference (but not SMC), to similarly find very low-dimensional dynamical
systems underlying neural data. These stochastic dynamical systems were
parameterized using neural differential equations [45]. While Eq. 2 can be
seen as a neural differential equation with one hidden layer [8], our
particular formulation allows us to find its fixed-points effectively and map
back to a regular, mechanistically interpretable RNN (Eq. 1) after fitting,
which enables additional investigations into neural population dynamics [21,
18, 20, 22].

The reason we can do the mapping between a low-rank RNN (Eq. 1) and a latent
dynamical system (Eq. 2) crucially relies on our assumption that samples from
the recurrent noise process are correlated, such that they lie within the
column-space of $\mathbf{M}$. [56] showed that for linear low-rank RNNs
arbitrary covariances in the full $N$ dimensional space can be used, when
increasing the dimensionality of the latent dynamics to twice the rank $R$ (to
the column space of both $\mathbf{M}$ and $\mathbf{N}$), this however does not
generalise to our non-linear setting. We do expect correlated recurrent noise
to be appropriate for modeling stochasticity arising from unobserved inputs or
from partial observations [56] —additionally, correlated noise constituted a
pragmatic choice that allows building an _stochastic_ model that can allow for
trial-by-trial variability while maintaining the tractability of low-rank
deterministic RNNs.

Still, future work can investigate training networks with more relaxed
assumptions on the recurrent noise models, including extensions to non-
Gaussian noise-processes. The latter could be of particular interest if more
biologically plausible (i.e., spiking) neurons were used in the recurrence
[36, 57].

Our results also open up further avenues to explore questions in neuroscience.
The relation between LFP and spike (phase) in the hippocampus has been of
great interest [58, 59, 47, 60]. While we performed some preliminary
investigation into the relation between the inferred latents and the local
field potential, further studies could perform a systematic investigation into
their relation, for instance, by using a multi-modal setup [13], or to
investigate multi-region temporal relationships and interactions [10].

Taken together, by inferring low-rank RNNs with variational SMC, we obtained
generative models of neural data whose trajectories match observed trial-to-
trial variability, and whose underlying latent dynamics are tractable.

## Acknowledgments

This work was supported by the German Research Foundation (DFG) through
Germany’s Excellence Strategy (EXC-Number 2064/1, PN 390727645) and SFB1233
(PN 276693517), SFB 1089 (PN 227953431) and SPP2041 (PN 34721065), the German
Federal Ministry of Education and Research (Tübingen AI Center, FKZ:
01IS18039; DeepHumanVision, FKZ: 031L0197B), the Else Kröner Fresenius
Stiftung (Project ClinbrAIn), and the European Union (ERC, DeepCoMechTome,
101089288). MP and MG are members of the International Max Planck Research
School for Intelligent Systems (IMPRS-IS). We thank Cornelius Schröder for
feedback on the manuscript, and all members of Mackelab for discussions
throughout the project.

## Code availability

Code to reproduce our results is available at
https://github.com/mackelab/smc_rnns.

## References

  * Churchland et al. [2007] M. M Churchland, B. M Yu, M Sahani, and K. V Shenoy.  Techniques for extracting single-trial activity patterns from large-scale neural recordings.  _Current Opinion in Neurobiology_ , 17(5):609–618, 2007.
  * Shenoy et al. [2013] K. V Shenoy, M Sahani, and M. M Churchland.  Cortical control of arm movements: A dynamical systems perspective.  _Annual Review of Neuroscience_ , 36(1):337–359, 2013.
  * Gallego et al. [2017] J Gallego, M Perich, L Miller, and S Solla.  Neural manifolds for the control of movement.  _Neuron_ , 94:978–984, 2017.
  * Vyas et al. [2020] S Vyas, M. D Golub, D Sussillo, and K. V Shenoy.  Computation through neural population dynamics.  _Annual Review of Neuroscience_ , 43(1):249–275, 2020.
  * Barack and Krakauer [2021] D. L Barack and J. W Krakauer.  Two views on the cognitive brain.  _Nature Reviews Neuroscience_ , 22(6):359–371, 2021.
  * Sussillo and Abbott [2009] D Sussillo and L. F Abbott.  Generating coherent patterns of activity from chaotic neural networks.  _Neuron_ , 63:544–557, 2009.
  * Pandarinath et al. [2018] C Pandarinath, D. J O’Shea, J Collins, R Jozefowicz, S. D Stavisky, J. C Kao, E. M Trautmann, M. T Kaufman, S. I Ryu, L. R Hochberg, J. M Henderson, K. V Shenoy, L. F Abbott, and D Sussillo.  Inferring single-trial neural population dynamics using sequential auto-encoders.  _Nature Methods_ , 15(10):805–815, 2018.
  * Hess et al. [2023] F Hess, Z Monfared, M Brenner, and D Durstewitz.  Generalized teacher forcing for learning chaotic dynamics.  In _Proceedings of the 40th International Conference on Machine Learning_ , ICML’23, 2023.
  * Durstewitz [2017] D Durstewitz.  A state space approach for piecewise-linear recurrent neural networks for identifying computational dynamics from neural measurements.  _PLOS Computational Biology_ , 13(6):1–33, 2017.
  * Perich et al. [2021] M. G Perich, C Arlt, S Soares, M. E Young, C. P Mosher, J Minxha, E Carter, U Rutishauser, P. H Rudebeck, C. D Harvey, and K Rajan.  Inferring brain-wide interactions using data-constrained recurrent neural network models.  bioRxiv:2020.12.18.423348, 2021.
  * Valente et al. [2022a] A Valente, J. W Pillow, and S Ostojic.  Extracting computational mechanisms from neural data using low-rank rnns.  In _Advances in Neural Information Processing Systems_ , volume 35, 2022a.
  * Dinc et al. [2023] F Dinc, A Shai, M Schnitzer, and H Tanaka.  CORNN: Convex optimization of recurrent neural networks for rapid inference of neural dynamics.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.
  * Brenner et al. [2024] M Brenner, F Hess, G Koppe, and D Durstewitz.  Integrating multimodal data for joint generative modeling of complex dynamics.  In _Forty-first International Conference on Machine Learning_ , 2024.
  * Barak [2017] O Barak.  Recurrent neural networks as versatile tools of neuroscience research.  _Current Opinion in Neurobiology_ , 46:1–6, 2017.  Computational Neuroscience.
  * Sussillo and Barak [2013] D Sussillo and O Barak.  Opening the Black Box: Low-Dimensional Dynamics in High-Dimensional Recurrent Neural Networks.  _Neural Computation_ , 25(3):626–649, 2013.
  * Seung [1996] H. S Seung.  How the brain keeps the eyes still.  _Proceedings of the National Academy of Sciences_ , 93(23):13339–13344, 1996.
  * Eliasmith and Anderson [2002] C Eliasmith and C. H Anderson.  _Neural Engineering (Computational Neuroscience Series): Computational, Representation, and Dynamics in Neurobiological Systems_.  MIT Press, Cambridge, MA, USA, 2002.
  * Mastrogiuseppe and Ostojic [2018] F Mastrogiuseppe and S Ostojic.  Linking connectivity, dynamics, and computations in low-rank recurrent neural networks.  _Neuron_ , 99(3):609–623.e29, 2018.
  * Schuessler et al. [2020] F Schuessler, A Dubreuil, F Mastrogiuseppe, S Ostojic, and O Barak.  Dynamics of random recurrent networks with correlated low-rank structure.  _Physical Review Research_ , 2(1):013111, 2020.
  * Beiran et al. [2021] M Beiran, A Dubreuil, A Valente, F Mastrogiuseppe, and S Ostojic.  Shaping Dynamics With Multiple Populations in Low-Rank Recurrent Networks.  _Neural Computation_ , 33(6):1572–1615, 2021.
  * Dubreuil et al. [2022] A Dubreuil, A Valente, M Beiran, F Mastrogiuseppe, and S Ostojic.  The role of population structure in computations through neural dynamics.  _Nature Neuroscience_ , 25(6):783–794, 2022.
  * Pals et al. [2024] M Pals, J. H Macke, and O Barak.  Trained recurrent neural networks develop phase-locked limit cycles in a working memory task.  _PLOS Computational Biology_ , 20(2):1–23, 2024.
  * Curto et al. [2019] C Curto, J Geneson, and K Morrison.  Fixed Points of Competitive Threshold-Linear Networks.  _Neural Computation_ , 31(1):94–155, 2019.
  * Brenner et al. [2022] M Brenner, F Hess, J. M Mikhaeil, L. F Bereska, Z Monfared, P.-C Kuo, and D Durstewitz.  Tractable dendritic RNNs for reconstructing nonlinear dynamical systems.  In _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , 2022.
  * Eisenmann et al. [2023] L Eisenmann, Z Monfared, N Göring, and D Durstewitz.  Bifurcations and loss jumps in rnn training.  In _Advances in Neural Information Processing Systems_ , volume 36, 2023.
  * Morrison et al. [2024] K Morrison, A Degeratu, V Itskov, and C Curto.  Diversity of emergent dynamics in competitive threshold-linear networks.  _SIAM Journal on Applied Dynamical Systems_ , 23(1):855–884, 2024.
  * Cunningham and Yu [2014] J. P Cunningham and B. M Yu.  Dimensionality reduction for large-scale neural recordings.  _Nature Neuroscience_ , 17(11):1500–1509, 2014.
  * Petreska et al. [2011] B Petreska, B. M Yu, J. P Cunningham, G Santhanam, S Ryu, K. V Shenoy, and M Sahani.  Dynamical segmentation of single trials from population neural data.  In _Advances in Neural Information Processing Systems_ , volume 24, 2011.
  * Macke et al. [2011] J. H Macke, L Buesing, J. P Cunningham, B. M Yu, K. V Shenoy, and M Sahani.  Empirical models of spiking in neural populations.  In _Advances in Neural Information Processing Systems_ , volume 24, 2011.
  * Linderman et al. [2017] S Linderman, M Johnson, A Miller, R Adams, D Blei, and L Paninski.  Bayesian learning and inference in recurrent switching linear dynamical systems.  In _Proceedings of the 20th International Conference on Artificial Intelligence and Statistics_ , volume 54 of _Proceedings of Machine Learning Research_ , pages 914–922, 2017.
  * Kim et al. [2023] T. D Kim, T. Z Luo, T Can, K Krishnamurthy, J. W Pillow, and C. D Brody.  Flow-field inference from neural data using deep recurrent networks.  bioRxiv:2023.11.14.567136, 2023.
  * Zhao et al. [2023] Y Zhao, J Nassar, I Jordan, M Bugallo, and I Park.  Streaming variational monte carlo.  _IEEE Transactions on Pattern Analysis & Machine Intelligence_, 45(01):1150–1161, 2023.
  * Le et al. [2018] T. A Le, M Igl, T Rainforth, T Jin, and F Wood.  Auto-encoding sequential monte carlo.  In _International Conference on Learning Representations_ , 2018.
  * Maddison et al. [2017] C. J Maddison, J Lawson, G Tucker, N Heess, M Norouzi, A Mnih, A Doucet, and Y Teh.  Filtering variational objectives.  In _Advances in Neural Information Processing Systems_ , volume 30, 2017.
  * Naesseth et al. [2018] C Naesseth, S Linderman, R Ranganath, and D Blei.  Variational sequential monte carlo.  In _Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics_ , volume 84 of _Proceedings of Machine Learning Research_ , 2018.
  * Sourmpis et al. [2023] C Sourmpis, C Petersen, W Gerstner, and G Bellec.  Trial matching: capturing variability with data-constrained spiking neural networks.  In _Advances in Neural Information Processing Systems_ , volume 36, 2023.
  * Doucet and Johansen [2011] A Doucet and A. M Johansen.  A tutorial on particle filtering and smoothing: Fifteen years later.  _The Oxford Handbook of Nonlinear Filtering_ , pages 656–704, 2011.
  * Kingma and Welling [2014] D. P Kingma and M Welling.  Auto-Encoding Variational Bayes.  In _2nd International Conference on Learning Representations, ICLR, Conference Track Proceedings_ , 2014.
  * Doya [1993] K Doya.  Bifurcations of recurrent neural networks in gradient descent learning.  _IEEE Transactions on Neural Networks_ , 1993.
  * Zenn and Bamler [2023] J Zenn and R Bamler.  Resampling gradients vanish in differentiable sequential monte carlo samplers.  In _The First Tiny Papers Track at ICLR 2023, Tiny Papers @ ICLR 2023_ , 2023.
  * Zaslavsky [1975] T Zaslavsky.  Facing up to arrangements: face-count formulas for partitions of space by hyperplanes.  _Memoirs of American Mathematical Society_ , 154:1–95, 1975.
  * Schalk et al. [2004] G Schalk, D McFarland, T Hinterberger, N Birbaumer, and J Wolpaw.  Bci2000: a general-purpose brain-computer interface (bci) system.  _IEEE Transactions on Biomedical Engineering_ , 51(6):1034–1043, 2004.
  * Moody et al. [2000] G Moody, R Mark, and A Goldberger.  Physionet: a research resource for studies of complex physiologic and biomedical signals.  _Computers in cardiology_ , 27:179–82, 2000.
  * Brunton et al. [2016] S. L Brunton, J. L Proctor, and J. N Kutz.  Discovering governing equations from data by sparse identification of nonlinear dynamical systems.  _Proceedings of the National Academy of Science_ , 113(15):3932–3937, 2016.
  * Chen et al. [2018] R. T. Q Chen, Y Rubanova, J Bettencourt, and D. K Duvenaud.  Neural ordinary differential equations.  In _Advances in Neural Information Processing Systems_ , volume 31, 2018.
  * Rusch et al. [2022] T. K Rusch, S Mishra, N. B Erichson, and M. W Mahoney.  Long expressive memory for sequence modeling.  In _International Conference on Learning Representations_ , 2022.
  * Mizuseki et al. [2009a] K Mizuseki, A Sirota, E Pastalkova, and G Buzsáki.  Theta oscillations provide temporal windows for local circuit computation in the entorhinal-hippocampal loop.  _Neuron_ , 64(2):267–280, 2009a.
  * Mizuseki et al. [2009b] K Mizuseki, A Sirota, E Pastalkova, and G Buzsáki.  Multi-unit recordings from the rat hippocampus made during open field foraging.  Database: CRCNS, 2009b.
  * O’Keefe [1976] J O’Keefe.  Place units in the hippocampus of the freely moving rat.  _Experimental Neurology_ , 51(1):78–109, 1976.
  * Grosmark and Buzsáki [2016a] A. D Grosmark and G Buzsáki.  Diversity in neural firing dynamics supports both rigid and learned hippocampal sequences.  _Science_ , 351(6280):1440–1443, 2016a.
  * Chen et al. [2016] Z Chen, A. D Grosmark, H Penagos, and M. A Wilson.  Uncovering representations of sleep-associated hippocampal ensemble spike activity.  _Scientific Reports_ , 6, 2016.
  * Grosmark and Buzsáki [2016b] L. J Grosmark, A.D. and G Buzsáki.  Recordings from hippocampal area ca1, pre, during and post novel spatial learning.  Database: CRCNS, 2016b.
  * Zhou and Wei [2020] D Zhou and X.-X Wei.  Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-vae.  In H Larochelle, M Ranzato, R Hadsell, M Balcan, and H Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 7234–7247. Curran Associates, Inc., 2020.
  * Santhanam et al. [2009] G Santhanam, B. M Yu, V Gilja, S. I Ryu, A Afshar, M Sahani, and K. V Shenoy.  Factor-analysis methods for higher-performance neural prostheses.  _Journal of Neurophysiology_ , 102(2):1315–1330, 2009.
  * Pei et al. [2021] F Pei, J Ye, D. M Zoltowski, A Wu, R. H Chowdhury, H Sohn, J. E O’Doherty, K. V Shenoy, M. T Kaufman, M Churchland, M Jazayeri, L. E Miller, J Pillow, I. M Park, E. L Dyer, and C Pandarinath.  Neural latents benchmark ’21: Evaluating latent variable models of neural population activity.  In _Advances in Neural Information Processing Systems (NeurIPS), Track on Datasets and Benchmarks_ , 2021.
  * Valente et al. [2022b] A Valente, S Ostojic, and J. W Pillow.  Probing the Relationship Between Latent Linear Dynamical Systems and Low-Rank Recurrent Neural Network Models.  _Neural Computation_ , 34(9):1871–1892, 2022b.
  * Cimeša et al. [2023] L Cimeša, L Ciric, and S Ostojic.  Geometry of population activity in spiking networks with low-rank structure.  _PLOS Computational Biology_ , 19(8):1–34, 2023.
  * O’Keefe and Recce [1993] J O’Keefe and M. L Recce.  Phase relationship between hippocampal place units and the eeg theta rhythm.  _Hippocampus_ , 3(3):317–330, 1993.
  * Buzsáki [2006] G Buzsáki.  _Rhythms of the Brain_.  Oxford University Press, 1 edition, 2006.
  * Liebe et al. [2022] S Liebe, J Niediek, M Pals, T. P Reber, J Faber, J Bostroem, C. E Elger, J. H Macke, and F Mormann.  Phase of firing does not reflect temporal order in sequence memory of humans and recurrent neural networks.  bioRxiv:2022.09.25.509370, 2022.
  * Schläfli [1901] L Schläfli.  _Theorie der vielfachen Kontinuität_.  Birkhäuser Basel, Basel, 1901.
  * Buck [1943] R. C Buck.  Partition of space.  _The American Mathematical Monthly_ , 50(9):541–544, 1943.
  * Stanley [2007] R Stanley.  An introduction to hyperplane arrangements.  _Geometric Combinatorics_ , 13:389–496, 2007.
  * Yu et al. [2008] B. M Yu, J. P Cunningham, G Santhanam, S Ryu, K. V Shenoy, and M Sahani.  Gaussian-process factor analysis for low-dimensional single-trial analysis of neural population activity.  In D Koller, D Schuurmans, Y Bengio, and L Bottou, editors, _Advances in Neural Information Processing Systems_ , volume 21, 2008.
  * Ye and Pandarinath [2021] J Ye and C Pandarinath.  Representation learning for neural population activity with neural data transformers.  _Neurons, Behavior, Data analysis, and Theory_ , 5(3), 2021.
  * Paszke et al. [2019] A Paszke, S Gross, F Massa, A Lerer, J Bradbury, G Chanan, T Killeen, Z Lin, N Gimelshein, L Antiga, A Desmaison, A Kopf, E Yang, Z DeVito, M Raison, A Tejani, S Chilamkurthy, B Steiner, L Fang, J Bai, and S Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In _Advances in Neural Information Processing Systems_ , volume 32, 2019.
  * Orvieto et al. [2023] A Orvieto, S. L Smith, A Gu, A Fernando, C Gulcehre, R Pascanu, and S De.  Resurrecting recurrent neural networks for long sequences.  In _Proceedings of the 40th International Conference on Machine Learning_ , ICML’23, 2023.
  * Liu et al. [2020] L Liu, H Jiang, P He, W Chen, X Liu, J Gao, and J Han.  On the variance of the adaptive learning rate and beyond.  In _International Conference on Learning Representations_ , 2020.

##  Appendix A Supplemental material

###  A.1 Proof of preposition 1

####  A.1.1 Problem definition

We are interested in finding all fixed points of the following equation:

$\displaystyle\tau\frac{d\mathbf{x}}{dt}=-\mathbf{x}(t)+\mathbf{J}\phi(\mathbf{x}(t)),$
(8)

with $\mathbf{x}(t)\in\mathbb{R}^{N}$, element-wise nonlinearity
$\phi(\mathbf{x}_{i})=\sum_{d}^{D}\mathbf{b}_{i}^{(d)}\mathsf{max}(\mathbf{x}_{i}-\mathbf{h}_{i}^{(d)})$
and low-rank matrix $\mathbf{J}=\mathbf{MN}^{\mathsf{T}}$, with
$\mathbf{M},\mathbf{N}\in\mathbb{R}^{N\times R}$ and $R\leq N$. Since $\tau$
only scales the speed of the dynamics, we will, for convenience and without
loss of generality, assume $\tau=1$.

####  A.1.2 Preliminaries: Fixed points in Piecewise-linear RNNs

First, we briefly repeat results from [9]. Assume $D=1$,
$\phi(\mathbf{x}_{i})=\mathsf{max}(\mathbf{x}_{i}-\mathbf{h}_{i})$. To find
all fixed points of Eq. 8, start by redefining $\phi$ by introducing a
diagonal indicator matrix:

$\displaystyle{\mathbf{D}_{\Omega}=\begin{bmatrix}d_{1}&&&\\\ &d_{2}&&\\\
&&\ddots&\\\ &&&d_{N}\end{bmatrix}},$ (9)

with $d_{i}=\begin{cases}1,&\text{if }\mathbf{x}_{i}>\mathbf{h}_{i}\\\
0,&\text{otherwise}\end{cases}$.

Then our RNN equation, for a given $\mathbf{x}$ and corresponding
$\mathbf{D}_{\Omega}$ reads:

$\displaystyle\frac{d\mathbf{x}}{dt}=-\mathbf{x}(t)+\mathbf{J}\mathbf{D}_{\Omega}\mathbf{x}(t)-\mathbf{J}\mathbf{D}_{\Omega}\mathbf{h}.$

Each of the $2^{N}$ configuration of $\mathbf{D}_{\Omega}$ corresponds to a
region in which the dynamics are linear. Thus, for each configuration, we can
solve:

$\displaystyle\mathbf{0}$
$\displaystyle=-\mathbf{x}+\mathbf{J}\mathbf{D}_{\Omega}\mathbf{x}-\mathbf{JD}_{\Omega}\mathbf{h},$
$\displaystyle\mathbf{x}^{*}$
$\displaystyle=(\mathbf{JD}_{\Omega}-\mathbf{I})^{-1}\mathbf{J}\mathbf{D}_{\Omega}\mathbf{h}.$

Next, we check whether the obtained $\mathbf{x}^{*}$ is consistent with the
assumed $\mathbf{D}_{\Omega}$ (Eq. 9). If so, we found a fixed point of the
RNN. We have to check, as the solution to the system of linear equations can
lie outside of the linear regions specified by $\mathbf{D}_{\Omega}$. Note
that if for some $\mathbf{D}_{\Omega}$ the matrix
$\mathbf{JD}_{\Omega}-\mathbf{I}$ is not invertible, then there is no single
fixed point, but we still can find a structure of interest (e.g., a direction
with eigenvalue 0 corresponds to marginal stability, i.e., a line attractor).

###  A.2 Preliminaries: Fixed points in Piecewise-linear low-rank RNNs

First, assume $\mathbf{x}(0)$ is in the subspace spanned by the columns of
$\mathbf{M}$. With the low-rank assumption, we can rewrite Eq. 8 for all
$t\in[0,\infty)$, by projecting it on $\mathbf{M}$ [18, 20, 21]:

$\displaystyle\frac{d\mathbf{z}}{dt}=-\mathbf{z}(t)+\mathbf{N}^{\mathsf{T}}\phi(\mathbf{Mz}(t)-\mathbf{h})$
(10)

with $\mathbf{x}(t)=\mathbf{Mz}(t)$.

Now assume $\mathbf{x}(0)$ contains some part $\mathbf{x}^{\perp}(0)$ not in
the subspace spanned by $\mathbf{M}$, i.e., we have
$\mathbf{x}(0)=\mathbf{Mz}(0)+\mathbf{x}^{\perp}(0)$. The dynamics of
$\mathbf{x}^{\perp}$(t) are simply given by
$\frac{\mathbf{x}^{\perp}}{dt}=-\mathbf{x}^{\perp}(t)$ which will decay to its
stable point at $\mathbf{0}$ irrespective of $\mathbf{z}(t)$, and can thus not
contribute additional fixed points.

Naively, using the same strategy as before to obtain all fixed points
$\mathbf{z}$, we would need to solve $2^{N}$ linear systems of $R$ equations
(again for all configurations of $\mathbf{D}_{\Omega}$):

$\displaystyle\mathbf{z}^{*}=(\mathbf{N}^{\mathsf{T}}\mathbf{D}_{\Omega}\mathbf{M}-\mathbf{I})^{-1}\mathbf{N}^{\mathsf{T}}\mathbf{D}_{\Omega}\mathbf{h},$
(11)

###  A.3 Preliminaries: Hyperplane arrangements

Supplementary Figure 1: Proof sketch including $\mathbf{D}_{\Omega}$’s. The
phase-space of an RNN with $N$ (here 2) units with activation
$\max(0,\mathbf{x}_{i}-\mathbf{h}_{i})$ is partitioned into $2^{N}$ (here 4)
regions in which the dynamics are linear, each corresponding to a
configuration of $\mathbf{D}_{\Omega}$. If dynamics are confined to the
$R$-dimensional subspace spanned by the columns of $\mathbf{M}$, only a subset
(here 3) can be reached. Each unit intersects the space spanned by the columns
of $\mathbf{M}$ with a hyperplane (the pink points in the Figure). The amount
of linear regions in $\mathbf{M}$, thus becomes equivalent to "how many
regions can we create in $R$-dimensional space with $N$ hyperplanes?"

In the subsequent section, we will turn to the question of how many equations
we need to solve to find all possible fixed points. Recall that it is possible
to calculate the fixed points analytically because piecewise linear
nonlinearities partition space into subregions in which dynamics are linear.
Each of the linear regions corresponds to a configuration of
$\mathbf{D}_{\Omega}$. For networks with low-rank connectivity, we have to
consider only a small subset of those, as only a small subset of all
configurations of $\mathbf{D}_{\Omega}$ correspond to $\mathbf{x}$’s within
the column space of $\mathbf{M}$ (See Fig. S1). To find out exactly how many
regions lie within the column space, we will need to answer the question: _in
how many regions can we divide $R$-dimensional space, with $N$ hyperplanes?_
To answer this question in general, we will need a theorem from the field of
hyperplane arrangements [61, 62, 41, 63]. Here we give a brief introduction.

##### Introduction to hyperplane arrangements:

A finite arrangements of hyperplanes is a set of $N$ affine subspaces
$\mathcal{A}=\\{a_{1},\dots,a_{N}\\}$ in some vector space $V=\mathbb{R}^{R}$.
Recall a hyperplane is a $R-1$ dimensional subspace defined by a linear
equation $a_{i}:=\\{\mathbf{v}\in V|\mathbf{m}^{\mathsf{T}}\mathbf{v}=h$} for
some $\mathbf{m}\in V,h\in\mathbb{R}$. Note that any linear system of
equations $\mathbf{M}\mathbf{v}=\mathbf{h}$ with
$\mathbf{M}\in\mathbb{R}^{N\times R}$ equivalents defines an arrangement of
$N$ hyperplanes in $R$ dimensional space. In Fig. S2a,c, we show arrangements
of $3$ hyperplanes in $\mathbb{R}^{2}$. In this case, a hyperplane is a line,
but there are infinitely many possibilities on how we can arrange these lines
in two-dimensional space. We are interested in

$\displaystyle\mathcal{N}(\mathcal{A}):=\text{ number of regions
}\mathcal{A}\text{ partitions }\mathbb{R}^{R},$

where regions correspond to the connected components of
$\mathbb{R}^{R}\setminus\mathcal{A}$. In this simple case, we can visually
verify that the arrangements in Fig. S 2a partitions the space into 7 regions,
whereas the arrangement in Fig. S2c partitions the space into only 6 regions.
Clearly, the number of regions $\mathcal{A}$ partitions space in is strongly
related to the number of unique intersections of lines. We have fewer regions
in Fig. S2c, simply because all lines intersect at the same point. If we can
wiggle the hyperplanes a little, and not change the number of regions (as we
can do in Fig. S2a, but not Fig. S2c), we call the hyperplanes in _general_
position (see Theorem 1 for a formal definition).

To count the amount of regions for any arrangement of hyperplanes, we can
leverage an algebraic construction called the intersections poset
$\mathcal{L}(\mathcal{A})$. This is the set of all nonempty intersections of
hyperplanes in $\mathcal{A}$ and includes $V$. Elements of this set are
generally referred to as _flats_. The flats are ordered by reverse inclusion
$x\leq y\iff x\supseteq y$ in the intersection poset. We visualized example
intersection posets of the previous examples (Fig. S2b,d). Here we organized
the flats by dimensionality (such a visualization is called a Hesse Diagram).
Importantly for any real arrangement $\mathcal{A}$, $\mathcal{N}(\mathcal{A})$
solely depends on $\mathcal{L}(\mathcal{A})$ (Corollary 2.1, [63]).

To calculate $\mathcal{N}(\mathcal{A})$ from $L(\mathcal{A})$, we need one
last construction, namely the Möbius function, recursively defined by

$\displaystyle\mu(\mathcal{X},s)=\begin{cases}1&\text{if $s=\mathcal{X}$}\\\
\text{$-\sum_{\mathcal{X}\supseteq s^{\prime}\supset
s}$$\mu(\mathcal{X},s^{\prime})$,}&\text{if
$s\subset\mathcal{X}$.}\end{cases}$ (12)

The numerical values for the example are shown in Fig. S2.

a$a_{1}$$a_{2}$$a_{3}$

b$a_{1}$$-1$$a_{2}$$-1$$a_{3}$$-1$$a_{1}\cap a_{2}$$1$$a_{2}\cap
a_{3}$$1$$a_{1}\cap a_{3}$$1$$\mathbb{R}^{2}$$1$

c$a_{1}$$a_{2}$$a_{3}$

d$a_{1}$$-1$$a_{2}$$-1$$a_{3}$$-1$$a_{1}\cap a_{2}\cap
a_{3}$2$\mathbb{R}^{2}$$1$

Supplementary Figure 2: a) An arrangement of 3 hyperplanes $a_{1},a_{2}$ and
$a_{3}$ in general position. b) the associated intersection poset of the
arrangement in a. c) An alternative arrangement with its associated
intersection poset. d). Blue numbers indicate the value of the Möbius
function.

######  Theorem 1 (Zaslavsky’s Theorem; [41, 63]).

Given a vector space $V=\mathbb{R}^{R}$ and a arrangement of $N$ hyperplanes
$\mathcal{A}=\\{a_{1},\dots,a_{N}\\}$ on $V$, then the number of regions
$\mathcal{A}$ partitions $V$ (denoted $N(\mathcal{A}$), can be expressed as
follows

$\mathcal{N}(\mathcal{A})=\sum_{s\in
L(\mathcal{A})}\mu(\mathbb{R}^{R},s)(-1)^{\text{dim}(s)}$

furthermore, it holds that

$\mathcal{N}(\mathcal{A})\leq\sum_{r=0}^{R}\binom{N}{r}$ (13)

with equality if and only if $\mathcal{A}$ is in general position i.e.
$\mathcal{A}$ must satisfy

  * (i)

$\\{a_{1},\dots,a_{p}\\}\subset\mathcal{A}\text{ and }p\leq
N\Rightarrow\text{dim}(\bigcap_{i=1}^{p}a_{i})=N-p$

  * (ii)

$\\{a_{1},\dots,a_{p}\\}\subset\mathcal{A}\text{ and
}p>N\Rightarrow\bigcap_{i=1}^{p}a_{i}=\emptyset$

One can verify this fact for the given example shown in Fig. S2. We refer to
Stanley [63] for an in-depth formal introduction to this topic. Fundamentally,
it is based on the following recursion that the number of regions for any
arrangement satisfies

$\mathcal{N}(\mathcal{A}\cup\\{a_{N+1}\\})=\mathcal{N}(\mathcal{A})+\mathcal{N}(\mathcal{A}^{a_{N+1}})$

where $\mathcal{A}^{a_{N+1}}:=\\{a_{N+1}\cap
a_{i}:a_{i}\in\mathcal{A},a_{N+1}\cap a_{i}\neq\emptyset,a_{N+1}\nsubseteq
a_{i}\\}$ (Lemma 2.1, Stanley [63]). Note that $a_{N+1}$ is itself an $R-1$
dimensional vector space, and each intersection $a_{N+1}\cap a_{i}$ is an
$R-2$ dimensional hyperplane within $a_{N+1}$ (e.g., the intersection of two
planes is a line within the planes). Hence $\mathcal{A}^{a_{N+1}}$ is itself
an arrangement of $N$ hyperplanes, but in an $R-1$ dimensional subspace. In
fact, the intersection poset exhaustively enumerates the elements of all
possible $\mathcal{A}^{a_{i}}$, and the Möbius function can be shown to
satisfy the above recursion.

If we choose $\phi(\mathbf{x}_{i})=\max(\mathbf{x}_{i}-\mathbf{h}_{i},0)$ i.e
$D=1$, each neuron would partition space by a single hyperplane
$\mathbf{x}_{i}=\mathbf{h}_{i}$ or equivalently the $R$ dimensional subspace
by the hyperplane $\mathbf{M}_{i}^{T}\mathbf{z}=\mathbf{h}_{i}$. Hence, the
hyperplane arrangement is determined by the matrix M and offset $\mathbf{h}$.
As these quantities are learned during training of the RNN, this arrangement
is often in a general position because it is simply numerically unlikely that
two hyperplanes are exactly parallel or intersect in exactly the same
"point".This does, however, change in the general case $D>1$, for which we
derive a tighter bound in the section below.

##### Arrangements of parallel families

For the general case
$\phi(\mathbf{x}_{i})=\sum_{d=1}^{D}\mathbf{b}_{i}^{d}\max(\mathbf{x}_{i}-\mathbf{h}_{i}^{d},0)$
each neurons will partition space with $D$ hyperplanes
$\mathbf{b}_{i}^{d}\mathbf{x}_{i}=\mathbf{b}_{i}^{d}\mathbf{h}_{i}^{(d)}\iff\mathbf{x}_{i}=\mathbf{h}_{i}^{(d)}$
as before; equivalently each neuron partitions the $R$ dimensional subspace
with hyperplanes $\mathbf{m}_{i}^{T}\mathbf{z}=\mathbf{h}_{i}^{d}$. Notably,
all the $D$ hyperplanes here will share the same row of $\mathbf{M}$, and thus
they are parallel. Clearly, any such arrangement cannot be in general
arrangement by definition.

The resulting arrangement will have a very specific structure. Let’s define

$A_{i}:=\\{a_{i1},\dots,a_{iD}\\}$

as a family of $D$ parallel hyperplanes. Any pair of hyperplanes
$a_{il},a_{im}\in A_{i}$ is parallel. A low-rank RNN with $N$ neurons and a
general activation function will thus lead to an arrangement consisting of $N$
families of $D$ parallel hyperplanes.

We can use this specific structure to obtain a tighter bound.

######  Lemma 1.

Let $\mathcal{A}=A_{1}\cup\dots\cup A_{N-1}$ be an arrangement of $N-1$
families of $D$ parallel lines then it satisfies the following recursion

$\mathcal{N}(\mathcal{A}\cup
A_{N})=\mathcal{N}(\mathcal{A})+\sum_{d=1}^{D}\mathcal{N}\left(\mathcal{A}^{a_{Nd}}\right)$

Furthermore, denote by $\mathcal{N}(N,R,D)$ the maximum number of regions
attainable by any arrangement of $N$ families of $D$ parallel hyperplanes in
$R$ dimensional space then

$\mathcal{N}(N,R,D)\leq\mathcal{N}(N-1,R,D)+D\cdot\mathcal{N}(N-1,R-1,D)$

###### Proof.

To add a $A_{n}$ to $\mathcal{A}$, we have to add $D$ new parallel
hyperplanes. We can do so by iteratively applying Lemma 2.1 [63]. We obtain

$\displaystyle\mathcal{N}(\mathcal{A}\cup\\{a_{N1},\dots,a_{ND}\\})$
$\displaystyle=\mathcal{N}(\mathcal{A}\cup\\{a_{N1},\dots,a_{N(D-1)}\\})+\mathcal{N}(\left(\mathcal{A}\cup\\{a_{N1},\dots,a_{N(D-1)}\\}\right)^{a_{ND}})$
$\displaystyle=\mathcal{N}(\mathcal{A})+\sum_{d=1}^{D}\mathcal{N}\left(\left[\mathcal{A}\cup\bigcup_{i=1}^{d-1}\\{a_{Ni}\\}\right]^{a_{Nd}}\right)$

Now note that $\mathcal{A}^{a_{Nj}}:=\\{a_{Nj}\cap
a_{lm}:a_{lm}\in\mathcal{A},a_{Nj}\cap a_{lm}\neq\emptyset,a_{Nj}\nsubseteq
a_{lm}\\}$, hence by definition only hyperplanes that do intersect with
$a_{Nj}$ are included in this set. As $a_{Nj}$ is parallel to any other
$a_{Ni}$ for all $i\neq j$, all $a_{Nj}\cap a_{Ni}$ cannot be in the set.
Hence for any $d$, we have that

$\mathcal{N}\left(\left[\mathcal{A}\cup\bigcup_{i=1}^{d-1}\\{a_{Ni}\\}\right]^{a_{Nd}}\right)=\mathcal{N}(\mathcal{A}^{a_{Nd}})$

which proves the first equation.

Recall that we define $\mathcal{N}(N,R,D)$ as the maximum number of regions
attainable by any arrangement. Notice that $\mathcal{A}$ by construction is an
arrangement of $N-1$ families of $D$ parallel hyperplanes in $R$ dimension.
Thus by definition $\mathcal{N}(\mathcal{A})\leq\mathcal{N}(N-1,R,D)$.

Furthermore, the intersection set of two hyperplanes in dimension $R$ is
itself hyperplanes of dimension $R-1$. Furthermore, the intersection sets of
$D$ parallel hyperplanes with $a_{Nd}$, remain parallel and hence
$\mathcal{A}^{a_{Nd}}$ is an arrangement of at most $N-1$ families of $D$
parallel hyperplanes in $R-1$ dimensions. Thus
$\mathcal{N}(\mathcal{A}^{a_{Nd}})\leq\mathcal{N}(N-1,R-1,D)$ leaving us with

$\mathcal{N}(\mathcal{A}\cup\\{a_{N1},\dots,a_{ND}\\})\leq\mathcal{N}(N-1,R,D)+D\cdot\mathcal{N}(N-1,R-1,D)$

As this holds for any arrangement, it also holds for the arrangement that has
$\mathcal{N}(N,R,D)$ regions (i.e., which maximizes the number of regions)
and, therefore, proves the second equation.

∎

######  Lemma 2.

Let $\mathcal{A}$ be an arrangement of $N$ families of $D$ parallel
hyperplanes. Then, it holds that

$N(\mathcal{A})\leq\sum_{r=0}^{R}D^{r}\binom{N}{r}$

with equality if each family is in a general position, i.e. that every
subarrangement $\\{a_{1j_{1}},\dots,a_{Nj_{N}}\\}$ for all $1\leq j_{i}\leq D$
is in general position.

###### Proof.

We will first construct an intersection poset $L(\mathcal{A})$ on the level of
families $A_{i}$ in general position. After all, the intersection properties
between these families is the same as between their elements, e.g. if $a_{i1}$
intersects $a_{j1}$ then also all lines in $A_{i}$ intersect all lines in
$A_{j}$.

The resulting intersection poset $L(\mathcal{A})$ can be clustered into the
corresponding families. We visualize the construction in Fig. S3.

At each rank $r$ (level from bottom to top), we can choose exactly
$\binom{N}{r}$ families of hyperplanes that intersect (exactly the case if we
just have $N$ hyperplanes in general position). To obtain a flat of dimension
$N-r$ we have to choose $r$ out of the $N$ hyperplane families without
replacement.

If, e.g., two families of parallel hyperplanes $A_{i},A_{j}$ intersect, then
any element $a_{ik}$ will intersect with any element $a_{jl}$ for all $1\leq
k,l\leq D$ leading to at most $D^{2}$ flats within each family (there can be
less as other families might intersect in the same "point"). In general, each
cluster of intersections of $r$ families will contain at most $D^{r}$ flats.

By construction of $L(\mathcal{A})$ and Theorem 1, the lemma follows directly.

To show that this construction indeed is an upper bound for all arrangements,
we can use Lemma 1. There, we established a recursion, which any such upper
bound must satisfy. Hence, assume
$\mathcal{N}(N,R,D)=\sum_{r=0}^{R}D^{r}\binom{N}{r}$. Notice that using
Pascal’s identity, we can rewrite

$\displaystyle\mathcal{N}(N,R,D)$
$\displaystyle=\sum_{r=0}^{R}D^{r}\binom{N}{r}$
$\displaystyle=\sum_{r=0}^{R}D^{r}\left(\binom{N-1}{r}+\binom{N-1}{r-1}\right)$
$\displaystyle=\sum_{r=0}^{R}D^{r}\binom{N-1}{r}+\sum_{r=0}^{R}D^{r}\binom{N-1}{r-1}$
$\displaystyle=\sum_{r=0}^{R}D^{r}\binom{N-1}{r}+\underbrace{D^{0}\binom{N-1}{-1}}_{:=0}+\sum_{r=1}^{R}D^{r}\binom{N-1}{r-1}$
$\displaystyle=\mathcal{N}(N-1,R,D)+\sum_{r=0}^{R-1}D^{r+1}\binom{N-1}{r}$
$\displaystyle=\mathcal{N}(N-1,R,D)+D\cdot\mathcal{N}(N-1,R-1,D)$

∎

$A_{1}\ |\ \\{a_{1i}|i<=D\\}$$\dots$$A_{N}\ |\ \\{a_{Ni}|i<=D\\}$$A_{1}\cap
A_{2}\ |\ \\{a_{1i}\cap a_{2j}|i,j\leq D\\}$$\dots$$A_{1}\cap A_{N}\ |\
\\{a_{1i}\cap a_{Nj}|i,j\leq
D\\}$$\mathbb{R}^{R}$$\vdots$$\vdots$$\vdots$$\bigcap_{k=1}^{R}A_{k}\ |\
\\{\bigcap_{k}a_{kd_{k}}|d_{1},\dots,d_{k}\leq D\\}$$\dots$$\dots$
Supplementary Figure 3: Construction of the intersection poset
$L(\mathcal{A})$ for a arrangement of $N$ families $A_{i}$ of $D$ parallel
hyperplanes in "general position".

###  A.4 Proof of proposition

Using the previously derived techniques, we will prove here the main
proposition. Furthermore, in Algorithm 1, pseudo-code is given to compute all
fixed points in practice.

######  Proposition 1.

Assume the RNN of Eq. 8 , with $\mathbf{J}$ of rank $R$ and piecewise-linear
activations:
$\phi(\mathbf{x}_{i})=\sum_{d}^{D}\mathbf{b}_{i}^{(d)}\mathsf{max}(\mathbf{x}_{i}-\mathbf{h}_{i}^{(d)},0)$.
For fixed rank $R$ and fixed number of basis functions $D$, we can find all
fixed points in the absence of noise, that is all $\mathbf{x}$ for which
$\frac{d\mathbf{x}}{dt}=0$, by solving at most $\mathcal{O}(N^{R})$ linear
systems of $R$ equations.

###### Proof.

By definition, each neuron partitions $\mathbb{R}^{N}$ in $D$ linear regions
with hyperplanes described by $\mathbf{x}_{i}^{(d)}=\mathbf{h}_{i}^{(d)}$, for
the $i$’th neuron. Using that in the columnspace of $\mathbf{M}$, we have
$\mathbf{x}=\mathbf{M}\mathbf{z}$., it follows that each neuron partitions the
$R$ dimensional subspace spanned by columns of $\mathbf{M}$, with $D$
hyperplanes described by
$\sum_{r}^{R}\mathbf{M}_{i,r}\mathbf{z}_{r}=\mathbf{h}_{i}^{(d)}$. Notice that
these hyperplanes are parallel, as they all share the same coefficients
$\mathbf{M}_{r}$ but have a different offset $\mathbf{h}_{i}^{(d)}$. Using
Lemma 2 we know that there can only be $\sum_{r=0}^{R}D^{r}\binom{N}{r}$ such
regions.

How do we find those regions? Let’s first consider the case of $D=1$, and
assume that the hyperplanes are in general position. We can find the
corresponding configurations of $\mathbf{D}_{\Omega}$ as follows. We first
obtain the set of all intersections of $R$ hyperplanes. For this we try to
solve ${N\choose R}$ systems of $R$ equations. Let
$\mathbf{M}_{R}\in\mathbb{R}^{R\times R}$ be the matrix obtained by choosing
$R$ different rows $1,\dots,R$ of $\mathbf{M}\in\mathbb{R}^{N\times R}$ (i.e.
picking $R$ neurons), then we may find the corresponding intersection of $R$
hyperplanes by solving the following linear system of $R$ equations

$\displaystyle\mathbf{z}_{\cap}=\mathbf{M}_{R}^{-1}\mathbf{h}_{R}\qquad\text{
and }\qquad\mathbf{x}_{\cap}=\mathbf{M}\mathbf{z}_{\cap}.$

which will always have a unique solution if all hyperplanes are in general
position, as then all $\mathbf{M}_{R}$ have rank $R$. Each $\mathbf{x}_{\cap}$
has $2^{R}$ possible bordering linear regions. We can find the corresponding
$\mathbf{D}_{\Omega}=\text{diag}([d_{1},\dots,d_{N})$’s matrices of each of
those subsections as follows. First $d_{i}=\mathbb{I}(\mathbf{x}_{\cap}<0)$
for all $i<=N$. By construction $1,\dots,R$ at $\mathbf{x}_{\cap}$ will be
exactly at the threshold, by moving away from it $d_{R}$ can become either
zero or one, depending on in which region we and up. Hence, the $2^{R}$
regions correspond to one in which either combination of neurons $1,\dots R$
is active (meaning that it is above the threshold). We thus just have to check
all combinations $d_{1},\dots,d_{R}\in\\{0,1\\}^{R}$. Using this, we will find
at most $\sum^{R}_{r=0}{N\choose r}$ unique configurations (as this is the
maximal number of regions possible for $D=1$). To find all the fixed points we
hence have to solve Eq. 11 for each configuration. We thus end up with solving
${N\choose R}$ systems of $R$ linear equations to find all regions, and
another $\sum^{R}_{r=0}{N\choose r}\in\mathcal{O}(N^{R})$ systems of $R$
linear equations to find all fixed points.

Let us now consider the case for $D>1$. Note that an RNN with $N$ units and
$D$ basis functions per unit, can be expanded to an RNN with $ND$ units with
activation
$\phi(\mathbf{x}_{i})=\mathsf{max}(\mathbf{x}_{i}-\mathbf{h}_{i},0)$ ([24],
Theorem 1). Any fixed point can then still be analytically computed using Eq.
11. We expand the network but keep track of all $\sum_{r}^{R}D^{r}{N\choose
R}$ possible intersections. It still holds that from each intersection, we can
reach $2^{R}$ regions. In total, we will now find at most
$\sum^{R}_{r=0}D^{r}{N\choose r}$ regions (Lemma 2). To find all the fixed
points, we hence have to solve
$\binom{N}{R}D^{r}+\sum^{R}_{r=0}\binom{N}{r}D^{r}$ systems of $R$ linear
equations, which for constant $D$ and $R$ has a cost of $\mathcal{O}(N^{R})$

Finally, let’s consider the case when hyperplanes are not in general position
(which is unlikely to happen when doing numerical optimization). If there are
intersections of more than $R$ hyperplanes, we proceed as before, but in case
the intersection of $R$ hyperplanes we are currently considering intersects
additional hyperplanes, set the diagonal elements of $\mathbf{D}_{\Omega}$
corresponding to these additional hyperplanes arbitrarily to $1$ (as
intersections including the additional hyperplanes are considered separately).
On the other hand, in case some hyperplanes are only part of intersections of
less than $R$ hyperplanes (because they became parallel), we proceed as
follows. Instead of considering only intersections of $R$ hyperplanes, we now
also consider all possible intersections of $r$ hyperplanes, with $1\leq r\leq
R$. For this, we solve no more than $\sum_{r}^{R}{N\choose r}$ systems of $r$
equations. Let $\mathbf{M}_{r}\in\mathbb{R}^{r\times R}$ be the matrix
obtained by choosing $r$ different linearly independent rows $1,\dots,r$ of
$\mathbf{M}\in\mathbb{R}^{N\times R}$; then we may find a point on the
corresponding intersection of $r$ hyperplanes (note that the intersection
itself can now also be a hyperplane) by to solving the following linear system
of $r$ equations

$\displaystyle\mathbf{z}_{\cap}=\mathbf{M}_{r}^{\mathsf{\dagger}}\mathbf{h}_{r}\qquad\text{
and }\qquad\mathbf{x}_{\cap}=\mathbf{M}\mathbf{z}_{\cap}.$

with $\dagger$ being the pseudoinverse. We here now end up with solving no
more than $\sum_{r}^{R}$ ${N\choose r}$ systems of $r$ linear equations to
find all regions, which has an equal cost in $N$ as the previous cases.

∎

We here provide pseudocode. For simplicity, we restrict ourselved to the case
of $D=1$ and assume that the arrangement specified by $\mathbf{M}$ and
$\mathbf{h}$ is in general position. This can be generalized to the general
setting as presented in the proof.

Data: $\mathbf{N}\in\mathbb{R}^{N\times R},\mathbf{M}\in\mathbb{R}^{N\times
R},\mathbf{h}\in\mathbb{R}^{N}$

Result: $z\\_set$ set of all fixpoints, $D\\_set$ the set of all relevant
$\mathbf{D}_{\Omega}$ configurations.

$D\\_set:=\\{\\}$;

$z\\_set:=\\{\\}$;

$idx$ = [$1,\dots,N$];

// Find feasible configurations

$idx\\_comb$ = all $\binom{N}{R}$ combinations of indices $idx$;

for  _$(i_{1},\dots,i_{R})$ in $idx\\_comb$_ do

$\mathbf{M}_{R}=\mathbf{M}[(i_{1},\dots,i_{R}),:]$;

$\mathbf{h}_{R}=\mathbf{h}[(i_{1},\dots,i_{R})]$;

// $\mathbf{M}_{R}$ is invertible as the arrangement is in general position

$\mathbf{z}_{\cap}=\text{solve}(\mathbf{M}_{R},\mathbf{h}_{R})$;

$\mathbf{x}_{\cap}=\mathbf{M}\mathbf{z_{\cap}}$;

$d\\_init=\mathbf{x}_{\cap}>\textbf{h}$;

for  _$(v_{1},\dots,v_{R})$ in $\\{0,1\\}^{R}$_ do

$d=d\\_init[(i_{1},\dots,i_{R})].set(v_{1},\dots,v_{R})$ ;

$\mathbf{D}_{\Omega}=\text{diag}(d)$ ;

$D\\_set=D\\_set\cup\\{\mathbf{D}_{\Omega}\\}$;

end for

end for

// Find fixed points, for the at most $\sum_{r=0}^{R}\binom{N}{r}$
configurations

for  _$D_{\Omega}$ in $D\\_set$_ do

$\mathbf{z}^{*}=\text{solve}(\mathbf{N}^{\mathsf{T}}\mathbf{D}_{\Omega}\mathbf{M}-\mathbf{I},\mathbf{N}^{\mathsf{T}}\mathbf{D}_{\Omega}\mathbf{h}$)
;

$z\\_set=z\\_set\cup\\{\mathbf{z}^{*}\\}$

end for

Algorithm 1 Improved exhaustive search for all fixedpoints

##  Appendix B Additional figures & tables

###  B.1 Additional statistics for Teacher-Student setups

Supplementary Figure 4: a-c) Pairwise correlations between units of the modes
for panel a-c) of Fig. 3, respectively. Note that c is computed over all
conditions.

###  B.2 Statistics for HPC-11

Supplementary Figure 5: a) We fit a rank-4 RNN to spikes recorded from rat
hippocampus [50, 51, 52], and generate new samples from the RNN (right). b)
Single neuron statistics. The mean rates and coefficient of variations of
interspike interval (ISI) distributions of a long trajectory of data generated
by the RNN (gen) match those of a held-out set of data (test). As a reference
we additionally computed the same statistics between the train and test set.
c) Population level statistics. The pairwise correlations between neurons for
generated data and the test data. d) The corresponding latents generated by
the RNN consists of 10Hz (fast theta) oscillations on top of slower
oscillations. e) Latents with further zooming in (on time), shown together
with the LFP signal. f) The power spectrum latents sampled from the RNN next
to that of the LFP.

###  B.3 Neural Latents Benchmark evaluation

We applied our method to the MC_Maze dataset of the Neural Latents Benchmark
(NLB) [55] at 20 millisecond bin size (Table 2). The benchmark evaluates
methods on a number of metrics: ‘co-bps’ (co-smoothing bits-per-spike)
assesses the quality of firing rate predictions for a set of held-out neurons
that are unobserved in the test data, evaluated with the Poisson likelihood of
the true spiking activity given the rate predictions. ‘vel R2’ evaluates how
well the model’s inferred firing rates can predict the subject’s hand
velocity. ‘PSTH R2’ evaluates how well peri-stimulus time histograms (PSTHs)
computed from model-inferred rates match empirical PSTHs from the data. We
found that our method outperforms classical methods (GPFA [64] and SLDS[30])
while certain state-of-the-art deep learning (LFADS [7], Neural Data
Transformer [65]) are slightly better than our method on this particular
benchmark. We do note that NLB metrics center around evaluating the quality of
smooth rates inferred from spikes, which is not the central focus of our
method. Rather, we aim to fit an RNN, from which — by design — we can sample
noisy latent trajectories that reproduce variability in the data.

While our method also has comparatively lower dimensionality than the other
deep learning approaches, a latent dimensionality of $36$ is still
considerably higher than all networks considered in the Main text. We reason
that we need a high number of latents, because the full MC_Maze dataset has a
large number of conditions (108), spanning multiple maze-configurations, which
may be difficult to fully model with autonomous low-dimensional latent
dynamics.

Table 2: Performance of our method on the MC_Maze dataset of the Neural Latents Benchmark, ‘dim’ refers to the dimensionality of the model’s underlying dynamics (where possible). method | dim | co-bps $\uparrow$ | vel R2 $\uparrow$ | PSTH R2 $\uparrow$
---|---|---|---|---
Spike smoothing | 137 | $0.2076$ | $0.6111$ | $-0.0005$
GPFA | 52 | $0.2463$ | $0.6613$ | $0.5574$
SLDS | 38 | $0.2117$ | $0.7944$ | $0.4709$
LFADS | 100 | $0.3554$ | $0.8906$ | $0.6002$
NDT | 274 | $0.3597$ | $0.8897$ | $0.6172$
Ours | 36 | $0.3210$ | $0.8571$ | $0.5902$

###  B.4 Stimulus-conditioning in monkey reaching task

For the experiment with stimulus-conditioned dynamics in the monkey reaching
task, we tested the performance of models with and without the conditioning
inputs. We found that the conditioning inputs allow the networks to perform
better on velocity decoding at lower dimensionalities.

Table 3: Performance benefits of conditioning for monkey reaching task. conditioning | dim | vel R2 $\uparrow$
---|---|---
w/o conditioning | 5 | $0.7897\pm 0.0687$
6 | $0.8944\pm 0.0039$
8 | $0.9085\pm 0.0048$
16 | $0.9196\pm 0.0041$
with conditioning | 5 | $0.8589\pm 0.0493$
6 | $0.9018\pm 0.0114$

Following the analysis in Fig. 7, we also further visualized the model’s match
to the spiking statistics, including mean and standard deviation (SD) of
spiking rate, and mean, SD, and coefficient of variation (CV) of inter-spike
intervals. We observed a good match to the mean and SD of the spiking rate
across all conditions. Match to ISI statistics is also quite reasonable given
the noise observed between estimates of the statistics from train and test.

Supplementary Figure 6: Spiking statistics of model-generated (teal) and train
data (brick red) compared against test data.

###  B.5 Comparison to approximate method for finding fixed points

Supplementary Figure 7: Repetition of the experiment of Fig. 8, but now with a
rank-2 RNN with 128 units. Again, we show the number of fixed points found as
a function of the number of matrix inverses computed, with errorbars denoting
the minimum and maximum amount of fixed points found over 20 independent runs
of the algorithm.

Recently an approximate method for finding fixed points in piece-wise linear
RNNs was proposed [25]. The method proceeds by randomly selecting a linear
region (a configuration of $\mathbf{D}_{\Omega}$, see Supplement A.1.2) and
calculating the corresponding fixed-point. If it is indeed a ’true’ fixed
point of the RNN (it is consistent with the assumed $\mathbf{D}_{\Omega}$), we
store it. If the fixed point was inconsistent with the assumed
$\mathbf{D}_{\Omega}$, we iteratively initialize $\mathbf{D}_{\Omega}$
according to the ’virtual’ fixed point found and calculate the new fixed point
corresponding to this $\mathbf{D}_{\Omega}$, until we either reach a ’true’
fixed point or reach a certain amount of iterations. Then, we reinitialize at
a randomly selected new configuration of $\mathbf{D}_{\Omega}$ and repeat the
procedure.

Under some conditions, the approximate method can be shown to converge in
linear time ($\lVert\mathbf{M}\tilde{\mathbf{N}}^{\mathsf{T}}\rVert+\lVert
a\mathbf{I}\rVert\leq 1$) [8], where it will be faster than our exact method —
however in general the convergence of the approximate method strongly depends
on the dynamics of the networks. In particular, there are reasonable settings
where the approximate method fails to find all fixed points, such as of a
rank-2, 128 unit RNN with 17 fixed points (trained similarly to the teacher
RNN of Fig. 3c; Fig. S7). While an in-depth study of the approximate method is
out of scope, we hypothesize that the failure to converge is because when
initializing with randomly selected $\mathbf{D}_{\Omega}$s out of $(D+1)^{N}$
possible configurations, the approximate method tends to converges to the same
set of $\mathbf{D}_{\Omega}$.

Our method is completely independent of the dynamics of the system and has a
fixed cost, after which one is guaranteed that all fixed points are found.
However, we do note that there can be scenarios where our exact method is
still too costly. In this scenario, we propose to use the approximate method,
with one adjustment - we first pre-compute the subset of
$\sum^{R}_{r}D^{r}{N\choose r}$ configurations that can contain fixed points,
and then initialize the approximate method using randomly selected
$\mathbf{D}_{\Omega}$ from this subset. Empirically, this leads to better
convergence in at least some scenarios (Fig. 8, Fig. S7)

For the approximate method, we used code from
https://github.com/DurstewitzLab/CNS-2023, which was released with the GNU
General Public License.

##  Appendix C Additional details of low-rank RNNs

###  C.1 Discretisation

Given

$\displaystyle\tau\frac{d\mathbf{z}}{dt}$
$\displaystyle=-\mathbf{z}(t)+\mathbf{N}^{\mathsf{T}}\phi(\mathbf{M}\mathbf{z}(t))+\Gamma_{\mathbf{z}}\xi(t),$

Using the Euler–Maruyama method with timestep $\Delta_{t}$:

$\displaystyle\mathbf{z}_{t+1}=(1-\frac{\Delta_{t}}{\tau})\mathbf{z}_{t}+\frac{\Delta_{t}}{\tau}\mathbf{N}^{\mathsf{T}}\phi(\mathbf{M}\mathbf{z}_{t})+\frac{\sqrt{\Delta_{t}}}{\tau}\Gamma_{\mathbf{z}}\epsilon_{t},$

and with $\epsilon_{t}\sim\mathcal{N}(0,\mathbf{I})$, define
$a=1-\frac{\Delta_{t}}{\tau}$,
$\tilde{\mathbf{N}}=\frac{\Delta_{t}}{\tau}\mathbf{N}$, and
$\Sigma_{\mathbf{z}}=\frac{\Delta_{t}}{\tau}^{2}\Gamma_{\mathbf{z}}\Gamma_{\mathbf{z}}^{\mathsf{T}}$,
we obtain the transition distribution used in our experiments.

###  C.2 Conditional generation

Given input weights $\mathbf{H}\in\mathbb{R}^{N\times N_{s}}$ and stimulus
$\mathbf{s}\in\mathbb{R}^{N_{s}}$, we define our model as

$\displaystyle\tau\frac{d\mathbf{x}}{dt}=-\mathbf{x}(t)+\mathbf{J}\phi(\mathbf{x}(t))+\mathbf{H}\mathbf{s}(t)+\xi_{\mathbf{x}}.$

Using the same assumptions as before, $\mathbf{x}$ can be described by
$R+N_{s}$ variables

$\displaystyle\tau\frac{d\mathbf{z}}{dt}$
$\displaystyle=-\mathbf{z}(t)+\mathbf{N}^{\mathsf{T}}\phi(\mathbf{M}\mathbf{z}(t)+\mathbf{H}\tilde{\mathbf{s}}(t))+\xi_{\mathbf{z}},$
$\displaystyle\tau\frac{d\tilde{\mathbf{s}}}{dt}$
$\displaystyle=-\tilde{\mathbf{s}}(t)+\mathbf{s}(t),$

with $\mathbf{x}=\mathbf{M}\mathbf{z}+\mathbf{H}\tilde{\mathbf{s}}$, and
$\begin{bmatrix}\mathbf{z}\\\
\tilde{\mathbf{s}}\end{bmatrix}=([\mathbf{M},\mathbf{H}]^{\mathsf{T}}[\mathbf{M},\mathbf{H}])^{-1}[\mathbf{M},\mathbf{H}]^{\mathsf{T}}\mathbf{x}$.

For constant input $\mathbf{s}$, $\tilde{\mathbf{s}}$ will converge to
$\mathbf{s}$, and we can ignore the additional $N_{s}$ variables, assuming
$\mathbf{x}(0)=\mathbf{Mz}(0)+\mathbf{s}$. Similarly if $\mathbf{s}$ varies on
a time scale slower than $\tau$, $\mathbf{s}\approx\tilde{\mathbf{s}}$ is a
good approximation [21]. Here, for all experiments, the input is either a
constant context signal or a rectangular pulse, so we always substitute
$\mathbf{s}$ for $\tilde{\mathbf{s}}$ and consider the $R$ dimensional system
described by $\mathbf{z}$ (which now has additional conditioning on
$\mathbf{s}$).

We can write the conditional sequential distribution generated by discretizing
our model as

$\displaystyle p(\mathbf{z}_{1:T},\mathbf{y}_{1:T}|\mathbf{s}_{1:T-1})$
$\displaystyle=p(\mathbf{z}_{1})\prod_{t=2}^{T}p(\mathbf{z}_{t}|\mathbf{s}_{t-1},\mathbf{z}_{t-1})\prod_{t=1}^{T}p(\mathbf{y}_{t}|\mathbf{z}_{t}),$
$\displaystyle p(\mathbf{z}_{t}|\mathbf{s}_{t-1},\mathbf{z}_{t-1})$
$\displaystyle=\mathcal{N}(F(\mathbf{s}_{t-1},\mathbf{z}_{t-1}),\Sigma_{\mathbf{z}}),\;p(\mathbf{z}_{1})=\mathcal{N}(\mu_{\mathbf{z}_{1}},\Sigma_{\mathbf{z}_{1}}),$

where the transition distribution is
$F(\mathbf{s}_{t},\mathbf{z}_{t})=a\mathbf{z}_{t}+\tilde{\mathbf{N}}^{\mathsf{T}}\phi(\mathbf{M}\mathbf{z}_{t}+\mathbf{Hs}_{t})$.

###  C.3 Linear transformations of the latent space and orthogonalisation

Given

$\displaystyle\mathbf{x}_{t+1}$
$\displaystyle=a\mathbf{x}_{t}+\mathbf{M}\tilde{\mathbf{N}}^{\mathsf{T}}\phi(\mathbf{x}_{t})+\epsilon_{\mathbf{x}}$
$\displaystyle\mathbf{z}_{t+1}$
$\displaystyle=a\mathbf{z}_{t}+\tilde{\mathbf{N}}^{\mathsf{T}}\phi(\mathbf{Mz}_{t})+\epsilon_{\mathbf{z}}$

with $\epsilon_{\mathbf{z}}\sim\mathcal{N}(0,\Sigma_{\mathbf{z}})$,
$\epsilon_{\mathbf{x}}\sim\mathcal{N}(0,\mathbf{M}\Sigma_{\mathbf{z}}\mathbf{M}^{\mathsf{T}}).$
We can do any linear transformation of the latent dynamics $\mathbf{z}$:
$\hat{\mathbf{z}}=\mathbf{Az}$, as long as $\mathbf{A}$ has rank $R$, without
changing the neuron activity $\mathbf{x}$. To see this, define
$\hat{\mathbf{M}}=\mathbf{M}\mathbf{A}^{-1}$,
$\hat{\mathbf{N}}=\mathbf{A}\tilde{\mathbf{N}}$, and
$\epsilon_{\hat{\mathbf{z}}}\sim\mathcal{N}(0,\mathbf{A}\Sigma_{\mathbf{z}}\mathbf{A}^{T})$,
giving us:

$\displaystyle\mathbf{x}_{t+1}$
$\displaystyle=a\mathbf{x}_{t}+\hat{\mathbf{M}}\hat{\mathbf{N}}^{\mathsf{T}}\phi(\mathbf{x}_{t})+\epsilon_{\mathbf{x}}$
$\displaystyle\hat{\mathbf{z}}_{t+1}$
$\displaystyle=a\hat{\mathbf{z}}_{t}+\hat{\mathbf{N}}^{\mathsf{T}}\phi(\hat{\mathbf{M}}\hat{\mathbf{z}}_{t})+\epsilon_{\hat{\mathbf{z}}},$

which will leave $\mathbf{x}$ unchanged, while our latents $\mathbf{z}$ are
expressed in a new basis. We typically got a more interpretable visualization
of the latents by orthonormalising the columns of $\mathbf{M}$. Thus we
applied for all visualisations after training
$\mathbf{A}=\mathbf{U}^{\mathsf{T}}\mathbf{M}$, with
$\hat{\mathbf{M}}=\mathbf{U}$, where $\mathbf{U}$ are the first $R$ left
singular vectors of $\mathbf{J}=\mathbf{MN}^{\mathsf{T}}$.

##  Appendix D Details of empirical experiments

###  D.1 Training details

####  D.1.1 Initialisation

Our models are (unless noted otherwise) initialized as follows:

$\displaystyle\tilde{\mathbf{N}}_{ij}$
$\displaystyle\sim\mathcal{U}_{[-\frac{1}{\sqrt{N}},\frac{1}{\sqrt{N}}]},$
$\displaystyle\mathbf{M}_{ij}$
$\displaystyle\sim\mathcal{U}_{[-\frac{1}{\sqrt{R}},\frac{1}{\sqrt{R}}]},$
$\displaystyle\mathbf{W}_{ij}$ $\displaystyle\sim\mathcal{N}(0,\frac{2}{R}),$
$\displaystyle\mathbf{H}_{ij}$
$\displaystyle\sim\mathcal{U}_{[-\frac{1}{\sqrt{N_{inp}}},\frac{1}{\sqrt{N_{inp}}}]},$
$\displaystyle\mathbf{h}_{i}$
$\displaystyle\sim\mathcal{U}_{[-\frac{1}{\sqrt{N}},\frac{1}{\sqrt{N}}]},$
$\displaystyle\mathbf{b}$ $\displaystyle\leftarrow\mathbf{0},$ $\displaystyle
a$ $\displaystyle\leftarrow.9,$ $\displaystyle\Sigma_{\mathbf{z}}$
$\displaystyle\leftarrow.01\mathbf{I},$ $\displaystyle\Sigma_{\mathbf{z}_{1}}$
$\displaystyle\leftarrow\mathbf{I},$ $\displaystyle\mu_{\mathbf{z}_{1}}$
$\displaystyle\leftarrow\mathbf{0},$

where $\mathbf{W}$ and $\mathbf{b}$ are the output weights and biases
respectively. For Gaussian observations we initialise
$\Sigma_{\mathbf{y}}\leftarrow.01\mathbf{I}$.

For experiments with Poisson observations, we jointly optimized a causal CNN
encoder as part of the proposal distribution. The CNN was conditioned on
observations and predicted the mean and $\log$ variance of a normal
distribution. It consisted of common initial layers consisting of 1D
convolutions, with a GeLU activation function, and a separate output
convolution for the predicted mean and ($\log$) variance. The CNN was
initialized to the Pytorch [66] defaults, except for the bias of the $\log$
variance output layer, to which we added a $\log(.01)$ term, such that the
output matches the initially predicted variance of the RNN. The exact number
of layers and channels are reported in the sections for each experiment.

For the teacher-student setups, we used as non-linearity
$\phi(\mathbf{x}_{i})=\mathsf{max}(\mathbf{x}_{i}-\mathbf{h}_{i},0)$ for both
the students and the teachers, and for all experiments with real-world data,
we used the ‘clipped’
$\phi(\mathbf{x}_{i})=\mathsf{max}(\mathbf{x}_{i}+\mathbf{h}_{i},0)-\mathsf{max}(\mathbf{x}_{i},0)$
[8].

####  D.1.2 Parameterisation

We constrain $a$ to be between $0$ and $1$ by instead optimising $\tilde{a}$
with the following (sigmoidal) parameterisation $a=\exp(-\exp(\tilde{a}))$
[67]. In experiments with the optimal proposal, we estimate the full
$\Sigma_{\mathbf{z}}$, which we constrain to be symmetric positive definite,
by optimizing a lower triangular matrix $\mathbf{C}$ such that
$\Sigma_{\mathbf{z}}=\mathbf{C}\mathbf{C}^{T}$, where we additionally
constrain the diagonal of $\mathbf{C}$ to be positive using
$\mathbf{C}_{ii}=\exp(\tilde{\mathbf{C}}_{ii}/2)$. For all diagonal
covariances, we parameterize the diagonal elements using
$\Sigma_{ii}=\exp(\tilde{\Sigma}_{ii})$. For Poisson observations, we apply a
$\mathsf{Softplus}$ function to rectify the predicted rate.

####  D.1.3 Optimisation

During training we minimise the variational SMC $\mathsf{ELBO}$ [33, 34, 35]
(Eq.7) with stochastic gradient descent, using the RAdam [68] optimiser in
Pytorch [66]. We generally use an exponentially decaying learning rate
(details under each experiment).

###  D.2 Teacher student experiments

####  D.2.1 Dataset description

We created datasets by first training ‘teacher’ RNNs to perform a task and
then generating observations by simulating the trained teacher RNNs.

For Fig. 3a,b we used code from [22]to train rank-2 RNNs to produce
oscillations, using a sine-wave with a periodicity of 50 time-steps as a
target and an additional L2 regularisation on the rates. After training, we
extracted the recurrent weights $\mathbf{M},\mathbf{N}$ and biases
$\mathbf{h}$, orthonormalized the columns of $\mathbf{M}$, and created a
dataset by simulating the model for 75 timesteps, with
$\Sigma_{\mathbf{z}}=.04\mathbf{I}$. For Fig. 3a we used $N=20$ units and
generated observations according to
$G=\mathcal{N}(\mathbf{Mz}_{t},\Sigma_{\mathbf{y}})$, with
$\Sigma_{\mathbf{y}}=.01\mathbf{I}$. Fig. 3b we used $N=40$ units and
generated observations according to
$G=\mathsf{Pois}(\mathsf{Softplus}(w\mathbf{Mz}_{t}-b)$, with $w=4$ and $b=3$.

For Fig. 3c, we followed a similar procedure but now trained the teacher RNN
on a task where it has to use input. After an initial period of $25$ time
steps, a stimulus was presented for $25$ timesteps consisting of
$[\sin(\theta)$,$\cos(\theta)]^{\mathsf{T}}$, where $\theta$ was randomly
selected every trial out of $8$ fixed angles. The RNN was tasked to produce
output that equals the transient stimulus for the next 100 time-steps. Here we
used $N=60$ units, $\Sigma_{\mathbf{z}}=.025\mathbf{I}$ and generated
observations according to
$G=\mathcal{N}(\mathbf{Mz}_{t},\Sigma_{\mathbf{y}})$, with
$\Sigma_{\mathbf{y}}=.01\mathbf{I}$. The training data for the student RNN was
included for each trial the corresponding stimulus.

####  D.2.2 Training details

The ‘student’ RNNs had $20,40,60$ units, respectively and rank $R=2$, matching
that of the teacher RNNs. For Fig. 3a, c. The observation model was a linear
Gaussian according to $G=\mathcal{N}(\mathbf{Mz}_{t},\Sigma_{\mathbf{y}})$,
and we used the optimal proposal distribution. For Fig. 3b we used
$G=\mathsf{Pois}(\mathsf{Softplus}(\mathbf{W}\mathbf{Mz}_{t}-\mathbf{B}))$,
with$\mathbf{W}$ a diagonal matrix (scaling the output of each unit
individually). For Fig. 3b, we used a causal CNN encoder as part of the
proposal distribution. It consisted of 3 layers, with kernel sizes
$(21,11,1)$, and channels $(64,64,2)$. We used (causal) circular padding.

For all three experiments, we used $k=64$ particles, batch-sizes of $10$, and
decreased the learning rate exponentially from $10^{-3}$ to $10^{-5}$. For
Fig. 3a,b, we used epochs of 400 trials and trained for 1000 and 1500 epochs,
respectively. For Fig. 3c we trained for $3000$ epochs of $96$ trials each. We
used a workstation with a NVIDIA GeForce RTX 3090 GPU for these runs. One
model took about 3 to 4 hours to finish training.

####  D.2.3 Evaluation setup

For Fig. 3 we generated long trajectories of $T=10000$ time-steps of data for
both the student and teacher RNNs. To facilitate visual comparisons between
student and teacher dynamics, we also orthonormalized the columns of the
students weights $\mathbf{M}$ after training, and for Fig. 3a,c picked signs
of the columns of $\mathbf{M}$ such that the student and teacher match (note
that after orthormalizing, the columns of $\mathbf{M}$ are equal to the non-
zero singular vectors of the full weight matrix $\mathbf{J}$, which are only
unique up to a sign flip). As noted before, this leaves the output of the
model unchanged. The autocorrelation in Fig. 3a was computed by convolving a
sequence of lag$=120$ steps of data with itself (with duration $2\times$lag),
and normalising such that lag=0 corresponds to a correlation of 1. We repeated
this for 80 sequences starting at different time-points of the whole
trajectory.

###  D.3 EEG data

####  D.3.1 Dataset description

We used openly accessible electroencephalogram (EEG) data from [42, 43] (
https://www.physionet.org/content/eegmmidb/1.0.0/, ODC-BY licence). The data
was recorded from a human subject sitting still with eyes open (session
S001R01), and was sampled at 160 Hz. Like [8], we used the full 1 minute of
recording, but unlike [8], we did not smooth the data (but just standardized
the data). Thus, to compare our performance to [8], who ran their evaluation
using the smoothed data, we smoothed our generated samples equivalently, using
a Hann filter with a window length of 15-time bins, so that we can also
compare our samples to the smoothed data.

####  D.3.2 Training details

We used $N=512$ units, and rank $R=3$. The observation model was a linear
Gaussian conditioned on the hidden state and we used the optimal proposal
distribution. We trained for 1000 epochs consisting of 50 batches of size of
10, and $k=10$ particles. The learning rate was decreased exponentially from
$10^{-3}$ to $10^{-6}$. Models were trained using NVIDIA RTX 2080 TI GPUs on a
compute cluster. A single model took between 4 and 5 hours to finish training.

####  D.3.3 Evaluation setup

We used our RNN to generate one long trajectory of $T=9760$ steps of data,
$\mathbf{y}_{t}$ (after discarding the first 2440 steps), which we compare to
the EEG data, $\hat{\mathbf{y}}_{t}$, using two evaluation measures from [8,
24] (using code from https://github.com/DurstewitzLab/GTF-shPLRNN, GNU General
Public License):

$\mathbf{D_{\mathsf{stsp}}}$: This is an estimate of the $\mathsf{KL}$
divergence between the ground truth and generated states. To compute this, we
obtained kernel density estimates of the probability density functions (over
states, not time), using a Gaussian kernel with standard deviation $\sigma=1$.
We get for the EEG data:
$\hat{p}(\mathbf{y})=\frac{1}{T}\sum_{t=1}^{T}\mathcal{N}(\hat{\mathbf{y}}_{t},\mathbf{I})$,
and for the generated data
$\hat{q}(\mathbf{y})=\frac{1}{T}\sum_{t=1}^{T}\mathcal{N}(\mathbf{y}_{t},\mathbf{I})$.
We then used the following Monte Carlo estimate of the $\mathsf{KL}$
divergence:
$D_{\mathsf{stsp}}\approx\frac{1}{n}\sum\log\frac{\hat{p}(\hat{\mathbf{y}}^{i})}{\hat{q}(\hat{\mathbf{y}}^{i})}$,
using $n=1000$ samples $\hat{\mathbf{y}}^{i}$ drawn randomly from the EEG
data.

$\mathbf{D_{H}}$: This is an estimate of the difference in power spectra
between the ground truth and generated states. We first computed for each data
dimension the spectra $\hat{\mathbf{y}}^{i}_{\omega}$,
$\mathbf{y}^{i}_{\omega}$ for the EEG and generated data, respectively. We
used a Fast Fourier Transform, smoothed the estimates with a Gaussian kernel
with standard deviation $\sigma=20$, and normalized the spectra so they sum to
1. We computed the mean of the Hellinger distances between the spectra:
$D_{H}=\frac{1}{64}\sum^{64}_{i}\frac{1}{\sqrt{2}}\lVert\sqrt{\hat{\mathbf{y}}^{i}_{\omega}}-\sqrt{\mathbf{y}^{i}_{\omega}}\rVert$.

###  D.4 Hippocampus HC-2

####  D.4.1 Dataset description

We used openly accessible neurophysiological data recorded from layer CA1 of
the right dorsal hippocampus [47, 48] (https://crcns.org/data-
sets/hc/hc-2/about-hc-2. Signals were recorded as the rats engaged in an open
field task, chasing drops of water or pieces of food that were randomly
placed. We used the session ec013.527 from rat ID ec13, which is approximately
1062 seconds long. From 37 units (neurons) we used 21 neurons that have
maximal spike counts, discarding the rest of the comparatively silent neurons.
We binned the spike data to 10ms. We used the first 80 percent of the data for
training, and the rest was saved for testing purposes.

####  D.4.2 Training details

We used $N$ = 512 units, and rank $R$ = 3 for the run that was used in our
Fig. 5. We used a causal CNN encoder as part of the proposal distribution,
which consisted of 3 layers with kernel sizes (150, 11, 1), with (64, 64, 3)
channels. During our study, we swept over multiple ranks and found that theta
oscillations consistently emerged from rank 3 onwards, after which
reconstruction accuracy was relatively stable. For each rank, we used three
different seeds and two different first layer sizes for the encoder, 25 or
150. The duration of a randomly sampled trial (sequence length) from the whole
data was 94 time steps when the first layer size was 25, and 219 when the
first layer size was 150. We, however, also found that the choice of the
duration did not affect the results much. We trained the model using 3000
epochs, each epoch consisting of 3000 trials with 64 batches and $k$ = 64
particles. The learning rate was decreased exponentially from $10^{-3}$ to
$10^{-6}$. A single model took approximately 21 hours to finish training on a
NVIDIA RTX 2080 TI GPU on a compute cluster.

####  D.4.3 Evaluation setup

We used our RNN to generate data that matches the duration of the test data,
which is 20810 time steps ($\sim$208 s) (after discarding the first 1000
steps). We compare different spike statistics of generated data with test
data, and for comparison purposes, we also compared the same statistics
measurements between train and test data as well. We calculated the mean
firing rate of each neuron, mean of ISI distributions, and pairwise
correlations. We used a band-pass filter 1-40 Hz for the latents and the LFP
signal before calculating the powerspectrogram (Fig.5e).

###  D.5 Hippocampus HC-11

####  D.5.1 Dataset description

We used openly accessible neurophysiological data recorded from hippocampal
CA1 region [50, 51, 52] (https://crcns.org/data-sets/hc/hc-11/about-hc-11). We
used the subset of the dataset called the maze epoch, where a rat was running
on a 1.6-meter linear track, with rewards located at each end (left and
right). Throughout this task, neural activity was recorded from 120 identified
pyramidal neurons. As in [13], we only used 60 neurons that had sufficient
activity and discarded rest of the units. We used code from [53]
(https://github.com/zhd96/pi-vae) to preprocess the spike data, and only use
data corresponding to the rat running and the location data being available.
We used 25ms bins.

####  D.5.2 Training details

We used $N$ = 512 units, and rank $R$ = 4. We used the causal CNN encoder with
zero padding, 3 layers (24, 11, 1), and (64, 64, 4) channels. The model is
trained for 3000 epochs, each epoch having 3000 trials with a sequence length
of 94 time bins (2.35 s), using batch size 64 and $k$ = 64 particles. The
learning rate was decreased exponentially from $10^{-3}$ to $10^{-6}$. A
single model took approximately 21 hours to finish training on a NVIDIA RTX
2080 TI GPU on a compute cluster.

####  D.5.3 Evaluation setup

We used our RNN to generate data that matches the duration of the test data,
which is 4289 time steps ($\sim$107 s) (after discarding the first 1000
steps). We calculated the mean firing rate of each neuron, coefficient of
variations of ISI distributions, and pairwise correlations. We fit a simple
linear regressor to posterior latents to predict the location data, and used
this regression model to predict the test data for the location.

###  D.6 Monkey Reach

####  D.6.1 Dataset description

We used the publicly available MC_Maze dataset from the Neural Latents
Benchmark (NLB) [55] (https://dandiarchive.org/dandiset/000128, CC-BY-4.0
licence). The data were recorded from a macaque performing a delayed center-
out reaching task with barriers, resulting in a variety of straight and curved
reaches. For simplicity, we took only the trials with no barriers and thus
straight reach trajectories, resulting in 592 training trials and 197 test
trials. We binned the data at 20 ms and aligned each trial from 250 ms before
to 450 ms after movement onset.

To create conditioning inputs for the model, we took the x and y coordinates
of the target position for each trial and scaled them to be between $-1$ and
$1$. We then provide this scaled target position as constant context input to
the RNN for the duration of the trial.

####  D.6.2 Training details

We ran a random search of 30 different models with rank $r\in{3,4,5,6}$ and
particle number $k\in{16,32,64}$. All models had $512$ units and used a causal
CNN encoder with kernel sizes $(14,4,2)$ and channels $(128,64,r)$. We used
(causal) reflect padding. We trained each model for up to 2000 epochs,
terminating training early if no improvement was seen for 50 epochs. Each
model took around $3$ to $4$ hours to train on an NVIDIA RTX 2080 TI GPU on a
compute cluster. Seeing that a rank of 5 was sufficient for velocity decoding
$R^{2}\approx 0.9$, we took the best-performing rank-5 model for subsequent
analyses.

####  D.6.3 Evaluation setup

For qualitative evaluation of replication of cross-condition differences, we
grouped the reach targets in the data into 7 conditions, one at each corner
and the midpoint of each edge of the rectangular reach plane, excluding the
midpoint directly at the bottom. We then generated data from the model RNN
using conditioning inputs from the test trials of the real data. Then, for the
test data and the model-generated data, we computed mean firing rate and
inter-spike interval for each neuron for each condition. We then computed
correlation distance ($1-r$, where $r$ is the Pearson correlation coefficient)
on the neuron statistics between conditions in the test data and model-
generated data.

For generation of data for Fig. 7d,e, we selected target locations by choosing
angles from $0$ to $360^{\circ}$, evenly spaced by $22.5^{\circ}$, and
determined the corresponding reach endpoint on a square spanning from
$(-1,-1)$ to $(1,1)$. We then constructed conditioning inputs similar to the
real data using these target locations and simulated the RNN with them. To
decode the reaches, we used a linear decoder trained from inferred firing
rates to reach velocity from the real data.

###  D.7 Neural Latents Benchmark

####  D.7.1 Dataset description

We again used the publicly available MC_Maze dataset from NLB (see Supplement
D.6.1). We resampled the data to 20 ms bin size and followed the standard data
preprocessing procedures for the benchmark, as described in [55].

####  D.7.2 Training details

We ran a random search of 30 different models with varying rank from 12 to 40
and particle number $k\in{16,32,64}$. All models had $512$ units and used a
causal CNN encoder with kernel sizes $(14,4,2)$ and channels $(128,64,36)$,
and reflect padding. We trained each model for up to 2000 epochs, terminating
training early if no improvement was seen for 50 epochs. Each model took
around $10$ to $12$ hours to train on an NVIDIA RTX 2080 TI GPU on a compute
cluster.

Because the primary task of the benchmark is co-smoothing, i.e., prediction of
held-out neuron firing rates from held-in neurons, we provide the encoder with
only the activity of held-in neurons. However, the observation likelihood
component of the ELBO is computed on all neurons, held-in and held-out.

After training, we selected the model with the best co-smoothing score on the
validation split and submitted its predictions to the benchmark for the final
evaluation.

####  D.7.3 Evaluation setup

Automated evaluation was performed on the benchmark platform, as described in
[55].

We used for the prediction at timestep $t$, the expected Poisson rate of held-
out neurons, conditioned on the activity of held-in neurons at the current and
previous timesteps, by making use of the filtering Posterior (Eq. 3). We
averaged over $32$ sets of trajectories with $192$ particles each.