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

Denys Ruban, Augusto Gerolin

Keywords: Optimal Transport (OT); chi-square regularization (chi-squared divergence); Primal–Dual Hybrid Gradient (PDHG); Primal-Dual algorithms; Transport plan; Wasserstein barycenters; estimation; statistical properties; Monge maps; Kantorovich potentials; statistical–computational trade-offs; GPU acceleration.

TL;DR: We present an adaptive PDHG solver (adaptive steps, restarts) for chi-square regularized OT, benchmarking vs Sinkhorn on Monge maps, Kantorovich potentials, and Wasserstein barycenter estimations.

Abstract: 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.

Submission Number: 45