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

May 23, 2026

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

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.

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

Ask PDF

Oct 13, 2026

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

This paper studies Wong--Zakai convergence and Fourier neural operator learning for stochastic Volterra integral equations (SVIEs) describing random systems with memory. A Wong--Zakai theorem establishes convergence of pathwise solutions driven by piecewise linear approximations of Brownian motion to the corresponding SVIE solution. This convergence provides a foundation for using mirror-padded Fourier neural operators to learn the solution operator mapping initial conditions and driving noise to entire solution trajectories. Numerical experiments demonstrate competitive prediction accuracy and generalization to temporal grids finer than those used for training, while assessing the influence of padding and computational cost. As a follow-up to \citet{LKP:26}, this study extends their operator-learning framework to SVIEs through separate convergence and approximation analyses tailored to convolutional memory.

Wonjae Lee, Jiuk Jang, Hyungbin Park

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.

Neural Operators Fourier Neural Operator Stochastic Differential Equations Neural Network Approximation Theory

Ask PDF

May 4, 2026

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

We introduce the Z-Domain Neural Operator (ZNO), a causal neural operator whose layers are stable low-rank multiple-input multiple-output (MIMO) rational filters parameterized directly in the

z

-plane. ZNO addresses a limitation of existing operator learning methods, many of which are primarily tailored for continuous-time problems, while a large class of system-identification problems is intrinsically discrete-time. The

z

-domain form expresses stability as a unit-disk pole constraint and makes learned discrete-time poles directly readable. The model combines low-rank channel mixing, smooth stable pole reparameterization, causal recurrence, and an optional short finite impulse response (FIR) branch in a single

z

-domain rational recurrent layer. Across controlled discrete system-identification experiments, ZNO's advantage is most evident when the target dynamics are stable rational systems with lightly damped poles near the unit circle. Under matched parameter budgets, ZNO is not uniformly dominant; however, with validation-selected configurations, the same architecture can achieve the lowest mean error across the controlled tasks. A five-bin difficulty sweep over near-unit-circle / long-memory dynamics shows that ZNO has the lowest mean error across memory regimes, from short (approximately 10 steps) to long (approximately 100-200 steps). On five public nonlinear system-identification benchmarks, ZNO is competitive with neural operator and state-space baselines, achieving the lowest mean error on benchmarks whose dynamics align with stable rational discrete-time filters, while classical or state-space baselines remain preferable on some systems. These results position ZNO as a strong model for stable rational discrete-time dynamics, especially in near-unit-circle and long-memory regimes, but not as a universal replacement for specialized system-identification methods.

Xianli Zhu, Jia Yin

School of Mathematical Sciences Fudan University Shanghai, China

System Identification Neural Operators

Ask PDF

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