Contents

Spatiotemporal Prediction of Unknown Areas Based on Structured State Space Diffusion and Resolution Adaptive Fixed Rank Kriging

Master's Thesis

Abstract

With the rapid advancement of sensing and remote sensing technologies, data collected in transportation, meteorology, and environmental monitoring have become increasingly characterized by high-frequency observations and cross-regional spatiotemporal dependencies. However, missing observations, irregular sampling, and prediction at unobserved locations remain major challenges for spatiotemporal forecasting. To address these issues, this study proposes a spatiotemporal forecasting framework, $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK, which integrates Structured State Space Diffusion Model with S4 Layers ($\text{SSSD}^{\text{S4}}$) and Multi-Resolution Thin-Plate Spline (MRTS) basis functions to jointly model long-range temporal dependencies and spatial structural patterns. MRTS is incorporated into the S4 layers as additional spatial conditioning information, enabling $\text{SSSD}^{\text{S4+MRTS}}$ to learn spatiotemporal dynamics conditioned on observed locations. During inference, the framework further integrates MRTS-based Resolution Adaptive Fixed Rank Kriging (AFRK), which preserves the spatial consistency of the predicted field while enabling interpolation at unobserved locations. Experiments conducted on the Weather2K, MERRA-2, and the dataset from The Second Competition on Spatial Statistics for Large Datasets are used to compare the proposed framework with Temporal Fusion Transformers (TFT), Vector Autoregression (VAR), Stochastic Variational Gaussian Process (SVGP), and Spatio-temporal DeepKriging (STDK). The results demonstrate that $\text{SSSD}^{\text{S4+MRTS}}$ generally improves Mean Squared Prediction Error (MSPE) under most experimental settings, with particularly notable gains in future forecasting at unobserved locations. Overall, the findings indicate that integrating temporal dynamic modeling with spatial statistical methods effectively enhances spatiotemporal forecasting performance under missing-data and spatial extrapolation scenarios.

Keywords: Structured State Space Diffusion Model, Multi-Resolution Thin-Plate Spline Basis Functions, Resolution Adaptive Fixed Rank Kriging, spatiotemporal forecasting, spatial interpolation, missing data imputation

1 Introduction

With the rapid advancement of sensing technologies and data infrastructures, domains such as traffic monitoring, hydrology, meteorological observation, and satellite remote sensing have become capable of collecting large-scale datasets at high temporal frequencies, across extensive spatial domains, and over long time horizons. These datasets consequently exhibit rich spatiotemporal structure (Gneiting, 2002; N. Cressie and Wikle, 2011). However, real-world observations are frequently accompanied by challenges such as missing measurements, inconsistent sampling frequencies, sparsely distributed stations, and regions with long-term absence of observations, resulting in highly incomplete datasets (Little and Rubin, 2002; Decorte et al., 2024). Such missingness not only undermines statistical inference but also degrades model robustness and generalization in both forecasting and spatial extrapolation.

In practice, missing observations may arise from various mechanisms, including missing-at-random behaviors and structural missingness caused by prolonged spatial gaps or entire periods without observations. When the research objective extends from simple imputation to predicting future states at unobserved locations, the problem inherently involves both temporal extrapolation and spatial extrapolation. The former requires models capable of capturing nonlinear, long-range temporal dependencies, while the latter demands preservation of spatial continuity and coherent spatial structure. Existing approaches typically address only one of these dimensions. Traditional time-series models such as Autoregressive Integrated Moving Average (ARIMA) and Vector Autoregression (VAR) can capture temporal dynamics but lack explicit spatial structure. Spatial statistical methods such as Kriging and Fixed Rank Kriging (FRK) characterize spatial dependence but are not designed for long-sequence temporal modeling. Deep learning models such as Long Short-Term Memory (LSTM) networks and Temporal Fusion Transformers (TFT) offer nonlinear representational capacity but still face instability and limited extrapolation ability in long-sequence, highly incomplete, or cross-space forecasting settings (Lim et al., 2020; Wu et al., 2022).

Recent advances in State Space Models (SSMs), particularly the Structured State Space Sequence Model (S4) family and their extensions into Structured State Space Diffusion Model (SSSD) (Gu, Dao, et al., 2020; Gu, Goel, and Ré, 2022; Alcaraz and Strodthoff, 2023), demonstrate strong performance in long-sequence modeling and missing-data imputation. Yet these architectures remain fundamentally time-centric and lack explicit spatial structure, which limits their ability to maintain geographic continuity when performing spatial extrapolation. In contrast, Resolution Adaptive Fixed Rank Kriging (AFRK) (N. Cressie and Johannesson, 2008; Tzeng and Huang, 2018) provides a computationally scalable low-rank approximation for spatial domains but is not designed to model long-range temporal dynamics or complex nonlinear sequences.

These observations highlight a persistent methodological gap. Existing techniques struggle to simultaneously model long-sequence temporal dynamics, preserve spatial structural consistency, and maintain predictive performance under severe data incompleteness. To address this challenge, this study proposes a unified spatiotemporal forecasting framework that integrates SSSD with AFRK. The core idea is to leverage the long-sequence modeling capacity of SSSD to overcome the limitations of deep models under high missingness and strong temporal dependency, while employing the spatial low-rank structure of AFRK to compensate for the lack of spatial continuity in purely temporal architectures. Through the integration of Multi-Resolution Thin-Plate Spline (MRTS) basis representations and diffusion-based temporal modeling, the proposed framework jointly captures long-range temporal dynamics, spatial structural information, and incomplete observational patterns, providing a unified framework for spatial prediction and long-term forecasting under incomplete observations.

Spatiotemporal forecasting methods can generally be classified into temporal models, spatial statistical models, and integrated spatiotemporal frameworks. Although significant progress has been achieved in each domain, effectively combining temporal dynamics and spatial dependence structures remains a challenging research problem. The following sections review representative methods in each category and discuss their advantages and limitations with respect to spatiotemporal forecasting.

2.1 Time Series Models

2.1.1 Vector Autoregression Model

The Vector Autoregression (VAR) model, introduced by Sims (1980), serves as an extension of the univariate Autoregressive (AR) model, designed to capture the dynamic interdependencies among multiple time-varying variables. Distinguishing itself from univariate models, VAR treats all variables within the system as endogenous variables and describes their lagged influences through a system of simultaneous equations.

For a $p$-order vector autoregression model with $k$ variables, denoted as $\text{VAR}(p)$, the mathematical expression is defined as: \begin{align} \boldsymbol{y}_t = \boldsymbol{c} + \sum_{i=1}^{p} \boldsymbol{\Phi}_i \boldsymbol{y}_{t-i} + \boldsymbol{\epsilon}_t, \end{align} where $\boldsymbol{y}_t \in \mathbb{R}^k$ represents the observation vector at time $t$, $\boldsymbol{c}$ is the vector of intercept terms, $\boldsymbol{\Phi}_i \in \mathbb{R}^{k \times k}$ denotes the lag operator matrix that quantifies the influence of variables at time $t-i$ on the current state, and $\boldsymbol{\epsilon}_t \sim \mathcal{N}(0, \boldsymbol{\Omega})$ is the noise vector.

The primary advantage of the VAR model lies in its ability to analyze causal transmission mechanisms and dynamic interactions through Impulse Response Functions (IRF) and Variance Decomposition (Sims, 1980). However, as the number of variables $k$ or the lag order $p$ increases, the number of parameters grows quadratically, which often leads to overfitting issues. Furthermore, traditional linear VAR models struggle to capture complex nonlinear structures and the time-varying nature of parameters (Primiceri, 2005).

2.1.2 Temporal Fusion Transformers

Given the limitations of traditional linear multivariate models in capturing complex nonlinear dynamics and high-dimensional feature correlations, recent research has increasingly shifted toward deep learning architectures to enhance predictive performance. To overcome the challenges faced by conventional recurrent architectures in handling long-range dependencies and to effectively integrate diverse information sources, including static covariates, known future inputs, and historical observations, Lim et al. (2020) proposed the Temporal Fusion Transformer (TFT). The TFT is a deep learning architecture specifically designed for multi-horizon time-series forecasting. Its design emphasizes the allocation of weights to input variables of different natures through specialized network components, thereby enhancing the interpretability of prediction results and addressing the inadequacies of traditional statistical models when processing large-scale heterogeneous data.

The implementation of TFT relies on Gated Residual Networks (GRN), which regulate information flow via Gated Linear Units (GLU). This component enables the model to automatically adjust the depth of nonlinear transformations based on data characteristics. For an input vector $\boldsymbol{a}$ and an optional context vector $\boldsymbol{c}$, the operation is as follows: \begin{align} \text{GRN}_\omega(\boldsymbol{a}, \boldsymbol{c}) = \text{LayerNorm}(\boldsymbol{a} + \text{GLU}_\omega(\boldsymbol{\eta}_1)), \end{align} where $\boldsymbol{\eta}_1$ is the feature vector transformed by a weight matrix. This mechanism ensures high flexibility when processing sequences of varying complexities, effectively preventing deep networks from overfitting on simpler datasets.

For spatiotemporal data containing numerous external factors, TFT introduces Variable Selection Networks (VSN) to identify key variables from a large pool of input features. By assigning a weight $\nu_t^{(i)}$ to each feature, the model can automatically ignore redundant information and focus on influential factors. The integrated feature vector is represented as: \begin{align} \tilde{\boldsymbol{\xi}}_t = \sum_{i=1}^{m} \nu_t^{(i)} \tilde{\boldsymbol{\xi}}_t^{(i)}, \end{align} where $\tilde{\boldsymbol{\xi}}_t^{(i)}$ is the processed feature vector. This design significantly enhances the model’s robustness when dealing with high-dimensional feature inputs and allows researchers to intuitively quantify the contribution of various variables to the prediction results.

Regarding the capture of temporal relationships, TFT utilizes a modified Temporal Self-Attention mechanism to handle long-term dependencies. Compared to the standard Transformer architecture, TFT incorporates gating layers for residual connections within the attention layer and integrates historical and future spatiotemporal context through a decoder to identify the most influential time steps for the current prediction. TFT not only demonstrates superior predictive accuracy but also grants deep learning models the ability to interpret the significance of specific time steps or features (Lim et al., 2020).

2.1.3 State Space Model

The State Space Model (SSM) is a class of mathematical models that describes dynamic systems or sequential data through a latent state vector. Initially proposed by Kalman (1960) within the fields of control theory and filtering, SSMs were designed to address optimal filtering and prediction problems for linear dynamic systems. Subsequently, researchers such as Gu, Goel, and Ré (2022) extended this concept to deep learning architectures for long-sequence time series modeling, demonstrating that SSMs outperform traditional RNNs and LSTMs in capturing long-range dependencies and maintaining stable gradients (Gu, Goel, and Ré, 2022).

