Bayesian Learning with Wasserstein Barycenters
This work addresses model selection for Bayesian practitioners, offering a parameter-free framework that is incremental in combining optimal transport with existing Bayesian methods.
The paper tackles the problem of model selection in Bayesian learning by introducing the Bayesian Wasserstein barycenter (BWB), a novel estimator based on optimal transport, which extends classic strategies and is shown to be statistically consistent and computable via stochastic gradient descent.
We introduce and study a novel model-selection strategy for Bayesian learning, based on optimal transport, along with its associated predictive posterior law: the Wasserstein population barycenter of the posterior law over models. We first show how this estimator, termed Bayesian Wasserstein barycenter (BWB), arises naturally in a general, parameter-free Bayesian model-selection framework, when the considered Bayesian risk is the Wasserstein distance. Examples are given, illustrating how the BWB extends some classic parametric and non-parametric selection strategies. Furthermore, we also provide explicit conditions granting the existence and statistical consistency of the BWB, and discuss some of its general and specific properties, providing insights into its advantages compared to usual choices, such as the model average estimator. Finally, we illustrate how this estimator can be computed using the stochastic gradient descent (SGD) algorithm in Wasserstein space introduced in a companion paper arXiv:2201.04232v2 [math.OC], and provide a numerical example for experimental validation of the proposed method.