ZENO neural operator solves complex filtering problems in dynamical systems

ZENO: Nonlinear filtering with Zakai Equation Neural Operators

Machine Learning

Summary

Filtering is a way to estimate hidden changing things from noisy observations, which is important in many fields like engineering and neuroscience. The authors developed ZENO, a method that uses advanced neural networks to solve difficult equations describing how these hidden states evolve over time. This method can learn both the hidden states and system parameters efficiently in complex systems, such as models for brain activity and physical oscillators. They tested ZENO on simulated systems and real brain imaging data to show it works better than existing methods.

What this means in practice

  • For engineering teams: Estimate hidden system states and parameters in complex nonlinear dynamical systems from noisy data using ZENO's operator-based filtering approach.
  • For neuroscience data analysts: Infer neural activity patterns and underlying dynamics from whole-brain imaging data with improved accuracy using ZENO's posterior distribution estimation.

Authors

Aaditya Prasad, Steven W. Flavell, Brian DePasquale

Abstract

State inference and parameter learning for nonlinear filters is a core computation for scientific machine learning with varied applications in engineering and neuroscience. The unnormalized filtering distribution is governed by the Zakai equation, a stochastic partial differential equation (SPDE) that combines the forward Kolmogorov operator (known as the Fokker-Planck equation) with observed likelihood updates. Although solving these high-dimensional SPDEs is often impossible analytically and intractable numerically, advances in operator learning and physics-informed machine learning have made it feasible. Harnessing these advances, we introduce the Zakai Equation Neural Operator (ZENO), a model that optimizes a neural operator against a Zakai Equation based physics-informed loss to approximate the exact unnormalized filtering distribution for latent stochastic dynamical systems. We use this ZENO-inferred posterior distribution in an expectation-maximization (EM) framework to perform parameter learning and state inference. We demonstrate ZENO's ability to accurately infer the posterior distribution and learn its corresponding SDE parameters in three simulated stochastic dynamical systems (Double-Well bistable system, Van Der Pol (VDP) oscillator, and Lorenz system) and compare its performance against established methods for latent SDE inference and learning. We show that ZENO recovers previously observed limit cycles representing forward and backward locomotion from a whole-brain imaging dataset of immobilized \textit{C. elegans} nematodes. ZENO is a novel approach to the problem of nonlinear filtering and dynamical inference through the lens of partial differential equations and operator theory.