Modeling Latent Neural Dynamics with Gaussian Process Switching Linear Dynamical Systems

Abstract

Understanding how the collective activity of neural populations relates to computation and ultimately behavior is a key goal in neuroscience. To this end, statistical methods which describe high-dimensional neural time series in terms of low-dimensional latent dynamics have played a fundamental role in characterizing neural systems. Yet, what constitutes a successful method involves two opposing criteria: (1) methods should be expressive enough to capture complex nonlinear dynamics, and (2) they should maintain a notion of interpretability often only warranted by simpler linear models. In this paper, we develop an approach that balances these two objectives: the Gaussian Process Switching Linear Dynamical System (gpSLDS). Our method builds on previous work modeling the latent state evolution via a stochastic differential equation whose nonlinear dynamics are described by a Gaussian process (GP-SDEs). We propose a novel kernel function which enforces smoothly interpolated locally linear dynamics, and therefore expresses flexible – yet interpretable – dynamics akin to those of recurrent switching linear dynamical systems (rSLDS). Our approach resolves key limitations of the rSLDS such as artifactual oscillations in dynamics near discrete state boundaries, while also providing posterior uncertainty estimates of the dynamics. To fit our models, we leverage a modified learning objective which improves the estimation accuracy of kernel hyperparameters compared to previous GP-SDE fitting approaches. We apply our method to synthetic data and data recorded in two neuroscience experiments and demonstrate favorable performance in comparison to the rSLDS.

1. Introduction

Computations in the brain are thought to be implemented through the dynamical evolution of neural activity. Such computations are typically studied in a controlled experimental setup, where an animal is engaged in a behavioral task with relatively few relevant variables. Consistent with this, empirical neural activity has been reported to exhibit many fewer degrees of freedom than there are neurons in the measured sample during such simple tasks [1]. These observations have driven the use of latent variable models to characterize low-dimensional structure in high-dimensional neural population activity [2, 3]. In this setting, neural activity is often modeled in terms of a low-dimensional latent state that evolves with Markovian dynamics [4–14]. It is thought that the latent state evolution is related to the computation of the system, and therefore, insights into how this evolution is shaped through a dynamical system can help us understand the mechanisms underlying computation [15–21].

In practice, choosing an appropriate modeling approach for a given task requires balancing two key criteria. First, statistical models should be expressive enough to capture potentially complex and nonlinear dynamics required to carry out a particular computation. On the other hand, these models should also be interpretable and allow for straightforward post-hoc analyses of dynamics. One model class that strikes this balance is the recurrent switching linear dynamical system (rSLDS) [8]. The rSLDS approximates arbitrary nonlinear dynamics by switching between a finite number of linear dynamical systems. This leads to a powerful and expressive model which maintains the interpretability of linear systems. Because of their flexibility and interpretability, variants of the rSLDS have been used in many neuroscience applications [19, 22–28] and are examples of a general set of models aiming to understand nonlinear dynamics using compact and interpretable components [29, 30].

However, rSLDS models suffer from several limitations. First, while the rSLDS is a probabilistic model, typical use cases do not capture posterior uncertainty over inferred dynamics. This makes it difficult to judge the extent to which particular features of a fitted model should be relied upon when making inferences about their role in neural computation. Second, the rSLDS often suffers from producing oscillatory dynamics in regions of high uncertainty in the latent space, such as boundaries between linear dynamical regimes. This artifactual behavior can significantly impact the interpretability and predictive performance of the rSLDS. Lastly, the rSLDS does not impose smoothness or continuity assumptions on the dynamics due to its discrete switching formulation. Such assumptions are often natural and useful in the context of modeling realistic neural systems.

In this paper, we improve upon the rSLDS by introducing the Gaussian Process Switching Linear Dynamical System (gpSLDS). Our method extends prior work on the Gaussian process stochastic differential equation (GP-SDE) model, a continuous-time method that places a Gaussian process (GP) prior on latent dynamics. By developing a novel GP kernel function, we enforce locally linear, interpretable structure in dynamics akin to that of the rSLDS. Our framework addresses the aforementioned modeling limitations of the rSLDS and contributes a new class of priors in the GP-SDE model class. Our paper is organized as follows. Section 2 provides background on GP-SDE and rSLDS models. Section 3 presents our new gpSLDS model and an inference and learning algorithm for fitting these models. In Section 4 we apply the gpSLDS to a synthetic dataset and two datasets from real neuroscience experiments to demonstrate its practical use and competitive performance. We review related work in Section 5 and conclude our paper with a discussion in Section 6.

2. Background

2.1. Gaussian process stochastic differential equation models