Given a one-dimensional input signal sequence $\boldsymbol{u}(t)$ and a one-dimensional output signal sequence $\boldsymbol{y}(t)$, the basic formulation of an SSM is \begin{align} \begin{aligned} \boldsymbol{x}’(t) & = \boldsymbol{A} \boldsymbol{x}(t) + \boldsymbol{B} \boldsymbol{u}(t); \\ \boldsymbol{y}(t) & = \boldsymbol{C} \boldsymbol{x}(t) + \boldsymbol{D} \boldsymbol{u}(t), \end{aligned} \end{align} where $\boldsymbol{x}(t) \in \mathbb{R}^N$ is an $N$-dimensional latent state that maps the input $\boldsymbol{u}(t)$, $\boldsymbol{x}’(t) = \frac{d}{dt}\boldsymbol{x}(t)$ denotes its time derivative, $\boldsymbol{A} \in \mathbb{R}^{N \times N}$ is the state matrix, and $\boldsymbol{B} \in \mathbb{R}^{N \times 1}$ and $\boldsymbol{C} \in \mathbb{R}^{1 \times N}$ are the input and output matrices, respectively, characterizing how the input influences the state and how the state maps to the output. The term $\boldsymbol{D} \in \mathbb{R}$ is the feedthrough matrix, allowing the input to directly affect the output, and is typically set to zero. In deep learning contexts, these parameters are generally learned via gradient descent.

https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/fig-SSM.svg
Figure 1: Typical State Space Model.

To address the practical issue in which the gradients of SSMs may increase or decrease exponentially with sequence length, Gu, Goel, and Ré (2022) introduced the HiPPO (High-order Polynomial Projection Operators) matrix (Gu, Dao, et al., 2020) to replace the original random matrix $\boldsymbol{A}$ in (4).

\begin{align} \boldsymbol{A}_{nk} = - \begin{cases} (2n + 1)^{1/2} (2k + 1)^{1/2}, & \text{if } n > k; \\ n + 1, & \text{if } n = k; \\ 0, & \text{if } n < k. \end{cases} \end{align}

As shown in (5), the HiPPO matrix dynamically evaluates the importance of each past time step as time evolves, enabling adaptive memory updates and retaining all historical information. By replacing the original random matrix $\boldsymbol{A}$ with the HiPPO matrix, the latent state $\boldsymbol{x}(t)$ can effectively store the historical information of the input sequence $\boldsymbol{u}(t)$ while avoiding gradient explosion or vanishing. Experimental results demonstrate that this design not only enhances computational stability but also significantly improves performance in long-sequence forecasting tasks (Gu, Goel, and Ré, 2022).

2.1.4 Structured State Space Sequence Model

Building upon the theory discussed in the previous section, Gu, Goel, and Ré (2022) proposed the Structured State Space Sequence Model (S4). This model aims to discretize the continuous-time State Space Model and embed it within deep learning frameworks for handling long sequential data.

S4 is grounded in the SSM formulation and discretizes (4) to enable its application to discrete input sequences. Let the step size be defined as $\Delta$, then the discretized SSM can be written as follows: \begin{align} \begin{aligned} \boldsymbol{x}_k & = \overline{\boldsymbol{A}} \boldsymbol{x}_{k - 1} + \overline{\boldsymbol{B}} \boldsymbol{u}_k; \\ \boldsymbol{y}_k & = \overline{\boldsymbol{C}} \boldsymbol{x}_k, \end{aligned} \end{align} where $\overline{\boldsymbol{A}}, \overline{\boldsymbol{B}}, \overline{\boldsymbol{C}}$ are the discrete approximations of $\boldsymbol{A}, \boldsymbol{B}, \boldsymbol{C}$, respectively. \begin{align} \begin{aligned} \overline{\boldsymbol{A}} & = (\boldsymbol{I} - \Delta / 2 \cdot \boldsymbol{A})^{-1} (\boldsymbol{I} + \Delta / 2 \cdot \boldsymbol{A}); \\ \overline{\boldsymbol{B}} & = (\boldsymbol{I} - \Delta / 2 \cdot \boldsymbol{A})^{-1} \Delta \boldsymbol{B}; \\ \overline{\boldsymbol{C}} & = \boldsymbol{C}. \end{aligned} \end{align}

To reduce the computational cost of matrix operations, S4 diagonalizes the discretized matrices, expressing them in an equivalent form under a different basis: \begin{align} \begin{aligned} (\boldsymbol{A}, \boldsymbol{B}, \boldsymbol{C}) \sim (\boldsymbol{V}^{-1} \boldsymbol{AV}, \boldsymbol{V}^{-1} \boldsymbol{B}, \boldsymbol{CV}), \end{aligned} \end{align} where $\boldsymbol{V}$ is the basis transformation matrix. Moreover, the discretized SSM can be reformulated in a convolutional form to enhance parallel computation efficiency: \begin{align} \boldsymbol{y} & = \overline{\boldsymbol{K}} * \boldsymbol{u}; \\ \overline{\boldsymbol{K}} & \in \mathbb{R}^L := (\overline{\boldsymbol{C}}\overline{\boldsymbol{B}}, \overline{\boldsymbol{C}}\overline{\boldsymbol{A}}\overline{\boldsymbol{B}}, \cdots, \overline{\boldsymbol{C}}\overline{\boldsymbol{A}}^{L - 1}\overline{\boldsymbol{B}}), \end{align} where $\overline{\boldsymbol{K}}$ denotes the SSM convolution kernel and $L$ represents the convolution length.

To further reduce computational complexity, S4 employs the Normal Plus Low-Rank (NPLR) parameterization to express the matrix $\overline{\boldsymbol{A}}$ as: \begin{align} \begin{aligned} \boldsymbol{A} = \boldsymbol{V \Lambda V}^* - \boldsymbol{PQ}^\top = \boldsymbol{V} (\boldsymbol{\Lambda} - (\boldsymbol{V}^*\boldsymbol{P})(\boldsymbol{V}^*\boldsymbol{Q})^*) \boldsymbol{V}^*, \end{aligned} \end{align} where $\boldsymbol{\Lambda}$ is a diagonal matrix, $\boldsymbol{P}, \boldsymbol{Q} \in \mathbb{R}^{N \times r}$ are low-rank matrices, and $\boldsymbol{V} \in \mathbb{C}^{N \times N}$ is a unitary matrix.

Based on the above theoretical components, S4 integrates the HiPPO matrix, discretization, diagonalization, convolution, and NPLR parameterization, enabling efficient computation and strong performance on long-sequence tasks (Gu, Goel, and Ré, 2022).

2.1.5 Diffusion Model

Diffusion models are a class of generative models that learn data distributions through a dual process consisting of a forward diffusion process and a reverse denoising process (Sohl-Dickstein et al., 2015). In the forward process, Gaussian noise is gradually injected into the original data until it approaches a standard normal distribution. The model is then trained to learn the reverse mapping of this process in order to reconstruct the data distribution. In recent years, diffusion models have been applied to time series imputation, where diffusion and denoising are performed only on missing segments to recover complete sequences under conditional observations (Alcaraz and Strodthoff, 2023).

Let $\boldsymbol{x}_0 \sim q(\boldsymbol{x}_0)$ denote an original data sample. The forward process is defined as a fixed-parameter Gaussian Markov chain that simulates progressively perturbed data generation: \begin{align} \begin{cases} q(\boldsymbol{x}_{1} | \boldsymbol{x}_0) = \prod_{t=1}^T q(\boldsymbol{x}_t | \boldsymbol{x}_{t-1}); \\ q(\boldsymbol{x}_t | \boldsymbol{x}_{t-1}) = \mathcal{N}(\boldsymbol{x}_t; \sqrt{1 - \beta_t} \boldsymbol{x}_{t-1}, \beta_t \boldsymbol{I}), \end{cases} \end{align} where $\beta_t$ is the variance schedule controlling noise intensity, and $\mathcal{N}$ denotes the normal distribution.

To reconstruct the data, the model must learn the reverse mapping of this process. The reverse process is defined as: \begin{align} \begin{cases} p_\theta(\boldsymbol{x}_{0}) = p(\boldsymbol{x}_T)\prod_{t=1}^T p_\theta(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t); \\ p_\theta(\boldsymbol{x}_{t-1} | \boldsymbol{x}_t) = \mathcal{N}(\boldsymbol{x}_{t-1}; \mu_\theta(\boldsymbol{x}_t, t), \Sigma_\theta(\boldsymbol{x}_t, t)), \end{cases} \end{align} where $p(\boldsymbol{x}_T) = \mathcal{N}(\boldsymbol{x}_T; 0, \boldsymbol{I})$ denotes a standard normal distribution, while $\mu_\theta$ and $\Sigma_\theta$ represent the mean vector and covariance matrix parameterized by a neural network with parameters $\theta$, respectively.

However, directly modeling the mean of the reverse process, $\mu_\theta$, is often difficult to optimize and may lead to unstable convergence in practice. To address this issue, Ho, Jain, and Abbeel (2020) proposed a parameterization known as the Denoising Diffusion Probabilistic Model (DDPM), which reparameterizes $p_\theta(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)$ as \begin{align} \mu_\theta(\boldsymbol{x}_t, t) & = \frac{1}{\sqrt{\alpha_t}} \Big(\boldsymbol{x}_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\, \boldsymbol{\epsilon}_{\theta}(\boldsymbol{x}_t,t)\Big), \\ \Sigma_\theta(\boldsymbol{x}_t, t) & = \sigma_t^2 \boldsymbol{I}, \qquad \sigma_t^2 = \beta_t \text{ or } \sigma_t^2 = \frac{1 - \bar{\alpha}_{t - 1}}{1 - \bar{\alpha}_t} \beta_t, \end{align} where $\alpha_t = 1 - \beta_t$ and $\bar{\alpha}_t = \prod_{s=1}^t \alpha_s$. Under this framework, $\boldsymbol{\epsilon}_{\theta}(\boldsymbol{x}_t,t)$ is used to estimate the random Gaussian noise added to $\boldsymbol{x}_t$ during the forward diffusion process, and $\boldsymbol{x}_{t-1}$ is reconstructed by removing the estimated noise component from $\boldsymbol{x}_t$.

This parameterization avoids the need to directly model complex high-dimensional data distributions, thereby substantially simplifying the training objective and improving numerical stability. Consequently, the sample at any diffusion step $t$ can be expressed as a linear combination of the original data $\boldsymbol{x}_0$ and Gaussian noise: \begin{align} \boldsymbol{x}_t = \sqrt{\bar{\alpha}_t} \boldsymbol{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}, \qquad \boldsymbol{\epsilon} \sim \mathcal{N}(0, \boldsymbol{I}), \end{align} which allows the model to randomly sample the time step $t$ and noise $\boldsymbol{\epsilon}$ during training without iteratively computing intermediate diffusion steps. Since this representation transforms the original problem of directly fitting the data distribution into the estimation of Gaussian noise, the training objective can be further simplified as \begin{align} \min_\theta \mathbb{E}_{t, \boldsymbol{x}_0, \boldsymbol{\epsilon}} \left[ \left\| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta \left(\sqrt{\bar{\alpha}_t} \boldsymbol{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}, t\right) \right\|^2 \right], \end{align} which corresponds to minimizing the Mean Squared Error (MSE) between the predicted noise and the injected noise.

For time series imputation, Alcaraz and Strodthoff (2023) further proposed a Conditional Diffusion Model that applies diffusion and denoising operations only to missing segments. During the training stage, the model receives partially observed sequences as conditional inputs to learn how to reconstruct the complete sequence. During the generation stage, the observed portions are fixed while the reverse denoising process is performed, allowing the missing segments to be reconstructed while preserving temporal consistency. By leveraging the stability and high-quality generation capability of diffusion models, this approach can effectively handle time series data with long-term dependencies or structurally missing segments.

