The GeometricKernels Package: Heat and Matérn Kernels for Geometric Learning on Manifolds, Meshes, and Graphs (2407.08086v1)

Abstract: Kernels are a fundamental technical primitive in machine learning. In recent years, kernel-based methods such as Gaussian processes are becoming increasingly important in applications where quantifying uncertainty is of key interest. In settings that involve structured data defined on graphs, meshes, manifolds, or other related spaces, defining kernels with good uncertainty-quantification behavior, and computing their value numerically, is less straightforward than in the Euclidean setting. To address this difficulty, we present GeometricKernels, a software package which implements the geometric analogs of classical Euclidean squared exponential - also known as heat - and Mat\'ern kernels, which are widely-used in settings where uncertainty is of key interest. As a byproduct, we obtain the ability to compute Fourier-feature-type expansions, which are widely used in their own right, on a wide set of geometric spaces. Our implementation supports automatic differentiation in every major current framework simultaneously via a backend-agnostic design. In this companion paper to the package and its documentation, we outline the capabilities of the package and present an illustrated example of its interface. We also include a brief overview of the theory the package is built upon and provide some historic context in the appendix.

• The paper introduces the GeometricKernels package that extends classical kernel methods using heat and Matérn kernels to non-Euclidean spaces.
• It supports diverse geometric spaces—including manifolds, meshes, and graphs—and integrates with major Gaussian process libraries via multi-backend design.
• It employs approximate finite-dimensional feature maps and automatic differentiation to enable efficient Gaussian process sampling and scalable computation.

An Expert Overview of "The GeometricKernels Package: Heat and Matérn Kernels for Geometric Learning on Manifolds, Meshes, and Graphs"

The paper "The GeometricKernels Package: Heat and Matérn Kernels for Geometric Learning on Manifolds, Meshes, and Graphs" introduces the GeometricKernels software package, which extends classical kernel methods to a variety of non-Euclidean geometric spaces such as manifolds, meshes, and graphs. This package addresses the complexities associated with defining and computing kernels in these settings, focusing on two primary kernels: heat and Matérn kernels.

Motivation and Background

Kernel methods, particularly Gaussian processes (GPs), have gained prominence in machine learning due to their powerful capability for uncertainty quantification. While kernels are well-defined and easily computable in Euclidean spaces, extending these methods to geometric spaces such as graphs, meshes, and manifolds poses significant challenges. This paper's motivation lies in developing a software solution—GeometricKernels—that simplifies the use of geometric kernels while ensuring robust support for automatic differentiation and GPU acceleration across multiple backends.

Capabilities of GeometricKernels

Classes for Geometric Spaces

The GeometricKernels package supports a wide array of geometric spaces categorized into:

• Compact Riemannian manifolds: Examples include the unit circle ($\mathbb{S}_1$), unit hyperspheres ($\mathbb{S}_n$), special orthogonal groups ($\mathrm{SO(n)}$), and special unitary groups ($\mathrm{SU}(n)$).
• Non-compact Riemannian manifolds: Including hyperbolic spaces ($\mathbb{H}_n$) and spaces of symmetric positive definite matrices ($\mathrm{SPD}(n)$).
• Meshes: Viewed as approximations of smooth two-dimensional surfaces.
• Graph node sets: For general undirected graphs with non-negative weights.
• Products of discrete spectrum spaces: These combinations can include the spaces mentioned above except for non-compact Riemannian manifolds.

Heat and Matérn Kernel Classes

The kernel classes provided in the package include:

• MaternGeometricKernel: A universal class covering most use cases, automatically dispatching based on the provided geometric space.
• ProductGeometricKernel: Facilitates constructing product kernels on product spaces, allowing independent hyperparameters for each factor, offering flexibility akin to automatic relevance determination (ARD).

Approximate Finite-dimensional Feature Maps

The package supports approximate finite-dimensional feature maps, critical for efficiently obtaining approximate samples from Gaussian processes. These maps allow: $k(x,x') \approx \v\phi(x)^\top \v\phi(x'),$ where $\v\phi$ is a functional representation into a finite-dimensional space. This approximation enables efficient operations without the cubic cost typically associated with Gaussian processes.

Multi-backend Design

A standout feature of GeometricKernels is its backend-agnostic design. It supports PyTorch, JAX, TensorFlow, and Numpy, facilitating seamless integration into various machine learning workflows. This multi-backend design is implemented via multiple dispatch and function/operator overloading, leveraging LAB library functionalities. This ensures that users can write code once and deploy it across different computational frameworks.

Integration with Existing Packages

GeometricKernels integrates with prominent Gaussian process libraries like GPyTorch, GPJax, and GPflow. Wrapper classes ensure compatibility and ease of use within these ecosystems, extending the package's utility.

Illustrative Example

The paper provides a worked example demonstrating how to construct a Matérn kernel matrix on the sphere using the package. Key steps involved include:

1. Initializing the hypersphere and kernel objects.
2. Setting kernel hyperparameters.
3. Evaluating the kernel on sample points to construct the kernel matrix.

This example underscores the package's ease of use and its flexibility in setting up and computing with geometric kernels.

Conclusions and Implications

GeometricKernels stands as a noteworthy contribution to the machine learning community, particularly in the context of geometric data. By extending classical kernel methods to manifold, mesh, and graph domains while ensuring robust computational support, it addresses a significant gap in existing tools. Practically, this package enables more reliable uncertainty quantification in complex geometric settings, which is critical in domains such as robotics, neuroscience, and beyond.

Future developments could expand the package's applicability to new types of geometric data, such as implicit geometries and other manifold types. This continued evolution could further enhance its utility and foster new research avenues in geometric learning.

In essence, GeometricKernels not only facilitates the practical application of geometric kernels but also serves as a springboard for future theoretical advancements and interdisciplinary applications in machine learning.