Neural Hamilton--Jacobi Characteristic Flows for Optimal Transport

TL;DR

This paper addresses the problem of learning optimal transport maps between probability distributions with general cost functions in a scalable and provably correct manner. It introduces Neural Characteristic Flow (NCF), a single implicit neural representation that exploits the Hamilton–Jacobi equation and its method of characteristics to obtain closed-form bidirectional OT maps without solving ODEs. The approach provides theoretical guarantees of consistency with OT and stability in Gaussian settings, and it extends naturally to class-conditional OT via per-class MMD alignment. Empirically, NCF achieves accurate, efficient transport across 2D toys, high-dimensional Gaussians, color transfer, and MNIST/Fashion-MNIST tasks, outperforming adversarial and dynamical baselines while handling a broad class of costs.

Abstract

We present a novel framework for solving optimal transport (OT) problems based on the Hamilton--Jacobi (HJ) equation, whose viscosity solution uniquely characterizes the OT map. By leveraging the method of characteristics, we derive closed-form, bidirectional transport maps, thereby eliminating the need for numerical integration. The proposed method adopts a pure minimization framework: a single neural network is trained with a loss function derived from the method of characteristics of the HJ equation. This design guarantees convergence to the optimal map while eliminating adversarial training stages, thereby substantially reducing computational complexity. Furthermore, the framework naturally extends to a wide class of cost functions and supports class-conditional transport. Extensive experiments on diverse datasets demonstrate the accuracy, scalability, and efficiency of the proposed method, establishing it as a principled and versatile tool for OT applications with provable optimality.

Figures (15)

• Figure 1: Swiss roll ($\mu$) $\rightleftarrows$ Double moons ($\nu$): The top row shows transport in the direction $\nu \rightarrow \mu$, and the bottom row shows $\mu \rightarrow \nu$. The leftmost column displays $\mu$ and $\nu$ for reference.
• Figure 2: Computational comparison. Training time (s/epoch), evaluation time (s/epoch), peak memory (MB) during training, and memory (MB) for storing bidirectional OT maps are reported.
• Figure 3: 2D class-conditional OT. The leftmost column displays $\mu$ (red) and $\nu$ (blue), with class labels indicated by distinct markers. In the remaining columns, blue dots denote transported samples, while solid black and dotted gray lines represent the learned transport maps for each class.
• Figure 4: Class-conditional OT between MNIST and Fashion MNIST. Left: Forward OT. Right: Backward OT. The first row shows the source data, while the second row presents the data generated by learned OT map.
• Figure 5: Checkerboard ($\mu$) $\rightleftarrows$ Eight Gaussians ($\nu$): The top row shows transport in the direction $\nu \rightarrow \mu$, and the bottom row shows $\mu \rightarrow \nu$, with $\mu$ and $\nu$ at the leftmost column.

Theorems & Definitions (25)

• Theorem 2.1: santambrogio2015optimal
• Proposition 4.1: Bidirectional OT Map
• Proposition 4.2
• proof
• Theorem 5.1: Consistency of loss
• Remark 5.2: On regularity assumption of $u$
• Theorem 5.4: Stability of loss
• Lemma A.1
• Lemma A.2
• proof