2.1.6 Structured State Space Diffusion Model with S4 Layers

The Structured State Space Diffusion (SSSD) model, proposed by Alcaraz and Strodthoff (2023), is designed to integrate the generative stability of the DiffWave-based diffusion architecture (Kong et al., 2021) with the long-range dependency modeling capability of SSM. The model is formulated as a conditional diffusion framework for time series imputation, in which noise is applied exclusively to the missing segments while the observed portions remain unperturbed. This strategy effectively prevents information leakage and preserves conditional consistency. By learning the reverse denoising process at each diffusion step, SSSD progressively reconstructs the missing data under fixed observations, achieving high-quality imputation and generation performance.

Building upon this foundation, Alcaraz and Strodthoff (2023) further introduce the Structured State Space Diffusion model with S4 layers ($\text{SSSD}^{\text{S4}}$) as an enhanced variant tailored for time series tasks. $\text{SSSD}^{\text{S4}}$ retains the conditional diffusion design of SSSD but replaces the Bidirectional Dilated Convolutional Layers in the original DiffWave architecture with S4 layers, thereby improving the model’s ability to capture long-range temporal dynamics. Experimental results demonstrate that this variant achieves more stable and accurate imputation across diverse missing patterns, outperforming diffusion models based on conventional convolutional or Transformer architectures (Alcaraz and Strodthoff, 2023).

2.2 Spatial Statistics

2.2.1 Kriging

Kriging originates from spatial statistics in geosciences and is a linear interpolation method used to estimate the values of a random field at unobserved spatial locations based on observed data (N. A. C. Cressie, 1993). The observed value $Z(\boldsymbol{s})$ at location $\boldsymbol{s}$ is modeled as: \begin{align} Z(\boldsymbol{s}) = Y(\boldsymbol{s}) + \varepsilon(\boldsymbol{s}), \qquad \boldsymbol{s} \in D \subset \mathbb{R}^D, \end{align} where $Y(\boldsymbol{s}) = \mu(\boldsymbol{s}) + \xi(\boldsymbol{s})$ represents a linear mean structure varying over space, and $\varepsilon(\boldsymbol{s})$ is a zero-mean random noise term that is uncorrelated with $Y(\boldsymbol{s})$. The covariance function of the noise is given by $C(\boldsymbol{s}, \boldsymbol{s}’) = \text{Cov}(\varepsilon(\boldsymbol{s}), \varepsilon(\boldsymbol{s}’))$, which may correspond to a non-stationary spatial covariance structure.

Traditional Kriging relies on the inversion of a full covariance matrix, leading to computational costs that scale cubically with the number of observations $n$, thereby creating a significant computational bottleneck as $n$ increases (N. Cressie and Johannesson, 2008). Fixed Rank Kriging (FRK), along with its adaptive extension Resolution Adaptive Fixed Rank Kriging (AFRK), was introduced in this context to alleviate the computational burden associated with large-scale spatial data analysis.

2.2.2 Fixed Rank Kriging

To reduce the computational burden of Kriging, N. Cressie and Johannesson (2008) proposed Fixed Rank Kriging (FRK), which represents the random field using a finite set of basis functions, thereby approximating high-dimensional spatial random effects with low-dimensional random coefficients: \begin{align} Y(\boldsymbol{s}) = \mu(\boldsymbol{s}) + \boldsymbol{f}(\boldsymbol{s})^\top \boldsymbol{w} + \xi(\boldsymbol{s}), \end{align} where $\boldsymbol{f}(\boldsymbol{s}) = (f_1(\boldsymbol{s}), \dots, f_K(\boldsymbol{s}))^\top$ is a pre-specified $K$-dimensional vector of basis functions with $K \le n$, $\boldsymbol{w} \sim N(\boldsymbol{0}, \boldsymbol{M})$ with $\boldsymbol{M}$ is an unknown nonnegative-definite matrix, and $\xi(\boldsymbol{s}) \sim \mathcal{N}(0, \sigma_\xi^2)$ represents fine-scale random noise. The corresponding covariance matrix can then be expressed as: \begin{align} \text{Cov}(Y(\boldsymbol{s}), Y(\boldsymbol{s}’)) = \boldsymbol{f}(\boldsymbol{s})^\top \boldsymbol{M} \boldsymbol{f}(\boldsymbol{s}’) + \sigma_\xi^2 \boldsymbol{I}(\boldsymbol{s} = \boldsymbol{s}’). \end{align}

Since the rank of $\boldsymbol{fMf}^\top$ is typically much smaller than the number of observation points $n$, FRK can significantly reduce the computational cost of covariance matrix inversion, making it particularly suitable for large-scale remote sensing and environmental monitoring data.

2.2.3 Resolution Adaptive Fixed Rank Kriging

Building upon FRK, Tzeng and Huang (2018) further proposed Resolution Adaptive Fixed Rank Kriging (AFRK). The core idea is to allow the resolution of the basis functions to automatically adapt to the spatial distribution of the data, thereby providing greater flexibility in capturing spatial heterogeneity. This adaptive mechanism enables the model to allocate appropriate spatial resolution across different regions according to the underlying spatial variability.

The basis functions adopted in AFRK are Multi-Resolution Thin-Plate Spline Basis Functions (MRTS), which are constructed from Thin-Plate Splines (TPS). TPS is a commonly used smoothing spline method that produces smooth functions by minimizing the sum of squared errors together with a smoothness penalty term (Wahba and Wendelberger, 1980; Green and Silverman, 1993). Based on TPS, Tzeng and Huang (2018) further constructed an ordered set of basis functions at multiple resolutions via eigen-decomposition, referred to as the MRTS.

To automatically adapt to the spatial distribution of the data and spatial heterogeneity, AFRK selects the number of basis functions according to the magnitude of the corresponding eigenvalues. Only the bases that explain most of the spatial variability are retained, allowing the model to capture the dominant spatial variations using a relatively small number of basis functions while improving computational efficiency.

In AFRK, the MRTS functions are defined as: \begin{align} f_k(\boldsymbol{s}) = \begin{cases} 1, & \text{if } k = 1; \\ x_{k-1}, & \text{if } k = 2, \ldots, d+1; \\ \lambda^{-1}_{k-d-1} \times \left \{ \phi(\boldsymbol{s}) - \boldsymbol{\Phi} \boldsymbol{X} (\boldsymbol{X}’\boldsymbol{X})^{-1} \boldsymbol{x} \right \}’ \boldsymbol{v}_{k-d-1}, & \text{if } k = d+2, \ldots, n, \end{cases} \end{align} where $f_k$ is the $k$-th basis in $\boldsymbol{f}$, $\boldsymbol{X} \in \mathbb{R}^{N \times (d+1)}$ is the design matrix, with each row corresponding to the intercept and coordinates of an observation location $\boldsymbol{s}_i$, and $\boldsymbol{x} = (1, \boldsymbol{s}’)’ = (1, x_1, \ldots, x_d)’$. The matrix $\boldsymbol{\Phi}$ satisfies \begin{align} J(f) = \boldsymbol{\alpha}’ \boldsymbol{\Phi \alpha}, \end{align} where $\boldsymbol{\Phi}$ is an $n \times n$ matrix with entries $\phi_j(\boldsymbol{s}_i)$, where $\phi(\boldsymbol{s})$ is defined by: \begin{align} \phi_i (\boldsymbol{s}) = \begin{cases} \frac{1}{12} \| \boldsymbol{s} - \boldsymbol{s}_i \|^3, & \text{if } d = 1; \\ \frac{1}{8\pi} \| \boldsymbol{s} - \boldsymbol{s}_i \|^2 \log(\| \boldsymbol{s} - \boldsymbol{s}_i \|), & \text{if } d = 2; \\ -\frac{1}{8} \| \boldsymbol{s} - \boldsymbol{s}_i \|, & \text{if } d = 3; \end{cases} \end{align} and $\boldsymbol{v}_{k}$ denotes the $k$-th row of matrix $\boldsymbol{V}$, where $\boldsymbol{V} \operatorname{diag}(\lambda_1, \ldots, \lambda_n) \boldsymbol{V}’$ is the eigendecomposition of $\boldsymbol{Q \Phi Q}$, with $\boldsymbol{Q} = \boldsymbol{I} - \boldsymbol{X} (\boldsymbol{X}’\boldsymbol{X})^{-1} \boldsymbol{X}’$. This approach adjusts the basis resolution according to data density, producing higher-resolution bases in densely sampled regions while maintaining smoother structures in sparse regions.

By retaining the computational advantages of low-rank approximation while introducing data-driven multi-resolution bases, AFRK effectively handles irregular sampling and regions with significant local variability. Compared to conventional FRK, AFRK demonstrates superior predictive performance on non-uniform and non-stationary spatial data.

2.3 Spatiotemporal Models

2.3.1 Gaussian Processes

A Gaussian Process (GP) is a non-parametric Bayesian model used to define a distribution over functions that map inputs to outputs (Rasmussen and Williams, 2006). A GP is fully specified by a mean function and a covariance function (also known as a kernel function). Within the Bayesian framework, GP provides both predictive means and uncertainty quantification for predictions. However, exact posterior inference requires inversion or Cholesky decomposition of an $N \times N$ covariance matrix, resulting in a computational complexity of $\mathcal{O}(N^3)$. This makes standard GP models infeasible for large-scale spatiotemporal datasets.

To address this limitation, Hensman, Fusi, and Lawrence (2013) proposed the Stochastic Variational Gaussian Process (SVGP). This approach introduces a set of $M \ll N$ inducing points to approximate the full posterior distribution. In the implementation framework of Gardner et al. (2021), the variational distribution is parameterized as a multivariate Gaussian with a full covariance matrix. To ensure positive definiteness during optimization, the covariance matrix is parameterized via its mean vector and a lower-triangular matrix, i.e., through Cholesky factorization (Hensman, Matthews, and Ghahramani, 2014).

Within the SVGP framework, the size of the variational distribution is determined by the number of inducing points. Specifically, the variational mean has dimension $M$, and the variational covariance matrix is of size $M \times M$. By approximating the full posterior using $M \ll N$ inducing points and optimizing the Evidence Lower Bound (ELBO), SVGP reduces the computational complexity to $\mathcal{O}(M^3)$ (Gardner et al., 2021).

2.3.2 Spatio-temporal DeepKriging

Spatio-temporal DeepKriging (STDK) is a recently proposed non-parametric approach that integrates deep learning with spatial statistics for large-scale spatiotemporal interpolation and probabilistic prediction (Nag, Sun, and Reich, 2023). In contrast to traditional Gaussian Processes, which require a pre-specified covariance function and suffer from $\mathcal{O}(N^3)$ computational complexity, STDK learns spatiotemporal dependencies directly from data in a data-driven manner, thereby improving scalability for large datasets.

STDK first embeds spatiotemporal coordinates $(\mathbf{s}, t)$ into a high-dimensional feature space. Let the basis function vector be defined as \begin{align} \boldsymbol{\phi}(\mathbf{s}, t) = \left[\phi_1(\mathbf{s}, t), \dots, \phi_K(\mathbf{s}, t)\right]^\top, \end{align} where $\{\phi_k(\cdot)\}_{k=1}^K$ are multi-resolution basis functions. Previous studies commonly adopt compactly supported Wendland functions or radial basis functions (RBFs) to capture spatial dependencies across multiple scales (Nag, Sun, and Reich, 2023).

