PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Bayes via Barycenter in Wasserstein Space

Sanvesh Srivastava, Cheng Li, David B. Dunson

arXiv:1508.05880v4stat.ME

TL;DR

Massive datasets make full-data Bayesian sampling costly, motivating methods that partition data and combine parallel subset inferences. This paper studies WASP, which combines rescaled subset posteriors through a Wasserstein barycenter, and shows that it can closely approximate the full posterior while enabling substantial computational speed-ups. The paper also identifies assumptions and practical trade-offs governing approximation quality and subset selection.

  • Problem

    Massive datasets create storage, communication, latent-variable, and repeated full-data-pass bottlenecks for Bayesian sampling algorithms.

  • Method

    WASP partitions data, rescales subset likelihoods, samples subset posteriors in parallel, and combines them through their Wasserstein barycenter.

  • Results

    WASP achieves nearly the full posterior’s convergence rate under suitable subset growth, and in a linear model n^1/2W2(Πn, Πn) → 0 while matching the full posterior’s asymptotic mean and variance.

  • Takeaways & Limitations

    WASP extends existing samplers to massive data with minimal modifications and can produce closely agreeing results with MCMC while enabling massive computational speed-ups.

  • Takeaways & Limitations

    WASP’s guarantees rely on subset posteriors being noisy approximations of the full posterior, and choosing too many subsets can increase statistical error.

Abstract

from arXiv · show

Divide-and-conquer based methods for Bayesian inference provide a general approach for tractable posterior inference when the sample size is large. These methods divide the data into smaller subsets, sample from the posterior distribution of parameters in parallel on all the subsets, and combine posterior samples from all the subsets to approximate the full data posterior distribution. The smaller size of any subset compared to the full data implies that posterior sampling on any subset is computationally more efficient than sampling from the true posterior distribution. Since the combination step takes negligible time relative to sampling, posterior computations can be scaled to massive data by dividing the full data into a sufficiently large number of data subsets. One such approach relies on the geometry of posterior distributions estimated across different subsets and combines them through their barycenter in a Wasserstein space of probability measures. We provide theoretical guarantees on the accuracy of approximation that are valid in many applications. We show that the geometric method approximates the full data posterior distribution better than its competitors across diverse simulations and reproduces known results when applied to a movie ratings database.

1 Introduction

The paper addresses computational bottlenecks in Bayesian sampling for massive datasets with a general divide-and-conquer method. WASP combines parallel subset posterior results through their Wasserstein barycenter and develops theoretical guarantees for its accuracy.

  • Approach: Divide-and-conquer methods partition data into subsets and perform posterior sampling in parallel before combining the results.The proposed technique is designed to combine outputs from any posterior sampling algorithm applied to data subsets.
  • Motivation: Massive datasets create storage, communication, latent-variable, and computational bottlenecks for existing MCMC and SMC algorithms.Hierarchical models may require O(n) latent variables, and full-data sampling can require repeated passes through all observations.
  • Related methods: Existing combination strategies include subsampling, modified dynamics, consensus averaging, and Wasserstein-based geometric centers.Consensus methods rely heavily on normality, while Wasserstein medians may ignore valuable information from some subset posteriors.
  • Contributions: WASP relaxes the iid assumption to inid data and can achieve nearly the full posterior’s convergence rate when the number of subsets is appropriately chosen.For linear models with location-scale errors, WASP and the full posterior have the same asymptotic mean and variance.
  • Contributions: The paper studies WASP theoretically and applies it across a variety of practical problems.The method is presented as general rather than as a new sampling algorithm.

2 Preliminaries

The preliminaries define Wasserstein geometry and explain how subset posteriors are constructed and combined. Stochastic approximation rescales subset likelihoods so their posteriors can approximate the full posterior while retaining standard sampling algorithms.

  • Wasserstein preliminaries: The order-2 Wasserstein space contains probability measures with finite second moment and is equipped with the order-2 Wasserstein distance.In the applications, the underlying metric is Euclidean, and the space and distance are referred to simply as Wasserstein space and distance.
  • Wasserstein preliminaries: A Wasserstein barycenter generalizes the Euclidean sample mean to a collection of probability measures.Although generally analytically intractable, its empirical version can reduce to a linear program.
  • WASP construction: WASP combines subset posterior distributions through their Wasserstein barycenter, which can be computed efficiently using a linear program.The barycenter is used as an approximation to the full-data posterior for inference and prediction.
  • Subset construction: Data are partitioned into k subsets, typically of equal size m, with n = km.Subset likelihoods are formed from the observations assigned to each subset.
  • Stochastic approximation: Stochastic approximation raises each subset likelihood to a power γ, using γ = k in the presented setting, to compensate for each subset’s smaller data fraction.The resulting subset posteriors have variances of the same order as the full posterior and can be sampled without explicitly replicating observations.
  • Illustration: In logistic regression, subset posterior contours have spreads similar to the full posterior but often exhibit bias and fail to concentrate at the true parameter.This illustrates why subset posteriors are treated as noisy approximations before combination.

