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

This paper introduces GeometricKernels, a Python package that implements heat and Matérn kernels for geometric learning on manifolds, meshes, and graphs, enabling uncertainty quantification and Fourier-feature expansions with automatic differentiation across major machine learning frameworks.

The Gist

Imagine you are a chef trying to bake the perfect cake. In the old days, you only baked cakes on flat, square baking sheets (this is Euclidean space, or normal 2D/3D space). You knew exactly how the heat would spread, how the batter would rise, and you had a perfect recipe (a kernel) that worked every time.

But now, you want to bake cakes on weird shapes: a sphere (like a planet), a twisted pretzel (a mesh), or a complex network of roads (a graph). Suddenly, your old flat-sheet recipes fail. The heat doesn't spread the same way; the math gets messy, and you can't easily predict how the cake will turn out.

This is the problem the GeometricKernels package solves.

The Problem: "Flat" Recipes Don't Work on Curved Worlds

In machine learning, we use "kernels" to measure how similar two things are. If you are trying to predict the weather on a mountain, you need to know how the temperature at the peak relates to the temperature in the valley.

• The Issue: In normal space, we have a "magic formula" (like the Squared Exponential or Matérn kernel) that tells us this similarity perfectly.
• The Catch: When you try to use that same formula on a sphere or a graph, it breaks. It's like trying to use a square map to navigate a globe; the distances get distorted, and the math stops making sense.

The Solution: A Universal "Shape-Shifting" Toolkit

The authors of this paper built a software package called GeometricKernels. Think of it as a universal adapter or a shape-shifting chef's tool.

1. It Understands Geometry: Whether you are working on a sphere, a hyperbolic space (a weird, saddle-shaped universe), a 3D mesh (like a video game character), or a social network graph, this package knows how to calculate similarity correctly. It doesn't just guess; it uses deep mathematical principles (like "Heat Kernels") to figure out how information flows over these strange shapes.
2. It Handles Uncertainty: In machine learning, knowing what you don't know is just as important as knowing what you do. This package is designed to give you a "confidence score" (uncertainty quantification). It's like a weather forecaster who doesn't just say "It will rain," but says, "It will rain, and I'm 90% sure, but if the wind shifts, it might be only 50%."
3. It Plays Nice with Everyone: One of the coolest features is that it works with all the major machine learning tools at once.
   • If you use PyTorch (like a Tesla), it works.
   • If you use JAX (like a Ferrari), it works.
   • If you use TensorFlow (like a Toyota), it works.
   • It even works with basic NumPy (like a reliable bicycle) if you just want to debug things without heavy machinery.
   • Analogy: It's like a universal power plug that fits into outlets in the US, UK, and Japan without needing an adapter.

How It Works (The Magic Behind the Curtain)

The package uses something called Fourier Features.

• The Analogy: Imagine a drum. When you hit it, it vibrates at specific frequencies (notes). Any sound on that drum can be broken down into a mix of these notes.
• The Math: The package breaks down complex shapes (like a sphere or a graph) into their own "vibrating notes" (eigenfunctions). It then uses these notes to reconstruct the "heat" or "similarity" across the shape. This allows the computer to calculate complex probabilities without getting stuck in a math loop.

Why Should You Care?

This isn't just for math nerds. This toolkit opens the door for AI to understand the real world better:

• Robotics: Robots moving in 3D space need to understand the curvature of the world, not just a flat grid.
• Neuroscience: The brain is a complex network of connections (a graph). This helps model how signals travel through it.
• Drug Discovery: Molecules are 3D structures. Understanding their shape helps predict if a drug will work.
• Climate Science: The Earth is a sphere. Predicting weather patterns requires understanding how heat moves over a globe, not a flat map.

The Bottom Line

The GeometricKernels package is like giving machine learning a pair of 3D glasses. Before, AI could only see the world as flat, simple squares. Now, with this tool, AI can see the curves, the twists, and the complex connections of the real world, all while knowing exactly how confident it is in its predictions. It makes the impossible math of curved spaces as easy to use as a standard calculator.

Technical Summary

1. Problem Statement

Kernel methods, particularly Gaussian Processes (GPs), are fundamental to modern machine learning, especially for tasks requiring principled uncertainty quantification (e.g., Bayesian optimization, active learning). However, these methods traditionally rely on Euclidean spaces where kernels (like the Squared Exponential or Matérn) are defined via closed-form analytic expressions based on Euclidean distance.

When data resides on structured, non-Euclidean spaces—such as graphs, meshes, Riemannian manifolds, or Lie groups—defining valid kernels becomes significantly more difficult:

• Positive Semi-Definiteness (PSD): Kernels must be PSD to be valid. Naive generalizations of Euclidean distance-based kernels (e.g., $e^{-d(x,y)^2}$) often fail to be PSD on non-Euclidean domains.
• Computational Complexity: Many geometric kernels are implicitly defined (e.g., via solutions to differential equations) rather than having closed-form expressions, making numerical computation complex and often incompatible with modern automatic differentiation (autodiff) frameworks.
• Lack of Unified Tools: There is a scarcity of software that implements these geometric kernels in a way that supports multiple backends (PyTorch, JAX, TensorFlow) and GPU acceleration simultaneously.

