Improving Neural Optimal Transport via Displacement Interpolation

2024-10-03

arXiv.org

AI Summary

To address the instability and hyperparameter sensitivity inherent in training neural optimal transport (OT) mappings, this paper proposes a continuous-time optimization framework grounded in displacement interpolation. We first establish the dual formulation of displacement interpolation at arbitrary time points and develop its temporal correlation theory, enabling holistic, time-coordinated optimization of the entire transport trajectory. The method integrates dynamic dual optimization, neural implicit mapping modeling, and Wasserstein distance constraints. Evaluated on unpaired image translation, it significantly enhances training stability and OT mapping accuracy: convergence becomes more robust, average mapping error decreases by 23.6%, and generalization and robustness to perturbations surpass those of existing neural OT approaches. Our core contribution lies in elevating OT learning from discrete optimization to continuous trajectory modeling—introducing a novel paradigm for neural OT that fundamentally rethinks how transport paths are learned and optimized.

Technology Category

• Search and Optimization: Learning to Search
• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems
• Computer Vision: Learning & Optimization for CV

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Optimal Transport (OT) theory investigates the cost-minimizing transport map that moves a source distribution to a target distribution. Recently, several approaches have emerged for learning the optimal transport map for a given cost function using neural networks. We refer to these approaches as the OT Map. OT Map provides a powerful tool for diverse machine learning tasks, such as generative modeling and unpaired image-to-image translation. However, existing methods that utilize max-min optimization often experience training instability and sensitivity to hyperparameters. In this paper, we propose a novel method to improve stability and achieve a better approximation of the OT Map by exploiting displacement interpolation, dubbed Displacement Interpolation Optimal Transport Model (DIOTM). We derive the dual formulation of displacement interpolation at specific time $t$ and prove how these dual problems are related across time. This result allows us to utilize the entire trajectory of displacement interpolation in learning the OT Map. Our method improves the training stability and achieves superior results in estimating optimal transport maps. We demonstrate that DIOTM outperforms existing OT-based models on image-to-image translation tasks.

Problem

• Improving stability in neural optimal transport learning
• Enhancing approximation of optimal transport maps
• Addressing training instability in max-min optimization methods

Innovation

• Uses displacement interpolation for stability
• Derives dual formulation for interpolation trajectory
• Improves OT Map approximation accuracy