Probabilistic Decomposed Linear Dynamical Systems for Robust Discovery of Latent Neural Dynamics

We thank Reviewer mAN7 for their thoughtful feedback and suggested improvements to our experiments. We are encouraged that they found our work to be "high-quality" and a "positive contribution to the modeling of neural dynamics".

W1: While p-dLDS does require more complex mathematical concepts, the main goal of this work is to develop a procedure that is easier and more reliable to work with than dLDS. Specifically, p-dLDS actually simplifies the fitting process by reducing the total number of hyperparameters for inference to two: the SBL-DF tradeoff $\xi$ and the offset window size $S$. In contrast, dLDS needs three hyperparameters ($\lambda_0$, $\lambda_1$, and $\lambda_2$) for determining the tradeoff between dynamic reconstruction, coefficient sparsity, and coefficient smoothness. In practice, manually balancing these tradeoffs between each of the Lagrange multipliers can be difficult and time-consuming. p-dLDS eliminates the need for manual tuning of these Lagrange multipliers through probabilistic inference. We estimate $\lambda_0$ with the inferred covariance $\Sigma_x$, $\lambda_1$ with $\gamma_{t,k}$'s from SBL-DF, and $\lambda_2$ with $\sigma^2_{t-1,k}$ also from SBL-DF. In practice, we find that our model is much easier to tune compared to the original dLDS model. As demonstrated in the results presented here, the outcomes are more accurate and robust for p-dLDS than dLDS.

Q1: We agree that our approach for the offset may be limited when the fixed point changes across different timescales. We did consider a cost-based dictionary learning approach for the offset term early on in our work, but found it challenging to prevent the learned dynamics from collapsing to the degenerate scenario described in Lemma 1. Often, our offset dictionary would span the same subspace as the direction of the dynamics, preventing the learning of meaningful structure in the DOs. Moreover, this approach introduces additional complexity with multiple hyperparameters (e.g. for dictionary size, coefficient structure, and more) and requires an expensive iterative solver to infer offset coefficients, further slowing down training. In contrast, our moving average approach is simple, requiring only a single hyperparameter (the window size), and can be efficiently computed in a single pass.

Q2: Yes, the p-dLDS hyperparameter selection procedure is very similar to the dLDS one. We use random search with a budget of 1000 samples to determine the values of $S$ and $\xi$ and fit a separate model for each set of hyperparameters. In our search, we uniformly sample an integer from 2 to the length of the time series $T$. For the real dataset, the optimal offset is $S=76$ which is smaller than the timescale of the p-dLDS coefficient switching (around 150 time points). This suggests that the same DO dynamics may persist even as the fixed point of the system fluctuates throughout the experiment. We appreciate the reviewer's suggestion for clarification and will include these additional details about the hyperparameters in the supplementary materials and the camera-ready version if accepted.

