[Paper Review] Doubly Regularized Entropic Wasserstein Barycenters

Lénaïc Chizat|arXiv (Cornell University)|Mar 21, 2023

Groundwater flow and contamination studies4 citations

TL;DR

This paper introduces doubly regularized entropic Wasserstein barycenters, a unified framework that combines inner (entropic) and outer (differential entropy) regularization to improve stability, smoothness, and approximation of unregularized Wasserstein barycenters. For the choice τ = λ/2, the method achieves a suboptimality gap of order λ², demonstrating debiasing, and enables grid-free optimization via noisy particle gradient descent with global convergence.

ABSTRACT

We study a general formulation of regularized Wasserstein barycenters that enjoys 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 $(λ,τ)$-barycenter, where $λ$ is the inner regularization strength and $τ$ the outer one. This formulation recovers several previously proposed EOT barycenters for various choices of $λ,τ\geq 0$ and generalizes them. First, in spite of -- and in fact owing to -- being \emph{doubly} regularized, we show that our formulation is debiased for $τ=λ/2$: the suboptimality in the (unregularized) Wasserstein barycenter objective is, for smooth densities, of the order of the strength $λ^2$ of entropic regularization, instead of $\max\{λ,τ\}$ in general. We discuss this phenomenon for isotropic Gaussians where all $(λ,τ)$-barycenters have closed form. Second, we show that for $λ,τ>0$, 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}$. And finally, this formulation lends itself naturally to a grid-free optimization algorithm: we propose a simple \emph{noisy particle gradient descent} which, in the mean-field limit, converges globally at an exponential rate to the barycenter.

Motivation & Objective

• To unify and generalize existing entropic optimal transport barycenters under a single framework with both inner and outer regularization.
• To establish theoretical properties such as regularity, stability, and convergence rates for the proposed barycenter formulation.
• To demonstrate that the (λ, λ/2)-barycenter achieves a debiased approximation of the unregularized Wasserstein barycenter with suboptimality O(λ²).
• To develop a grid-free optimization algorithm—noisy particle gradient descent—that converges globally to the barycenter in the mean-field limit.
• To show that the barycenter has a smooth density and converges in relative entropy at rate n⁻¹ᐟ² under sampling from the input measures.

Proposed method

• The (λ, τ)-barycenter is defined as the minimizer of Fλ,τ(μ) = Gλ(μ) + τH(μ), where Gλ(μ) is the sum of entropic optimal transport costs and H(μ) is the differential entropy.
• Inner regularization (λ) ensures finite and well-behaved EOT costs via relative entropy with respect to μ⊗ν, while outer regularization (τ) promotes smoothness and stability.
• The formulation is shown to be equivalent to a standard EOT barycenter with a modified reference measure σref = [(dμ/dx)^α dx] ⊗ [(dν/dy)^α dy], where α = 1 − τ/λ.
• A noisy particle gradient descent (NPGD) algorithm is proposed for grid-free optimization, which converges globally in the mean-field limit at an exponential rate.
• Theoretical analysis leverages PDEs and mean-field limits to establish convergence and stability, with a focus on the critical case τ = λ/2.
• Numerical validation uses gradient ascent on the dual problem and compares with Sinkhorn divergence and unregularized barycenters on 1D and 2D examples.

Figure 1 : Wasserstein distance $d^{-1}W_{2}(\mu^{*}_{\lambda,\tau},\mu^{*}_{0,0})=(\sqrt{b}-\sqrt{a})^{2}$ between $\mu^{*}_{\lambda,\tau}$ and $\mu^{*}_{0,0}$ from ( 21 ), with $a=1.0$ . The white line shows the best debiasing choice of $\tau^{*}(\lambda)$ from ( 22 ) for which this distance is $0

Experimental results

Research questions

• RQ1 How does combining inner and outer entropic regularization affect the approximation quality of Wasserstein barycenters?
• RQ2 Can the (λ, τ)-barycenter achieve a suboptimality gap of O(λ²) instead of O(max{λ, τ}) when τ = λ/2, indicating debiasing?
• RQ3 Does the doubly regularized formulation ensure smoothness and strong stability of the barycenter under perturbations of the input measures?
• RQ4 Can a grid-free optimization method like noisy particle gradient descent achieve global convergence to the barycenter?
• RQ5 What is the statistical convergence rate of the empirical barycenter to the population barycenter under i.i.d. sampling?

Key findings

• For τ = λ/2, the (λ, τ)-barycenter achieves a suboptimality gap of O(λ²) in the unregularized Wasserstein barycenter objective, demonstrating debiasing for smooth densities.
• When λ, τ > 0, the barycenter has a smooth density and is strongly stable under perturbations of the input measures.
• Given n i.i.d. samples from each input measure, the empirical barycenter converges to the population barycenter in relative entropy at rate n⁻¹ᐟ².
• The (λ, τ)-barycenter can be interpreted as a standard EOT barycenter with a modified reference measure σref involving power-weighted densities.
• Noisy particle gradient descent converges globally to the barycenter in the mean-field limit at an exponential rate, escaping local minima.
• Numerical experiments confirm that τ = λ/2 yields the best approximation of the unregularized Wasserstein barycenter, outperforming Sinkhorn divergence barycenters which exhibit oscillatory behavior.

(a) Densities $(\nu_{k})_{k=1}^{3}$ and $\mu^{*}_{0,0}$