PaperlayerView exact revision on arXiv

Source-linked AI summary

ZENO: Nonlinear filtering with Zakai Equation Neural Operators

Aaditya Prasad, Steven W. Flavell, Brian DePasquale

arXiv:2610.32926v1cs.LGq-bio.NC

TL;DR

Numerically solving the Zakai equation for the unnormalized filtering distribution can require grid-dependent finite-difference methods that scale poorly. ZENO trains a neural operator with a Zakai-equation physics-informed loss to approximate the unnormalized filtering distribution, then uses that posterior in an expectation-maximization-like framework for state inference and parameter learning. ZENO achieves latent-inference accuracy similar to optimal nonlinear filters and parameter-learning performance comparable to state-of-the-art methods on simulated systems, while recapitulating latent structure in whole-brain C. elegans recordings.

  • Problem

    Numerically solving the Zakai equation for the unnormalized filtering distribution can require grid-dependent finite-difference methods that scale poorly.

    Source: §2.1.2 · p. 3

  • Method

    ZENO trains a neural operator with a Zakai-equation physics-informed loss to approximate the unnormalized filtering distribution, then uses that posterior in an expectation-maximization-like framework for state inference and parameter learning.

    Source: §1 · p. 2

  • Results

    ZENO achieves latent-inference accuracy similar to optimal nonlinear filters and parameter-learning performance comparable to state-of-the-art methods on simulated systems, while recapitulating latent structure in whole-brain C. elegans recordings.

    Source: §1 · p. 2

  • Takeaways & Limitations

    On whole-brain C. elegans recordings, ZENO recovers previously observed latent dynamical structure and produces behaviorally informative latents, with slightly lower reconstruction R2 than baselines.

    Source: §5 · p. 9

    ZENO does not outperform established methods in all cases, and the benchmarked methods estimate smoothing distributions that provide a superior marginal target to filtering distributions.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

The paper addresses inference of latent nonlinear dynamics from noisy observations, for which exact filtering equations become intractable in high dimensions. It introduces ZENO, a physics-informed neural operator for learning Zakai-equation solutions, and reports accurate latent inference and parameter learning across simulated systems.

1 Introduction

  • Filtering estimates p(zt|X[0:t]) from observations up to time t, whereas smoothing estimates p(zt|X[0:T]) using the full observation record.The Zakai equation governs the unnormalized filtering density; its solution gives the posterior after normalization.Source: §1 · p. 2
  • The Zakai and Kushner–Stratonovich equations provide exact posteriors, but solving them is often analytically impossible and numerically intractable for high-dimensional systems.Existing tractable alternatives make assumptions or approximations about the posterior; the Kalman–Bucy filter is exact under linear dynamics and Gaussian observation noise.Source: §1 · p. 2
  • Physics-informed neural operators use equation residuals to train across PDE instances without labeled solutions, addressing the mismatch between single-instance PINNs and filtering’s path-specific Zakai equations.This approach combines operator learning across instances with equation-based supervision when observations of the solution are unavailable.Source: §1 · p. 2
  • ZENO learns a global solution over observation histories by minimizing a physics-informed Zakai loss, while estimating dynamics and observation-model parameters in an expectation-maximization-like setup.Its loss matches posterior derivatives to the Fokker–Planck operator and incorporates likelihood updates at observation times.Source: §1 · p. 2
  • In simulated systems, ZENO achieves latent-inference accuracy similar to optimal nonlinear filters, parameter-learning performance comparable to state-of-the-art baselines, and recapitulates latent structure in a neuroscience dataset.The reported evaluations include the Double-Well, Van der Pol, and Lorenz systems, and whole-brain recordings.Source: §1 · p. 2

2 Methods

The method formulates nonlinear filtering as inference over a latent stochastic process from noisy observations, then uses the Zakai equation to evolve its unnormalized filtering density. ZENO learns this solution operator with physics-informed losses and alternates posterior inference with parameter fitting.

