Density-Reweighted Entropic Optimal Transport: Decoupling Geometry from Sampling Density

Dataset alignment is a fundamental challenge across science and engineering, where the goal is to establish reliable correspondences between observations across different datasets. Entropic Optimal Transport (EOT) provides a computationally tractable way to match datasets by solving an entropy-regularized optimal transport problem. However, standard EOT relies on uniform empirical marginals and squared Euclidean distances, causing it to match points based on relative sampling density rather than true geometric proximity when datasets exhibit density disparities. This leads to geometrically misleading correspondences. The authors aim to design a principled framework that can discount sampling density influences, interpolating smoothly between standard density-sensitive EOT and alignment driven purely by underlying geometry.

The authors propose a density-reweighted EOT framework governed by a density discounting factor theta ranging from zero to one. When theta equals zero, the method recovers standard EOT; when theta equals one, the population limit is completely independent of sampling densities, yielding a geometry-driven alignment. The approach reweights the transport cost, the entropic regularizer, and the marginal constraints using local density estimates. To compute this in practice without knowing the intrinsic data dimension, the authors utilize a mean-normalized kernel density estimator. The resulting density-adjusted kernel matrix is then scaled efficiently via the Sinkhorn algorithm to recover the final density-reweighted transport plan.

The study establishes rigorous theoretical convergence guarantees for the proposed framework under mild regularity conditions. Specifically, the authors prove that when the reweighted EOT plan is computed using known sampling densities, its entries converge uniformly with high probability to a family of population-level EOT plans as sample sizes grow. Furthermore, the paper provides explicit convergence rates matching standard sample-to-population bounds up to logarithmic factors. Through simulations on datasets with severe sampling density disparities, the authors demonstrate that their approach successfully recovers faithful geometric correspondences and outperforms standard EOT and Unbalanced Optimal Transport.