Geometry-Aware Probabilistic Circuits via Voronoi Tessellations

TL;DR

This paper introduces a geometry-aware approach to probabilistic circuits using Voronoi tessellations, enhancing their ability to model data manifolds while maintaining tractability through novel inference methods and relaxations.

Contribution

It proposes a new method integrating Voronoi tessellations into probabilistic circuits, with solutions for maintaining tractability and enabling gradient-based learning.

Findings

Provides guaranteed bounds for approximate inference.

Identifies structural conditions for exact inference.

Empirically improves density estimation performance.

Abstract

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\"ively 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.