Q3: Great question! While we did not directly explore modeling $\gamma_t$'s independently, previous works [1,2,3] have shown that the introduction of a dynamically informed prior leads to lower error and faster convergence when compared to their independent counterparts in both probabilistic and cost-based inference approaches. In particular, Figures 8 and 10 in [1] demonstrates that SBL-DF (with dynamically informed $\gamma_t$'s) shows an improvement in error and convergence time over the static SBL (with independent $\gamma_t$'s). Additionally, we highlight that our inference procedure is computationally efficient, and significantly reduces the training time required for decomposed models (see additional pdf for runtime experiments).

[1] O’Shaughnessy et al. "Sparse Bayesian learning with dynamic filtering for inference of time-varying sparse signals." 2019.

[2] Charles et al. "Dynamic Filtering of Time-Varying Sparse Signals via $\ell _1$ Minimization." 2016.

[3] Charles et al. "Convergence of basis pursuit de-noising with dynamic filtering." 2014.

Q4: Thank you for your suggestions! We originally did not consider using the continuous latent state to classify because we wanted to show that the learned dynamics contained information about the reach angle. However, following your suggestion, we found that the classification accuracy based on the continuous latent state $x_t$ outperformed the classification accuracy from previous experiments in all models. Similar to our previous classification set up, features are generated from the continuous latent states by averaging $x_t$'s over all points in time. We highlight that classification based on the continuous states improved p-dLDS's top-1 accuracy to 57.60% and top-3 accuracy to 94.87% (see additional pdf).

Minor. Thank you for your careful review! We will update these typos.

We thank Reviewer 3hfN for taking their time to review our work.

W1: We report the standard deviation across independent runs in Tables 3, 4 and 5 in the supplementary materials. We omit these values in the main text due to limited space.

W2: While we do not perform an ablation study for the offset term, we discuss theoretically in section 3.1 why the lack of an offset term in dLDS leads to a representation that does not generalize in systems that contain multiple fixed points. Namely, that any parameter setting in the original dLDS dynamics model reduces to a linear dynamical system (LDS) which is characterized by a single fixed point centered around the origin. The Lorenz experiment (section 4.1) illustrates how this limitation leads to a representation that does not align with the attractor lobes, inappropriately segments the system radially with respect to the origin (Fig. 2D), and leads to decreased performance in our quantitative metrics (Table 1). In contrast, we show that the offset term in p-dLDS alleviates these problems and enables the learning of a representation that aligns with the multiple fixed points of the Lorenz attractor.

Q1: We include additional runtime experiments that sweep across a range of dictionary sizes and latent dimensions (see attached pdf). We observe that in general, p-dLDS is significantly faster than dLDS, and even matches rSLDS and SLDS in certain parameter settings. Our approach adopts the worst-case computational complexity of the SBL inference procedure which is $\mathcal{O}(TK^3)$ in time and $\mathcal{O}(TK^2)$ in memory due to a matrix inversion required to compute the posterior coefficient covariance [1]. However, as noted by [2], the introduction of an dynamically informed hyperprior reduces the number of overall iterations required for convergence and typically reduces the actual runtime by an order of magnitude.

[1] Tipping. "Sparse Bayesian learning and the relevance vector machine." 2001.

[2] O’Shaughnessy et al. "Sparse Bayesian learning with dynamic filtering for inference of time-varying sparse signals." 2019.

Q2: In general, we observe that the posterior mean of the coefficients in p-dLDS will be different from the deterministic analog in dLDS. In Figures 2B, 2D, 4C and 4F, we illustrate these differences by plotting the inferred coefficients from both models side by side. We find that without the probabilistic inference procedure and the time-varying offset term, the inferred coefficients of dLDS can produce switching patterns that may not align with the ground truth (Figures 2B and 2D) and can oscillate unpredictably (Figures 4C and 4F).

Thank you for your careful reading of our submission and are encouraged that you find our work to be "well-written" and "adds value to the recently developed dLDS".

W1. Sweeping dynamics noise. Advantage of p-dLDS in a realistic setting.

We agree that it's important to understand our model's behavior under different noise conditions. In our experiments, we focus on covering a range of neurally plausible sources of noise and demonstrate that p-dLDS improves significantly over existing models on a variety of metrics. Moreover, we emphasize that our clinical neurophysiological experiment showcases exactly “the advantage of p-dLDS over other models in a realistic setting”. We use real data with real noise (e.g. noisy trial structure, unexpected head movements, etc.), and demonstrate that p-dLDS can capture meaningful structure where other models cannot.

As discussed in our introduction, dynamical noise can arise from numerous sources, making it challenging to thoroughly explore all scenarios within the scope of this paper. Without a specific end application in mind, it is difficult to construct a experiment that fully addresses this point.

W2. Selecting number of discrete states for rSLDS via cross validation.

Great point! In our original experiments, we focus on approximately matching the model parameters to control for the model complexity. However, we include an additional experiment where we sweep the number of discrete states for rSLDS (see additional pdf). While more discrete states did improve rSLDS performance, it still underperforms with respect to p-dLDS at a lower parameter count, highlighting the inherent limitations of using a switched formulation to model a continuum of signals.

W3. How are "discrete" states defined for pdLDS? Parsimonious and interpretable segmentations.

As noted in the captions of Figures 2B and 2D, “discrete” states for dLDS models are colored according to the “dominant coefficient” which we define as the dynamic operator (DO) state with the largest coefficient magnitude. We plan on making our definition more clear in the main text in the camera-ready version of the paper, if accepted.

For NASCAR, we respectfully disagree that pdLDS "leads to incorrect segmentations". Due to the pdLDS model formulation, the learned model components are capturing the notion that different track segments can have their movement captured by the same DO (i.e., using the same movement but backward). While this results in fewer segments than the ground truth switched model the data was generated with, it correctly identifies redundancies in the switched segments that may be important to capture in many applications. The ability to identify these redundancies is a key advantage of pdLDS over switched models. While you certainly may want to distinguish these cases in a specific application, this is easily done by looking at the sign of the coefficients. We emphasize here again that pdLDS recovers the track segments despite noise (unlike other methods), leading to a variety of quantitative performance improvements.

In the Lorenz experiment, while rSLDS segmentations may seem simpler and more interpretable, it inappropriately under-segments the true system by assigning the same state to both fast and slow segments, which loses important dynamics information about speed. This behavior can be problematic for signals that vary across a continuum, such as those from the reaching experiment. Moreover, our quantitative assessment (Table 1) shows that p-dLDS's representation leads to improved performance across various metrics, supporting its value despite apparent complexity.

Q1: While we agree that understanding the limitations of the time-varying offset term is important, fully characterizing the model's failure modes without a specific application in mind is challenging. There are numerous ways to study when the model breaks due to the offset terms, and the most relevant ones depend on the specific use case.

Regarding the covariances, lemma 2 allows users to implicitly specify the sum of both covariances $\Sigma_l + \Sigma_b$ through the selection of the smoothing window hyperparameter $S$. Given a particular $S$, our fitting procedure estimates the value of the covariances $\Sigma_l + \Sigma_b$.

Q2: We use PCA to determine the size of the latent space due to its simplicity, computational efficiency, and ability to control for parameter count across different models. In computational neuroscience, PCA is a widely used and well established method that has provided numerous scientific insights [citations available upon request]. In contrast, cross-validation can be costly and may select different latent space sizes for different models, which can confound differences due to parameter count with differences in the latent dynamics model. By selecting the latent dimension across all models using PCA, we control for model complexity and can more directly compare their learned representations.

Q3: We believe this results from the training dynamics introduced by using SGD and Momentum, combined with the challenging non-convex optimization landscape defined by our ELBO objective. Generally speaking, variational inference is only guaranteed to converge to a local optimum.

Q4: Hyperparameters (HPs) include $M$, $N$, and $K$ for model dimensionality; $\xi$ and $S$ for coefficient inference; and any SGD HPs (e.g., learning rate, momentum). In general, HPs can be estimated using either domain knowledge or cross-validation. While p-dLDS has a more complex parameter sweep than rSLDS, it is easier to fit than dLDS since it does not require the manual balancing of the BPDN-DF Lagrange multipliers. This greatly simplifies the fitting process for decomposed models.

I would like to thank the authors for their response to my comments and clarifications. I would like to raise my score from 6 (weak accept) to 7 (accept).

Thank you for taking the time to review our paper thoroughly. We are encouraged that you find our work to be "original", "a pleasure to read", "well-conceived", and of "interest to the computational neuroscience community".

Q1 and Q2: Factorization of the coefficient transition and Definition of $\sigma_t$

We clarify that the transition density $p(c_{t,k} | c_{t-1,k}, \gamma_{t,k} )$ over the coefficients $c_t$ has the following functional form:

In the first line, we propose a density that captures the constraints of sparsity and smoothness for the inferred coefficients $c_{t,k}$. A small variance around 0, $\gamma_{t,k}$, promotes sparsity while a small variance around the previous coefficient, $\sigma^2_{t-1,k}$, promotes smoothness. The second line factorizes this density into two separate terms and the third line shows that each term is proportional to a normal distribution with their respective parameters.

While the idea of combining two shrinkage effects in a single density has been explored in previous works [1,2,3,4], those approaches generally require manual balancing of the two penalties. In contrast, we estimate these parameters using SBL-DF results. Specifically, SBL-DF produces estimates for the sparsity variance $\gamma_{t,k}$ through the estimated hyperparameter posterior $q(\gamma_{t,k} )$ and the smoothness variance $\sigma^2_{t-1,k}$ through the approximate coefficient posterior $q(c_{t,k} ) = N(c_{t,k} ; \hat{c}\_{t,k} ,\sigma_{t,k}^2 )$.

We thank the reviewer for pointing out that our discussion on the coefficient transition densities could be improved and will incorporate this discussion into the camera-ready version if accepted.

[1] Irie, Kaoru. "Bayesian dynamic fused LASSO." 2019.

[2] Casella et al. "Penalized regression, standard errors, and Bayesian lassos." 2010.

[3] Li and Lin. "The Bayesian elastic net." 2010.

[4] Kakikawa et al. "Bayesian fused lasso modeling via horseshoe prior." 2023.

Q3: SGD over covariance matrices.

We clarify that we assume a diagonal structure for the posterior covariances. Estimating general covariance matrices $\Sigma$ is challenging both mathematically and computationally due to the requirements for a valid $\Sigma$ (e.g., symmetry, positive semidefiniteness). Vanilla SGD over individual matrix entries does not inherently respect these constraints and does not guarantee that the solution is a valid covariance matrix. Typically, solving for $\Sigma$ involves a semidefinite program (SDP) requiring specialized methods like ellipsoid or interior point solvers [1]. However, we simplify this problem by considering only when $\Sigma$ is diagonal (i.e., $\Sigma = {\rm diag} (\sigma_1 ,\dots, \sigma_n)$ ). This allows us to simplify general SDP optimization to an unconstrained problem by operating in the log space of the diagonal variance terms. We find that this approach works well in our experiments and is computationally efficient (see additional pdf for runtimes).

[1] Vandenberghe and Boyd. ``Semidefinite programming.'' SIAM review 1996.

Q4: Derivation of the approximate posterior.

Thank you for pointing out that equation 8 (following line 187) could be improved in clarity. It is the variational optimum that results from the model structure assumed in line 157. We derive this result below. First, we acknowledge two typos:

1. The joint distribution of pdLDS (equation 6) is missing emission densities and will be updated to,

2. The formula for the optimal coordinate ascent variational update (equation 7) should be the exponentiated log of the joint distribution [1] and will be update to,

We proceed by substituting equation 1 above into equation 2, dropping out constant terms,

\\mathbb{E}\_{q(c, \\gamma)} \\left[ \\log p(x, y, c, \\gamma | \\theta) \\right] = \\mathbb{E}\_{q(c, \\gamma)} \\left[ \\log p(x_1) + \sum_{t=1}^ T \\log p(y_t | x_t) + \\sum_{t=1} ^ {T-1} \\log p(x_{t+1} | x_{t}, c_{t}) \\right] + {\\rm const.} \tag{3}

Next, our assumed decomposition in line 157 gives us that $p(x_{t+1} | x_{t}, c_t) = p(x_{t+1} = l_{t+1} + b_{t+1})$. The distribution over $x_{t+1}$ is computed as,

where we have defined the convolution operator $\star$ and $\Sigma_x =\Sigma_b + \Sigma_l$. Substituting the above equations 5 and 3 into the above equation 2 gives us the optimal variational update,

where $\hat{F}\_t$ is estimated from samples from $\hat{c}\_{t} \sim q(c,\gamma)$. We note that the additional emissions term in the above equation 1 leads to an additional term in our above equation 8.

Again, we thank the reviewer for bringing this to our attention. We plan on correcting the typos and will include this discussion into the camera-ready version if accepted.

[1] Blei et al. "Variational inference: A review for statisticians." 2017.