Scaling Gaussian Process Regression with Full Derivative Observations

TL;DR

This paper introduces DSoftKI, a scalable Gaussian Process method that efficiently incorporates full derivative observations, enabling accurate predictions in high-dimensional settings and extending kernel interpolation techniques.

Contribution

The paper extends SoftKI to include local temperature vectors, allowing scalable kernel interpolation with derivatives and facilitating extensions like Deep Kernel Learning.

Findings

DSoftKI accurately predicts with full derivatives on synthetic and real datasets.

It scales to high-dimensional molecular force field prediction (up to 1000 dimensions).

Demonstrates improved scalability over previous GP methods with derivative observations.

Abstract

We present a scalable Gaussian Process (GP) method called DSoftKI that can fit and predict full derivative observations. It extends SoftKI, a method that approximates a kernel via softmax interpolation, to the setting with derivatives. DSoftKI enhances SoftKI's interpolation scheme by replacing its global temperature vector with local temperature vectors associated with each interpolation point. This modification allows the model to encode local directional sensitivity, enabling the construction of a scalable approximate kernel, including its first and second-order derivatives, through interpolation. Moreover, the interpolation scheme eliminates the need for kernel derivatives, facilitating extensions such as Deep Kernel Learning (DKL). We evaluate DSoftKI on synthetic benchmarks, a toy n-body physics simulation, standard regression datasets with synthetic gradients, and high-dimensional molecular force field prediction (100-1000 dimensions). Our results demonstrate that DSoftKI is accurate and scales to larger datasets with full derivative observations than previously possible.