Assumed Density Filtering and Smoothing with Neural Network Surrogate Models

Simon Kuang, Xinfan Lin
Proceedings of The 8th Annual Learning for Dynamics and Control Conference, PMLR 331:1741-1790, 2026.

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. Code available at https://github.com/simontheflutist/analytic-moments.

Cite this Paper

BibTeX
@InProceedings{pmlr-v331-kuang26b, title = {Assumed Density Filtering and Smoothing with Neural Network Surrogate Models}, author = {Kuang, Simon and Lin, Xinfan}, booktitle = {Proceedings of The 8th Annual Learning for Dynamics and Control Conference}, pages = {1741--1790}, year = {2026}, editor = {Sukhatme, Gaurav and Lindemann, Lars and Tu, Stephen and Wierman, Adam and Atanasov, Nikolay}, volume = {331}, series = {Proceedings of Machine Learning Research}, month = {17--19 Jun}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v331/main/assets/kuang26b/kuang26b.pdf}, url = {https://proceedings.mlr.press/v331/kuang26b.html}, 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. Code available at https://github.com/simontheflutist/analytic-moments.} }
Endnote
%0 Conference Paper %T Assumed Density Filtering and Smoothing with Neural Network Surrogate Models %A Simon Kuang %A Xinfan Lin %B Proceedings of The 8th Annual Learning for Dynamics and Control Conference %C Proceedings of Machine Learning Research %D 2026 %E Gaurav Sukhatme %E Lars Lindemann %E Stephen Tu %E Adam Wierman %E Nikolay Atanasov %F pmlr-v331-kuang26b %I PMLR %P 1741--1790 %U https://proceedings.mlr.press/v331/kuang26b.html %V 331 %X 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. Code available at https://github.com/simontheflutist/analytic-moments.
APA
Kuang, S. & Lin, X.. (2026). Assumed Density Filtering and Smoothing with Neural Network Surrogate Models. Proceedings of The 8th Annual Learning for Dynamics and Control Conference, in Proceedings of Machine Learning Research 331:1741-1790 Available from https://proceedings.mlr.press/v331/kuang26b.html.

Related Material