# Expressive dynamics models with nonlinear injective readouts enable reliable
recovery of latent features from neural activity

Christopher Versteeg1 Andrew R. Sedler1,2 Jonathan D. McCart1,2
Chethan Pandarinath1,2

1 Wallace H. Coulter Department of Biomedical Engineering
Emory University and Georgia Institute of Technology
Atlanta, GA, USA

2 Center for Machine Learning
Georgia Institute of Technology
Atlanta, GA, USA

###### Abstract

The advent of large-scale neural recordings has enabled new approaches that
aim to discover the computational mechanisms of neural circuits by
understanding the rules that govern how their state evolves over time. While
these neural dynamics cannot be directly measured, they can typically be
approximated by low-dimensional models in a latent space. How these models
represent the mapping from latent space to neural space can affect the
interpretability of the latent representation. We show that typical choices
for this mapping (e.g., linear or MLP) often lack the property of injectivity,
meaning that changes in latent state are not obligated to affect activity in
the neural space. During training, non-injective readouts incentivize the
invention of dynamics that misrepresent the underlying system and the
computation it performs. Combining our injective Flow readout with prior work
on interpretable latent dynamics models, we created the Ordinary Differential
equations autoencoder with Injective Nonlinear readout (ODIN), which learns to
capture latent dynamical systems that are nonlinearly embedded into observed
neural activity via an approximately injective nonlinear mapping. We show that
ODIN can recover nonlinearly embedded systems from simulated neural activity,
even when the nature of the system and embedding are unknown. Additionally, we
show that ODIN enables the unsupervised recovery of underlying dynamical
features (e.g., fixed points) and embedding geometry. When applied to
biological neural recordings, ODIN can reconstruct neural activity with
comparable accuracy to previous state-of-the-art methods while using
substantially fewer latent dimensions. Overall, ODIN’s accuracy in recovering
ground-truth latent features and ability to accurately reconstruct neural
activity with low dimensionality make it a promising method for distilling
interpretable dynamics that can help explain neural computation.

##  1 Introduction

Recent evidence has shown that when artificial recurrent neural networks are
trained to perform tasks, the rules that govern how the internal activity
evolves over time (i.e., the network dynamics) can provide insight into how
the network performs the underlying computation [1, 2, 3, 4]. Given the
conceptual similarities between artificial neural networks and biological
neural circuits, it may be possible to apply these same dynamical analyses to
brain activity to gain insight into how neural circuits perform complex
sensory, cognitive, and motor processes [5, 6, 7]. However, unlike in
artificial networks, we cannot easily interrogate the dynamics of biological
neural circuits and must first estimate them from observed neural activity.

Fortunately, advances in recording technology have dramatically increased the
number of neurons that can be simultaneously recorded, providing ample data
for novel population-level analyses of neural activity [8, 9, 10]. In these
datasets, the activity of hundreds or thousands of neurons can often be
captured by relatively low-dimensional subspaces [11], orders-of-magnitude
smaller than the total number of neurons. Neural activity in these latent
spaces seems to evolve according to consistent sets of rules (i.e., latent
dynamics) [12, 6]. Assuming no external inputs, these rules can be expressed
mathematically as:

$\displaystyle\mathbf{z}_{t+1}$
$\displaystyle=\mathbf{z}_{t}+f(\mathbf{z}_{t})$ (1)
$\displaystyle\mathbf{y}_{t}$ $\displaystyle=\exp g(\mathbf{z}_{t})$ (2)
$\displaystyle\mathbf{x}_{t}$
$\displaystyle\sim\text{Poisson}(\mathbf{y}_{t})$ (3)

where $\mathbf{z}_{t}\in\mathbb{R}^{D}$ represents the latent state at time
$t$, $f(\mathbf{\cdot}):\mathbb{R}^{D}\rightarrow\mathbb{R}^{D}$ is the vector
field governing the dynamical system, $\mathbf{y}_{t}\in\mathbb{R}^{N}$
denotes the firing rates of the $N$ neurons,
$g(\cdot):\mathbb{R}^{D}\rightarrow\mathbb{R}^{N}$ maps latent activity into
log-firing rates, and $\mathbf{x}_{t}\in\mathbb{R}^{N}$ denotes the observed
spike counts at time $t$, assuming the spiking activity follows a Poisson
distribution with time-varying rates given at each moment $t$ by
$\mathbf{y}_{t}$.

Unfortunately, any latent system can be equivalently described by many
combinations of dynamics $f$ and embeddings $g$, which makes the search for a
unique latent system futile. However, versions of a latent system’s dynamics
$f$ and embedding $g$ that are less complex and use fewer latent dimensions
can be easier to interpret than alternative representations that are more
complex and/or higher-dimensional. Models of latent dynamics that can discover
simple and low-dimensional representations will make it easier to link latent
dynamics to neural computation.

A popular approach to estimate neural dynamics [13, 14, 15] is to use neural
population dynamics models (NPDMs), which model neural activity as a latent
dynamical system embedded into neural activity. We refer to the components of
an NPDM that learn the dynamics and embedding as the generator $\hat{f}$ and
the readout $\hat{g}$, respectively. When modeling neural activity, the
generator and readout are jointly trained to infer firing rates
$\mathbf{\hat{y}}$ that maximize the likelihood of the observed neural
activity $\mathbf{x}$.

Using NPDMs to estimate underlying dynamics and embedding implicitly assumes
that good reconstruction performance (i.e.,
$\hat{\mathbf{x}}\approx\mathbf{x}$) implies interpretable estimates of the
underlying system (i.e., $\hat{\mathbf{z}}\approx\mathbf{z}$, $\hat{f}\approx
f$, $\hat{g}\approx g$). However, recent work has shown that when the state
dimensionality of the generator $\hat{D}$ is larger than a system’s latent
dimensionality $D$, high reconstruction performance may actually correspond to
estimates of the latent system that are overly complex or misleading and
therefore harder to interpret [15]. At present, reconstruction performance is
seemingly an unreliable indicator for the interpretability of the learned
dynamics.

This vulnerability to learning overly complex latent features might emerge
from the fact that, without constraints on the readout $\hat{g}$, changes in
the latent state are not obligated to have an effect on predicted neural
activity. Thus, NPDMs can be rewarded for inventing latent activity that
boosts reconstruction performance, even if that latent activity has no direct
correspondence to neural activity. A potential solution is to make $\hat{g}$
injective, which obligates all latent activity to affect neural
reconstruction. This would penalize any latent activity that is not reflected
in the observed neural activity, thereby putting pressure on the generator
$\hat{f}$ and readout $\hat{g}$ to learn a more interpretable (i.e., simpler
and lower dimensional) representation of the underlying system.

In addition, most previously used readouts $\hat{g}$ were not expressive
enough to model diverse mappings from latent space to neural space, assuming
the embedding $g$ to be a relatively simple (often linear) transformation
(though there are exceptions [16, 17, 18]). Capturing nonlinear embeddings is
important because neural activity often lives on a lower-dimensional manifold
that is nonlinearly embedded into the higher-dimensional neural space [7].
Therefore, assumptions of linearity are likely to prevent NPDMs from capturing
dynamics in their simplest and lowest-dimensional form, making them less
interpretable than the latent features learned by NPDMs that can approximate
these nonlinearities.

To address these challenges, we propose a novel architecture called the
Ordinary Differential equation autoencoder with Injective Nonlinear readout
(ODIN), which implements $\hat{f}$ using a Neural ODE (NODE [19]) and
$\hat{g}$ using a network inspired by invertible ResNets [20, 21, 22, 19, 23].
ODIN approximates an injective nonlinear mapping between latent states and
neural activity, obligating all latent state variance to appear in the
predicted neural activity and penalizing the model for using excessively
complex or high-dimensional dynamics to model the underlying system. On
synthetic data, ODIN learns representations of the latent system that are more
interpretable, with simpler and lower-dimensional latent activity and
dynamical features (e.g., fixed points) than alternative readouts. ODIN’s
interpretability is also more robust to overestimates of latent dimensionality
and can recover the nonlinear embedding of synthetic data that evolves on a
simulated manifold. When applied to neural activity from a monkey performing a
reaching task with obstacles, ODIN reconstructs neural activity comparably to
state-of-the-art recurrent neural network (RNN)-based models while requiring
far fewer latent state dimensions. In summary, ODIN estimates interpretable
latent features from synthetic data and has high reconstruction performance on
biological neural recordings, making it a promising tool for understanding how
the brain performs computation.

##  2 Related Work

Many previous models have attempted to understand neural activity through the
lens of neural dynamics. Early efforts limited model complexity by
constraining both $\hat{f}$ and $\hat{g}$ to be linear [24, 25, 26]. While
these models were relatively straightforward to analyze, they often failed to
adequately explain neural activity patterns [27].

Other approaches increased the expressiveness of the modeled dynamics
$\hat{f}$. RNNs can learn to approximate complex nonlinear dynamics, and have
been shown to substantially outperform linear dynamics models in
reconstructing neural activity [27]. Unfortunately, RNNs implicitly couple the
capacity of the model to the latent state dimensionality, meaning their
ability to model complex dynamics relies on having a high-dimensional latent
state. In contrast, NODEs can model arbitrarily complex dynamics of embedded
dynamical systems at the dimensionality of the system [19, 15]. On synthetic
data, NODEs have been shown to recover dynamics more accurately than RNN-based
methods [28, 15]. In contrast to our approach, previous NODE-based models used
a linear readout $\hat{g}$ that lacks injectivity. This can make the accuracy
of estimated latent activity vulnerable to overestimates of the latent
dimensionality (i.e., when $\hat{D}>D$) and/or fail to capture potential
nonlinearities in the embedding $g$.

Early efforts to allow greater flexibility in $\hat{g}$ preserved linearity in
$\hat{f}$, using feed-forward neural networks to nonlinearly embed linear
dynamical systems in high-dimensional neural firing rates [16]. More recently,
models have used Gaussian processes to approximate nonlinear mappings from
latent state to neural firing with tuning curves [17]. Other models have
combined nonlinear dynamics models and nonlinear embeddings for applications
in behavioral tracking [29] and neural reconstruction [18]. Additional
approaches extend these methods to incorporate alternative noise models that
may better reflect the underlying firing properties of neurons [16, 30]. While
nonlinear, the readouts of these models lacked injectivity in their mapping
from latent activity to neural activity.

Many alternative models seek to capture interpretable latent features of a
system from observations. One popular approach uses a sparsity penalty on a
high-dimensional basis set to derive a sparse symbolic estimate of the
governing equations for the system [31]. However, it is unclear whether such
sparse symbolic representation is necessarily a benefit when modeling dynamics
in the brain. Another recent model uses contrastive loss and auxiliary
behavioral variables to learn low-dimensional representations of latent
activity [32]. This approach does not have an explicit dynamics model,
however, so is not amenable to the dynamical analyses performed in this
manuscript.

Normalizing flows – a type of invertible neural network – have recently become
a staple for generative modeling and density estimation [20, 23]. Some latent
variable models have used invertible networks to approximate the mapping from
the latent space to neural activity [33] or for generative models of visual
cortex activity [34]. To allow this mapping to change dimensionality between
the latent space and neural activity, some of these models used a zero-padding
procedure similar to the padding used in this manuscript (see Section 3.3.1),
which makes the transformation injective rather than invertible [33, 23].
However, these previous approaches did not have explicit dynamics models,
making our study, to our knowledge, the first to test whether injective
readouts can improve the interpretability of neural population dynamics
models.

##  3 Methods

###  3.1 Synthetic Neural Data

To determine whether different models can distill an interpretable latent
system from observed population activity, we first used reference datasets
that were generated using simple ground-truth dynamics $f$ and embedding $g$.
Our synthetic test cases emulate the empirical properties of neural systems,
specifically low-dimensional latent dynamics observed through noisy spiking
activity [13, 35, 36, 37]. We sampled latent trajectories from the Arneodo
system ($f$, $D=3$) and nonlinearly embedded these trajectories into neural
activity via an embedding $g$ . We consider models that can recover the
dynamics $f$ and embedding $g$ used to generate these data as providing an
interpretable description of the latent system and its relation to the neural
activity. Additional detail on data generation, models, and metrics can be
found in the Supplementary Material.

