ZENO: Nonlinear filtering with Zakai Equation Neural Operators

Aaditya Prasad; Steven W. Flavell; Brian DePasquale

cs.LGOct 23, 2026

ZENO: Nonlinear filtering with Zakai Equation Neural Operators

Authors: Aaditya Prasad, Steven W. Flavell, Brian DePasquale

Organizations: Picower Institute for Learning & Memory, Department of Brain & Cognitive Sciences, Massachusetts Institute of Technology, Cambridge, MA 02139, USA · Howard Hughes Medical Institute, Picower Institute for Learning & Memory, Department of Brain & Cognitive Sciences, Massachusetts Institute of Technology, Cambridge, MA 02139, USA · Department of Biomedical Engineering, Boston University, Boston, MA 02215, USA

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 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.

Figures & tables

Figure 1: Architecture Diagram. A) Generative Model. B) Causal GRU encodes observations into context trajectory. C) Context trajectory is fed into the branch of a Deep Operator Network along with substeps to produce basis functions. Trunk network can be queried at any

\mathbf{z}

and outputs

\log\rho_{\theta}

, the unnormalized log density. IC loss between the learned posterior and an IC prior

\mathcal{N}(0,\mathbf{I})

combined with Fokker Planck Residual Loss between observations and Jump Loss at observations D) Posterior Expectation

\bar{z}=\mathbb{E}_{p(z|x)}[z]

computed through MALA based MCMC sampling. E) Dynamics and Observation Parameters fit by regressing against

\bar{\mathbf{z}}

Figure 2: Latent Inference of Simulated Dynamical Systems. A) Exact filtering distribution for the double well system vs filtering posterior distribution estimated by ZENO. B) Latent Posterior Mean Estimate for the Van Der Pol System. C) Latent Trajectory Inference for the Lorenz Attractor. D) Comparison of ZENO’s normalized root mean squared error with the exact optimal filter for each system.

Figure 3: ZENO Drift Estimation Performance. A) True cubic drift of double well (dashed line) overlayed with model estimate. B) True Van Der Pol Vector field (Left) vs Model’s infered Vector Field. C) Lorenz Attractor Dynamics plotted in the $xz$ plane. D) Model’s NRMSE performance compared to baselines. Dots represent seeds and error bars are 1 SD.

Figure 4: Whole-brain calcium dynamics in C. elegans A) Example ZENO reconstruction of RIML neuron. B-C) ZENO recovers limit cycles. D-E) Inferred latents are consistent with known behavioral labels. Model performance compared to baselines on F) reconstruction $R^{2}$ , G) behavior decoding accuracy, and H) forecasting RMSE. Points represent individual worms, separated into NoStim and Stim conditions. Error bars are 1SD.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

	Double well	Van der Pol	Lorenz
Latent dimension $d$	1	2	3
Drift parameters	$a=4$	$\tau=15,\ \mu=2$	$(10,28,8/3)$
Diffusion $G$	$1$	$I_{2}$	$0.3\,I_{3}$
Initial state $z_{0}$	$\mathcal{N}(0,1)$	$\mathrm{Unif}([-3,3]^{2})$	$\mathcal{N}(0,I_{3})$
Duration $T$ (s)	3	3	2
Observations per trajectory	301	301	100

Appendix

Table 1: Simulated datasets.

d

is the latent dimension and

D

the number of observed channels.

G

multiplies

dW

. The observation noise SD

\epsilon

is in physical units for the double well and Van der Pol and in units of the standardized clean signal for Lorenz. Trajectory counts are train / validation / test.

Component	Hyperparameter	Value
Context encoder	causal GRU hidden / context dimension	64 / 64
DeepONet	basis size $p$	64
	branch MLP hidden width	128
	trunk MLP width $\times$ depth, activation	$64\times 3$ , softplus
Drift $f_{\phi}$	MLP width $\times$ depth, activation	$64\times 3$ , tanh
Diffusion	scalar $g$ : initial value / lower bound	1.0 / 0.05

Appendix

Table 2: ZENO hyperparameters. Shared values apply to every dataset; the last four rows are the only per-dataset settings.

Shared hyperparameter	Value
Kernel	smoothly switching linear (SSL)
Variational EM iterations	50
Natural-gradient E-steps per iteration	15
Adam M-steps per iteration (kernel hyperparameters)	50 (Kato: 10)
Drift-free warm E-steps per trial	15
Held-out inference iterations (parameters frozen)	15

Appendix

Table 3: gpSLDS hyperparameters. The upper block is shared; the lower block is per dataset. Kato lists the primary configuration (9 of 12 animals; the alternative used for the other three is given in the text).

Hyperparameter	Value
Prior drift network	MLP $d\to 64\to 64\to d$ , tanh, last layer zero-initialized
Prior diffusion	constant diagonal, softplus-parameterized, initialized at 0.3
Encoder	backward GRU, hidden size 64
Marginal network $(m,\log s)$	MLP $(2\cdot 64+1)\to 64\to 64\to 2d$ , tanh
Initial prior $p(z_{0})$	learned diagonal Gaussian
Emission $C,d$	affine, initialized by PCA of the training observations

Appendix

Table 4: SDE Matching hyperparameters (identical for all datasets unless noted).

		Reconstruction		NRMSE $\downarrow$
Dataset	Method	$R^{2}$ $\uparrow$	RMSE $\downarrow$	Latent	Dynamics	Forecast
Double well	ZENO	$0.988\pm 0.004$	$0.027\pm 0.003$	$0.107\pm 0.013$	$0.342\pm 0.085$	$0.644\pm 0.163$
	gpSLDS	$0.992\pm 0.002$	$0.023\pm 0.003$	$0.090\pm 0.010$	$0.321\pm 0.188$	$0.574\pm 0.127$
	SDE Matching	$0.984\pm 0.008$	$0.030\pm 0.005$	$0.106\pm 0.024$	$0.488\pm 0.201$	$0.605\pm 0.141$
