🤖 machine learning

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

The paper proposes Deep ZakaiJ, a neural encoder-decoder framework that embeds the Zakai nonlinear filtering equation to effectively forecast partially observed jump-diffusion time series by recursively updating latent state beliefs through interpretable propagation and innovation steps.

Original authors: Yan Leng, Thibaut Mastrolia, Hao Wang

Published 2026-05-26

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yan Leng, Thibaut Mastrolia, Hao Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict the weather, but the atmosphere is chaotic. Sometimes the wind blows steadily, but occasionally, a sudden, violent storm hits out of nowhere. In the world of data, this is called a jump-diffusion system: a mix of smooth, predictable movement ("diffusion") and sudden, shocking jumps ("jumps").

The problem is that we can't see the "engine" driving these changes. We only see the weather (the data), not the hidden atmospheric pressure systems (the latent states) causing the storms. If you try to predict the future based only on what you see, you'll be caught off guard by the next sudden jump.

Enter Deep ZakaiJ, a new AI model designed to be a "super-forecasting detective" for these chaotic systems. Here is how it works, broken down into simple concepts:

1. The Detective's Notebook (The Encoder)

Most AI models just look at the past data and guess the future. Deep ZakaiJ is different. It maintains a "belief" about what is happening behind the scenes.

Think of this belief as a detective's notebook. Every time a new piece of data arrives (like a temperature reading), the detective updates their notebook. But instead of just scribbling notes, Deep ZakaiJ uses a very specific, mathematically rigorous method called the Zakai equation.

To make this complex math work inside a neural network, the authors broke the update process into three distinct, easy-to-understand steps (like a recipe):

Step A: The Guess (Prior Propagation): The detective makes a prediction about what the hidden state should be doing next, based on how it usually behaves.
Step B: The Smooth Check (Diffusion Innovation): The detective looks at the data. If the weather changed smoothly (like a gentle breeze), they adjust their belief slightly to match this smooth movement.
Step C: The Shock Check (Jump Innovation): This is the special part. If the data shows a massive, sudden spike (like a hurricane), the detective doesn't just tweak the numbers; they radically rewrite their belief to account for the possibility that a "jump" just happened.

By separating these steps, the model can react sharply to sudden shocks without getting confused by normal fluctuations.

2. The Crystal Ball (The Decoder)

Once the detective has updated their notebook (the "filtered belief"), they pass this information to the Decoder.

Think of the Decoder as a crystal ball that is now informed. Because it knows the hidden state (e.g., "a storm is brewing"), it doesn't just guess a single number for the future. Instead, it generates a full probability map. It can say, "There's a 90% chance the price stays steady, but a 10% chance it crashes."

Crucially, this crystal ball is "structured." It doesn't just throw darts at a board; it respects the rules of physics and finance. It knows that smooth changes and sudden jumps are different things, and it models them separately.

3. The Results: Why It Matters

The authors tested Deep ZakaiJ on three types of "weather":

Synthetic Data: A made-up world where they knew the exact rules. Deep ZakaiJ figured out the hidden rules perfectly.
Gold Prices (XAU/USD): Real financial data where prices jump due to news. The model was better at predicting the range of possible outcomes (uncertainty) than other top models, even if its single-point guess was similar to others.
Ocean Waves (NDBC): Real data from buoys where waves can suddenly spike during storms. Again, the model provided much more reliable "confidence intervals" (predictive ranges) than competitors.

The Bottom Line

Deep ZakaiJ is like giving a weather forecaster a pair of X-ray glasses. While other models only see the surface data (the rain), Deep ZakaiJ infers the hidden storm systems causing the rain. By mathematically separating "smooth changes" from "sudden shocks," it builds a more accurate and reliable picture of the future, especially when things are about to go wrong.

What the paper doesn't claim:
The paper does not claim this model can predict stock market crashes with 100% certainty, nor does it claim to solve climate change. It simply states that for time series with hidden states and sudden jumps, this specific mathematical approach leads to better, more calibrated predictions than current methods.

Technical Summary: Deep ZakaiJ

Problem Statement

Time series data in domains such as finance and oceanography frequently exhibit abrupt jump discontinuities driven by unobserved latent states (e.g., macroeconomic shocks or storm regimes). Predicting these trajectories is challenging because the timing and magnitude of jumps cannot be inferred solely from observed history.

Existing approaches face significant limitations:

