Doubly Regularized Entropic Wasserstein Barycenter

Abstract We study a general formulation of regularized Wasserstein barycenters that enjoy favorable regularity, approximation, stability and (grid-free) optimization properties. This barycenter is defined as the unique probability measure that minimizes the sum of entropic optimal transport (EOT) costs with respect to a family of given probability measures, plus an entropy term. We denote it the $$(\lambda ,\tau )$$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mrow> <mml:mo>(</mml:mo> <mml:mi>λ</mml:mi> <mml:mo>,</mml:mo> <mml:mi>τ</mml:mi> <mml:mo>)</mml:mo> </mml:mrow> </mml:math> -barycenter, where $$\lambda $$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mi>λ</mml:mi> </mml:math> is the inner regularization strength and $$\tau $$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mi>τ</mml:mi> </mml:math> the outer one. This formulation recovers several previously proposed EOT barycenters for various choices of $$\lambda ,\tau \ge 0$$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mrow> <mml:mi>λ</mml:mi> <mml:mo>,</mml:mo> <mml:mi>τ</mml:mi> <mml:mo>≥</mml:mo> <mml:mn>0</mml:mn> </mml:mrow> </mml:math> and generalizes them. First, we show that, as $$\lambda , \tau \rightarrow 0$$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mrow> <mml:mi>λ</mml:mi> <mml:mo>,</mml:mo> <mml:mi>τ</mml:mi> <mml:mo>→</mml:mo> <mml:mn>0</mml:mn> </mml:mrow> </mml:math> , regularizing doubly can decrease the approximation error compared to a single regularization. More specifically, we show that for smooth densities and the quadratic cost, the leading order term of the suboptimality in the (unregularized) Wasserstein barycenter objective cancels when $$\tau \sim \frac{\lambda }{2}$$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mrow> <mml:mi>τ</mml:mi> <mml:mo>∼</mml:mo> <mml:mfrac> <mml:mi>λ</mml:mi> <mml:mn>2</mml:mn> </mml:mfrac> </mml:mrow> </mml:math> . We discuss also this phenomenon for isotropic Gaussian distributions where all $$(\lambda ,\tau )$$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mrow> <mml:mo>(</mml:mo> <mml:mi>λ</mml:mi> <mml:mo>,</mml:mo> <mml:mi>τ</mml:mi> <mml:mo>)</mml:mo> </mml:mrow> </mml:math> -barycenters have closed-form. Second, we show that for $$\lambda ,\tau &gt;0$$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mrow> <mml:mi>λ</mml:mi> <mml:mo>,</mml:mo> <mml:mi>τ</mml:mi> <mml:mo>&gt;</mml:mo> <mml:mn>0</mml:mn> </mml:mrow> </mml:math> , this barycenter has a smooth density and is strongly stable under perturbation of the marginals. In particular, it can be estimated efficiently: given n samples from each of the probability measures, it converges in relative entropy to the population barycenter at a rate $$n^{-1/2}$$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:msup> <mml:mi>n</mml:mi> <mml:mrow> <mml:mo>-</mml:mo> <mml:mn>1</mml:mn> <mml:mo>/</mml:mo> <mml:mn>2</mml:mn> </mml:mrow> </mml:msup> </mml:math> . Finally, this formulation is amenable to a grid-free optimization algorithm: we propose a simple Noisy Particle Gradient Descent method which, in the mean-field limit, converges globally at an exponential rate to the $$(\lambda ,\tau )$$ <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"> <mml:mrow> <mml:mo>(</mml:mo> <mml:mi>λ</mml:mi> <mml:mo>,</mml:mo> <mml:mi>τ</mml:mi> <mml:mo>)</mml:mo> </mml:mrow> </mml:math> -barycenter.