Unless otherwise noted, we generated activations for $N$ neurons ($N=12$) by
projecting the simulated latent trajectories $\mathbf{Z}$ through a $3\times
N$ matrix whose columns were random encoding vectors with elements sampled
from a uniform distribution $U[-0.5,0.5]$ (Fig. 1A, left). We standardized
these activations to have zero mean and unit variance and applied a different
scaled sigmoid function to each neuron, yielding a matrix of non-negative
time-varying firing rates $\mathbf{Y}$. The scaling of each sigmoid function
was evenly spaced on a logarithmic scale between $10^{0.2}$ and $10$. This
process created a diverse set of activation functions ranging from quasi-
linear to nearly step-function-like behavior (Fig. 1A, Activation Functions).
For one experiment, we used the standard linear-exponential activation
function, as described in previous work [15], instead of the scaled sigmoid.

Figure 1:  A) Synthetic neural data generation (left to right). Trajectories
from the Arneodo system are projected onto random encoding vectors to compute
activations at each timepoint. A scaled sigmoid nonlinearity is applied to
convert the activations into firing rates. B) Zero-padded latent dynamics
(green) are reversibly warped into higher-dimensional neural activity space
(blue). C) The Flow readout maps from latent space to neural space by applying
a sequence of $K$ small updates (parameterized by an MLP, bottom). The reverse
pass of the Flow maps from neural space to latent space and is implemented by
serial subtraction of updates from the same MLP.

We simulated spiking activity $\mathbf{X}$ by sampling from inhomogeneous
Poisson processes with time-varying rate parameters equal to the firing rates
$\mathbf{Y}$ of the simulated neurons (Fig. 1A, right). We randomly split
70-point segments of these trials into training and validation datasets
(training and validation proportions were 0.8 and 0.2, respectively).

###  3.2 Biological Neural Data

We evaluated how well our model could reconstruct biological neural activity
on a well-characterized dataset [38] included in the Neural Latents Benchmark
(NLB) [27]. This dataset is composed of single-unit recordings from primary
and pre-motor cortices of a monkey performing a visually-guided reaching task
with obstacles, referred to as the Maze task. Trials were trimmed to the
window [-250, 350] ms relative to movement onset, and spiking activity was
binned at 20 ms. To compare the reconstruction performance of our model
directly against the benchmark, we split the neural activity into held-in and
held-out neurons, comprising 137 and 35 neurons, respectively, using the same
sets of neurons as were used to assess models for the NLB leaderboard.

###  3.3 Model Architecture

We used three sequential autoencoder (SAE) variants in this study, with the
main difference being the choice of readout module, $\hat{g}(\cdot)$. In
brief, a sequence of binned spike counts $\mathbf{x}_{1:T}$ was passed through
a bidirectional GRU encoder, whose final hidden states were converted to an
initial condition $\hat{\mathbf{z}}_{0}$ via a mapping $\phi(\cdot)$. A
modified NODE generator unrolled the initial condition into time-varying
latent states $\hat{\mathbf{z}}_{1:T}$. These were subsequently mapped to
inferred rates via the readout
$\hat{g}(\cdot)\in\\{\text{Linear},\text{MLP},\text{Flow}\\}$. All models were
trained for a fixed number of epochs to infer firing rates
$\hat{\mathbf{y}}_{1:T}$ that minimize the negative Poisson log-likelihood of
the observed spikes $\mathbf{x}_{1:T}$.

$\displaystyle\mathbf{h}_{T}=\left[\mathbf{h}_{fwd}\big{|}\mathbf{h}_{bwd}\right]=\text{BiGRU}(\mathbf{x}_{1:T})$
(4) $\displaystyle\hat{\mathbf{z}}_{0}=\phi(\mathbf{h}_{T})$ (5)
$\displaystyle\hat{\mathbf{z}}_{t+1}=\hat{\mathbf{z}}_{t}+\alpha\cdot\text{MLP}(\hat{\mathbf{z}}_{t})$
(6) $\displaystyle\hat{\mathbf{y}}_{t}=\exp\hat{g}(\hat{\mathbf{z}}_{t})$ (7)

For models with Linear and MLP readouts, $\phi(\cdot)$ was a linear map to
$\mathbb{R}^{\hat{D}}$. For models with Flow readouts, $\phi(\cdot)$ was a
linear map to $\mathbb{R}^{N}$ followed by the reverse pass of the Flow (see
Section 3.3.1). We unrolled the NODE using Euler’s method with a fixed step
size equal to the bin width and trained using standard backpropagation for
efficiency. A scaling factor ($\alpha=0.1)$ was applied to the output of the
NODE’s MLP to stabilize the dynamics during early training. Readouts were
implemented as either a single linear layer (Linear), an MLP with two 150-unit
ReLU hidden layers (MLP), or a Flow readout (Flow) which contains an MLP with
two 150-unit ReLU hidden layers. We refer to these three models as Linear-
NODE, MLP-NODE, and ODIN, respectively.

####  3.3.1 Flow Readout

The Flow readout resembles a simplified invertible ResNet [23]. Flow learns a
vector field that can reversibly transform data between latent and neural
representations (Figure 1B). The Flow readout has three steps: first, we
increase the dimensionality of the latent activity $\mathbf{z}_{t}$ to match
that of the neural activity by padding the latent state with zeros. This
corresponds to an initial estimate of the log-firing rates,
$\log\hat{\mathbf{y}}_{t,0}$. Note that zero-padding makes our mapping
injective rather than fully invertible (see [23, 33]). The Flow network then
uses an MLP to iteratively refine $\log\hat{\mathbf{y}}_{t,k}$ over $K$ steps
($K=20$) after which we apply an exponential to produce the final firing rate
predictions, $\hat{\mathbf{y}}_{t}$. A scaling factor ($\beta=0.1)$ was
applied to the output of the Flow’s MLP, which prevents the embedding from
becoming unstable during the early training period.

$\displaystyle\log\hat{\mathbf{y}}_{t,0}=[\hat{\mathbf{z}}_{t}|\mathbf{0}]^{T}$
(8)
$\displaystyle\log\hat{\mathbf{y}}_{t,k+1}=\log\hat{\mathbf{y}}_{t,k}+\beta\cdot\text{MLP}(\log\hat{\mathbf{y}}_{t,k})$
(9)
$\displaystyle\hat{g}\left(\hat{\mathbf{z}}_{t}\right)=\log\hat{\mathbf{y}}_{t,K}=\log\hat{\mathbf{y}}_{t}$
(10)

We also use a reverse pass of the Flow to transform the output of the encoders
to initial conditions in the latent space via $\phi(\cdot)$, approximating the
inverse function $\hat{g}^{-1}$. Our method subtracts the output of the MLP
from the state rather than adding it as in the forward mode (Fig 1C), a
simplified version of the fixed-point iteration procedure described in [23].
We then trim the excess dimensions to recover $\hat{z}\in\mathbb{R}^{\hat{D}}$
(in effect, removing the zero-padding dimensions).

$\displaystyle\log\hat{\mathbf{y}}_{t,k-1}=\log\hat{\mathbf{y}}_{t,k}-\beta\cdot\text{MLP}(\log\hat{\mathbf{y}}_{t,k})$
(11)
$\displaystyle\hat{g}^{-1}\left(\log\hat{\mathbf{y}}_{t}\right)=[\log\hat{y}_{t,0,1},\dots,\log\hat{y}_{t,0,\hat{D}}]^{T}=\hat{\mathbf{z}}_{t}$
(12)

The Flow mapping is only guaranteed to be injective if changes in the output
of the MLP are sufficiently small relative to changes in the input (i.e.,
Lipschitz constant for the MLP that is strictly less than 1) [23]. The model
can be made fully injective by either restricting the weights of the MLP
(e.g., spectral norm [39]), or using a variable step-size ODE solver that can
prevent crossing trajectories (e.g., continuous normalizing flows [19]). In
practice, we found that using a moderate number of steps allows Flow to
preserve approximate injectivity of the readout at all tested dimensionalities
(Supp. Fig. S2).

###  3.4 Metrics and characterization of dynamics

We assessed model performance in five domains: 1) reconstruction performance,
2) latent accuracy, 3) dynamical accuracy, 4) embedding accuracy, and 5)
readout injectivity. All metrics were evaluated on validation data.
Critically, on biological data without a ground-truth system, only the
reconstruction performance and readout injectivity can be assessed, since all
the other metrics rely on full observability of the underlying system.
Therefore, we need models for which good performance on the observable metrics
(reconstruction, injectivity) implies good performance on the unobservable
metrics (latent, dynamical, and embedding accuracy).

Reconstruction performance for the synthetic data was assessed using two key
metrics. The first, spike negative log-likelihood (Spike NLL), was defined as
the Poisson NLL employed during model training. The second, Rate $R^{2}$, was
the coefficient of determination between the inferred and true firing rates,
averaged across neurons. We used Spike NLL to assess how well the inferred
rates explain the spiking activity, while Rate $R^{2}$ reflects the model’s
ability to find the true firing rates. These metrics quantify how well the
model captures the embedded system’s dynamics (i.e., that $\hat{f},\hat{g}$
captures the system described by $f,g$), but give no indication of the
interpretability of the learned latent representation (i.e., that the learned
$\hat{f},\hat{g}$ are simple and low-dimensional).

For the biological neural data, we measured model performance using two
metrics from the Neural Latents Benchmark (NLB) [27], co-smoothing bits-per-
spike (co-bps) and velocity decoding performance on predicted firing rates
(Vel $R^{2}$). co-bps is a measure of reconstruction performance that
quantifies how well the model predicts the spiking of the held-out neurons,
while Vel $R^{2}$ quantifies how well the denoised rates can predict the
monkey’s hand velocity during the reach. We have no way to directly assess
embedding, latent, or dynamical accuracy because they are unobserved in most
biological datasets.

To determine whether a model’s inferred latent activity contains features that
are not in the simulated latent activity, we used a previously published
metric called the State $R^{2}$ [15]. State $R^{2}$ is defined as the
coefficient of determination ($R^{2}$) of a linear regression from simulated
latent trajectories $\mathbf{z}$ to the inferred latent trajectories
$\hat{\mathbf{z}}$. State $R^{2}$ will be low if the inferred latent
trajectories contain features that cannot be explained by an affine
transformation of the true latent trajectories. Importantly, State $R^{2}$
alone cannot ensure latent accuracy. This is because a model can achieve high
State $R^{2}$ trivially if the inferred latent activity $\hat{\mathbf{z}}$ is
a low-dimensional projection of the simulated activity $\mathbf{z}$.
Therefore, only models that have _both_ good reconstruction performance (Spike
NLL, Rate $R^{2}$) and State $R^{2}$ can be said to accurately reflect the
simulated latent dynamics without extra features that make the model harder to
interpret (i.e., $\hat{\mathbf{z}}\approx\mathbf{z}$).

As a direct comparison of the estimated dynamics $\hat{f}$ to the simulated
dynamics $f$, we extracted the fixed-point (FP) structure from our trained
models and compared it to the FP structure of the underlying system. We used
previously published FP-finding techniques [40] to identify regions of the
generator’s dynamics where the magnitude of the vector field was close to
zero, calling this set of locations the putative FPs. We linearized the
dynamics around the FPs and computed the eigenvalues of the Jacobian of
$\hat{f}$ to characterize each FP. Capturing FP location and character gives
an indication of how closely the estimated dynamics resemble the simulated
dynamics (i.e., $\hat{f}\approx f$).

To determine how well our embedding $\hat{g}$ captures the simulated embedding
$g$, we projected the encoding vectors used to generate the synthetic neural
activity from the ground-truth system into our model’s latent space using the
same affine transformation from ground-truth latent activity to inferred
latent activity that was used to compute State $R^{2}$. We projected the
inferred latent activity onto each neuron’s affine-transformed encoding vector
to find the predicted activation of each synthetic neuron. We then related the
predicted firing rates of each neuron to its corresponding activations to
derive an estimate of each neuron’s activation function. Because the inferred
latent activity is arbitrarily scaled/translated relative to the true latent
activity, we fit an affine transformation from the predicted activation
function to the ground-truth activation function. The coefficient of
determination $R^{2}$ of this fit quantifies how well our models were able to
recover the synthetic warping applied to each neuron (i.e., $\hat{g}\approx
g$).

