ZENO: Neural Operator Implementation of Nonlinear Filtering Based on the Zakai Equation

Nonlinear filtering is a core computational problem in scientific machine learning and in engineering and neuroscience. Traditional methods often face the dilemma of being analytically intractable and numerically infeasible in high-dimensional settings. The recently proposed ZENO (Zakai Equation Neural Operator) model offers a new approach to approximating the exact unnormalized filtering distribution by combining neural operators with physics-informed loss. This method has not only been validated on multiple simulated dynamical systems but also successfully applied to whole-brain imaging data from Caenorhabditis elegans, revealing movement-related limit cycles.

Background and Challenges: Nonlinear Filtering and the Zakai Equation

The goal of nonlinear filtering is to infer the state and parameters of a latent stochastic dynamical system from noisy observations. The core distribution is described by the Zakai equation—a stochastic partial differential equation (SPDE) that combines the forward Kolmogorov operator (i.e., the Fokker-Planck equation) with observation likelihood updates. However, analytical solutions to high-dimensional SPDEs are nearly impossible, and numerical solutions are also prohibitive due to the curse of dimensionality.

In recent years, advances in operator learning and physics-informed machine learning have made it possible to address such high-dimensional PDEs. ZENO leverages these techniques by embedding the Zakai equation as a physical constraint in the training of a neural operator, thereby directly approximating the exact unnormalized filtering distribution.

The Zakai equation is the core stochastic partial differential equation that describes the evolution of the unnormalized posterior distribution in nonlinear filtering. Unlike the well-known Kushner-Stratonovich equation, the Zakai equation targets the unnormalized filtering density, and its advantage lies in the fact that the equation is linear (with respect to the filtering distribution), thereby avoiding the mathematical difficulties caused by the nonlinear terms in the Kushner equation. The Zakai equation treats the observation process as a continuous-time diffusion process and uses the Girsanov theorem for a change of measure, transforming the filtering problem into solving a linear SPDE for the unnormalized density. The equation consists of the forward Kolmogorov operator (describing state transitions) and the observation likelihood term, and can be formally written as: ( d\rho_t = \mathcal{L}^* \rho_t dt + h(x) \rho_t dY_t ), where (\mathcal{L}^*) is the Fokker-Planck operator, (h(x)) is the observation function, and (Y_t) is the observation process. Solving the Zakai equation yields the unnormalized posterior, which upon normalization gives the filtering distribution.

ZENO Architecture: Physics-Informed Neural Operator

The core of ZENO is to optimize a neural operator so that it satisfies the physics-informed loss derived from the Zakai equation. This loss function ensures that the distribution output by the network conforms to the evolution law of the Zakai equation, thereby learning the underlying filtering distribution without the need for large amounts of labeled data. Specifically, the model trains the neural operator by minimizing the Zakai equation residual, enabling it to generalize to different initial conditions and parameter settings.

After obtaining the posterior distribution, ZENO embeds it into an expectation-maximization (EM) framework to achieve parameter learning and state inference. The EM framework alternates between the E-step (using current parameter estimates to compute the posterior distribution) and the M-step (maximizing the complete-data log-likelihood to update parameters), and the accurate posterior approximation provided by ZENO makes the EM iterations more stable and accurate.

Neural operators are a class of deep learning architectures for learning mappings between infinite-dimensional function spaces, representative examples of which include the Fourier Neural Operator (FNO) and DeepONet. Unlike traditional neural networks that learn mappings from finite-dimensional vectors to vectors, neural operators directly learn mappings from functions to functions, and thus possess resolution invariance and the ability to generalize to different discretizations. In ZENO, the neural operator is designed to approximate the mapping from the initial distribution and observation sequence to the filtering distribution. The physics-informed loss draws on the idea of physics-informed neural networks (PINNs), incorporating the residual of the Zakai equation as a regularization term into the loss function, so that the network output satisfies physical constraints during training, thereby reducing dependence on labeled data and improving generalization.

Experimental Validation: From Simulated Systems to Real Neural Data

To evaluate the performance of ZENO, the researchers tested it on three simulated stochastic dynamical systems: the double-well bistable system, the Van der Pol oscillator, and the Lorenz system. These systems cover typical nonlinear phenomena such as multistability, limit cycles, and chaos. Compared with existing latent SDE inference methods, ZENO demonstrates advantages in both accurately recovering the posterior distribution and learning SDE parameters, being particularly more precise in distribution shape and parameter consistency.

More strikingly, ZENO was applied to the whole-brain imaging dataset of Caenorhabditis elegans (C. elegans). This dataset records the neural activity of immobilized worms, and ZENO successfully recovered the previously observed limit cycles, which correspond to the forward and backward locomotion of the worm. This indicates that ZENO can extract biologically meaningful dynamical patterns from high-dimensional neural data, validating its practicality in real scientific scenarios.

The expectation-maximization (EM) framework is used for parameter estimation in latent variable models. In the context of ZENO, the E-step requires computing the posterior distribution of the latent state, and since nonlinear filtering problems are typically analytically intractable, the neural operator approximation provided by ZENO becomes key to the E-step. The M-step updates parameters by maximizing the complete-data log-likelihood. The EM algorithm guarantees that the likelihood function is monotonically non-decreasing, but it may converge to a local optimum. ZENO's accurate posterior approximation makes the E-step more precise, thereby improving the overall convergence and parameter estimation quality of EM. This combination of neural operators with EM provides an effective iterative optimization strategy for parameter learning in complex dynamical systems.

Significance and Outlook

ZENO transforms the nonlinear filtering problem into an operator learning problem from a PDE perspective, providing a new paradigm for dynamical system inference. The design of its physics-informed loss avoids dependence on large amounts of labeled data while maintaining scalability to complex high-dimensional systems. In the future, this method is expected to be further extended to a broader range of scientific machine learning tasks, such as climate modeling, financial time series analysis, and brain-computer interfaces.

Although ZENO performs well in experiments, its computational cost and dependence on physical equations remain factors to be considered in practical deployment. In addition, how to extend ZENO to non-Gaussian noise or time-varying parameter systems is also a direction worth exploring.

Caenorhabditis elegans (C. elegans) is an important model organism in neuroscience, with a nervous system of only 302 neurons whose connectome has been fully mapped. Whole-brain imaging techniques can simultaneously record the activity of a large number of neurons, providing unprecedented data for studying neural dynamics. However, inferring latent dynamics from high-dimensional neural data is extremely challenging. ZENO successfully recovered movement-related limit cycles from C. elegans data; these limit cycles correspond to the neural activity patterns of forward and backward locomotion, indicating that neural population activity can be described by low-dimensional dynamical systems. This application not only validates the effectiveness of ZENO but also provides a new tool for understanding the principles of neural computation.

One-line Summary

ZENO learns the physical constraints of the Zakai equation via a neural operator to achieve nonlinear filtering and parameter inference.

Paragraph Summary

ZENO is a new neural-operator-based method for nonlinear filtering that directly approximates the unnormalized posterior distribution by embedding the Zakai equation as a physics-informed loss in neural operator training. The method combines an EM framework to achieve parameter learning and state inference, and outperforms existing latent SDE inference methods on the double-well system, the Van der Pol oscillator, and the Lorenz system. More importantly, ZENO has been successfully applied to whole-brain imaging data from Caenorhabditis elegans, recovering movement-related limit cycles and demonstrating the ability to extract biologically meaningful dynamical patterns from high-dimensional neural data. ZENO provides a scalable new paradigm for high-dimensional filtering problems in scientific machine learning.