Optimal Transport for Probabilistic Circuits

2024-10-16

arXiv.org

AI Summary

Existing probabilistic circuits (PCs) lack a computationally tractable definition of the Wasserstein distance, hindering their use in metric learning and distributional optimization. Method: We propose the first optimal transport (OT) framework for PCs, introducing a PC-specific Wasserstein distance where the OT coupling is explicitly encoded as a structured PC. We identify structural circuit conditions ensuring polynomial-time solvability and support end-to-end parameter optimization to minimize empirical Wasserstein distance. Our approach integrates OT theory, linear programming, and PC structural analysis. Results: The method enables both efficient distance computation and direct analytical recovery of the optimal transport plan. Empirical evaluation on real-world data demonstrates significant improvements in distributional fidelity. Crucially, our framework unifies computational tractability, interpretability (via explicit coupling representation), and differentiability for gradient-based learning—thereby filling a fundamental gap in metric learning for probabilistic circuits.

Technology Category

• Machine Learning: Probabilistic Circuits and Graphical Models
• Search and Optimization: Learning to Search
• Reasoning under Uncertainty: Probabilistic Programming

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Web performance, measurement, and characterization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Introduces optimal transport framework for probabilistic circuits.
• Proposes Wasserstein-type distance for distributions represented by PCs.
• Develops algorithm to compute distance via linear programs.

Innovation

• Introduces Wasserstein-type distance for probabilistic circuits.
• Develops algorithm using small linear programs.
• Estimates PC parameters via efficient iterative algorithm.