Table of Contents
Fetching ...

STACI: Spatio-Temporal Aleatoric Conformal Inference

Brandon R. Feng, David Keetae Park, Xihaier Luo, Arantxa Urdangarin, Shinjae Yoo, Brian J. Reich

TL;DR

STACI tackles uncertainty quantification for non-stationary spatio-temporal fields by blending a variational Bayesian neural network that implements a spectral, non-stationary GP with a latent dimension-expansion, with a conformal inference layer that guarantees valid prediction intervals. The approach achieves scalable interpolation on large ST datasets by employing random Fourier features and Stein variational gradient descent, while preserving interpretability of covariance structure through latent factors. A ST-conformal step adapts intervals to local ST geometry, ensuring nominal coverage despite approximations. Empirical results on synthetic MSS and real AOD data demonstrate improved predictive accuracy and calibrated, efficient uncertainty quantification relative to scalable GP and deep-learning baselines, particularly when using a FFNP latent model.

Abstract

Fitting Gaussian Processes (GPs) provides interpretable aleatoric uncertainty quantification for estimation of spatio-temporal fields. Spatio-temporal deep learning models, while scalable, typically assume a simplistic independent covariance matrix for the response, failing to capture the underlying correlation structure. However, spatio-temporal GPs suffer from issues of scalability and various forms of approximation bias resulting from restrictive assumptions of the covariance kernel function. We propose STACI, a novel framework consisting of a variational Bayesian neural network approximation of non-stationary spatio-temporal GP along with a novel spatio-temporal conformal inference algorithm. STACI is highly scalable, taking advantage of GPU training capabilities for neural network models, and provides statistically valid prediction intervals for uncertainty quantification. STACI outperforms competing GPs and deep methods in accurately approximating spatio-temporal processes and we show it easily scales to datasets with millions of observations.

STACI: Spatio-Temporal Aleatoric Conformal Inference

TL;DR

STACI tackles uncertainty quantification for non-stationary spatio-temporal fields by blending a variational Bayesian neural network that implements a spectral, non-stationary GP with a latent dimension-expansion, with a conformal inference layer that guarantees valid prediction intervals. The approach achieves scalable interpolation on large ST datasets by employing random Fourier features and Stein variational gradient descent, while preserving interpretability of covariance structure through latent factors. A ST-conformal step adapts intervals to local ST geometry, ensuring nominal coverage despite approximations. Empirical results on synthetic MSS and real AOD data demonstrate improved predictive accuracy and calibrated, efficient uncertainty quantification relative to scalable GP and deep-learning baselines, particularly when using a FFNP latent model.

Abstract

Fitting Gaussian Processes (GPs) provides interpretable aleatoric uncertainty quantification for estimation of spatio-temporal fields. Spatio-temporal deep learning models, while scalable, typically assume a simplistic independent covariance matrix for the response, failing to capture the underlying correlation structure. However, spatio-temporal GPs suffer from issues of scalability and various forms of approximation bias resulting from restrictive assumptions of the covariance kernel function. We propose STACI, a novel framework consisting of a variational Bayesian neural network approximation of non-stationary spatio-temporal GP along with a novel spatio-temporal conformal inference algorithm. STACI is highly scalable, taking advantage of GPU training capabilities for neural network models, and provides statistically valid prediction intervals for uncertainty quantification. STACI outperforms competing GPs and deep methods in accurately approximating spatio-temporal processes and we show it easily scales to datasets with millions of observations.

Paper Structure

This paper contains 22 sections, 2 theorems, 23 equations, 3 figures, 3 tables.

Key Result

Theorem 1

The prior mean of the spatiotemporal covariance function of the discrete process in (eq: brff_st) equals the $\hbox{Mat$\acute{\hbox{e}}$rn}$ correlation with distance defined as in (eq: de) for all $J$, and the point-wise prior variance decreases at rate $J$.

Figures (3)

  • Figure 1: STACI algorithm pipeline. Data from coordinates across space and time are first fed into our approximate spatio-temporal GP architecture and trained using SVGD. A conformal inference step is then fit to provide customized prediction intervals for valid uncertainty quantification.
  • Figure 2: Predicted AOD surface values. Top row: predicted surface. Red indicates higher AOD values.; Bottom row: interval widths for Bayesian and conformal (STACI-C) uncertainty on AOD data. Darker shades denote narrower intervals.
  • Figure 3: Ablation Studies for STACI

Theorems & Definitions (2)

  • Theorem 1
  • Theorem 2