[Paper Review] Bayesian Federated Learning over Wireless Networks

Seunghoon Lee, Chanho Park+3 more|arXiv (Cornell University)|Dec 31, 2020

Privacy-Preserving Technologies in Data34 references16 citations

TL;DR

This paper proposes Bayesian Federated Learning (BFL) for wireless networks, where a server optimally aggregates one-bit quantized gradients by leveraging prior distributions, quantizer functions, and channel state information to minimize mean-squared error (MSE). It further introduces scalable-BFL (SBFL), a low-complexity variant that transmits only two scalar parameters per device, achieving superior convergence over signSGD on MNIST under heterogeneous wireless conditions.

ABSTRACT

Federated learning is a privacy-preserving and distributed training method using heterogeneous data sets stored at local devices. Federated learning over wireless networks requires aggregating locally computed gradients at a server where the mobile devices send statistically distinct gradient information over heterogenous communication links. This paper proposes a Bayesian federated learning (BFL) algorithm to aggregate the heterogeneous quantized gradient information optimally in the sense of minimizing the mean-squared error (MSE). The idea of BFL is to aggregate the one-bit quantized local gradients at the server by jointly exploiting i) the prior distributions of the local gradients, ii) the gradient quantizer function, and iii) channel distributions. Implementing BFL requires high communication and computational costs as the number of mobile devices increases. To address this challenge, we also present an efficient modified BFL algorithm called scalable-BFL (SBFL). In SBFL, we assume a simplified distribution on the local gradient. Each mobile device sends its one-bit quantized local gradient together with two scalar parameters representing this distribution. The server then aggregates the noisy and faded quantized gradients to minimize the MSE. We provide a convergence analysis of SBFL for a class of non-convex loss functions. Our analysis elucidates how the parameters of communication channels and the gradient priors affect convergence. From simulations, we demonstrate that SBFL considerably outperforms the conventional sign stochastic gradient descent algorithm when training and testing neural networks using MNIST data sets over heterogeneous wireless networks.

Motivation & Objective

• To address the challenge of optimal gradient aggregation in federated learning over heterogeneous wireless networks with one-bit quantization and fading channels.
• To jointly model local gradient priors, quantization effects, and channel fading for improved estimation accuracy.
• To reduce computational and communication overhead in large-scale federated learning by proposing a scalable variant (SBFL).
• To analyze convergence of SBFL for non-convex loss functions under realistic wireless constraints.
• To demonstrate superior performance over signSGD in terms of test accuracy and convergence speed on MNIST.

Proposed method

• Proposes Bayesian Federated Learning (BFL) that uses Bayesian linear minimum mean-square error (BLMMSE) estimation to optimally reconstruct the global gradient from noisy, faded, one-bit quantized gradients.
• Incorporates three key components: prior distributions of local gradients, the one-bit quantizer function, and channel state distributions to minimize estimation MSE.
• Introduces scalable-BFL (SBFL), which assumes a simplified Gaussian prior per device and transmits only two scalar parameters (mean and variance) instead of full gradients.
• Employs BLMMSE estimation at the server using the received signals and the known statistical parameters to reconstruct the global gradient.
• Uses adaptive learning rate $\gamma^t = \frac{1}{t+1}$ in the convergence analysis to ensure convergence for non-convex objectives.
• Derives theoretical bounds on the MSE of the reconstructed gradient and connects it to the expected improvement per iteration in the optimization process.

Experimental results

Research questions

• RQ1 How can one-bit quantized gradients be optimally aggregated in a wireless federated learning system with heterogeneous devices and fading channels?
• RQ2 What is the impact of gradient priors and channel quality on the performance of gradient reconstruction in federated learning?
• RQ3 Can a low-complexity variant of Bayesian gradient aggregation be designed that maintains high performance with minimal communication overhead?
• RQ4 How does the proposed SBFL algorithm compare to signSGD in terms of convergence speed and final test accuracy on non-convex models?
• RQ5 What theoretical guarantees can be provided for the convergence of the proposed algorithm under non-convex loss functions?

Key findings

• SBFL achieves significantly better test accuracy than signSGD when training a neural network on the MNIST dataset over heterogeneous wireless networks.
• The convergence of SBFL is theoretically guaranteed for non-convex loss functions under appropriate learning rate scheduling.
• The MSE of the reconstructed gradient in BFL is minimized by jointly exploiting gradient priors, quantizer models, and channel state distributions.
• SBFL reduces communication cost by transmitting only two scalar parameters per device, while maintaining performance close to full-precision aggregation.
• The convergence rate of SBFL depends on the interplay between channel quality, gradient variance, and the learning rate, with better channel conditions improving convergence.
• Theoretical analysis shows that the expected improvement per iteration is bounded by the norm of the true gradient and the estimation error variance, with the latter controlled by channel and prior parameters.