Gaussian processes (GPs) define nonparametric distributions over functions. They are a popular choice in machine learning due to their ability to capture nonlinearities and encode reasonable prior assumptions such as smoothness and continuity [31]. Here, we review the GP-SDE, a Bayesian generative model that leverages the expressivity of GPs for inferring latent dynamics [10].

Generative model

In a GP-SDE, the evolution of the latent state $x\in {ℝ}^K$ is modeled as a continuous-time SDE which underlies observed neural activity $y\left(t_i\right)\in {ℝ}^D$ at time-points $t_i\in [0,T]$. Mathematically, this is expressed as

$$dx=f\left(x\right)dt+{\text{Σ}}^{\frac{1}{2}}dw,\mathbb{E}\left[y\left(t_i\right)|x\right]=g\left(Cx\left(t_i\right)+d\right).$$ (1)

The drift function $f:{ℝ}^K\to {ℝ}^K$ describes the system dynamics, $\Sigma$ is a noise covariance matrix, and $dw\sim 𝓝(0,dtI)$ is a Wiener process increment. Parameters $C\in {ℝ}^{\mathit{D\times K}}$ and $d\in {ℝ}^D$ define an affine mapping from latent to observed space, which is then passed through a pre-specified inverse link function $g\left(·\right)$.

A GP prior is used to model each output dimension of the dynamics $f\left(·\right)$ independently. More formally, if $f\left(·\right)={[f_1\left(·\right),\mathrm{...},f_K(·)]}^{\top }$, then

$$f_k\left(\cdot \right)\stackrel{\text{iid}}{\sim }𝓖𝓟\left(0,{\kappa }^{\Theta }\left(\cdot ,\cdot \right)\right),\text{for}k=1,\dots ,K,$$ (2)

where ${\kappa }^{\Theta }\left(\cdot ,\cdot \right)$ is the kernel for the GP with hyperparameters $\Theta$.

Interpretability

GP-SDEs and their variants can infer complex nonlinear dynamics with posterior uncertainty estimates in physical systems across a variety of applications [32–34]. However, one limitation of using this method with standard GP kernels, such as the radial basis function (RBF) kernel, is that its expressivity leads to dynamics that are often challenging to interpret. In Duncker et al. [10], this was addressed by conditioning the GP prior of the dynamics $f\left(·\right)$ on fixed points $f(x^{\ast })=0$ and their local Jacobians $J\left(x^{*}\right)=\frac{\partial }{\partial x}f\left(x\right){|}_{x=x^{*}}$, and subsequently learning the fixed-point locations $x^{*}$ and the locally-linearized dynamics $J\left(x^{*}\right)$ as model parameters. This approach allows for direct estimation of key features of $f\left(·\right)$. However, due to its flexibility, it is also prone to finding more fixed points than those included in the prior conditioning, which undermines its overall interpretability.

2.2. Recurrent switching linear dynamical systems

The rSLDS models nonlinear dynamics by switching between different sets of linear dynamics [8]. Accordingly, it retains the simplicity and interpretability of linear dynamical systems while providing much more expressive power. For these reasons, variations of the rSLDS are commonly used to model neural dynamics [19, 22–29].

Generative model

The rSLDS is a discrete-time generative model of the following form:

$$x_t\sim 𝓝\left({\mathit{\text{A}}}_{s_t}{x}_{t-1}+b_{s_t},Q_{s_t}\right),\mathbb{E}\left[y_t|x_t\right]=g\left(C{x}_t+d\right)$$ (3)

where dynamics switch between $J$ distinct linear systems with parameters ${\left\{A_j,b_j,Q_j\right\}}_{j=1}^J$. This is controlled by a discrete state variable $s_t\in \left\{1,\dots ,J\right\}$, which evolves via transition probabilities modeled by a multiclass logistic regression,

$$p\left(s_t|s_{t-1},x_{t-1}\right)\propto \exp \left(w^{\text{T}}{x}_{t-1}+r_{{s_t}_{-1}}\right).$$ (4)

The “recurrent” nature of this model comes from the dependence of eq. (4) on latent space locations $x_t$. As such, the rSLDS can be understood as learning a partition of the latent space into $J$ linear dynamical regimes seprated by linear decision boundaries. This serves as important motivation for the parametrization of the gpSLDS, as we describe later.

Interpretability

While the rSLDS has been successfully used in many applications to model nonlinear dynamical systems, it suffers from a few practical limitations. First, it often produces unnatural artifacts of modeling nonlinear dynamics with discrete switches between linear systems. For example, it may oscillate between discrete modes with different discontinuous dynamics when a trajectory is near a regime boundary. Next, common fitting techniques for rSLDS models with non-conjugate observations typically treat dynamics as learnable hyperparameters rather than as probabilistic quantities [26], which prevents the model from being able to capture posterior uncertainty over the learned dynamics. Inferring a posterior distribution over dynamics is especially important in many neuroscience applications, where scientists often draw conclusions from discovering key features in latent dynamics, such as fixed points or line attractors.

