Scaling Gaussian Process Regression with Full Derivative Observations

TL;DR

This work tackles the scalability barrier of Gaussian Processes with full derivative observations (GPwD) by introducing DSoftKI, a derivative-aware extension of SoftKI that constructs a scalable approximate kernel through learned, directional interpolation. By augmenting the SoftKI interpolation with per-point directional temperatures and avoiding explicit kernel derivatives, DSoftKI achieves a posterior inference complexity of $O(ndm^2)$ and supports first and second-order derivatives, enabling accurate gradient modeling in high-dimensional settings. The approach is validated on synthetic benchmarks and a high-dimensional molecular force-field task (MD22), demonstrating improved derivative-aware accuracy and the ability to scale to large $n$ and $d$ where prior GPwD methods struggle. These results suggest DSoftKI as a practical option for GP regression with full derivative information in scientific domains, with potential extensions to deeper kernels and multi-GPU implementations.

Abstract

We present a scalable Gaussian Process (GP) method that can fit and predict full derivative observations called DSoftKI. It extends SoftKI, a method that approximates a kernel via softmax interpolation from learned interpolation point locations, to the setting with derivatives. DSoftKI enhances SoftKI's interpolation scheme to incorporate the directional orientation of interpolation points relative to the data. This enables the construction of a scalable approximate kernel, including its first and second-order derivatives, through interpolation. We evaluate DSoftKI on a synthetic function benchmark and high-dimensional molecular force field prediction (100-1000 dimensions), demonstrating that DSoftKI is accurate and can scale to larger datasets with full derivative observations than previously possible.

Figures (10)

• Figure 1: Comparison of GPwD regression with vanilla GPwD, DSVGP, DDSVGP, and DSoftKI on the Branin surface (2D). We plot the contours of the surface, the gradient with respect to the first argument ($\nabla_1$), and the gradient with respect to the second argument ($\nabla_2$). Vanilla GPwD is accurate but intractable for sizable $n$ and/or $d$. DSVGP forms a nice approximation of the original surface but is intractable for large $d$. DDSVGP with uses $p = 2$ directional derivatives to enhance scalability but loses fidelity in modeling the surface. DSoftKI provides an accurate and scalable approximation of the surface.
• Figure 2: Comparison of learned interpolation points locations versus learned inducing point locations (only those in unit square shown) on the Branin surface. Interpolation point locations encode geometric structure in the data that is most useful for interpolation whereas inducing point locations reflect the normal priors placed on their associated inducing variables.
• Figure 3: Reconstructed Styblinski Tang surface using a shared temperature (original SoftKI scheme) versus individual temperature (proposed scheme) during interpolation for DSoftKI. We also overlay the learned interpolation points (scaled and translated to fit the unit square) in green to illustrate the differences in learned interpolation points.
• Figure 4: Test RMSE per atom vs. test gradient RMSE per component obtained by various methods on MD22 dataset. Bottom left is best.
• Figure 5: Effect of noise.