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

Authors: yenho  chen
noga  mudrik
kyle a johnsen
sankaraleengam  alagapan
adam s charles
christopher j rozell
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 06:05:25

Abstract: Time-varying linear state-space models are powerful tools for obtaining mathematically interpretable representations of neural signals. For example, switching and decomposed models describe complex systems using latent variables that evolve according to simple locally linear dynamics. However, existing methods for latent variable estimation are not robust to dynamical noise and system nonlinearity due to noise-sensitive inference procedures and limited model formulations. This can lead to inconsistent results on signals with similar dynamics, limiting the model's ability to provide scientific insight. In this work, we address these limitations and propose a probabilistic approach to latent variable estimation in decomposed models that improves robustness against dynamical noise. Additionally, we introduce an extended latent dynamics model to improve robustness against system nonlinearities. We evaluate our approach on several synthetic dynamical systems, including an empirically-derived brain-computer interface experiment, and demonstrate more accurate latent variable inference in nonlinear systems with diverse noise conditions. Furthermore, we apply our method to a real-world clinical neurophysiology dataset, illustrating the ability to identify interpretable and coherent structure where previous models cannot.

Main paper: 
 introduction: 
 A central goal in computational neuroscience is to develop models capable of discovering latent structure within noisy, high-dimensional neural signals. By identifying hidden relationships within neural recordings, we can begin to understand, predict, and control the behaviors of the underlying systems. Modeling neural time-series is challenging due to the range of temporal dynamics present. For example, there may be gradual short-term fluctuations, abrupt shifts in response to external stimuli, and long-term global drifts resulting from changes in baseline activity levels [36, 43, 39, 11, 7] .Although black-box approaches based on deep learning are available [34, 20, 38] , their complexity often obscures the relationships learned from the data, making it difficult to extract scientific insights from these models. As a result, practitioners may favor time-varying linear state-space models which offer mathematically interpretable representations by approximating complex dynamics with simple locally linear regimes [26] . However, obtaining latent variable estimates that are robust to dynamical noise and system nonlinearity in these state-space models is challenging. When applied to neural time-series, latent variable estimation may become unstable due to inflexible model formulations or noise-sensitive inference procedures. This can incorrectly produce disparate results for signals generated from the same underlying system. For example, switching linear dynamical systems (SLDS) and related models [42, 44] segment time-series into discrete linear dynamical states, providing a piecewise linear approximation of the underlying system while highlighting coherent groups of activity. However, the assumption of discrete components can be a poor modeling choice for neural signals that contain continuous-valued fluctuations, such as gradual or random changes to the system speed as seen during neural ramping activity [31] or Lévy walk dynamics in the cerebral cortex [27] . As demonstrated in [30] and our experiments in Section 4.3, when applied to real-world datasets, inference of the switching variables can result in rapid, random oscillations between the discrete modes, indicating that the model is unable to identify meaningful structure in the data.To address the limitations of discrete states, the decomposed linear dynamical systems (dLDS) model [30] learns a dictionary set of linear dynamical regimes, referred to as dynamic operators (DOs), that can be modified and combined through a linear combination of sparse coefficients. By allowing coefficients to be time-varying and continuous-valued, dLDS naturally captures both gradual changes by adjusting the coefficient magnitudes and abrupt shifts by changing the set of active DOs over time. Inference is accomplished by optimizing over a cost function that encourages data reconstruction while also constraining the structure of the dynamic coefficients to be sparse and temporally smooth.Unfortunately, there are two critical shortcomings that prevent the robust inference of the latent variables in dLDS. First, the cost-based inference procedure is sensitive to noise, because of the regularization term encouraging temporal smoothness. This term sequentially propagates errors from noisy coefficient estimates over the length of the time series. Consequently, the model may produce inconsistent coefficient estimates on similar signals and have poor multi-step inference performance, indicating that the learned dynamics are unable to generalize well beyond a single time step. Second, the original latent dynamics model lacks a method for accurately representing systems with multiple fixed points, causing DO coefficients to oscillate or switch arbitrarily in a way that may not align with the underlying process. To learn an effective decomposed model in practice, we require a strategy that provides robust estimates of latent variables despite the presence of noise and system nonlinearity.In this work, we address these limitations by introducing the probabilistic decomposed linear dynamical systems (p-dLDS) model. Our approach improves robustness of latent variable estimation while maintaining the richness of a decomposed dynamical systems model. First, we propose a probabilistic inference procedure that reduces the model's sensitivity to temporal noise by accounting for uncertainty in the latent variable estimates over time. Namely, we introduce time-informed hierarchical variables that encourage both sparse and smooth model coefficients. We devise a variational expectation maximization (vEM) procedure to perform inference and learning over this probabilistic structure. Second, we incorporate a time-varying offset term to model systems that orbit multiple fixed points. While we analytically identify model degeneracies with this formulation, we propose an additive decomposition strategy that prevents convergence to trivial solutions.Through several synthetic examples, we demonstrate how these contributions lead to improved accuracy and robustness of latent variable estimation despite difficult noise conditions. We extend these results to an empirically-derived brain-computer interface experiment [8] , showcasing robustness to highly nonlinear observation functions and the ability to extract meaningful insights from the learned latent variables. Finally, we illustrate how our method effectively identifies interpretable and coherent structure in a clinical neurophysiology dataset where previous models are unsuccessful.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: We propose a probabilistic treatment and an extended dynamics formulation in decomposed models (Section 3). We demonstrate that these changes reduce estimation errors and finds coherent structure where previous models fail in many challenging synthetic examples, and a noisy real-world example (Section 4). Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
