Source-linked AI summary

Large-Scale Stochastic Sampling from the Probability Simplex

Jack Baker, Paul Fearnhead, Emily B Fox, Christopher Nemeth

arXiv:1806.07137v2stat.COcs.LGstat.ML

TL;DR

SGMCMC struggles with constrained spaces, particularly sparse simplex settings where many components are near zero. The paper proposes SCIR, using the Cox-Ingersoll-Ross process to avoid discretization error; it reports improved performance over other SGMCMC methods on sparse simplex sampling.

  • Problem

    SGMCMC struggles on constrained spaces, and many large-scale models use sparse simplex parameters with components close to zero.

    Source: §1 · p. 1, §2.3 · p. 3
  • Method

    SCIR uses the Cox-Ingersoll-Ross process to simulate gamma variables without discretization error and transforms them into simplex samples.

    Source: §1 · p. 2, §3 · p. 4
  • Results

    The authors report improved performance over other SGMCMC methods for sampling from sparse simplex spaces.

    Source: §6 · p. 9
  • Takeaways & Limitations

    Because many important large-scale models are sparse, the authors describe SCIR's improved sparse-simplex sampling as an important contribution.

    Source: §6 · p. 9

    For the Bayesian nonparametric mixture model, stochastic-gradient methods do not perform well against exact Gibbs sampling, apparently because they get stuck in local stationary points.

    Source: §5.2 · p. 9

Abstract

from arXiv · show

Stochastic gradient Markov chain Monte Carlo (SGMCMC) has become a popular method for scalable Bayesian inference. These methods are based on sampling a discrete-time approximation to a continuous time process, such as the Langevin diffusion. When applied to distributions defined on a constrained space the time-discretization error can dominate when we are near the boundary of the space. We demonstrate that because of this, current SGMCMC methods for the simplex struggle with sparse simplex spaces; when many of the components are close to zero. Unfortunately, many popular large-scale Bayesian models, such as network or topic models, require inference on sparse simplex spaces. To avoid the biases caused by this discretization error, we propose the stochastic Cox-Ingersoll-Ross process (SCIR), which removes all discretization error and we prove that samples from the SCIR process are asymptotically unbiased. We discuss how this idea can be extended to target other constrained spaces. Use of the SCIR process within a SGMCMC algorithm is shown to give substantially better performance for a topic model and a Dirichlet process mixture model than existing SGMCMC approaches.

1 Introduction

SGMCMC scales Bayesian inference through minibatch gradients, but stochastic-gradient and time-discretization errors make its target approximate. The paper focuses on constrained, often sparse simplex parameters and proposes SCIR as a discretization-free alternative.

1 Introduction

  • SGMCMC replaces costly full-data gradients with minibatch estimates, while discretizing an otherwise intractable continuous process; both introduce error into the posterior target.In the non-stochastic-gradient setting, Metropolis–Hastings can correct discretization error, but this is not generically feasible in the stochastic-gradient setting.Source: §1 · p. 1
  • Constrained spaces challenge SGMCMC because naive updates can propose invalid or inaccurate values, with errors especially pronounced near boundaries.Source: §1 · p. 1
  • Large-scale models including topic, network, mixture, and mixed-membership models use simplex parameters that often become sparse, placing many components near the boundary.Source: §1 · p. 2
  • Although SGRLD improved simplex sampling, its discretization error can cause asymptotic bias and substantial inaccuracy at the boundary.Source: §1 · p. 2
  • SCIR uses the Cox-Ingersoll-Ross process, whose known transition equation enables direct gamma-variable simulation without discretization error, followed by a transformation to the simplex.The paper also characterizes its non-asymptotic moment generating function, mean, and variance, and discusses extension to spaces such as (0, ∞).Source: §1 · p. 2
  • The paper demonstrates SCIR across a broad model class, including a scalable Dirichlet-process sampler built from a minibatched slice sampler.The authors note that SGMCMC for Bayesian nonparametric models had not, to their knowledge, been explored.Source: §1 · p. 2

2 Stochastic Gradient MCMC on the Probability Simplex

The section introduces minibatch-gradient sampling and simplex parameterizations, then examines why SGRLD can misrepresent sparse simplex posteriors. It traces the boundary problem to discretization bias, whose relative impact grows near zero.

