Quantization bounds for Wasserstein metrics

Jonathan Bobrutsky, Amit Moscovich

Keywords: optimal transport, earth mover’s distance, bilevel, cryo-electron microscopy

TL;DR: Fast methods to bound Wasserstein metrics in 2D and 3D

Abstract: The Wasserstein metric is becoming increasingly important in many machine learning applications such as generative modeling, image retrieval and domain adaptation. Despite its appeal, it is often too costly to compute. This has motivated approximation methods like entropy-regularized optimal transport, downsampling, and subsampling, which trade accuracy for computational efficiency. In this paper, we consider the challenge of computing efficient approximations to the Wasserstein metric that also serve as strict upper or lower bounds, as these are essential components of branch-and-bound, A$^*$ path finding, and heuristic search techniques in tasks such as trajectory inference, alignment, and clustering. Focusing on discrete measures on regular grids, our approach involves formulating and exactly solving a Kantorovich problem on a coarse grid using a quantized measure with a tailored cost matrix, followed by an upscaling and correction stage. This is done either in the primal or dual space to obtain valid upper and lower bounds on the Wasserstein metric of the full-resolution inputs. We evaluate our methods on the DOTmark optimal transport images benchmark as well as alignment tasks on volumetric dataset of macromolecules, demonstrating a 10×–100× speedup compared to entropy-regularized OT while keeping the approximation error well below 5\% at 2D, and 30\% width bounding regions at 3D.

Primary Area: other topics in machine learning (i.e., none of the above)

Submission Number: 24510