Scaling Gaussian Process Regression with Full Derivative Observations

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.

💡 Research Summary

The paper introduces DSoftKI, a scalable Gaussian Process (GP) method that can train and predict with full derivative (gradient) observations. Traditional GP regression scales as O(n³) in the number of data points n, but when derivative information is included the cost becomes O(n³ d³), where d is the input dimensionality, making it infeasible for high‑dimensional or large‑scale problems. Existing scalable approaches such as Stochastic Variational GP (SVGP) extended to derivatives (DSVGP) reduce the cost to O(m³ d³) using m ≪ n inducing points, while directional‑derivative variants (DDSVGP) further lower it to O(m³ p³) with p ≪ d directional inducing directions. However, DSVGP still requires explicit kernel derivatives, and DDSVGP sacrifices the ability to predict full gradients.

Soft Kernel Interpolation (SoftKI) previously approximated a kernel by softmax interpolation over a small set of learned interpolation points, using a global temperature vector T to control length‑scale per dimension. While SoftKI provides O(m² n) inference for standard GP regression, it cannot directly handle derivative data because the interpolation scheme does not encode directional sensitivity.

DSoftKI solves this by replacing the global temperature vector with a local temperature vector t_j for each interpolation point z_j. The softmax weight becomes σ_jz(x)=exp(−‖x⊘t_j−z_j‖) / Σ_k exp(−‖x⊘t_k−z_k‖). This local scaling allows the interpolation matrix Σₓz to contain both the scalar weights σ_jz(x) and their gradients ∇σ_jz(x) in a single block (Equation 18). Consequently, the approximate kernel with derivatives, ˜K_DSoftKI≈˜Σₓz K_zz ˜Σ_zx, can be constructed without ever computing analytical derivatives of the base kernel. The method therefore sidesteps the O(d³) cost associated with kernel Hessians and enables the use of arbitrary learned kernels, including deep kernels, without deriving their gradients.

Training proceeds by maximizing a Hutchinson stochastic trace estimator of the marginal log‑likelihood (the “pseudoloss”), which has per‑mini‑batch cost O(b² + m³). Both the interpolation point locations and the local temperature vectors are learned jointly with kernel hyper‑parameters (length‑scale ℓ, output scale γ, noise variances β_v² and β_g²). Posterior prediction follows the same algebra as Sparse GP Regression (SGPR) but with interpolated cross‑covariances, yielding O(m² n) complexity for the predictive mean and variance.

The overall computational complexity of DSoftKI is O(m² n d): linear in the data dimension d, quadratic in the number of interpolation points m, and linear in the dataset size n. This is dramatically lower than the O(m³ d³) of DSVGP and O(m³ p³) of DDSVGP, while still providing full gradient predictions.

Empirical evaluation covers four settings:

1. Synthetic functions with analytically known gradients, demonstrating that DSoftKI matches or exceeds the accuracy of exact GPwD while scaling to larger n and d.
2. The 2‑D Branin surface with full gradient fields, where DSoftKI reproduces both the scalar surface and the two component gradient fields accurately, unlike DDSVGP which loses fidelity.
3. A toy n‑body physics simulation, showing that the method can learn complex interaction forces from position‑gradient pairs.
4. High‑dimensional molecular force‑field datasets (d≈100–1000), where DSoftKI predicts both potential energy and forces. Competing methods either run out of memory or suffer severe accuracy loss, whereas DSoftKI remains tractable and competitive.

The authors also discuss extensions: because DSoftKI never requires kernel derivatives, it can be combined with Deep Kernel Learning (DKL) to learn expressive kernels from data, and the local temperature vectors provide a form of automatic relevance detection that adapts to local geometry.

In summary, DSoftKI makes four key contributions: (i) it enables full‑gradient GP regression at scale, (ii) it eliminates the need for explicit kernel derivative calculations, (iii) it introduces local temperature vectors to capture directional sensitivity, and (iv) it achieves O(m² n d) inference complexity. This advances the practical applicability of Gaussian Processes to high‑dimensional scientific problems where gradient information is abundant, such as molecular dynamics, fluid simulations, and physics‑based surrogate modeling. Future work may explore richer interpolation schemes, hierarchical temperature structures, and integration with deep learning pipelines for even larger and more complex datasets.