[Paper Review] Extreme Stochastic Variational Inference: Distributed and Asynchronous

Jiong Zhang, Parameswaran Raman+4 more|arXiv (Cornell University)|May 31, 2016

Gaussian Processes and Bayesian Inference20 references3 citations

TL;DR

This paper proposes Extreme Stochastic Variational Inference (ESVI), a distributed, asynchronous, and lock-free algorithm that enables scalable variational inference for large-scale mixture models by achieving simultaneous data and model parallelism. ESVI outperforms VI and SVI in both wall-clock time and solution quality, particularly on massive datasets like UMBC-3B with 3 billion tokens and 3 million vocabulary, while also introducing a TOPK variant to accelerate computation and reduce memory usage for high-topic models.

ABSTRACT

Stochastic variational inference (SVI), the state-of-the-art algorithm for scaling variational inference to large-datasets, is inherently serial. Moreover, it requires the parameters to fit in the memory of a single processor; this is problematic when the number of parameters is in billions. In this paper, we propose extreme stochastic variational inference (ESVI), an asynchronous and lock-free algorithm to perform variational inference for mixture models on massive real world datasets. ESVI overcomes the limitations of SVI by requiring that each processor only access a subset of the data and a subset of the parameters, thus providing data and model parallelism simultaneously. We demonstrate the effectiveness of ESVI by running Latent Dirichlet Allocation (LDA) on UMBC-3B, a dataset that has a vocabulary of 3 million and a token size of 3 billion. In our experiments, we found that ESVI not only outperforms VI and SVI in wallclock-time, but also achieves a better quality solution. In addition, we propose a strategy to speed up computation and save memory when fitting large number of topics.

Motivation & Objective

• To address the scalability limitations of traditional variational inference (VI) and stochastic variational inference (SVI) when handling datasets with billions of documents and parameters that exceed single-machine memory capacity.
• To enable both data and model parallelism in mixture models by allowing processors to update only subsets of local and global variables, rather than all parameters in a synchronized manner.
• To design a lock-free, asynchronous algorithm that avoids costly synchronization overhead in multi-core and distributed environments, improving computational efficiency.
• To demonstrate the effectiveness of ESVI on real-world large-scale datasets such as UMBC-3B, PubMed, and NY Times, where existing methods fail due to memory and scalability constraints.

Proposed method

• ESVI introduces a novel update strategy where each processor accesses only a subset of data and a subset of parameters, enabling simultaneous data and model parallelism.
• It employs an owner-computes paradigm to support asynchronous and lock-free updates, eliminating the need for global synchronization and reducing communication overhead.
• The algorithm allows global parameters to 'nomadically move' across workers, ensuring parameter mixing and convergence even with partial updates.
• It uses a TOPK strategy to store only the top-k topics per worker, significantly reducing memory usage and accelerating computation for large-topic models.
• The method is applied to Gaussian Mixture Models (GMM) and Latent Dirichlet Allocation (LDA), with empirical evaluation on real-world datasets including UMBC-3B.
• It leverages the separable structure of local variational updates, enabling stochastic optimization over subsets of data and parameters.

Experimental results

Research questions

• RQ1 Can variational inference be scaled to datasets with billions of tokens and millions of vocabulary terms where parameters exceed single-machine memory capacity?
• RQ2 How can both data and model parallelism be efficiently combined in stochastic variational inference without incurring synchronization bottlenecks?
• RQ3 Can an asynchronous, lock-free algorithm maintain convergence and solution quality while achieving high throughput on distributed systems?
• RQ4 What performance gains can be achieved by selectively updating only a subset of topics and parameters in large-scale LDA models?
• RQ5 How does the TOPK strategy improve memory efficiency and computational speed in high-topic LDA inference?

Key findings

• On the UMBC-3B dataset (3 billion tokens, 3 million vocabulary), ESVI-LDA achieved a better ELBO than distributed VI and was significantly faster, demonstrating scalability to industry-scale data.
• In multi-machine, multi-core experiments on PubMed and UMBC-3B, ESVI-LDA outperformed distributed VI in both convergence speed and final ELBO, with ESVI-LDA-TOPK showing the best performance.
• On Enron and NY Times datasets, ESVI-LDA reached comparable test perplexity to VI but in substantially less wall-clock time, especially with increasing core counts.
• The TOPK variant reduced memory usage and accelerated computation, particularly benefiting large-topic models, and achieved the fastest convergence in all multi-core and multi-machine settings.
• In single-machine settings, ESVI-GMM outperformed both VI and SVI in ELBO improvement over time, even with moderate data sizes.
• The asynchronous, lock-free design of ESVI enabled high throughput and scalability across 32 nodes and 512 cores, confirming its practicality for extreme-scale inference.