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

Summary

This paper develops a novel model of latent neural dynamics that builds on existing work in switching and time-varying linear state space models. In particular, the authors propose a full probabilistic formulation of decomposed LDS models and provide a variational EM algorithm to do inference in these models. They apply their method to various synthetic and real neural datasets, showing improved performance over the primary time-varying linear alternative methods.

Strengths

Overall, the proposed model and method are original and are a positive contribution to the modeling of neural dynamics. The authors correctly point out shortcomings of both the SLDS and dLDS approaches, and develop a new model that addresses shortcomings. In particular, the probabilistic formulation of dLDS dramatically improves the models performance, and the flexibility given by the time-varying linear combination of DOs helps p-dLDS perform well in situations that the SLDS is not well suited for. The evaluation is high-quality, as the paper compares relevant methods in four different tasks. Finally, the approach is well-motivated, and more often than not the methods and model are clearly described.

Weaknesses

The model and fitting method have two added complexities that may make it harder than needed to reliably fit the model parameters and make inferences about the data. In particular, the update step sparsity coefficients requires a multi-step approximate procedure and the time-varying offset term depends on a hyperparameter that can take on a wide number of values. These points are discussed in more detail in the questions.

Questions

- The approach for determining `b_t` via a moving average accomplishes the goals of learning a changing fixed point. However, it appears to be suboptimal for cases where the fixed point can change both rapidly or slowly (i.e. across different timescales) because of the fixed window size. Additionally, the requirement to search over a wide number of window sizes pose schallenges for optimization. I'd encourage the authors to consider other approaches. Could the bias term also be determined via a sum over a dictionary with shared coefficients `c`? - Relatedly, can the authors provide more details of the hyperparameter selection scheme for the window size? The paper states that for the Lorenz experiment, the window size was sampled uniformly from 2 to the length of the time series. Does the procedure follow the dLDS hyperparameter selection, in which 1000 random choices of `S` and `\xi` were generated and a separate model was fit given each setting of hyperparameters? What was the optimal offset values on the real datasets, and does this relate to the timescales of the p-dLDS coefficient switching? - The parameterization and update step for the sparsity coefficients appears to work well on the simulated and real datasets. However, it does require a multistep approximate procedure for the update. I'm wondering if the authors have considered simplifications to the model that may allow for using simpler inference in this step while still achieving the goals of learning sparse coefficients? For example, does it work to model each $\gamma_t$ independently? This appears to be very similar to the original SBL-DF algorithm and the first step of the sparsity coefficient update which initializes $q(c) q(\gamma)$ using SBL-DF. How much improved performance does the proposed model & methods have relative to this baseline? - In the reaching experiment, the reach angles are classified using the discrete states or p-dLDS coefficients and not the continuous states. However, I would expect classification using the continuous states to outperform this in both cases. I would suggest the authors further justify that comparison in the text and compare to classification based on the continuous states. Minor comments - It appears that eq 5 describes the joint distribution $p(c_t, \gamma_t | c_{t-1})$, not just $p(c_t \mid c_{t-1})$.

Rating

7

Confidence

5

Soundness

4

Presentation

4

Contribution

3

Limitations

The authors address the limitations in the text.

Summary

The Authors propose a probabilistic extension to the "Decomposed Linear Dynamical Systems (dDLS)" method by Mudrik et al. (2024). The proposed p-dLDS belongs to a family of models specifically designed to describe neural activity from high-dimensional dynamical data. The effectiveness of p-dLDS is evaluated on a set of synthetic benchmarks and a real dataset. As the primary use-case of p-dLDS -- modeling neural data -- is not my expertise, my evaluation capability for this paper is quite limited, as reflected in my confidence assessment.

Strengths

On top of giving accurate predictions of future states, the proposed method recovers a sensible switching behavior for every one of the examples reported.

Weaknesses

- The experimental results do not report the variance across independent runs. As most of the datasets are synthetic, these should be fairly straightforward to produce. - Apart from adding a probabilistic structure to the sparse coefficients of dLDS, p-dLDS adds a slow-fast decomposition of the latent state (Eq. 4). Yet, no ablation study assessing the importance of this modeling choice is reported.

Questions

- What is the running time of the different baselines? What is the computational complexity of p-dLDS? - From the experimental evaluation, I can see that dLDS and p-dLDS behave quite differently. Have you observed any relationship between the posterior mean of the coefficients $\mathbf{c}_{t}$ in p-dLDS and the deterministic analog in dLDS?

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

N / A

Summary

The paper proposes the probabilistic decomposed linear dynamical systems (p-dLDS) model, which extends the existing dLDS model. With the time-varying offset term and probabilistic formulation, p-dLDS improves upon the dLDS model in terms of robustness to noise and the ability to capture systems that orbit multiple fixed points. The authors show the advantage of the p-dLDS model over SLDS and dLDS through simulated and real datasets.

Strengths