We compared the injectivity of the Flow readout to Linear and MLP readouts
using effective rank [41] and cycle-consistency, respectively. Effective rank
quantifies the number of significant singular values in a Linear readout,
while cycle-consistency quantifies how well the inferred latent activity
$\hat{\mathbf{z}}$ can be recovered from the predicted log-firing rates
$\log\hat{\mathbf{y}}$.

Figure 2:  ODIN recovers latent activity more accurately than alternative
models and is robust to overestimates of latent dimensionality. A) Diagram of
models tested, including Linear-NODE (green), MLP-NODE (orange), ODIN (red).
B) Inferred latent activity of representative model at each state
dimensionality $\hat{D}$. True latent activity (affine-transformed to overlay
inferred latent activity) shown in light blue. C) All: Model metrics as a
function of $\hat{D}$. Shaded areas represent one standard deviation around
the mean. Dashed vertical line indicates $\hat{D}=3$ Top: Spike NLL, Middle:
Rate $R^{2}$, Bottom: State $R^{2}$.

##  4 Results

###  4.1 Finding interpretable latent activity across state dimensionalities
with ODIN

As the latent dimensionality $D$ is unknown for biological datasets, we wanted
to test how robust each model was to choices of state dimensionality
$\hat{D}$. We trained Linear/MLP -NODE, and ODIN (Fig 2A) to reconstruct
synthetic neural activity from the Arneodo system [42] and compared
reconstruction performance (i.e. Spike NLL and Rate $R^{2}$) and latent
recovery (i.e. State $R^{2}$) as functions of the dimensionality $\hat{D}$ of
the state space. We trained 5 different random seeds for each of the 3 model
types and 5 state dimensionalities (75 total models, model hyperparameters in
Supp. Table 1, representative hyperparameter sweeps in Supp. Fig. S1).

First, we observed that latent activity inferred by Linear-NODE did not
closely resemble the simulated latent activity, with all tested
dimensionalities performing worse than either ODIN or the MLP-NODE at
$\hat{D}$ = 3 (Fig 2B,C, mean State $R^{2}$ = 0.70 for Linear-NODE vs. 0.89,
0.93 for MLP-NODE, ODIN respectively). We also found that Linear-NODE required
many more dimensions to reach the peak reconstruction performance (Fig 2C,
Rate $R^{2}$). These results demonstrate that models that are unable to
account for nonlinear embeddings are vulnerable to learning more complex and
higher dimensional dynamics than those learned by models with nonlinear
readouts.

Next, we compared ODIN to MLP-NODE and found that at the correct
dimensionality ($\hat{D}=3$), these models had similar performance for both
reconstruction and latent recovery. However, as the dimensionality increased
beyond the true dimensionality ($\hat{D}>3$), the latent recovery of the MLP-
NODE degraded rapidly while ODIN’s latent recovery remained high (Fig 2C, as
$\hat{D}>3$). As the true latent dimensionality $D$ is usually unknown, NPDMs
with non-injective readouts (like MLPs) may be predisposed to learning
misleading latent activity that can make it more difficult to interpret
biological datasets.

###  4.2 Common readouts learn non-injective mappings from latent activity to
firing rates

We then sought to assess the injectivity of different readouts. First, we used
effective rank [41] to quantify the injectivity of our Linear readouts. We
trained 5 Linear-NODE models at a range of state dimensionalities
($\hat{D}=3,5,8,10$) to reconstruct simulated neural activity from Arneodo
that was linearly embedded into 12D neural space. We found that while
reconstruction performance was optimal when $\hat{D}>3$ (Supp. Fig. S3), the
effective rank of these best-reconstructing models never exceeded 4 (mean
erank = 3.74 at $\hat{D}=10$). This means that for the largest Linear-NODE
models, around 6 of 10 latent dimensions had no effect on reconstructed log-
rates. The fact that linear readouts learn mappings with low effective rank,
coupled with improved reconstruction performance when $\hat{D}>3$ suggests
that the Linear readouts utilize non-injectivity to improve reconstruction at
the expense of latent accuracy.

Figure 3: Linear- and MLP-NODEs tend towards non-injectivity A) Effective rank
of Linear readout as a function of state dimensionality $\hat{D}$. Each point
represents one randomly instantiated model. B) Cycle-consistency $R^{2}$ for
ODIN and MLP-NODE as a function of noise corruption.

Next, we used a cycle consistency metric to show that MLP readouts also have a
tendency to become non-injective. Cycle consistency quantifies how well inputs
to a function can be recovered from the function’s outputs. We trained a
separate MLP to predict inferred latents $\hat{\mathbf{z}}$ from predicted
log-firing rates $\log\hat{\mathbf{y}}$ for 10D MLP-NODE and ODIN models shown
in Figure 2. We found that the cycle consistency of the ODIN model was
consistently higher than for MLP-NODE (Fig. 3B, Noise Level = 0). It is
possible that models may learn to compress latent activity to arbitrarily
small firing rate changes while still remaining technically injective. This
failure mode could potentially be invisible to the standard cycle-consistency.
To address this concern, we added Gaussian noise to the log-firing rates
$\log\hat{\mathbf{y}}$ and tried to recover the inferred latent activity from
these noise corrupted log-rates. Consistent with ODIN’s bias towards
injectivity, we found that ODIN’s cycle consistency was more robust to the
addition of noise than MLP-NODE (Fig. 3B, Noise Level > 0).

To demonstrate that injectivity was the critical feature that allowed ODIN to
outperform other models, we tested an alternative injective readout, an
Invertible Neural Network (INN). INN implementation differs significantly from
Flow, but they share the property of injectivity. We found that INN-NODE
qualitatively reproduced ODIN’s performance in Figure 2C (Supp. Fig. S4),
suggesting that the injectivity is the critical feature for recovering
interpretable latent activity. We describe the advantages of ODIN over INN-
NODE in the Supplementary material.

###  4.3 Recovering fixed point structure with ODIN

Figure 4:  ODIN recovers fixed point properties accurately at the correct
dimensionality. A,B) Representative latent activity and fixed-points from the
true (blue, $\circ$), ODIN (red, $\times$), and Linear-NODE (green, $+$)
systems. Each fixed point is labeled with reference to C. C) Plots of the real
vs. imaginary part of the eigenvalues of the Jacobian evaluated at each fixed
point. Unit circle in the complex plane (black curve) shows boundary between
attractive and repulsive behavior (the attractive and repulsive sides of the
boundary are indicated by inset).

A common method to examine how well dynamics models capture the underlying
dynamics from synthetic data is to compare the character and structure of the
inferred fixed points (FPs) to the FPs of the ground-truth system [15]. At a
high-level, FPs enable a concise description of the dynamics in a small region
of state-space around the FP, and can collectively provide a qualitative
picture of the overall dynamical landscape. To obtain a set of candidate FPs,
we searched the latent space for points at which the magnitude of the vector
field $\lVert\hat{f}\rVert$ is minimized (as in [1, 40]). We computed the
eigenvalues of the Jacobian of $\hat{f}$ at each FP location. The real and
imaginary components of these eigenvalues identify each FP as attractive,
repulsive, etc.

We found that 3D ODIN models and 3D Linear-NODEs were both able to recover
three fixed points that generally matched the location of the three fixed
points of the Arneodo system (Fig 4A), However, while ODIN was also able to
capture the eigenspectra of all three FPs (Fig. 4B, red $\times$), the Linear-
NODE failed to capture the rotational dynamics of the central FP (Fig 4B,
middle column, green $+$). Both models were able to approximately recover the
eigenspectra of outermost FPs of the system (Fig. 4B, left, right columns). We
found that the MLP-NODE was also able to find FPs with similar accuracy to
ODIN at 3D. These results show that the inability to model the nonlinear
embedding can lead to impoverished estimates of the underlying dynamics
$\hat{f}$.

###  4.4 Recovering simulated activation functions with ODIN

Figure 5: ODIN can recover nonlinear activation functions of neurons. A) True
encoding vectors (numbered lines over true latent activity (blue)) were
affine-transformed into a representative model’s latent space. B) Inferred
activation function for two example neurons (columns), color coded by readout
type (Linear-NODE = green, MLP-NODE = orange, ODIN = red, True = black). Plots
show the predicted firing rate vs. the activation of the selected neuron. C)
Comparison of the $R^{2}$ values of the fits across all neurons for models
with $\hat{D}=3$.

While obtaining interpretable dynamics is our primary goal, models that allow
unsupervised recovery of the embedding geometry may provide additional insight
about the computations performed by the neural system [43, 7]. For this
section, we considered a representative model from each readout class with the
correct number of latent dimensions ($D=3$). We performed an affine
transformation from the ground truth encoding vectors into the modeled latent
space and computed the projection of the modeled latent activity onto the
affine-transformed encoding vectors (Fig 5A). From this projection, we derived
an estimate of the activation function for each neuron, and compared this
estimate to the ground-truth activation function.

We found, as expected, that Linear-NODE was unable to approximate the
sigmoidal activation function of individual neurons (Fig 5B, green). On the
other hand, both ODIN and MLP-NODE were able to capture activation functions
ranging from nearly linear to step function-like in nature (Fig 5B, red,
orange). Across all simulated neurons for models with $D=3$, we found that
ODIN more accurately estimated the activation function of individual neurons
compared to both Linear- and MLP-NODEs (Fig 5C), suggesting that ODIN’s
injectivity allows more accurate estimation of nonlinear embeddings (two-sided
paired t-test, p-val for ODIN vs. Linear-, MLP-NODE < 1e-10).

###  4.5 Modeling motor cortical activity with ODIN

To validate ODIN’s ability to fit neural activity from a biological neural
circuit, we applied ODIN to the Maze dataset from the Neural Latents
Benchmark, composed of recordings from the motor and pre-motor cortices of a
monkey performing a reaching task (Fig. 6A). After performing hyperparameter
sweeps across regularization parameters and network size (Supp. Table 2), we
trained a set of ODIN and Linear-NODE models to reconstruct the neural
activity with a range of state dimensionalities $\hat{D}$. We visualized the
top 3 PCs of the condition-averaged latent trajectories and predicted single-
neuron firing rates for example models from each readout type. We found no
visually obvious differences in the inferred latent trajectories (Fig. 6B),
but when we computed condition-averaged peri-stimulus time histograms (PSTHs)
of single neuron firing rates, we found that ODIN typically produced firing
rate estimates that more closely resembled the empirical PSTHs than those from
the Linear-NODE (Fig. 6C).

Without access to a ground truth dynamics $f$ and embedding $g$ that generated
these biological data, the dimensionality required to reconstruct the neural
activity was our primary measure of interpretability. We computed co-bps –a
measure of reconstruction performance on held-out neurons– for each model and
found that 10D ODIN models substantially outperformed Linear-NODE models, even
when the Linear-NODE had more than twice as many dimensions (10D ODIN: 0.333,
vs 25D Linear: 0.287). This suggests that ODIN’s injective non-linear readout
is effective at reducing the state dimensionality required to capture the data
relative to a simple linear readout.

We also compared ODIN to alternative models including AutoLFADS, GPFA, and
MLP-NODE [27] at the same state dimensionalities. Trained AutoLFADS and GPFA
models had lower co-bps at all tested state dimensionalities. In particular,
co-bps was substantially higher for 10D ODIN compared to the 10D AutoLFADS or
GPFA models (0.333 vs. 0.237, 0.204, respectively). As expected, MLP-NODE (not
shown) performed similarly to ODIN; however, without a known state
dimensionality, the MLP readout may incentivize the MLP-NODE to invent latent
activity that is not reflected in the dataset. Of note, increasing AutoLFADS
to a very high state dimensionality ($\hat{D}=100)$ allowed it to outperform
ODIN in co-bps. However, as we have shown in Figures 2 and 3, improved
reconstruction performance often comes at the expense of accuracy in latent
recovery. Together, these results suggest that ODIN is effective at reducing
the state dimensionality needed for good neural reconstruction, which may
provide more interpretable latent representations than alternative models.

