Conditional Normalizing Flows for Forward and Backward Joint State and Parameter Estimation

TL;DR

The paper tackles nonlinear, non-Gaussian state estimation by introducing conditional normalizing flows conditioned with either Transformer or Mamba-SSM embeddings to capture complex posteriors in forward and backward time. It augments the NF training with an optimal-transport–inspired kinetic term to regularize layer transitions, and demonstrates effectiveness on autonomous vehicle dynamics and SIR-based epidemiology, including rollout to real-world COVID-19 data and joint state-parameter estimation. Key findings show that conditional NF can accurately represent multimodal distributions and that Mamba-SSM conditioning performs particularly well for temporally evolving, ODE-like dynamics, while the KE term improves sampling efficiency and stability. The work advances practical, uncertainty-aware state estimation for time-series in domains like autonomous systems and epidemiology, with potential for broader real-world deployment and integration with physics-based models.

Abstract

Traditional filtering algorithms for state estimation -- such as classical Kalman filtering, unscented Kalman filtering, and particle filters - show performance degradation when applied to nonlinear systems whose uncertainty follows arbitrary non-Gaussian, and potentially multi-modal distributions. This study reviews recent approaches to state estimation via nonlinear filtering based on conditional normalizing flows, where the conditional embedding is generated by standard MLP architectures, transformers or selective state-space models (like Mamba-SSM). In addition, we test the effectiveness of an optimal-transport-inspired kinetic loss term in mitigating overparameterization in flows consisting of a large collection of transformations. We investigate the performance of these approaches on applications relevant to autonomous driving and patient population dynamics, paying special attention to how they handle time inversion and chained predictions. Finally, we assess the performance of various conditioning strategies for an application to real-world COVID-19 joint SIR system forecasting and parameter estimation.

Figures (24)

• Figure 1: Sensor data for (a) autonomous vehicle with random switching, and (b) noisy trajectories from the dynamical epidemiological SIR model.
• Figure 2: Generation of conditional embeddings using either a transformer ot mamba-SSM architecture.
• Figure 3: Unconditional normalizing flow transform trained without (top) and with (bottom) kinetic loss term. Each plot is obtained by extracting the distribution of the output of each layer in the forward mapping $\boldsymbol{f(\cdot, \theta)}$ from the target density $\boldsymbol{p_X(X)}$ to the base distribution $\boldsymbol{p_Z(Z)}$. Note the effect of input parameter swap due to a permutation being performed between successive layers.
• Figure 4: Example of training losses for different normalizing flow conditioning operators over 3,000 iterations.
• Figure 5: Confidence regions corresponding to 1, 2, and 3 base distribution standard deviations for forward (top) and backward (bottom) state estimation, for NF with transformer-based conditioning. Predicted density refers to three locations along the trajectory (left to right).