[Paper Review] Neural Jump Stochastic Differential Equations

Junteng Jia, Austin R. Benson|arXiv (Cornell University)|May 24, 2019

Model Reduction and Neural Networks48 references76 citations

TL;DR

Neural Jump Stochastic Differential Equations (JSDEs) learn both continuous latent dynamics and abrupt event-driven jumps in hybrid systems, enabling modeling of continuous flows plus discrete events in temporal point processes. The approach extends Neural ODEs with stochastic jumps and applies to various real-world datasets.

ABSTRACT

Many time series are effectively generated by a combination of deterministic continuous flows along with discrete jumps sparked by stochastic events. However, we usually do not have the equation of motion describing the flows, or how they are affected by jumps. To this end, we introduce Neural Jump Stochastic Differential Equations that provide a data-driven approach to learn continuous and discrete dynamic behavior, i.e., hybrid systems that both flow and jump. Our approach extends the framework of Neural Ordinary Differential Equations with a stochastic process term that models discrete events. We then model temporal point processes with a piecewise-continuous latent trajectory, where the discontinuities are caused by stochastic events whose conditional intensity depends on the latent state. We demonstrate the predictive capabilities of our model on a range of synthetic and real-world marked point process datasets, including classical point processes (such as Hawkes processes), awards on Stack Overflow, medical records, and earthquake monitoring.

Motivation & Objective

• Motivate and model systems that evolve continuously but are interrupted by stochastic events.
• Propose a data-driven framework that jointly learns continuous latent dynamics and jump events via neural networks.
• Demonstrate the model on classical point processes and real-world datasets with discrete and real-valued event features.

Proposed method

• Represent the system with a latent state z(t) that flows continuously via Neural ODE dynamics f(z(t), t; θ).
• Introduce a stochastic point process with intensity λ(z(t)) governing event arrivals and a jump function w(z(t), k(t), t; θ) that updates the latent state at events.
• Model event types with p(k|z(t)) (discrete types via probabilities or real-valued features via Gaussian mixtures).
• Train using a likelihood objective that includes the log-intensity terms and an integral of λ(z(t)) over time, using adjoint methods adapted for discontinuities.
• Architect latents with c(t) (internal state) and h(t) (event memory) and parameterize f, w, λ, p with neural networks.
• Handle jumps in the adjoint equations at event times via derived jump conditions to enable backpropagation.

Experimental results

Research questions

• RQ1 Can Neural JSDEs accurately recover the conditional intensity functions of classical point processes (Poisson, Hawkes with exponential and power-law kernels, self-correcting)?
• RQ2 How well do Neural JSDEs perform on discrete event type prediction on social and medical datasets, compared to state-of-the-art models?
• RQ3 Can Neural JSDEs handle real-valued event features (e.g., earthquake locations) and provide predictive distributions for event attributes?

Key findings

• Neural JSDEs achieve better mean absolute percentage error in learned intensities than RNN baselines and traditional point process models across Poisson, Hawkes (Exponential), Hawkes (Power-Law), and Self-Correcting processes.
• On Stack Overflow and MIMIC2 datasets, Neural JSDEs attain competitive or superior discrete event type prediction accuracy compared to baselines.
• For synthetic real-valued feature data (e.g., inter-event times) and earthquake location data, Neural JSDEs accurately predict event features and outputs probabilistic embeddings.
• The model demonstrates interpretable latent dynamics with a continuous flow interleaved with event-driven jumps, and can model marked point processes with both discrete and real-valued event information.