Parallelised Differentiable Straightest Geodesics for 3D Meshes

TL;DR

This paper introduces a parallel GPU implementation of differentiable straightest geodesics on 3D meshes, enabling improved learning and optimization on non-Euclidean surfaces through new methods and applications.

Contribution

It presents a novel parallelized framework for differentiable geodesics on meshes, including two differentiation methods and multiple applications in geometric learning.

Findings

Parallel GPU implementation achieves high performance and accuracy.

Differentiable exponential map enhances learning pipelines on meshes.

New geodesic convolutional layer and flow matching method demonstrate versatility.

Abstract

Machine learning has been progressively generalised to operate within non-Euclidean domains, but geometrically accurate methods for learning on surfaces are still falling behind. The lack of closed-form Riemannian operators, the non-differentiability of their discrete counterparts, and poor parallelisation capabilities have been the main obstacles to the development of the field on meshes. A principled framework to compute the exponential map on Riemannian surfaces discretised as meshes is straightest geodesics, which also allows to trace geodesics and parallel-transport vectors as a by-product. We provide a parallel GPU implementation and derive two different methods for differentiating through the straightest geodesics, one leveraging an extrinsic proxy function and one based upon a geodesic finite differences scheme. After proving our parallelisation performance and accuracy, we demonstrate how our differentiable exponential map can improve learning and optimisation pipelines on general geometries. In particular, to showcase the versatility of our method, we propose a new geodesic convolutional layer, a new flow matching method for learning on meshes, and a second-order optimiser that we apply to centroidal Voronoi tessellation. Our code, models, and pip-installable library (digeo) are available at: circle-group.github.io/research/DSG.