The embedded features are then passed into a deep neural network to model nonlinear mappings, which can be expressed as \begin{align} Z(\mathbf{s}, t) = f_{\boldsymbol{\theta}}\big(\boldsymbol{\phi}(\mathbf{s}, t)\big) + \epsilon, \end{align} where $f_{\boldsymbol{\theta}}(\cdot)$ denotes a deep neural network parameterized by $\boldsymbol{\theta}$, and $\epsilon$ represents a random noise term.

For probabilistic forecasting, STDK adopts a quantile loss function to estimate different quantiles of the conditional distribution, thereby constructing predictive intervals. Compared with mean squared error (MSE)-based approaches, this enables explicit uncertainty quantification.

By combining basis function embeddings with deep neural networks, STDK avoids explicit covariance matrix decomposition and scales efficiently to large datasets, making it a competitive baseline for large-scale spatiotemporal modeling tasks.

3 Methodology

This study presents a spatiotemporal forecasting framework that combines deep time-series modeling and spatial statistical modeling. By incorporating temporal dynamics and spatial correlations within a unified framework, the proposed approach is designed to reconstruct missing observations and generate forecasts for future time steps.

3.1 Spatiotemporal Forecasting Task

Consider a spatial domain consisting of a set of locations $\mathcal{S}$, which can be partitioned into the subset of observed locations $\mathcal{S}_{\mathrm{observed}}$ with available measurements and the subset of unobserved locations $\mathcal{S}_{\mathrm{unobserved}}$ with no recorded observations. For each observed location $\boldsymbol{s} \in \mathcal{S}_{\mathrm{observed}}$, the target variable is fully observed over the temporal horizon $t \in \{1, \dots, T\}$, and the corresponding observations are denoted as $y_t(\boldsymbol{s})$. In contrast, for each unobserved location $\boldsymbol{s}^\ast \in \mathcal{S}_{\mathrm{unobserved}}$, no historical measurements exist within this time interval. The objective of this study is to construct a predictive model that leverages the historical spatiotemporal features from the observed locations, denoted by $\boldsymbol{Y}_{1:T}(\mathcal{S}_{\mathrm{observed}})$, in order to perform temporal extrapolation and spatial estimation of the target variable $\hat{y}_{t}(\boldsymbol{s}^\ast)$ at unobserved locations for future time points $t > T$.

Table 1 summarizes the primary notations used throughout this study.

Table 1: Notation Summary.

SymbolDescription
$\mathcal{S}$Spatial domain.
$\mathcal{S}_{\mathrm{observed}}$Set of observed locations.
$\mathcal{S}_{\mathrm{unobserved}}$Set of unobserved locations.
$\boldsymbol{s}$Observed spatial location.
$\boldsymbol{s}^{\ast}$Unobserved spatial location.
$N$Number of observed locations.
$T$Length of the observed temporal sequence.
$y_t(\boldsymbol{s})$Observed value at location $\boldsymbol{s}$ and time $t$.
$\boldsymbol{Y}_{1:T}(\mathcal{S}_{\mathrm{observed}})$Collection of observations from all observed locations over the time interval $1,\ldots,T$.
$\tilde{y}^{(T)}_t(\boldsymbol{s})$Temporally standardized observation at location $\boldsymbol{s}$ and time $t$.
$\hat{\tilde{y}}^{(T)}_t(\boldsymbol{s})$Prediction produced by $\text{SSSD}^{\text{S4+MRTS}}$ on the temporally standardized scale.
$\tilde{y}^{(S)}_t(\boldsymbol{s})$Spatially standardized prediction at observed location $\boldsymbol{s}$ and time $t$.
$\hat{\tilde{y}}^{(S)}_t(\boldsymbol{s}^{\ast})$AFRK prediction on the spatially standardized scale at unobserved locations.
$\hat{y}_t(\boldsymbol{s})$Prediction at observed location $\boldsymbol{s}$ on the original measurement scale.
$\hat{y}_t(\boldsymbol{s}^{\ast})$Prediction at unobserved location $\boldsymbol{s}^{\ast}$ on the original measurement scale.
$\mu(\boldsymbol{s})$Temporal mean of the time series at location $\boldsymbol{s}$.
$\sigma(\boldsymbol{s})$Temporal standard deviation of the time series at location $\boldsymbol{s}$.
$\bar{y}_t$Spatial mean of predicted values over all observed locations at time $t$.
$\tau_t$Spatial standard deviation of predicted values over all observed locations at time $t$.
$\boldsymbol{\phi}(\boldsymbol{s})$MRTS basis feature vector associated with location $\boldsymbol{s}$.
$\boldsymbol{\Phi}$MRTS basis matrix constructed from all observed locations.
$\tilde{\boldsymbol{\Phi}}$MRTS basis tensor expanded along the temporal dimension.
$D_{\mathrm{MRTS}}$Number of MRTS basis functions.
$\boldsymbol{m}_t(\boldsymbol{s})$Binary mask indicating whether an observation is available at location $\boldsymbol{s}$ and time $t$.
$\mathbf{c}_t(\boldsymbol{s})$Conditional input vector supplied to the S4 layer.
$\theta$Trainable parameters of the diffusion network.
$g_\theta(\cdot)$Noise estimation network used in the reverse diffusion process.
$\boldsymbol{x}_t$Latent variable at diffusion step $t$.
$\boldsymbol{\epsilon}$Gaussian noise injected during the forward diffusion process.
$\hat{\boldsymbol{\epsilon}}_\theta$Predicted diffusion noise generated by the network $g_\theta$.
$\alpha_t$Variance schedule parameter at diffusion step $t$.
$\bar{\alpha}_t$Cumulative product of diffusion coefficients up to step $t$.
$K$Number of channels in the original feature representation.
$C$Number of output channels in the conditional convolution layer.
$\boldsymbol{f}(\boldsymbol{s})$Spatial basis vector used by AFRK.
$\hat{\boldsymbol{w}}_t$Estimated AFRK basis coefficients at time $t$.
$\hat{\xi}_t(\boldsymbol{s})$Estimated AFRK residual component at location $\boldsymbol{s}$.

The objective can be formulated as learning a mapping function \begin{align} \hat{y}_{t}(\boldsymbol{s}^\ast) = f \left (\boldsymbol{Y}_{1 : T}(\mathcal{S}_{\mathrm{observed}}), \boldsymbol{s}^\ast \right ), \qquad \boldsymbol{s}^\ast\in\mathcal{S}_{\mathrm{unobserved}}, ~ t > T, \end{align} where $\boldsymbol{Y}_{1 : T}(\mathcal{S}_{\mathrm{observed}})$ denotes the collection of observed sequences from all known locations up to time $T$, and $\hat{y}_{t}(\boldsymbol{s}^\ast)$ represents the estimated value at an unknown location $\boldsymbol{s}^\ast$ for a future time step $t$.

3.2 Spatiotemporal Modeling Approach

To achieve the aforementioned mapping objective, this study integrates the S4 layer of the $\text{SSSD}^{\text{S4}}$ model with the MRTS basis functions, hereafter referred to as $\text{SSSD}^{\text{S4+MRTS}}$. This design enables the model to account for spatial dependence while performing convolution-based training on temporal sequences. During the inference stage, the predicted values at observed locations generated by $\text{SSSD}^{\text{S4+MRTS}}$, together with the spatial coordinates of unobserved sites, are subsequently provided to AFRK. By jointly incorporating information from both spatial and temporal dimensions, AFRK estimates the target values at unobserved locations.

3.2.1 Temporal Modeling Based on $\text{SSSD}^{\text{S4}}$

To effectively capture the temporal dependencies inherent in the data, the $\text{SSSD}^{\text{S4}}$ model is adopted as the temporal feature extractor. Its objective is to learn latent temporal features that simultaneously encode long-term dependencies and local dynamics from incomplete or noisy time series, leveraging the combination of the diffusion model and the S4 layers.

Using the observations at the observed locations $\mathcal{S}_{\mathrm{observed}}$ as the training basis, the data are first standardized as \begin{align} \tilde{y}^{(T)}_t(\boldsymbol{s}) = \frac{y_t(\boldsymbol{s}) - \mu(\boldsymbol{s})}{\sigma(\boldsymbol{s})}, \qquad \boldsymbol{s} \in \mathcal{S}_{\mathrm{observed}}, \end{align} where $\mu(\boldsymbol{s})$ and $\sigma(\boldsymbol{s})$ denote the mean and standard deviation of the time series at location $\boldsymbol{s}$, defined as \begin{align} \mu(\boldsymbol{s}) & = \frac{1}{T} \sum_{t=1}^{T} y_t(\boldsymbol{s}), \\ \sigma(\boldsymbol{s}) & = \sqrt{\frac{1}{T}\sum_{t=1}^{T} \left(y_t(\boldsymbol{s})-\mu(\boldsymbol{s})\right)^2 }. \end{align} The standardized input sequences $\tilde{y}^{(T)}_t(\boldsymbol{s})$ are then fed into the $\text{SSSD}^{\text{S4}}$ model to learn the temporal dependency structure.

To effectively incorporate spatial structural information into the temporal modeling process, MRTS is introduced as one of the conditional inputs and is jointly used with observational features and masking information to form the conditional channels of the S4 layer. For each spatial location $\boldsymbol{s} \in \mathcal{S}_{\mathrm{observed}}$, its MRTS feature vector is defined as $\boldsymbol{\phi}(\boldsymbol{s}) \in \mathbb{R}^{D_{\mathrm{MRTS}}}$. The MRTS basis matrix composed of all observed locations is given by \begin{align} \boldsymbol{\Phi} = \begin{bmatrix} \boldsymbol{\phi}(\boldsymbol{s}_1)^\top\\ \boldsymbol{\phi}(\boldsymbol{s}_2)^\top\\ \vdots\\ \boldsymbol{\phi}(\boldsymbol{s}_N)^\top \end{bmatrix} \in \mathbb{R}^{N \times D_{\mathrm{MRTS}}}, \end{align} where the MRTS feature dimension is specified as \begin{align} D_{\mathrm{MRTS}} = \min \left( N, \operatorname{round}\left( 10\sqrt{N} \right) \right), \end{align} and $N$ denotes the number of observed locations. This heuristic setting follows the practical recommendation adopted in AFRK-related implementations (Tzeng, Huang, Wang, Nychka, et al., 2021). This design provides a practical balance between spatial representation capability and computational efficiency while preventing an excessive number of basis functions when the number of observed locations is large.

To integrate the MRTS basis into the temporal model, $\boldsymbol{\Phi}$ is treated as a spatial basis that does not vary over time. It is then broadcast along the temporal dimension so that it aligns with the time-series length $T$. The resulting expanded tensor is \begin{align} \tilde{\boldsymbol{\Phi}} \in \mathbb{R}^{N \times D_{\mathrm{MRTS}} \times T}. \end{align}

