Targeted Neural Dynamical Modeling

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 . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 1 Introduction . Recent progress in high-density , microelectrode array technology now allows for recording from hundreds to thousands of neurons with the precision of single spikes [ 11 ] . Despite the apparent high dimensionality of these datasets , neural activity is often surprisingly well-explained by lowdimensional latent dynamics [ 4 , 24 , 6 , 8 ] . Extracting these dynamics from single trials is crucial for understanding how neural activity relates to a behavioural task or stimulus [ 17 ] . Latent variable models ( LVMs ) are a natural choice for capturing low-dimensional structure from neural activity as 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 as principal components analysis ( PCA ) [ 5 ] to complex state-space models such as LFADS [ 17 ] . 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 , 10 ] . 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 , 19 , 31 ] . These methods , 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 [ 25 ] . When applied to neural activity recorded in the premotor cortex ( PMd ) and primary motor cortex ( M1 ) of a monkey during a 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 can not capture the nonlinear dynamics which are thought to underlie phenomena such as rhythmic motor patterns [ 22 , 9 ] or decision making [ 21 ] . 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 . In this work , we introduce Targeted Neural Dynamical Modeling ( TNDM ) , a 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 at 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 ) [ 13 , 29 ] 1 . 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 a monkey during a 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 x ∈ NN×T be the observed spike counts and let y ∈ RB×T be the observed behaviour during a single-trial.2 We define the unobserved latent factors in a single trial as z ∈ RM×T where M < N . For TNDM , as with PSID , it is important to distinguish between behaviourally relevant zr and behaviourally irrelevant zi latent factors . The behaviourally relevant latent factors 1The code for running and evaluating TNDM on real data can be found at https : //github.com/HennigLab/tndm_paper . We also provide a Tensorflow2 re-implemention of TNDM at https : //github.com/HennigLab/tndm . It is important to 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 . 2In 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 to discrete/non-temporal behaviours ( e.g . reach direction ) . zr summarize the variability in the neural activity associated with the observed behaviour while the behaviourally irrelevant latent factors zi explain everything else in the neural data ( Figure 1a ) . We assume that each of the unobserved , single-trial factors can be partitioned into these relevant and irrelevant factors z : = { zr , zi } . State-space models for neural data . There are 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 , 20 ] 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 : LFADS and PSID . LFADS , or latent factor analysis via dynamical systems , is a state-of-the-art nonlinear state-space model for neural data . In LFADS , the latent dynamics are generated by sampling high-dimensional initial conditions g0 from some distribution pg0 and then evolving g0 with a deterministic RNN fθ . A linear mapping Wz is then applied to the high-dimensional dynamics gt to transform them into the low-dimensional ‘ dynamical ’ factors z . These dynamical factors are transformed into spike counts by mapping each time point to a rate parameter r of a Poisson distribution using a weight matrix Wr followed by an exponential nonlinearity . The generative process is defined as : g0 ∼ pg0 , gt = fθ ( gt−1 ) , zt = Wz ( gt ) , rt = exp ( Wr ( zt ) ) , xt ∼ Poisson ( xt|rt ) . The initial conditions g0 are inferred from x with an RNN encoder network qφ . 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 . PSID is a linear dynamical model that partitions the latent dynamics into behaviourally relevant and behaviourally irrelevant dynamics z : = { zi , zr } . The dynamical evolution of z is defined by a transition matrix A along with a Gaussian noise term wz . z is transformed into the observed firing rates x by mapping each time point to the mean of a Gaussian distribution with a weight matrix Wx and noise term wx . The behaviourally relevant dynamics zr at each time point are transformed into the observed behaviour y with a weight matrix Wy . The state space model for PSID is then defined as : zt ∼ N ( zt|A ( zt−1 ) , wz ) , xt ∼ N ( x|Wx ( zt ) , wx ) , yt = Wy ( zr ) . PSID uses a novel two-stage 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 a good fit to nonlinear activity patterns or behaviours ( e.g . multi-directional reaches ) . Second , the relevant dynamics at each time step zrt must be mapped one-to-one to the behaviour yt during training ( i.e . no time lag ) . This imposes a 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 a 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 1a . We assume that the observed neural activity x and behaviour y in each trial are generated by two sets of latent factors zi and zr . The generative process of TNDM is defined below : gi0 ∼ pgi0 , gr0 ∼ pgr0 , git = fθi ( git−1 ) , grt = fθr ( grt−1 ) zit = Wiz ( git ) , zrt = Wrz ( grt ) , rt = exp ( Wr ( zit , zrt ) ) ) xt ∼ Poisson ( xt|rt ) , y ∼ N ( y|Cy ( zr ) , I ) ( 1 ) In the above equation , pgi0 and pgr0 are the distributions over the initial conditions of the behaviourally irrelevant and behaviourally relevant dynamics , respectively ( assumed to be Gaussian ) . Similarly to LFADS , we parameterize the nonlinear transition functions fθi and fθr using RNNs 3 . As git and grt can have arbitrarily high-dimensionality in our model ( defined by the number of units in the RNN ) , we utilize two weight matrices , Wiz and Wrz , to project these high-dimensional dynamics into a low-dimensional subspace , giving rise to the relevant and irrelevant dynamical factors zit and zrt 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 a 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 ] . It is important to 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 at each time step to the behaviour at each time step , i.e yt = Wy ( zrt ) . This parameterization does not allow for modeling any latency , long-term dependencies or 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 to replicate the behaviour of interest . These results are summarized in Supplement 1 . 3To implement TNDM , we primarily adapt the original Tensorflow [ 1 ] implementation of LFADS from https : //github.com/lfads/models ( 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 Cy with dimensionality nzrT ×BT where nzr 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 zr into behaviours using Cy , we concatenate each dimension of zr in time to form a 1D vector Zr with length nzrT and then perform the operation Y = CyZr where Y is the resulting concatenated behaviour . As Y is a 1D vector with length BT , we can reshape Y to recover the reconstructed behaviour ŷ . Importantly , we do not allow acausal connections in Cy , i.e . the lower triangular components of each of the dynamics to behaviour blocks are set to zero during training . For an example weight matrixCy , see Figure 4a . 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 to the neural reconstruction than when using the no time lag decoder . Inference To extract the latent dynamics zr and zi from the neural activity x , we first learn to approximate the posterior over the initial conditions of the dynamics gr0 and gi0 . Then , we learn a deterministic mapping from the initial conditions to the latent dynamics zr and zi . To approximate the true posterior p ( gr0 , gi0|x , 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 a shared encoder network eφ0 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 : qΦr ( gr0|x ) qΦi ( gi0|x ) = N ( µφr1 ( eφ0 ( x ) ) , σ 2 φr2 ( eφ0 ( x ) ) ) · N ( µφi1 ( eφ0 ( x ) ) , σ 2 φi2 ( eφ0 ( x ) ) ) ( 2 ) The inference networks for the behaviourally relevant and the behaviourally irrelevant initial conditions are parameterized by Φr = { φ0 , φr1 , φr2 } and Φi = { φ0 , φi1 , φi2 } where φ0 are the parameters of the shared RNN . It 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 are evolved using separate decoder RNNs to produce the behaviourally relevant grt and irrelevant high-dimensional dynamics git . The high-dimensional dynamics at each time-step are projected into a low-dimensional subspace to recover the low-dimensional dynamical factors zrt and zit . The neural activity x and behaviour y are generated from the latent factors z as shown in Equation 1 . The ELBO for the observed data from a single trial is therefore defined as : ELBO ( x , y ) = −KL [ qΦr ‖ pgr0 ] − KL [ qΦi ‖ pgi0 ] + EqΦrqΦi [ log pθ1 ( x|gi , gr ) pθ2 ( y|gr ) ] ( 3 ) where pθ1 and pθ2 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 ( g0i , g0r|x , y ) can not be factorized ; that is , zi and zr ( which are deterministic transforms of g0i and g0r ) are conditionally dependent given the observed data . This means that zi and zr will be statistically dependent . To reduce sharing of information and redundancy between these two sets of dynamics , we introduce a novel disentanglement penalty on the two initial condition distributions . For this penalty , we take inspiration from the domain of unsupervised disentangled representation learning where it is standard to introduce additional penalties that encourage disentanglement or independence of the latent representations [ 15 , 3 , 12 ] . 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 g0r and g0i . Importantly , this cross-correlation penalty is applied in such a way that the final objective is still a valid lower bound of the log-likelihood ( the cross-correlation penalty is always negative ) . We refer to this penalty as Q ( qΦr , qΦi ) and we adjust its weight with a hyperparameter λQ ( see Supplement 4 for an ablation study of this penalty ) . The final objective function for TNDM is then : J ( x , y ) =− KL [ qΦr ‖ pgr0 ] − KL [ qΦi ‖ pgi0 ] + EqΦrqΦi [ log pθ1 ( x|gi , gr ) ] + λbEqΦr [ log pθ2 ( y|gr ) ] + λQQ ( qΦr , qΦi ) ( 4 ) where λb is an additional hyperparameter introduced to 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 is 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 .

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.