Our work “Scaling Probabilistic Circuits via Data Partitioning” has just been accepted at the 41st Conference on Uncertainty in Artificial Intelligence Conference (UAI). Congratulations to first-author Jonas Seng and all co-authors!
In the paper we provide a unified framework for federated learning across different settings of heterogeneous data. To this end, we introduce federated circuits to bridge horizontal (clients’ data distributions differ), vertical (clients hold different features), and hybrid (both of the former) federated learning. Federated circuits reframe the federated learning problem as aggregation of distributions learned by each client on its data partition, rather than aggregation of model parameters. In an instantiation using probabilistic circuits, we show that sum nodes help us address horizontal federated settings through mixture distributions, whereas vertical federated settings are addressed through mixtures of products of independent clusters. Check out the full paper for the precise framework, training algorithm, and application examples on images and tabular datasets (based on credit fraud and breast cancer detection)
For more information, read the full paper. The abstract is provided below:
Probabilistic circuits (PCs) enable us to learn joint distributions over a set of random variables and to perform various probabilistic queries in a tractable fashion. Though the tractability property allows PCs to scale beyond non-tractable models such as Bayesian Networks, scaling training and inference of PCs to larger, real-world datasets remains challenging. To remedy the situation, we show how PCs can be learned across multiple machines by recursively partitioning a distributed dataset, thereby unveiling a deep connection between PCs and federated learning (FL). This leads to federated circuits (FCs)—a novel and flexible federated learning (FL) framework that (1) allows one to scale PCs on distributed learning environments (2) train PCs faster and (3) unifies for the first time horizontal, vertical, and hybrid FL in one framework by re-framing FL as a density estimation problem over distributed datasets. We demonstrate FC’s capability to scale PCs on various large-scale datasets. Also, we show FC’s versatility in handling horizontal, vertical, and hybrid FL within a unified framework on multiple classification tasks.