State and Flow: Differentiable Switching State Space Models for Dynamical Systems Reconstruction

Dynamical systems reconstruction (DSR) leverages neural surrogate models to recover underlying dynamics from noisy, indirect observations. Complex real-world systems often exhibit both continuous nonlinear flows and discrete regime switching. However, existing DSR approaches tackle these challenges in isolation: Neural Ordinary Differential Equations (NODEs) capture expressive nonlinearities but lack discrete states, whereas recurrent Switching Linear Dynamical Systems (rSLDS) identify discrete regimes but restrict per-regime dynamics to linear transformations. Moreover, while DSR surrogates are intended to be simulated and perturbed, current training and evaluation typically reward reconstruction rather than autonomous generation. To bridge these gaps, we introduce a latent autoencoder framework that emphasizes autonomous rollout generation during training, along with two novel architectures: recurrent differentiable SLDS (rdSLDS), an end-to-end trainable rSLDS, and recurrent Switching Neural ODE (rSNODE), which combines discrete state inference with nonlinear NODE flows. Both models infer discrete states in an unsupervised manner. To evaluate models on systems exhibiting coupled nonlinear and switching behavior, we propose the switching Lorenz system as a new benchmark. Across chaotic, piecewise-linear switching, and chaotic switching benchmarks, rdSLDS recovers the ground truth linear regime structure but rarely converges under chaotic switching, where rSNODE alone reconstructs both the attractor geometry and the correct two-state segmentation. Applied to task-fMRI data, rSNODE generates complete regional time series from short initial windows, preserves their temporal and task-activation structure, and identifies states aligned with broad experimental conditions without using task labels. Combining interpretable switching regimes with the expressive power of NODEs yields a scalable, versatile tool for DSR in neuroimaging and beyond.