Convex Potential Learning for Optimal Transport Between Medical Images
Authors: Hongyu Kan, Ivan Medri, Natasha Ironside, Kristofor Pas, Richard Corp, Shinjini Kundu, Gustavo Kunde Rohde
Organizations: George Washington University, Washington, DC 20052, USA · University of Illinois Chicago, Chicago, IL 60607, USA · University of Virginia, Charlottesville, VA 22903, USA · Washington University in St. Louis, St. Louis, MO 63110, USA
Computing optimal transport maps between medical images plays a fundamental role in morphological analysis. However, existing methods are either difficult to scale to large 3D medical images or rely on regularized formulations that provide only approximate transport maps. To improve the accuracy of downstream medical image analysis, we propose a novel approach for estimating optimal transport maps by parameterizing the convex potential function. By construction, the resulting transport map is guaranteed to be curl-free, consistent with Brenier's theorem. We evaluate the proposed method on both 2D and 3D medical imaging datasets. Experimental results demonstrate that our approach achieves lower transport estimation errors and produces higher-quality image reconstructions than existing methods.
Figures & tables
Figure 1: An example of optimal transport between two Gaussian distributions with different centers. I0 (left) and I1 (second from left) denote the source and target distributions, respectively. The estimated optimal transport map pushes the mass of I0 onto I1 . The reconstructed source image (second from the right) obtained via formula ( 1 ) and the estimated transport map f , closely matches I0 , demonstrating accurate push-forward consistency. The rightmost panel shows the estimated transport vector field, where the red and blue contours represent I0 and I1 , respectively, and the white arrows indicate the direction of the optimal transport map computed by the proposed method.
Cell dataset reconstruction error
ICH dataset recontruction error
Method
L1
L2
L1
L2
S. Kundu et al. [ 10 ]
0.334 ± 0.187
(21.04 ± 16.06) ∗10−5
1.098 ± 0.398
(7.89 ± 6.24) ∗10−5
Convex Potential
0.239 ± 0.068
(13.74 ± 7.88) ∗10−5
0.287 ± 0.139
(1.31 ± 1.72) ∗10−5
Table 1: L1 and L2 distances comparison
Figure 2: Visual comparison of image reconstruction using the estimated transport maps. The reconstructed images are compared with the reference image I0 according to det(Df(x))I1(f(x)) . The left panel shows representative samples from the cell dataset, while the right panel presents the middle slice of representative samples from the ICH dataset.
Figure 3: Comparison of the evolution of the L1 error, L2 error, transport cost, and curl over optimization iterations in cell dataset. Curves of the same color correspond to the same image pair.
Figure 4: Comparison of the evolution of the L1 error, L2 error, transport cost, and curl over optimization iterations in ICH dataset. Curves of the same color correspond to the same image pair.
Optimal Transport (OT) provides a principled framework for learning transformations between probability distributions from unpaired samples. In many applications, however, a single transformation must map several source distributions to a common target distribution. For example, image restoration might require handling different types of degradation without knowing the degradation of each input at inference time. Simple approaches of pooling the source distributions only encourage alignment with the target at the aggregate level and may leave individual sources misaligned. In our paper, we consider the simultaneous OT problem which formalizes the task of learning a shared transport map that minimizes the average transport cost while aligning each source distribution with a prescribed target. We propose a neural method for solving the simultaneous OT problem by learning a shared transport map that minimizes the average transport cost while aligning each source distribution with a prescribed target. We derive a max-min formulation for learning this map. We illustrate its application to image restoration, where a single model handles multiple degradation types using a common collection of clean target images.
Milena Gazdieva, Kirill Sokolov, Jiawei Chen +2
Applied AI Institute AXXX Moscow, Russia · Lomonosov Moscow State University Moscow, Russia · MIRIAI Moscow, Russia
The optimal transport (OT) map provides a geometric transformation for aligning probability distributions and has become a useful tool in machine learning. However, existing estimators of the OT map still exhibit a gap between sharp statistical guarantees and practical parametric estimation based on stable training objectives. Theoretical estimators achieve minimax optimal convergence rates, but they are typically nonparametric and can incur demanding implementation design or inference costs. Practical estimators are parametric and scalable, but their statistical guarantees remain underexplored, and their min-max, adversarial-like training objectives can be sensitive to optimization algorithms. We propose BROT (Barycentric Regression for OT), a simple two-step method that first computes the unregularized OT plan and then fits a deep neural network (DNN) to the induced barycentric targets by least-squares regression. Under standard regularity conditions, we prove that the DNN estimator of BROT attains the minimax convergence rate, when the ground-truth OT map is Lipschitz. Numerical studies on synthetic datasets and an image dataset show that BROT provides accurate map estimates, strong target distribution matching, and competitive transport costs, compared to existing estimation methods. Experiments on two downstream tasks, single-cell perturbation prediction and unsupervised domain adaptation, further suggest that the accurate estimation of BROT can translate into stronger task performance.
We propose an implicit neural formulation of optimal transport that eliminates adversarial min--max optimization and multi-network architectures commonly used in existing approaches. Our key idea is to parameterize a single potential in the Kantorovich dual and reformulate the associated c-transform as a proximal fixed-point problem. This yields a stable single-network framework in which dual feasibility is enforced exactly through proximal optimality conditions rather than adversarial training. Despite the inner fixed-point computation, gradients can be computed without differentiating through the fixed-point iterations, enabling efficient training without requiring implicit differentiation. We further establish convergence of stochastic gradient descent. The resulting framework is efficient, scalable, and broadly applicable: it simultaneously recovers forward and backward transport maps and naturally extends to class-conditional settings. Experiments on high-dimensional Gaussian benchmarks, physical datasets, and image translation tasks demonstrate strong transport accuracy together with improved training stability and favorable computational and memory efficiency.
Yesom Park, Eric Gelphman, Stanley Osher +1
Department of Mathematics, University of California, Los Angeles · Department of Applied Mathematics and Statistics, Colorado School of Mines