3. Gaussian process switching linear dynamical systems

To address these limitations of the rSLDS, we propose a new class of models called the Gaussian Process Switching Linear Dynamical System (gpSLDS). The gpSLDS combines the modeling advantages of the GP-SDE with the structured flexbility of the rSLDS. We achieve this balance by designing a novel GP kernel function that defines a smooth, locally linear prior on dynamics. While our main focus is on providing an alternative to the rSLDS, the gpSLDS also contributes a new prior which allows for more interpretable learning of dynamics and fixed points than standard priors in the GP-SDE framework (e.g., the RBF kernel). Our implementation of the gpSLDS is available at: https://github.com/lindermanlab/gpslds.

3.1. The smoothly switching linear kernel

The core innovation of our method is a novel GP kernel, which we call the Smoothly Switching Linear (SSL) kernel. The SSL kernel specifies a GP prior over dynamics that maintains the switching linear structure of rSLDS models, while allowing dynamics to smoothly interpolate between different linear regimes.

For every pair of locations $x$, $x\prime \in {ℝ}^K$, the SSL kernel with $J$ linear regimes is defined as,

$${\kappa }_{\text{ssl}}\left(x,x\prime \right)=\sum _{j=1}^J\underset{{\kappa }_{\text{lin}}^{\left(j\right)}\left(x,x\prime \right)}{\underbrace{\left({\left(x-c_j\right)}^{\text{T}}M\left(x\prime -c_j\right)+{\sigma }_0^2\right)}}\underset{{\kappa }_{\text{part}}^{\left(j\right)}\left(x,x\prime \right)}{\underbrace{{\pi }_j\left(x\right){\pi }_j\left(x\prime \right)}}$$ (5)

where $c_j\in {ℝ}^K$, $M\in {ℝ}^{\mathit{K\times K}}$ is a diagonal positive semi-definite matrix, ${\sigma }_0^2\in {ℝ}_{+}$, and ${\pi }_j\left(x\right)\ge 0$ with ${\sum }_{j=1}^J{\pi }_j\left(x\right)=1$. To gain an intuitive understanding of the SSL kernel, we will separately analyze each of the two terms in the summands.

The first term, ${\kappa }_{\text{lin}}^{\left(j\right)}\left(x,x\prime \right)$, is a standard linear kernel which defines a GP distribution over linear functions [31]. The superscript $j$ denotes that it is the linear prior on the dynamics in regime $j$. $\mathit{\text{M}}$ controls the variance of the function’s slope in each input dimension, and $c_j$ is such that the variance of the function achieves its minimum value of ${\sigma }_0^2$ at $x=c_j$. We expand on the relationship between the linear kernel and linear dynamical systems in Appendix A.

The second term is what we define as the partition kernel, ${\kappa }_{\text{part}}^{\left(j\right)}\left(x,x\prime \right)$, which gives the gpSLDS its switching structure. We interpret $\pi (x)={[{\pi }_1(x)\dots {\pi }_J(x)]}^{\top }$ as parametrizing a categorical distribution over $J$ linear regimes akin to the discrete switching variables in the rSLDS. We model $\pi \left(x\right)$ as a multiclass logistic regression with decision boundaries $\{w_j^{\top }\varphi (x)=0|j=1,\dots ,J-1\}$, where $\varphi (x)$ is any feature transformation of $x$. This yields random functions which are locally constant and smoothly interpolate at the decision boundaries. More formally,

$${\pi }_j\left(x\right)=\frac{\exp \left(w_j^{\top }\varphi \left(x\right)/\tau \right)}{1+{\sum }_{j=1}^{J-1}\exp \left(w_j^{\top }\varphi \left(x\right)/\tau \right)},j=1,\dots ,J$$ (6)

where $w_J=0$. The hyperparameter $\tau \in {ℝ}_{+}$ controls the smoothness of the decision boundaries. As $\tau \to 0$, $\pi \left(x\right)$ approaches a one-hot vector which produces piecewise constant functions with sharp boundaries, and as $\tau \to \infty$ the boundaries become more uniform. While we focus on the parametrization in eq. (6) for the experiments in this paper, we note that in principle any classification method can be used, such as another GP or a neural network.

