Targeted Neural Dynamical Modeling
Authors: Cole Lincoln Hurwitz
Akash Srivastava
Kai Xu
Justin Jude
Matt Perich
Lee E. Miller
Matthias H. Hennig
Venue: NeurIPS 2021 Poster (2021)
Keywords: neuroscience
latent variable model
recurrent neural network
behavioural neuroscience
motor neuroscience
dynamics

Abstract: Latent dynamics models have emerged as powerful tools for modeling and interpreting neural population activity. Recently, there has been a focus on incorporating simultaneously measured behaviour into these models to further disentangle sources of neural variability in their latent space. These approaches, however, are limited in their ability to capture the underlying neural dynamics (e.g. linear) and in their ability to relate the learned dynamics back to the observed behaviour (e.g. no time lag). To this end, we introduce Targeted Neural Dynamical Modeling (TNDM), a nonlinear state-space model that jointly models the neural activity and external behavioural variables. TNDM decomposes neural dynamics into behaviourally relevant and behaviourally irrelevant dynamics; the relevant dynamics are used to reconstruct the behaviour through a flexible linear decoder and both sets of dynamics are used to reconstruct the neural activity through a linear decoder with no time lag. We implement TNDM as a sequential variational autoencoder and validate it on simulated recordings and recordings taken from the premotor and motor cortex of a monkey performing a center-out reaching task. We show that TNDM is able to learn low-dimensional latent dynamics that are highly predictive of behaviour without sacrificing its fit to the neural data.