By concatenating the spatial features with the observational data $\tilde{y}^{(T)}_t(\boldsymbol{s})$ and the missing-data mask $\boldsymbol{m}_t(\boldsymbol{s})$, the conditional channels of the S4 layer are extended from the original form $\left[ \tilde{y}^{(T)}_t(\boldsymbol{s}), \boldsymbol{m}_t(\boldsymbol{s}) \right]$ to \begin{align} \mathbf{c}_t(\boldsymbol{s}) = \left[ \tilde{y}^{(T)}_t(\boldsymbol{s}), \boldsymbol{m}_t(\boldsymbol{s}), \tilde{\boldsymbol{\Phi}}(\boldsymbol{s}) \right]. \end{align}

This conditional vector $\mathbf{c}_t(\boldsymbol{s})$, together with the noise input from the diffusion process, is supplied to the S4 layer to model temporal state transitions. Consequently, the number of conditional channels in the S4 layer is expanded from the original $2K$ to \begin{align} (2 + D_{\mathrm{MRTS}}) K, \end{align} and the corresponding conditional convolution layer is modified to \begin{align} \mathrm{Conv}\left((2 + D_{\mathrm{MRTS}})K, 2C\right), \end{align} where $K$ denotes the number of input channels in the original feature representation. This design allows the MRTS spatial basis features to be effectively incorporated into the subsequent state-space transformation processes within the S4 layer.

MRTS provides spatial basis information to the S4 layers through the conditional-channel mechanism, enabling spatial information to be incorporated into the diffusion process and allowing the model to capture temporal dependencies and spatial structural characteristics within a unified generative framework. After the reverse diffusion process, the model produces the temporally standardized prediction \begin{align} \hat{\tilde{y}}^{(T)}_t(\boldsymbol{s}), \qquad \boldsymbol{s}\in\mathcal S_{\mathrm{observed}}, \end{align} which is subsequently transformed back to the original measurement scale. Since the model is trained and inferred on standardized data, the final prediction must be transformed back to the original measurement scale using the inverse standardization: \begin{align} \hat{y}_{t}(\boldsymbol{s}) = \sigma(\boldsymbol{s}) \hat{\tilde{y}}^{(T)}_{t}(\boldsymbol{s}) + \mu(\boldsymbol{s}), \end{align} where $\mu(\boldsymbol{s})$ and $\sigma(\boldsymbol{s})$ denote the mean and standard deviation of the time series at location $\boldsymbol{s}$, respectively. This transformation restores the model outputs to the original measurement scale, facilitating subsequent analyses and spatial interpolation applications.

The SSSD model is trained using the MSE as the loss function, which iteratively updates the parameters $\theta$ to ensure that the extracted temporal features are stable and predictive. For tasks such as missing value imputation and multi-step future forecasting, the temporal representations learned by $\text{SSSD}^{\text{S4}}$ provide discriminative sequential embeddings, while the MRTS-based spatial representations provide additional spatial information that may improve spatiotemporal estimation performance.

3.2.2 Spatial Modeling Based on AFRK

Within the proposed integrated framework, AFRK treats the temporal predictions at observed locations generated by $\text{SSSD}^{\text{S4+MRTS}}$, denoted as $\hat{y}_{t}(\boldsymbol{s})$ for $\boldsymbol{s} \in \mathcal{S}_{\mathrm{observed}}$, as input information along the temporal dimension. To normalize the spatial distribution of predictions at each time point, the predictions are standardized across observed locations: \begin{align} \tilde{y}^{(S)}_t(\boldsymbol{s}) = \frac{\hat{y}_t(\boldsymbol{s}) - \bar{y}_t}{\tau_t}, \qquad \boldsymbol{s} \in \mathcal{S}_{\mathrm{observed}}, \end{align} where $\bar{y}_t$ and $\tau_t$ denote the spatial mean and spatial standard deviation at time $t$ over all observed locations $\mathcal{S}_{\mathrm{observed}}$, defined as \begin{align} \bar{y}_t & = \frac{1}{N} \sum_{\boldsymbol{s} \in \mathcal{S}_{\mathrm{observed}}} \hat{y}_t(\boldsymbol{s}), \\ \tau_t & = \sqrt{\frac{1}{N} \sum_{\boldsymbol{s} \in \mathcal{S}_{\mathrm{observed}}} \left(\hat{y}_t(\boldsymbol{s}) - \bar{y}_t\right)^2}. \end{align}

The standardized values $\tilde{y}^{(S)}_t(\boldsymbol{s})$ are then used as input to AFRK to perform conditional estimation at unobserved locations $\boldsymbol{s}^\ast \in \mathcal{S}_{\mathrm{unobserved}}$. On the standardized scale, the predictions at unobserved locations can be expressed as \begin{align} \hat{\tilde{y}}^{(S)}_t(\boldsymbol{s}^\ast) = \boldsymbol{f}(\boldsymbol{s}^\ast)^\top \hat{\boldsymbol{w}}_t + \hat{\xi}_t(\boldsymbol{s}^\ast), \end{align} where $\hat{\boldsymbol{w}}_t$ and $\hat{\xi}_t(\boldsymbol{s}^\ast)$ are the spatial conditional distribution parameters estimated by AFRK based on the observed location information $\tilde{y}^{(S)}_t(\boldsymbol{s})$.

Finally, the spatial predictions are transformed back to the original measurement scale via inverse standardization: \begin{align} \hat{y}_t(\boldsymbol{s}^\ast) = \tau_t \hat{\tilde{y}}^{(S)}_t(\boldsymbol{s}^\ast) + \bar{y}_t. \end{align} This step completes the reconstruction and forecasting of the global spatiotemporal field.

3.3 Model Training and Inference Algorithms

This section summarizes the operational workflow of the proposed spatiotemporal integration framework. Algorithm 1 illustrates the training procedure of $\text{SSSD}^{\text{S4+MRTS}}$, while Algorithm 2 presents the complete inference pipeline in which $\text{SSSD}^{\text{S4+MRTS}}$ is integrated with AFRK.

Algorithm 1 Training
Require: $y_t(\boldsymbol{s}), \boldsymbol{s} \in \mathcal{S}_{\mathrm{observed}}$; Max steps $T$; Learning rate $\eta$; Model parameters $\theta$
1. $\tilde{y}^{(T)}_t(\boldsymbol{s}) = \left(y_t(\boldsymbol{s}) - \mu(\boldsymbol{s})\right) / \sigma(\boldsymbol{s})$
2. $\tilde{\boldsymbol{\Phi}} \in \mathbb{R}^{N \times D_{\mathrm{MRTS}} \times T}$
3. repeat
4. $t \sim \text{Uniform}(\{1, \dots, T\}), \quad \boldsymbol{\epsilon} \sim \mathcal{N}(0, \boldsymbol{I})$
5. $\boldsymbol{x}_t = \sqrt{\bar{\alpha}_t} \tilde{y}^{(T)}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}$
6. $\mathbf{c}_t = \left[\tilde{y}^{(T)}_t(\boldsymbol{s}), \boldsymbol{m}_t(\boldsymbol{s}), \tilde{\boldsymbol{\Phi}}(\boldsymbol{s})\right]$
7. $\hat{\boldsymbol{\epsilon}}_{\theta} = g_{\theta}(\boldsymbol{x}_t, \mathbf{c}_t)$
8. $\mathcal{L} = \| \boldsymbol{\epsilon} - \hat{\boldsymbol{\epsilon}}_{\theta} \|^2$
9. $\theta \leftarrow \theta - \eta \nabla_{\theta} \mathcal{L}$
10. until convergence
Algorithm 2 Inference
Require: $y_t(\boldsymbol{s}), \boldsymbol{s} \in \mathcal{S}_{\mathrm{observed}}$; $\boldsymbol{s}^* \in \mathcal{S}_{\mathrm{unobserved}}$; Trained parameters $\theta$
1. $\tilde{y}^{(T)}_t(\boldsymbol{s}) = \left(y_t(\boldsymbol{s}) - \mu(\boldsymbol{s})\right) / \sigma(\boldsymbol{s})$
2. $\tilde{\boldsymbol{\Phi}} \in \mathbb{R}^{N \times D_{\mathrm{MRTS}} \times T}$
3. $\boldsymbol{x}_T \sim \mathcal{N}(0, \boldsymbol{I})$
4. for $t = T, \dots, 1$ do
5. $\mathbf{c}_t(\boldsymbol{s}) = \left[\tilde{y}^{(T)}_t(\boldsymbol{s}), \boldsymbol{m}_t(\boldsymbol{s}), \tilde{\boldsymbol{\Phi}}(\boldsymbol{s})\right]$
6. $\hat{\boldsymbol{\epsilon}}_{\theta} = g_{\theta}(\boldsymbol{x}_t, \mathbf{c}_t)$
7. $\boldsymbol{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( \boldsymbol{x}_t - \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}} \hat{\boldsymbol{\epsilon}}_{\theta} \right) + \sqrt{1-\alpha_t} \mathbf{z}, \quad \mathbf{z} \sim \mathcal{N}(0, \boldsymbol{I})$
8. end for
9. $\hat{\tilde{y}}^{(T)}_t(\boldsymbol{s}) = \boldsymbol{x}_0$
10. $\hat{y}_t(\boldsymbol{s}) = \sigma(\boldsymbol{s}) \hat{\tilde{y}}^{(T)}_t(\boldsymbol{s}) + \mu(\boldsymbol{s}), \quad \boldsymbol{s} \in \mathcal{S}_{\mathrm{observed}}$
11. $\tilde{y}^{(S)}_t(\boldsymbol{s}) = (\hat{y}_t(\boldsymbol{s}) - \bar{y}_t) / \tau_t, \quad \boldsymbol{s} \in \mathcal{S}_{\mathrm{observed}}$
12. $\hat{\tilde{y}}^{(S)}_t(\boldsymbol{s}^*) = \text{AFRK}(\tilde{y}^{(S)}_t(\boldsymbol{s}) \mid \boldsymbol{s}^*), \quad \forall \boldsymbol{s}^* \in \mathcal{S}_{\mathrm{unobserved}}$
13. $\hat{y}_t(\boldsymbol{s}^*) = \tau_t \hat{\tilde{y}}^{(S)}_t(\boldsymbol{s}^*) + \bar{y}_t$
14. return Full field $\hat{y}_t(\mathcal{S}_{\mathrm{observed}} \cup \mathcal{S}_{\mathrm{unobserved}})$

4 Experiments

This chapter presents the experimental design and setup for the spatiotemporal integration framework, aiming to evaluate its effectiveness in forecasting unobserved spatiotemporal values under heterogeneous observational settings. The experiments are conducted using three primary datasets, with detailed descriptions of model configurations, training and inference procedures, as well as the characteristics and partitioning of the datasets, in order to demonstrate the framework’s performance under different data conditions.

4.1 Datasets

This study adopts two representative meteorological datasets, covering in-situ ground station observations and global reanalysis data, as well as one synthetic dataset. These datasets are selected to evaluate spatiotemporal modeling performance under heterogeneous observational settings, including irregular station-based measurements, globally gridded reanalysis fields, and controlled synthetic spatial processes.

4.1.1 Weather2K

