Scalable Knothe--Rosenblatt-like Heuristic Transportation Plans for Imaging Problems

In this paper, we introduce a novel formalism for computing the Wasserstein Distance between any pair of probability distributions, μ and ν. Standard approaches require solving a matching problem between two discrete distributions, which becomes computationally expensive as the dimensionality increases. To address this challenge, we propose a new family of heuristic transportation plans that extend the classic Knothe-Rosenblatt transport plan. Each heuristic plan is associated with a method for combining the two original measures into an intermediate measure, significantly reducing the number of variables required to characterise any transportation plan. Specifically, if the probability measures μ and ν have supports consisting of N and M points, respectively, our approach reduces the number of variables from N ×M to min{N, M}. We demonstrate that our method is particularly well-suited for defining a neural network to solve the optimal transport problem and validate our model through extensive numerical experiments.