Van der Pol	ZENO	$0.978\pm 0.004$	$0.048\pm 0.002$	$0.154\pm 0.006$	$0.271\pm 0.030$	$0.655\pm 0.067$
	gpSLDS	$0.987\pm 0.002$	$0.038\pm 0.002$	$0.127\pm 0.005$	$0.192\pm 0.047$	$0.766\pm 0.026$

Appendix

Table 5: Simulated benchmarks: mean

\pm

SD over 5 seeds. Reconstruction

R^{2}

and RMSE are against the noiseless signal in standardized units; latent, dynamics and forecasting errors are normalized as defined in Section D.3 (lower is better for all but

R^{2}

). ZENO scores are filtering estimates; the baselines are smoothers.

Figure 5: Example results for a NoStim worm from the Kato Dataset. See Fig 4 A-E

			ZENO			gpSLDS			SDE Matching
Worm	$N$	$T$ (s)	$R^{2}$ $\uparrow$	Decode $\uparrow$	Forecast $\downarrow$	$R^{2}$ $\uparrow$	Decode $\uparrow$	Forecast $\downarrow$	$R^{2}$ $\uparrow$	Decode $\uparrow$	Forecast $\downarrow$
NoStim 0	109	1080	0.876	0.617	0.613	0.891	0.403	0.658	0.883	0.232	0.652
NoStim 1	135	1080	0.862	0.537	0.761	0.893	0.250	0.731	0.884	0.255	0.701
NoStim 2	131	1080	0.801	0.577	0.782	0.840	0.364	0.770	0.833	0.481	0.774
NoStim 3	125	1080	0.894	0.492	0.730	0.906	0.499	0.806	0.891	0.464	0.628
NoStim 4	129	1080	0.850	0.641	0.838	0.856	0.726	div.	0.859	0.631	0.813

Appendix

Table 6: Kato et al. (2015) per-worm results.

N

is the number of recorded neurons and

T

the recording length in seconds. For each method: reconstruction

R^{2}

(

\uparrow

), behavioral decoding balanced accuracy (

\uparrow

; chance

1/7

NoStim,

1/4

Stim) and forecast RMSE relative to persistence (

\downarrow

; below 1 beats persistence). “div.” marks a diverged forecast rollout.

Latent Dynamics Modeling

Latent dynamics modeling learns compact hidden-state representations and transition rules to predict, explain, and control evolving systems.

See all 254 papers

Deep ZakaiJ: Structured Filtering for Jump-Diffusion Time Series Forecasting
May 23, 2026Yan Leng, Thibaut Mastrolia, Hao WangStochastic Differential Equations Latent Dynamics Modeling
University of Texas at Austin · University of California, Berkeley
Wong--Zakai Convergence and FNO Learning for Stochastic Volterra Equations
Oct 13, 2026Wonjae Lee, Jiuk Jang, Hyungbin ParkNeural Operators Fourier Neural Operator
InnoCore AI Meta-Scientist Research Group, KAIST, Daejeon 34141, South Korea. · Department of Mathematical Sciences, Seoul National University, Seoul 08826, South Korea. · Department of Mathematical Sciences and Research Institute of Mathematics, Seoul National University, 1, Gwanak-ro, Gwanak-gu, Seoul 08826, South Korea.
ZNO: Stable Rational Neural Operators in the Z-Domain for Discrete-Time Dynamics
May 4, 2026Xianli Zhu, Jia YinSystem Identification Neural Operators
School of Mathematical Sciences Fudan University Shanghai, China

May 23, 2026J/K move · Enter open · S save

Deep ZakaiJ: Structured Filtering for Jump-Diffusion Time Series Forecasting

Yan Leng, Thibaut Mastrolia, Hao Wang

University of Texas at Austin · University of California, Berkeley

Stochastic Differential Equations Latent Dynamics Modeling Time Series Forecasting Probabilistic Time Series Forecasting

Open paper AskPDF arXiv

Time series driven by unobserved latent states frequently exhibit abrupt jump discontinuities whose timing and magnitude cannot be predicted from observed history alone. Classical jump-diffusion models offer a principled mathematical framework but assume rigid parametric forms, while recent neural jump models operate on fully observed trajectories without inferring the hidden states that govern the dynamics. We propose Deep ZakaiJ, a latent-state model for partially observed jump-diffusion systems that embeds the Zakai nonlinear filtering equation into a neural encoder--decoder architecture. The encoder recursively updates a belief over the latent state via Strang splitting into three interpretable substeps: prior propagation, diffusion innovation, and jump innovation, yielding a differentiable, first-order-accurate approximation of the exact filtering evolution. The decoder is a structured jump-diffusion model explicitly conditioned on the filtered belief, preserving the separation between continuous dynamics and discontinuous shocks. On synthetic, financial, and oceanographic datasets, Deep ZakaiJ improves distributional forecasts while remaining competitive in point accuracy, achieving calibrated predictive intervals and recovering interpretable latent structure in synthetic and qualitative case studies.

Abstract

Latent Dynamics Modeling

Explore similar work

Deep ZakaiJ: Structured Filtering for Jump-Diffusion Time Series Forecasting

Wong--Zakai Convergence and FNO Learning for Stochastic Volterra Equations

ZNO: Stable Rational Neural Operators in the Z-Domain for Discrete-Time Dynamics