Efficient Sliced Wasserstein Distance Computation via Adaptive Bayesian Optimization

Manish Acharya, David Hyde

Keywords: Sliced Wasserstein Distance, Wasserstein Distance, Bayesian Optimization, Bayesian Quadrature, Quasi-Monte Carlo

TL;DR: We develop algorithms that involve using Bayesian optimization techniques for computing sliced Wasserstein distances (SW), and we show our algorithms can achieve state-of-the-art on this task.

Abstract: The sliced Wasserstein distance (SW) reduces optimal transport on $\mathbb{R}^d$ to a sum of one-dimensional projections, and thanks to this efficiency, it is widely used in geometry, generative modeling, and registration tasks. Recent work shows that quasi-Monte Carlo constructions for computing SW (QSW) yield direction sets with excellent approximation error. This paper presents an alternate, novel approach: learning directions with Bayesian optimization (BO), particularly in settings where SW appears inside an optimization loop (e.g., gradient flows). We introduce a family of drop-in selectors for projection directions: **BOSW**, a one-shot BO scheme on the unit sphere; **RBOSW**, a periodic-refresh variant; **ABOSW**, an adaptive hybrid that seeds from competitive QSW sets and performs a few lightweight BO refinements; and **ARBOSW**, a restarted hybrid that periodically relearns directions during optimization. Our BO approaches can be composed with QSW and its variants (demonstrated by ABOSW/ARBOSW) and require no changes to downstream losses or gradients. We provide numerical experiments where our methods achieve state-of-the-art performance, and on the experimental suite of the original QSW paper, we find that ABOSW and ARBOSW can achieve convergence comparable to the best QSW variants with modest runtime overhead. We release code with fixed seeds and configurations to support faithful replication (see supplementary material).

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

Submission Number: 1375