background and related work: 
 State Space Models. Our goal is to accurately describe the evolution of high-dimensional time-series data y t ∈ R M with the following state-space equations, y t = Dx t + d + ϵ yt , ϵ yt ∼ N (0, Σ y ), (Observations)(1)EQUATIONwhere x t ∈ R N is the latent state, f t (•) is the dynamics function, and D ∈ R M ×N and d ∈ R M describe a linear observation function. Our work focuses on the case when N < M , which compresses high-dimensional signals into a low-dimensional latent space. By choosing f t to be a time-varying linear operator, we can approximate complex nonlinear dynamics with simple locallylinear components, balancing expressivity with mathematical interpretability. However, learning a time-varying linear operator from data can be challenging, and typically requires additional constraints on the underlying generative model to identify meaningful representations.Switching Linear Dynamical System (SLDS). SLDS approximates nonlinear systems by introducing a discrete switching variable z t = {1, . . . , K} into the time-varying linear dynamics equation,x t+1 = x t + F zt x t + b zt + ϵ xt .At each time step t, the latent state x t evolves according to the z t -th linear regime defined by F zt ∈ R N ×N and b zt ∈ R N while the switching variables evolve according to a Markov matrix.Inference is performed through a vEM algorithm, where the approximate posterior of the latent variables is estimated through coordinate ascent updates over tractable subgraphs. There are many extensions of SLDS, such as rSLDS [26] which modifies its generative behavior by informing the transitions of z t with x t-1 . However, switching models are inherently limited when describing complex signals due to their discrete formulation. For instance, a switched representation is unable to learn that a dynamic regime may exhibit a range of variations. In neural systems, these variations may arise from random spiking processes [37, 40] or systems with randomly distributed speeds [41, 12] . In SLDS, each variation is learned as a separate discrete state, thus obscuring that the learned states are related. Furthermore, the switching formulation cannot adapt the learned system to unseen variations (i.e. different levels of random speeds). This can produce unstable inference behavior, where the switching state oscillates unpredictably or collapses to a single uninformative state.Decomposed Linear Dynamical Systems (dLDS). dLDS [30] relaxes the discrete formulation by approximating nonlinear and nonstationary signals with a time-varying mixture of linear dynamical systems (LDS) defined by the following equations and constraints,EQUATIONEvery transition F t is decomposed as a linear combination of sparse coefficients c t ∈ R K and a dictionary of K DOs f k ∈ R N ×N . Figure 1A shows the corresponding graphical model. Inference of the latent variables is accomplished by solving the Basis Pursuit Denoising with Dynamic Filtering (BPDN-DF) [6] objective sequentially for all t and λ 0 , λ 1 , λ 2 > 0,x t , c t = arg min xt,ct ∥y t -Dx t ∥ 2 2 + λ 0 ∥x t -x t-1 -F t x t-1 ∥ 2 2 + λ 1 ∥c t ∥ 1 + λ 2 ∥c t -c t-1 ∥ 2 2 .This produces a point estimate of x t and c t that matches the likelihood function resulting from Equations ( 1) and (2) . In this objective, the dynamic coefficients are encouraged to be sparse through the ℓ 1 penalty and temporally smooth through the ℓ 2 penalty centered around the previous coefficient estimate. However, this approach is sensitive to noise because inference relies on propagating noisy point estimates of c t-1 over time. As a result, BPDN-DF may accumulate errors that can lead to significantly different coefficient estimates on signals sampled from the same generative process. Furthermore, the lack of robustness to noise can degrade multi-step inference performance, causing the inferred system to quickly diverge from the true system. This suggests that the inferred latent variables only capture the local activity narrowly and are unable to accurately represent the dynamics beyond a single time-step. Another drawback of dLDS arises from the dynamics model in equation ( 2) which implicitly assumes that the observed dynamics contain a single fixed point that revolves around the origin. This limits dLDS's ability to model systems that cannot be easily mean-centered such as those with multiple fixed points or nonstationary drifts.Sparse Bayesian Learning with Dynamic Filtering. Sparsity is achieved in probabilistic models through hierarchical scale-mixture priors [10, 2, 4] . To integrate dynamical information into probabilistic sparse signal inference, previous work [33] proposes the Sparse Bayesian Learning with Dynamic Filtering (SBL-DF) framework where the following hierarchical model is defined, Here, p(x t |c t ) = N (Φc t , λ t I) specifies the likelihood, where Φ is a measurement matrix. Sparsity is encouraged through the zero-mean Gaussian priors p(c t,k |γ t,k ) = N (0, γ t,k ) independently placed on each element of the sparse vector. EQUATIONQuestion: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: A limitations section is provided in Appendix H due to space constraints. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
