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.

💡 Research Summary

This paper addresses a fundamental limitation of probabilistic circuits (PCs): the mixture weights at sum nodes are globally fixed and therefore cannot adapt to the local geometry of the data manifold. To endow PCs with geometry‑aware routing, the authors propose to replace constant sum‑node weights with a Voronoi tessellation (VT) over the input space. Each Voronoi cell is associated with a child subcircuit (expert), and an input is routed to the expert whose centroid is nearest. This deterministic, interpretable gating mechanism naturally captures spatial locality and can be viewed as a mixture‑of‑experts formulation embedded within a PC.

However, naïvely inserting VT into PCs destroys the two structural properties—smoothness and decomposability—that guarantee exact, linear‑time inference. Even when each expert factorizes completely across variables, the integration required for the partition function (\displaystyle Z=\sum_k \pi_k\int_{V_k}p_k(\mathbf{x})d\mathbf{x}) involves integrating over convex polyhedra with oblique facets. Such integrals do not factor into products of one‑dimensional integrals and are #P‑hard. Consequently, a single VT‑gated sum node already makes exact inference intractable, and the problem compounds in deep circuits where multiple VT gates intersect.

To reconcile geometry with tractable inference, the authors develop two complementary solutions.

1. Certified Approximate Inference
   They replace each Voronoi cell with an axis‑aligned hyper‑rectangle that tightly bounds the cell from above and below. Because rectangles are separable across dimensions, the required integrals decompose, allowing the propagation of lower and upper bounds through the circuit. The resulting algorithm yields provable intervals (