Optimal Transport for Probabilistic Circuits

Abstract

We introduce 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. We propose a Wasserstein-type distance that restricts the coupling measure 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 study the empirical Wasserstein distance between a PC and a dataset, and show that we can estimate the PC parameters to minimize this distance through an efficient iterative algorithm.

Citation

@inproceedings{CiotingaUAI25,
  author    = {Ciotinga, Adrian and Choi, YooJung},
  title     = {Optimal Transport for Probabilistic Circuits},
  booktitle = {In the 41st Conference on Uncertainty in Artificial Intelligence (UAI)},
  month     = {July},
  year      = {2025},
}

Preliminary version appeared in the International OPT Workshop on Optimization for Machine Learning at NeurIPS 2024.