Geometric Probabilistic Models

Part I

Alexander Terenin

Part I: background, motivation, technical fundamentals

1. Model-based decision-making
2. Probabilistic modeling and uncertainty
3. Geometric learning: working with structured data

Part II: geometric probabilistic models

1. Three types of probabilistic models: an example problem
2. Deep dive: geometric Gaussian processes
3. Emerging applications
4. Software
5. Summary

References

Neural network baseline

Gaussian process: squared exponential kernel

Gaussian process: polynomial kernel

Neural network ensemble

Bayesian neural network using MCMC

Gaussian process: stationary kernel

Gaussian process: polynomial kernel

Ensemble

Bayesian neural network

Vague idea: regulate uncertainty using symmetries

Non-linear regression and classification: $$ \htmlClass{fragment}{ \min_{f\in\c{F}} \sum_{i=1}^N L(y_i, f(x_i)) } $$ where

• $(x_i,y_i)$: training data
• $L$: loss function, such as square loss $\norm{y_i - f(x_i)}^2$
• $\c{F}$: some space of functions $f : X \to Y$

Spaces $X$ and/or $Y$ carry geometric structure

Probabilistic non-linear regression and classification: $$ \htmlClass{fragment}{ y_i \~ p_{y\given f}(\.\given f(x_i)) } $$ where

• $(x_i,y_i)$: training data
• $p_{y\given f}$: likelihood, such as Gaussian $y_i \~\f{N}(f(x_i), \sigma^2)$
• Bayesian approach: place a prior $p_f$ over functions $f : X \to Y$

Most models can be formulated in both ways

Spaces $X$ and $Y$ carry geometric structure

Meshes

Graphs

Manifolds

Vector Fields

And many more!

Different orderings represent same graph $\leadsto$ permutation invariance

Two ways of representing points on a sphere:

• Points in $\R^3$: $(x,y,z) \in \R^3$ with $x^2 + y^2 + z^2 = 1$
• Angles: $(\theta,\phi)$ representing angle relative to a chosen point

Different numbers: same point on sphere

Predictions should depend on our data, not on how we store it numerically

$ \large\x\, $

$ \large\x\, $

$ \large= $

Same vector field: represented differently in different frames

Predictions should depend on our data, not on the choice of frame

Need models that respect representation's symmetries

Invariance

Equivariance

$ \x\, $

$ \x\, $

$ = $

Challenge: model design

Design principle: models which respect the space's symmetries

• Symmetry-based extrapolation
• Symmetry-based uncertainty

..unless data indicates otherwise?

Just as relevant in classical Euclidean setting

Geometric learning: right technical language for studying this

Graph neural networks

• Geometric analogs of ConvNets
• Most-common model by far

Idea:

1. Pixel-space convolution: discretization of Euclidean convolution
2. Replace Euclidean convolution with non-Euclidean analog
3. Discretize: obtain graph neural network architecture

Challenge: what is the right notion of convolution?

Graph neural network: predicted antibiotic activity in halicin, prev. studied for diabetes

Shown in mice to have broad-spectrum antibiotic activity

Figure and results: Stokes et al. (Cell, 2020)

Graph neural networks: used to improve complementary product recommendation

Improved performance in production compared to Amazon's prior approach

Figure and results: Hao et al. (CIKM 2020)

Graph neural networks: Google Maps traffic ETA prediction

Improved accuracy in production compared to prior approach

Figure and results: Derrow-Pinion et al. (CIKM 2021)

Geometric Probabilistic Models

Part II

Viacheslav (Slava) Borovitskiy

Part I: background, motivation, technical fundamentals

1. Model-based decision-making
2. Probabilistic modeling and uncertainty
3. Geometric learning: working with structured data

Part II: geometric probabilistic models

1. Three types of probabilistic models: an example problem
2. Deep dive: geometric Gaussian processes
3. Emerging applications
4. Software
5. Summary

References

Dataset details: Borovitskiy et al. (AISTATS 2021)

Deep Ensemble

Bayesian Graph CNN

Gaussian Process

This example: easy-to-use and visualize benchmark for regression on geometric data

Geometric Gaussian processes:

• Reliable off-the-shelf probabilistic models
• Solid baselines to benchmark other models against

