Scan-Parallel Variational Inference for event-driven Latent Jump Models

Vansh Bansal; Min Chen; James G. Scott

stat.MLOct 9, 2026

Scan-Parallel Variational Inference for event-driven Latent Jump Models

Authors: Vansh Bansal, Min Chen, James G. Scott

Organizations: Department of Statistics and Data Science UT Austin

Abstract

Continuous-time latent state-space models are widely used for irregularly observed time series, typically representing uncertainty through a random initial state followed by smooth ordinary differential equation (ODE) dynamics, or Brownian perturbations accumulated by a stochastic differential equation (SDE). While these constructions are natural for smooth or diffusive systems, they do not directly encode a finite set of unobserved shocks whose timing, magnitude, and effect on memory are themselves uncertain. We instead introduce latent jump models (LJMs): the latent state follows closed-form deterministic dynamics between Poisson-driven event times and undergoes a learned instantaneous jump at each event, with state-dependent intensity and mark distributions. This defines a time-homogeneous Markov process whose jumps are intrinsic state transitions rather than numerical integration steps. For training, we derive a tractable event-augmented path-space evidence lower bound (ELBO) and introduce Scan-Parallel Variational Inference (SPVI): an observation-conditioned variational posterior samples the complete latent event configuration before the state recurrence is evaluated using triangular affine scans. This enables fast and scalable training for LJMs in logarithmic temporal depth. Across synthetic stochastic systems and real irregular time series, LJMs remain competitive with continuous-time latent baselines in forecast accuracy, often with fewer latent updates than numerical solver evaluations, while training one to two orders of magnitude faster than the sequential baselines.

Figures & tables

Figure 1: Comparison across four event-driven and smooth systems. All models use identical history-encoder and emission architectures. Solver step sizes are coarsened for baselines at evaluation to approach LJM event budget; NFE labels show the number of dynamics evaluations. LJM shows better reconstruction and forecast quality at fewer NFE. Latent SDE diverges on the regular PDMP, while Latent ODE and Neural Stochastic Flow (solver-free) exhibit errors in several examples.

Figure 2: (a) Generative prior : the current latent state determines the intensity and mark distribution of subsequent events, requiring chronological simulation. (b) SPVI posterior : observations determine the initial-state distribution and event laws independently of the sampled latent state. The complete event configuration can therefore be sampled before constructing the trajectory through successive blockwise scans. Both processes share the closed-form flow

\Phi_{\delta}

and jump map

F_{\theta}

System	CRPS $\downarrow$	RMSE $\downarrow$	NFE $\downarrow$	Method	CRPS $\downarrow$	RMSE $\downarrow$	NFE $\times$	Train time $\times$
	LJM (ours)			Best baseline at a similar budget
Cyclic jumps	$\underset{\text{\tiny$ (0.0032) $}}{\mathbf{0.0679}}$	$\underset{\text{\tiny$ (0.0036) $}}{\mathbf{0.1261}}$	89	Latent ODE	$\underset{\text{\tiny$ (0.0010) $}}{0.0729}$	$\underset{\text{\tiny$ (0.0015) $}}{0.1276}$	2.0 $\times$	7 $\times$
Irregular PDMP	$\underset{\text{\tiny$ (0.0012) $}}{\mathbf{0.2244}}$	$\underset{\text{\tiny$ (0.0017) $}}{\mathbf{0.5210}}$	19	Latent SDE	$\underset{\text{\tiny$ (0.0036) $}}{0.2410}$	$\underset{\text{\tiny$ (0.0118) $}}{0.5467}$	2.6 $\times$	10 $\times$
Regular PDMP	$\underset{\text{\tiny$ (0.0009) $}}{\mathbf{0.2509}}$	$\underset{\text{\tiny$ (0.0022) $}}{0.5995}$	19	SDE Matching	$\underset{\text{\tiny$ (0.0009) $}}{0.2567}$	$\underset{\text{\tiny$ (0.0041) $}}{\mathbf{0.5945}}$	2.3 $\times$	1 $\times$
Damped rotation	$\underset{\text{\tiny$ (0.0056) $}}{0.1012}$	$\underset{\text{\tiny$ (0.0225) $}}{0.1693}$	205	NCDSSM-NL	$\underset{\text{\tiny$ (0.0017) $}}{\mathbf{0.0743}}$	$\underset{\text{\tiny$ (0.0027) $}}{\mathbf{0.1274}}$	1.7 $\times$	86 $\times$