2.1 Stochastic Gradient MCMC

  • SGMCMC reduces per-iteration data and memory costs by using subsets of observations, but approximates continuous dynamics with Euler–Maruyama and replaces the full gradient with an unbiased estimate.The resulting chain approximately targets the posterior.Source: §2.1 · pp. 2–3

2.2 SGMCMC on the Probability Simplex

  • Normalizing independent gamma variables by their sum produces a Dirichlet-distributed vector, providing a transformation between gamma variables and the simplex.Source: §2.2 · p. 3
  • For categorical observations with a Dirichlet prior, the simplex posterior is Dirichlet and can be calculated exactly in the simple case.The motivating setting instead has changing latent variables and large N, where standard MCMC is prohibitively slow and subsampled latent data can propose moves to the simplex parameter.Source: §2.2 · p. 3

2.3 SGRLD on Sparse Simplex Spaces

  • SGRLD addresses simplex constraints by transforming parameters to R^d and using a Riemannian SGLD variant, but discretization still causes boundary problems.Its parameterization uses a mirrored gamma prior so the simplex parameter retains a Dirichlet prior.Source: §2.2 · p. 3, §2.3 · p. 3
  • Sparse simplex parameters have many components close to zero; applications such as LDA and sparse-network modeling commonly exhibit this boundary setting.Source: §2.3 · p. 3
  • In log-scale boxplots, SGRLD rarely reaches small fifth-component values in the sparse Dirichlet experiment, yielding a poor posterior approximation against exact draws.Figure 1 compares 1,000-iteration SGRLD and SCIR samples with 1,000 exact independent samples; the experiment sets α to 0.1 and leaves components 4–10 unobserved.Source: §2.3 · p. 4, Figure 1 · p. 4
  • Near the boundary, discretization bias has leading order O(h), and its fixed absolute error becomes large relative to small θ, with upward bias from the positivity constraint.A good approximation can then require prohibitively small h; the next method is designed to make this asymptotic bias zero for any stepsize.Source: §2.3 · p. 4

3 The Stochastic Cox-Ingersoll-Ross Algorithm

SCIR uses the exact transition of a Cox-Ingersoll-Ross process to sample gamma variables and transform them into simplex samples without discretization error. With minibatches, it performs better than SGRLD on sparse simplexes and similarly on dense ones, at the same per-iteration cost.

3 The Stochastic Cox-Ingersoll-Ross Algorithm

  • Modeling each θ_j independently as a gamma posterior preserves the desired Dirichlet posterior, and the CIR process's exact transition enables simulation without discretization error.The CIR process is nonnegative with probability one, so SCIR avoids the absolute-value correction required by SGRLD.Source: §3 · pp. 4–5

3.1 Adapting for Large Datasets

  • With minibatches, SCIR replaces the CIR parameter with an unbiased estimate and targets the desired gamma distribution, while minibatch-estimation noise adds error.The transformation linking CIR to a Langevin diffusion supports this minibatch construction; SCIR has the same per-iteration computational cost as SGRLD.Source: §3.1 · p. 5

3.2 SCIR on Sparse Data

  • On sparse Dirichlet posteriors, SCIR handles small components more readily and achieves better dKS results than SGRLD; on dense posteriors, the samplers perform similarly.At larger minibatch sizes, SCIR converges toward the exact sampler.Source: §3.2 · p. 6, Figure 2 · p. 6

3.3 Extensions

  • SCIR can extend to most Dirichlet constructions with unintegrated latent variables and to positive gamma-constrained spaces; other tractable-transition diffusions can support related samplers.Geometric Brownian motion is given as a route to a stochastic sampler for lognormal distributions on (0, ∞).Source: §3.3 · p. 6

4 Theoretical Analysis

The theoretical analysis compares SCIR with an exact CIR process targeting Gamma(a, 1), deriving the approximate process's moment generating function and its bias and variance. The results establish asymptotic unbiasedness while noting variance inflation relative to exact CIR.

4 Theoretical Analysis

  • The analysis targets Gamma(a, 1), with a = α + Σ_i z_i, and compares an M-step SCIR run of stepsize h with exact CIR run for time Mh.The comparison derives the moment generating function of the SCIR sample in terms of the exact process.Source: §4 · p. 6
  • The SCIR moment generating function is expressed using the exact CIR process's moment generating function plus a product-form error term.The comparison couples the two processes from the same starting value and uses the same noise.Source: §4 · pp. 6–7
  • SCIR is asymptotically unbiased, and its non-asymptotic variance is the exact CIR variance plus a term involving the variance of the estimate ˆa.The SCIR variance is somewhat higher than the exact CIR variance.Source: §4 · p. 7

