How to choose weights when averaging distributions via Barycenters

[conorhassan]

Hi there,

I am trying to average two distributions (Gaussian for now), informed by different data points, using the Bures functionality.

Currently, I am setting the weights to be proportional to the number of observations used to inform each distribution, but this feels relatively ad-hoc... I would like to know whether there is any functionality in ott-jax to solve for the weights or whether anyone can point me to literature that provides alternative guidance on how to choose the weights.

Thanks!

[marcocuturi]

Hi @conorhassan ! thanks for your interest in ott-jax!

I would say that in the absence of more prior information, choosing weights as (0.5, 0.5) (which would be the default) makes more sense. The standard thinking would be that each distribution (no matter how many points were used to estimate its mean/cov parameters) is a single observation. You are therefore averaging 2 objects. If you want to compute their mean (aka isobarycenter) then it probably makes sense to use 1/2 and 1/2. If for some reason you want to tilt the average towards the first, then you can consider 3/4 1/4 etc... But intuitively it feels that the averaging should not depend on the number of points (unless you would like to factor in a notion of "confidence", but even then this would not be completely orthodox! if all samples come from the same distribution you're probably better off aggregating them). What would be your use case?

[conorhassan]

Hi @marcocuturi, thanks for your response! Unfortunately, I missed this reply. We recently released a preprint (https://arxiv.org/abs/2302.03314) on federated learning of latent variable models, where in algorithm 2, we updated a subset of "global" or "shared" parameters via a Wasserstein Barycenter of Gaussians, thanks to ott-jax!

In the preprint, we went with equal weights for the Barycenter but weighted the log-likelihood component that went into each Gaussian via estimators used in the subsampling MCMC literature.

If you were interested or had ideas on whether it could expand to an approximate algorithm for non-Gaussian variational families (e.g., implicit distributions), I would be interested to discuss more!

Thanks again, Marco!