The bottleneck
Neural optimal transport methods learn a transport map between probability distributions. In existing solvers, much of the difficulty comes from finding a near-exact approximation of the conjugate operator, also called the c-transform.
That approximation is commonly obtained either through non-convex max–min optimization or through computationally intensive fine-tuning of an initial prediction. The former can be unstable; the latter makes training slow.
The key idea
ENOT introduces a theoretically justified loss based on expectile regularization. The regularizer places binding conditions on the learned dual potentials and provides an upper-bound estimate over the distribution of possible conjugate potentials.
The practical goal is simple: keep the neural approximation close to the true conjugate potential without repeatedly solving an expensive inner optimization problem.
Expectiles summarize a random variable through asymmetric least squares. By penalizing overly optimistic approximations of the conjugate Kantorovich potential, the regularizer encourages the learned potential to stay close to the true conjugate.
What changes in practice
The resulting learning process is stable and eliminates the need for extensive additional fine-tuning. On established Wasserstein-2 benchmark tasks, the paper reports up to a three-fold improvement in quality and up to a ten-fold improvement in runtime over prior state-of-the-art approaches.
The work also evaluates ENOT with several cost functions and tasks, including image generation, to demonstrate the method’s generalizability and robustness. An implementation is available through the OTT-JAX library tutorial.
Read the paper and try ENOT
For the full derivation, experiments, and visual comparisons, use the original project page and paper resources.