The SSL kernel in eq. (5) naturally combines aspects of the linear and partition kernels via sums and products of kernels, which has an intuitive interpretation [35]. The product kernel ${\kappa }_{\text{lin}}^{\left(j\right)}\left(x,x\prime \right){\kappa }_{\text{part}}^{\left(j\right)}\left(x,x\prime \right)$ enforces linearity in regions where ${\pi }_j\left(x\right)$ is close to 1. Summing over $J$ regimes then enforces linearity in each of the $J$ regimes, leading to a prior on locally linear functions with knots determined by $\pi \left(x\right)$. We note that our kernel is reminiscent of the one in Pfingsten et al. [36], which uses a GP classifier as a prior for $\pi \left(x\right)$ and applies their kernel to a GP regression setting. Here, our work differs in that we explicitly enforce linearity in each regime and draw a novel connection to switching models like the rSLDS.

Figure 1A depicts 1D samples from each kernel. Figure 1B shows how a SSL kernel with $J=2$ linear regimes separated by decision boundary $x_1^2+x_2^2=4$ (top) produces a structured 2D flow field consisting of two linear systems, with $x_1$- and $x_2$- directions determined by 1D function samples (bottom).

Figure 1:

SSL kernel and generative model. A. 1D function samples, plotted in different colors, from GPs with five kernels: two linear kernels with different hyperparameters, partition kernels for each of the two regimes, and the SSL kernel. B. (top) An example $\pi \left(x\right)$ in 2D and (bottom) a sample of dynamics from a SSL kernel in 2D with $\pi \left(x\right)$ as hyperparameters. The $x_1$- and $x_2$- directions of the arrows are given by independent 1D samples of the kernel. C. Schematic of the generative model. Simulated trajectories follow the sampled dynamics. Each trajectory is observed via Poisson process or Gaussian observations.

3.2. The gpSLDS generative model

The full generative model for the gpSLDS incorporates the SSL kernel in eq. (5) into a GP-SDE modeling framework. Instead of placing a GP prior with a standard kernel on the system dynamics as in eq. (2), we simply plug in our new SSL kernel so that

$$f_k\left(\cdot \right)\stackrel{\text{iid}}{\sim }𝓖𝓟\left(0,{\kappa }_{\text{ssl}}^{\Theta }\left(\cdot ,\cdot \right)\right),\text{for}k=1,\dots ,K,$$ (7)

where the kernel hyperparameters are $\Theta =\left\{M,{\sigma }_0^2,{\left\{c_j\right\}}_{j=1}^J,{\left\{w_j\right\}}_{j=1}^{J-1},\tau \right\}$. We then sample latent states and observations according to the GP-SDE via eq. (1). A schematic of the full generative model is depicted in fig. 1C.

Incorporating inputs

In many modern neuroscience applications, we are often interested in how external inputs to the system, such as experimental stimuli, influence latent states. To this end, we also consider an extension of the model in eq. (1) which incorporates additive inputs of the form,

$$dx=\left(f\left(x\right)+Bv\left(t\right)\right)dt+{\text{Σ}}^{\frac{1}{2}}dw,$$ (8)

where $v\left(t\right)\in {ℝ}^I$ is a time-varying, known input signal and $B\in {ℝ}^{\mathit{K\times I}}$ maps inputs linearly to the latent space. The latent path inference and learning approaches presented in the following section can naturally be extended to this setting, with updates for $B$ available in closed form. Further details are provided in Appendix B.3–B.6.

3.3. Latent path inference and parameter learning

For inference and learning in the gpSLDS, we apply and extend a variational expectation-maximization (vEM) algorithm for GP-SDEs from Duncker et al. [10]. In particular, we propose a modification of this algorithm that dramatically improves the learning accuracy of kernel hyperparameters, which are crucial to the interpretability of the gpSLDS. We outline the main ideas of the algorithm here, though full details can be found in Appendix B.

As in Duncker et al. [10], we consider a factorized variational approximation to the posterior,

$$q\left(x,f,u\right)=q\left(x\right)\prod _{k=1}^{K}p\left(f_k|u_k,\Theta \right)q\left(u_k\right),$$ (9)

where we have augmented the model with sparse inducing points to make inference of $f$ tractable [37].

The inducing points are located at ${\left\{z_m\right\}}_{m=1}^M\subset {ℝ}^K$ and take values ${\left(f_k\left(z_1\right),\dots ,f_k\left(z_M\right)\right)}^{\top }=u_k$. Standard vEM maximizes the evidence lower bound (ELBO) to the marginal log-likelihood $\text{log}p\left(y|\Theta \right)$ by alternating between updating the variational posterior $q$ and updating model hyperparameters $\Theta$ [38]. Using the factorization in eq. (9), we will denote the ELBO as $𝓛(q\left(x\right)$, $q\left(u\right)$, $\Theta$).

