About the project
Objective
The aim is to establish a mathematical and computational framework for gradient-flow variational inference over families of probabilistic circuits.
Gradient flows in the Bures–Wasserstein and Wasserstein–Fisher–Rao (Hellinger–Kantorovich) geometries are well developed, with particularly explicit results for Gaussian distributions. However, it remains an open question how these frameworks extend to more expressive model classes such as probabilistic circuits, which generalise many tractable models while still supporting exact and efficient inference. In this project, we will investigate how Bures–Wasserstein and Wasserstein–Fisher–Rao gradient flows can be extended to probabilistic circuits. This will yield principled, geometry-aware algorithms for fitting expressive yet tractable variational approximations, with applications in robotics, AI for science, and computational epidemiology.
Background
Approximating intractable distributions, such as Bayesian posteriors, with an optimised tractable surrogate through variational inference is a cornerstone of probabilistic machine learning, with applications in robotics, AI for Science, healthcare, and computational epidemiology. Despite its popularity, modern variational inference faces a dilemma: mathematically rigorous frameworks exist, but they are limited to restrictive families of surrogates, such as Gaussians, which in turn limits approximation quality.
In parallel, probabilistic circuits have found increasing attention in the tractable probabilistic modelling community as a unifying framework. They offer a promising avenue to represent richer variational families, as they can represent complex, multimodal distributions while guaranteeing exact, polynomial-time density evaluation, marginalisation and moment computations. However, the geometry of families of probabilistic circuits within Wasserstein space is largely unexplored, and how this geometry interacts with structural conditions which enable tractable computations is not yet understood.
This project unites these two complementary research directions to develop mathematically rigorous and computationally effective gradient-flow variational inference techniques over families of probabilistic circuits.
Crossdisciplinary collaboration
The project established a collaboration between Trapp’s group in the Computational and Learning Systems department and Zhu’s group in the Mathematics department at KTH Royal Institute of Technology and brings together expertise on tractable probabilistic modelling with probabilistic circuits (Trapp) and mathematical foundations on gradient-flow variational inference (Zhu).
