MeshFlow#

The full implementation of MeshFlow is available on the DSG-Applications repository.

../_images/duck_distrib.svg

Visual representation of MeshFlow. Noise samples, \(\mathbf{p}_i\), are displaced via our \(K\)-step exponential map according to the directions dictated by our learnable static vector field, \(\mathbf{v}_\theta\). The end-points, \(\hat{\mathbf{q}}_i\), are compared with the OT-coupled training samples, \(\mathbf{q}_{\sigma(i)}\), using a squared biharmonic distance function \(d_{BH}^2\).#

By utilizing the differentiable exponential map, MeshFlow learns a stationary vector field to transport points from a base noise distribution to a target data distribution entirely on the manifold.

The method relies on the Geodesic Finite Differences (GFD) differentiation scheme to enable back-propagation through both point positions and vectors. Since the straightest geodesics algorithm naturally follows straight paths, MeshFlow breaks the trajectory into multiple discrete steps to allow for curved paths on the manifold.

Mathematical Formulation#

The objective of MeshFlow is to optimize a time-invariant vector field by comparing transported samples against data samples using point-wise correspondences.

\[\mathcal{L}(\theta) = \frac{1}{B}\sum_{i=1}^{B} d_{BH}^{2}\big( \operatorname{Exp}_{p_{i}}^{\bigcirc K}(v_{\theta}), q_{\sigma(i)} \big)\]

Where:

  • \(B\) is the batch size.

  • \(p_{i}\) represents initial samples drawn from a uniform base distribution over the mesh \(\mathcal{U}(\mathcal{M})\).

  • \(q_{i}\) represents samples drawn from the target distribution \(\mathcal{Q}(\mathcal{M})\).

  • \(d_{BH}\) represents the biharmonic distance, used as a computationally efficient surrogate for the true geodesic distance.

  • \(\sigma\) is a permutation of \([1, B]\) that represents the Optimal Transport (OT) coupling for the batch.

  • \(\operatorname{Exp}^{\bigcirc K}\) represents the \(K\)-step exponential map, which iteratively breaks the trajectory into multiple steps to enable non-straightest curves on the mesh \(\mathcal{M}\).

  • \(v_{\theta}\) is the learned static vector field parameterized by a neural network.

Implementation#

To implement this, the neural network uses a Multi-Layer Perceptron (MLP) that takes the 3D position of the samples as input and outputs the corresponding tangent vector. Because the differentiable exponential map operates intrinsically, it allows points to move directly along the surface of the mesh without introducing any projection errors during the transport process.