Koopman Generative Operators for Efficient Probabilistic Time-Series Forecasting

AI summary

Overview

Research area: Probabilistic time-series forecasting, generative modeling, and operator-theoretic (Koopman) dynamical systems.

Technical level: Advanced — the paper assumes familiarity with flow matching, diffusion sampling, Koopman operator theory, matrix exponentials, and the CRPS/NMAE evaluation protocol.

Scope: The paper introduces KGO (Koopman Generative Operator), a framework that casts probabilistic forecasting as the evolution of structured uncertainty in a learned Koopman latent space, enabling single-step generation rather than iterative sampling.

What This Paper Is About

Modern probabilistic forecasters face a three-way tension: latent dynamical models capture temporal structure but use restrictive generation mechanisms, while diffusion and flow-matching models are expressive but require many iterative sampling steps. The authors argue that forecasting should instead be treated as the evolution of structured uncertainty over time — temporal dynamics and uncertainty should be decoupled, so that each can be modeled efficiently. KGO is their instantiation of this idea: a model that extrapolates latent dynamics with a learned linear operator and then transports noise to future trajectories in closed form, in one step.

Key Contributions

1. A reframing of probabilistic forecasting. The paper formulates prediction as uncertainty transport governed by temporal structure, rather than as static distribution prediction or iterative simulation, and derives three requirements from it: efficient generation, structured temporal evolution, and adaptive uncertainty.

2. Koopman Patch Embedding (KoPE). A learned linear operator K (normalized by its spectral norm to ≤ 1 for rollout stability) that generates future latent states as z̃_{N+i} = K^{i-1} z_N. Crucially, KoPE does not generate forecasts — it only produces a conditioning latent trajectory, decoupling temporal evolution from the generative process. This differs from prior Koopman work such as K²VAE, where latent evolution is the generative mechanism.

3. Koopman Flow Matching (KoFM). A generative head that replaces the time-dependent velocity field of standard flow matching with a time-independent field in latent space, admitting the closed-form affine transport z̃^{(1)} = e^{(1-τ)A} z̃^{(τ)} + (1-τ)b. At inference τ = 0, so a single cached matrix e^A yields a one-step matrix-product transport.

4. Adaptive Uncertainty Gate (AUG). A lightweight encoder–router that outputs per-(variate, horizon) gating coefficients g^h ∈ [0,1] mixing a deterministic ResMLP decoder with the KoFM stochastic branch, giving a learned decomposition of aleatoric uncertainty without manual tuning.

Main Findings

• Leading accuracy on ProbTS: KGO achieves the lowest CRPS on 13/17 evaluated settings and the lowest NMAE on 16/17, which the authors summarize as state-of-the-art on 29/34 configurations with CRPS improvements of up to 27%.
• Strongest in long-horizon forecasting: first place on 7/9 datasets in CRPS and 9/9 in NMAE (K²VAE, the closest competitor, manages only 2/9 in CRPS and none in NMAE). On ETTh2-L, CRPS drops by more than 27% and NMAE by more than 13% relative to K²VAE; on Weather-L, CRPS by roughly 26% and NMAE by roughly 18%.
• Also strongest at short horizons: 6/8 first-place CRPS and 7/8 NMAE results; the largest relative gains versus K²VAE are on ETT (about 23% on ETTm1-S, 25% on ETTm2-S, 19% on ETTh2-S in CRPS). Notable exception: on Solar-S, KGO (0.456 CRPS) trails K²VAE (0.367).
• Fast inference: 11 ms/window versus 27 ms/window for K²VAE — a 2.5× speedup over 36 dataset–horizon pairs — and the paper claims at least a 25× reduction relative to iterative generative models. Against Sundial-base-128m at its 50-step Euler-ODE budget, KGO reports roughly 13% lower CRPS, roughly 7% lower NMAE, and 7.8 ms/window versus 387 ms/window (about 50×).
• Memory is a trade-off, not a win: average peak GPU memory is 0.11 GB, higher than K²VAE (0.07 GB) and CSDI/MAF baselines (0.03–0.05 GB), though lower than TimeGrad and GRU NVP (0.12 GB) and TSDiff (0.34 GB). Batching many trajectories in parallel increases activation memory.
• Beats zero-shot foundation models in the full-shot setting: KGO ranks first on 5/6 datasets for both CRPS and NMAE against TiRex, Chronos-2, and Sundial; the single exception is Exchange, where Chronos-2 and TiRex slightly outperform it.
• Ablations confirm each component: average CRPS rises by 16.9% when KoPE is replaced by a last-patch representation (0.141 vs 0.165; context pooling ties at 0.165), by 4.7% when sLSTM is swapped for a Transformer (0.141 vs 0.148; a Linear RNN gives 0.145), and by 9.4% when KoFM is replaced by a normalizing-flow head (a single-step VAE head gives +7.2%).
• AUG's margin is small but qualitative: 0.141 average CRPS versus 0.144 for conditioning KoFM on the deterministic prediction and 0.142 for fixed residual addition — so its value lies in adaptive, interpretable uncertainty allocation rather than raw accuracy.
• Learned dynamics are interpretable: the dominant eigenvalues of the learned KoPE operator recover the daily and weekly cycles in the data.

Methodology in Plain English

The pipeline runs in four stages. First, the input series is de-trended with reversible instance normalization (RevIN) and cut into non-overlapping patches, which are embedded and passed through an sLSTM encoder to produce a sequence of latent states; the last state is treated as a compact summary of the system. Second, KoPE repeatedly applies a single learned matrix to that summary state to "roll forward" a trajectory of future latent states, one per future patch, with the matrix rescaled whenever its spectral norm exceeds 1 so the rollout does not blow up.