2.1.1 Problem Formulation

  • Filtering estimates the latent-state posterior given noisy observations up to the current time, with latent dynamics represented by an SDE and observations by a readout model.The latent state evolves continuously while observations may be discrete or continuous.Source: §2.1.1 · pp. 2–3

2.1.2 PDEs and SPDEs: Fokker-Planck, Kushner-Stratonovich, and Zakai Equations

  • The Zakai equation gives an unnormalized filtering distribution; with discrete observations it evolves by the Fokker–Planck equation between observations and receives likelihood updates at observation times.Grid-dependent finite-difference methods scale poorly for solving it.Source: §2.1.2 · p. 3

2.1.3 Physics Informed Neural Operators (PINOs)

  • A physics-informed neural operator learns the mapping from an observation path to its estimated unnormalized filtering density, using the Zakai equation residual rather than labeled posterior solutions.The method minimizes the discrepancy between posterior derivatives and the Fokker–Planck operator at collocation points.Source: §2.1.3 · pp. 3–4

2.2 Model Architecture

  • A causal GRU encodes observation history, and a DeepONet uses that context and within-interval time to produce queryable log density at latent points.The trunk network can be queried at any latent state, making the model mesh-free.Source: §2.2 · pp. 4–5

2.3 Fitting Procedure

  • ZENO alternates learning the unnormalized filtering density with dynamics and observation parameters fixed, then refitting those parameters with the density frozen.Posterior means are estimated with MALA; the readout, drift, and scalar diffusion are then fit in sequence.Source: §2.3 · p. 5

3 Results

ZENO is evaluated on simulated latent inference and dynamics learning, then on whole-brain calcium recordings from immobilized C. elegans. Results cover agreement with optimal filters, comparisons with established methods, and recovery of behaviorally relevant neural structure.

3.1 Simulated Data

  • The simulated benchmarks use three systems with known latents and observations formed by a random linear projection plus Gaussian noise.The projection matrix entries are sampled from a standard normal distribution.Source: §3.1 · p. 5

3.1.1 Latent Inference

  • ZENO reaches the optimal-filter ceiling on double well and Van Der Pol, but trails it by about 0.005 on Lorenz.Figure 2 compares the filtering distributions and normalized root mean squared error across the simulated systems.Source: Figure 2 · p. 6, §3.1.1 · p. 6

3.1.2 Dynamics Parameter Learning

  • ZENO learns the double-well cubic drift and reconstructs Van Der Pol and Lorenz dynamics, though it estimates slower edge velocities for Van Der Pol and misestimates some Lorenz flow directions and speeds.The drift and vector-field comparisons are shown in Figure 3A–C.Source: §3.1.2 · p. 6, Figure 3 · p. 7
  • For dynamics accuracy, ZENO matches gpSLDS on double well and SDE Matching on Van der Pol.The reported comparison is against the established baselines shown in Figure 3D.Source: §3.1.2 · pp. 6–7, Figure 3 · p. 7

3.2 Analysis of whole-brain calcium dynamics in C. elegans

  • On immobilized C. elegans whole-brain calcium data, ZENO recovers reported limit cycles and behaviorally relevant latents, with comparable reconstruction R2 and forecasting RMSE to established methods.The study also reports accurate reconstruction of neural observations for an example worm.Source: §3.2 · p. 7

4 Related Work

Related work spans switching linear dynamical systems, latent SDE methods, and physics-informed PDE solvers. The paper positions ZENO as jointly learning a full filtering solution operator and unknown dynamics.

Switching Linear Dynamical Systems

  • Switching linear dynamical systems approximate nonlinear dynamics with mixtures of linear systems, balancing interpretability and tractability against expressivity; their variational posteriors may diverge from the true posterior, especially for multimodal or asymmetric distributions.The paper contrasts these approximations with ZENO’s direct approximation of the posterior.Source: §4 · pp. 7–8

