Federated Learning with Neural Graphical Models

TL;DR

This paper introduces FedNGMs, a federated learning framework for Neural Graphical Models that preserves data privacy, maintains constant model size, and effectively handles data heterogeneity and client-specific variables.

Contribution

The paper proposes FedNGMs, a novel federated learning approach for NGMs that avoids parameter explosion and personalizes models with a stitching algorithm.

Findings

FedNGMs effectively learn from distributed data without sharing raw data.

The framework maintains a constant global model size during training.

FedNGMs demonstrate robustness to data heterogeneity and communication constraints.

Abstract

Federated Learning (FL) addresses the need to create models based on proprietary data in such a way that multiple clients retain exclusive control over their data, while all benefit from improved model accuracy due to pooled resources. Recently proposed Neural Graphical Models (NGMs) are Probabilistic Graphical models that utilize the expressive power of neural networks to learn complex non-linear dependencies between the input features. They learn to capture the underlying data distribution and have efficient algorithms for inference and sampling. We develop a FL framework which maintains a global NGM model that learns the averaged information from the local NGM models while keeping the training data within the client's environment. Our design, FedNGMs, avoids the pitfalls and shortcomings of neuron matching frameworks like Federated Matched Averaging that suffers from model parameter explosion. Our global model size remains constant throughout the process. In the cases where clients have local variables that are not part of the combined global distribution, we propose a `Stitching' algorithm, which personalizes the global NGM models by merging the additional variables using the client's data. FedNGM is robust to data heterogeneity, large number of participants, and limited communication bandwidth. We experimentally demonstrated the use of FedNGMs for extracting insights from CDC's Infant Mortality dataset and discuss interesting future applications.