Spatiotemporal System Forecasting with Irregular Time Steps via Masked Autoencoder

Kewei Zhu; Yanze Xin; Jinwei Hu; Xiaoyuan Cheng; Yiming Yang; Sibo Cheng

cs.LGMar 26, 2026

Spatiotemporal System Forecasting with Irregular Time Steps via Masked Autoencoder

Authors: Kewei Zhu, Yanze Xin, Jinwei Hu, Xiaoyuan Cheng, Yiming Yang, Sibo Cheng

Organizations: Department of Chemical Engineering, University College London, UK · Department of Computing, Imperial College London, UK · Department of Computer Science, University of Liverpool, UK · Dynamic Systems Lab, University College London, UK ·

Abstract

Predicting high-dimensional dynamical systems with irregular time steps presents significant challenges for current data-driven algorithms. These irregularities arise from missing data, sparse observations, or adaptive computational techniques, reducing prediction accuracy. To address these limitations, we propose a novel method: a Physics-Spatiotemporal Masked Autoencoder. This method integrates convolutional autoencoders for spatial feature extraction with masked autoencoders optimised for irregular time series, leveraging attention mechanisms to reconstruct the entire physical sequence in a single prediction pass. The model avoids the need for data imputation while preserving physical integrity of the system. Here, 'physics' refers to high-dimensional fields generated by underlying dynamical systems, rather than the enforcement of explicit physical constraints or PDE residuals. We evaluate this approach on multiple simulated datasets and real-world ocean temperature data. The results demonstrate that our method achieves significant improvements in prediction accuracy, robustness to nonlinearities, and computational efficiency over traditional convolutional and recurrent network methods. The model shows potential for capturing complex spatiotemporal patterns without requiring domain-specific knowledge, with applications in climate modelling, fluid dynamics, ocean forecasting, environmental monitoring, and scientific computing.

Figures & tables

Figure 1: Comparison of Sequence-to-Sequence(Seq2Seq) prediction methods in dynamical systems of irregular time steps. Left : Traditional RNN-based models feature step-wise rolling out with necessary data imputation for handling missing steps, which may introduce biases and cumulative errors. Right : Our model performs element-wise predictions in the latent space by adaptive attention mechanism to reconstruct the complete sequence in a single pass.

Panel of Figure 2 — Figure 2: Architecture of the proposed P-STMAE framework. (a) Convolutional encoder compresses physical states into latent representations. Positional encodings are added, and a masked transformer captures temporal dependencies in latent space. (b) Learnable masking tokens are padded at missing and future time steps. Transformer blocks process the sequence, and the convolutional decoder reconstructs the complete physical fields. (c) Each transformer block consists of layer normalization, multi-head self-attention, and a feedforward network. Self-attention operates only on observed latent states.

Dataset	Model	MSE	SSIM	PSNR
Shallow Water	CAE	$5.32\times 10^{-5}$	$0.9596$	$44.64$
	P-STMAE	$\bm{6.16\times 10^{-5}}$	$\bm{0.9538}$	$\bm{43.90}$
	ConvRAE	$9.86\times 10^{-5}$	$0.9394$	$42.47$
	ConvLSTM	$1.82\times 10^{-4}$	$0.9231$	$40.72$
Diffusion-Reaction	CAE	$3.66\times 10^{-5}$	$0.9887$	$48.78$
	P-STMAE	$\bm{5.99\times 10^{-5}}$	$0.9870$	$44.20$

Table 1: Test metrics in the full physics space across all three datasets with a missing ratio of 0.5. Metrics for Shallow Water are averaged over all

h

u

, and

v

variables; metrics for Diffusion-Reaction are averaged over both

u

and

v

variables. Bold values indicate the best performance among the three forecasting models (excluding CAE).

Figure 3: Ground truth (top) and error maps of P-STMAE, ConvRAE, and ConvLSTM for forecasting the variable $u$ in the shallow water dataset with a sampling dilation of 3. Columns represent successive forecasting steps. Among the models, P-STMAE yields the smallest errors, indicating predictive accuracy.

Panel of Figure 4 — Figure 4: Robustness analysis of P-STMAE on the shallow water dataset. (a) Performance comparison under varying numbers of missing steps in the input sequence with a length of 10. Each model is trained and evaluated with missing steps ranging from 1 to 6. P-STMAE demonstrates consistent performance and robustness, while the RNN-based models, especially ConvLSTM, show higher sensitivity to increasing missing steps. (b) Test performance comparison regarding the sampling dilations of data sequences. All models are separately trained on the shallow water dataset of different dilations.

Figure 5: Ground truth (top) and error maps of P-STMAE, ConvRAE, and ConvLSTM for forecasting the variable $u$ in the diffusion-reaction dataset with a sampling dilation of 5. Columns represent successive forecasting steps. The results show that P-STMAE consistently achieves lower errors than the baselines, confirming its advantage.

Figure 6: Error maps of P-STMAE, ConvRAE, and ConvLSTM predictions for the forecasting steps in the SST dataset.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Layer Type	Output Shape	Activation
Input	(128, 128, 3)
Conv 2D	(128, 128, 8)	GELU
Conv 2D	(64, 64, 16)	GELU
Conv 2D	(32, 32, 32)	GELU
Conv 2D	(16, 16, 64)	GELU
Conv 2D	(8, 8, 128)	GELU

Appendix