Latent SDEs

  • gpSLDS forecasting results were inconsistent, likely because the hyperparameter tuning did not accurately reflect the method.The authors qualify this as a likely explanation.Source: §4 · p. 7
  • Latent SDE methods use neural prior and posterior processes, but ELBO balancing can cause posterior collapse and simulation-based training can be costly; ZENO instead uses a physics-based loss rather than reconstruction/ELBO.The related-work discussion also describes simulation-free SDE-Matching and Neural Stochastic Flows.Source: §4 · pp. 8–9

Physics-informed Neural Networks for the Fokker-Planck and Zakai Equations

  • Earlier Fokker–Planck and Zakai approaches either use grid-based one-dimensional inverse fitting with noisy PDFs, assume known drift and diffusion, or propagate solutions stepwise; ZENO claims joint full-operator inference and unknown-dynamics learning with EM.The paper characterizes this combination as, to its knowledge, the first of its kind.Source: §4 · p. 9

5 Discussion

ZENO performs competitively on synthetic inference and parameter learning, and on C. elegans recordings yields behaviorally informative latents alongside comparable forecasting but slightly lower reconstruction R2. The discussion situates this approach among operator-learning and physics-informed methods and identifies limits and possible extensions.

5 Discussion

  • On C. elegans recordings, ZENO was comparable to baselines on forecasting and produced more behaviorally informative latents, with slightly lower reconstruction R2.The authors expect relaxing the Stiefel constraint on the observation model could address the reconstruction trade-off.Source: §5 · p. 9
  • ZENO matched or came close to optimal numerical filters for synthetic latent inference and performed on par with state-of-the-art methods for parameter learning.The authors note it does not outperform established methods in every case; those methods estimate smoothing distributions using observations both forward and backward in time.Source: §5 · p. 9

AI use statement

  • The authors describe generative-AI use for synthetic data, methodology feedback, implementation, figure modification, experimental parameters, code, and brainstorming, with review of AI-assisted work.They report not using these tools to propose or refine hypotheses, analyze qualitative data, interpret results, or formulate mathematical claims.Source: AI use statement · pp. 9–10

Reproducibility statement

A.1 Formulation

  • Operator learning approximates a map between function spaces, with training pairs of input and output functions used to fit a parametric operator.A useful architecture should remain meaningful across discretizations, since functions are represented computationally through finite samples and high-resolution data are expensive.Source: §A.1 · p. 15

A.2 Deep Operator Networks (DeepONet)

  • DeepONet combines input-dependent branch coefficients with trunk basis functions to evaluate an output function at query locations.Its outputs can be queried off-mesh, but its branch network expects inputs at fixed sensor locations.Source: §A.2 · p. 15

B Primer on Physics Informed Neural Networks (PINNs)

  • PINNs represent differential-equation solutions with neural networks and train them to satisfy the governing equation, using automatic differentiation to evaluate derivatives.For the stated general formulation, the operator Aλ includes parameters and B is a boundary operator.Source: B Primer on Physics Informed Neural… · p. 15, §B.1 · p. 16

B.2 Training

The training discussion presents losses as penalties that constrain a network where data are absent. ZENO alternates between fitting its filtering density and refitting dynamical parameters.

B.2 Training

  • PINN training minimizes a composite loss over sampled domain points, with terms for each information source and the differential equation acting as a soft constraint where data are absent.Source: §B.2 · p. 16

C ZENO Fitting Procedure

  • ZENO alternates between learning the unnormalized log-density with dynamical and observation parameters fixed, and refitting those parameters while the density is frozen.The procedure follows an expectation-maximization framework; the observation mask indicates whether each observation is present.Source: C ZENO Fitting Procedure · p. 16

C.1 E step

The E step fits ZENO’s unnormalized log-density using Zakai residual and observation-update losses, then estimates posterior means for parameter refitting. The accompanying procedure also specifies the M step, initialization, and experimental software setup.

