Federated Learning from Heterogeneous Data via Controlled Air Aggregation with Bayesian Estimation

Tomer Gafni, Kobi Cohen, Yonina C. Eldar

Research output: Contribution to journalArticlepeer-review

Abstract

Federated learning (FL) is an emerging machine learning paradigm for training models across multiple edge devices holding local data sets, without explicitly exchanging the data. Recently, over-the-air (OTA) FL has been suggested to reduce the bandwidth and energy consumption, by allowing the users to transmit their data simultaneously over a Multiple Access Channel (MAC). However, this approach results in channel noise directly affecting the optimization procedure, which may degrade the accuracy of the trained model. In this paper we jointly exploit the prior distribution of local weights and the channel distribution, and develop an OTA FL algorithm based on a Bayesian approach for signal aggregation. Our proposed technique, dubbed Bayesian Air Aggregation Federated learning (BAAF), is shown to effectively mitigate noise and fading effects induced by the channel. To handle statistical heterogeneity of users data, which is a second major challenge in FL, we extend BAAF to allow for appropriate local updates by the users and develop the Controlled Bayesian Air Aggregation Federated-learning (COBAAF) algorithm. In addition to using a Bayesian approach to average the channel output, COBAAF controls the drift in local updates using a judicious design of correction terms. We analyze the convergence of the learned global model using BAAF and COBAAF in noisy and heterogeneous environments, showing their ability to achieve a convergence rate similar to that achieved over error-free channels. Simulation results demonstrate the improved convergence of BAAF and COBAAF over existing methods in machine learning tasks.

Original languageEnglish
Pages (from-to)1928-1943
Number of pages16
JournalIEEE Transactions on Signal Processing
Volume72
DOIs
StatePublished - 1 Jan 2024

Keywords

  • Bayesian inference
  • Federated learning
  • stochastic gradient descent
  • wireless communication

ASJC Scopus subject areas

  • Signal Processing
  • Electrical and Electronic Engineering

Fingerprint

Dive into the research topics of 'Federated Learning from Heterogeneous Data via Controlled Air Aggregation with Bayesian Estimation'. Together they form a unique fingerprint.

Cite this