Learning Mixtures of Linear Dynamical Systems (MoLDS) via Hybrid Tensor-EM Method

TL;DR

MoLDS learning from heterogeneous neural time-series is addressed with a hybrid Tensor-EM framework that first uses tensor moment-based initialization to obtain globally consistent estimates of mixture weights and LDS parameters and then refines all parameters with a Kalman-EM procedure. The tensor stage relies on an impulse-response reformulation and Simultaneous Matrix Diagonalization (SMD) to recover $p_k$, $A_k$, $B_k$, $C_k$, $D_k$, while EM updates $Q_k$ and $R_k$ via residual covariances and closed-form LDS updates. Synthetic experiments show improved robustness and accuracy over purely tensor or randomly initialized EM, especially in complex MoLDS settings, and real neural data (Area2 and PMd) demonstrate interpretable, direction-specific dynamical subsystems learned in an unsupervised manner. The results suggest Tensor-EM as a practical and reliable tool for modeling heterogeneous neural dynamics and uncovering latent dynamical motifs across trials and conditions.

Abstract

Mixtures of linear dynamical systems (MoLDS) provide a path to model time-series data that exhibit diverse temporal dynamics across trajectories. However, its application remains challenging in complex and noisy settings, limiting its effectiveness for neural data analysis. Tensor-based moment methods can provide global identifiability guarantees for MoLDS, but their performance degrades under noise and complexity. Commonly used expectation-maximization (EM) methods offer flexibility in fitting latent models but are highly sensitive to initialization and prone to poor local minima. Here, we propose a tensor-based method that provides identifiability guarantees for learning MoLDS, which is followed by EM updates to combine the strengths of both approaches. The novelty in our approach lies in the construction of moment tensors using the input-output data to recover globally consistent estimates of mixture weights and system parameters. These estimates can then be refined through a Kalman EM algorithm, with closed-form updates for all LDS parameters. We validate our framework on synthetic benchmarks and real-world datasets. On synthetic data, the proposed Tensor-EM method achieves more reliable recovery and improved robustness compared to either pure tensor or randomly initialized EM methods. We then analyze neural recordings from the primate somatosensory cortex while a non-human primate performs reaches in different directions. Our method successfully models and clusters different conditions as separate subsystems, consistent with supervised single-LDS fits for each condition. Finally, we apply this approach to another neural dataset where monkeys perform a sequential reaching task. These results demonstrate that MoLDS provides an effective framework for modeling complex neural data, and that Tensor-EM is a reliable approach to MoLDS learning for these applications.

Figures (7)

• Figure 1: Comparison of SMD-Tensor and RTPM-Tensor methods for MoLDS. (a,b) Mean Markov parameter estimation errors for the SMD method decrease as the number of trajectories as $N$ and/or $T$ increase; (c) (c) Heatmap of the best trial result across all $(N,T)$ configurations. (d–f) Corresponding results for the RTPM method.(g) Difference in mean Markov errors between RTPM and SMD (positive values indicate SMD performs better). (h) Scatter plot comparing mean Markov errors of RTPM vs. SMD across configurations, with SMD outperforming in $91\%$ of cases. (i) Example recovery for $N=T=1280$, where SMD recovers both mixture weights and Markov parameters more accurately.
• Figure 2: Performance comparison of pure tensor, Tensor-EM, and random-initialized EM on a simulated MoLDS. (a) Relative accuracy radar plot across metrics of Markov parameter accuracy, weight accuracy, and aggregate parameter accuracy. (b-c) Tensor-EM achieves the lowest Markov parameter and weight errors, while random EM performs the worst and shows high variability. (d) reports aggregate parameter errors (geometric mean of $A,B,C$ errors). (e) Tensor–EM converges in far fewer iterations than random EM, highlighting efficiency. (f) Error-efficiency plot shows that Tensor-EM combines low error with moderate iteration cost, yielding robust and accurate recovery.
• Figure 3: Overview of MoLDS and its applications: (a) Schematic of the MoLDS model structure; MoLDS models observed trajectories as emanating from a set of LDS's. (b) The Area2 Dataset contains neural trajectories from monkey primary somatosensory cortex; reach trajectories and firing rate from an example neuron. (c) The PMd Dataset contains recordings from monkey dorsal premotor cortex; hand speed and neurons' rasters.
• Figure 4: MoLDS application on Area2 Dataset: (a) Validation metrics for different $K$. (b) One-step predictions for two example trials using corresponding LDS components of the MoLDS selected by the lowest validation metrics. (c) Agreement between MoLDS trial assignments (outer ring) and per-direction single-LDS clusters (inner ring); the SLDS-based method (middle ring) cannot provide correct cross-trial clusters. (d) Per-direction impulse response (Markov parameter) curves for assigned MoLDS component vs. the corresponding single-LDS (titles report cosine similarity). (e) MoLDS component usage fractions on held-out test trials.
• Figure 5: MoLDS application on PMd Dataset: (a) Validation metrics for different $K$ (NLL, RMSE, BIC). (b) Test one-step predictions from the validation-chosen MoLDS. (c) Angle bins: dominant component; PD arrows (responsibility-weighted). (d) Impulse response magnitude of MoLDS components. (e) MoLDS component usage fractions on held-out test trials.