Posterior sampling of word embeddings

Abstract We study Bayesian inference for probabilistic word embedding models, including skip-gram with negative sampling (SGNS) and continuous bag-of-words (CBoW). These models are widely used for inference, but their posteriors are non-identifiable under general linear transformations, which prevents meaningful Bayesian uncertainty quantification. We provide a simple and principled identifiability constraint that yields a well-defined posterior without altering the data model. Building on this, we develop scalable, fully Bayesian posterior sampling algorithms for SGNS and CBoW using Pólya–Gamma augmentation. Our block Gibbs samplers yield exact samples from the posterior and scale well to large corpora due to their straightforward parallelisation. We discuss the efficient implementation of the Gibbs sampler, and provide a fast TensorFlow reference implementation. In a combination of simulations and empirical experiments, we demonstrate that the sampler is both accurate and computationally efficient. It matches the accuracy of NUTS-HMC at a fraction of the computational cost and avoids the variance underestimation of mean-field variational inference. On large real datasets, we show that posterior means yield improved predictive performance, while posterior credible intervals provide meaningful uncertainty quantification for semantic similarity.