Figure 6:  ODIN can reconstruct cortical activity with low-dimensional
dynamics A) Top: Schematic of task [38] Bottom: example hand trajectories and
condition-averaged firing rates aligned to move onset. B) Example condition-
averaged latent activity from ODIN and Linear-NODE models applied to neural
activity recorded during the Maze task. C) Example single-neuron peri-stimulus
time histograms for ODIN and Linear-NODE models across conditions. D) Effects
of latent state dimensionality $\hat{D}$ on reconstruction (top, co-bps) and
decoding (bottom, Vel $R^{2}$) performance. Plot shows mean (point) and
standard deviation (shading) of 5 randomly initialized ODIN and Linear-NODE
models at each $\hat{D}$. GPFA and AutoLFADS were a single run, or the best
performing model from an adaptive hyperparameter search, respectively.
Horizontal lines represent peak performance by AutoLFADS with $\hat{D}=100$.

##  5 Discussion

Dynamics models have had great success in reproducing neural activity patterns
and relating brain activity to behavior [44, 27, 45]. However, it has been
difficult to use these models to investigate neural computation directly. If
neural population models could be trusted to find interpretable
representations of latent dynamics, then recent techniques that can uncover
computation in artificial networks could help to explain computations in the
brain [1, 40, 46]. In this work, we created a new model called ODIN that can
overcome major barriers to learning interpretable latent dynamical systems. By
combining Neural ODE generators and approximately injective nonlinear
readouts, ODIN offers significant advantages over the current state-of-the-
art, including lower latent dimensionality, simpler latent activity that is
robust to the choice of latent dimensionality, and the ability to model
arbitrary nonlinear activation functions.

Circuits in the brain are densely interconnected, and so a primary limitation
of this work is that ODIN is not yet able to account for inputs to the system
that may be coming from areas that are not directly modeled. Thus ODIN is
currently only able to model the dynamics of a given population of neurons as
an autonomous system. Inferring inputs is difficult due to ambiguity in the
role and timecourse of inputs compared to internal dynamics for driving the
state of the system. While some RNN-based models have methods for input
inference [44], more work is needed to develop solutions for NODE-based
models. Injective readouts are an important step towards addressing the
fundamental difficulties of input inference, as models without injective
readouts can be incentivized to imagine latent features that are actually the
result of inputs.

Interpretable dynamics derived from neural population recordings could answer
critical scientific questions about the brain and help improve brain-machine
interface technology. A potential negative consequence is that human neural
interfaces combined with an understanding of neural computation might make it
possible and profitable to develop strategies that are effective at
influencing behavior. Future researchers should focus on applications of this
research that are scientific and medical rather than commercial or political.

##  6 Acknowledgements

The authors would like to acknowledge Timothy D. Kim and Carlos Brody for
helpful discussions that further developed the ideas in this manuscript.

This work was supported by NSF NCS 1835364, NIH-NINDS/OD DP2NS127291, NIH
BRAIN/NIDA RF1 DA055667, and the Alfred P. Sloan Foundation (CP), NIH
BRAIN/NINDS F32 RFA-MH-23-110 (CV), the Simons Foundation as part of the
Simons-Emory International Consortium on Motor Control (CP, CV), and NSF
Graduate Research Fellowship DGE-2039655 (ARS).

