Table of Contents
Fetching ...

A network based approach for unbalanced optimal transport on surfaces

Jiangong Pan, Wei Wan, Yuejin Zhang, Chenlong Bao, Zuoqiang Shi

TL;DR

This work develops a mesh-free neural PDE solver for dynamic unbalanced OT on surfaces by formulating MUOT on manifolds via a Hamiltonian-flow (KKT) framework. It approximates the density and potential with neural networks on point clouds and enforces the MUOT KKT conditions through a composite loss, enabling MAP-free transport without meshing. The approach demonstrates accurate transport on various manifolds, robustness to moderate noise, and applicability to shape-transfer with mass imbalance, indicating broad potential for geometry processing and surface transport problems. By leveraging the Wasserstein-Fisher-Rao (WFR) perspective and a simplified loss structure, the method offers a scalable, geometry-aware tool for unbalanced transport on complex surfaces.

Abstract

In this paper, we present a neural network approach to address the dynamic unbalanced optimal transport problem on surfaces with point cloud representation. For surfaces with point cloud representation, traditional method is difficult to apply due to the difficulty of mesh generating. Neural network is easy to implement even for complicate geometry. Moreover, instead of solving the original dynamic formulation, we consider the Hamiltonian flow approach, i.e. Karush-Kuhn-Tucker system. Based on this approach, we can exploit mathematical structure of the optimal transport to construct the neural network and the loss function can be simplified. Extensive numerical experiments are conducted for surfaces with different geometry. We also test the method for point cloud with noise, which shows stability of this method. This method is also easy to generalize to diverse range of problems.

A network based approach for unbalanced optimal transport on surfaces

TL;DR

This work develops a mesh-free neural PDE solver for dynamic unbalanced OT on surfaces by formulating MUOT on manifolds via a Hamiltonian-flow (KKT) framework. It approximates the density and potential with neural networks on point clouds and enforces the MUOT KKT conditions through a composite loss, enabling MAP-free transport without meshing. The approach demonstrates accurate transport on various manifolds, robustness to moderate noise, and applicability to shape-transfer with mass imbalance, indicating broad potential for geometry processing and surface transport problems. By leveraging the Wasserstein-Fisher-Rao (WFR) perspective and a simplified loss structure, the method offers a scalable, geometry-aware tool for unbalanced transport on complex surfaces.

Abstract

In this paper, we present a neural network approach to address the dynamic unbalanced optimal transport problem on surfaces with point cloud representation. For surfaces with point cloud representation, traditional method is difficult to apply due to the difficulty of mesh generating. Neural network is easy to implement even for complicate geometry. Moreover, instead of solving the original dynamic formulation, we consider the Hamiltonian flow approach, i.e. Karush-Kuhn-Tucker system. Based on this approach, we can exploit mathematical structure of the optimal transport to construct the neural network and the loss function can be simplified. Extensive numerical experiments are conducted for surfaces with different geometry. We also test the method for point cloud with noise, which shows stability of this method. This method is also easy to generalize to diverse range of problems.
Paper Structure (11 sections, 1 theorem, 24 equations, 12 figures, 8 tables)

This paper contains 11 sections, 1 theorem, 24 equations, 12 figures, 8 tables.

Key Result

Theorem 3.1

Suppose $(\hat{\rho}, \hat{\boldsymbol{v}}, \hat{g}, \hat{\phi}, \hat{\phi_1}, \hat{\phi_2})$ is a saddle point of (Lagrangian-fun), Then, solution $(\hat{\rho}, \hat{\phi})$ (where $\hat{\boldsymbol{v}}=\nabla_{\Gamma}\hat{\phi}$ and $\hat{g} = \frac{\eta}{2}\hat{\phi}$) of (dyKKT-MUOT) is the optimal solution for (dyMUOT) and (dyMUOTC). And we refer to (dyKKT-MUOT) as dynamic KKT-MUOT:

Figures (12)

  • Figure 1: 2D KKT-UOT test. First row: Simple gaussian transport. Second row: Mixed gaussian split transport.
  • Figure 2: Variable Coefficients UOT: Numerical results for different $\eta$, from top to bottom each line $\eta=10^2, 1, 10^{-6}$.
  • Figure 3: OT tests on point cloud.
  • Figure 4: OT tests on point cloud in four-dimensional space.
  • Figure 5: MUOT test on Sphere.
  • ...and 7 more figures

Theorems & Definitions (2)

  • Theorem 3.1
  • proof