Parallelised Differentiable Straightest Geodesics for 3D Meshes

1. Problem Statement

Machine learning on 3D meshes (non-Euclidean domains) has lagged behind due to the lack of efficient, differentiable Riemannian operators. While closed-form solutions exist for simple manifolds (e.g., spheres, $SO(3)$), complex mesh surfaces require numerical methods to compute geodesics.

• The Bottleneck: The standard method for computing the exponential map (Exp) and tracing geodesics on meshes is the Straightest Geodesics algorithm (Polthier & Schmies, 2006). However, existing implementations are:
  1. Non-differentiable: They rely on discrete combinatorial choices (e.g., selecting edges/faces), preventing backpropagation.
  2. Sequential: They run on CPUs, making them too slow for modern deep learning pipelines requiring batch processing.
  3. Inefficient: They struggle with optimization tasks where gradients must flow through the geodesic tracing process.

2. Methodology

The authors propose a framework to compute the Exponential Map, Geodesic Tracing, and Parallel Transport on triangular meshes that is both GPU-parallelized and fully differentiable.

A. Core Algorithm: Straightest Geodesics

The method traces a curve $\gamma$ on a mesh such that the angle of incidence equals the angle of emergence at every point (preserving "straightness").

• Discrete Formulation: Instead of solving a continuous ODE, the algorithm iteratively moves a point across faces, edges, and vertices.
• Handling Boundaries: A novel modification allows the algorithm to traverse mesh boundaries (holes) by temporarily relaxing the straightness criteria, tracing along the boundary, and resuming the geodesic, preventing premature termination.

B. Differentiation Schemes

To enable backpropagation through the non-differentiable tracing process, the authors derive two distinct methods:

1. Extrinsic Proxy (EP):

   • Concept: Uses a proxy function $\phi(p, v)$ that emulates the geodesic result in Euclidean space. It defines a "proxy endpoint" based on a rigid rotation and translation of the input vector.
   • Mechanism: It approximates the Jacobian with respect to the initial vector $v$ using a rotation matrix derived from the start and end frames of the geodesic.
   • Limitation: It assumes the starting point $p$ is fixed; gradients with respect to $p$ are zeroed out as the proxy does not accurately model position shifts.
   • Use Case: Fast, suitable for applications where only the vector direction is optimized.

2. Geodesic Finite Differences (GFD):

   • Concept: A finite difference scheme adapted to the local geometry of the mesh.
   • Mechanism: Perturbs the input point $p$ and vector $v$ along local basis vectors (barycentric coordinates for $p$, tangent frame for $v$). It computes the exponential map for these perturbed inputs and projects the resulting differences back to the local tangent space.
   • Advantage: Provides accurate gradients for both $p$ and $v$.
   • Trade-off: Computationally more expensive (requires multiple geodesic traces) but highly accurate.

C. GPU Parallelization

• Strategy: The authors implement a "one-thread-per-geodesic" CUDA kernel.
• Efficiency: Since geodesic paths are independent, thousands of traces can be computed simultaneously. The implementation handles variable path lengths and complexities without complex synchronization, achieving massive speedups over CPU-based sequential solvers.

3. Key Contributions

1. First Differentiable & Parallel Exp Map: A fully differentiable, GPU-accelerated implementation of the exponential map and parallel transport for general 3D meshes.
2. Software Library (digeo): A PyTorch-compatible library with C++/CUDA kernels, enabling seamless integration of discrete Riemannian geometry into learning pipelines.
3. Adaptive Geodesic Convolutions (AGC): A new convolutional layer that dynamically learns the patch size (receptive field) for each channel and layer, outperforming fixed-patch methods.
4. MeshFlow: A flow-matching generative model for meshes. It uses the differentiable Exp map to transport distributions without requiring ODE solvers or projection steps, drastically reducing inference time and memory.
5. Mesh-LBFGS: A second-order Riemannian optimizer (L-BFGS) for meshes, applied to solve the Centroidal Voronoi Tessellation (GCVT) problem with faster convergence than Lloyd's algorithm.

4. Results & Performance

Accuracy & Speed

• Speed: The GPU implementation is 3 orders of magnitude faster than CPU-based state-of-the-art (SotA) libraries (e.g., pp3d, geometry-central) for large batches. It scales linearly with mesh face count.
• Accuracy: The forward pass matches the accuracy of pp3d (mean error $\approx 10^{-5}$ to $10^{-9}$ depending on precision).
• Differentiation:
  • GFD achieves near-perfect cosine similarity ($>0.999$) with closed-form gradients on spheres.
  • EP provides reliable gradients for vector $v$ but ignores gradients for position $p$.
  • Compared to Projection Integration (PI) methods, the proposed methods are $10^5$ to $10^6$ times faster on high-resolution meshes.

Application Performance

• AGC (Segmentation): On human body part segmentation, AGC achieved 92.3% accuracy (using HKS inputs), outperforming GCNN (86.4%) and MDGCNN (89.5%).
• MeshFlow (Generative Modeling):
  • Speed: Inference is 16,000x faster than Riemannian Flow Matching (RFM).
  • Memory: Uses 97% less GPU memory than RFM.
  • Quality: Achieves superior KL Divergence and Biharmonic Chamfer Distance metrics compared to RFM.
• Mesh-LBFGS (Optimization): Converges to optimal Centroidal Voronoi Tessellations significantly faster than Lloyd's algorithm, particularly when initial seeds are clustered.

5. Significance

This work bridges the gap between theoretical Riemannian geometry and practical deep learning on 3D shapes. By providing the first efficient, differentiable, and parallelized operators for the exponential map on meshes, it enables:

• End-to-End Learning: Models can now optimize geometric properties directly on the surface without manual pre-computation.
• Scalability: The GPU parallelization makes Riemannian operations feasible for large-scale datasets and real-time applications.
• Versatility: The framework is agnostic to the specific forward algorithm, meaning it can be applied to other geodesic tracing methods in the future.

The release of the digeo library and the demonstrated success in convolution, generative modeling, and optimization mark a significant step forward for geometric deep learning, moving beyond simple point clouds to complex surface-based learning.