Scaling Tractable Probabilistic Circuits: A Systems Perspective

TL;DR

This paper introduces PyJuice, a GPU system that significantly accelerates training and reduces memory usage of Probabilistic Circuits, enabling larger models and improved performance on image and language tasks.

Contribution

PyJuice provides a novel GPU implementation for Probabilistic Circuits, achieving 1-2 orders of magnitude speedup and 2-5x memory reduction over prior systems, facilitating large-scale model training.

Findings

PyJuice is 1-2 orders of magnitude faster than existing systems.

PyJuice consumes 2-5x less GPU memory, enabling larger models.

PyJuice improves state-of-the-art PCs on image and language datasets.

Abstract

Probabilistic Circuits (PCs) are a general framework for tractable deep generative models, which support exact and efficient probabilistic inference on their learned distributions. Recent modeling and training advancements have enabled their application to complex real-world tasks. However, the time and memory inefficiency of existing PC implementations hinders further scaling up. This paper proposes PyJuice, a general GPU implementation design for PCs that improves prior art in several regards. Specifically, PyJuice is 1-2 orders of magnitude faster than existing systems (including very recent ones) at training large-scale PCs. Moreover, PyJuice consumes 2-5x less GPU memory, which enables us to train larger models. At the core of our system is a compilation process that converts a PC into a compact representation amenable to efficient block-based parallelization, which significantly reduces IO and makes it possible to leverage Tensor Cores available in modern GPUs. Empirically, PyJuice can be used to improve state-of-the-art PCs trained on image (e.g., ImageNet32) and language (e.g., WikiText, CommonGen) datasets. We further establish a new set of baselines on natural image and language datasets by benchmarking existing PC structures but with much larger sizes and more training epochs, with the hope of incentivizing future research. Code is available at https://github.com/Tractables/pyjuice.