Assumed Density Filtering and Smoothing with Neural Network Surrogate Models
Simon Kuang, Xinfan Lin
TL;DR
This work develops an analytic moment-propagation framework for neural-network surrogates within Kalman filtering and RTS smoothing. By treating neural dynamics with a layer-wise Gaussian approximation and an identity-augmentation coupling, it achieves accurate uncertainty propagation through nonlinear models, outperforming standard nonlinear KF variants. The authors advocate cross-entropy as a key metric for calibration and illustrate superior performance on stochastic Lorenz systems and LTI dynamics with nonlinear outputs, including near-optimal closed-loop LQR performance. The approach enables risk-aware decision making in complex nonlinear systems and broadens the applicability of Kalman-like estimation to neural-dynamics models. The results emphasize the importance of calibrated uncertainty over pure point accuracy in nonlinear filtering tasks.
Abstract
The Kalman filter and Rauch-Tung-Striebel (RTS) smoother are optimal for state estimation in linear dynamic systems. With nonlinear systems, the challenge consists in how to propagate uncertainty through the state transitions and output function. For the case of a neural network model, we enable accurate uncertainty propagation using a recent state-of-the-art analytic formula for computing the mean and covariance of a deep neural network with Gaussian input. We argue that cross entropy is a more appropriate performance metric than RMSE for evaluating the accuracy of filters and smoothers. We demonstrate the superiority of our method for state estimation on a stochastic Lorenz system and a Wiener system, and find that our method enables more optimal linear quadratic regulation when the state estimate is used for feedback.