Geometry-Aware Probabilistic Circuits via Voronoi Tessellations

Probabilistic circuits (PCs) enable exact and tractable inference but employ data independent mixture weights that limit their ability to capture local geometry of the data manifold. We propose Voronoi tessellations (VT) as a natural way to incorporate geometric structure directly into the sum nodes of a PC. However, naïvely introducing such structure breaks tractability. We formalize this incompatibility and develop two complementary solutions: (1) an approximate inference framework that provides guaranteed lower and upper bounds for inference, and (2) a structural condition for VT under which exact tractable inference is recovered. Finally, we introduce a differentiable relaxation for VT that enables gradient-based learning and empirically validate the resulting approach on standard density estimation tasks.

Most AI systems today are good at making predictions, but it is often difficult to know how confident they truly are or to ask them reliable follow-up questions about uncertainty. Probabilistic circuits are a family of models designed to answer such questions efficiently, but they usually use fixed routing rules that do not adapt well to the local shape of complex data. In this work, we make probabilistic circuits more geometry-aware. The key idea is to divide the data space into regions, similar to drawing boundaries around nearby groups of points, so that different parts of the model can specialize in different regions. We use Voronoi tessellations for this purpose. However, adding these geometric regions creates a new challenge: the model may lose the efficient reasoning guarantees that make probabilistic circuits useful. We show exactly why this difficulty occurs and propose two ways to address it. One method keeps the full geometric flexibility while producing guaranteed lower and upper bounds on answers. The other restricts the geometry so that exact and efficient reasoning is preserved. Our experiments show that this approach can improve density modeling while keeping inference reliable, making it a step toward more interpretable and trustworthy probabilistic AI systems.