Targeted Neural Dynamical Modeling
Cole Hurwitz
Akash Srivastava MIT-IBM Watson Al Lab Cambridge, MA 02142, Akash.Srivan @rivastavaa bm com
School of Informatics University of Edinburgh Edinburgh, Scotland, EH8 9AB colehurwitz@gmail. com
Kai Xu School of Informatics University of Edinburgh Edinburgh, Scotland, EH8 9AB xuka19211108gmal com
Justin Jude School of Informatics University of Edinburgh Edinburgh, Scotland, EH8 9AB justin, justin.jude@ed.. judeded. ac. uk
Matthew G. Perich
Lee E. Miller Feinberg School of Medicine
Icahn School of Medicine at Mount Sinai
New York, NY 10029
Northwestern
mperich@gmail.com com
Chicago, IL 60611
im@nortwwesterc
Matthias H. Hennig School of Informaties University Edinburgh Edinburgh, Scotland, EH8 9AB m. beenn ac. uk
Abstract
Latent dynamics models have emerged as powerful tools for modeling and inter- preting neural population activity. Recently, there has been a focus incorporating simultaneously measured behaviour into these models to further disentangle sources of neural variability in their latent space. These approaches, however, are limited in their ability to capture the underlying neural dynamics (e.g. linear) and in their ability to relate the learned dynamics back to the observed behaviour (e.g. time lag). To this end, we introduce Targeted Neural Dynamical Modeling (TNDM), nonlinear state-space model that jointly models the neural activity and external behavioural variables. TNDM decomposes neural dynamics into behaviourally relevant and behaviourally irrelevant dynamics; the relevant dynamics to reconstruct the behaviour through flexible linear decoder and both sets dynamics are used to reconstruct the neural activity through linear decoder with no time lag. We implement TNDM as a sequential variational autoencoder and validate it simulated recordings and recordings taken from the premotor and motor cortex of a monkey performing a center-out reaching task. We show that TNDM is able learn low-dimensional latent dynamics that highly predictive of behaviour without sacrificing its fit to the neural data.
used
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Introduction
Recent progress in high-density, microelectrode arry technology now allows for recording from
hundreds to thousands of neurons with the precision of single spikes [I]. Despite the apparent high dimensionality of these datasets, neural activity is often surprisingly well-explained by low- dimensional latent dynamics [ |2 6 8). Extracting these dynamics from single trials is crucial for understanding how neural activity relates a behavioural task or stimulus
Latent variable models (LVMs) are a natural choice for capturing nww-dimensional structure from neural activity they can learn to map a few latent variables to arbitrarily complicated response structure in the activity. Already, there exist a number of LVMs that have been successfully applied to neural data ranging from simple non-temporal models such principal components analysis (PCA) [5] to complex state-space models such as LFADS [7]. In these models, the goal is to learn a set of latent factors that best explain neural variability. As such, there is no guarantee that the different sources of variability present in the population activity will be disentangled in the latent space (e.g.
behaviour, arousal, thirst, etc.) 25 [0]
To better partition sources of neural variability in the latent space, some LVMs have been developed that incorporate an external behaviour into the generative process [14 however, do not model temporal dependencies between the latent states. Recently, a novel state-space model termed preferential subspace identification (PSID) was developed that jointly models neural activity and behaviour with a shared set of dynamics 225 When applied to neural activity recorded in the premotor cortex (PMd) and primary motor cortex (M1) of monkey during 3D reaching task, PSID was shown to extract latent factors that were more predictive of behaviour than the factors extracted by other approaches. Despite the strength and simplicity of this approach, it suffers from two main drawbacks. First, PSID is a linear state-space model and cannol capture the nonlinear dynamics which are thought to underlie phenomena such as rhythmic motor patterns [22] g or decision making [2 Second, PSID assumes that behaviourally relevant dynamics explain both the neural activity and behaviour with no time lag. This limits the ability of PSID to capture more complex temporal relationships between the latent dynamics and the behaviour.
19 3T). These methods,
In this work, we introduce Targeted Neural Dynamical Modeling (TNDM), nonlinear state-space model that jointly models neural activity and behaviour. Similarly to PSID, TNDM decomposes neural activity into behaviourally relevant and behaviourally irrelevant dynamics and uses the relevant dynamics to reconstruct the behaviour and both sets of dynamics to reconstruct the neural activity. Unlike PSID, TNDM does not constrain the latent dynamics each time step to explain behaviour at each time step and instead allows for any linear relationship (constrained to be causal in time) between the relevant dynamics and the behaviour of interest. We further encourage partitioning of the latent dynamics by imposing a disentanglement penalty on the distributions of the initial conditions of the relevant and irrelevant dynamics. To perform efficient inference of the underlying nonlinear dynamics, TNDM is implemented as a sequential variational autoencoder (VAE) [3 29 We
compare TNDM to PSID and to LFADS, a nonlinear state-space model that only models neural activity, to illustrate that TNDM extracts more behaviourally relevant dynamics without sacrificing its fit to the neural data. We validate TNDM on simulated recordings and neural population recordings taken from the premotor and motor cortex of monkey during center-out reaching task. In this analysis, we find that the behaviourally relevant dynamics revealed by TNDM are lower dimensional than those of other methods while being more predictive of behaviour.
2 Background/Related work
Notatation. Let E NNTT be the observed spike counts and let y R be the observed behaviour during single-trial? We define the unobserved latent factors in a single trial as E R MxT where M N. For TNDM, as with PSID, it is important to distinguish between behaviourally relevant 2, and behaviourally irrelevant 2; latent factors. The behaviourally relevant latent factors
3
'The code for running and evaluating TNDM on real data can be found at paper We also provide a Tensorflow2 re- re-implemention of TNDM at https://githubb. com/HHennigLab/tddnn It is important note that all reportraed results for the real datasets use the old model and not the re-implementation. For the synthetic dataset results, we use the re-implementation. 'In this work, we assume that behaviour is temporal and has the same time length as recorded neural activity (e.g. hand position). TNDM can be extended discrete/non temporal behaviours (e.g. reach direction).

---

2, summarize the variability in the neural activity associated with the observed behaviour while the behaviourally irrelevant latent factors 2, explain everything else in the neural data (Figure [l). We assume that each of the unobserved, single-trial factors can be partitioned into these relevant and irrelevant factors = == (i).
State-space models for neural data. There a number of state-space models that have been developed and applied to neural population activity. The expressivity of these models range from simple linear dynamical models [27 2 (00] to more complex nonlinear models where the latent dynamics are parameterized by recurrent neural networks (RNNs) 17, 26
For this work, there are
two state-space models that are most relevant: LEADS and PSID.
LFADS, or latent factor analysis via dynamical systems, is a sfate-of-the-ard nonlinear state-space model for neural data. In LFADS, the latent dynamics are generated by sampling high-dimensional initial conditions go from some distribution Pao and then evolving 9o with deterministic RNN fe. A linear mapping W₂ is then applied to the high-dimensional dynamics di to transform them into the low-dimensional 'dynamical" factors 2. These dynamical factors are transformed into spike counts by mapping each time point to a rate parameter r of Poisson distribution using a weight matrix W, followed by an exponential nonlincarity. The generative process is defined as: Do p It. fo(gt-1; 2 (g(g)), r. exp(W,(z4))+ conditions Go are inferred from & with an RNN encoder network 4s- Utilizing the reparameterization trick, the model is trained using gradient descent and by optimizing the evidence lower-bound (ELBO) of the marginal log-likelihood While LFADS provides an excellent fit to the neural data, it inevitably mixes different sources of neural variability in the latent dynamics z as there is no constraints imposed to disentangle these dynamics.
Poisson (rr.rr)) The initial
PSID is linear dynamical model that partitions the latent dynamics into behaviourally relevant and behaviourally irrelevant dynamics 2 == {24, }. The dynamical evolution of 2 is defined by a transition matrix A along with Gaussian noise term Wg. 2 is transformed into the observed firing rates x by mapping each time point to the mean of a Gaussian distribution with a weight matrix Wz and noise term Wzz The behaviourally relevant dynamics 2, each time point are transformed into the observed behaviour y with a weight matrix Wy. The state space model for PSID is then defined -N(x (x:[W,(2>)) wx), Vt W₂ (2,).. PSID uses a novel two-stage
as: 2, N(=]A(==-), (=))) subspace identification approach to learn the parameters of their model. In the first stage, PSID extracts the behaviourally relevant dynamics through an orthogonal projection of future behaviour onto past neural activity. The irrelevant dynamics are then extracted through an additional orthogonal projection of residual neural activity onto past neural activity. In comparison to LFADS, PSID was shown to extract latent states that are better able to predict behaviour when using a Kalman filter. Despite the analytical simplicity of this approach, it suffers from a few drawbacks. First, it can only model linear dynamics which may not provide good fit to nonlinear activity patterns or behaviours (e.g. multi-directional reaches). Second, the relevant dynamies at each time step 2,z must be mapped one-to-one to the behaviour Y/ during training (i.e. no time lag). This imposes strong structural constraint on the relevant dynamics which hampers their ability to explain neural variability.
3 Model
In this work, we introduce Targeted Neural Dynamical Modeling (TNDM). TNDM is nonlinear state-space model which jointly models neural activity and an observed behaviour. Crucially. TNDM learns to reconstruct both the population activity and behaviour by disentangling the behaviourally relevant and behaviourally irrelevant dynamics that underlie the neural activity.
Generative model. A plate diagram of TNDM's graphical model is shown in Figure la We assume that the observed neural activity x and behaviour y in each trial are generated by two sets of latent factors 2 and Zy.
The generative process of TNDM is defined below:
gio Paio:9ro ~ Pare Pie = So gre -So(9re-1)
2/ = Wiz(gu), 2,t (gr( re (eie
(1)
Tt ~ Poisson(xe(1r), y y~ Nyy ((yC 1Z)
3

---

Relevant Si
Relevent
Relevent
g0>
90r
Spikes
Z
Zr
Newses
Rates
Spikes
Penalty
Enooder RNN
K
Inslevant
(a) TNDM graphical model
(b) TNDM architecture
Figure I: (a) The latent space of TNDM is partitioned into irrelevant and relevant high-dimensional initial conditions go; and go,- These initial conditions are deterministically transformed to recover the latent factors 2; and Z,. which give rise to the jointly observed neural activity and behaviour y- We assume there are K trials in the dataset. (b) TNDM utilizes sequential variational autoencoding approach ammortize inference of the relevant and irrelevant initial conditions Irn and 9to- The initial conditions are passed through two separate RNNs to generate the behaviourally relevant and irrelevant dynamics 9i and 9. which are then projected into a low-dimensional subspace recover the dynamical factors and 2,- These factors are used to reconstruct the neural activity and behaviour. The behaviour is reconstructed from the relevant factors using flexible linear decoder which can capture complex temporal relationships (see the paragraph on Behaviour Decoding).
In the above equation, Paio and Pg.a are the distributions over the initial conditions of the behaviourally irrelevant and behaviourally relevant dynamics, respectively (assumed be Gaussian). Similarly to LFADS, we parameterize the nonlinear transition functions J'; and fe, using RNN 5T As Jir and gr: can have arbitrarily dimenssinnii in our model (defined by the number of units in the RNN), we utilize two weight matrices, Wi, and Wr.. to project these nigh-dimenssonal dynamics into a low-dimensional subspace, giving rise the relevant and irrelevant dynamical factors zi, and Zr: at each time step. These factors are then used to reconstruct both the observed neural activity and behaviour using linear decoders. An essential feature of our generative model is that although neural activity is reconstructed from the latent dynamics at each time step (i.e. no time lag), we let the relevant factors reconstruct the behaviour with more flexible linear decoder Cy that allows for time lags (explained in the Behaviour Decoding paragraph). We fix the variance of the Gaussian observation model for the behaviour (i.e. we utilize mean squared error) as naively inferring the variance in the decoder can lead to numerical instabilities (23).
is important understand that although the dimensionality of the dynamics in TNDM (and LFADS) can be arbitrarily high, the dimensionality of the subspace that gives rise to neural activity and behaviour will be low due to the projection. Therefore, our model can be used to examine the number of latent variables (i.e. activity patterns) that are needed to characterize the population response and corresponding behaviour. As this is the primary goal when fitting LVMs to neural data [5], we compare all LVMs in this paper (TNDM, LFADS, and PSID) by the dimensionality of this subspace rather than the dimensionality of the dynamics.
Behaviour Decoding. As mentioned above, PSID utilizes a linear weight matrix that maps the relevant latent dynamics each time step to the behaviour at each time step, ie yr W, (2,,)) This parameterization does not allow for modeling any latency, long-term dependencies correlations, therefore, it severely limits the ability of the relevant dynamics to simultaneously explain neural activity and behaviour. To demonstrate the drawbacks of the no time lag behaviour decoder, we show that while training TNDM using this decoder leads to accurate behaviour prediction, the reconstruction of the neural activity noticeably decreases. This issue gets exacerbated in models with nonlinear dynamics as the expressivity of the underlying RNNs, along with the inflexible one-to-one behaviour mapping, leads the relevant dynamics to simply learn replicate the behaviour of interest. These results are summarized in Supplement 1.
'To implement TNDM, we primarily adapt the original Tensorflow [ implementation of LFADS from https://github.. com wwt.dlaaaaaamaaaa httpm//laaa//ddlleaae///////sses (Apache License 2.0).

---

To overcome this limitation we instead allow the relevant latent dynamics to reconstruct the behaviour through any learned linear causal relationship. To this end, we introduce a linear weight matrix C, with dimensionality n=..7' ? x BT where nz,,, is the number of relevant dynamics, B is the number of behaviour dimensions, and T is the time length of a single trial. To transform the relevant factors E, into behaviours using Cg. we concatenate each dimension of 2, in time to form 1D vector Z, with length nz, I and then perform the operation Y = Cr where Y is the resulting concatenated behaviour. As Y is a ID vector with length BT, we can reshape Y to recover the reconstructed behaviour V.. Importantly, we do not allow acausal connections in Cy, i.c. the lower triangular components of each of the dynamics to behaviour blocks set zero during training. For an example weight matrix Cy. see Figure In comparison to a simple no time lag mapping, we find that our flexible, causal linear decoder allows the relevant latent dynamics to both reconstruct the measured behaviour and capture neural variability. This is shown in Supplement 1 where the behaviourally relevant factors learned by TNDM with the full causal decoder contribute more meaningfully the neural reconstruction than when using the no time lag decoder.
Inference To extract the latent dynamics 2, and 2; from the neural activity 2, first learn to approximate the posterior over the initial conditions of the dynamics Gro and gio- Then, we learn a deterministic mapping from the initial conditions the latent dynamics zr and zę- To approximate the true posterior p(Gro- 9:g xr, y). we implement TNDM as a sequential VAE and define our variational posterior as the product of two independent multivariate Gaussian distributions with diagonal covariances. The variational parameters of each Gaussian are computed with shared encoder network Ed,, followed by separate linear transformations. This shared encoder is implemented as a bi-directional RNN such that the initial conditions are inferred using information from the full trial. The equation for inferring the variational posterior is shown below:
go (o,(2rol (Mss,, (cpo (x)), 0² oer (co0 (x))) (uso, (2)) oi (Rop (1)))
The inference networks for the behaviourally relevant and the behaviourally irrelevant initial con- ditions are parameterized by T, = {du, Φri Φr,) } and $; = { di, Φ where do are the parameters of the shared RNN. is important to note that TNDM's variational posterior only depends on the neural activity x. This approximation forces the learned initial conditions to come from the observed activity and allows for decoding of unseen behaviours after training. The reparameterization trick is used to sample from each initial condition distribution and the sampled initial conditions evolved using separate decoder RNNs to produce the behaviourally relevant Gr, and irrelevant digh-dimensional dynamics 9it. The high-dimensiona dynamics at each time-step are projected into a oow-iimesiional subspace recover the low-dimensional dynamical factors 2yt and 2,,- The neural activity and behaviour y are generated from the latent factors z as shown in Equation]
The ELBO for the observed data from a single trial is therefore defined as:
ELBO(z,y) KI. go, Pa-o]
KI. [go,
Pora] E.g,, ro. [log po, (xlgi fr po, (u19)) (3)
where pa, and pa, are the observation models for the neural activity and behaviour, respectively.
Disentangling the latent dynamics Despite factorising the variational posterior, the true posterior over the latent variables P( Ix, y) cannot be factorized; that is, 2, and z,. (which are determin- istic transforms of go, and 90,)) are conditionally dependent given the observed data. This means that 24 and z, will be statistically dependent. To reduce sharing of information and redundancy between these two sets of dynamics, we introduce novel disentanglement penalty on the two initial condition distributions. For this penalty, we take inspiration from the domain of unsupervised disen- tangled representation learning where standard introduce additional penalties that encourage disentanglement or independence of the latent representations [15] 3 As mutual information is hard to compute for high-dimensional variables and the experimental data often has a very limited number of trials to estimate these distributions reliably, we instead penalize the mean of the sample cross-correlations between 9o, and go.- Importantly, this cross-correlation penalty is applied in such way that the final objective is still valid lower bound of the log- likeeloodoo (the cross-correlation
5

---

(a)
(b)
(c)
0.9
Poisson log nog-likiihood
-1750
TNDM
RMSE firing rates
0.1650
LFADS PSID
Behaviour R2
0.8
-1780
0.1625
0.7
-1770
Relevant
0.1600
factors
0.6
-1780
0.1575
2
0.5
-1790
0.1550
4
2
Factors
Factors
Factors
Figure 2: Summary of the behaviour and activity reconstruction accuracy for TNDM, LFADS, and PSID fit to neural recordings from the monkey primary motor cortex (M1) during a center-out reaching task. Each plot shows performance averaged over five fits with different initialisation (random seeds) and different random training-test data splits. Error bars show standard error of the mean. (a) Coefficient of determination (R for measured and reconstructed behaviour (hand position). For TNDM, the reconstruction performed by the behaviour decoder (with only the relevant factors) while for LFADS, ridge regression had to be used to decode the hehaviour ex post factor. For PSID the reconstruction by the model was additionally Kalman smoothed. trial for TNDM and LFADS. (c) Root mean square eгror (RMSE) between the predicted and actual ground-truth firing rates. Averaging was performed across all trials with the same movement direction. Behaviour reconstruction and log-likelihood were computed on held out test data, and the firing rate RMSE on the whole data set allow for more reliable averaging.
function of the total number of latent factors,
(b) Poisson log-likellne li-elikoohood for the activity reconstruction per single
penalty is always negative). We refer to this penalty as Qga qa:) and we adjust its weight with hyperparameter λo (see Supplement for ablation study of this penalty). The final objective function for TNDM is then:
J (z, y) KL ger Paral KL [99; Pa.ol + Egg, [log pe, (al gi g+)]
Ag.. [log Po, (ulg+)] + AqQ((o, qo.)
(4)
where Ås is an additional hyperparameter introduced balance the behavioural likelihood with the neural likelihood (see Supplement 3 for hyperparameter details). While TNDM is not the first VAE to jointly model two observed variables with a partitioned latent space [30], it distinguished by its unique objective function and penalties, its RNN-based architecture, its causal linear decoder, and its novel application to neural activity and behaviour.
Experiments
4.1 Simulated Data
We evaluate TNDM on synthetic spike trains generated from a Lorenz system, a common benchmark for state space models of neural activity. For detailed description of the simulated data and evaluation, we refer the reader to Appendix 6.
4.2 M1 neural recordings during reach
We apply TNDM to data gathered from a previously published monkey reaching experiment 77]. The monkey is trained to perform center-out reaching task with eight outer targets. On go cue, the monkey moves a manipulandum along 2D plane to the presented target and receives a liquid reward upon success. Spiking activity from MI and PMd along with the 2D hand position are recorded during each trial. We train each model on single-session data gathered from one of the six trained monkeys. The data consist two paired datasets: PMd activity paired with hand position and MI activity paired with hand position. We show results for the MI recordings in the main text and the results for the PMd recordings in Supplement 2. The neural activity counted in 10ms bins and the behaviour is also measured every 10ms. We align the behaviour to the spikes for both datasets

---

by taking the activity starting during movement onset. We set the length of the neural activity to be the minimum time until the target is reached across all trials. As one of our baselines, PSID cannot model spike count data, we smooth the spike counts with a Gaussian kernel smoother (with standard deviation 50ms) before applying PSID. Out of the 176 trials from the experiment, we use 80% for training (136 trials). We hold out the remaining 34 trials to test the models.
Models/Evaluattos For all models, we perform a sweep over the number latent factors. For TNDM and PSID, we train models with all combinations of 1-5 relevant latent factors and 1-5 irrelevant factors (c.g. 3 relevant and 2 irrelevant). For LFADS, we train models with the number of latent factors ranging from 2-10. As TNDM and LFADS are both implemented as sequential variational autoencoders, we fix the architectures to be same for the two methods (64 units in the generators and encoder). We fix all shared hyperparameters to be the same between the two methods except for the dropout (TNDM requires more regularization due to the behavioural loss). For detailed explanation of the hyperparameters and architectures used in these experiments, see Supplement All results reported here are based on five fits of each model with different random seeds and data splits.
To compare TNDM to LFADS, we first evaluate their neural reconstruction using the test data Poisson log- leeliikei and the root mean square eπror (RMSE) between the predicted and actual fround-truth firing rates. To calculate the ground-truth firing rates, we average the neural data across all trials with the same movement direction and used both the training and test sets to get more robust estimates of the rates from the experimental data. To evaluate the behaviour reconstruction of LFADS, we perform an ex post facto regression from the extracted latent factors to the behaviour in the training set. This regression is linear and is from all time steps of the factors to all time steps of the behaviour Note that this approach for regressing the LFADS factors is more flexible than the decoder in TNDM which is also linear but constrained to be causal. We then compute the coefficient of determination (R2) between the decoded and ground-truth behaviour for each model on the test data.
To compare TNDM PSID, we evaluate the neural reconstruction by computing the RMSE between the predicted and actual ground-truth firing rates. For behaviour reconstruction, we compute the R between the decoded and ground-truth behaviours for each model on the test data As both TNDM and LFADS use information from the whole trial to infer the latent factors (which is inherently acausal), we use a standard Kalman smoother for PSID to make the state estimation comparable.
5 Results
5.1 Simulated Data
For detailed discussion of TNDM's results on synthetic spike trains generated from a Lorenz system, we refer the reader to Appendix 6.
5.2
M1 neural recordings during reach
Fit to behaviour The behaviour reconstruction of TNDM, LFADS and PSID is summarized in Figure 2h. LFADS behavioural reconstruction saturates at around eight factors (R² zs 0.89) and with just four factors yields a respectable behavioural fit (R² z 0.86). This indicates that interpretable in terms of
the LFADS factors, which are constrained only by neural activity, externally measured variables. In comparison, TNDM achieves saturating performance with just three latent factors (R2 R 0.90) where only two are constrained to be relevant for behavioural reconstruction. In fact, all TNDM models with three or more factors (where at least two are constrained to be relevant) have similar behaviour reconstruction accuracy. In comparison TNDM,
"We utilize standard ridge regression from scikit-lcam 18 with default parameters.
' potential concern when comparing TNDM and PSID on behaviour reconstruction is that TNDM has more parameters in its behaviour decoder than PSID does. This is because TNDM decodes the behavioural variable at time step using all past time steps of the latent factors while PSID only uses the current time step f. As shown in Supplement I, however, TNDM achieves equally high behaviour reconstruction using the time lag decoder as does using the proposed casual decoder, therefore, the number parameters TNDM's behaviour decoder is not a confounding factor for this evaluation. Also, while it is possible to remap PSID's 's learned latent factors to the behaviour using a higher parameter regression, this would be equivalent to changing PSID's 's generative model and, therefore, would no longer be valid comparison to PSID.
7

---

LFADS achieves behaviour reconstruction of just (R² ~ 0.68) for three latent factors. TNDM also has much more accurate behaviour reconstruction on the test data than PSID. For three latent factors, where two are constrained to be relevant, PSID achieves a behavioural fit of (R² x 0.82). PSID's behavioural reconstruction saturates at six latent factors where five are constrained to be relevant R² ~ 0.88). Overall, TNDM's behaviour decoding performance implies that the dimensionality of the behaviourally relevant dynamics for this 2D task are lower-dimensional than previously predicted by other latent variable modeling approaches.
Fit neural activity Do the additional constraints and penalties of TNDM affect the accuracy of neural activity reconstruction? This is an important question to answer as the learned latent dynamics are only meaningful if they also explain the observed activity well. Surprisingly, we find that TNDM's and LFAD's Poisson log-likelihoods on the test data are very close (Figure 2b). This indicates that the partitioning of the latent dynamics and the additional constraints imposed by TNDM have very small effect on its neural reconstruction for this dataset. Instead, TNDM and LFADS both show a gradual improvement of neural activity reconstruction as function of the number of factors. The only deviation from this trend is TNDM with one relevant and one irrelevant factor. This is not surprising, however, as much of the neural variability is explained by the behaviour of interest; only allowing one latent factor to explain the neural variabiliity related to behaviour (while simultaneously enforcing disentanglement between the relevant and irrelevant dynamics) will cause TNDM's s neural activity reconstruction to suffer. Perhaps more surprisingly, TNDM achieves lower firing rate RMSE than LFADS with the same number of factors (Figure 2c). While on the surface this result seems counterintuitive. it may be because the RMSE is computed for the average firing rate over all trials of the same movement direction. While TNDM and LFADS have a very similar Poisson likelihood on single trials, TNDM can better distinguish trials by movement direction since it is explicitly modeling behaviour, hence the firing rate prediction split by trial type is improved. PSID provides a worse fit to the neural data than TNDM which is expected given that it is constrained learn linear dynamics.
PSID failure mode Although the neural reconstruction is fairly good for PSID, we find an unex- pected result when analyzing PSID's learned model parameters. Specifically, we find that PSID's state-transition matrix A, which characterizes the behaviourally relevant neural dynamics 25 is approximately the identity matrix for this dataset and is non-informative about the neural activity. We expand upon this analysis of PSID in Supplement where we show that PSID recovers the same frae-transition matrix A when we shuffle the neural data by trial or by time. We provide further evidence that PSID is unable to find informative linear dynamics for this dataset because the behaviour is inherently nonlinear across trials (i.e. multi-directional reaches). Therefore, on this dataset, we conclude that PSID's performance on neural reconstruction is mainly determined by the behaviourally irrelevant factors and its performance on behaviour reconstruction is completely determined by the Kalman gain during decoding.
Interpretation of the learned dynamics In Figure we visualize each stage of the generative process of TNDM (2 relevant and 2 irrelevant factors) and LFADS (4 factors) after training both models on the MI dataset. As can be seen in the figure, there appears to be a clear separation between the relevant and irrelevant initial condition distributions in TNDM that is less apparent in the mixed latent space of LFADS. In fact, the relevant initial conditions of TNDM seem to precisely capture the reach direction of the behaviour. Despite the noticeable differences between the dynamics of LFADS and TNDM, their ability to infer the underlying firing rates from this dataset are nearly identical.
Looking at the learned dynamical factors for TNDM, one can see that the relevant dynamics are more clearly distinguished by reach condition and there is much less variance in the learned trajectories than those of LFADS. At the same time, the relevant TNDM factors do not trivially re-capitulate the behaviour dynamics, indicating that the dual constraint of behaviour and neural activity unmasks a more complicated relationship between the two. This relationship can be analysed by visualizing the learned weights of TNDM's behaviour decoder as shown in Figure 4a (for two relevant factors and two irrelevant factors). In this weight matrix, each time point of the behaviour receives contributions from a broad time interval of preceding factor activity. This corresponds to a temporal integration of the factors and suggests that the relevant factors represent information about movement velocity. Indeed, velocity can be decoded well from the relevant factors using simple ridge regression (both for TNDM and LFADS, Figure g The learned coefficients of this ridge regression for TNDM have a diagonally banded structure that corresponds delayed identity transformation (Figure pp).
8

---

TNDM
Initial cond. rel.
rel factor
rel factor
neural firing rates (RMSE)
0.06
0.08
behaviour
reconstruction
0.06
irrel.
irrel.
T-SNE dim
0.11
0.09
0.12
X
-10
2
500 Time (ms)
T-SNE
Time (ms)
LFADS
factor
factor
neural firing rates (RMSE)
0.06
0.08
0.12
initial conditions
T-SNE 2
0.06
factor
factor
-5
1
-5
0
5
X
T-SNE dim
500
0
500
Time (ms)
Time (ms)
We visualize each component of the generative process for TNDM (top) and LFADS (bottom) after training. On the far left, visualize the inferred initial conditions for each method after reducing the dimension to 2 with T-SNE. As can be seen, TNDM's inferred initial conditions show clear distinction between behaviourally relevant and behaviourally irrelevant information whereas LFADS inferred initial conditions mix this information together. Next, we show the condition- averaged inferred latent dynamical factors (along with the single-trial factors) for each method demonstrate that there is clear distinction between the behaviourally relevant and behaviourally irrelevant factors in TNDM but not LFADS. Finally, we show neural activity reconstruction (numbers are RMSE between data and prediction) and behaviour reconstructions (linear regression for LFADS) for both methods to illustrate that TNDM provides an excellent fit the neural data despite the partitioned latent space and behavioural prediction.
Figure 3:
which is not visible for the LFADS factors (not illustrated). Taken together, these results suggest that MI neural dynamics are related to velocity of the hand in this task. Interestingly, we find that velocity decoding peaks at two relevant factors for TNDM and is less discernible when this number is increased, indicating that the addition of more relevant factors may spread this velocity information across the factors in nonlinear way which cannot be recovered by the ridge regression (not visualized in Figure 4. also illustrates that the TNDM's behaviour prediction saturation point (two relevant factors) has perhaps the most interpretable latent space of all the trained TNDM models.
The irrelevant factors TNDM show task-aligned dynamics that do not depend strongly on the task type, but are rather homogeneous (see Figure 3). For instance, over the course of each trial, irrelevant factor has large initial fluctuation followed by a steady increase its absolute value over time until around 600ms where it tapers off (around when the monkey reaches the target destination). As this factor is agnostic to the reach direction, this may reflect dynamics associated with exccution of these movements more generally.
Discussion
In this work, we introduce TNDM, nonlinear state-space model designed to disentangle
behaviourally relevant and behaviourally irrelevant dynamics underlying neural activity. We evaluated
TNDM on synthetic data and on neural population recordings from PMd and M1 paired with a center-out reach behaviour. We showed that TNDM was able to extract low- low-dimensional latent

---

(a)
(b)
(c)
factor position
factor= velocity
factor 1
factor 1
Velocity
0.5
TNDM
fastor 2
factor 2
LFADS
0.4
behavlour
behavlour
behaviour
behaviour
0.3
Factors
Figure 4: (a) Visualization of the weights of the TNDM behaviour decoder C. that transforms two relevant latent factors into behaviour (hand position). The behaviours are aligned horizontally and the factors vertically. The upper-triangular structure reflects causal decoding, i.e. factors can only influence future behaviour. The model had two relevant and two irrelevant factors. (b) Weights obtained by using ridge regression to predict movement velocity (x and y components) from the relevant factors. A diagonally banded structure can be observed indicating a (delayed) identity transformation. Unlike in (a), this matrix is not constrained to be causal. (c) Decoding accuracy for velocity obtained using ridge regression for LFADS and TNDM with two relevant factors and varying number of irrelevant factors.
dynamics that were much more predictive of behaviour than those of current oftte-of-theaar state- space models without sacrificing its fit the neural activity. This led to the interpretation that the dimensionality of the neural activity associated with the 2D reaching task is potentially lower than previously thought and may be associated with the velocity of the hand.
Although the initial results presented for TNDM are quite promising, the method has a few limitations that should be addressed. First, we find that some hyperparameter settings combined with certain random initialisations can cause biologically implausible oscillations in the learned latent dynamics. While more work needs be done to understand this, it could be related to the weighting between the behavioural and neural likelihoods or to the capacity of the model. A second limitation of TNDM is whether the disentanglement penalty between the relevant and irrelevant dynamics is sufficient. Although the covariance penalty works well in practice (on the presented datasets), disentangling sources of variation using deep generative models is still an open problem [16]. Similarly to PSID, TNDM could be implemented in multiple separate stages which may allow for better disentanglement of the relevant and irrelevant dynamics (88] Third, the linear causal decoder we introduce for behaviour reconstruction is parameter inefficient: the number parameters scales quadratically with time and dynamics/behaviour dimension. Lastly, it can be challenging to determine the 'correct' latent dimensionalities for TNDM. For the datasets in this paper, we found that performing wide sweep over a number of relevant and irrelevant dimensionalities and then choosing the relevant dimensionality where the behaviour prediction saturates is a potential recipe for finding an interpretable latent space.
In future work, we plan to train TNDM with higher-dimension behaviours such as joint angle clectromyography (EMG) recordings. We also plan to extend TNDM non-temporal/ddisrett hehaviours which are of interest behavioural neuroscience (e.g. decision making). Finally, we hope extend TNDM such that it can model dynamics with external inputs from another brain region, i.e. non-autonomous dynamics.
7 Broader Impact
In this work, we develop an expressive and interpretable model of neural population activity. As such, we imagine that TNDM will be useful for answering important questions about neural function (e.g. how the motor cortex gives rise to behaviour). We also imagine that the ability of TNDM to accurately model both the neural activity and the behaviour of interest will of interest for the brain-computer interface community. We believe that TNDM (or the principles behind it) can be used to improve behaviour decoding from neural activity. We also hope that TNDM inspires more research into deep generative models of neural activity that incorporate in external variables of interest. A
10

---

possible negative societal impact of TNDM is that, like all deep neural network models, it requires relatively large amount of compute and has noticeable carbon footprint.

---

Acknowledgements
We thank Alessandro Facchin and Nina Kudryashova for the code contributions and for the insightful discussions. We also thank the reviewers for their thoughtful critiques and suggestions.
References
[1] Abudi Martín, Agarwal Ashish, Barham Paul, Brevdo Eugene, Chen Zhifeng, Cirr Craig, Corrado Greg S., Davis Andy, Dean Jeffrey, Devin Matthieu. Ghemawat Sanjay, Goodfellow Ian, Harp Andrew, Irving Geoffrey, Isard Michael Jia Yangqing, Jozefowicz Rafal, Kaiser Lukasz, Kudlur Manjunath, Levenberg Josh, Mané Dandelion, Monga Rajat, Moore Sherry, Murray Derek, Olah Chris, Schuster Mike, Shlens Jonathon, Steiner Benoit, Sutskever Ilya, Talwar Kunal, Tucker Paul, Vanhoucke Vincent, Vasudevan Vijay, Viégas Fernanda, Vinyals Oriol, Warden Pete, Wattenberg Martin, Wicke Martin, Yu Yuan, Zheng Xiaoqiang. TensorFlow: Large-Scale Machine Learning Heterogeneous Systems. 2015. Software available from tensorflow.org.
[2] Buesing Lars, Macke Jakob H, Sahani Meneesh. Spectral learning of linear dynamics from generalised-linear observations with application to neural population data // Advances in Neural Information Processing Systems 25: 26th Conference on Neural Information Processing Systems (NIPS 2012). 2013. 1691-1699.
[3] Chen Ricky TQ, Li Xuechen, Grosse Roger, Duvenaud David. Isolating sources of disentangle- ment in variational autoencoders arXiv preprint 4XXi:11802.04942 2018.
[4] Churchland Mark M., Cunningham John P., Kaufmman Matthew T., Foster Justin D., Nuyujukian Paul, Ryu Stephen L. Shenoy Krishna V. Neural Population Dynamics during Reaching // Nature.
VII 2012. 487, 7405. 51-56.
[5] Cunningham John P, Byron M Yu. Dimensionality reduction for large-scale neural recordings // Nature Neuroscience. 2014. 17, 11. 1500-1509.
[6] Elsayed Gamaleldin F, Cunningham John P. Structure in neural population recordings: an expected byproduct of simpler phenomena? // Nature Neuroscience. 2017. 20, 9. 1310.
[7] Gallego Juan A, Perich Matthew G, Chowdhury Raeed H, Solla Sara A, Miller Lee Long-term stability of cortical population dynamics underlying consistent behavior // Nature Neuroscience. 2020. 23, 260-270.
[8] Gallego Juan A, Perich Matthew G. Miller Lee E, Solla Sara A. Neural manifolds for the control of movement Neuron. 2017. 94, 978 778-984.
[9] Hall Thomas M, Carvalho Felipe de, Jackson Andrew. A common structure underlies low- frequency cortical dynamics in movement, sleep, and sedation // Neuron. 2014. 83, 5. 1185- 1199.
[10] Hurwitz C., Kudryashova N., Onken A., Hennig M. Building population models for large-scale neural recordings: Opportunities and pitfalls // Current Opinion in Neurobiology. 2021. 70. 64-73.
[11] Jun James J, Steinmetz Nicholas A, Siegle Joshua H, Denman Daniel J, Bauza Marius, Barbarits Brian, Lee Albert K, Anastassiou Costas A, Andrei Alexandru, Aydin Çağatay, others Fully integrated silicon probes for high-density recording of neural activity // Nature. 2017. 551, 7679. 232-236.
[12] Kim Hyunjik, Mnih Andriy. Disentangling by factorising // International Conference on Machine Learning. 2018. 2649-2658.
[13] Kingma Diederik P, Welling Max. Auto-encoding variational bayes 1111 2013.
arXiv preprint
12

---

[14] Kobak Dmitry, Brendel Wieland, Constantinidis Christos, Feierstein Claudia E. Kepecs Adam, Mainen Zachary F, Qi Xie-Lian, Romo Ranulfo, Uchida Naoshige, Machens Christian K. Demixed principal component analysis of neural population data // Elife. 2016. 5. e10989.
[15] Kumar Abhishek, Sattigeri Prasanna, Balakrishman Avinash. Variational inference of dis- entangled latent concepts from unlabeled observations arXiv preprint arXiv:1711.0084. 2017.
[16] Locatello Francesco, Bauer Stefan, Lucic Mario, Raetsch Gunnar, Gelly Sylvain, Schölkopf Bernhard, Bachem Olivier. Challenging common assumptions in the unsupervised learning of disentangled representations // international conference on machine learning. 2019. 4114-4124.
[17] Pandarinath Chethan, O'Shea Daniel J, Collins Jasmine, Jozefowicz Rafal, Stavisky Sergey D, Kao Jonathan C, Trautmann Eric M, Kaufiman Matthew T, Ryu Stephen I, Hochberg Leigh R, others. Inferring single-trial neural population dynamics using sequential auto-encoders I| Nature Methods. 2018. 15, 10. 805-815.
[18] Pedregosa F., Varoquaux G. Gramfon A., Michel V., Thirion B., Grisel O., Blondel M., Pretten-
hofer P., Weiss R., Dubourg V., Vanderplas J., Passos A., Cournapeau D., Brucher M., Perrot M., Duchesnay E. Scikit-learn: Machine Learning in Python Journal of Machine Learning Research. 2011. 12. 2825-2830.
[19] Perich Matthew G, Conti Sara, Badi Marion, Bogaard Andrew, Barra Beatrice, Wurth Sophie, Bloch Jocelyne, Courtine Gregoire, Micera Silvestro, Capogrosso Marco, others Motor cortical dynamics are shaped by multiple distinct subspaces during naturalistic behavior // BioRxiv. 2020.
[20] Pfau David, Pnevmatikakis Eftychios A, Paninski Liam. Robust learning of low-dimensional dynamics from large neural ensembles. // Advances in Neural Information Processing Systems. 2013.
[21] Rabinovich Mikhail I, Huerla Ramón, Varona Pablo, Afraimovich Valentin S. Transient cognitive dynamics, metastability, and decision making // PLoS Computational Biology. 2008 4, 5. e1000072.
[22] Russo Abigail A, Khajeh Ramin, Bittner Sean R, Perkins Sean M, Cunningham John P. Abbott Laurence F. Churchland Mark M. Neural trajectories in the supplementary motor area and motor cortex exhibit distinct geometries, compatible with different classes of computation // Neuron. 2020. 107, 4. 745-758.
[23] Rybkin Oleh, Daniilidis Kostas, Levine Sergey. Simple and effective VAE training with calibrated decoders // International Conference on Machine Learning. 2021. 9179-91999999 9189.
[24] Sadtler Patrick T, Quick Kristin M, Golub Matthew D. Chase Steven M, Ryu Stephen I, Tyler- Kabara Elizabeth C, Byron M Yu, Batista Aaron P. Neural constraints on learning Nature. 2014. 512, 7515. 423-426 4266
[25] Sani Omid G., Abbaspourazad Hamidreza. Wong Yan T., Pesaran Bijan. Shanechi Maryam M. Modeling Behaviorally Relevant Neural Dynamics Enabled by Preferential Subspace Identifica- tion // Nature Neuroscience. XI 2020. 1-10.
[26] She Qi. Wu Angi Neural dynamics discovery via Gaussian process recurrent neural networks // Uncertainty in Artificial Intelligence. 2020. 464.
[27] Smith Anne C, Brown Emery N. Estimating state-space model from point process observations I/ Neural Computation. 2003. 15, 5. 965-9999999 991.
[28] Srivastava Akash, Bansal Yamini, Ding Yukun, Hurwitz Cole, Xu Kai, Egger Bernhard, Sat- tigeri Prasanna, Tenenbaum Josh, Cox David D, Gutfreund Dan. Improving the Reconstruc- tion of Disentangled Representation Learners via Multi-Stage Modelling // arXiv preprint arXiv:2010.11 13 137 2020.
[29] Sussillo David, Jozefowicz Rafal, Abbott LF, Pandarinath Chethan. L.addslattett factor analysis via dynamical systems // arXiv preprint 1xXiv:1608.066 2016.
13

---

[30] Whiteway Matthew R. Biderman Dan, Friedman Yoni. Dipoppa Mario, Buchanan E Kelly, Wu Angi, Zhou John, Noel Jean-Paul R, Cunningham John P. Paninski Liam, others. Partitioning variability in animal behavioral videos using semi-supervised variational autoencoders // bioRxiv. 2021.
[31] Zhou Ding, Wei Xin- Learning identifiable and interpretable latent models of high- dimensional neural activity using vA-V // arXiv preprint arXiv:201 1.04999 2020.
14


Reviews:
[{"review_id": "xugDbKEiwI5", "reviewer": "Reviewer_AfQj", "summary": "This paper details a dimensionality reduction and dynamical modeling technique that partitions the neural variability into behaviorally-relevant and behaviorally-irrelevant factors. This idea was introduced and explored in preferential subspace identification (PSID), a reformulation of SSID (subspace identification), which models the latent dynamics as evolving linearly. This paper extends that to the nonlinear regime, using recurrent neural networks to model the evolution of the latent variables. The authors use the LFADS framework to model the temporal evolution, while partitioning the state into behaviorally-relevant and behaviorally-irrelevant factors, and not allowing for the behaviorally-irrelevant state to influence the behaviorally-relevant state.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors may have severely misrepresented PSID. From Equation (4) of the Methods in Sani et al. (2021), PSID actually seems to allow the behaviorally-irrelevant states as having an effect on the behaviorally relevant states, and not in fact a block diagonal transition matrix A as the authors have mentioned (Question: is this why their implementation of PSID fails?). TNDM does not allow this influence to occur, which is a significant departure from the assumptions in PSID, and in fact, the assumptions for TNDM require a strong motivation. The assumption that the behaviorally-irrelevant dynamics are completely disjoint from the behaviorally-relevant dynamics implies that there are two completely independent processes occurring in the neural activity that require a separate set of dynamics to explain. On the contrary, papers such as Sussillo et al. (2015) and Russo et al. (2018) show that if one trains a network to produce behavioral outputs in the context of motor control, the network states can actually recapitulate neural activity, traditionally thought of as behaviorally-irrelevant, but that are in fact necessary to produce behavior. Moreover, the authors mention that PSID is hindered since it cannot incorporate lags between the neural activity and the behavior. However, the state z is a latent representation that can in fact accumulate past neural activity to lead to current behavior, and thus can incorporate lags between the two representations.\n\nThe authors find that the reconstruction of the firing rates using TNDM is dependent almost purely on the number of total factors that are used, regardless of how many of them are constrained to be behaviorally-relevant (Figure 2C). Please discuss. It would be interesting to include in Figure 2A,C the behavioral and neural predictions for # of behaviorally-irrelevant factors = 0.\n\nPlease explain why TNDM achieves a lower firing rate RMSE than LFADS with the same number of factors (Figure 2C) for # total factors > 4.\n\nPlease show the post-hoc behavior reconstruction using the LFADS factors in Figure 3 using the 4 factors. The behavior R^2 was high, so the reconstructions should be visually comparable to TNDM.\n\nPlease discuss and / or show the post-hoc partition of LFADS factors into behaviorally-relevant and behaviorally-irrelevant by decoding behavior from the factors (as done for PSID in this paper).\n\nThe paper would benefit from a better discussion of the structure and role of behaviorally-relevant vs. behaviorally-irrelevant dynamics.\n\nThe authors refer to the behaviorally-irrelevant factors and networks as simply ‘irrelevant’. I would encourage them to use the full term ‘behaviorally-irrelevant’ instead.\n\nOriginality: The paper is an extension of LFADS using ideas in PSID. \nQuality: The paper could lead to interesting insights into the structure of behavioral vs. neural activity. However, in its current form, it is unclear what it is adding to the realm of different dimensionality reduction models of neural activity.\nClarity: The paper is written very clearly, and the figures are of high quality.\nSignificance: While the paper explores an interesting idea, it would be much improved with proper motivation of their model structure. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "msU2Jv1fZBg", "reviewer": "Reviewer_wP6X", "summary": "In this work, the authors propose a state-space model and inference algorithm that aims to recover the latent nonlinear dynamics from sequential data with particular interests in neuroscience. The proposed model inherits the latent factor analysis via dynamical systems (LFADS) and incorporates the advantage of disentangling behavior relevant/irrelevant dynamics in the latent space as preferential subspace identification (PSID) does. It overcomes the limitation of linear dynamics and more faithfully captures the behavior. The proposed method was tested on the data from a monkey reaching task.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: The proposed model uses the same structure as LFADS with extra conditional Gaussian observation model for behavior variables of the similar idea as that of PSID. The objective function modifies the ELBO with terms that regularize the entanglement and behavior likelihood. Previous work are discussed and compared. The contributions are clearly stated.\n\nQuality: There're more have to be done to support the claims.\n- Only one empirical result on the reaching task was shown. It isn't adequate to evaluate the strength and weakness of the method. Synthetic data and other tasks are needed.\n- The work is focusing on the reaching task overly. TNDM seems to be tightly tailored for the task. These particular settings prohibits the method from generalizing to other task and weaken its easiness for use. E.g. the choice of initial conditions is needed. It is also mentioned that a random initialization might lead to weird dynamics; The causal relationship is not always true for behavior and neural activity, considering feedback scenarios. \n- \"to learn low-dimensional dynamics\" is ambiguous. In TNDM, the dynamical system is modeled by a high-D RNN as well as LFADS. Only the latent factor $z$ are low-D. \n- The power of SSM like LFADS is not only to recover the latent trajectories but also to learn the dynamical system (velocity/phase portrait/simulated trajectories). This should be also evaluated. \n- It is expected that TNDM outperforms LFADS on behavior encode at small number factors, because of imperfect encoding of behavior variables and other internal dynamics. \n- It is also important to evaluate the disentanglement, e.g. the decoding performance of irrelevant dynamics.\n- PSID used Gaussian filtered spikes. This made the data and likelihood are both different from that of TNDM. Since TNDM is not limited to Poisson for spikes. It's more fair to fit TNDM to the same data and likelihood as PSID. \n\nClarity: The submission is well organized. The writing can be improved.\n- The authors spent a lot for the details of LFADS and PSID but not say much about how PSID separates dynamics which is a key point to the work. \n- In Eq(2), the encoder network is only a function of the spike $x$ without behavior $y$? and why not?\n- It is unclear that how the initial conditions $\\Phi_r$ and $\\Phi_i$ are parametrized in Eq(2).\n- Where are $\\mathbf{g}_{0r}$ and $\\mathbf{g}_{0i}$ defined? These bold symbols have different indexing order from their plain counterparts. Are they vectors?\n- How is the mean of the sample cross-correlations between $\\mathbf{g}_{0r}$ and $\\mathbf{g}_{0i}$ computed?\n\nSignificance: The results show  that TNDM can capture, from neural recordings, the latent trajectories that faithfully reconstructs the neural activity and behavior. It could help to separate the latent factors into behavior-related and internal dynamics. The demonstration could be improved though.\n\nMinor comments:\n- Why to trim the trials to minimum? Varying length trials should not be a prolem.\n- Why the covariance of behavior is fixed at the identity matrix?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors discussed the limitations and broader impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10"}, "extra_text": {}}, {"review_id": "ahTYzQa8bI", "reviewer": "Reviewer_6ehA", "summary": "In this paper authors provide a novel architecture for modelling neural dynamics with the aim of separating dynamics that are related to behaviour from the ones that are not. Using amortized inference, observed neural activities are used to set the initial state of two RNNs, one for relevant dynamics and one for irrelevant dynamics, which in turn predict behavioural and neural data. The results show that compared to baseline models the model is indeed able to predict behavioural responses, while maintaining prediction ability for the neural data.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is clearly written and the model, contributions, and the results are well explained. The model itself is also well designed for the purpose and potentially provides important contributions to the literature. \n\nI have two concerns as explained below.\n\nI am rather confused about why neural dynamics at time t are modelled to affect future behaviours at c > t directly (in addition to the behaviour at time t). Normally, I would assume that current neural dynamics at time t influence neural activity at time c > t, which in turn affect behaviour at time c. But the current setup implies that neural activities at time t had to directly affect behaviour at time c without mediation of neural responses at time c. It seems to imply that the RNN network was unable to capture the whole neural dynamics over time, which seems a serious limitation. \n\nFrom the results, it's hard to evaluate whether irrelevant RNN is contributing at all to the prediction of neural activities; i.e., it could be the case that the whole neural dynamics are explained by the relevant RNN; is there any evidence supporting the role of relevan RNN in predicting neural activities? Similarly, how much is the contribution of the irrelevant RNN in predicting neural activities? Is there any evidence that not all the neural dynamics is explained by the irrelevant RNN?\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "They are adequately discussed.", "review_text": "", "extra_scores": {"ethical_concerns": "I suggest presenting the ethics approval related to the used data.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "SmUxzDviWG", "reviewer": "Reviewer_1AfS", "summary": "The authors introduce a new sequential VAE that simultaneously models behavioral dynamics and neural activity for experimental neuroscience data. The model includes partitioned latents which map to either neural activity and behavior, or just behavior.  The model uses RNNs to describe the evolution of each of these latent states (similar to LFADS), and the model is trained by optimizing an objective function which contains separate variational distributions for behaviorally relevant and irrelevant latents, as well as a penalty to prevent entanglement of the latent states. The authors test their model on neural data recorded from Monkey V1 during a reaching task.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think this work is an interesting application and extension of sequential VAEs for neuroscience. The writing is clear, the model is novel, and I believe the work would be of interest to the neuroscience community. My primary concern is that this paper, in its current form, may not be appropriate for NeurIPS.  My reason for this is that there is not a lot of mathematical or theoretical development -- specifically, the objective function is weakly motivated, and there is minimal to no attention given to inference. I elaborate on each of these pieces below.\n\nWeak motivation of an objective function: The introduction of the ELBO in (3) feels like it naturally follows from the previous explanation and description of the model.  However, immediately after introducing this, the authors say that the variational distributions can become entangled and thus introduce the cross-correlation penalty (4). Though they show some qualitative comparisons with and without this penalty in the appendix, the motivation feels ad hoc and quickly thrown in. As this is a fundamentally new objective, I think elaboration is warranted.\n\nI'm curious if the authors could bake this type of regularization into the model itself, as opposed to just modifying the objective function and exploring the qualitative changes. Perhaps one could impose a prior on the latent state(s)? \n\nAdditionally, the authors introduce a parameter balancing the behavioral and neural objectives (4) with minimal justification or motivation.\n\nLack of attention given to inference: The current inference section is mostly devoted to specifying the objective function, not actually how inference is done. I am not familiar with inference procedure in sequential VAEs. Does it use the reparameterization trick? Is there anything specific in this model that can be done to perform inference quickly or efficiently?\n\nI do think that overall this work is interesting and well-done.  However, given that the work is a somewhat natural combination pieces of existing models, additional attention should be devoted to expanding on partitioned seq VAEs in general, some more fleshing out of inference, or better theoretical justification/exploration of the objective function, especially for NeurIPS. \n\nA few other smaller concerns:\n\nThe authors fail to mention statespace models that more have flexible than linear dynamics - namely GPFA. See \n\nZhao, Yuan, and Il Memming Park. \"Variational latent gaussian process for recovering single-trial dynamics from population spike trains.\" Neural computation 29.5 (2017): 1293-1316.\n\nKeeley, Stephen, et al. \"Efficient non-conjugate Gaussian process factor models for spike count data using polynomial approximations.\" International Conference on Machine Learning. PMLR, 2020\n\nThe authors talk about the importance of time-lags in the model, but there are no results supporting this. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "None", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a new sequential VAE that simultaneously models behavioral and neural time series. The reviewers agree that it provides meaningful contribution to neural data analysis. The discussion among the reviewers and the authors was very constructive and productive, and the overall writing of the final version should reflect the nuances raised through the process and results from the additional experiments.
Decision: Accept (Poster)