C.1 E step

  • At each site, collocation points are split between a component near the data-driven latent estimate and a broad defensive component; at unobserved frames the first component is also broad and origin-centered.Source: §C.1 · p. 16
  • The residual loss applies the log-space Fokker–Planck equation over normalized observation intervals, with a scale-normalized residual and a stop-gradient denominator to avoid reducing loss by inflating its magnitude.The filtering marginal is estimated with self-normalized importance sampling, and posterior means are then estimated by MALA sampling.Source: §C.1 · pp. 17–18
  • The E-step objective combines a weighted Zakai residual with jump and initial-condition losses, balancing posterior spreading against sharpening from measurement updates.It uses repeated gradient updates with resampled collocation points and residual intervals, accumulating gradients over trajectory chunks to bound memory.Source: §C.1 · pp. 16–17

C.2 M step

  • In the M step, ZENO refits the observation map and, once its latent coordinates stabilize, refits drift, diffusion, and observation noise from posterior means and observation data.The drift uses a second-order Runge–Kutta midpoint to remove the O(∆t) regression bias described for forward increments.Source: §C.2 · pp. 18–19

C.3 Initialization

  • Initialization uses a future-block Hankel matrix and its leading singular vectors to initialize the observation map, then fits linear dynamics before warming-starting the drift network.A bootstrap M step uses projected data-driven latent estimates in place of posterior means before operator training.Source: §C.3 · pp. 19–20

D Experimental Details

  • Experiments were trained on one GPU, with a new simulated dataset per seed and the same dataset shared across methods using that seed.The reported software stack includes Python and JAX; ZENO uses Equinox and Optax.Source: D Experimental Details · p. 20

D.1 Datasets

The simulated benchmarks use three latent SDE systems with noisy observations generated under system-specific settings. The C. elegans dataset comprises recordings from immobilized worms, with behavioral annotations kept out of model fitting and selection.

D.1 Datasets

  • The simulated benchmarks use double-well, Van der Pol, and Lorenz latent SDEs, following prior setups and using additive diffusion for Lorenz.Trajectories are simulated by Euler–Maruyama with steps at most 1 ms and observations at regular-grid times.Source: §D.1 · p. 20
  • Double-well and Van der Pol observations use random Gaussian sensors, while Lorenz observations directly measure the latent state and add noise after clean-signal standardization.Train, validation, and test trajectories are independent draws sharing the sensor; observations are standardized using training data only.Source: §D.1 · p. 20, Table 1 · p. 21
  • The whole-brain C. elegans dataset includes 12 immobilized worms, with spontaneous recordings and oxygen-stimulus recordings treated as an unconditional benchmark.Recordings use corrected fluorescence at native frame rate without smoothing, clipping, or derivative preprocessing; models are fit and scored on every frame, so scores are in-sample.Source: §D.1 · p. 20
  • Human-annotated behavioral states accompany the worm recordings but are never used to fit or select a dynamical model.Source: §D.1 · p. 20

D.2 Training

Training uses a shared evaluation protocol, with method-specific procedures for ZENO, gpSLDS, and SDE Matching. On Kato, gpSLDS configurations were selected per animal, unlike the fixed configurations used by the other methods.

D.2 Training

  • All methods use the same standardized data, observation grid, dimensions, and validation cadence; simulated benchmarks use six seeds, while Kato fits one model per worm.Checkpoints are selected without ground truth, and test splits are scored once after selection.Source: §D.2 · p. 21

D.2.1 ZENO

  • ZENO alternates E-steps with M-steps every 2,000 E-steps, refitting the sensor, drift, diffusion, and observation noise from posterior summaries and increments.Its operator is trained against the Fokker–Planck residual at sampled time and state collocation points.Source: §D.2.1 · p. 21

D.2.2 gpSLDS

  • gpSLDS uses SING with a smoothly switching linear kernel; on Kato, three animals used an alternate configuration selected by lower in-sample forecast RMSE.Thus, gpSLDS uses a per-animal best-of configuration on Kato, while ZENO and SDE Matching use one fixed configuration for all worms.Source: §D.2.2 · pp. 21–22