3 Wasserstein Posterior (WASP): The general framework

WASP combines subset posterior distributions through their barycenter in Wasserstein space, replacing full-data posterior inference for massive datasets. The paper establishes approximation guarantees for independent non-identically distributed data and identifies conditions under which WASP matches full-posterior behavior.

  • Definition and estimation: WASP combines subset posterior distributions through their barycenter in P2(Θ), a geometric center computable by linear programming.The subset posteriors can be represented by empirical measures from samples generated by HMC, MCMC, SMC, or importance sampling.
  • Approximation error for inid data: For the weighted linear model, WASP and the full-data posterior have the same posterior inference up to o(n−1) terms, while computation is reduced by a factor of k.The divide-and-conquer WASP computation has complexity O(km^2), compared with the full-data posterior computation.
  • Approximation error for inid data: Theorem 3.1 shows that WASP and the true posterior have the same posterior mean and variance, with differences only in higher-order terms of n.The result also applies to block-dependent data and error distributions with mean zero and covariance Σ.
  • General convergence rates: With regular models and slowly increasing k, WASP reaches a near-optimal W2 convergence rate, differing from the parametric rate by a logarithmic factor.For α = 1 and k = O(log^c n), the rate is Op(n−1/2 log^(c/2+1) n), compared with Op(n−1/2) for the true posterior.
  • General convergence rates: The guarantees require subset posteriors to be noisy approximations of the full-data posterior; when assumptions fail, WASP can approximate it poorly.In the rare-events example, many subsets contain only zeros for θ0 = 10−5 or 10−6, severely violating assumption (A1) and causing WASP failure.

4 Experiments

The experiments evaluate WASP and competing posterior approximations across Gaussian mixtures, linear mixed-effects models, and additional simulated and real-data settings. WASP generally provided accurate approximations, often outperforming competitors for random-effects variances, covariances, and joint covariance distributions, while retaining scalable subset-based computation.

  • Finite mixture of Gaussians: In the Gaussian-mixture experiment, CMC, SDP, and WASP accurately approximated posterior distributions of correlations, while VB underestimated their uncertainty.For the credible bands of g(x), CMC, VB, and WASP were accurate, whereas SDP failed because of a numerical matrix-inversion error.
  • Finite mixture of Gaussians: WASP’s accuracy for nonlinear functions of mixture-model parameters was unaffected by the form of the parameters used in combination.This empirically verified Corollary 3.5; comparable theoretical guarantees were unavailable for CMC and SDP.
  • Linear mixed effects model: In the linear mixed-effects experiment, CMC, SA, SDP, and WASP accurately approximated marginal posteriors for fixed effects and random-effects covariances.All methods except SGLD were significantly faster than full-data posterior sampling, with SA being the fastest.
  • Linear mixed effects model: WASP’s accuracy remained stable as the number of subsets varied, whereas CMC and SDP accuracy depended on the subset count and could deteriorate with smaller subsets.The combination step took negligible time relative to sampling, supporting larger datasets through sufficiently many subsets.
  • MovieLens ratings data: WASP outperformed competing methods for posterior variances, covariances, and covariance pairs of random effects in the movie-ratings analysis.ADVI, SA, and SGLD underperformed for fixed effects and the random-effects covariance matrix, while CMC and SDP were less accurate for some covariance structures.

5 Discussion

WASP provides a computationally efficient approximation to full-data posterior distributions, with theoretical convergence guarantees and empirical agreement with MCMC. Its practical scope is bounded by assumptions needed for stochastic approximation and uncertainty quantification, and by unresolved choices about the number of subsets.

  • Contributions: WASP approximates posterior distributions for parameters and their functions in large-sample settings while enabling parallel computation.The method replaces inefficient full-data posterior computation with subset posterior computations and combination.
  • Theory: WASP's convergence rate in W2 matches the optimal parametric rate up to a logarithmic factor when the number of subsets grows slowly with the full dataset.
  • Empirical results: The discussion reports accuracy evaluations for approximate covariance and two-dimensional joint posterior distributions across repeated replications.Tables 8 and 9 describe averaged accuracies and Monte Carlo errors for covariance-related posteriors.
  • Empirical results: WASP and MCMC agree closely across several diverse examples, while WASP yields massive computational speed-ups.
  • Limitations: Within-subset independence is required for part of the convergence proof and the justification of stochastic approximation.
  • Limitations: The optimal practical choice of subset count k remains unclear because larger k can improve computation but may increase statistical error.

