Accelerated Spatio-Temporal Bayesian Modeling for Multivariate Gaussian Processes

TL;DR

This paper introduces DALIA, a scalable GPU-accelerated framework for Bayesian inference on high-dimensional spatio-temporal multivariate Gaussian processes, significantly improving computational efficiency and enabling detailed pollution data analysis.

Contribution

DALIA employs a sparse inverse covariance formulation, hierarchical parallelism, and GPU acceleration to vastly improve scalability for multivariate GPs in spatial-temporal applications.

Findings

Surpasses state-of-the-art weak scaling by two orders of magnitude.

Achieves three orders of magnitude speedup on 496 superchips.

Provides refined spatial pollution data analysis over Italy.

Abstract

Multivariate Gaussian processes (GPs) offer a powerful probabilistic framework to represent complex interdependent phenomena. They pose, however, significant computational challenges in high-dimensional settings, which frequently arise in spatial-temporal applications. We present DALIA, a highly scalable framework for performing Bayesian inference tasks on spatio-temporal multivariate GPs, based on the methodology of integrated nested Laplace approximations. Our approach relies on a sparse inverse covariance matrix formulation of the GP, puts forward a GPU-accelerated block-dense approach, and introduces a hierarchical, triple-layer, distributed memory parallel scheme. We showcase weak scaling performance surpassing the state-of-the-art by two orders of magnitude on a model whose parameter space is 8$\times$ larger and measure strong scaling speedups of three orders of magnitude when running on 496 GH200 superchips on the Alps supercomputer. Applying DALIA to air pollution data from northern Italy over 48 days, we showcase refined spatial resolutions over the aggregated pollutant measurements.