A Framework for Nonstationary Gaussian Processes with Neural Network Parameters

TL;DR

This paper introduces a flexible framework for nonstationary Gaussian processes where kernel parameters vary across the feature space and are modeled by neural networks, improving accuracy and interpretability on various datasets.

Contribution

The paper proposes a novel joint training approach for nonstationary Gaussian processes with neural network parameterization, enhancing model flexibility and scalability.

Findings

Achieves better accuracy and log-score than stationary and hierarchical models.

Effectively recovers nonstationary parameters in spatial datasets.

Compatible with large-scale approximation methods.

Abstract

Gaussian processes have become a popular tool for nonparametric regression because of their flexibility and uncertainty quantification. However, they often use stationary kernels, which limit the expressiveness of the model and may be unsuitable for many datasets. We propose a framework that uses nonstationary kernels whose parameters vary across the feature space, modeling these parameters as the output of a neural network that takes the features as input. The neural network and Gaussian process are trained jointly using the chain rule to calculate derivatives. Our method clearly describes the behavior of the nonstationary parameters and is compatible with approximation methods for scaling to large datasets. It is flexible and easily adapts to different nonstationary kernels without needing to redesign the optimization procedure. Our methods are implemented with the GPyTorch library and can be readily modified. We test a nonstationary variance and noise variant of our method on several machine learning datasets and find that it achieves better accuracy and log-score than both a stationary model and a hierarchical model approximated with variational inference. Similar results are observed for a model with only nonstationary variance. We also demonstrate our approach's ability to recover the nonstationary parameters of a spatial dataset.