For inference of $q\left(x\right)$, we follow the approach first proposed by Archambeau et al. [39] and extended by Duncker et al. [10]. Computing the ELBO using this approach requires computing variational expectations of the SSL kernel, which we approximate using Gauss-Hermite quadrature as they are not available in closed form. Full derivations of this step are provided in Appendix B.3. For inference of $q\left(u\right)$, we follow Duncker et al. [10] and choose a Gaussian variational posterior for $q\left(u_k\right)=𝓝(u_k|m_u^{k*},S_u^{k*})$, which admits closed-form updates for the mean $m_u^{k*}$ and covariance $S_u^{k*}$ given $q\left(x\right)$ and $\Theta$. Duncker et al. [10] perform these updates before updating $\Theta$ via gradient ascent on the ELBO in each vEM iteration.

In our setting, this did not work well in practice. The gpSLDS often exhibits strong dependencies between $q\left(u\right)$ and $\Theta$, which makes standard coordinate-ascent steps in vEM prone to getting stuck in local maxima. These dependencies arise due to the highly structured nature of our GP prior; small changes in the decision boundaries ${\left\{w_j\right\}}_{j=1}^{J-1}$ can lead to large (adverse) changes in the prior, which prevents vEM from escaping suboptimal regions of parameter space. To overcome these difficulties, we propose a different approach for learning $\Theta$: instead of fixing $q\left(u\right)$ and performing gradient ascent on the ELBO, we optimize $\Theta$ by maximizing a partially optimized ELBO,

$${\Theta }^{*}=\underset{\Theta }{\arg \max }\left\{\underset{q\left(u\right)}{\max }𝓛\left(q\left(x\right),q\left(u\right),\Theta \right)\right\}.$$ (10)

Due to the conjugacy of the model, the inner maximization can be performed analytically. This approach circumvents local optima that plague coordinate ascent on the standard ELBO. While other similar approaches exploit model conjugacy for faster vEM convergence in sparse variational GPs [37, 40] and GP-SDEs [41], our approach is the first to our knowledge that leverages this structure specifically for learning the latent dynamics of a GP-SDE model. We empirically demonstrate the superior performance of our learning algorithm in Appendix C.

3.4. Recovering predicted dynamics

It is straightforward to obtain the approximate posterior distribution over $f^{\ast }:=f(x^{\ast })$ evaluated at any new location $x^{\ast }$. Under the assumption that $f^{\ast }$ only depends on the inducing points, we can use the approximation $q\left(f^{*}\right)={\prod }_{k=1}^K\int p\left(f_k^{*}|u_k,\Theta \right)q\left(u_k\right)\text{d}{u}_k$ which can be computed in closed-form using properties of conditional Gaussian distributions. For a batch of points ${\left\{x_i^{*}\right\}}_{i=1}^N$, this can be computed in $O\left(N{M}^2\right)$ time. The full derivation can be found in Appendix B.4.

This property highlights an appealing feature of the gpSLDS over the rSLDS. The gpSLDS infers a posterior distribution over dynamics at every point in latent space, even in regions of high uncertainty. Meanwhile, as we shall see later, the rSLDS expresses uncertainty by randomly oscillating between different sets of most-likely linear dynamics, which is much harder to interpret.

4. Results

4.1. Synthetic data

We begin by applying the gpSLDS to a synthetic dataset consisting of two linear rotational systems, one clockwise and one-counterclockwise, which combine smoothly at $x_1=0$ (fig. 2A). We simulate 30 trials of latent states from an SDE as in eq. (1) and then generate Poisson process observations given these latent states for $D=50$ output dimensions (i.e. neurons) over T = 2.5 seconds (Fig. 2B). To initialize $\mathit{\text{C}}$ and $\mathit{\text{d}}$, we fit a Poisson LDS [4] to data binned at 20ms with identity dynamics. For the rSLDS, we also bin the data at 20ms. We then fit the gpSLDS and rSLDS models with $J=2$ linear states using 5 different random initializations for 100 vEM iterations, and choose the fits with the highest ELBOs in each model class.

Figure 2:

Synthetic data results. A. True dynamics and latent states used to generate the dataset. Dynamics are clockwise and counterclockwise linear systems separated by x₁ = 0. Two latent trajectories are shown on top of a kernel density estimate of the latent states visited by all 30 trials. B. Poisson process observations from an example trial. C. True vs. inferred latent states for the gpSLDS and rSLDS, with 95% posterior credible intervals. D. Inferred dynamics (pink/green) and two inferred latent trajectories (gray) corresponding to those in Panel A from a gpSLDS fit with 2 linear regimes. The model finds high-probability fixed points (purple) overlapping with true fixed points (stars). E. Analogous plot to D for the GP-SDE model with RBF kernel. Note that this model does not provide a partition of the dynamics. F. rSLDS inferred latents, dynamics, and fixed points (pink/green dots). G. (top) Sampled latents and corresponding dynamics from the gpSLDS, with 95% posterior credible intervals. (bottom) Same, but for the rSLDS. The pink/green trace represents the most likely dynamics at the sampled latents, colored by discrete switching variable. H. MSE between true and inferred latents and dynamics for gpSLDS, GP-SDE with RBF kernel, and rSLDS while varying the number of trials in the dataset. Error bars are ±2SE over 5 random initializations.

