Scaling Up Probabilistic Circuits by Latent Variable Distillation
Anji Liu, Honghua Zhang, Guy Van den Broeck
TL;DR
The paper tackles the plateau observed when scaling Probabilistic Circuits (PCs) to large, high-dimensional datasets by introducing latent variable distillation (LVD), which injects semantics-aware supervision from deep generative models into the LV space of PCs. The method materializes and assigns latent variables via clustering neural embeddings, optimizes a tractable lower bound on the joint likelihood, and then fine-tunes on the original objective, enabling large PCs to fully exploit their capacity. The approach yields substantial performance gains on image and language benchmarks, and demonstrates competitive results against flow-based models and VAEs, while significantly reducing training time. Overall, LVD provides a principled, scalable avenue for combining tractable probabilistic modeling with neural representations for high-dimensional generative tasks.
Abstract
Probabilistic Circuits (PCs) are a unified framework for tractable probabilistic models that support efficient computation of various probabilistic queries (e.g., marginal probabilities). One key challenge is to scale PCs to model large and high-dimensional real-world datasets: we observe that as the number of parameters in PCs increases, their performance immediately plateaus. This phenomenon suggests that the existing optimizers fail to exploit the full expressive power of large PCs. We propose to overcome such bottleneck by latent variable distillation: we leverage the less tractable but more expressive deep generative models to provide extra supervision over the latent variables of PCs. Specifically, we extract information from Transformer-based generative models to assign values to latent variables of PCs, providing guidance to PC optimizers. Experiments on both image and language modeling benchmarks (e.g., ImageNet and WikiText-2) show that latent variable distillation substantially boosts the performance of large PCs compared to their counterparts without latent variable distillation. In particular, on the image modeling benchmarks, PCs achieve competitive performance against some of the widely-used deep generative models, including variational autoencoders and flow-based models, opening up new avenues for tractable generative modeling. Our code can be found at https://github.com/UCLA-StarAI/LVD.