Source-linked AI summary
Neural Jump Stochastic Differential Equations
Junteng Jia, Austin R. Benson
TL;DR
Many systems combine continuous flows with stochastic discrete jumps, but their governing dynamics and event effects are often unknown. Neural JSDEs learn both components through a piecewise-continuous latent trajectory, and experiments show strong performance across classical and real-world marked point processes.
Problem
Existing continuous models do not capture systems whose trajectories are abruptly changed by stochastic events, motivating data-driven modeling of both dynamics.
Method
Neural JSDEs use Neural ODEs for continuous latent flow and neural networks for event intensity, event embeddings, and jump effects.
Results
The model learns classical point-process intensities, matches state-of-the-art discrete event prediction, and supports events with discrete or real-valued features.
Takeaways & Limitations
Neural JSDEs provide a flexible framework for modeling event sequences alongside interpretable latent continuous dynamics.
Abstract
from arXiv · showhide
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.
1 Introduction
Neural JSDEs provide a data-driven framework for hybrid systems whose latent states flow continuously and jump at stochastic events. They extend Neural ODEs to model event effects and support temporal point-process prediction.
- Motivation: Hybrid systems combine continuous trajectories with abrupt changes caused by discrete events.Examples include user reputation evolving over time and badges altering subsequent activity, or physical trajectories changing after collisions.
- Approach: Neural JSDEs encode system state in a latent vector that flows continuously until a random event causes an abrupt trajectory jump.The continuous flow uses Neural ODEs, while event intensity and jump influence are neural functions of the latent state.
- Approach: Neural ODEs provide neural parameterizations of continuous latent dynamics and support adjoint-based gradient computation with constant memory.The adjoint method integrates an auxiliary equation backward through time.
- Contribution: Neural JSDEs extend continuous models to discrete events by handling jumps in both the latent state and the adjoint vector.The latent state also determines event conditional intensity, enabling the model to generate event sequences stochastically.
- Evaluation: The model learns classical point-process intensities and achieves state-of-the-art performance for predicting discrete event labels across synthetic and real-world datasets.Evaluations include Hawkes and self-correcting processes, Stack Overflow awards, medical records, and earthquake monitoring.
2 Background, Motivation, and Challenges
Temporal point processes represent stochastic event sequences through event histories and conditional intensities, while Neural ODEs model continuous-time dynamics. Their inability to represent abrupt event effects motivates adding jumps while retaining adjoint-based training.
- Motivation and Challenge: Neural ODEs suit real-valued timestamps but cannot model sudden event effects that immediately alter future dynamics.Time-dependent flow parameterizations make event effects begin continuously, whereas abrupt latent jumps represent immediate shocks.
- Temporal Point Processes: A temporal point process generates discrete event sequences, represented by event times and a counting function recording events before each time.The event history is denoted H = {τj}, and N(t) counts events before t.
- Temporal Point Processes: The conditional intensity λ(t) gives the probability density of an event in the next infinitesimal interval conditioned on prior events.It depends on the history Ht of events before t.
- Classical Models: Poisson processes use history-independent intensity, whereas Hawkes processes model event-triggered intensity increases that decay over time.The Hawkes formulation includes a baseline intensity, a positive excitation parameter, and a kernel such as an exponential or delaying power-law kernel.
- Classical Models: Self-correcting processes have intensity that grows exponentially with time, while events suppress future events.This model has been used for earthquake modeling after aftershocks are removed.
- Marked Processes: Marked point processes augment event times with embeddings representing discrete labels or real-valued features.Discrete types use one-hot encodings, while real-valued features use real-valued embeddings.
- Neural ODEs: Neural ODEs transform latent variables continuously with neural-network-defined dynamics and compute derivatives using the adjoint method without memory overhead.The latent trajectory can be recomputed backward together with the adjoint variables.
- Resolution: The proposed extension incorporates jumps into Neural ODEs while maintaining the simplicity of adjoint-based training.This provides a route to model abrupt event effects in hybrid systems.
3 Neural Jump Stochastic Differential Equations
Neural JSDEs represent hybrid systems with a latent state that flows continuously and jumps at stochastic events. The model learns event timing, event marks, continuous dynamics, jump effects, and their gradients with neural networks and adjoint-based training.
- Latent Dynamics and Stochastic Events: Events occur stochastically according to a conditional intensity λ(z(t)), while event embeddings are sampled from p(k|z(t)).Both the intensity and event-embedding distribution are parameterized by neural networks; real-valued marks use a Gaussian mixture model.
- Latent Dynamics and Stochastic Events: The latent dynamics combine neural ODE flow with event-driven jumps through dz(t) = f(z(t), t; θ) · dt + w(z(t), k(t), t; θ) · dN(t).The neural networks f and w control the continuous flow and jump, respectively, with left-continuous time-dependent variables.
- Latent Dynamics and Stochastic Events: The model can simulate hybrid trajectories by integrating the jump differential equation forward with an adaptive-step ODE solver, although experiments focus on prediction.The event process interrupts continuous latent evolution at stochastic timestamps.
- Learning the Hybrid System: Adjoint differentiation must account for discontinuities in both the latent state and adjoint vectors at event times.Computing the required Jacobian uses the left-limit latent state recorded during forward integration, and adjoints are lifted from right to left limits when integrating backward.
- Network Architectures: The latent state is split into internal state c(t) and event memory h(t), with constrained continuous dynamics and event-dependent memory jumps.The internal-state derivative is orthogonal to c(t), memory decays at a positive softplus-parameterized rate, and events change h(t) without directly changing c(t).
- Network Architectures: MLPs parameterize flow, jumps, total intensity, and event-mark distributions, using direct event-type intensities for discrete marks and Gaussian-mixture parameters for real-valued marks.The architecture includes CELU MLPs, element-wise operations, and a projection that orthogonalizes the internal-state dynamics against c.
4 Experimental Results
The experiments evaluate Neural JSDEs on synthetic point processes and real-world event sequences with discrete or real-valued features. Across these settings, the model learns conditional intensities, predicts event types competitively, and substantially improves synthetic feature prediction over a simple baseline.
- Experimental scope: The experiments cover classical point processes, discrete event types in Stack Overflow and medical records, and real-valued event features including earthquake locations.Synthetic datasets use simulated event sequences, while real-world tasks include awards, clinical visits, and earthquakes.
- Conditional intensity: Neural JSDE achieves lower conditional-intensity mean absolute percentage error than the RNN and other point-process baselines across the evaluated generating processes.The ground-truth model is the exception because it represents the performance expected with perfect process knowledge.
- Conditional intensity: The learned intensity tracks ground-truth trends and captures the delaying effect of a power-law kernel through interaction between event memory and internal state.The intensity peaks after the event-memory jump, when the internal state reaches its maximum.
- Discrete event types: For discrete event-type prediction, Neural JSDE achieves performance similar to RNN- and LSTM-based neural point-process models.The comparison uses Stack Overflow badge histories and MIMIC2 clinical visit histories with held-out event types.
- Real-valued features: 0.353 mean absolute error is obtained for predicting synthetic real-valued event features, versus 3.654 for predicting the running mean.The features are time intervals since the previous event in sequences generated by an exponential-kernel Hawkes process.
- Real-valued features: The earthquake experiment predicts event times and longitude-latitude features for magnitude-above-4 earthquakes in 2007–2018 from historical data spanning 1970–2006.The learned conditional-intensity contours are shown for the prediction period alongside earthquake locations.
5 Related Work
The related work contrasts classical point-process assumptions and event-only recurrent models with Neural JSDEs' unified treatment of event history and continuous latent-state evolution. It also situates the method among data-driven differential-equation learning and stochastic jump-diffusion models.
- Modeling point processes: Classical point-process models impose strong assumptions on how event history affects future dynamics, while RNNs and LSTMs model events without latent evolution between events.A prior approach used separate recurrent networks for event sequences and evenly spaced time intervals; Neural JSDEs provide a unified approach.
- Learning differential equations: Neural ODEs support continuous-time modeling but cannot represent sudden event effects, motivating their extension with discontinuities for hybrid dynamics.The paper places this extension alongside broader efforts to learn differential-equation parameters from data and stochastic jump-diffusion models.
6 Discussion
Neural Jump Stochastic Differential Equations provide a general framework for modeling temporal event sequences by learning continuous latent dynamics and abrupt event effects from data. They retain Neural ODE simplicity and memory efficiency while extending the model to flexible event types and interpretable latent dynamics.
- Neural Jump Stochastic Differential Equations learn both latent continuous dynamics and abrupt event effects from data.
- The framework maintains the simplicity and memory efficiency of Neural ODEs while using a similar adjoint method for learning.
- Neural JSDEs model intensity functions and discrete or continuous event types while providing interpretable latent-space dynamics.
A.1 Algorithm for Simulating Hybrid System with Stochastic Events
The simulation algorithm alternates adaptive ODE integration with stochastic event simulation. When an event occurs within a proposed solver step, the algorithm shortens the step, advances to the event, records it, and applies the jump.
- The algorithm initializes time, an event counter, an empty event sequence, and the initial latent state before iterating until the end time.
- At each iteration, it proposes an adaptive ODE step and simulates the next event by sampling an exponential distribution.
- When the next event occurs within the proposed step, the algorithm records the event and shrinks the step to its timestamp.
- The state is first advanced continuously to the event and then updated using the jump function before the loop continues.
A.2 Adjoint Sensitivity Analysis at Discontinuities
The adjoint analysis handles discontinuities by relating left and right limits of latent and sensitivity variables at each event. It extends the augmented Neural ODE formulation so parameter and time sensitivities also receive jump updates.
- At event timestamp τ_j, the method relates the left and right limits of latent variables across the discontinuity.
- Adjoint sensitivity variables likewise satisfy a discontinuity relation, with time-dependent variables treated as left continuous.
- The augmented latent representation includes the adjoint, parameter sensitivity, and time sensitivity variables for discontinuity handling.
- The augmented jump function at τ_j specifies how these variables are updated at an event, and the resulting limits can be verified within the Neural ODE framework.
A.3 Algorithm for Adjoint Method with Discontinuities
The discontinuous adjoint algorithm computes the loss derivatives by integrating backward through continuous intervals and applying explicit backward jump updates at recorded events. It shrinks backward steps to reach each preceding event exactly.
- The algorithm takes model parameters, time bounds, an initial state, and an event sequence as input, and returns the loss and its derivatives.
- During the forward pass, it retrieves the next recorded event and advances continuously unless the event falls within the proposed step.
- The backward pass initializes terminal adjoint quantities and integrates the state and sensitivities backward using an adaptive solver.
- When backward integration reaches an event, it shrinks the step to the event time and applies the backward jump update to the state and adjoint variables.