## References

  * Sussillo and Barak [2013] David Sussillo and Omri Barak.  Opening the black box: low-dimensional dynamics in high-dimensional recurrent neural networks.  _Neural Computation_ , 25(3):626–649, March 2013\.  ISSN 1530-888X.  doi: 10.1162/NECO_a_00409.
  * Mante et al. [2013] Valerio Mante, David Sussillo, Krishna Shenoy, and William Newsome.  Context-dependent computation by recurrent dynamics in prefrontal cortex.  _Nature_ , 503:78–84, November 2013.  doi: 10.1038/nature12742.
  * Remington et al. [2018] Evan D. Remington, Devika Narain, Eghbal A. Hosseini, and Mehrdad Jazayeri.  Flexible Sensorimotor Computations through Rapid Reconfiguration of Cortical Dynamics.  _Neuron_ , 98(5):1005–1019.e5, June 2018.  ISSN 1097-4199.  doi: 10.1016/j.neuron.2018.05.020.
  * Maheswaranathan et al. [2019] Niru Maheswaranathan, Alex Williams, Matthew Golub, Surya Ganguli, and David Sussillo.  Reverse engineering recurrent networks for sentiment classification reveals line attractor dynamics.  In _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/d921c3c762b1522c475ac8fc0811bb0f-Abstract.html.
  * Vyas et al. [2020] Saurabh Vyas, Matthew D. Golub, David Sussillo, and Krishna V. Shenoy.  Computation Through Neural Population Dynamics.  _Annual Review of Neuroscience_ , 43(1):249–275, 2020.  doi: 10.1146/annurev-neuro-092619-094115.  URL https://doi.org/10.1146/annurev-neuro-092619-094115.  _eprint: https://doi.org/10.1146/annurev-neuro-092619-094115.
  * Shenoy et al. [2013] Krishna V. Shenoy, Maneesh Sahani, and Mark M. Churchland.  Cortical control of arm movements: a dynamical systems perspective.  _Annual Review of Neuroscience_ , 36:337–359, July 2013\.  ISSN 1545-4126.  doi: 10.1146/annurev-neuro-062111-150509.
  * Jazayeri and Ostojic [2021] Mehrdad Jazayeri and Srdjan Ostojic.  Interpreting neural computations by examining intrinsic and embedding dimensionality of neural activity.  Technical Report arXiv:2107.04084, arXiv, August 2021.  URL http://arxiv.org/abs/2107.04084.  arXiv:2107.04084 [q-bio] type: article.
  * Stevenson and Kording [2011] Ian H. Stevenson and Konrad P. Kording.  How advances in neural recording affect data analysis.  _Nature Neuroscience_ , 14(2):139–142, February 2011.  ISSN 1546-1726.  doi: 10.1038/nn.2731.
  * Steinmetz et al. [2021] Nicholas A Steinmetz, Cagatay Aydin, Anna Lebedeva, Michael Okun, Marius Pachitariu, Marius Bauza, Maxime Beau, Jai Bhagat, Claudia Böhm, Martijn Broux, Susu Chen, Jennifer Colonell, Richard J Gardner, Bill Karsh, Fabian Kloosterman, Dimitar Kostadinov, Carolina Mora-Lopez, John O’Callaghan, Junchol Park, Jan Putzeys, Britton Sauerbrei, Rik J J van Daal, Abraham Z Vollan, Shiwei Wang, Marleen Welkenhuysen, Zhiwen Ye, Joshua T Dudman, Barundeb Dutta, Adam W Hantman, Kenneth D Harris, Albert K Lee, Edvard I Moser, John O’Keefe, Alfonso Renart, Karel Svoboda, Michael Häusser, Sebastian Haesler, Matteo Carandini, and Timothy D Harris.  Neuropixels 2.0: A miniaturized high-density probe for stable, long-term brain recordings.  _Science_ , 372(6539), April 2021.
  * Demas et al. [2021] Jeffrey Demas, Jason Manley, Frank Tejera, Kevin Barber, Hyewon Kim, Francisca Martínez Traub, Brandon Chen, and Alipasha Vaziri.  High-speed, cortex-wide volumetric recording of neuroactivity at cellular resolution using light beads microscopy.  _Nature Methods_ , 18(9):1103–1111, September 2021.  ISSN 1548-7105.  doi: 10.1038/s41592-021-01239-8.  URL https://www.nature.com/articles/s41592-021-01239-8.  Number: 9 Publisher: Nature Publishing Group.
  * Gao and Ganguli [2015] Peiran Gao and Surya Ganguli.  On simplicity and complexity in the brave new world of large-scale neuroscience.  _Current Opinion in Neurobiology_ , 32:148–155, June 2015\.  ISSN 0959-4388.  doi: 10.1016/J.CONB.2015.04.003.  URL https://www.sciencedirect.com/science/article/pii/S0959438815000768.  Publisher: Elsevier Current Trends.
  * Duncker and Sahani [2021] Lea Duncker and Maneesh Sahani.  Dynamics on the manifold: Identifying computational dynamical activity from neural population recordings.  _Current Opinion in Neurobiology_ , 70:163–170, October 2021.  ISSN 0959-4388.  doi: 10.1016/j.conb.2021.10.014.  URL https://www.sciencedirect.com/science/article/pii/S0959438821001264.
  * Sussillo et al. [2016] David Sussillo, Rafal Jozefowicz, L. F. Abbott, and Chethan Pandarinath.  LFADS - Latent Factor Analysis via Dynamical Systems.  Technical Report arXiv:1608.06315, arXiv, August 2016.  URL http://arxiv.org/abs/1608.06315.  arXiv:1608.06315 [cs, q-bio, stat] type: article.
  * Schimel et al. [2021] Marine Schimel, Ta-Chu Kao, Kristopher T. Jensen, and Guillaume Hennequin.  iLQR-VAE : control-based learning of input-driven dynamics with applications to neural data.  Technical report, bioRxiv, October 2021.  URL https://www.biorxiv.org/content/10.1101/2021.10.07.463540v1.  Section: New Results Type: article.
  * Sedler et al. [2023] Andrew R. Sedler, Christopher Versteeg, and Chethan Pandarinath.  Expressive architectures enhance interpretability of dynamics-based neural population models, February 2023.  URL http://arxiv.org/abs/2212.03771.  arXiv:2212.03771 [cs, q-bio].
  * Gao et al. [2016] Yuanjun Gao, Evan Archer, Liam Paninski, and John P. Cunningham.  Linear dynamical neural population models through nonlinear embeddings.  Technical Report arXiv:1605.08454, arXiv, October 2016.  URL http://arxiv.org/abs/1605.08454.  arXiv:1605.08454 [q-bio, stat] type: article.
  * Wu et al. [2017] Anqi Wu, Nicholas A. Roy, Stephen Keeley, and Jonathan W Pillow.  Gaussian process based nonlinear latent structure discovery in multivariate spike train data.  In _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://papers.nips.cc/paper_files/paper/2017/hash/b3b4d2dbedc99fe843fd3dedb02f086f-Abstract.html.
  * Zhao and Park [2020] Yuan Zhao and Il Memming Park.  Variational Online Learning of Neural Dynamics.  _Frontiers in Computational Neuroscience_ , 14, 2020.  ISSN 1662-5188.  URL https://www.frontiersin.org/article/10.3389/fncom.2020.00071.
  * Chen et al. [2019] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud.  Neural Ordinary Differential Equations.  Technical Report arXiv:1806.07366, arXiv, December 2019.  URL http://arxiv.org/abs/1806.07366.  arXiv:1806.07366 [cs, stat] type: article.
  * Dinh et al. [2014] Laurent Dinh, David Krueger, and Yoshua Bengio.  Nice: Non-linear independent components estimation.  _arXiv preprint arXiv:1410.8516_ , 2014.
  * Kingma and Dhariwal [2018] Durk P Kingma and Prafulla Dhariwal.  Glow: Generative flow with invertible 1x1 convolutions.  _Advances in neural information processing systems_ , 31, 2018.
  * Ardizzone et al. [2019] Lynton Ardizzone, Jakob Kruse, Sebastian Wirkert, Daniel Rahner, Eric W. Pellegrini, Ralf S. Klessen, Lena Maier-Hein, Carsten Rother, and Ullrich Köthe.  Analyzing Inverse Problems with Invertible Neural Networks.  Technical Report arXiv:1808.04730, arXiv, February 2019.  URL http://arxiv.org/abs/1808.04730.  arXiv:1808.04730 [cs, stat] type: article.
  * Behrmann et al. [2019] Jens Behrmann, Will Grathwohl, Ricky T. Q. Chen, David Duvenaud, and Joern-Henrik Jacobsen.  Invertible Residual Networks.  In _Proceedings of the 36th International Conference on Machine Learning_ , pages 573–582. PMLR, May 2019.  URL https://proceedings.mlr.press/v97/behrmann19a.html.  ISSN: 2640-3498.
  * Macke et al. [2011] Jakob H Macke, Lars Buesing, John P Cunningham, Byron M Yu, Krishna V Shenoy, and Maneesh Sahani.  Empirical models of spiking in neural populations.  In _Advances in Neural Information Processing Systems_ , volume 24. Curran Associates, Inc., 2011.  URL https://papers.nips.cc/paper/2011/hash/7143d7fbadfa4693b9eec507d9d37443-Abstract.html.
  * Archer et al. [2015] Evan Archer, Il Memming Park, Lars Buesing, John Cunningham, and Liam Paninski.  Black box variational inference for state space models, November 2015\.  URL http://arxiv.org/abs/1511.07367.  arXiv:1511.07367 [stat].
  * Pfau et al. [2013] David Pfau, Eftychios A Pnevmatikakis, and Liam Paninski.  Robust learning of low-dimensional dynamics from large neural ensembles.  In _Advances in Neural Information Processing Systems_ , volume 26. Curran Associates, Inc., 2013.  URL https://papers.nips.cc/paper_files/paper/2013/hash/47a658229eb2368a99f1d032c8848542-Abstract.html.
  * Pei et al. [2022] Felix Pei, Joel Ye, David Zoltowski, Anqi Wu, Raeed H. Chowdhury, Hansem Sohn, Joseph E. O’Doherty, Krishna V. Shenoy, Matthew T. Kaufman, Mark Churchland, Mehrdad Jazayeri, Lee E. Miller, Jonathan Pillow, Il Memming Park, Eva L. Dyer, and Chethan Pandarinath.  Neural Latents Benchmark ’21: Evaluating latent variable models of neural population activity.  Technical Report arXiv:2109.04463, arXiv, January 2022.  URL http://arxiv.org/abs/2109.04463.  arXiv:2109.04463 [cs, q-bio] type: article.
  * Kim et al. [2021] Timothy D Kim, Thomas Z Luo, Jonathan W Pillow, and Carlos Brody.  Inferring latent dynamics underlying neural population activity via neural differential equations.  In _International Conference on Machine Learning_ , pages 5551–5561. PMLR, 2021.
  * Johnson et al. [2017] Matthew J. Johnson, David Duvenaud, Alexander B. Wiltschko, Sandeep R. Datta, and Ryan P. Adams.  Composing graphical models with neural networks for structured representations and fast inference, July 2017.  URL http://arxiv.org/abs/1603.06277.  arXiv:1603.06277 [stat].
  * Stevenson [2016] Ian H. Stevenson.  Flexible models for spike count data with both over- and under- dispersion.  _Journal of Computational Neuroscience_ , 41(1):29–43, August 2016.  ISSN 1573-6873.  doi: 10.1007/s10827-016-0603-y.  URL https://doi.org/10.1007/s10827-016-0603-y.
  * Brunton et al. [2016] Steven L Brunton, Joshua L Proctor, and J Nathan Kutz.  Discovering governing equations from data by sparse identification of nonlinear dynamical systems.  _Proceedings of the national academy of sciences_ , 113(15):3932–3937, 2016.
  * Schneider et al. [2023] Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis.  Learnable latent embeddings for joint behavioural and neural analysis.  _Nature_ , 617(7960):360–368, May 2023.  ISSN 1476-4687.  doi: 10.1038/s41586-023-06031-6.  URL https://www.nature.com/articles/s41586-023-06031-6.  Number: 7960 Publisher: Nature Publishing Group.
  * Zhou and Wei [2020] Ding Zhou and Xue-Xin Wei.  Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-VAE, November 2020.  URL http://arxiv.org/abs/2011.04798.  arXiv:2011.04798 [cs, q-bio, stat].
  * Bashiri et al. [2021] Mohammad Bashiri, Edgar Walker, Konstantin-Klemens Lurz, Akshay Jagadish, Taliah Muhammad, Zhiwei Ding, Zhuokun Ding, Andreas Tolias, and Fabian Sinz.  A flow-based latent state generative model of neural population responses to natural images.  In _Advances in Neural Information Processing Systems_ , volume 34, pages 15801–15815. Curran Associates, Inc., 2021.  URL https://proceedings.neurips.cc/paper/2021/hash/84a529a92de322be42dd3365afd54f91-Abstract.html.
  * Smith et al. [2021] Jimmy T. H. Smith, Scott W. Linderman, and David Sussillo.  Reverse engineering recurrent neural networks with Jacobian switching linear dynamical systems.  Technical Report arXiv:2111.01256, arXiv, November 2021.  URL http://arxiv.org/abs/2111.01256.  arXiv:2111.01256 [cs] type: article.
  * Hurwitz et al. [2021] Cole Hurwitz, Akash Srivastava, Kai Xu, Justin Jude, Matthew Perich, Lee Miller, and Matthias Hennig.  Targeted Neural Dynamical Modeling.  In _Advances in Neural Information Processing Systems_ , volume 34, pages 29379–29392. Curran Associates, Inc., 2021.  URL https://papers.nips.cc/paper_files/paper/2021/hash/f5cfbc876972bd0d031c8abc37344c28-Abstract.html.
  * Jensen et al. [2021] Kristopher Jensen, Ta-Chu Kao, Jasmine Stone, and Guillaume Hennequin.  Scalable Bayesian GPFA with automatic relevance determination and discrete noise models.  In _Advances in Neural Information Processing Systems_ , volume 34, pages 10613–10626. Curran Associates, Inc., 2021.  URL https://proceedings.neurips.cc/paper/2021/hash/58238e9ae2dd305d79c2ebc8c1883422-Abstract.html.
  * Churchland et al. [2010] Mark M. Churchland, John P. Cunningham, Matthew T. Kaufman, Stephen I. Ryu, and Krishna V. Shenoy.  Cortical preparatory activity: representation of movement or first cog in a dynamical machine?  _Neuron_ , 68(3):387–400, November 2010.  ISSN 1097-4199.  doi: 10.1016/j.neuron.2010.09.015.
  * Miyato et al. [2018] Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida.  Spectral Normalization for Generative Adversarial Networks, February 2018.  URL http://arxiv.org/abs/1802.05957.  arXiv:1802.05957 [cs, stat].
  * Golub and Sussillo [2018] Matthew D. Golub and David Sussillo.  Fixedpointfinder: A tensorflow toolbox for identifying and characterizing fixed points in recurrent neural networks.  _Journal of Open Source Software_ , 3(31):1003, 2018.  doi: 10.21105/joss.01003.  URL https://doi.org/10.21105/joss.01003.
  * Roy and Vetterli [2007] Olivier Roy and Martin Vetterli.  The Effective Rank: a Measure of Effective Dimensionality.  _European Association for Signal Processing_ , 2007.
  * Arneodo et al. [1980] A Arneodo, P Coullet, and C Tresser.  Occurence of strange attractors in three-dimensional Volterra equations.  _Physics Letters A_ , 79(4):259–263, October 1980\.  ISSN 0375-9601.  doi: 10.1016/0375-9601(80)90342-4.  URL https://www.sciencedirect.com/science/article/pii/0375960180903424.
  * Gardner et al. [2021] Richard J. Gardner, Erik Hermansen, Marius Pachitariu, Yoram Burak, Nils A. Baas, Benjamin A. Dunn, May-Britt Moser, and Edvard I. Moser.  Toroidal topology of population activity in grid cells.  Technical report, bioRxiv, February 2021.  URL https://www.biorxiv.org/content/10.1101/2021.02.25.432776v1.  Section: New Results Type: article.
  * Pandarinath et al. [2018] Chethan Pandarinath, Daniel J. O’Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D. Stavisky, Jonathan C. Kao, Eric M. Trautmann, Matthew T. Kaufman, Stephen I. Ryu, Leigh R. Hochberg, Jaimie M. Henderson, Krishna V. Shenoy, L. F. Abbott, and David Sussillo.  Inferring single-trial neural population dynamics using sequential auto-encoders.  _Nature Methods_ , 15(10):805–815, October 2018\.  ISSN 1548-7105.  doi: 10.1038/s41592-018-0109-9.  URL https://www.nature.com/articles/s41592-018-0109-9.  Number: 10 Publisher: Nature Publishing Group.
  * Smith et al. [2023] Jimmy T. H. Smith, Andrew Warrington, and Scott W. Linderman.  Simplified State Space Layers for Sequence Modeling, March 2023\.  URL http://arxiv.org/abs/2208.04933.  arXiv:2208.04933 [cs].
  * Driscoll et al. [2022] Laura Driscoll, Krishna Shenoy, and David Sussillo.  Flexible multitask computation in recurrent networks utilizes shared dynamical motifs, August 2022.  URL https://www.biorxiv.org/content/10.1101/2022.08.15.503870v1.  Pages: 2022.08.15.503870 Section: New Results.
  * Gilpin [2021] William Gilpin.  Chaos as an interpretable benchmark for forecasting and data-driven modelling.  _Advances in Neural Information Processing Systems_ , 2021.  URL http://arxiv.org/abs/2110.05266.
  * Maynard et al. [1997] Edwin M. Maynard, Craig T. Nordhausen, and Richard A. Normann.  The Utah Intracortical Electrode Array: A recording structure for potential brain-computer interfaces.  _Electroencephalography and Clinical Neurophysiology_ , 102(3):228–239, March 1997.  ISSN 0013-4694.  doi: 10.1016/S0013-4694(96)95176-0.  URL https://www.sciencedirect.com/science/article/pii/S0013469496951760.
  * Rübel et al. [2022] Oliver Rübel, Andrew Tritt, Ryan Ly, Benjamin K Dichter, Satrajit Ghosh, Lawrence Niu, Pamela Baker, Ivan Soltesz, Lydia Ng, Karel Svoboda, Loren Frank, and Kristofer E Bouchard.  The Neurodata Without Borders ecosystem for neurophysiological data science.  _eLife_ , 11:e78362, October 2022.  ISSN 2050-084X.  doi: 10.7554/eLife.78362.  URL https://doi.org/10.7554/eLife.78362.  Publisher: eLife Sciences Publications, Ltd.
  * Keshtkaran et al. [2022] Mohammad Reza Keshtkaran, Andrew R. Sedler, Raeed H. Chowdhury, Raghav Tandon, Diya Basrai, Sarah L. Nguyen, Hansem Sohn, Mehrdad Jazayeri, Lee E. Miller, and Chethan Pandarinath.  A large-scale neural network training framework for generalized estimation of single-trial population dynamics.  _Nature Methods_ , 19(12):1572–1577, December 2022.  ISSN 1548-7105.  doi: 10.1038/s41592-022-01675-0.  URL https://www.nature.com/articles/s41592-022-01675-0.  Number: 12 Publisher: Nature Publishing Group.
  * Willett et al. [2021] Francis R. Willett, Donald T. Avansino, Leigh R. Hochberg, Jaimie M. Henderson, and Krishna V. Shenoy.  High-performance brain-to-text communication via handwriting.  _Nature_ , 593(7858):249–254, May 2021.  ISSN 0028-0836.  doi: 10.1038/s41586-021-03506-2.  URL https://www.ncbi.nlm.nih.gov/pmc/articles/PMC8163299/.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  PyTorch: An Imperative Style, High-Performance Deep Learning Library.  Technical Report arXiv:1912.01703, arXiv, December 2019.  URL http://arxiv.org/abs/1912.01703.  arXiv:1912.01703 [cs, stat] type: article.
  * Liaw et al. [2018] Richard Liaw, Eric Liang, Robert Nishihara, Philipp Moritz, Joseph E Gonzalez, and Ion Stoica.  Tune: A research platform for distributed model selection and training.  _arXiv preprint arXiv:1807.05118_ , 2018.
  * Pedregosa et al. [2011] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay.  Scikit-learn: Machine learning in Python.  _Journal of Machine Learning Research_ , 12:2825–2830, 2011\.

Expressive dynamics models with nonlinear injective readouts enable reliable
recovery of latent features from neural activity
Supplementary Material

##  G Datasets

###  G.1 Simulated neural data

####  G.1.1 Latent trajectories

We used the Arneodo system [42] to generate synthetic data because it exhibits
mildly chaotic behavior (Lyapunov exponent equal to 0.243), it has a low-
dimensional state space, and the regions around its fixed points are well-
sampled by trajectories of the system. As demonstrated by [15], these
properties allow recovery of latent dynamics in the absence of a nonlinear
embedding. The Arneodo system is described by the following system of
equations

