Scaling Probabilistic Circuits via Monarch Matrices

TL;DR

This paper introduces a novel sparse Monarch matrix parameterization for Probabilistic Circuits, enabling significant scalability improvements and state-of-the-art generative performance on large benchmarks with reduced computational costs.

Contribution

It presents a new sparse structured parameterization for PCs using Monarch matrices, combining theoretical insights with practical scalability benefits.

Findings

Achieves state-of-the-art results on Text8, LM1B, and ImageNet.

Reduces memory and computation costs significantly.

Demonstrates superior scaling with less FLOPs during training.

Abstract

Probabilistic Circuits (PCs) are tractable representations of probability distributions allowing for exact and efficient computation of likelihoods and marginals. Recent advancements have improved the scalability of PCs either by leveraging their sparse properties or through the use of tensorized operations for better hardware utilization. However, no existing method fully exploits both aspects simultaneously. In this paper, we propose a novel sparse and structured parameterization for the sum blocks in PCs. By replacing dense matrices with sparse Monarch matrices, we significantly reduce the memory and computation costs, enabling unprecedented scaling of PCs. From a theory perspective, our construction arises naturally from circuit multiplication; from a practical perspective, compared to previous efforts on scaling up tractable probabilistic models, our approach not only achieves state-of-the-art generative modeling performance on challenging benchmarks like Text8, LM1B and ImageNet, but also demonstrates superior scaling behavior, achieving the same performance with substantially less compute as measured by the number of floating-point operations (FLOPs) during training.