Classical Jump-Diffusion Models: While mathematically principled (e.g., Merton's models), they rely on rigid parametric assumptions and struggle to capture complex, non-linear latent structures.
Neural Jump Models: Recent deep learning approaches (e.g., Neural Jump SDEs) often operate on fully observed trajectories, failing to infer the hidden states that govern the dynamics.
Latent-State Models: Methods like Latent ODEs or Neural CDEs treat latent states as implicit representations learned via amortized variational inference. They lack a principled Bayesian update mechanism for handling the totally inaccessible stopping times characteristic of jump processes.
Kalman Filtering Extensions: Models like CRU and NCDSSM embed Kalman filtering into deep architectures but are restricted to linear-Gaussian dynamics, which cannot model the non-linearities inherent in jump-diffusion systems.

There is a need for a framework that performs nonlinear filtering over latent regimes in partially observed jump-diffusion systems while remaining differentiable and trainable.

Methodology: Deep ZakaiJ

The authors propose Deep ZakaiJ, a latent-state model that embeds the Zakai nonlinear filtering equation into a neural encoder-decoder architecture. The framework is designed for partially observed systems where an observed process $X_t$ is driven by an unobserved latent state $\Theta_t$ .

1. Theoretical Foundation

The model is built upon the Zakai equation, which describes the evolution of the unnormalized conditional law (filter) of the latent state given observations. By changing the measure to a reference measure $Q$ where the observation process is driftless, the authors derive a weak form of the Zakai equation (Equation 3) that evolves the unnormalized density $q_t(\phi)$ .

2. The Zakai Encoder

The core innovation is a Strang-splitting discretization of the Zakai equation, decomposing the belief update into three interpretable, differentiable substeps:

A-step (Prior Propagation): Propagates the latent prior using a numerical propagator (Markov transition kernel) derived from the latent generator. A learned residual correction $R^A_\psi$ is added to account for model mismatch and numerical bias.
B-halfstep (Diffusion Innovation): Updates the belief based on the continuous component of the observation. It performs a Bayes reweighting using a local Gaussian approximation of the continuous increment.
C-halfstep (Jump Innovation): Updates the belief based on the jump component. It employs an "at-most-one-jump" approximation valid for small time steps, reweighting the posterior against a mixture of "no-jump" and "jump-occurred" hypotheses. This allows the filter to react sharply to large dislocations.

The encoder recursively updates the unnormalized posterior $q_k$ over a fixed latent grid, extracting a low-dimensional belief feature $\beta_k$ for forecasting. The authors prove that this discrete encoder achieves first-order global accuracy relative to the exact Zakai evolution.

3. Structured Jump-Diffusion Decoder

The decoder is a structured model that parameterizes the conditional distribution of future increments given the current observation, the filtered belief $\beta_k$ , and a candidate latent state $\theta$ .

It explicitly disentangles continuous dynamics (drift $\mu$ , volatility $\sigma$ ) from discontinuous shocks (jump intensity $\lambda$ , jump-size map $\gamma$ ).
Neural networks condition these coefficients on the observation and the belief feature.
Small-Jump Approximation: To handle infinite-activity jumps within a discrete architecture, the decoder absorbs micro-jumps (magnitude $<\epsilon$ ) into effective continuous coefficients, treating only macroscopic jumps explicitly in the filtering step.

4. Joint Training

The model is trained end-to-end by maximizing a stepwise objective that combines:

Forecasting Likelihood: The log-likelihood of future increments under the learned jump-diffusion law.
Filtering Regularization: A KL-divergence term encouraging the filtered posterior to remain close to the prior (preventing collapse).

Gradients are backpropagated through the Strang splitting operations and normalization steps.

Key Contributions

Zakai Encoder: A discretization of the Zakai equation via Strang splitting into three interpretable steps (prior, diffusion innovation, jump innovation). It provides a differentiable, first-order-accurate approximation of the exact filtering evolution for partially observed jump-diffusion systems.
Structured Decoder: A jump-diffusion model explicitly conditioned on the filtered latent belief, preserving the separation between continuous and discontinuous dynamics and learning a predictable compensator.
Empirical Validation: Demonstrated improvements in distributional forecast quality (calibration and sharpness) across synthetic, financial, and oceanographic datasets, while maintaining competitive point accuracy.

Experimental Results

The authors evaluated Deep ZakaiJ on:

Synthetic Data: A partially observed jump-diffusion process with known ground truth.
Financial Data: XAU/USD (Gold vs. USD) 10-minute closing prices.
Oceanographic Data: NOAA/NDBC wave height (WVHT) observations.

Performance:

Distributional Quality: Deep ZakaiJ consistently outperformed baselines (including Latent SDEs, Neural Jump SDEs, and foundation models like Chronos) in terms of CRPS (Continuous Ranked Probability Score) and Log-Likelihood.
Calibration: The model achieved well-calibrated predictive intervals, with 90% coverage (Cov90) close to the nominal level (e.g., 89.0% on synthetic, 92.9% on XAU/USD, 85.3% on NDBC). In contrast, many baselines showed significant under- or over-coverage.
Point Accuracy: The model remained competitive in MAE and RMSE, often matching or slightly improving upon state-of-the-art point predictors.
Latent Recovery: In synthetic and qualitative case studies, the inferred filter successfully tracked latent state changes and reacted to regime shifts (e.g., trend reversals in gold prices), demonstrating the extraction of interpretable latent structure.

Significance and Claims

The paper claims that Deep ZakaiJ bridges the gap between principled stochastic filtering and modern deep learning. By embedding the Zakai equation directly into a neural architecture, the model:

Enables structured inference for partially observed systems with jumps, a capability missing in current neural jump models.
Provides interpretable latent representations that separate continuous dynamics from discrete shocks.
Delivers superior uncertainty quantification, producing calibrated predictive intervals even under structural shifts where other models fail.

The authors position this work as a step toward more robust time series forecasting in environments characterized by latent regimes and abrupt discontinuities, such as financial risk management and environmental early warning systems. They acknowledge limitations, noting the current formulation is univariate and relies on a fixed latent grid, suggesting future work in high-dimensional and adaptive representations.