We find that the gpSLDS accurately recovers the true latent trajectories (fig. 2C) as well as the true rotation dynamics and the decision boundary between them (fig. 2D). We determine this decision boundary by thresholding the learned $\pi \left(x\right)$ at 0.5. In addition, we can obtain estimates of fixed point locations by computing the posterior probability ${\prod }_{k=1}^K{\mathbb{P}}_{q\left(f\right)}\left(\left|f_k\left(x\right)\right|<ϵ\right)$ for a small $ϵ>0$; the locations $x$ with high probability are shaded in purple. This reveals that the gpSLDS finds high-probability fixed points which overlap significantly with the true fixed points, denoted by stars. In comparison, both the rSLDS and the GP-SDE with RBF kernel do not learn the correct decision boundary nor the fixed points as accurately (fig. 2E–F). Of particular note, the RBF kernel incorrectly extrapolates and finds a superfluous fixed point outside the region traversed by the true latent states.

Figure 2G illustrates the differences in how the gpSLDS and the rSLDS express uncertainty over dynamics. To the left of the dashed line, we sample latent states starting from $x_0=\left(7,0\right)$ and plot the corresponding true dynamics. To the right, we simulate latent states $x_{\text{samp}}$ from the fitted model and plot the true dynamics (in gray) and the learned most likely dynamics (in color) at $x_{\text{samp}}$. For a well-fitting model, we would expect the true and learned dynamics to overlap. We see that the gpSLDS produces smooth simulated dynamics that match the true dynamics at $x_{\text{samp}}$ (top). By contrast, the rSLDS expresses uncertainty by oscillating between the two linear dynamical systems, hence producing uninterpretable dynamics at $x_{\text{samp}}$ (bottom). This region of uncertainty overlaps with the $x_1=0$ boundary, suggesting that the rSLDS fails to capture the smoothly interpolating dynamics present in the true system.

Next, we perform quantitative comparisons between the three competing methods (fig. 2H). We find that both continuous-time methods consistently outperform the rSLDS on both metrics, suggesting that these methods are likely more suitable for modeling Poisson process data. Moreover, the gpSLDS better recovers dynamics compared to the RBF kernel, illustrating that the correct inductive bias can yield performance gains over a more flexible prior, especially in a data-limited setting.

Lastly, we note that the gpSLDS can achieve more expressive power than the rSLDS by learning nonlinear decision boundaries between linear regimes, for instance by incorporating nonlinear features into $\varphi (x)$ in eq. (6). We demonstrate this feature for a 2D limit cycle in Appendix D.

4.2. Application to hypothalamic neural population recordings during aggression

In this section, we revisit the analyses of Nair et al. [27], which applied dynamical systems models to neural recordings during aggressive behavior in mice. To do this, we reanalyze a dataset which consists of calcium imaging of ventromedial hypothalamus neurons from a mouse interacting with two consecutive intruders. The recording was collected from 104 neurons at 15 Hz over ∼343 seconds (i.e. 5140 time bins). Nair et al. [27] found that an rSLDS fit to this data learns dynamics that form an approximate line attractor corresponding to an aggressive internal state (fig. 3A). Here, we supplement this analysis by using the gpSLDS to directly assess model confidence about this finding.

Figure 3:

Results on hypothalamic data from Nair et al. [27]. In each of the panels A-C, flow field arrow widths are scaled by the magnitude of dynamics for clarity of visualization. A. rSLDS inferred latents and most likely dynamics. The presumed location of the line attractor from [27] is marked with a red box. B. gpSLDS inferred latents and most likely dynamics in latent space. Background is colored by posterior standard deviation of dynamics averaged across latent dimensions, which adjusts relative to the presence of data in the latent space. C. Posterior probability of slow points in gpSLDS, which validates line-attractor like dynamics, as marked by a red box. D. Comparison of in-sample forward simulation R² between gpSLDS, rSLDS, and GP-SDE with RBF kernel. To compute this, we choose initial conditions uniformly spaced 100 time bins apart in both trials, simulate latent states k steps forward according to learned dynamics (with k ranging from 100–1500), and evaluate the R² between predicted and true observations as in Nassar et al. [23]. Error bars are ±2SE over 5 different initializations.

