Optimal Transport for Probabilistic Circuits

This work introduces a novel optimal transport framework for probabilistic circuits (PCs). While it has been shown recently that divergences between distributions represented as certain classes of PCs can be computed tractably, to the best of our knowledge, there is no existing approach to compute the Wasserstein distance between probability distributions given by PCs. In this work, we propose a Wasserstein-type distance that restricts the coupling mea- sure of the associated optimal transport problem to be a probabilistic circuit. We then develop an algorithm for computing this distance by solving a series of small linear programs and derive the circuit conditions under which this is tractable. Furthermore, we show that we can easily retrieve the optimal transport plan between the PCs from the solutions to these linear programs. Lastly, we explore approaches to parameter learning for PCs that minimize the empirical Wasserstein distance between a PC and a dataset, and provide two approaches that minimize this distance.