D.2.3 SDE Matching

  • SDE Matching uses a backward-GRU smoother to produce Gaussian posterior marginals and trains with an initial KL, path energy, and observation likelihood, without solving an SDE during training.Its forecasts propagate the learned prior from the posterior at the end of the observed prefix.Source: §D.2.3 · p. 22

D.3.1 Simulated Data

The simulated-data evaluation defines observation reconstruction, aligned latent and drift errors, and prefix-based forecasting. Latent comparisons account for affine coordinate ambiguity without using test data to fit the alignment.

D.3.1 Simulated Data

  • Latent states are identifiable only up to an affine transformation, so posterior means are aligned to true latents on validation data after checkpoint selection, then frozen for test scoring.Observation-space scores do not use latent alignment.Source: §D.3.1 · p. 22
  • Reconstruction R2 and RMSE compare the posterior-mean reconstruction with the noiseless signal in standardized units, with channel means computed over the test split.The observation-space scores do not use latent alignment.Source: §D.3.1 · p. 24
  • Latent NRMSE uses the aligned posterior mean and scales each latent coordinate by its standard deviation over training trajectories.Dynamics NRMSE compares learned and true drift at ground-truth test states, normalized by the RMS true drift on training states.Source: §D.3.1 · p. 24
  • Forecasting NRMSE hides the second half of each test trajectory, infers from the first half, and scores the mean of 128 prior-SDE forecast paths against the clean hidden signal.The RMSE is normalized by the pooled standard deviation of the per-channel-centered clean test signal.Source: §D.3.1 · p. 24

D.3.2 Kato Dataset

The Kato evaluation uses observation-space and behavioral scores because latent states, clean signals, and drift are unknown. Forecasting is tested against persistence using held-out frames from each window.

D.3.2 Kato Dataset

  • Kato scores are observation-space and behavioral because the dataset has no ground-truth latent states, clean signal, or drift.Source: §D.3.2 · p. 24
  • Reconstruction R2 pools the comparison between reconstructed fluorescence and observed standardized fluorescence across neurons.Source: §D.3.2 · p. 24
  • Forecasting hides the second 100 frames of each 200-frame window and compares the mean forecast with persistence; ratios below 1 beat that worm's persistence floor.The prior is propagated across the hidden half using 128 samples.Source: §D.3.2 · p. 24
  • Behavioral decoding fits a class-balanced multinomial logistic regression from posterior-mean latents to annotated states and reports held-out balanced accuracy.Chance is 1/7 for NoStim and 1/4 for Stim worms.Source: §D.3.2 · p. 24

D.4 Additional Results

The additional results compare simulated benchmarks and report per-worm Kato reconstruction, decoding, and forecast scores. The Kato results include divergent gpSLDS forecasts, while the simulated benchmark table summarizes five-seed results after excluding each method-dataset pair's worst dynamics-NRMSE seed.

D.4 Additional Results

  • The simulated benchmark reports mean ± SD over five seeds after dropping the worst dynamics-NRMSE seed for each method-dataset pair; ZENO filters, while the baselines smooth.Forecasting is the one metric that is prefix-only for every method.Source: §D.4 · pp. 24–25, Table 5 · p. 25
  • Figure 5 presents example results for a NoStim worm from the Kato dataset and points readers to Fig. 4A–E.Source: Figure 5 · p. 25
  • Mean reconstruction R2: 0.830 (ZENO), 0.857 (gpSLDS), 0.858 (SDE Matching); decoding balanced accuracy: 0.624, 0.524, 0.482, respectively; forecast ratio: 0.78, 1.95, 0.71, respectively.The gpSLDS forecast mean excludes three diverged animal rollouts.Source: §D.4 · p. 25, Table 6 · p. 26
Loading 2610.32926v1…