$$ \htmlData{class=fragment fade-out,fragment-index=6}{ \footnotesize \mathclap{ k_{\nu, \kappa, \sigma^2}(x,x') = \sigma^2 \frac{2^{1-\nu}}{\Gamma(\nu)} \del{\sqrt{2\nu} \frac{\abs{x-x'}}{\kappa}}^\nu K_\nu \del{\sqrt{2\nu} \frac{\abs{x-x'}}{\kappa}} } } \htmlData{class=fragment d-print-none,fragment-index=6}{ \footnotesize \mathclap{ k_{\infty, \kappa, \sigma^2}(x,x') = \sigma^2 \exp\del{-\frac{\abs{x-x'}^2}{2\kappa^2}} } } $$

$\nu\to\infty$: RBF kernel (Gaussian, heat, diffusion)

$\nu = 1/2$

$\nu = 3/2$

$\nu = 5/2$

$\nu = \infty$

Geometry-aware, but...

Manifolds: not well-defined unless the manifold is isometric to a Euclidean space

Feragen et al. (CVPR 2015)

Graphs: not well-defined unless nodes can be isometrically embedded into a Hilbert space

Schoenberg (Trans. Am. Math. Soc. 1938)

Spaces of graphs: what is a space of graphs?

Need a different generalization

$$ \htmlData{class=fragment,fragment-index=0}{ \ubr{\del{\frac{2\nu}{\kappa^2} - \Delta}^{\frac{\nu}{2}+\frac{d}{4}} f = \c{W}}{\t{Matérn}} } $$

SPDE turns into a stochastic linear system: solution has kernel $$ \htmlData{fragment-index=2,class=fragment}{ k_{\nu, \kappa, \sigma^2}(i, j) = \frac{\sigma^2}{C_{\nu, \kappa}} \sum_{n=0}^{\abs{V}-1} \Phi_{\nu, \kappa}(\lambda_n) \v{f_n}(i)\v{f_n}(j) } $$

$\lambda_n, \v{f_n}$: eigenvalues and eigenvectors of the Laplacian matrix

$$ \htmlData{fragment-index=4,class=fragment}{ \Phi_{\nu, \kappa}(\lambda) = \begin{cases} \htmlData{fragment-index=5,class=fragment}{ \del{\frac{2\nu}{\kappa^2} - \lambda}^{-\nu} } & \htmlData{fragment-index=5,class=fragment}{ \nu\lt\infty : \t{Matérn} } \\ \htmlData{fragment-index=6,class=fragment}{ e^{-\frac{\kappa^2}{2} \lambda} } & \htmlData{fragment-index=6,class=fragment}{ \nu = \infty : \t{heat (sq. exp., RBF)} } \end{cases} } $$

Compact manifolds: $$ \htmlData{class=fragment}{ \begin{aligned} k_{\nu, \kappa, \sigma^2}(x,x') &= \frac{\sigma^2}{C_{\nu, \kappa}} \sum_{n=0}^\infty \Phi_{\nu, \kappa}(\lambda_n) f_n(x) f_n(x') \\ &= \frac{\sigma^2}{C_{\nu, \kappa}} \sum_{l=0}^\infty \Phi_{\nu, \kappa}(\lambda_l) \underbrace{\sum_{s=1}^{d_s} f_{l, s}(x) f_{l, s}(x')}_{G_l(x, x')} \end{aligned} } $$

Infinite summation where $\lambda_n, f_n$ are Laplace–Beltrami eigenpairs

Addition theorems: more efficient $G_l(x, x')$-based expressions

Non-compact manifolds:

• No summation: instead, integral approximated via Monte Carlo

Spaces of graphs:

• Geometry: make a set of graphs to be a node set of a larger graph
• Then, heat and Matérn kernels are automatically defined
• Larger graph is huge: need clever computational techniques

Joint postures: $\mathbb{T}^d$

Rotations: $\mathbb{S}_3$, $\operatorname{SO}(3)$

Manipulability: $\operatorname{SPD(3)}$

Jaquier et al. (CoRL 2019, CoRL 2021), figures by the authors

Prediction

Uncertainty

Matérn Gaussian processes on meshes (independently proposed)

Coveney et al. (TBE 2019, PTRSA 2020), figures by the authors

Based on ensembles of geodesic CNNs

Neural Concept, video from neuralconcept.com

https://geometric-kernels.github.io/

Example code:

>>> # Define a space (2-dim sphere).
>>> hypersphere = Hypersphere(dim=2)

>>> # Initialize kernel.
>>> kernel = MaternGeometricKernel(hypersphere)
>>> params = kernel.init_params()

>>> # Compute and print out the 3x3 kernel matrix.
>>> xs = np.array([[0., 0., 1.], [0., 1., 0.], [1., 0., 0.]])
>>> print(kernel.K(params, xs))

[[1.    0.356     0.356]
	[0.356 1.        0.356]
	[0.356 0.356 1.       ]]

Multi-backend: Numpy, JAX, PyTorch, TensorFlow

Neural networks beyond the graph setting: various small libraries

Geometric machine learning: increasingly popular

• Waves: kernels on structured data, geometric deep learning

Geometric probabilistic models: emerging research area

• Geometric Gaussian processes: solid off-the-shelf baselines
• Increasingly available but actively developing: join us!