$\displaystyle\dot{x}$ $\displaystyle=y$ (13) $\displaystyle\dot{y}$
$\displaystyle=z$ (14) $\displaystyle\dot{z}$ $\displaystyle=-ax-by-cz+dx^{3}$
(15)

where $a=-5.5$, $b=4.5$, $c=1.0$, and $d=-1.0$ [42].

The system was simulated using the dysts Python package, which offered well-
reasoned standards for initial conditions, integration steps, and resampling
frequency [47]. Initial conditions had been determined by running the model
until the moments of the autocorrelation function were stationary. Integration
steps had been chosen based on the highest significant frequency observed in
the power spectrum. After integration, trajectories were resampled to contain
35 points per period, where period was based on the dominant frequency in the
power spectrum.

####  G.1.2 Embedding low-dimensional trajectories on a nonlinear manifold

We simulated neural activity by nonlinearly embedding the Arneodo trajectories
as firing rates in the neural space. First, the trajectories were linearly
projected into the neural space via a set of encoding vectors
$\mathbf{\gamma}_{i}$ and standardized for each neuron (see Methods). These
activations $\mathbf{a}_{i}$ were passed through a sigmoid with input scaling
$\mathbf{\eta}_{i}$ and output scaling $b=2$ to produce reasonable firing
rates as follows:

$\displaystyle\mathbf{\eta_{i}}=10^{0.8\cdot\frac{i-1}{N-1}+0.2},$ (16)
$\displaystyle\mathbf{y}_{i}=\psi_{i}(\mathbf{a}_{i})=b\times\sigma(\mathbf{\eta}_{i}\times\mathbf{a}_{i}),\quad
i=1,2,\dots,N.$ (17)

where $\sigma(\cdot)$ denotes the sigmoid function. This resulted in a set of
activation functions $\psi_{i}(\cdot)$ ranging from quasi-linear to step-like.
The resulting rates $\mathbf{y}_{i}$ were used to parameterize a Poisson
process, which was sampled to obtain spiking data for $N$ neurons ($N=12$).

####  G.1.3 Embedding low-dimensional trajectories onto linear manifold

For Figure 2A, we tested whether Linear-NODEs fit to linearly-embedded data
would find non-injective readouts when $\hat{D}>D$. We simulated an
alternative dataset with the same procedure as above, except instead of
passing the activations $\mathbf{a}_{i}$ through the sigmoidal non-linearity,
we exponentiated them to find the rate parameter $\mathbf{y}_{i}$ of a Poisson
process, which was sampled to obtain spiking data for $N$ neurons ($N=12$).
These data were used only in Figure 2A.

###  G.2 Real neural data

The maze dataset was previously collected from the motor cortex of a monkey
performing a reaching task [38]. This dataset has been widely used to
characterize the dynamics of motor cortical activity [38, 27, 44]. In
particular, these data are well-modeled by autonomous dynamics [44].

The monkey was trained to perform a delayed reaching task in which it had to
maintain its hand at the center of a 2D maze displayed on a screen while a
target was shown somewhere within the maze. After a randomly-timed delay, a
go-cue was issued which prompted the monkey to move its hand from the center
of the screen to the indicated target. Each trial also had a set of obstacles
(i.e., the walls of a maze) with various configurations that required the
monkey to produce reaches with varied trajectories, even when they were
directed towards the same target. A total of 108 of these maze configurations
(i.e., target and obstacle combinations) are included in this dataset.

Neural activity was recorded using two Utah arrays [48], one in the dorsal
premotor (PMd) cortex and one in the primary motor cortex (M1) [38]. Threshold
crossings were sorted offline. The dataset contained 182 neurons in total, of
which 137 were included in the held-in set and the remaining 45 were part of
the held-out set. The held-out neurons were used to calculate the co-smoothing
bits-per-spike metric (K.2.1). The monkey’s hand and cursor positions were
recorded during the experiment (K.2.2).

These data were downloaded from the Distributed Archives for Neurophysiology
Data Integration (DANDI, [49]). We binned spike counts at 20 ms and trialized
and aligned the data to 250 ms before and 450 ms after movement onset. Further
details can be found in [38, 27].

##  H Model training

###  H.1 Simulated neural data

All weights were initialized from $\mathcal{U}(-\sqrt{k},\sqrt{k})$, where
$k=1/\mathrm{in\\_features}$ for linear layers and
$k=1/\mathrm{hidden\\_size}$ for the GRU encoder weights. Dropout layers
($p=0.05$) were inserted before and after the initial condition linear
projection during training. We used the average Poisson negative log-
likelihood (NLL) across neurons and time points as our training objective.
Models were trained incrementally to improve the stability of training: rather
than compute loss on the whole trajectory, we added groups of 5 new time steps
every 75 epochs, up to the max of 70 steps. Models were trained by stochastic
gradient descent using Adam for 3000 epochs. A single learning rate was shared
for the optimizer of the encoder, generator, and readout weights for each
model. Each generator was a NODE that contained an MLP with six hidden layers,
each with 128 ReLU units.

Figure S1: Example hyperparameter sweeps for ODIN and MLP-NODE

We performed initial hyperparameter sweeps to determine ranges that resulted
in good reconstruction performance as measured by Spike NLL (see Methods), and
used the same hyperparameter setting for models across state dimensionalities.
Two example hyperparameter sweeps testing the effect of readout capacity (100
model initializations with readout hidden sizes in [60,200] and number of
hidden layers in [1,3]) and weight decay (100 model initializations with
weight decay drawn log-uniformly from [1e-8, 1e-4], Supp. Fig. S1). We found
that across readout capacities, good reconstruction performance implied good
latent recovery for ODIN but not MLP-NODE. Additionally, we found that
increasing weight decay on MLP-NODE tended to degrade rather than improve
latent recovery. Across all HPs tested, we found no hyperparameter settings
for which ODIN had good reconstruction performance but poor latent recovery.

HPs for models trained on the Arneodo system are given in Table S1.

Table S1: Training hyperparameters (Synthetic Data) |  | Arneodo |
---|---|---|---
| Linear | MLP | ODIN
Batch Size | 650 | 650 | 650
Learning Rate | 2e-3 | 1.88e-4 | 1.88e-4
Encoder Hidden Size | 100 | 100 | 100
Dropout | 0.05 | 0.05 | 0.05
NODE Hidden Layers | 6 | 6 | 6
NODE Hidden Size | 128 | 128 | 128
Readout Hidden Layers | 0 | 2 | 2
Readout Hidden Size | - | 150 | 150

###  H.2 Real neural data

The weight initialization procedure and dropout settings were the same as for
the models trained on Arneodo. In addition to Poisson NLL, we also added
regularization terms ($L_{2}$ norm on weights) and used different learning
rates for the encoder, generator, and readout modules. We trained these models
using Adam for 1500 epochs with the loss function given by Equation 18:

$\displaystyle
L(\mathbf{x},\mathbf{\hat{y}},\mathbf{\theta}_{E},\mathbf{\theta}_{G},\mathbf{\theta}_{R})=\text{PoissonNLL}(\mathbf{x}|\mathbf{\hat{y}})+\lambda_{E}||\mathbf{\theta}_{E}||_{2}^{2}+\lambda_{G}||\mathbf{\theta}_{G}||_{2}^{2}+\lambda_{R}||\mathbf{\theta}_{R}||_{2}^{2}$
(18)

where $\mathbf{x}$ and $\mathbf{\hat{y}}$ represent the observed spiking
activity and the predicted firing rates, respectively, and
$\lambda_{E},\lambda_{G},\lambda_{R}$ represent the regularization
coefficients for the $L_{2}$ regularization penalty applied to the model
weights $\mathbf{\theta}_{E},\mathbf{\theta}_{G},\mathbf{\theta}_{R}$ of the
encoder, generator, and readout, respectively. To improve training stability,
we also used different learning rates for each component of the model
($\alpha_{E},\alpha_{G},\alpha_{R}$). Specific parameters for models trained
on the Maze dataset are given in Table S2.

####  H.2.1 AutoLFADS

We trained AutoLFADS models of varying latent dimensionalities as a point of
reference for ODIN’s performance [50]. Notably, we used the autonomous version
of LFADS and fixed the initial condition, generator, and factors
dimensionality to $\hat{D}$ for these experiments. The batch size was 512 and
the encoder hidden size was 100. Population-Based Training was used with a
population of 20 workers to search initial learning rate (init: 1e-2, range:
loguniform; 1e-5, 5e-2), dropout rate (init: 5e-2, range: uniform; 0.0, 0.6),
coordinated dropout rate (init: 0.3, range: uniform; 0.01, 0.7), initial
condition KL (range: loguniform; 1e-10, 1e-3), generator L2 scale (range:
loguniform; 1e-10, 1e0), and encoder L2 scale (range: loguniform; 1e-10, 1e0).
Linear ramp-up of KL and L2 penalties occurred over the first 80 epochs and
the population was subjected to binary tournament and perturbation every 25
epochs for a total of 1000 training epochs.

Table S2: Training hyperparameters (Maze Data) |  Maze |
---|---|---
| Linear | ODIN
Batch Size | 64 | 64
$\lambda_{E}$ | 1.6e-5 | 2.2e-6
$\lambda_{G}$ | 1.6e-5 | 1.35e-9
$\lambda_{R}$ | 1.6e-5 | 4.2e-6
$\alpha_{E}$ | 5e-3 | 4e-4
$\alpha_{G}$ | 5e-3 | 7e-4
$\alpha_{R}$ | 5e-3 | 1.4e-4
Encoder Hidden Size | 100 | 100
Dropout | 0.05 | 0.05
NODE Hidden Layers | 6 | 6
NODE Hidden Size | 128 | 128
Readout Hidden Layers | 0 | 3
Readout Hidden Size | - | 128
Number of Flow Steps | - | 25

##  I Injectivity estimation

Figure S2: Injectivity of the Flow readout across state dimensionalities. Each
bar indicates the mean value of 5 randomly initialized ODIN models for each
state dimensionality. Results from individual models are plotted as points.

To demonstrate the approximate injectivity of the Flow readout, we tested
whether the readout could be inverted to recover the inferred latent activity.
The readout mapping $\hat{g}$ should satisfy the following equations

$\displaystyle\tilde{\mathbf{z}}_{t}$
$\displaystyle=\hat{g}^{-1}(\hat{g}(\hat{\mathbf{z}}_{t}))$ (19)
$\displaystyle\tilde{\mathbf{z}}_{t}$
$\displaystyle\approx\hat{\mathbf{z}}_{t}$ (20)

where $\hat{\mathbf{z}}_{t}$ is the inferred latent activity and
$\tilde{\mathbf{z}}_{t}$ is the latent activity recovered by the reverse pass
of the Flow.

We computed the $R^{2}$ between the inferred and recovered $\mathbf{z}_{t}$
for these models and found that our mappings were able to recover the inferred
$\hat{\mathbf{z}}_{t}$ with average $R^{2}$ values across randomly initialized
models of 0.997, 0.996, 0.990, and 0.988 at $\hat{D}=3,5,8,10$, respectively
(Supplementary Figure S2).

###  I.1 Effective Rank

To assess the injectivity of the Linear readout, we used a previously
published method that determines the approximate number of significant
singular values of a given matrix $A$ [41]. Let $A$ be a complex-valued, non-
all-zero matrix of size $N\times\hat{D}$, where $N>\hat{D}$ that acts as the
weight matrix of a readout from inferred latents $\hat{\mathbf{z}}$ to
predicted log-rates $\log\hat{\mathbf{y}}$ in the equation
$\log\hat{\mathbf{y}}=A\hat{\mathbf{z}}+\mathbf{b}$. We perform a singular
value decomposition (SVD) on $A$, such that $A=U\Delta V$, where $U$ and $V$
are unitary matrices of size $N\times N$ and $\hat{D}\times\hat{D}$,
respectively, and $\Delta$ is an $N\times\hat{D}$ rectangular diagonal matrix
containing the real non-negative singular values
$\sigma_{1}\geq\sigma_{2}\geq\ldots\geq\sigma_{\hat{D}}\geq 0$.

