OT-Net: A Reusable Neural Optimal Transport Solver
This work addresses a bottleneck in optimal transport for researchers and practitioners by providing a more efficient and reusable solver that can handle discontinuous maps, which is incremental but offers specific gains in applications like generative models.
The paper tackles the problem of efficiently computing optimal transport (OT) maps, especially for discontinuous target distributions, by introducing OT-Net, a reusable neural solver that learns Brenier's height representation to derive the OT map. The method demonstrates empirical success in applications like image generation, color transfer, and domain adaptation, with improved efficiency and the ability to handle sharp boundaries.
With the widespread application of optimal transport (OT), its calculation becomes essential, and various algorithms have emerged. However, the existing methods either have low efficiency or cannot represent discontinuous maps. A novel reusable neural OT solver OT-Net is thus presented, which first learns Brenier's height representation via the neural network to obtain its potential, and then gained the OT map by computing the gradient of the potential. The algorithm has two merits, 1) it can easily represent discontinuous maps, which allows it to match any target distribution with discontinuous supports and achieve sharp boundaries. This can well eliminate mode collapse in the generated models. 2) The OT map can be calculated straightly by the proposed algorithm when new target samples are added, which greatly improves the efficiency and reusability of the map. Moreover, the theoretical error bound of the algorithm is analyzed, and we have demonstrated the empirical success of our approach in image generation, color transfer, and domain adaptation.