- Clarity: The paper is presented clearly, with easy-to-understand figures and descriptions for the model, inference procedure, and experimental results. The background and related work sections are also well-written, making the paper easy to understand for readers who are not experts in the field. - Reproducibility: The paper includes code and hyperparameter settings to reproduce its experiments, which is crucial for the ML community. - Significance: The probabilistic extension, in addition to the time-varying offset term, adds value to the recently developed dLDS, which suffers from dynamics noise.

Weaknesses

- Experiments Although existing experimental results are exciting, I have a couple of suggestions. One is that, while the paper claims the robustness of p-dLDS to dynamics noise, it does not have experiments that show how robust p-dLDS is to noise. In other words, as you sweep the experiments with different dynamics noises, when do SLDS, dLDS, and p-dLDS show similar/different results? When does p-dLDS fail? Showing the advantage of p-dLDS over other models in a realistic setting of dynamics noise would be important. In addition, the experiments set the number of DOs and discrete states the same for p-dLDS and rSLDS (e.g., Section 4.2). For a fair comparison, I think that the number of discrete states for rSLDS should be set separately on its own via cross-validation. It could be that rSLDS needs more discrete states to perform as well as p-dLDS on e.g., linear classification of reach directions. - Interpretation of discrete states of p-dLDS To my understanding, the paper does not explain how the learned p-dLDS parameters are used to segment the data to infer "discrete" states, as in Figures 2 and 4. In addition, for the NASCAR experiment, p-dLDS leads to fewer segmentations than four true segments. I am unsure whether we could say that this representation is more "parsimonious" while it seems like it leads to incorrect segmentations. In the Lorenz attractor experiment, it also seems like the rSLDS segmentation is more interpretable than p-dLDS segmentation.

Questions

- When does the model break due to the inclusion of time-varying offset terms? In other words, are there cases or hyperparameter settings when the model fails? How should the users avoid this? Are the covariances in line 158 learned or user-specified? - The paper resorts to using PCA to determine the dimensionality of the model. Is there a reason why the dimensionality is not chosen via cross-validation? - Minor: In Section E.1 Figure 6 panel A, it seems like the ELBO peaks at the first few iterations, then drops, and finally converges to a value smaller than the peak. I wonder if there's an explanation for this. - How many hyperparameters are there for p-dLDS? How are these chosen, and how simple or complex is the model selection process for p-dLDS compared to the hyperparameter sweep for rSLDS?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

- As the authors noted, one limitation of p-dLDS is that it assumes Gaussian observation noise.

Summary

The paper presents a probabilistic version of the dLDS model, first presented by Mudrik et al (2024). The primary impetus for this model was to present a version of the dLDS that was more robust to noise, although the authors here also included a slowly-evolving offset term meant to capture evolving fixed points. Inference is conducted by a variational objective which is optimized with a combination of ad-hoc but reasonable methods. The paper paves the way for many other probabilistic state-space models including those with non-Gaussian measurements.

Strengths

Originality - The model is original in the sense that it has not been presented previously. Quality - The paper is mostly well written and well reasoned. I see virtually no major conceptual flaws and only minor problems. Both the simulation experiments and analysis demonstration are well conceived and of interest to the computational neuroscience community. The model is well conceived and represents a natural extension of the dLDS model. Clarity - Very well organized, easy to follow, and a pleasure to read. Significance - I wouldn’t exactly call this a weakness but the extension to dLDS model is so natural as to be virtually obvious. The contribution itself, while seemingly inevitable, was nonetheless executed by the authors first (props to them) and with high quality. It serves as a strong rung in what is sure to be a dLDS ladder.

Weaknesses

The primary weakness is in some technical confusion that i have the with paper. I detail below.

Questions

I am a bit confused about the decomposition of the coefficient prior defined on lines 168-169 following equation (5). Here the authors claim p(c_t|c_t-1, \gamma_t) is defined by a factorization N(c_t-1,\sigma_t-1)N(0,\gamma_t). The authors state that the first of these factors encourages smoothness while the second encourages sparsity. However, it is not at all clear to me that this makes sense. First, the prior is a conditional distribution over c_t for a single coefficient, so what is the factorization over exactly? Did the authors mean to present a joint distribution? Over which variables? I can see that the authors would like c_t to have both properties (sparsity, smoothness) but this particular specification doesn’t make sense. Second, \sigma_t and its inference is undefined anywhere in the paper. Third, on lines 215-216 the authors state that parameters can be learned by SGD “which is possible when we assume that the covariance matrices have diagonal structure.” Do the authors specifically mean the posterior covariances (eq 8), the prior covariances over the latent variables (line 158)? Either way, what is the limitation on SGD learning the parameters when these covariance matrices are not diagonal? Lastly, the authors state on line 187 “the approximate posterior becomes…[some Gaussians]”. It is not clear from the text if the Gaussian form of this variational posterior is a choice of the authors or if this is a variational optimum that drops out of the model structure and factorized structure of q. Can the authors please clarify?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

2

Limitations

No major conceptual limitations but technical details should be clarified