Weather2K is a recently proposed multivariate ground-based observation benchmark dataset, comprising measurements from thousands of meteorological stations across China, with a temporal resolution of 3 hours and covering near-surface meteorological variables such as air temperature, air pressure, humidity, and wind speed (Zhu et al., 2023). The open-source version, Weather2K-R (hereafter Weather2K), contains 1,866 stations and 13,632 consecutive time steps, featuring complete and regularly sampled time series without missing values. The dataset also provides constant location information (latitude, longitude, and elevation), which facilitates spatiotemporal modeling.

The Weather2K dataset exhibits heterogeneous spatial distribution and uneven station density, encompassing various geographical environments and climatic conditions, including urban areas, plains, and mountainous regions, thus providing rich spatiotemporal variation signals. These characteristics make Weather2K suitable for spatiotemporal interpolation, short-term sequence forecasting, and evaluations of model generalization and robustness.

This study utilizes observational sequences from the Weather2K-R dataset spanning from 00:00 on July 1, 2021, to 21:00 on August 31, 2021. Air Temperature is selected as the target variable. Detailed descriptions of the variables are provided in Appendix A.

4.1.2 MERRA-2

MERRA-2 (Modern-Era Retrospective Analysis for Research and Applications, Version~2) is a global atmospheric reanalysis dataset provided by the NASA Goddard Earth Sciences Data and Information Services Center (GES DISC). It reconstructs the state of the global atmosphere since 1980 through data assimilation techniques that integrate numerical weather prediction models with multi-source observations, primarily from satellites (GMAO, 2015). In this study, we use the hourly, single-level, instantaneous assimilation diagnostic product M2I1NXASM (Version 5.12.4) as the analysis dataset.

MERRA-2 is archived and managed by the Distributed Active Archive Center (DAAC) at NASA Goddard Space Flight Center, providing globally consistent reanalysis data. Its data assimilation process integrates satellite, ground-based, and remote sensing observations, with significant improvements over its predecessor MERRA in representing physical processes such as aerosols, radiative balance, and the hydrological cycle. It has been widely applied in climate analysis, atmospheric diagnostics, and model evaluation. The M2I1NXASM product offers hourly instantaneous data with high resolution and representative climate signals, making it suitable for short-term spatiotemporal forecasting and statistical feature analysis.

This study selects data from the M2I1NXASM (Version 5.12.4) product of the MERRA-2 dataset, spanning from 00:00 on December 11, 2023 to 23:00 on December 31, 2023. The Surface Skin Temperature variable is used as the primary data source. Detailed descriptions of the variables are provided in Appendix B. The experimental domain is defined as a regional subset extracted from the global reanalysis field, bounded by $26.0^\circ$ to $48.0^\circ$N latitude and $70.0^\circ$ to $123.0^\circ$W longitude.

4.1.3 KAUST Spatial Statistics Competition Dataset

The KAUST Spatial Statistics Competition Dataset originates from the 2022 KAUST Spatial Statistics Competition, which was designed to evaluate prediction accuracy and computational efficiency in large-scale spatial and spatiotemporal data analysis (Abdulah, Alamri, Nag, et al., 2022). The dataset consists of simulated spatial observations generated under controlled settings and was specifically designed to evaluate the interpolation performance of complex spatial and spatiotemporal models under high-dimensional and large-scale scenarios with intricate spatial dependence structures.

In this study, Configuration 2a-7 of the competition dataset is adopted for evaluation (Abdulah, Alamri, Ltaief, et al., 2022). A detailed description of the dataset and simulation settings is provided in Appendix C.

4.2 Environment and Computational Resources

To ensure the feasibility and reliability of training and inference experiments of the proposed spatiotemporal integration framework on large-scale datasets, a unified computational environment was established. Multiple existing software packages were integrated to support model development, training, and evaluation.

The SSSD model has been implemented in Python by its original authors (AI4HealthUOL, 2023), which is capable of effectively capturing long-range dependency structures in time-series data. AFRK was implemented in the R programming language by Wen-Ting Wang and released as the autoFRK package (Tzeng, Huang, Wang, Nychka, et al., 2021), providing stable and scalable spatial interpolation capabilities. To integrate these functionalities, this study reimplemented and encapsulated autoFRK as a Python package (Tzeng, Huang, Wang, and Hsu, 2025). The algorithm was further integrated with the original SSSD implementation, enabling a complete spatiotemporal modeling workflow within the PyTorch framework in Python.

All experiments, including model training, validation, and inference, were conducted on the Taiwania 2 high-performance computing platform (NCHC, 2018). Taiwania 2 provides high-performance GPU computing resources, along with large-capacity memory and high-speed storage systems. These resources enable efficient processing of high-resolution datasets and long time series while ensuring computational stability and reproducibility of the experimental results.

4.3 Experimental Design

To clearly present the configurations of $\text{SSSD}^{\text{S4+MRTS}}$ and autoFRK, the following sections summarize the hyperparameters for training and inference as shown in Table 2.

Table 2: Model Hyperparameter Settings.

ModelHyperparameterValue
Training Configuration
Batch size40
Learning rate0.001
Only generate missingTrue
MaskingForecast
Missing $k$
SSSDS4
WaveNetInput channels1
Output channels1
Residual layers32
Residual channels64
Skip channels64
Diffusion step embeddingInput dimension128
Hidden dimension256
Output dimension256
S4Max sequence length
State dimension128
Dropout0.1
BidirectionalTrue
Layer normalizationTrue
DiffusionDiffusion steps ($T$)100
$\beta_0$0.0001
$\beta_T$0.05
MRTSUse MRTSTrue
autoFRK
MethodFast

Table 3 summarizes the experimental parameters and dataset configurations used to compare the performance of different models. In all experiments, 200 coordinate points are randomly sampled from the spatial extent of each dataset and split into observed and unobserved subsets containing 160 and 40 locations, respectively. Since the datasets differ in temporal resolution, with MERRA-2 providing hourly observations and Weather2K providing 3-hourly observations, the input sequence lengths and forecasting horizons are determined based on complete-day units rather than enforcing identical temporal proportions across datasets. Consequently, the temporal partitions are selected to remain approximately consistent with a 9:1 training-testing ratio while aligning the sequence boundaries with complete-day intervals under the native temporal resolution of each dataset. This design avoids fractional-day prediction windows and facilitates a more meaningful comparison across datasets with heterogeneous sampling frequencies. The spatial sampling regions of the datasets are illustrated in Figure 2, Figure 4, and Figure 6, respectively.

Table 3: Experimental parameter settings and dataset configurations (Weather2K / MERRA-2 / 2a-7).

ParameterValue
Number of iterations500
Observed locations160
Unobserved locations40
Input sequence length448 / 456 / 90
Prediction horizon $k$48 / 48 / 10

https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/Stations%20of%20Weather2k%203-Hourly%20Data.png
Figure 2: Spatiotemporal observation distribution of the air temperature variable in the Weather2K dataset. The figure presents temperature variations at observation stations over five consecutive time steps starting from 00:00 on July 1, 2021, and demonstrates the heterogeneous and irregular spatial sampling patterns across different geographical regions.
https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/Weather2k/sample_locations/sample_locations.png
Figure 3: Representative observations from two spatial locations in the Weather2K dataset. The blue segments denote historical inputs used for modeling, while the red segments indicate the forecasting horizon. The figure highlights location-dependent temporal dynamics and variability in the observed processes.
https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/Stations%20of%20MERRA-2%20Hourly%20Data.png
Figure 4: Spatiotemporal distribution of surface skin temperature observations in the MERRA-2 dataset. The figure illustrates the reanalysis fields over five consecutive time steps starting from 00:00 on December 11, 2023, together with the regular spatial grid structure adopted for spatiotemporal reconstruction and forecasting tasks.
https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/MERRA2/sample_locations/sample_locations.png
Figure 5: Representative surface temperature trajectories from two spatial grid locations in the MERRA-2 dataset. Historical observations (blue) and future targets (red) are shown, illustrating structured temporal evolution driven by large-scale atmospheric dynamics.
Figure 6: Spatiotemporal observation distribution of the KAUST 2a-7 simulated dataset, presenting the simulated field values over the first five time steps and the controlled yet irregular spatial sampling pattern of observation locations for spatiotemporal evaluation.
https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/2a7/sample_locations/sample_locations.png
Figure 7: Representative time series from two spatial locations in the KAUST 2a-7 dataset. The blue and red segments denote historical observations and evaluation horizon respectively, highlighting stochastic temporal fluctuations with weak deterministic structure.

We evaluate the impact of incorporating MRTS-based spatial information by comparing $\text{SSSD}^{\text{S4}}$ and $\text{SSSD}^{\text{S4+MRTS}}$. Specifically, we compare the baseline $\text{SSSD}^{\text{S4}}$ model with the MRTS-enhanced variant $\text{SSSD}^{\text{S4+MRTS}}$ to assess whether MRTS strengthens spatial dependency modeling and improves the imputation of unobserved locations during inference. In addition, this study includes several representative baselines, namely TFT, VAR, SVGP, and STDK, to provide a comprehensive comparison across different model architectures and learning paradigms for spatiotemporal forecasting tasks.

5 Experimental Results

This study adopts the Mean Squared Prediction Error (MSPE) as the primary evaluation metric to measure predictive accuracy. Table 4 reports the results of an integrated spatiotemporal framework, in which models trained with $\text{SSSD}^{\text{S4+MRTS}}$ and $\text{SSSD}^{\text{S4}}$ are both combined with AFRK for spatial field imputation during inference.

The reported results are averaged over 30 independent runs, using random seeds ranging from 42 to 71 across different datasets. The table also includes several baseline models for comparison. Since some baselines are limited to time-series forecasting, AFRK is applied to impute values at unobserved spatial locations, enabling a consistent evaluation of their performance in spatiotemporal prediction tasks.

Table 4: Average MSPE performance of different models under various datasets and forecasting scenarios.

Model$\text{SSSD}^{\text{S4+MRTS}}$$\text{SSSD}^{\text{S4}}$TFTVARSVGPSTDK
Spatial ImputationAFRKAFRKAFRKAFRK
Weather2K
Unobserved & Future15.163918.572820.194223.428446.679335.1823
Unobserved & Past7.06737.07447.08777.091034.448327.0631
Observed & Future11.085814.565018.343419.256429.902234.9767
MERRA-2
Unobserved & Future12.907413.295716.5911389.56263.0083e4111.0262
Unobserved & Past6.47876.48116.48676.49371.6750e4111.2048
Observed & Future7.45127.468511.9694459.52762.8654e4107.9238
2a-7
Unobserved & Future0.89950.89870.89900.92710.97831.1991
Unobserved & Past0.90910.90930.90920.90980.99031.2128
Observed & Future1.02871.07521.59842.17030.91651.1642

Table 4 shows that, after spatial imputation using AFRK, $\text{SSSD}^{\text{S4}}$ already outperforms most baseline models. When MRTS is incorporated, $\text{SSSD}^{\text{S4+MRTS}}$ combined with AFRK achieves lower MSPE than $\text{SSSD}^{\text{S4}}$ + AFRK in most forecasting scenarios, particularly on the Weather2K and MERRA-2 datasets. These results suggest that incorporating MRTS can provide additional spatial information that improves predictive accuracy, especially for forecasting tasks involving unobserved locations.

Table 5: Mean and standard deviation of MSPE over 30 independent runs for the proposed model and its counterpart without MRTS.

