Primal-dual hybrid algorithms for chi-squared regularized Optimal Transport: statistical-computational trade-offs and applications to Wasserstein Barycenters

We investigate the interplay between optimization and statistics in regularized Optimal Transport problems---the standard approach for computational Optimal Transport. While regularization parameter improves optimization and tractability, it inevitably introduces bias. In this paper, we design a computational algorithm based on chi-square regularized Primal-Dual Hybrid Gradient (PDHG) that are not only efficient and robust, but also capable of accurately recovering key statistical properties, including Monge maps, Kantorovich potentials and Wasserstein barycenters. Our implementation is JAX-based, GPU-paralellizable, and incorporates adaptive step sizes and restart strategies, the latter ensuring a linear convergence rate.