probabilistic decomposed linear dynamical systems: 
 We build upon dLDS and propose a probabilistic decomposed linear dynamical systems (p-dLDS) model. Rather than propagating noisy point estimates of the latent variables during inference, we improve robustness by marginalizing over uncertainty with respect to time. Additionally, we propose a tractable method for extending the dynamics model to systems with multiple fixed points.Note that for any parameter setting, the dLDS local dynamics (eq. ( 2)) reduces to a linear dynamical system (LDS) that is characterized by a single fixed point centered around the origin. Yet, real-world dynamical systems often consist of much more complicated behaviors. Nonlinearities can cause a signal to navigate through multiple fixed points throughout its trajectory, while nonstationarities may change the behavior of the system entirely with new fixed points emerging or disappearing. Simple preprocessing measures, such as mean-centering the data, are inadequate to account for these behaviors. To enable robustness against these behaviors, we introduce a time-varying offset term b t ∈ R N into Equation (2) as a flexible way to account for dynamics not readily captured by the original dLDS latent dynamics model. Lemma 1. Let the transition between any two state vectors x t , x t+1 ∈ R N be defined by the linear dynamics matrix F t ∈ R N ×N and the dynamics offset b t ∈ R N . For any λ > 0, the objective,arg min Ft,bt ∥x t+1 -x t -F t x t -b t ∥ 2 2 + λ∥F t ∥ 2 2 ,is minimized whenF t = 0 and b t = x t+1 -x t .This result, proven in Appendix B.1, reveals that introducing a time-varying offset term makes inference of the dynamics a degenerate problem. While the solution in Lemma 1 minimizes the objective, it fails to capture any meaningful structure in F t as the result of b t being unconstrained. To prevent the convergence to these trivial solutions, we decompose the latent state space as,EQUATIONwhere l t captures fast dynamics and b t captures slow-varying trend behavior. These latent variables follow the dynamic equations l t+1 = l t + F t l t + ϵ lt+1 and b t+1 = b t + ϵ bt+1 where ϵ l , ϵ b ∈ R N represent noise sampled from ϵ b ∼ N (0, Σ b ) and ϵ l ∼ N (0, Σ l ) respectively.In decomposed models, we aim to achieve two goals simultaneously: sparsity and smoothness of coefficients over time. Motivated by this, we incorporate dynamics-informed probabilistic structure. First, we assume that each coefficient evolves independently of the others. Second, we introduce a hierarchical variance parameter γ t,k that controls the sparsity for each c t,k . Moreover, we introduce dynamics information during sparse inference by encouraging a similar active support set in consecutive time slices through the variance hyperpriors. Put together, the resulting coefficient transition density in p-dLDS becomes,EQUATIONWe define the first term on the right-hand side with the following functional form,EQUATIONThis density captures the constraints of sparsity and smoothness for the inferred coefficients c t,k . When the variance around zero γ t,k is small, this structure promotes sparsity by shrinking coefficient values towards zero. Conversely, when the variance around the previous time step σ 2 t-1,k is small, it encourages smooothness by shrinking coefficients towards the previous value. While the idea of combining two shrinkage effects in a single density has been explored in previous works [15, 5, 24, 18] , those approaches generally require manual balancing of the two penalties. In contrast, we devise a procedure in the following section that estimates these variance parameters automatically during inference and learning.The second density on the right-hand side from equation ( 5) is defined similarly to the hyperprior in SBL-DF. (i.e., p(γ t,k |c t-1,k ) = IG(ξ, ξc 2 t-1,k ) where ξ weighs the influence of the dynamics when estimating γ t,k ). The resulting graphical model is shown in Figure 1B . We note that since the value of the previous coefficient is squared, the overall prior placed on the inverse gamma density follows a χ 2 distribution.The joint distribution of p-dLDS is given by, p(x, y, c, γ|θ) = p(x 1 )EQUATIONwhere we denote x = x 1:T for brevity. Exact posterior inference is intractable due to the nonconjugacy introduced by incorporating time-informed sparsity-inducing structure into the graphical model. As a result, we devise a variational expectation maximization (vEM) procedure where the approximate posterior is factorized as p(x, c, γ|y, θ) ≈ q(x)q(c, γ).Here, the parameters are given by θ ∈ {f 1:K , D, d, Σ y , Σ x , Σ c }. Our approach contrasts with BPDN-DF, which estimates latent variables through separate ℓ 1 problems at each point in time. Instead, we preserve the time-dependence structure within each class of latent variables and leverage efficient inference algorithms that marginalize over uncertainty with respect to time. In general, we seek to maximize the variational lower bound, L q (θ) = E q(x)q(c,γ) [log p(y, x, c, γ|θ) -log q(x)q(c, γ)],with coordinate ascent updates on the latent state posterior, the dynamics coefficients posterior, and the model parameters.Updating Latent State Posterior. The optimal coordinate ascent variational update is given by,EQUATIONOur assumed decomposition in equation ( 4) allows us to define the latent state transition density as p(x t+1 |x t , c t ) = p(x t+1 = l t+1 + b t+1 ) = N (x t+1 ; l t + F t l t + b t ). Substituting this into equations ( 7) and ( 8), we get that the optimal coordinate ascent approximate posterior becomes,EQUATIONwhere, µ 1 and Σ 1 are the mean and covariance of the initial state.Lemma 2. Let l, b ∈ R N be independent random variables such that l ∼ p(l) and b ∼ p(b). Their sum x = l + b is distributed according to N (µ l + µ b , Σ l + Σ b ) when 1) p(b) = N (µ b , Σ b ) and p(l) = N (µ l , Σ l ) and when 2) p(b) = δ(b -µ b ) and p(l) = N (µ l , Σ l + Σ b ).We leverage Lemma 2 to reparameterize the state space trajectories into a deterministic and a stochastic component. The deterministic component captures the slow-moving offset density q(b)= T t=1 δ(b t -b t ), where b t is estimated using a moving average of window size S which can be efficiently parallelized. The remaining dynamics are captured in the stochastic component. We define the family of variational distributions to be the class of linear Gaussian state space models, such thatq(l) = N (l 1 , Σ 1 ) T t=2 N (l t-1 + F t-1 l t-1 , Σ x ).Conditioned on estimates b and samples of c, the optimal coordinate ascent variational update for q(l) is efficiently computed using the Kalman Smoother [35] . We provide a full derivation of the update rule and Lemma 2 in Appendix B.2.Updating Dynamics Coefficient Posterior. Sparse probabilistic representations introduce non-Gaussian factors which prevent closed-form message passing inference. Specifically, nonconjugacy arises from the inverse gamma term p(γ t+1,k |c t,k ) since it is parameterized by c 2 t,k ∼ χ 2 . Moreover, the posterior distribution over the coefficients is highly multi-modal as a result of the implicit tdistribution in our hierarchical model. To update the coefficient posteriors, we propose a three-step procedure, where we factorize q(c, γ) = q(c)q(γ). First, we obtain an initial estimate of the variational distributions using SBL-DF. Second, we update q(c t,k ) = N (c * t,k , γ t,k ) using stochastic gradient descent (SGD) overEQUATIONwhere we have estimated the expectation in the optimal coordinate ascent update rule using samples from γ ∼ q(γ) and l ∼ q(l). To retain coefficient sparsity, we only update coefficients within the active support set. In our work, this is defined as coefficients that have an initial estimate of |c t,k | > η where η = 10 -4 . Finally, we update q(γ) based on closed form conjugacy rules.Update Parameters. Given our updated posteriors of the latent variables, we proceed to update the model parameters based on the ELBO,EQUATIONwhere we estimate the expectation with samples from our variational distributions and drop terms not dependent on θ. We use SGD to update all model parameters, which is possible when we assume that the covariance matrices have diagonal structure.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: We provide functions to generate the synthetic datasets from our experiments. Unfortunately, we are unable to release the neurophysiological dataset due to request from our collaborators.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