For simplicity, let us define
$\sigma=(\sigma_{1},\sigma_{2},\ldots,\sigma_{\hat{D}})^{T}$. We then compute
the singular value distribution $p_{k}$, for $k=1,2,\ldots,\hat{D}$, as

$p_{k}=\frac{\sigma_{k}}{\|\sigma\|_{1}},$ (21)

where $\|\sigma\|_{1}$ is the $L_{1}$-norm. Using this singular value
distribution, we can calculate the Shannon entropy $H$ as

$H(p_{1},p_{2},\ldots,p_{\hat{D}})=-\sum_{k=1}^{\hat{D}}p_{k}\log(p_{k}).$
(22)

The authors in [41] define the effective rank of the matrix $A$, denoted as
$\text{{erank}}(A)$, using the Shannon entropy $H$ as follows:

$\text{{erank}}(A)=\exp\left(H(p_{1},p_{2},\ldots,p_{\hat{D}})\right).$ (23)

The effective rank gives us a measure of the number of significant singular
values in $A$. As traditional rank counts a matrix as being “full-rank” even
if it has negligibly small but non-zero singular values, the effective rank
provides a more informative assessment of the matrix’s rank when used as the
readout from a NPDM. We assessed the effective rank of the linear readout for
5 Linear-NODE models (with state dimensionality of $\hat{D}=2,3,5,8,10$,
respectively) trained on synthetic neural data generated by linearly embedding
trajectories from the Arneodo system (Section G.1.1) into log-firing rates,
and found that while the reconstruction performance improved as $\hat{D}$
increased, the effective rank plateaued at erank $\approx 4$ (Fig 2A, Supp.
Fig. S3).

Figure S3: Linear-NODE trained on synthetic neural activity from linearly-
embedded Arneodo system

###  I.2 Cycle Consistency

To directly compare injectivity of the Flow readout versus the MLP, we
quantified how well each model’s inferred latent activity could be recovered
from the reconstructed log-rates. To do this, we took our fully-trained 10D
ODIN and MLP-NODE models (shown in Fig. 2C, $\hat{D}=10$) and obtained the
inferred latent activity $\hat{\mathbf{z}}$ and predicted log-firing rates
$\log\hat{\mathbf{y}}$ from the Arneodo dataset. Then, we trained a separate
network $h:\log\mathbf{Y}\rightarrow\mathbf{Z}$ to minimize the mean squared
error between its output $\tilde{\mathbf{z}}$ and the model-inferred latent
activity $\hat{\mathbf{z}}$ (see Table S3 for hyperparameters).

$\displaystyle\tilde{\mathbf{z}}=h\left((\log\hat{\mathbf{y}})\right)$ (24)

We computed the coefficient of determination between the re-generated latent
activity $\tilde{\mathbf{z}}$ and inferred latent activity $\hat{\mathbf{z}}$.
If this performance is high, the inferred latents can be recovered from the
log-rates suggesting that the readout is approximately injective.

$\displaystyle\text{Cycle
Consistency}=R^{2}(\tilde{\mathbf{z}},\hat{\mathbf{z}})$ (25)

Table S3: Training hyperparameters (Cycle-Consistency MLP, $h$) Parameter | Value
---|---
Batch Size | 2048
Learning Rate | 1e-3
Hidden Layers | 3
Hidden Size | 128
Epochs | 1000

It is possible for a readout to be fully injective (i.e., that $\hat{g}^{-1}$
exists), but still compress some features of latent activity into negligibly
small contributions to the predicted firing rates, making the readout
effectively, if not technically, non-injective. We reasoned that if this were
the case, the inverse mapping $h$, in order to properly invert the warping
applied by $\hat{g}$, would be highly sensitive to noise. We expect that such
noise perturbations would be warped by $h$ into large changes in the predicted
latents. Using the models trained without noise, we computed the $R^{2}$ of
re-generated latents $\tilde{\mathbf{z}}$ compared to the inferred latents
$\hat{\mathbf{z}}$. We therefore consider both the noise-free and noise-
corrupted cycle consistency scores as indicators of the approximate
injectivity of each readout, taking into consideration undue distortion
applied in the process of learning the injective mapping.

$\begin{gathered}\tilde{\mathbf{z}}_{\sigma}=h(\log\hat{\mathbf{y}}+\epsilon_{\sigma}),\quad\epsilon_{\sigma}=\mathcal{N}(0,\,\sigma),\
\ \sigma\in[0.01,0.05,0.1,0.2,0.5])\\\ \end{gathered}$ (26)
$\begin{gathered}ccR^{2}_{\sigma}=R^{2}(\tilde{\mathbf{z}}_{\sigma},\hat{\mathbf{z}})\end{gathered}$
(27)

###  I.3 Alternative injective readout

As an additional confirmation that injectivity was the critical addition to
non-linear readouts that made latent recovery more robust, we tested an
alternative injective architecture — an invertible neural network (INN) [20].
We found that using a 6-layer INN in place of the Flow readout had comparable
Rate $R^{2}$ and State $R^{2}$ to ODIN, and that, like ODIN, State $R^{2}$ was
stable as $\hat{D}$ increased beyond $D=3$ (see Supp. Fig. S4). Each INN layer
was composed of coupling, permutation and affine transformations. Additional
training parameters are noted in S4. This result further supports our claims
that injective networks empirically promote robust latent recovery.

Unfortunately, the INN hidden layer size is obligated to be the size of either
the input or output dimensionalities, whichever is larger. Therefore, in
realistic biological datasets where the number of neurons can be highly
variable across datasets, the capacity of the INN readout is intrinsically
linked to the number of recorded neurons. For this reason, we chose to use the
Flow readout, which decouples the computational capacity of the injective
transformation from the dimensionality of the neural space.

Table S4: Training hyperparameters INN (Synthetic Arneodo Data) Parameter | Value
---|---
Batch Size | 650
Learning Rate | 1.88e-4
Encoder Hidden Size | 100
Dropout | 0.05
NODE Hidden Layers | 6
NODE Hidden Size | 128
Readout Hidden Layers | 6
Readout Hidden Size | 12
Figure S4: Invertible Neural Network readouts produce qualitatively similar
results to Flow readout models. Data shown is the same as Fig. 2C, except
overlaid with INN readout model (purple)

##  J Fixed point finding and characterization

For each model (Linear-NODE, MLP-NODE and ODIN), we located fixed points (FPs)
by finding the positions in the latent space that minimized the norm of the
vector field via the objective $q=\frac{1}{2}\lVert\hat{f}\rVert_{2}^{2}$ [1,
40]. We initialized our search with 1024 randomly sampled initial states from
along inferred latent trajectories. We used Adam with a learning rate of 5e-2
to minimize the $q$-value for each point independently over 10,000 iterations.
Candidate points that did not achieve a $q$-value less than a magnitude of
7e-3 were excluded. As more than one candidate can approach the same FP, we
combined candidate points that were within a specified distance, $\epsilon$ =
1, from one another. In practice, points that were excluded had much larger
$q$-values than the putative fixed points. We then linearized the dynamics
around each FP and computed the system Jacobian to determine the stability and
rotational character of the system around these FPs.

##  K Metrics

###  K.1 Synthetic data metrics

####  K.1.1 Rate reconstruction (Rate $R^{2}$)

We computed the coefficient of determination between true ($\mathbf{Y}$) and
predicted ($\hat{\mathbf{Y}}$) rates for each neuron, and reported the average
value across neurons.

$\text{Rate
}R^{2}=R^{2}(\mathbf{Y},\hat{\mathbf{Y}})=\frac{1}{N}\sum_{i=0}^{N}1-\frac{\sum(\mathbf{y}_{i}-\hat{\mathbf{y}}_{i})^{2}}{\sum(\mathbf{y}_{i}-\bar{\mathbf{y}}_{i})^{2}}$

####  K.1.2 Latent state reconstruction (State $R^{2}$)

To compute State $R^{2}$, we concatenated a vector of ones with the true
latent states ($\mathbf{Z_{1}}$), then used the pseudoinverse to find the
optimal affine transformation from the true latents to the inferred latents
($\hat{\mathbf{Z}}$) (i.e., optimal linear estimation). We computed the
coefficient of determination ($R^{2}$) between the true and inferred latent
activity with the same equation as in K.1.1.

$\displaystyle\mathbf{W}_{z}$
$\displaystyle=\mathbf{Z}_{1}^{\dagger}\hat{\mathbf{Z}}$ (28)
$\displaystyle\text{State}\>R^{2}$
$\displaystyle=R^{2}(\hat{\mathbf{Z}},\mathbf{Z}_{1}\mathbf{W}_{z})$ (29)

####  K.1.3 Activation function comparison

We developed a method for deriving an estimate of the inferred activation
functions $\hat{\mathbf{\psi}}_{i}(\cdot)$ for a comparison to the true
activation functions $\mathbf{\psi_{i}(\cdot)}$ (see Equation 17). We
projected the true encoding vectors $\boldsymbol{\gamma}_{i}$ into the latent
space of the model via the affine transformation $\mathbf{W}_{z}$ (see section
K.1.2). We then used these encoding vectors
$\hat{\boldsymbol{\gamma}}_{i}\in\mathbb{R}^{\hat{D}}$ to convert inferred
latent states $\mathbf{\hat{Z}}\in\mathbb{R}^{T\times\hat{D}}$ into an
activation $\hat{\mathbf{a}}_{i}\in\mathbb{R}^{T}$ for each neuron.

$\displaystyle\hat{\boldsymbol{\gamma}_{i}}$
$\displaystyle=\boldsymbol{\gamma}_{1,i}\mathbf{W}_{z},\ \text{for
}i=1,2,\cdots,N$ (30) $\displaystyle\hat{\mathbf{a}}_{i}$
$\displaystyle=\hat{\mathbf{Z}}\cdot\hat{\boldsymbol{\gamma}_{i}}$ (31)

To estimate the activation function for a given neuron $i$, we need pairs of
inferred activations $\hat{\mathbf{a}}_{i}$ and firing rates
$\hat{\mathbf{y}}_{i}$. For each neuron, we split firing rates into 20
quantiles and computed the corresponding median activation
$\hat{\mathbf{a}}_{i,1:20}^{med}$ and firing rate
$\hat{\mathbf{y}}_{i,1:20}^{med}$ within each quantile.

$\displaystyle\hat{\mathbf{y}}_{i,1:20}^{med},\hat{\mathbf{a}}_{i,1:20}^{med}=\text{Quantize}(\mathbf{\hat{y}}_{i},\mathbf{\hat{a}}_{i},20)$
(32)

We represented the inferred activation function
$\mathbf{\hat{\psi}}_{i}(\cdot)$ using these activation-firing rate pairs. We
then performed the same procedure on the true rates and activations to find a
similar representation of the true activation function
$\mathbf{\psi}_{i}(\cdot)$ for each neuron. To compare the true activation
function $\psi(\cdot)$ to the estimated activation function
$\hat{\psi}(\cdot)$, we combined the activations of each neuron $i$ and its
corresponding firing rate as the columns of the matrices:

$\mathbf{\hat{\Psi}}_{i}=\begin{pmatrix}\mathbf{\hat{a}}_{i}^{med}&\mathbf{\hat{y}}_{i}^{med}\\\
\end{pmatrix},\quad\mathbf{\Psi}_{i}=\begin{pmatrix}\mathbf{a}_{i}^{med}&\mathbf{y}_{i}^{med}\\\
\end{pmatrix}$

Because the inferred latent activity can be scaled and translated arbitrarily
with respect to the true latent activity, we found the optimal affine
transformation between $\mathbf{\hat{\Psi}}_{i}$ and $\mathbf{\Psi}_{i}$. We
used the $R^{2}$ of this mapping to quantify the correspondence between the
two activation functions $\mathbf{\hat{\psi}}_{i}(\cdot)$ and
$\mathbf{\psi}_{i}(\cdot)$ for each neuron.

###  K.2 Neural Latents Benchmark metrics

####  K.2.1 Co-smoothing bits-per-spike (co-bps)