A.1 Proof of Theorem 3.1

The proof derives asymptotic orders for posterior mean and covariance terms, then combines them to establish the theorem's stated asymptotic result.

  • A.1 Proof of Theorem 3.1: The proof analyzes the expectation and covariance contributions to the relevant posterior-distance expression.
  • A.1 Proof of Theorem 3.1: Because p is fixed, several displayed equalities simplify the asymptotic expressions used in the proof.
  • A.1 Proof of Theorem 3.1: The covariance matrices of the full-data and subset posteriors are bounded through matrix inequalities and trace calculations.
  • A.1 Proof of Theorem 3.1: The proof concludes by combining the asymptotic mean and variance expressions.

A.2 Proof of Theorem 3.4

The proof bounds subset-posterior behavior under the stated assumptions using metric comparisons, probability bounds, and a subset-size condition.

  • A.2 Proof of Theorem 3.4: The proof decomposes the W2 distance between a subset posterior and the Dirac measure at the true parameter.
  • A.2 Proof of Theorem 3.4: The argument replaces the rho metric with a pseudo Hellinger distance under condition (A2).
  • A.2 Proof of Theorem 3.4: Bounds for the numerator and denominator yield a high-probability control of the subset-posterior distance.
  • A.2 Proof of Theorem 3.4: The second term in the bound dominates the other terms when the subset size satisfies m ≲n.
  • A.2 Proof of Theorem 3.4: The resulting explicit bound has a constant depending on the assumptions and auxiliary lemma constants.
  • A.2 Proof of Theorem 3.4: The proof invokes assumption (A5), the Cauchy-Schwarz inequality, and the subset bound (12).

B Univariate density estimation

The univariate density-estimation example uses a Dirichlet process mixture model and data augmentation to make subset-posterior sampling feasible for WASP.

  • B Univariate density estimation: The example models scalar observations with a hierarchical stick-breaking representation of Dirichlet process mixtures.
  • B Univariate density estimation: Direct use of existing samplers is unclear when the prior is raised to the power 1/k for competing subset methods.
  • B Univariate density estimation: Data augmentation enables sampling from the subset posterior density despite the powered prior.
  • B Univariate density estimation: Stochastic approximation raises the augmented likelihood to the power γ.
  • B Univariate density estimation: The resulting full conditional densities are tractable in terms of standard distributions, including the latent-variable posterior.

C Linear program

The linear program constrains each Tj entry to [0,1] and can be solved with several standard solvers.

  • Each entry (Tj)uv is constrained between 0 and 1 across the specified indices.
  • The program can be solved in Matlab or R using several linear-programming algorithms.Examples include algorithms by Cuturi and Doucet (2014) and Srivastava et al. (2015).

1 Technical Lemmas

The technical development establishes concentration and entropy bounds for subset likelihood processes, then derives consequences for posterior concentration and Wasserstein barycenters.

  • Technical Lemmas: A generalized Wong–Shen concentration inequality is introduced to support the main theorem under the independent, non-identically distributed setup.The proof uses lower-truncated log likelihood ratios and generalized bracketing entropy.
  • Technical Lemmas: The proof selects constants and radii to obtain uniform probability bounds for all subsets when the subset size is sufficiently large.The construction specifies choices such as q1 = 0.5 and q2 = 1/230.
  • Technical Lemmas: A separate lemma bounds the W2 distance from a barycenter of posterior measures to a target point through optimal transport maps.The proof uses the barycenter representation and properties of the ρ metric.

2 Experiments

The experiments evaluate subset-based posterior approximations for mixture and mixed-effects models, including comparisons on simulated data and MovieLens.

  • Experiments: For the Gaussian-mixture experiment, observations are clustered and split so every subset represents the full-data mixture proportions.Subset posteriors use stochastic approximation by raising each subset likelihood to the power γ.
  • Experiments: The Gaussian-mixture model uses Dirichlet, normal, and inverse-Wishart priors, with analytically tractable Gibbs-sampling conditionals.The sampler updates mixture weights, component means, covariances, and latent allocations.
  • Experiments: The linear mixed-effects analysis obtains subset posterior samples using a stochastic-approximation likelihood implemented in Stan.The implementation uses Stan's increment_log_prob function and default priors for β and Σ.
  • Experiments: Accuracy comparisons include ADVI, CMC, SA, SDP, SGLD, and WASP for simulated mixed-effects models and MovieLens fixed effects.The reported accuracies are averaged over simulation replications, with Monte Carlo errors in parentheses.
  • Experiments: Computation time is compared across linear mixed-effects simulations, probabilistic PARAFAC simulations, and MovieLens analysis.The comparison includes methods such as CMC, SDP, and WASP.
Loading 1508.05880v4…