results: 
 We demonstrate p-dLDS in a variety of synthetic examples, highlighting improved robustness to noise and system nonlinearity. Additionally, we apply our model to a clinical neurophysiology dataset, revealing interpretable patterns where previous methods fail. We compare our method against SLDS, rSLDS, and dLDS as described in Section 2. All datasets are split 50:50 for training and testing. Due to space constraints, we provide full descriptions of the simulation setup in Appendix C and metric definitions in Appendix D. NASCAR with Random Speeds. We evaluate our inference procedure on the NASCAR dataset [26, 32, 23] . Since this system is easily mean-centered, we can isolate the effect of our proposed inference procedure as offset terms are not necessary. To make this dataset more realistic, we introduce speed variability into the ground truth dynamics as opposed to having a perfect constant speed at all time. Specifically, whenever a trajectory enters a new segment of the track, the system experiences a random change in speed. We trained all models on 30 trials, each consisting of 1000 time steps (Fig. 2A ) with randomly sampled initial points, and a randomly constructed 10-dimensional linear observation matrix (see Appendix C.1 for full details). Performance is evaluated on 30 held-out trials where the ground truth switching states are defined by the different segments of the track. In all models, we set the M = 10, N = 2 and K = 4 for DOs or switching states. In our experiments, we define the "discrete states" for decomposed models as the DO state with the largest coefficient magnitude.Figure 2B shows that changes in the system speed mask the true transition behavior between segments of the track in rSLDS. Moreover, dLDS identifies coherent segments, but inappropriately learns a different switching pattern for outer and inner edges of the track. In contrast, p-dLDS identifies a switching pattern most consistent with the true track segments despite the presence of noise and randomness in the system's speed. We note that while there are four true segments, decomposed models form a more parsimonious representation by identifying similar behaviors in different track segments such as in both edges and curves.Table 1 summarizes our quantitative evaluations on three metrics: 1) the mean squared error (MSE) between the learned and ground truth latent dynamics, 2) the MSE between the inferred and true switch rate to determine agreement of the discrete switching behavior, and 3) the 100-step inference R 2 to demonstrate that the learned system generalizes beyond a single step on held-out data. (See Appendix D for mathematical definitions). We see that p-dLDS broadly outperforms existing methods in all metrics and significantly improves inference for decomposed models.Lorenz System with Random Ramping. Next, we consider the Lorenz system, a chaotic nonlinear system with multiple fixed points, exploring the effect of the offset term. The system is described by the differential equation,ẋ = [σ(x 2 -x 1 ), x 1 (ρ -x 3 ) -x 2 , x 1 x 2 -βx 3 ]⊤ where the parameters ρ = 28, β = 8/3, and σ = 10 define a chaotic attractor with two opposing lobes (Fig. 2D ). We introduce continuous fluctuations in the underlying dynamics by randomly ramping the system's speed throughout each trajectory. This is accomplished by adjusting the evaluation time intervals given to an ODE solver. Similar to before, we randomly construct a linear observation function with M = 10 and train models on 30 randomly constructed trials with 1000 time points each (see Appendix C.2). Furthermore, we define the ground truth switch events as the time points when the signal transitions between the two lobes in addition to the moments when a ramping period concludes. All models are trained with a latent space of N = 3 and K = 4 states or DOs.In figure 2D , we see that rSLDS does not distinguish between the different speeds along the outer and inner sections of the attractor. Instead, the discrete states obscure the continuum of speeds by incorrectly grouping all activity in each lobe into a single regime. Furthermore, we observe that dLDS is limited without an offset term, unable to accurately represent multiple fixed points. Instead of aligning with the two attractor lobes, transitions in the dominant coefficients occur radially relative to the origin and fail to reconstruct the two orbiting fixed points. Conversely, p-dLDS's offset term enables learning a system where coefficients better match the true geometry. This representation correctly recovers differences between the outer and inner sections of the attractor while also accurately reconstructing the two orbiting fixed points. Moreover, this leads to improved estimation of latent dynamics, a switching rate that agrees with the true system, and improved multistep inference performance as shown in Table 1 . We now turn to an empirically-derived synthetic experiment related to brain-computer interfaces, where the dynamics and observation functions are nonlinear and derived from analysis of neural data. Our focus is on the reaching task, a neuroscience experiment designed to study motor control in nonhuman primates [19, 22] . In this experiment, the subjects are trained to reach towards visually cued targets, while neural activity is recorded from motor-related areas such as electromyography (EMG) data from arm muscles. Each trial consists of two distinct phases: preparation and movement. In the preparation phase, the subject plans its movement while keeping their arm still. In the movement phase, the subject physically reaches towards the target. The goal of this experiment is to decode reach intention from neural data. We construct a dataset by first simulating a spiking neural network with known latent factors [8] trained to reproduce empirical EMG signals from the center-out reach task in [22] . Spikes are then converted to 50-channel local field potentials (LFP) recordings via a weighted, delayed sum of synaptic currents (see Figure 3A and B ) [29, 17, 16] . Our dataset contains 150 6-second trials sampled at 250 hz, where each trial represents one out of eight reach directions visually cued at a random start time. PCA identifies that three components captures 98% of the variance. Thus, we set M = 50, N = 3, and K = 4 DOs or discrete states.Figure 3C shows the trial-averaged DO coefficients from p-dLDS, which change smoothly and cyclically according to the true reach angle. Additionally, the DOs appear to differentiate between the two distinct dynamical regimes, where the activity of f 1 and f 3 localize to the preparatory and movement phase respectively. Quantitatively, we compute the linear classification accuracy of the reach angles using the average state activity over time as features (see Appendix D.5). Figure 3D shows that classifiers built from SLDS features fail to capture the full continuum of reach angles. This limitation occurs because the discrete switching states are unable to efficiently capture the diversity of activity present in the LFPs, which arise from randomness inherent in the spike sampling process. Consequently, the inferred features from SLDS generalize poorly to held-out data. In contrast, the p-dLDS classifier predictions recover the full spectrum of reach angles since features are naturally continuous and the inferred coefficients can adjust the learned DOs to accurately capture the activity in the held out data. Table 2 shows that p-dLDS outperforms all other models in state and dynamics reconstruction as well as the top-1 and top-3 reach classification accuracy. The learned patterns broadly generalize to the held-out data.We demonstrate p-dLDS on LFP recordings from the subcallosal cingulate cortex (SCC) in patients with treatment-resistant depression (ClinicalTrials.gov identifier NCT01984710). Subjects are asked to watch videos with different emotional content (positive, negative, and neutral), describe the videos, and then discuss how the video made them feel. SCC dynamics have been previously shown to provide a quantitative signal for the presence of emotional content [14] and depression recovery [1] . Thus, we hypothesize that the underlying dynamics may provide information about emotional changes throughout the experiment. We apply p-dLDS to a single patient's LFP spectrogram data (Fig. 4B, E ) within the 0-40 Hz frequency range (M = 40). PCA indicates that the first 7 components explain 90% of the variance. Therefore, we train a model with N = 7 latent dimension and K = 4 DOs.In Figure 4 , rSLDS and dLDS produces a high degree of state oscillations making it difficult to identify time intervals with consistent emotional content. In contrast, p-dLDS infers coherent structure that corresponds to changes in emotional content in the trial. For example, f 4 (red) coincides with resting, f 2 (orange) with positive videos, and f 1 and f 3 (blue and green) to negative videos (Fig. 4C ). Importantly, this structure persists even on held out data from the second half of the session (Fig. 4 F). We note this preliminary analysis on a single subject isn't intended to make a claim about specific neurophysiological responses to emotional content in this brain region, but generally highlights that p-dLDS identifies meaningful dynamical modes where previous models are unable to. 