For our experiments, we z-score and then split the data into two trials, one for each distinct intruder interacting with the mouse. Following Nair et al. [27], we choose $J=4$ linear regimes to compare the gpSLDS and rSLDS. We choose $K=2$ latent dimensions to aid the visualization of the resulting model fits; we find that even in such low dimensions, both models still recover line attractor-like dynamics. We model the calcium imaging traces as Gaussian emissions on an evenly spaced time grid and initialize $C$ and $d$ using principal component analysis. We fit models with 5 different initializations for 50 vEM iterations and display the runs with highest forward simulation accuracy (as described in the caption of fig. 3).

In fig. 3A–B, we find that both methods infer similar latent trajectories and find plausible flow fields that are parsed in terms of simpler linear components. We further demonstrate the ability of the gpSLDS to more precisely identify the line attractor from Nair et al. [27]. To do this, we use the learned $q(f)$ to compute the posterior probability of slow dynamics on a dense (80×80) grid of points in the latent space using the procedure in Section 3.4. The gpSLDS finds a high-probability region of slow points corresponding to the approximate line attractor found in Nair et al. [27] (fig. 3C). This demonstrates a key advantage of the gpSLDS over the rSLDS: by modeling dynamics probabilistically using a structured GP prior, we can validate the finding of a line attractor with further statistical rigor. Finally, we compare the gpSLDS, rSLDS, and GP-SDE with RBF kernel using an in-sample forward simulation metric (fig. 3D). All three methods retain similar predictive power 500 steps into the future. After that, the gpSLDS performs slightly worse than the RBF kernel; however, it gains interpretability by imposing piecewise linear structure while still outperforming the rSLDS.

4.3. Application to lateral intraparietal neural recodings during decision making

In neuroscience, there is considerable interest in understanding how neural dynamics during decision making tasks support the process of making a choice [26, 34, 42–47]. In this section, we use the gpSLDS to infer latent dynamics from spiking activity in the lateral intraparietal (LIP) area of monkeys reporting decisions about the direction of motion of a random moving dots stimulus with varying degrees of motion coherence [42, 43]. The animal indicated its choice of net motion direction (either left or right) via a saccade. On some trials, a 100ms pulse of weak motion, randomly oriented to the left or right, was also presented to the animal. Here, we model the dynamics of 58 neurons recorded across 50 trials consisting of net motion coherence strengths in {−.512, −.128, 0.0, .128, .512}, where the sign corresponds to the net movement direction of the stimulus. We only consider data 200ms after motion onset, corresponding to the start of decision-related activity.

To capture potential input-driven effects, we fit a version of the gpSLDS described in eq. (8) with K = 2 latent dimensions and J = 2 linear regimes over 50 vEM iterations. We encoded the input signal as ±1 with sign corresponding to pulse direction. In Figure 4A, we find that not only does the gpSLDS capture a distinct visual separation between trials by motion coherence, but it also learns a precise separating decision boundary between the two linear regimes. Our finding is consistent with previous work on a related task, which found that average LIP responses can be represented by a 2D curved manifold [48], though here we take a dynamical systems perspective. Additionally, our model learns an input-driven effect which appears to define a separating axis. To verify this, we project the inferred latent states onto the 1D subspace spanned by the input effect vector. Figure 4B shows that the latent states separate by coherence (top) and by choice (bottom), further suggesting that the pulse input relates to meaningful variation in evidence accumulation for this task. Fig. 4C shows an example latent trajectory aligned with pulse input; during the pulse there is a noticeable change in the latent trajectory. Lastly, in Fig. 4D we find that the gpSLDS expresses high confidence in learned dynamics where latent trajectories are present and low confidence in areas further from this region.

Figure 4:

Results on LIP spiking data from a decision-making task in Stine et al. [42]. A. gpSLDS inferred latents colored by coherence, inferred dynamics with background colored by most likely linear regime, and the learned input-driven direction depicted by an orange arrow. B. Projection of latents onto the 1D input-driven axis from Panel A, colored by coherence (top) and choice (bottom). C. Inferred latents with 95% credible intervals and corresponding 100ms pulse input for an example trial. D. Posterior variance of dynamics produced by the gpSLDS.

5. Related work

There are several related approaches to learning nonlinear latent dynamics in discrete or continuous time. Gaussian process state-space models (GP-SSMs) [49–53] can be considered a discrete-time analogue to GP-SDEs. In a GP-SSM, observations are assumed to be regularly sampled and latent states evolve according to a discretized dynamical system with a GP prior. Wang et al. [49] and Turner et al. [50] learned the dynamics of a GP-SSM using maximum a posteriori estimation. Frigola et al. [51] and Eleftheriadis et al. [52] employed a variational approximation with sparse inducing points to infer the latent states and dynamics in a fully Bayesian fashion. In our work, we use a continuous-time framework that more naturally handles irregularly sampled data, such as point-process observations commonly encountered in neural spiking data. Neural ODEs and SDEs [54, 55] use deep neural networks to parametrize the dynamics of a continuous-time system, and have emerged as prominent tools for analyzing large datasets, including those in neuroscience [33, 56–58]. While these methods can represent flexible function classes, they are likely to overfit to low-data regimes and may be difficult to interpret. In addition, unlike the gpSLDS, neural ODEs and SDEs do not typically quantify uncertainty of the learned dynamics.