Table 1: Synthetic forecasting benchmarks. The reported numbers are mean and range over three independent runs of the whole study. NFE is dynamics evaluations per context trajectory; in the columns marked

\times

, a value

K\times

means LJM is

K

times cheaper than the baseline.

Context	Method	CRPS $/s_{h}\downarrow$	RMSE (W) $\downarrow$	$\text{E-CRPS}_{\text{ \tiny(kWh)}}$ $\downarrow$	90% cov.	Width (W)
512	LJM (ours)	0.0495 ±0.0004	551.7 ±2.9	0.0630 ±0.0005	0.925 ±0.005	858 ±33
	SDE Matching	0.0544 ±0.0011	566.9 ±1.2	0.0705 ±0.0007	0.876 ±0.017	804 ±120
	Latent NSF †	0.0588 ±0.0018	601.5 ±21.3	0.0779 ±0.0031	0.877 ±0.012	918 ±322
1024	LJM (ours)	0.0492 ±0.0004	551.4 ±0.2	0.0628 ±0.0005	0.926 ±0.008	926 ±40
	SDE Matching	0.0568 ±0.0006	573.7 ±2.0	0.0730 ±0.0010	0.883 ±0.003	937 ±161
	Latent NSF †	0.0599 ±0.0015	589.3 ±9.2	0.0794 ±0.0030	0.872 ±0.013	740 ±135

Table 2: REFIT aggregate load forecasting using 256 predictive trajectories per test window. Context is the number of observed points. CRPS

/s_{h}

normalises by the household scale. Energy CRPS scores total consumption over the 30-minute forecast horizon.

Model	Forecast MSE $\downarrow$	CRPS $\downarrow$	90% coverage	NFE $\downarrow$	Train time/fold $\downarrow$
LJM (ours)	$\underline{0.360}_{\pm 0.069}$	$\mathbf{0.258}_{\pm 0.024}$	$0.837_{\pm 0.030}$	461	7.5 min
NCDSSM-NL	$\mathbf{0.352}_{\pm 0.055}$	$0.305_{\pm 0.013}$	$0.968_{\pm 0.006}$	6740	3.6 h
Latent SDE	$0.414_{\pm 0.078}$	$0.306_{\pm 0.040}$	$0.592_{\pm 0.127}$	1685	21 h
Latent ODE	$0.451_{\pm 0.107}$	$0.371_{\pm 0.042}$	$0.284_{\pm 0.033}$	6740	11.5 h
SDE Matching	$0.433_{\pm 0.106}$	$\underline{0.259}_{\pm 0.030}$	$0.958_{\pm 0.018}$	1685	33 min
Latent NSF	$0.829_{\pm 0.162}$	$0.477_{\pm 0.058}$	$0.605_{\pm 0.049}$	$\mathbf{422}$	19 min

Table 3: USHCN test forecasts: The numbers show mean

\pm

sd over five station folds with 128 predictive paths.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Value
Latent dimension / ordered block sizes	$8$ / $(4,4)$
Mark dimension / observation dimension	$2$ / $2$
Training batch size / observations per window	$512$ / $52$
Context / forecast observations	$44$ / $8$
Predictive samples / evaluation batch size	$128$ / $16$
Posterior rate range	$(10^{-3},10)$

Appendix

Table 4: Synthetic implementation settings. The KL weight takes values in the listed set; the event-regularization weight is swept separately.

Dataset	Event times	Jump / innovation	$\sigma$	Grid	Per-traj. randomness
Deterministic rotation	none	none	$0.10$	regular	$\theta$ , $z_{0}$
Damped rotation	none	Gaussian, $\varsigma{=}0.02$	$0.10$	regular	$\alpha,\omega$ , $z_{0}$
Cyclic jumps	every $5$ steps	$4$ -cycle, fixed dirs.	$0.10$	regular	$a_{1},a_{2}$ , $c_{1},c_{2}$ , $z_{0}$
Cyclic jumps, fixed decay	every $5$ steps	$4$ -cycle, fixed dirs.	$0.10$	regular	$c_{1},c_{2}$ , $z_{0}$
Regular PDMP	Poisson, $\lambda{=}0.3$	$W_{1}\tanh(\cdot)$ , Eq. ( 28 )	$0.15$	regular	events, marks
Irregular PDMP	Poisson, $\lambda{=}0.3$	$W_{1}\tanh(\cdot)$ , Eq. ( 28 )	$0.15$	random	events, marks, grid