5 Experiments

The experiments compare SCIR and SGRLD on LDA and a Bayesian nonparametric mixture, evaluating held-out predictive performance. SCIR shows better predictive behavior in both comparisons, while stochastic-gradient methods remain weaker than exact Gibbs sampling on the mixture model.

5 Experiments

  • The evaluation measures held-out predictive performance over five seeds, with tuned stepsizes and hyperparameters and a fixed minibatch size for each experiment.The LDA perplexity is evaluated every five iterations on 1,000 held-out documents; mixture inference uses minibatches of 1,000 users and a fixed stepsize.Source: §5 · p. 7, §5.1 · p. 8, §5.2 · p. 9

5.1 Latent Dirichlet Allocation

  • LDA is tested on scraped Wikipedia documents, where sparse topic-word distributions make it a challenging simplex-sampling task.The vocabulary contains approximately 8,000 words, and each minibatch contains 50 documents.Source: §5.1 · p. 8

5.1 Latent Dirichlet Allocation; 5.2 Bayesian Nonparametric Mixture Model

  • SCIR outperforms SGRLD on both tested models: it keeps improving on LDA as SGRLD converges, and reaches a lower test log-predictive score on the mixture.The figure plots LDA perplexity and mixture-model test log predictive across the two comparisons.Source: §5.1 · p. 8, Figure 3 · p. 8, §5.2 · p. 9

5.2 Bayesian Nonparametric Mixture Model

  • For categorical data, the authors develop a scalable, truncation-free SGMCMC sampler for Dirichlet processes using a minibatched slice sampler.The method is applied to an anonymous Microsoft user dataset with about 105 instances, 30,000 users, and 294 categories; the data are expected to be sparse.Source: §5.2 · pp. 8–9
  • On the complicated mixture model, stochastic-gradient methods perform poorly against exact Gibbs sampling, apparently because they become stuck in local stationary points.The comparison is reported in the Supplementary Material; SCIR still outperforms SGRLD within the stochastic-gradient comparison.Source: §5.2 · p. 9

6 Discussion

The discussion presents SCIR as a discretization-free, asymptotically unbiased SGMCMC method for simplex spaces. It highlights theoretical properties, empirical performance on sparse simplex spaces, and possible extension to other constrained spaces.

6 Discussion

  • SCIR is presented as an SGMCMC method for simplex spaces with no discretization error and asymptotically unbiased samples.The discussion says experiments show improved performance over other SGMCMC methods on sparse simplex spaces.Source: §6 · p. 9
  • The paper derives non-asymptotic variance and moment-generating-function properties for the sampler.Source: §6 · p. 9
  • The authors discuss extending the sampler to other constrained spaces and identify reducing its non-asymptotic variance with control variates as future work.Source: §6 · p. 9

A.1 Proof of Proposition 2.1

The proof checks that the identity test function preserves SGLD's local weak error order. Smoothness of the identity and agreement of weak Taylor terms support the O(h) conclusion.

A.1 Proof of Proposition 2.1

  • The identity test function is infinitely differentiable, so it is sufficiently smooth for the local weak-error argument.Source: §A.1 · p. 12
  • The proof uses the weak Taylor expansion of the Langevin diffusion under the stated assumptions and shows that the terms agree up to O(h).It concludes that SGLD retains local weak error O(h) even with the identity test function.Source: §A.1 · p. 12

A.2 Proof of Theorem 3.1

The appendix connects gamma-posterior sampling to the CIR process through a transformed Langevin diffusion, then derives the SCIR moment-generating function using the CIR transition law.

A.2 Proof of Theorem 3.1

  • The generalized-gamma target is the approximate target of the discretization-free SGLD algorithm, whose update is identified as an SCIR process.The proof relates this target to a gamma posterior and uses an unbiased estimate of a in the SGLD construction.Source: §A.2 · pp. 12–13
  • Transforming the Langevin diffusion for a generalized-gamma variable back to the gamma variable yields the integral form of the CIR process.The transformed variable is θ_t = U_t^2/4.Source: §A.2 · pp. 12–13

A.3 Proof of Theorem 4.1

  • The SCIR moment-generating function is derived from the CIR transition moment-generating function and recursive expectations.Source: §A.3 · p. 13

A.4 Proof of Theorem 4.2