Scenario$\text{SSSD}^{\text{S4+MRTS}}$ + AFRK$\text{SSSD}^{\text{S4}}$ + AFRK
Weather2K
Unobserved & Future15.1639 $\pm$ 3.043418.5728 $\pm$ 3.7072
Unobserved & Past7.0673 $\pm$ 2.79487.0744 $\pm$ 2.7995
Observed & Future11.0858 $\pm$ 0.649414.5650 $\pm$ 1.4134
MERRA-2
Unobserved & Future12.9074 $\pm$ 2.330113.2957 $\pm$ 3.0640
Unobserved & Past6.4787 $\pm$ 1.94636.4811 $\pm$ 1.9348
Observed & Future7.4512 $\pm$ 0.62537.4685 $\pm$ 1.0986
2a-7
Unobserved & Future0.8995 $\pm$ 0.07750.8987 $\pm$ 0.0753
Unobserved & Past0.9091 $\pm$ 0.03860.9093 $\pm$ 0.0390
Observed & Future1.0287 $\pm$ 0.04131.0752 $\pm$ 0.0684

To further evaluate the stability of the proposed framework, the mean and standard deviation of MSPE over 30 independent runs were computed for $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK and $\text{SSSD}^{\text{S4}}$ + AFRK. The results are summarized in Table 5, where lower standard deviations indicate more consistent predictive performance under different random sampling configurations.

As shown in Table 5, $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK generally achieves lower average MSPE than $\text{SSSD}^{\text{S4}}$ + AFRK across most experimental scenarios, while maintaining comparable or smaller standard deviations. This observation suggests that incorporating MRTS-based spatial representations may improve predictive accuracy without introducing additional variability across different spatial sampling configurations.

The model abbreviations used across all subsequent figures are defined as follows: A denotes $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK; B denotes $\text{SSSD}^{\text{S4}}$ + AFRK; C denotes TFT + AFRK; D denotes VAR + AFRK; E denotes SVGP; and F denotes STDK. These abbreviations are consistently adopted in the following figures to facilitate comparison among models.

Figure 8, Figure 9, and Figure 10 present boxplots of MSPE distributions for different models under various forecasting tasks on the Weather2K dataset. In the future prediction task at unobserved locations, $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK exhibits a notably more concentrated distribution with lower error values. Its median and lower quartile are consistently better than those of $\text{SSSD}^{\text{S4}}$ + AFRK and other baseline models. The lower prediction errors suggest that the MRTS-based spatial representations provide useful information for modeling spatially dependent forecasting patterns.

In addition, for the future prediction task at observed locations presented in Figure 10, $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK also yields lower MSPE compared to $\text{SSSD}^{\text{S4}}$ + AFRK. This suggests that incorporating MRTS may also provide benefits for forecasting tasks at observed locations.

https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/Weather2k/Unobs_and_Future/Unobs_and_Future_SSSD_S4+MRTS+AFRK_SSSD_S4+AFRK_TFT+AFRK_VAR+AFRK_SVGP_STDK_MSPE_BoxGrid1xn_nickname.png
Figure 8: Boxplots of MSPE distributions for different models in the future prediction task at unobserved locations on the Weather2K dataset.
https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/Weather2k/Unobs_and_Past/Unobs_and_Past_SSSD_S4+MRTS+AFRK_SSSD_S4+AFRK_TFT+AFRK_VAR+AFRK_SVGP_STDK_MSPE_BoxGrid1xn_nickname.png
Figure 9: Boxplots of MSPE distributions for different models in the past prediction task at unobserved locations on the Weather2K dataset.
https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/Weather2k/Obs_and_Future/Obs_and_Future_SSSD_S4+MRTS+AFRK_SSSD_S4+AFRK_TFT+AFRK_VAR+AFRK_SVGP_STDK_MSPE_BoxGrid1xn_nickname.png
Figure 10: Boxplots of MSPE distributions for different models in the future prediction task at observed locations on the Weather2K dataset.

As shown in Figure 11, Figure 12, and Figure 13, the differences in MSPE across models become even more pronounced on the MERRA-2 dataset under various forecasting tasks. Models such as VAR + AFRK, SVGP, and STDK exhibit substantially poorer performance on MERRA-2, particularly in the future prediction task at unobserved locations, where their MSPE distributions display extremely large error values. This indicates that these models are unable to effectively capture the complex spatial structures and dynamic variability inherent in large-scale, high-dimensional climate data, resulting in predictions that deviate severely from the true values.

In contrast, although $\text{SSSD}^{\text{S4}}$ + AFRK yields higher MSPE on MERRA-2 compared to Weather2K, it still maintains relatively stable predictive performance. Incorporating MRTS further improves the model, $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK achieves a noticeably reduced MSPE range and a lower median. This demonstrates that MRTS substantially enhances the model’s ability to capture large-scale spatial structures, enabling more accurate predictions when dealing with complex climate datasets.

https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/MERRA2/Unobs_and_Future/Unobs_and_Future_SSSD_S4+MRTS+AFRK_SSSD_S4+AFRK_TFT+AFRK_VAR+AFRK_SVGP_STDK_MSPE_BoxGrid1xn_nickname.png
Figure 11: Boxplots of MSPE distributions for different models in the future prediction task at unobserved locations on the MERRA-2 dataset.
https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/MERRA2/Unobs_and_Past/Unobs_and_Past_SSSD_S4+MRTS+AFRK_SSSD_S4+AFRK_TFT+AFRK_VAR+AFRK_SVGP_STDK_MSPE_BoxGrid1xn_nickname.png
Figure 12: Boxplots of MSPE distributions for different models in the past prediction task at unobserved locations on the MERRA-2 dataset.
https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/MERRA2/Obs_and_Future/Obs_and_Future_SSSD_S4+MRTS+AFRK_SSSD_S4+AFRK_TFT+AFRK_VAR+AFRK_SVGP_STDK_MSPE_BoxGrid1xn_nickname.png
Figure 13: Boxplots of MSPE distributions for different models in the future prediction task at observed locations on the MERRA-2 dataset.

Figure 14, Figure 15, and Figure 16 present boxplots of the MSPE distributions for three models, $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK, $\text{SSSD}^{\text{S4}}$ + AFRK, and TFT + AFRK, on the MERRA-2 dataset. Across all forecasting tasks, the MSPE distribution of $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK is clearly more concentrated within a lower error range, with all quartiles outperforming the other two models. This pattern indicates that MRTS substantially strengthens the model’s ability to capture spatial structure.

Moreover, in the future prediction task at observed locations, $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK demonstrates an even more pronounced reduction in MSPE compared to both $\text{SSSD}^{\text{S4}}$ + AFRK and TFT + AFRK, further confirming the effectiveness of MRTS in enhancing predictive accuracy under complex spatiotemporal conditions.

https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/MERRA2/Unobs_and_Future/Unobs_and_Future_SSSD_S4+MRTS+AFRK_SSSD_S4+AFRK_TFT+AFRK_MSPE_BoxGrid1xn_nickname.png
Figure 14: Boxplots of MSPE distributions for $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK, $\text{SSSD}^{\text{S4}}$ + AFRK, and TFT + AFRK in the future prediction task at unobserved locations on the MERRA-2 dataset.
https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/MERRA2/Unobs_and_Past/Unobs_and_Past_SSSD_S4+MRTS+AFRK_SSSD_S4+AFRK_TFT+AFRK_MSPE_BoxGrid1xn_nickname.png
Figure 15: Boxplots of MSPE distributions for $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK, $\text{SSSD}^{\text{S4}}$ + AFRK, and TFT + AFRK in the past prediction task at unobserved locations on the MERRA-2 dataset.
https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/MERRA2/Obs_and_Future/Obs_and_Future_SSSD_S4+MRTS+AFRK_SSSD_S4+AFRK_TFT+AFRK_MSPE_BoxGrid1xn_nickname.png
Figure 16: Boxplots of MSPE distributions for $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK, $\text{SSSD}^{\text{S4}}$ + AFRK, and TFT + AFRK in the future prediction task at observed locations on the MERRA-2 dataset.

Figure 17, Figure 18, and Figure 19 present the boxplots of MSPE distributions for all models across different forecasting tasks on the 2a-7 dataset. In the prediction tasks at unobserved locations, the MSPE distributions of the models are highly similar, indicating that under this simulated dataset, characterized by relatively simple and regular spatial structures, $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK does not exhibit a pronounced advantage, yet remains comparable in accuracy to other spatial modeling approaches.

Furthermore, in the future prediction task at observed locations, the MSPE distribution of $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK is slightly higher than that of SVGP. This outcome may be attributed to the limited benefit MRTS provides in capturing spatial structure within this dataset, potentially coupled with an increased risk of overfitting. As a result, the predictive performance at observed locations shows a marginal decline.

https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/2a7/Unobs_and_Future/Unobs_and_Future_SSSD_S4+MRTS+AFRK_SSSD_S4+AFRK_TFT+AFRK_VAR+AFRK_SVGP_STDK_MSPE_BoxGrid1xn_nickname.png
Figure 17: Boxplots of MSPE distributions for different models in the future prediction task at unobserved locations on the 2a-7 dataset.
https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/2a7/Unobs_and_Past/Unobs_and_Past_SSSD_S4+MRTS+AFRK_SSSD_S4+AFRK_TFT+AFRK_VAR+AFRK_SVGP_STDK_MSPE_BoxGrid1xn_nickname.png
Figure 18: Boxplots of MSPE distributions for different models in the past prediction task at unobserved locations on the 2a-7 dataset.
https://josh-test-lab.github.io/posts/Spatiotemporal%20Prediction%20of%20Unknown%20Areas%20Based%20on%20Structured%20State%20Space%20Diffusion%20and%20Resolution%20Adaptive%20Fixed%20Rank%20Kriging/results/2a7/Obs_and_Future/Obs_and_Future_SSSD_S4+MRTS+AFRK_SSSD_S4+AFRK_TFT+AFRK_VAR+AFRK_SVGP_STDK_MSPE_BoxGrid1xn_nickname.png
Figure 19: Boxplots of MSPE distributions for different models in the future prediction task at observed locations on the 2a-7 dataset.

By integrating the results in Table 4 with the MSPE boxplots for each dataset, it is evident that the proposed $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK does not necessarily achieve the lowest MSPE across all datasets and forecasting scenarios. Nevertheless, its overall performance remains highly competitive, consistently exhibiting stable and strong predictive capability under most experimental configurations. In particular, for real-world datasets such as Weather2K and MERRA-2, which feature complex spatial structures, $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK generally maintains lower errors and reduced variability. This observation indicates that MRTS effectively enhances the model’s ability to capture spatial dependencies and cross-regional variations.

In contrast, the performance gap between $\text{SSSD}^{\text{S4}}$ + AFRK and $\text{SSSD}^{\text{S4+MRTS}}$ + AFRK on the 2a-7 dataset is notably small. This suggests that when the dataset’s spatial structure is inherently simple and regular, the benefits introduced by MRTS may be less pronounced. Additionally, SVGP achieves the best performance in the future prediction task at observed locations on the 2a-7 dataset, highlighting that different model architectures may exhibit distinct suitability depending on the data distribution and forecasting conditions.