A common failure mode of many dynamics models is to find latent activity that
can accurately reconstruct the firing rates of neurons seen by the encoders,
but fails to reconstruct neural activity of held-out neurons. To avoid this
pitfall, we used a previously developed metric called co-smoothing bits-per-
spike which evaluates reconstruction performance on a set of held-out neurons
not visible to the encoders [27]. At a high-level, this metric quantifies how
well the firing rates of the held-out neurons can be predicted from the
spiking of the held-in neurons (see G.2). This metric is defined by Equation
33 for each held-out neuron.

$\text{co-bps}=\frac{1}{n_{s}\log
2}(\mathcal{L}(\mathbf{\hat{y}}_{n,t};\mathbf{x}_{n,t})-\mathcal{L}(\mathbf{\bar{y}}_{n,:};\mathbf{x}_{n,t}))$
(33)

where $\mathbf{\bar{y}}_{n,:}$ is the mean firing rate for neuron $n$ across
time, $n_{s}$ is the total number of spikes for that neuron,
$\mathbf{\hat{y}}_{n,t}$ is the predicted firing rate from the model at time
$t$, $\mathbf{x}_{n,t}$ represents the observed spiking of that neuron at time
$t$, and $\mathcal{L}$ represents the Poisson log-likelihood. More information
can be found in [27].

####  K.2.2 Velocity decoding $\mathbf{R}^{2}$

A common metric of performance is how well inferred firing rates can be used
to predict behavioral variables, as this can be used downstream for decoding
intent in clinical applications like brain-computer interfaces [51]. For the
Maze dataset, hand velocity has been shown to be highly correlated with the
neural firing in motor cortices. We compute this metric using the method from
[27], in which a ridge regression model is trained to predict the observed
hand velocity from inferred firing rates. The coefficient of determination
($R^{2}$) was then evaluated on validation data that was not used to train the
ridge regression velocity decoder.

##  L Compute resources

We used an internal computing cluster with a total of 30 Nvidia GeForce RTX
2080 Ti GPUs for model training. Each model trained on simulated neural data
took approximately 3 hours to train, while each model trained on real
biological data took approximately 1.5 hours to train. With 2 models training
on each GPU, the 100 models included in Figs. 2, 3, 4, and 5 took
approximately 150 GPU-hours and the 50 NODE-based models included in Fig. 6
took approximately 37.5 GPU-hours. FP finding was fast, requiring 1 minute for
each model.

##  M Open-source packages used

  * •

torch [52] (BSD license): Deep learning framework providing layer definitions,
GPU acceleration, automatic differentiation, optimization, and more.

  * •

pytorch_lightning (Apache 2.0 license): Lightweight wrappers for model
training.

  * •

ray.tune [53] (Apache 2.0 license): Distributed hyperparameter tuning.

  * •

dysts [47] (Apache 2.0 license): Implementations for modeled dynamical
systems.

  * •

fixed_point_finder [40] (Apache 2.0 license): Inspiration for torch-based
fixed point finder.

  * •

FrEIA (MIT license): Implementation of alternative Invertible Neural Network
architecture.

  * •

scikit-learn[54] (BSD License): Implementations of linear regression models
and principal component analysis.


# Function Approximation with Quantum Circuit

Alberto Delgado  Department of Electrical and Electronics Engineering,
Universidad Nacional de Colombia, Bogotá<EMAIL_ADDRESS>
(August 27, 2024)

###### Abstract

A mathematical proposition with a trainable pair, operator and quantum
circuit, are introduced to approximate functions expressed as cubic Taylor
polynomials, numerical simulations illustrate three cases.

The motivations behind this paper are twofold, to explore applications of
small quantum circuits [1] and to show that sigmoid functions can be
approximated with quantum circuits, it is important to mention that sigmoid
functions are the building blocks of neural networks [2]. Physical quantum
systems include parameterized quantum gates [3], these gates offer the
opportunity to change a quantum state with external data or even train quantum
circuits to estimate unknown probability distributions [4].

The paper is organized as follows. Section one, introduces a parameterized
quantum circuit and formulates a mathematical proposition to prove its
function approximating capabilities. Section two, presents the simulation
results to approximate three functions, quadratic, Gaussian, and sigmoid.
Finally, section three is the summary.

##  1 Expectation Value as Polynomial

A mathematical proposition proves that the expectation value of an operator G,
for a two qubits parameterized quantum circuit, approximates a cubic
polynomial with tunable coefficients.

###  1.1 Quantum Circuit

Consider the two qubits parameterized quantum circuit, Fig. 1, where x is an
independent variable and $(\theta_{1},\theta_{2})$ are parameters. The
expectation value $\hat{f}(x)\hskip 2.84526pt=\hskip 2.84526pt<G>$ can
approximate a function $f(x)$, represented with a cubic polynomial in Taylor
series, by tuning parameters $(\theta,G)$ (1) to minimize a performance index
J (2) over a set of N samples in the domain $x\in[-x_{0},+x_{0}]$.

$\displaystyle\theta=\begin{bmatrix}\theta_{1}\\\
\theta_{2}\end{bmatrix};\hskip 5.69054ptG=\begin{bmatrix}g_{0}&0&0&0\\\
0&g_{1}&0&0\\\ 0&0&g_{2}&0\\\ 0&0&0&g_{3}\end{bmatrix}$ (1) $\displaystyle
J=\sum\limits_{k=1}^{N}(f(x_{k})-\hat{f}(x_{k}))^{2}$ (2)

Figure 1: Function approximation with quantum circuit, $\hat{f}(x)\hskip
2.84526pt=\hskip 2.84526pt<G>$. Parameters $(\theta_{1},\theta_{2})$ and
operator G are adjusted so $\lvert f(x)-\hat{f}(x)\lvert\hskip
2.84526pt\leq\hskip 2.84526pt\epsilon$, $x\in[-x_{0},+x_{0}]$.

###  1.2 Cubic Polynomial

It is shown that the expectation value is equivalent to a cubic polynomial by
using Taylor series, the coefficients depend on parameters (1).

###### Theorem 1

Consider the quantum state $\lvert\psi(\theta_{1},\theta_{2},x)>$ and operator
G,

$\displaystyle\lvert\psi(\theta_{1},\theta_{2},x)>\hskip 2.84526pt=\hskip
2.84526ptU\lvert 00>$ (3) $\displaystyle U=R(x-\theta_{2})\otimes
R(x-\theta_{1})(H\otimes H)$ (4) $\displaystyle
R(x-\theta_{i})=\begin{bmatrix}\cos(\frac{x-\theta_{i}}{2})&-\sin(\frac{x-\theta_{i}}{2})\\\
\\\ \sin(\frac{x-\theta_{i}}{2})&\cos(\frac{x-\theta_{i}}{2})\end{bmatrix}$
(5) $\displaystyle G=\begin{bmatrix}g_{0}&0&0&0\\\ 0&g_{1}&0&0\\\
0&0&g_{2}&0\\\ 0&0&0&g_{3}\end{bmatrix}$ (6)

The expectation value has the structure of a cubic polynomial,

$\displaystyle\hat{f}(x)\hskip 2.84526pt=\hskip
2.84526pt<\psi(\theta_{1},\theta_{2},x)\lvert
G\lvert\psi(\theta_{1},\theta_{2},x)>$ (7)
$\displaystyle\hat{f}(x)=a_{0}+a_{1}x+a_{2}x^{2}+a_{3}x^{3}$ (8)

the coefficients $a_{i}$ are function of parameters
$(\theta_{1},\theta_{2},g_{0},g_{1},g_{2},g_{3})$.

Proof:

  1. 1.

Calculate $U_{1}=R(x-\theta_{2})\otimes R(x-\theta_{1})$

  2. 2.

Multiply $U_{1}$ by $(H\otimes H)\lvert 00>$

  3. 3.

Replace $\cos(\frac{x-\theta_{i}}{2})$ and $\sin(\frac{x-\theta_{i}}{2})$ with
the first terms of their Taylor series.

The amplitudes of the quantum state are quadratic polynomials, calculating the
expectation value (7) results in (8).

##  2 Simulations

A bioinspired training algorithm known as chemotaxis [5, 6] and implemented in
octave [7] was used to minimize the performance index (2) by tuning the set
(1). The number of samples is N = 30 in the interval with $x_{0}=1.5$.

###  2.1 Quadratic function

Fig. 2 shows a quadratic function $f(x)=x^{2}$ (red) and its approximation
$\hat{f}(x)$ (black). The final performance index was J = 0.03, parameters (9)
and (10).

$\displaystyle\theta$ $\displaystyle=\begin{bmatrix}1.373\\\
1.770\end{bmatrix}$ (9) $\displaystyle G$
$\displaystyle=\begin{bmatrix}-0.081&0&0&0\\\ 0&2.260&0&0\\\ 0&0&2.272&0\\\
0&0&0&4.954\end{bmatrix}$ (10)

Figure 2: Quadratic function $f(x)=x^{2}$ (red), 5000 iterations, final J =
0.03.

###  2.2 Gaussian function

The second function is Gaussian $f(x)=e^{-x^{2}}$ shown in Fig. 3 (red) and
its approximation $\hat{f}(x)$ (black). After training, the final set of
parameters is given in (11) and (12) with a performance index (2) J = 0.005.
Notice the small differences between the actual function and the
approximation, there is a compromise between quantum circuit complexity and
the final performance index.

$\displaystyle\theta$ $\displaystyle=\begin{bmatrix}0.497\\\
-0.498\end{bmatrix}$ (11) $\displaystyle G$
$\displaystyle=\begin{bmatrix}-0.088&0&0&0\\\ 0&1.152&0&0\\\ 0&0&1.711&0\\\
0&0&0&-0.089\end{bmatrix}$ (12)

Figure 3: Gaussian function $f(x)=e^{-x^{2}}$ (red), after 5000 iterations,
the final J = 0.005.

###  2.3 Sigmoid function

Fig. 4 shows the sigmoid function $f(x)=tanh(x)$ (red) and its approximation
$\hat{f}(x)$ (black).

Figure 4: Sigmoid function $f(x)=tanh(x)$ (red), final J = 0.006, 5000
iterations.

The final performance index was J = 0.006, parameters (13) and (14).

$\displaystyle\theta$ $\displaystyle=\begin{bmatrix}0.266\\\
0.069\end{bmatrix}$ (13) $\displaystyle G$
$\displaystyle=\begin{bmatrix}-0.885&0&0&0\\\ 0&0.055&0&0\\\ 0&0&0.466&0\\\
0&0&0&0.931\end{bmatrix}$ (14)

##  3 Summary

A mathematical proposition has been proved to show that the expectation value
of an operator for a two qubits parameterized quantum circuit can approximate
a function expressed as a third degree Taylor polynomial. Simulations for
three functions, quadratic, Gaussian, and sigmoid, agree with the theory. The
approximation error is related to the complexity of the quantum circuit.

Neural networks are trainable universal maps, here an operator with a
parameterized quantum circuit are used to approximate a sigmoid function, the
building block of neural networks.

## References

  * [1] Preskill, J.: Quantum computing in the NISQ era and beyond (2018-01-27).
  * [2] Norgaard, M., Ravn, O., Poulsen, N.K., Hansen, L.K.: Neural Networks for Modelling and Control of Dynamic Systems, Springer, London, UK, 2001.
  * [3] Q Composer IBM Quantum Experience.
  * [4] Benedetti, M., Garcia-Pintos, D., Nam, Y., Perdomo-Ortiz, A.: A generative modeling approach for benchmarking and training shallow quantum circuits (2018-01-23).
  * [5] Bremermann H.J., Anderson R.W. (1991) How the Brain Adjusts Synapses—Maybe. In: Boyer R.S. (eds) Automated Reasoning. Automated Reasoning Series, vol 1. Springer, Dordrecht. DOI: 10.1007/978-94-011-3488-0-6.
  * [6] Delgado, A.: Control of Nonlinear Systems Using a Self-Organising Neural Network, NCA (2000) 9:113. DOI: 10.1007/s005210070022.
  * [7] John W. Eaton, David Bateman, Søren Hauberg, Rik Wehbring (2014). GNU Octave version 3.8.1 manual: a high-level interactive language for numerical computations. CreateSpace Independent Publishing Platform. ISBN 1441413006.