Structured Temporal Inference in State-Space Models

We propose a framework for structured temporal inference in nonlinear state-space models (SSMs) with hybrid latent dynamics that mix discrete and continuous variables. Our method follows a two-stage inference: continuous states are estimated via Kalman inspired updates, while discrete variables are sampled by a neural model conditioned on these states, avoiding explicit Markov assumptions. To handle instabilities arising from recurrent dynamics, we introduce stabilization approach, and train all components jointly using surrogate gradient estimators that support REINFORCE-style updates. This design achieves SOTA results across synthetic and real-world datasets, in state estimation, regime detection, and imputation under noise and partial observability.