Overall, the proposed method demonstrates superior robustness and generalization capability on high-dimensional and highly variable datasets, while maintaining competitive performance comparable to other models on datasets with simpler structural characteristics.

6 Conclusion

This study addresses the problem of spatiotemporal forecasting under incomplete spatial observations by proposing an integrated framework that combines $\text{SSSD}^{\text{S4}}$, MRTS, and AFRK to enhance prediction capability at unobserved locations and future time points. Experimental results demonstrate that, on real-world datasets including Weather2K and MERRA-2, the proposed method effectively reduces prediction errors while maintaining more stable MSPE distributions, indicating its capability to model complex spatiotemporal dependencies. In comparison, conventional baseline models are more likely to suffer from error amplification and unstable prediction distributions when dealing with high-dimensional data with strong spatial heterogeneity, particularly in future forecasting tasks at unobserved locations.

Furthermore, the results reveal that the proposed integrated framework is not universally optimal for all types of spatiotemporal data, and its effectiveness depends on the characteristics of the underlying spatial structures. When the data exhibit strong spatial nonlinearity and cross-scale variations, MRTS effectively enhances the spatial representation capability of the S4 architecture, leading to improved forecasting performance. However, for datasets with smoother spatial structures or relatively simple variations, the benefits introduced by MRTS may become limited, and the proposed method may not necessarily outperform alternative approaches in certain scenarios. In addition, although AFRK effectively estimates spatial information at unobserved locations, its imputation capability relies on the spatial correlations captured from available observations. Therefore, when observations become increasingly sparse, missing ratios increase, or observed samples fail to adequately represent the overall spatial variability, the prediction performance may still be affected.

Similar to other data-driven approaches, this study assumes that training and testing data share similar statistical characteristics. Consequently, if future data distributions undergo significant changes due to extreme events, environmental changes, or other factors, model performance may degrade, requiring additional retraining or adaptation strategies to maintain prediction quality. Moreover, compared with the original $\text{SSSD}^{\text{S4}}$, the proposed framework introduces additional computational costs due to the integration of MRTS and AFRK during model construction and inference. Therefore, practical deployment requires careful consideration of the trade-off between predictive performance and computational efficiency. The selection of an appropriate model should ultimately depend on the spatial characteristics of the data, observation conditions, and application requirements.

This study demonstrates that deep state-space models combined with multi-scale spatial representation and inference-stage spatial imputation mechanisms can effectively improve the accuracy and robustness of spatiotemporal forecasting under incomplete spatial observations. Future research may further investigate the applicability of the proposed framework to multivariate datasets, different spatial resolutions, higher missing ratios, and more diverse environmental variables. In addition, developing spatiotemporal modeling approaches that achieve both high predictive performance and computational efficiency remains an important direction for improving generalization capability and practical applicability in environmental monitoring, weather forecasting, and geoscientific analysis.

References

Appendix

A Weather2K Dataset

The Weather2K-R dataset includes a set of meteorological variables. In this study, Air Temperature is selected as the primary variable for analysis.

Table 6: Variable List of Weather2K. Description of variables, their abbreviations, and measurement units.
VariableShort NameUnit
Latitudelatdegrees east
Longitudelondegrees north
Altitudealtm
Air PressureaphPa
Air Temperaturet°C
Maximum Temperaturemxt°C
Minimum Temperaturemnt°C
Relative Humidityrh%
Precipitation in 3hp3mm
Wind Directionwddegrees
Wind Speedwsm s$^{-1}$
Maximum Wind Directionmwddegrees
Maximum Wind Speedmwsm s$^{-1}$

B MERRA-2 Dataset

The MERRA-2 dataset includes a set of atmospheric variables. In this study, Surface Skin Temperature is selected as the primary variable for analysis.

Table 7: Variable List of MERRA-2. Description of variables, their abbreviations, and measurement units.
VariableShort NameUnit
Longitudelondegrees east
Latitudelatdegrees north
Timetimeminutes since 2024-06-01 00:00:00
2-Meter Air Temperaturet2mK
Total Precipitable Liquid Watertqlkg m$^{-2}$
Total Column Odd Oxygentoxkg m$^{-2}$
2-Meter Eastward Windu2mm s$^{-1}$
Surface PressurepsPa
Tropopause Temperature Using Blended TROPP EstimatetroptK
Northward Wind at 50 Metersv50mm s$^{-1}$
Zero Plane Displacement Heightdisphm
Total Column Ozoneto3Dobsons
Surface Skin TemperaturetsK
10-Meter Air Temperaturet10mK
Tropopause Pressure Based on Thermal EstimatetropptPa
Total Precipitable Ice Watertqikg m$^{-2}$
Sea Level PressureslpPa
Tropopause Pressure Based on Blended EstimatetroppbPa
Total Precipitable Water Vaportqvkg m$^{-2}$
2-Meter Northward Windv2mm s$^{-1}$
Tropopause Specific Humidity Using Blended TROPP Estimatetropqkg kg$^{-1}$
10-Meter Northward Windv10mm s$^{-1}$
Eastward Wind at 50 Metersu50mm s$^{-1}$
10-Meter Eastward Windu10mm s$^{-1}$
2-Meter Specific Humidityqv2mkg kg$^{-1}$
Tropopause Pressure Based on EPV EstimatetroppvPa
10-Meter Specific Humidityqv10mkg kg$^{-1}$

C Dataset of The Second Competition on Spatial Statistics for Large Datasets

This dataset originates from the 2022 KAUST Competition on Spatial Statistics for Large Datasets, which provides large-scale simulated datasets designed to evaluate the performance of spatial statistical methods under a unified experimental framework. The dataset is generated using the ExaGeoStat high-performance statistical computing framework and is based on reproducible Gaussian process (GP) simulations, offering standardized and comparable benchmarking conditions (Abdulah, Alamri, Nag, et al., 2022; Abdulah, Alamri, Ltaief, et al., 2022).

In Sub-competition 2a and 2b, a non-separable and stationary Gaussian process model is adopted, where the covariance structure follows the formulation proposed in Gneiting (2002). For any two spatial locations $s \in [0, 1]^2$ and temporal lag $t \in \mathbb{R}$, the covariance function is defined as: \begin{equation} C(\boldsymbol{h},u ; \boldsymbol{\theta}) = \frac{\sigma^{2}}{a_t |u|^{2\alpha} + 1} \, M_{\nu}\!\left( \frac{\lVert \boldsymbol{h} \rVert / a_s}{\left(a_t |u|^{2\alpha} + 1\right)^{\beta/2}} \right), \end{equation} where $\boldsymbol{h}$ denotes the spatial lag and $u$ denotes the temporal lag. The parameter $\sigma^{2} > 0$ represents the variance, $\nu > 0$ and $\alpha \in [0,1]$ control smoothness, $a_s, a_t > 0$ are spatial and temporal scaling parameters, and $\beta \in (0,1]$ governs the interaction strength between space and time. $M_{\nu}(\cdot)$ denotes the Matérn correlation function.

According to the competition protocol, Sub-competitions 2a and 2b generate a total of 18 datasets, covering multiple configurations of spatial scales (weak, moderate, and strong), spatial resolutions (1K and 10K locations), and temporal lengths (100 and 1000 time steps). Three types of missingness mechanisms are considered for prediction tasks: random spatial removal (RS), random spatiotemporal removal (RST), and systematic removal of the last 10 time steps (T10). All configurations are summarized in Table 1 of Abdulah, Alamri, Nag, et al. (2022). The 2a-7 dataset corresponds to a specific configuration within this benchmark, defined by a particular combination of spatial resolution, temporal length, and model parameters.

D Other Model Configurations

This chapter describes the main hyperparameter settings for the other baseline models used in the experiments, including TFT, VAR, SVGP, and STDK, for reference and replication of the experimental results. The implementation details can be found at https://github.com/Josh-test-lab/SSSDS4-AFRK.

D.1 TFT

This study implements the Temporal Fusion Transformer (TFT) model using the Darts library, with the main hyperparameter settings shown in Table 8.

Table 8: Hyperparameter configuration of the TFT model.

HyperparameterValue
Batch size40
Learning rate0.001
Epochs500
Add relative indexTrue
Input chunk length80
Hidden size32
LSTM layers1
Dropout0.1
Output size1

D.2 VAR

This study implements the Vector Autoregression (VAR) model using the statsmodels library, with the main hyperparameter settings shown in Table 9.

Table 9: Hyperparameter configuration of the VAR model.

HyperparameterValue
Maximum lag order240
Information criterionaic
Channel dimension1

D.3 SVGP

This study implements the Sparse Variational Gaussian Process (SVGP) model using the GPyTorch library (https://pypi.org/project/gpytorch/), with the main hyperparameter settings shown in Table 10.

Table 10: Hyperparameter configuration of the SVGP model.

HyperparameterValue
Learning rate0.01
Epochs500
OptimizerAdam
LikelihoodGaussian
DeviceCUDA / CPU
Kernel typeRBF Kernel
ARD dimensions3
Mean functionConstant Mean
Variational distributionCholesky Variational
Variational strategyLearn inducing points
Inducing points1024
Loss functionVariational ELBO
Likelihood noiseLearned Gaussian noise

D.4 STDK

This study implements the Spatio-temporal DeepKriging (STDK) model using the da-stdk library (https://pypi.org/project/da-stdk/), with the main hyperparameter settings shown in Table 11.

Table 11: Hyperparameter configuration of the STDK model.

HyperparameterValue
DeviceCUDA
Split methodRandom
Train ratio0.8
Calibration ratio0.2
Calibration split methodRandom
Observation methodSite-wise
Observation patternUniform
Observation ratio0.1
Spatial intensity10.0
Provider splitFalse
Learning rate0.001
Basis learning rate ratio0.05
Epochs500
Batch size40
Weight decay5e-4
Patience50
Gradient clipping10.0
SchedulerCosine
Warmup epochs10
Basis unfreeze epoch10
Basis LR ramp-up epochs10
Number of workers0
Hidden dimensions128, 128, 128
Spatial basis functionWendland
Spatial centers25, 81, 121
Temporal centers10, 15, 45
Spatial learnableFalse
Spatial init methodUniform
Delta reparameterizationFalse
Layer normalizationTrue
Gradient dampingTrue
Damping threshold0.0
Damping strength5.0
Domain penalty weight0.01
Movement penalty weight0.0
Sparsity typeNone
Sparsity L1 weight0.0
Sparsity group weight0.0
Spatial sparsity appliedTrue
Temporal sparsity appliedFalse
Sparsity threshold ratio0.01
Non-crossing lambda0.0
Non-crossing L1True
Input covariates0
Regression typeMulti-quantile
Quantile levels0.5
Conformal modeBoth
Conformal alpha0.1
Conformal center sourceTrained
Conformal cluster min N30
Save quantile predictionsTrue
Target normalizationTrue

See Also

1
2
3
4
5
6
7
8
@mastersthesis{Hsu2026SSSD_AFRK,
    author      = {HSU, YAO-CHIH},
    title       = {Spatiotemporal Prediction of Unknown Areas Based on Structured State Space Diffusion and Resolution Adaptive Fixed Rank Kriging},
    school      = {National Dong Hwa University},
    type        = {Master's thesis},
    year        = {2026},
    url         = {https://hdl.handle.net/11296/n96vkz},
}