Appendix

Table 5: The six synthetic datasets. All have latent and observation dimension

2

and

52

observations per trajectory, of which

44

are context

Figure 3: Forecast quality against latent update cost on the six synthetic systems. Rows are systems, columns are CRPS, forecast RMSE, horizon CRPS, joint LPD of the central 90% predictive interval, and joint log predictive density per scalar. The horizontal axis is the number of dynamics-network evaluations per trajectory over the context window, on a logarithmic scale. LJM is shown at each of

\beta_{\text{sparse}}\in\{0,0.03,0.1,0.3,1\}

; points of one family are joined in order of cost.

	RMSE		90% coverage		Future events/path
$\beta_{\mathrm{sparse}}$	Full	Floor	Full	Floor	Full
Cyclic jumps
$0$	$0.124$	$1.312$	$0.88$	$0.37$	$50.6$
$0.03$	$0.126$	$1.225$	$0.88$	$0.38$	$35.2$
$0.1$	$0.126$	$0.808$	$0.88$	$0.39$	$22.7$
$0.3$	$0.127$	$0.473$	$0.88$	$0.56$	$14.7$

Appendix

Table 6: Future-event suppression across all six synthetic systems and five event-count penalties. Full uses the learned prior intensity; Floor holds it at its positive minimum during forecasting. Context inference and model parameters are unchanged. Each condition uses 1,000 validation trajectories and 128 predictive paths. The last column is the mean accepted future event count under Full. Coverage is rounded to two decimal places.

$\beta_{\mathrm{sparse}}$	$N_{\mathrm{ev}}\downarrow$	CRPS $/s_{h}\downarrow$	RMSE (W) $\downarrow$	Cov. $\to$ 0.90	Width (W)	Energy CRPS $\downarrow$
0	205.4	0.0491	550.3	0.923	830	0.0627
0.01	200.6	0.0496	549.8	0.914	858	0.0633
0.03	176.8	0.0508	563.8	0.929	864	0.0664
0.1	185.0	0.0507	559.8	0.928	876	0.0649
0.3	147.5	0.0517	563.5	0.919	855	0.0668
1	98.4	0.0516	565.4	0.926	920	0.0666

Appendix

Table 7: Event-count penalty sweep for the latent jump model on REFIT, 18 households, context 512, test split.

Model	$B$	$L$	Seq.	Scan	Seq.	Scan	Speed-up	Seq.	Scan
			Forward (ms)		Forward + backward (ms)			Peak memory
Cyclic	16	$1024^{\dagger}$	985	5.0	3,025	24	$125\times$	155 MB	174 MB
Cyclic	128	1024	1,019	10.9	3,053	32	$96\times$	1.20 GB	1.30 GB
Cyclic	128	4096	4,120	41.8	12,506	93	$134\times$	4.81 GB	5.20 GB

Appendix

Table 8: Matched-model posterior state construction on one NVIDIA A40 (float32; median times). Seq. is the chronological flow–gated-jump loop; Scan is the pure-PyTorch scan used in training, without a fused kernel.

B

is batch size and

L

is the processed proposal length.

\dagger

marks training padding lengths. Speed-up refers to the forward–backward pass and are computed from unrounded timing measurements; timing entries are rounded. Peak-memory units are shown in each entry.

Explore similar work

Deep ZakaiJ: Structured Filtering for Jump-Diffusion Time Series Forecasting
May 23, 2026Yan Leng, Thibaut Mastrolia, Hao WangStochastic Differential Equations Latent Dynamics Modeling
University of Texas at Austin · University of California, Berkeley
Latent Laplace Diffusion for Irregular Multivariate Time Series
May 19, 2026Zinuo You, Jin Zheng, John CartlidgeIrregular Time-Series Modeling Multivariate Time Series Forecasting
School of Computer Science, University of Bristol · School of Engineering Mathematics and Technology, University of Bristol
Variational Inference for Lévy Process-Driven SDEs via Neural Tilting
May 11, 2026Yaman Kindap, Manfred Opper, Benjamin Dupuis +2Stochastic Differential Equations Bayesian Inference
1Imperial College London, UK · 2Technical University of Berlin, Germany · 3INRIA, CNRS, Département d’Informatique de l’Ecole Normale Supérieure / PSL, France

May 23, 2026J/K move · Enter open · S save

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

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

Latent Laplace Diffusion for Irregular Multivariate Time Series

Variational Inference for Lévy Process-Driven SDEs via Neural Tilting