The proof establishes its result by tracking moments through the minibatch noise, using the moment-generating function and the unbiasedness and independence of minibatch estimates.

A.4 Proof of Theorem 4.2

  • The proof begins from Theorem 4.1 and differentiates the moment-generating function to obtain the required moment relations.It first notes that the relevant functions equal 1 at zero.Source: §A.4 · p. 14
  • Unbiased minibatch estimates are used to express the moments through the accumulated minibatch noise.The proof labels the noise through iteration M as B_M and uses E[â_i] = a.Source: §A.4 · p. 14
  • The proof obtains the second moment by differentiating again, using independence of distinct minibatch estimates, and evaluating the moment-generating function.Source: §A.4 · pp. 14–15

B.1 Proof of Lemma A.1

The block proves two lemmas by induction, specifies a convenient CIR parameterization, and develops the Dirichlet-process slice-sampling construction used by the stochastic sampler.

B.1 Proof of Lemma A.1

  • The proof of Lemma A.1 proceeds by induction and concludes that the result holds for every n in N.Source: §B.1 · p. 15

B.2 Proof of Lemma A.2

  • The proof proceeds by induction: it establishes the result at n = 1, assumes it for all n ≤ k, uses Lemma A.1 to prove n = k + 1, and concludes it holds for all n ∈ N.Source: §B.2 · p. 15

C CIR Parameter Choice

  • Setting 2b = σ^2 gives the CIR process a Gamma(a, 1) stationary distribution, and b is fixed to 1 because b and h are unidentifiable.Source: C CIR Parameter Choice · p. 16

D.1 Dirichlet Processes

  • A Dirichlet process is parameterized by a base distribution G0 and concentration α, and its stick-breaking construction uses V_j distributed as Beta(1, α).Source: §D.1 · p. 16

D.2 Slice sampling Dirichlet process mixtures

  • The slice sampler adds latent slice variables so only finitely many weights need updating, making its conditional-independence structure suitable for stochastic sampling.It is based on stick-breaking rather than the sequential Pólya-urn formulation.Source: §D.2 · p. 17

D.3 Stochastic Sampler

The stochastic slice sampler exploits conditionally independent updates to use minibatches while preserving the posterior as the invariant distribution for the marginal updates.

D.3 Stochastic Sampler

  • Conditional independence in the slice sampler permits minibatch updates of z and u while leaving the posterior as the invariant distribution.Source: §D.3 · p. 17
  • For the current minibatch, the sampler recomputes Z* and S, with Z* the maximum allocation and S the selected data indices.Only n allocation values change, so Z* and S can be recalculated in O(n) time.Source: §D.3 · p. 17
  • The sampler updates stick breaks v_j with SCIR and component parameters θ_j with SGMCMC, using minibatch-based allocation counts and likelihood contributions.Source: §D.3 · p. 18
  • It samples slice variables and allocations for minibatch observations, then adds components and associated parameters as needed to reach k*.Source: §D.3 · p. 18

E Experiments

Experiments compare SCIR and SGRLD on synthetic simplex targets and two models, including comparisons with exact Gibbs samplers. SCIR is competitive with exact Gibbs for LDA, while both stochastic methods trail Gibbs for the mixture model.

E.1 Comparison with Gibbs

  • SCIR is competitive with exact collapsed Gibbs for LDA, although the large dataset prevented running Gibbs for many iterations.Source: §E.1 · p. 18
  • For the Bayesian nonparametric mixture, SCIR outperforms SGRLD but both lag converged Gibbs, with stochastic methods getting stuck in local stationary points.The authors identify large-scale Bayesian nonparametric inference as a continuing challenge for stochastic samplers.Source: §E.1 · p. 18

E.2 Synthetic

  • The synthetic experiments compare tuned SGRLD and SCIR against exact sampling on sparse and dense Dirichlet posteriors using Kolmogorov-Smirnov distance.The sparse case examines a component near zero; performance is measured across five seeds.Source: §E.2 · p. 19

E.3 Latent Dirichlet Allocation

  • For LDA, the experiments use a decreasing stepsize schedule that benefits both SGRLD and SCIR, with hyperparameters detailed for the experiment.Source: §E.3 · p. 19

E.4 Bayesian Nonparametric Mixture

  • The Bayesian nonparametric mixture experiment reports its sampler hyperparameter settings, distinguishing stepsizes for θ from those for the Dirichlet-process weights ω.Source: §E.4 · p. 20
Loading 1806.07137v2…