Federated Learning with Uncertainty via Distilled Predictive Distributions

TL;DR

This paper introduces a federated learning framework that captures and distills predictive uncertainty into a single neural network, enabling accurate uncertainty estimation without extensive communication or computation, and improves performance in classification, active learning, and OOD detection.

Contribution

The proposed method efficiently encodes client predictive distributions into neural networks, avoiding the need to transmit full posterior distributions and simplifying uncertainty-aware predictions in federated learning.

Findings

Outperforms existing federated learning baselines in classification tasks.

Effectively detects out-of-distribution samples.

Enhances active learning by incorporating uncertainty estimates.

Abstract

Most existing federated learning methods are unable to estimate model/predictive uncertainty since the client models are trained using the standard loss function minimization approach which ignores such uncertainties. In many situations, however, especially in limited data settings, it is beneficial to take into account the uncertainty in the model parameters at each client as it leads to more accurate predictions and also because reliable estimates of uncertainty can be used for tasks, such as out-of-distribution (OOD) detection, and sequential decision-making tasks, such as active learning. We present a framework for federated learning with uncertainty where, in each round, each client infers the posterior distribution over its parameters as well as the posterior predictive distribution (PPD), distills the PPD into a single deep neural network, and sends this network to the server. Unlike some of the recent Bayesian approaches to federated learning, our approach does not require sending the whole posterior distribution of the parameters from each client to the server but only the PPD in the distilled form as a deep neural network. In addition, when making predictions at test time, it does not require computationally expensive Monte-Carlo averaging over the posterior distribution because our approach always maintains the PPD in the form of a single deep neural network. Moreover, our approach does not make any restrictive assumptions, such as the form of the clients' posterior distributions, or of their PPDs. We evaluate our approach on classification in federated setting, as well as active learning and OOD detection in federated settings, on which our approach outperforms various existing federated learning baselines.