Generative forecasting with joint probability models

TL;DR

Chaotic dynamical systems pose fundamental limits on deterministic forecasting due to sensitivity to initial conditions and unresolved processes. The authors propose a fully generative approach that learns the joint distribution over short state windows, $\hat{p}_\theta(x_t, x_{t-\Delta t}, \ldots, x_{t-(n-1)\Delta t})$, and obtains forecasts by marginalization, enabling richer temporal dependencies and trajectory segments. The framework is model-agnostic and provides intrinsic uncertainty quantification from joint samples via ensemble variance, short-horizon autocorrelation, and Wasserstein drift, without ground-truth data. Evaluations on Lorenz--63 and KS demonstrate improved short-term predictive skill, preservation of attractor geometry, and substantially better long-range statistics compared to conventional next-step models, with UQ metrics offering predictive insight into forecast reliability and potential extremes.

Abstract

Chaotic dynamical systems exhibit strong sensitivity to initial conditions and often contain unresolved multiscale processes, making deterministic forecasting fundamentally limited. Generative models offer an appealing alternative by learning distributions over plausible system evolutions; yet, most existing approaches focus on next-step conditional prediction rather than the structure of the underlying dynamics. In this work, we reframe forecasting as a fully generative problem by learning the joint probability distribution of lagged system states over short temporal windows and obtaining forecasts through marginalization. This new perspective allows the model to capture nonlinear temporal dependencies, represent multistep trajectory segments, and produce next-step predictions consistent with the learned joint distribution. We also introduce a general, model-agnostic training and inference framework for joint generative forecasting and show how it enables assessment of forecast robustness and reliability using three complementary uncertainty quantification metrics (ensemble variance, short-horizon autocorrelation, and cumulative Wasserstein drift), without access to ground truth. We evaluate the performance of the proposed method on two canonical chaotic dynamical systems, the Lorenz-63 system and the Kuramoto-Sivashinsky equation, and show that joint generative models yield improved short-term predictive skill, preserve attractor geometry, and achieve substantially more accurate long-range statistical behaviour than conventional conditional next-step models.

Figures (11)

• Figure 1: Schematic overview of the proposed generative joint forecasting framework for chaotic dynamical systems. (A) Inherent challenges of forecasting high-dimensional chaotic systems, characterized by extreme sensitivity to initial conditions and complex multi-scale dynamics. (B) Proposed core methodology: modelling the joint probability distribution of temporal sequences, e.g., $p({x}_{t_n}, {x}_{t_1}, {x}_{t_2},\ldots)$. The forecast for the current state, ${x}_{t_n}$, is then obtained through marginalization of this learned joint distribution. (C) Forecasts are obtained by sieving a joint ensemble using components overlapping with an observed segment. (D) Joint samples enable intrinsic uncertainty quantification through ensemble variance, autocorrelation, and Wasserstein drift metrics.
• Figure 2: Visualization of a single inference step with Algorithm \ref{['alg:inference']} for $n=5$. The trajectory formed by the orange points (top) serves as the reference for isolating the closest-matching tail points, $\hat{x}_{t-(n-1)\Delta t}$ through $\hat{x}_{t-\Delta t}$, from the joint point cloud (bottom). The corresponding head point $\hat{x}_t$ (purple star) is the prediction, which can then be appended to the input set (orange) to perform a successive autoregressive step.
• Figure 3: Single trajectory of Lorenz--63, 10 time units forward from $t=2.5\times10^4$. Top row shows each models' $x$ in time series, with the corresponding $(y, z)$-phase space view below.
• Figure 4: Schematic overview of the specific model implementation throughout Section \ref{['sec:results']}. The transformers (blue) are configured to have an input dimension of 256 (8 reserved for sinusoidal position encoding), feedforward dimension of 1024, 4 layers, and 4 attention heads. They are configured as encoder--decoder architectures for the conditional models and decoder-only for the unconditional model; the input and latent sequences are inputs for their respective transformers' decoders. Red outlines denote the conditioning inputs passed into the encoder portion of the transformers (when present).
• Figure 5: Mean absolute error (MAE) across 500 different initial conditions for Lorenz--63; for each individual component in the first 3 rows, and for the mean across all components below. Each plot redundantly overlays the other columns' MAEs for comparison.