Fast multitask Gaussian processes, with application to surrogate modeling of the quark-gluon plasma

Abstract

Gaussian processes (GPs) are broadly used for the surrogate modeling of computer experiments with reliable uncertainty quantification. Our motivating application comes from the study of the quark-gluon plasma (QGP), an extreme state of nuclear matter that filled the universe shortly after the Big Bang. To reliably infer properties of the QGP, multiple surrogate models need to be constructed for related particle collision simulation systems (i.e., multiple "tasks"). While there is a body of literature on multitask GPs (MTGPs), such models can be computationally expensive with large datasets: they require $\mathcal{O}(N^3)$ work and $\mathcal{O}(N^2)$ storage for model training, where $N$ is the total number of samples over all tasks. To address this, we propose a new fast MTGP approach, which pairs structured low-discrepancy design points, such as Sobol' points, with special kernel forms for efficient and exact model fitting. This pairing of a kernel and potentially different design points of different sizes across tasks provides a structured Gram matrix, e.g., a circulant block matrix, which we exploit via a novel algorithm for efficient Gram matrix inversion and determinant computation. Our algorithm reduces training costs to $\mathcal{O}(N \log N)$ work and $\mathcal{O}(N)$ storage in the case of equal sample sizes for each tasks. In the worst case of severely unbalanced sample sizes across tasks, our algorithm may require up to $\mathcal{O}(N^2)$ work and storage, with continuous interpolation between these best and worst cases depending on the balance of sample sizes across tasks. An open-source Python implementation is made available in the FastGPs package (https://alegresor.github.io/fastgps/). We demonstrate the effectiveness of our fast MTGP approach on a range of simulation experiments and on our motivating QGP application.