GPU-Accelerated Vecchia Approximations of Gaussian Processes for Geospatial Data using Batched Matrix Computations

Abstract: Gaussian processes (GPs) are commonly used for geospatial analysis, but they suffer from high computational complexity when dealing with massive data. For instance, the log-likelihood function required in estimating the statistical model parameters for geospatial data is a computationally intensive procedure that involves computing the inverse of a covariance matrix with size n X n, where n represents the number of geographical locations. As a result, in the literature, studies have shifted towards approximation methods to handle larger values of n effectively while maintaining high accuracy. These methods encompass a range of techniques, including low-rank and sparse approximations. Vecchia approximation is one of the most promising methods to speed up evaluating the log-likelihood function. This study presents a parallel implementation of the Vecchia approximation, utilizing batched matrix computations on contemporary GPUs. The proposed implementation relies on batched linear algebra routines to efficiently execute individual conditional distributions in the Vecchia algorithm. We rely on the KBLAS linear algebra library to perform batched linear algebra operations, reducing the time to solution compared to the state-of-the-art parallel implementation of the likelihood estimation operation in the ExaGeoStat software by up to 700X, 833X, 1380X on 32GB GV100, 80GB A100, and 80GB H100 GPUs, respectively. We also successfully manage larger problem sizes on a single NVIDIA GPU, accommodating up to 1M locations with 80GB A100 and H100 GPUs while maintaining the necessary application accuracy. We further assess the accuracy performance of the implemented algorithm, identifying the optimal settings for the Vecchia approximation algorithm to preserve accuracy on two real geospatial datasets: soil moisture data in the Mississippi Basin area and wind speed data in the Middle East.

• The paper introduces a GPU-accelerated Vecchia approximation for Gaussian processes that achieves up to 1380X speedups while maintaining accuracy comparable to exact MLE.
• It leverages batched matrix computations and the KBLAS library to efficiently handle large-scale geospatial datasets on NVIDIA GPUs such as GV100, A100, and H100.
• The study evaluates different reordering strategies and demonstrates improved memory efficiency and scalability for modeling real-world climate data like soil moisture and wind speed.

GPU-Accelerated Gaussian Process Approximation via Batched Vecchia

This paper introduces a GPU-accelerated implementation of the Vecchia approximation for Gaussian processes (GPs), leveraging batched matrix computations to enhance computational efficiency in geospatial data analysis. The implementation targets the efficient estimation of statistical model parameters for large-scale climate and weather applications. By utilizing the KBLAS library and batched linear algebra operations, the authors demonstrate significant speedups on NVIDIA GPU architectures, including GV100, A100, and H100, while maintaining accuracy comparable to exact maximum likelihood estimation (MLE).

Key Contributions

The paper makes several notable contributions to the field of computational statistics and high-performance computing:

• A GPU-accelerated implementation of the Vecchia approximation algorithm is presented, designed for efficient parameter estimation in climate and weather applications.
• The use of the KBLAS library and batched linear algebra operations accelerates the implementation on modern NVIDIA GPUs, including GV100, A100, and H100.
• The accuracy of the proposed implementation is assessed through numerical studies and real-world datasets, including soil moisture data from the Mississippi Basin and wind speed data from the Middle East. The study identifies optimal settings for the Vecchia algorithm to achieve performance comparable to exact MLE as implemented in the ExaGeoStat software.
• Performance evaluations on NVIDIA GPUs show significant speedups, reaching up to 700X on GV100, 833X on A100, and 1380X on H100, compared to exact MLE.
• The implementation accommodates larger problem sizes within the same GPU memory, enabling improved modeling for high-resolution geospatial data.

Batched Vecchia Approximation Framework

The proposed framework leverages batched matrix operations to accelerate the Vecchia approximation algorithm for Gaussian fields. The method involves reordering the location set and selecting the nearest neighbors for each location. The batched Vecchia approximated likelihood algorithm is described in detail (Figure 1), which involves replacing the high-dimensional joint distribution of the GP with a product of univariate conditional distributions:

Figure 1: Batched Vecchia algorithm description. $\bm \Sigma_{m:n}$ are constructed by the nearest neighbors of $\mathbf y^{\tau}_{m:n}$. The batched POTRF routine is applied to these matrices. After this decomposition, the resulting outputs are utilized as inputs for the batched TRSV operation with $\mathbf v_{m:n}$ and $\mathbf y^{\tau}_{\mathbf J_{m:n}$, separately.

For each spatial location, the algorithm computes a covariance matrix of its nearest neighbors and a cross-covariance vector between the location and its nearest neighbors. The approach applies uniform operations to batches of small matrices to leverage the underlying GPU accelerators. The implementation details include the use of a batched CUDA kernel for covariance matrix generation, batched Cholesky decomposition (POTRF), and triangular linear solver (TRSV) routines from the KBLAS library.

Reordering Strategies

The paper explores the impact of different reordering strategies on the accuracy of the log-likelihood approximation. The choice of reordering method is crucial for selecting nearest-neighbor points for each location. The study compares random ordering and Morton ordering, highlighting that random ordering generally outperforms Morton ordering for large-scale problems (Figure 2).

Figure 2: The example of random and Morton ordering on locations $20 \times 20$. (First row) The 45th and 250th locations (red stars) in the random ordering are marked with their nearest neighbors (orange circle); (Second row) The 45th and 250th locations (red stars) in the Morton order algorithm are marked with their nearest neighbors (orange circle). Blue circles indicate past locations for a given ordering algorithm.

Random ordering retains nearest neighbors immediately surrounding a target location, whereas Morton ordering may initially sacrifice proximity accuracy for earlier-ordered locations.

Performance and Accuracy Assessment

The paper includes a comprehensive performance and accuracy assessment of the batched Vecchia algorithm. Numerical studies using Kullback-Leibler (KL) divergence compare the approximation to exact MLE. The results indicate that the complexity of the approximation increases with the range or smoothness parameters of the Matérn kernel. Also, as the problem size increases, more conditioning points are required to maintain approximation accuracy.

Real-world datasets, including soil moisture and wind speed data, are used to evaluate the accuracy of the batched Vecchia algorithm in modeling and prediction tasks. The estimated parameter vectors closely align with those obtained via ExaGeoStat (exact MLE), particularly as the number of conditioning neighbors increases (Figure 3).

Figure 3: The estimated parameter vectors using Vecchia approximation with different conditioning sizes compared to {\it ExaGeoStat} (exact MLE). The first row is the parameter vector for soil moisture, and the second for wind speed.

The performance evaluation demonstrates significant speedups on NVIDIA GPUs, with the Vecchia method outperforming ExaGeoStat-GPU in single likelihood estimation time. The batched Vecchia approximation can handle larger problem sizes, up to 1 million locations, on a single GPU.

Memory Footprint and Arithmetic Complexity

The paper analyzes the memory footprint and arithmetic complexity of the batched Vecchia implementation compared to exact MLE (Figure 4).

Figure 4: Comparison of Arithmetic complexity: Vecchia algorithm versus Exact MLE.

The Vecchia algorithm significantly reduces memory requirements and computational complexity, making it suitable for large-scale geospatial data analysis.

Conclusions and Future Directions

The paper demonstrates the effectiveness of a GPU-accelerated Vecchia approximation for Gaussian processes, achieving significant speedups while maintaining accuracy. The use of batched matrix computations and the KBLAS library enables efficient parameter estimation for large-scale geospatial datasets. The study identifies optimal settings for the Vecchia algorithm and showcases its applicability to real-world problems in climate and weather modeling. This work paves the way for future research in scalable Gaussian process approximations and their application to various scientific domains. The source code for this work is publicly available.