2. Methodology

The authors present GeometricKernels, a Python package designed to implement geometric analogs of classical Euclidean kernels. The methodology relies on a rigorous mathematical framework and a flexible software architecture.

A. Mathematical Foundation

The package implements two primary classes of kernels derived from the Heat Kernel ($k_\infty$) and the Matérn Kernel ($k_\nu$):

1. Heat Kernel: Defined as the solution to the heat equation $\frac{\partial P}{\partial t} = \Delta P$ on a manifold/graph. It represents the transition probability of a diffusion process.
2. Matérn Kernel: Defined as a time-integral of the heat kernel:
   $$k_{\nu, \kappa, \sigma^2}(x, x') = \frac{\sigma^2}{C_{\nu, \kappa}} \int_0^\infty u^{\nu - 1 + n/2} e^{-\frac{2\nu}{\kappa^2}u} P(u, x, x') \, du$$
   This formulation generalizes the Euclidean Matérn kernel to arbitrary spaces where a heat kernel exists.

Computation via Spectral Expansions:
Instead of solving differential equations numerically for every query, the package utilizes Fourier-feature-type expansions (Mercer expansions).

• For discrete spectrum spaces (compact manifolds, graphs, meshes), the kernel is expanded using Laplacian eigenfunctions ($\phi_n$) and eigenvalues ($\lambda_n$):
  $$k(x, x') \approx \sum_{n} \Phi(\lambda_n) \phi_n(x) \phi_n(x')$$
  The package handles cases with repeated eigenvalues efficiently by grouping them into "levels" (e.g., using Gegenbauer polynomials for spheres or characters for Lie groups), avoiding the need to compute individual eigenfunctions.
• For non-compact spaces (e.g., hyperbolic space), the package uses Monte Carlo approximations of intractable integrals, ensuring the resulting approximations remain PSD.

B. Software Architecture

• Backend Agnosticism: The package uses the LAB (Linear Algebra Backend) library to support multiple numerical backends simultaneously: PyTorch, JAX, TensorFlow, and NumPy.
• Automatic Differentiation: All computations are fully differentiable, enabling gradient-based optimization of kernel hyperparameters within GP pipelines.
• Modular Design:
  • Spaces: Classes representing specific geometries (e.g., Hypersphere, SO(n), Graph, Mesh, Hyperbolic).
  • Kernels: A unified MaternGeometricKernel class that dispatches to the correct internal implementation based on the input space.
  • Feature Maps: Approximate finite-dimensional feature maps ($\phi: X \to \mathbb{R}^\ell$) are provided to enable efficient sampling and linearized GP inference (avoiding $O(N^3)$ costs).

3. Key Contributions

1. Unified Implementation: The first software package to provide a unified, plug-and-play implementation of Heat and Matérn kernels across a wide variety of geometric spaces (compact/non-compact manifolds, Lie groups, meshes, and graphs).
2. Multi-Backend Support: Seamless integration with PyTorch, JAX, and TensorFlow, preserving array types and supporting GPU acceleration and autodiff across all frameworks without code duplication.
3. Theoretical Generalization: Formalizes the definition of Matérn kernels on non-Euclidean spaces via the heat kernel integral, providing a consistent theoretical framework that subsumes previous ad-hoc definitions.
4. Efficient Approximation: Implements optimized spectral expansions and Monte Carlo methods that guarantee positive semi-definiteness while allowing for scalable computation via feature maps.
5. Ecosystem Integration: Provides wrapper classes for integration with major GP libraries (GPyTorch, GPJax, GPflow).

4. Results and Capabilities

The paper demonstrates the package's capabilities through:

• Interface Usability: A simple API allows users to define a space (e.g., a 2D hypersphere), initialize a kernel, set hyperparameters (lengthscale, $\nu$), and compute kernel matrices with minimal code.
• Visualization: The paper includes visualizations of GP samples generated using the heat kernel on the Stanford Bunny mesh, the unit sphere ($S^2$), and hyperbolic space ($H^2$), as well as Matérn kernels on graphs.
• Performance: The implementation supports batching and GPU acceleration, making it suitable for large-scale geometric learning tasks.
• Flexibility: The ProductGeometricKernel allows constructing kernels on product spaces (e.g., combining a graph and a manifold) with independent hyperparameters for each factor, facilitating automatic relevance determination.

5. Significance

The GeometricKernels package addresses a critical bottleneck in geometric machine learning: the difficulty of applying uncertainty-aware methods (like GPs) to structured data.

• Scientific Impact: It lowers the barrier to entry for researchers in robotics, neuroscience, and physics who need to model data on manifolds or graphs with rigorous uncertainty quantification.
• Practical Utility: By supporting multiple backends and GPUs, it enables the deployment of geometric kernels in modern deep learning pipelines, bridging the gap between classical kernel methods and deep learning.
• Future Directions: The framework is extensible, with the authors noting potential for future support of implicit manifolds, vector fields, and higher-order combinatorial structures (cellular complexes).

In summary, this work provides a robust, theoretically sound, and computationally efficient toolset for performing kernel-based learning on non-Euclidean domains, significantly advancing the state of the art in geometric Gaussian processes.