Federated Learning from Heterogeneous Data via Controlled Bayesian Air Aggregation

Tomer Gafni,Kobi Cohen,Yonina C. Eldar

2023-03-30

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 algorithm, 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 environment, 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 algorithms in machine learning tasks.

Signal Processing

What problem does this paper attempt to address?

### Problems Addressed by the Paper

The paper primarily aims to address two core challenges in Federated Learning (FL):

1. **Model Training Accuracy Degradation Due to Channel Noise**:
- In Federated Learning, devices transmit data through wireless channels, which are susceptible to channel noise, thereby affecting the accuracy of model training. The paper proposes a Bayesian-based signal aggregation algorithm—Bayesian Air Aggregation Federated Learning (BAAF), which effectively mitigates the noise and fading effects introduced by the channel.

2. **Statistical Heterogeneity of User Datasets**:
- Another major challenge in Federated Learning is the statistical heterogeneity among user datasets. To address this issue, the paper extends the BAAF algorithm and proposes the Controlled Bayesian Air Aggregation Federated Learning (COBAAF) algorithm. COBAAF controls user drift through appropriate local update strategies and reduces inter-user bias using control variate techniques.

### Algorithm Contributions

- **BAAF Algorithm**: By jointly utilizing the prior distribution of local weights and the channel distribution, a new OTA FL algorithm is designed, which can effectively reduce the impact of channel noise.
- **COBAAF Algorithm**: Building on BAAF, it further addresses the issue of statistical heterogeneity by using control variate techniques to reduce user drift and improve model convergence speed.

### Conclusion

The BAAF and COBAAF algorithms proposed in this paper demonstrate superior performance in both theoretical analysis and numerical simulations, achieving convergence rates similar to those under ideal channel conditions in noisy and heterogeneous environments.