Table A1: Network structures of CAE encoder (left) and decoder (right) used in the P-STMAE. All of the convolutions and transpose convolutions use the kernel size of

3\times 3

and the same padding. The input dimension and channel number can vary depending on the dataset.

Parameter	Symbol	Min Value	Max Value
Initial bump centre (x)	$p_{x}$	54.00	74.00
Initial bump centre (y)	$p_{y}$	54.00	74.00
Bump height	$h$	0.05	0.20
Bump radius	$r$	8.94	12.65
Friction coefficient	$b$	0.02	2.00
Snapshot interval (steps)	–	60.00	100.00

Appendix

Table B1: Ranges of computational parameters used in the shallow water equation simulations. All parameters are uniformly sampled.

$\lambda$	MSE	SSIM	PSNR
0.01	$5.85\times 10^{-5}$	0.9276	41.02
0.02	$5.25\times 10^{-5}$	0.9317	41.52
0.05	$4.72\times 10^{-5}$	0.9264	40.82
0.10	$5.20\times 10^{-5}$	0.9315	41.60
0.20	$6.10\times 10^{-5}$	0.9177	39.56
0.50	$8.00\times 10^{-5}$	0.9356	42.07

Appendix

Table B2: Quantitative ablation study on

\lambda

Figure B1: Visualisation of the trained CAE model performance on the variable $h$ of shallow water samples.

Figure B2: Validation MSEs on the shallow water dataset with the sampling dilation of 3. Left: Full space MSE curves of the P-STMAE, ConvRAE, and ConvLSTM. The dotted line shows the performance of the trained CAE. Right: Latent space MSE curves of the P-STMAE and ConvRAE. Note that the ConvLSTM does not use latent representations.

Figure B3: Model predictions and error maps for the forecasting steps of the variable $u$ in the shallow water dataset with a sampling dilation of 3. The first row shows the ground truth input, rows 2–4 show the model-predicted outputs of the three models, and rows 5–7 show the pixel-level absolute differences between the two.

Figure B4: Model error maps for different values of $\lambda$ across forecasting steps in the shallow water dataset.

Figure C1: Validation MSEs on the diffusion reaction dataset with the sampling dilation of 5.

Figure C2: Model predictions and error maps for the forecasting steps of the variable $u$ in the diffusion reaction dataset with a sampling dilation of 5.

Figure D1: Model predictions and error maps for the forecasting steps in the SST dataset.

Spatiotemporal Forecasting

Spatiotemporal forecasting predicts how quantities change across locations and time, from traffic and weather to disease and environmental systems.

See all 423 papers

Enhancing Irregular Time Series Forecasting with Continuous-Time Modeling Framework
Jul 30, 2026Tianen Shen, Zhengyu Li, Yutong Li +4Irregular Time-Series Modeling Multivariate Time Series Forecasting
MASCIT: A Mask-Aware State Space Classifier for Naturally Irregular Time Series
Sep 28, 2026Yoo-Min Jung, Hyeon-Gi Kim, Jonghun ParkIrregular Time-Series Modeling Selective SSMs
Department of Industrial Engineering Seoul National University, Seoul, Republic of Korea
LLMODE: Aligning ODEs with LLMs via Gated Token Injection for Irregular Spatio-Temporal Forecasting
Aug 30, 2026Di Zhang, Jingyang Zhang, Ziqian Wang +4Irregular Time-Series Modeling Zero-Shot Time Series Forecasting
Beihang University Beijing, China

Jul 30, 2026J/K move · Enter open · S save

Enhancing Irregular Time Series Forecasting with Continuous-Time Modeling Framework

Tianen Shen, Zhengyu Li, Yutong Li, Xiangfei Qiu+3

Irregular Time-Series Modeling Multivariate Time Series Forecasting Flow Matching Time Series Modeling Time Series Transformers

Open paper AskPDF arXiv

Irregular multivariate time series are widely encountered in applications such as healthcare monitoring, human activity recognition, and environmental sensing. Their core challenges stem from asynchronous observations, non-uniform sampling intervals, and the fact that temporal patterns themselves carry critical dynamic information. Existing approaches either rely on discretization-based preprocessing (e.g., interpolation, imputation, or aggregation), which disrupts the underlying continuous-time semantics, or adopt continuous-time modeling via ODE-based frameworks, which typically require specialized architectures and incur substantial computational overhead due to numerical solvers. To address these limitations, we propose WrapFlow, a continuous-time modeling framework for irregular time series forecasting. On the input side, WrapFlow introduces Continuous-Time Tokenization, which directly encodes raw observation events and explicitly models long unobserved intervals via gap-aware tokens. The resulting continuous-time tokens are then processed by a standard Transformer backbone to capture long-range temporal dependencies. On the output side, we develop a simulation-free training paradigm for Residual Flow Matching, which learns conditional residual vector fields around base predictions while avoiding numerical-solver simulation and backpropagation during training. This design enables high-quality continuous forecasting using only a small number of fixed rollout steps at inference. Extensive experiments on multiple real-world datasets demonstrate that WrapFlow achieves state-of-the-art performance.

Abstract

Spatiotemporal Forecasting

Explore similar work

Enhancing Irregular Time Series Forecasting with Continuous-Time Modeling Framework

MASCIT: A Mask-Aware State Space Classifier for Naturally Irregular Time Series

LLMODE: Aligning ODEs with LLMs via Gated Token Injection for Irregular Spatio-Temporal Forecasting