conclusion: 
 In this work, we present a probabilistic decomposed linear dynamical systems model that can be used to discover meaningful representations in neural signals. By marginalizing over uncertainty in latent variable estimates and incorporating an offset into the dynamics, we enhance robustness and improve a variety of performance metrics. Some areas of future work includes exploiting structure in the offsets to automatically identify window size and extending the probabilistic model to include more complicated emissions distributions, such as the Poisson likelihood commonly used to model neural spiking data [28] .A Appendix / supplemental material A.1 p-dLDS Algorithm Algorithm 1 describes the proposed inference algorithm. In our experiments, we set n = 1 and η = 10 -4 and observe that the model converges. Below we use the notation hat notation for latent variable estimates or samples and the variable itself to represent the parameters of the variational distributions.Algorithm 1 Variational EM for Probabilistic dLDS Require: M observation dimension, N latent state dimension, K number of dynamic operators, S moving average window size, ξ SBL-DF trade-off parameter, n number of samples to estimate expectations, η sparsity threshold, θ model parameters.// Initialize parameters c t ← 0 D i,j ∼ N (0, σ 2 ) f k,i,j ∼ N (0, σ 2 ) x t ← D + y t ▷← MovingAverage S ( x 1:T ) c t ∼ q(c t ) F t ← K k=1 f k c k,t l 1:T , Σ x ← KalmanSmoother(y 1:T , b 1:T , F 1:T , θ)// Update Coefficient Posterior Initialize q(c) and q(γ) jointly with SBL-DF. Update q(c t,k ) with SGD over equation (10) for densities where |c t,k | > η.Update q(γ t ) ← IG(ξEQUATION) // Update Parameters Update θ with SGD over equation (11) . end while B Latent Variable Inference B.1 Lemma 1 Derivation Lemma 1. Let the transition between any two state vectors x t , x t+1 ∈ R N be defined by the linear dynamics matrix F t ∈ R N ×N and the dynamics offset b t ∈ R N . For any λ > 0, the objective,arg min Ft,bt ∥x t+1 -x t -F t x t -b t ∥ 2 2 + λ∥F t ∥ 2 2 ,is minimized when F t = 0 and b t = x t+1 -x t .Proof. Let r t = x t+1 -x t . We can rewrite the reconstruction objective in the following form, 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We specify that all datasets are split 50:50 for train and test in the main paper. Moreover, we provide optimizer and hyperparameter settings in Appendix F.1.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: We provide standard deviations in Appendix E.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our method does not generate synthetic content that is at risk of abuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification:Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [Yes] Justification: We include details on where we obtained our data (either synthetically generated, or part of an existing dataset) in the main text and appendix. We also point to the commonly package used for SLDS and rSLDS in the appendix. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)? Answer: [Yes] Justification: We provide the IRB identifier in Appendix G which includes the informed consent form with experiment details. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [Yes] Justification: We include the IRB identifier number in Appendix G, but removed any information related to the institution conducting the experiments. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Ft,bt∥r t -F t x t -b t ∥ 2 2 + λ∥F t ∥ 2 2 .This objective is identical to the standard ridge regression with an unpenalized intercept term [13] .The solution is obtained by first centering the data, and then solving for the parameters using the solution for the standard Tikhonov regression. Below, we define the centered data as xt and rt for inputs and outputs respectively. Finally, we can use these values to obtain the following estimates of the parameters,b t = µ t F t = (x ⊤ t xt + λI) -1 x⊤ t rtHowever, when there is only a single datapoint, we get that xt = 0, which results in F = 0.This result arises from having only a single observation for any dynamic transition, which leads to a singular design matrix. Although we can improve our estimate of F t by collecting more samples along a given trajectory, this is impractical when dealing with naturalistic time-series. For instance, it may be infeasible to collect more data from the exact same initial condition in a naturalistic environment due to noise in the experimental setup. In chaotic systems, minor deviations can lead to drastically different outcomes over long time horizons. Even if it were possible to precisely control for the initial condition of the signal, the presence of dynamical noise can cause initially aligned time series to quickly drift out of alignment. Consequently, it is not uncommon to observe a single transition between any two time points, as it is not guaranteed that events across multiple trials will be well-aligned. Proof. Case 1. Let l ∼ N (µ l , Σ l ) and b ∼ N (µ b , Σ b ). The sum of normal random variables follows a distribution that results from convolving their individual distributions,q(x) = q(l + b) = q(l) * q(b) = N (µ l , Σ l ) * N (µ l , Σ b ) = N (µ l + µ b , Σ l + Σ b )This is a standard result from probability theory.Case 2. Now let l ∼ N (µ l , Σ l + Σ b ) and b ∼ δ(b -µ b ). Similarly, the distribution of the sum of these variables is distributed according to their convolution,q(x) = q(l) * q(b) = N (µ l , Σ l + Σ b ) * δ(b -µ b ) = ∞ -∞ N (x -τ ; µ l , Σ l + Σ b )δ(τ -µ b )dτ = N (x + µ b ; µ l , Σ l + Σ b ) = N (x; µ l + µ b , Σ l + Σ b ),where the fourth line is the result of the sifting property of delta distributions. Since the final distribution in Case 1 and Case 2 are identical, we complete the proof.The optimal coordinate ascent variational update is given by the following equation,EQUATIONConditioned on estimates of b 1:T and samples of c 1:T , the factor graph of equation ( 12) corresponds exactly to a time-varing Linear Gaussian State Space Model. Thus we can leverage the efficient inference algorithms such as the Kalman filter and RTS smoother when computing the marginals of the variational distribution of l 1:T .Our dataset is constructed from publicly available data and code from the center-out reach task in [22, 8] . We obtain latent factors from spiking networks that are trained to reproduce empirically measured EMG signals, given a 3-dimensional input that specifies the go input and the reach angle.In our experiments, these factors are considered ground truth. Our trained factor-based spiking network then generates spiking activity for 1200 neurons. Synaptic currents are used as inputs into the Weighted Sum of synaptic currents LFP proxy method (WSLFP) [29] , as implemented in the wslfp Python package [17, 16] . As WSLFP is a function of the relative location of neurons and electrodes, we place neurons randomly within a 5 mm by 10 mm by 1 mm region and electrodes in a grid centered in this region. The result is a multi-channel LFP dataset with nonlinear dynamics and measurements characteristic of systems neuroscience. The multi-step inference performance is computed with the following R-squared metric,EQUATION)where k is the number of steps from the initial condition, ȳ is the mean estimator for each trajectory and y t+k is the model prediction after applying the inferred dynamics for k steps. When testing, model parameters such as the dynamics and observation matrices are frozen, while specific latent variables are estimated based on the held-out data. In Table 1 , we show results for k = 100.We measure the accuracy of the latent dynamics with the mean squared error (MSE) of the inferred speed, defined as,EQUATION)where the true speed ẋt = x t+1 -x t is computed from the denoised ground truth latent state, and the predicted speed ẋt = x t+1 -x t is computed using the model's 1-step prediction. Since latent trajectories are only identifiable up to a linear transformation, we align the inferred trajectories with the true trajectories using a least squares fit before computing this score. More specifically, we find the optimal linear transformation U ∈ R N ×N between the estimated and true states across all trajectories by solving,EQUATIONSimilarly, we measure the accuracy of the latent state space by computing the MSE after a linear alignment between trajectories from the inferred and true state space. We use this metric only for the reaching example, since the true observation function is a complex nonlinear function,EQUATION)The linear alignment U ∈ R N ×N between the estimated and true states across all trajectories is computed by solving the least squares problem in equation (15) .Evaluating the accuracy of the switching behavior is a more difficult task. In fact, developing a procedure that matches predicted switch times with true switch times can lead to a complicated optimal transport procedure. To simplify the evaluation of switching times, we marginalize over time, and compare only the MSE of the switch rate defined as,EQUATION)where m is the number of trials, r i is the true switch rate for the ith trajectory, andr i = 1 T T t=1 1{z t ̸ = z t-1} is the predicted switch rate. Intuitively, r i is the number of times that the state or dominant DO changes between consecutive time points normalized by the length of the interval T . In switching models, switch events are defined as a time point where the current inferred dynamical state differs from the state in the previous time step. Similarly in decomposed models, switch events are defined as time points where the active set of DOs change from the previous time step.In the NASCAR example, r i is defined with the number of transitions between ground truth segments. In the Lorenz example, r i is defined by the number of times that the trajectory switches between the two lobes in addition to the number of ramping periods.We quantitatively evaluate the reaching experiment with a classification task. Here, we want to determine whether the learned systems can be used to distinguish between different reach directions. Recall that switched models infer a switching variable for each time point where z t ∈ {1, . . . , K} while decomposed models infer a coefficient vector c t ∈ R K . Rather than viewing z t as an index, we can equivalently view it as a one-hot encoded vector z t ∈ {0, 1} K which describes whether a particular switching state is active at any given time. This matches the dimensionality of the variables in both switched and decomposed systems.For simplicity, we focus on linear logistic regression classifiers in our experiment. If we let the inputs be z t and c t directly, then our classifiers quickly overfits since there are many more input features than trials. Specifically, the number of features scales linearly with the number of time points and systems O(T K). Instead, we marginalize over time and compute features from the estimated latent variables by averaging state activity over time. In switched models, this is the average one-hot encoding value over time. Similarly, this is the average coefficient value in decomposed models. However, for each dynamical state, we compute separate features for positive and negative coefficient values to prevent interference between them. In this setup, the input (feature) dimensionality scales according to O(K) while the output dimensionality of the linear classifiers are the reaching directions. For all classifiers, we perform a grid search over the values {10 i } 4 i=-4 to identify an appropriate amount of L2 regularization. Top-k accuracies are a standard metric in machine learning [21, 3] and computed using the estimated class probabilities from the logistic regression classifier.Figure 6A demonstrates that our inference procedure converges to a local optimium while Figure 6B shows a full sweep of the multi-step inference metric. Tables 1 in the main paper reports the final value. For completeness, we include Tables 3 and 4 which reports the means across 5 seeds of each model, and includes the standard deviations in parenthesis. 3 and 4 report the final values. For each model, we visualize the trial-averaged dynamic regime activity of each reach direction (Fig. 7 ). In SLDS, this is visualized by considering the discrete states as a one hot vector over time. When a dynamic regime is active, that state will have a value of 1 while the unactive states will have a value of 0. Thus the trial averaged value of each state must have a value in the interval [0, 1]. In dLDS, we plot the inferred coefficient value without any modification.Although SLDS correctly identifies preparatory and movement phases using states 4 and 3 respectively, it fails to differentiate dynamics occurring outside of these expected phases, incorrectly grouping unrelated regions together. Furthermore, the discrete formulation produces very similar activity patterns across all reach angles, obscuring any differences that are present. In dLDS, we observe that the features change smoothly and cyclically with the reach angle. However, the dynamic operator activity do not localize to the preparatory and movement phases due to a limited inference procedure.Table 5 : Inference performance for the reaching experiment (see Figure 3 ) on a held-out test set.Top-1 and Top-3 accuracies are obtained by predicting reach directions from latent variable features using linear classifiers. State and Dynamics MSE are computed with respect to true latent variables. We report standard deviations in parenthesis across 5 seeds. For switching models, we rely on the ssm package which allows for efficient Bayesian inference for a variety of state space models [25] . We set the variational posterior to structured_meanfield, and the fitting procedure to laplace_em as recommended by the developers. Additionally, we set the distributional form of the dynamics and emissions matrices to Gaussian.The hyperparameters of dLDS primarily consists of the lagrange multipliers in the BPDN-DF objective including λ 0 , λ 1 , λ 2 . We find the optimal value of these hyperparameters using a random search with a fixed budget of 1000 evaluations. For each hyperparameter, we uniformly sample over the log of the interval [10 -3 , 10 3 ] and evaluate it against the BPDN-DF objective. For the NASCAR experiment, we found that λ 0 = 1.044, λ 1 = 0.254, and λ 2 = 0.023 resulted in the best performance.For the Lorenz experiment, we found that λ 0 = 0.628, λ 1 = 2.010, and λ 2 = 0.0124 yielded the best performance.For p-dLDS, the relevant hyperparameters consists of the SBL-DF dynamics tradeoff ξ, and the offset window size S. We use random search with a budget of 1000 samples to determine the values of S and ξ and fit a separate model for each set of hyperparameters. In the NASCAR experiment, we isolate the effect of the probabilistic inference procedure by setting S = T , removing the influence of the time-varying offset term. For ξ, we perform a random search by uniformly sample over the log of the interval [10 -3 , 10 3 ] and found that ξ = 0.945 was optimal. For the Lorenz experiment, we also optimize for the window size S by uniformly sample a discrete index on the interval {2, . . . , T }.For the Lorenz experiment, the optimized hyperparameters are S = 85 and ξ = 8.928. For the real dataset, the optimal offset is S = 76 which is smaller than the timescale of p-dLDS coefficient switching (around 150 time points), suggesting that the same DO dynamics may persist even as the fixed points of the system fluctuates throughout the experiment.We perform hyperparameter sweep on our institution's HPC cluster using small-scale CPU resources which consists of Dual Intel Xeon Gold 6226 CPUs. Once hyperparameters have been optimized, it is possible to run each experiment within approximately 2 hours on the 2020 edition of the M1 Macbook Pro.Data was collected as part of a study investigating deep brain stimulation for treatment-resistant depression (TRD). The study is pre-registered in ClinicalTrials.gov (identifier NCT04106466). The study protocol was approved by the IRB (identifier IRB00066843). Informed consent was obtained from participants before participation in the trial. Patients receive no monetary compensation, but instead have their DBS electrodes and Summit RC+S IPG device provided free of charge. The analysis focused on LFP signals from a single participant with all personally identifiable information removed.While our proposed method demonstrates strong performance in our experiments, there are many limitations. For instance, our approach does not have a strong mechanism for generating future unseen coefficients. Our assumed coefficient transition model is primarily motivated by our desire to obtain smooth coefficients over time. However, we believe that they may be more complex transition models that can both capture persistent activity in challenging systems while also being an accurate forecaster, such as a deep learning based transition model. Another limitation of our approach is that our method assumes smoothness in the latent space. However, we do not explore the possibility of having sparse structure in the latent space which can be easily accomplished in BPDN-DF by adding an L1 penalty over x.NeurIPS Paper ChecklistCode is available at: https://github.com/siplab-gt/probabilistic-decomposed-linear-dynamical-systems 


Decision: This is a nice extension of recent work on decomposed LDS models. The novel contribution is a probabilistic formulation of the model and a variational EM framework for state and parameter estimation. I was particularly impressed with the improvements in state segmentation on real datasets. Please incorporate the changes and clarifications described during the discussion period. I think this will be a solid contribution to NeurIPS.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: The primary weakness is in some technical confusion that i have the with paper. I detail below.
questions: 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?
limitations: limitations: No major conceptual limitations but technical details should be clarified
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: limitations: - As the authors noted, one limitation of p-dLDS is that it assumes Gaussian observation noise.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: limitations: N / A
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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})$.
limitations: limitations: The authors address the limitations in the text.

Comments: comment: 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 your thoughtful and informative responses to my questions and for exploring classifying based on the continuous state for the reach experiment. I think it be worthwhile to incorporate these points into the paper. 

I remain convinced that this paper should be accepted. I intend to keep my score the same, and I encourage the other reviewers who scored lower than I did to consider raising their score.