Third, KoFM handles uncertainty. Standard flow matching trains a velocity field that gradually pushes noise into data over many ODE steps. Here the transport is constrained to be linear and time-independent in latent space, which means the whole path has an exact solution via a matrix exponential plus a learned offset. At sampling time this collapses to a single matrix multiply and bias add — that is where the speedup comes from. Training interpolates between noise and data with a transport time τ, and τ is set to 0 with probability 0.5 per sample to improve optimization. In parallel, a residual MLP decodes the KoPE latents into a deterministic point forecast.

Fourth, AUG learns a gate per variable and per horizon that blends the deterministic and probabilistic outputs, so predictable signals lean on the point forecast and uncertain ones lean on the generative branch. The whole model is trained end-to-end with a CRPS-based objective combining an Energy Score term (an L1 accuracy term plus a repulsion term that prevents sample collapse, implemented by shuffling samples) and a CRPS quantile (pinball) loss over a set of quantile levels.

Evaluation uses the ProbTS benchmark: 17 data configurations spanning ETT, Electricity, Exchange, Traffic, Weather, Solar, and ILI, against eight baselines (GRU MAF, Trans MAF, GRU NVP, TimeGrad, TSDiff, CSDI, K²VAE, D³VAE), with CRPS and NMAE reported as averages over five independent runs. Short-term settings use context equal to horizon (24–30); long-term settings fix context at 96 and average over horizons {96, 192, 336, 720}, with ILI using context 36 and horizons {24, 36, 48, 60}.

Why This Matters

Impact on research. The paper challenges the presumed expressivity-versus-efficiency trade-off in generative forecasting by showing that restricting the transport to a time-independent linear field in a learned latent space — a finite-dimensional approximation of the Koopman generator — still produces a stronger predictive distribution than one-shot VAE or normalizing-flow heads (9.4% and 7.2% worse respectively). It also argues that a full-shot model trained on target data can still beat zero-shot foundation models on probabilistic metrics, which is a useful counterweight to the current scaling narrative. The decoupling of latent evolution from generation is a clean architectural idea that other latent dynamical models could adopt.

Real-world applications (as identified by the authors and benchmark domains):

• Energy systems — load and solar generation forecasting where miscalibrated uncertainty drives either failures or over-conservative provisioning.
• Finance — risk-sensitive forecasting where under- and over-estimation of outcome ranges carry direct cost.
• Healthcare — demand and patient-flow forecasting where calibrated predictive distributions support decisions.
• Traffic and weather — the benchmark domains used, where forecasts must be produced repeatedly under latency constraints.

Industry relevance. Single-step inference (11 ms/window, roughly 50× faster than a 50-step flow-matching baseline in the reported comparison) matters for real-time and large-scale deployment. The memory cost (0.11 GB average peak) is a caveat for edge deployment, but it is below TimeGrad, GRU NVP, and TSDiff. Code is released at github.com/azencot-group/KGO.

Future Directions

• Unifying physical and transport time. The authors explicitly call for a single Koopman framework that models evolution in both the time axis of the signal and the transport time of the generative process, rather than treating them separately.
• Reducing the memory cost of batching. Since latency gains come partly from processing all trajectories in parallel, and peak activation memory grows with the number of simultaneous samples, better memory management or adaptive batching is an open engineering question.
• Understanding AUG's small accuracy margin. AUG's CRPS advantage over fixed coupling rules is narrow (0.141 vs 0.142–0.144); whether more expressive gating or calibrated gating objectives widen the gap, and how the gate behaves under distribution shift, is unresolved.
• Extending beyond the evaluated benchmark suite. Whether the KoPE rollout remains stable and the operator's eigenvalues remain interpretable on much longer horizons, higher-dimensional variates, or non-cyclical dynamics is not established by the 17 ProbTS configurations.

Target Audience

Graduate students and researchers in time-series forecasting, generative modeling, and dynamical systems who are already comfortable with flow matching and latent-variable models; practitioners building latency-sensitive probabilistic forecasting systems for energy, finance, traffic, or healthcare who want a single-step alternative to diffusion; and readers interested in applying Koopman operator theory to machine learning. Readers new to CRPS, Koopman operators, or ODE-based generative models will need background reading first — the paper assumes these as prerequisites and defers implementation detail to its appendices.

Authors’ abstract

Probabilistic time-series forecasting requires models that simultaneously capture structured temporal dynamics, expressive uncertainty, and efficient inference, yet existing approaches fall short of this goal: latent dynamical models impose structure but rely on restrictive generation mechanisms, while modern generative methods such as diffusion and flow matching achieve flexibility at the cost of iterative and computationally expensive sampling. We introduce the Koopman Generative Operator (KGO), a probabilistic forecasting method that conceptualizes prediction as the evolution of structured uncertainty. KGO integrates three core components: (1) Koopman Patch Embedding (KoPE) for temporally consistent latent trajectory extrapolation; (2) Koopman Flow Matching (KoFM), which enables fast, single-step generation through a closed-form matrix exponential in a Koopman latent space, bypassing iterative sampling; and (3) an Adaptive Uncertainty Gate (AUG), which provides probabilistic forecasts by adapting uncertainty per-variable and per-horizon, corresponding to a learned decomposition of aleatoric uncertainty without manual tuning. By unifying these components, KGO achieves state-of-the-art accuracy on 13/17 ProbTS benchmarks in CRPS and 16/17 in NMAE. Further, by eliminating iterative sampling, KGO delivers at least a 25$\times$ reduction in inference time compared to iterative generative models. These results position KGO as a practical and principled framework that advances the accuracy-efficiency trade-off in probabilistic time-series forecasting, combining leading accuracy on the evaluated benchmarks with fast, single-step generation.