In the context of dynamical mixture models, Köhs et al. [59, 60] proposed a continuous-time switching model in a GP-SDE framework. This model assumes a latent Markov jump process over time which controls the system dynamics by switching between different SDEs. The switching process models dependence on time, but not location in latent space. In contrast, the gpSLDS does not explicitly represent a latent switching process and rather models switching probabilities as part of the kernel function. The dependence of the kernel on the location in latent space allows for the gpSLDS to partition the latent space into different linear regimes.

While our work has followed the inference approach of Archambeau et al. [61] and Duncker et al. [10], other methods for latent path inference in nonlinear SDEs have been proposed [32, 41, 62, 63]. Verma et al. [41] parameterized the posterior SDE path using an exponential family-based description. The resulting inference algorithm showed improved convergence of the E-step compared to Archambeau et al. [39]. Course and Nair [32, 62] proposed an amortization strategy that allows the variational update of the latent state to be parallelized over sequence length. In principle, any of these approaches could be applied to inference in the gpSLDS and would be an interesting direction for future work.

6. Discussion

In this paper, we introduced the gpSLDS to infer low-dimensional latent dynamical systems from high-dimensional, noisy observations. By developing a novel kernel for GP-SDE models that defines distributions over smooth locally-linear functions, we were able to relate GP-SDEs to rSLDS models and address key limitations of the rSLDS. Using both simulated and real neural datasets, we demonstrated that the gpSLDS can accurately infer true generative parameters and performs favorably in comparison to rSLDS models and GP-SDEs with other kernel choices. Moreover, our real data examples illustrate the variety of potential practical uses of this method. On calcium imaging traces recorded during aggressive behavior, the gpSLDS reveals dynamics consistent with the hypothesis of a line attractor put forward based on previous rSLDS analyses [27]. On a decision making dataset, the gpSLDS finds latent trajectories and dynamics that clearly separate by motion coherence and choice, providing a dynamical systems view consistent with prior studies [42, 48].

In our experiments, we demonstrated the ability of the gpSLDS to recover ground-truth dynamical systems and key dynamical features using fixed settings of hyperparameters: the latent dimensionality K and the number of regimes J. For simulated data, we set hyperparameters to their true values; for real data, we chose hyperparameters based on prior studies and did not further optimize these values. However, for most real neural datasets, we do not know the true underlying dimensionality or optimal number of regimes. To tune these hyperparameters, we can resort to standard techniques for model comparison in the neural latent variable modeling literature. Two common evaluation metrics are forward simulation accuracy [23, 27] and co-smoothing performance [6, 64, 65].

While these results are promising, we acknowledge a few limitations of the gpSLDS. First, the memory cost scales exponentially with the size of the latent dimension due to using quadrature methods to approximate expectations of the SSL kernel, which are not available in closed form. This computational limitation renders it difficult to fit the gpSLDS with many (i.e. greater than 3) latent dimensions. One potential direction for future work would be to instead use Monte Carlo methods to approximate kernel expectations for models with larger latent dimensionality. In addition, while both the gpSLDS and rSLDS require choosing a discretization timestep for solving dynamical systems, in practice we find that the gpSLDS requires smaller steps for stable model inference. This allows the gpSLDS to more accurately approximate dynamics with continuous-time likelihoods, at the cost of allocating more time bins during inference. Finally, we acknowledge that traditional variational inference approaches – such as those employed by the gpSLDS– tend to underestimate posterior variance due to the KL-divergence-based objective [38]. Carefully assessing biases introduced by our variational approximation to the posterior would be an important topic for future work.

Overall, the gpSLDS provides a general modeling approach for discovering latent dynamics of noisy measurements in an intepretable and fully probabilistic manner. We expect that our model will be a useful addition to the rSLDS and related methods on future analyses of neural data.

Funding Statement

This work was supported by grants from the NIH BRAIN Initiative (U19NS113201, R01NS131987, & RF1MH133778) and the NSF/NIH CRCNS Program (R01NS130789), as well as fellowships from the Simons Collaboration on the Global Brain, the Wu Tsai Neurosciences Institute, the Alfred P. Sloan Foundation, and the McKnight Foundation. We thank Barbara Engelhardt, Julia Palacios, and the members of the Linderman Lab for helpful feedback throughout this project. The authors have no competing interests to declare.