Distribution-Level AirComp for Wireless Federated Learning under Data Scarcity and Heterogeneity

TL;DR

This paper introduces a Bayesian federated learning framework for wireless edge networks that effectively handles data scarcity and heterogeneity by using distribution-level aggregation and a tailored AirComp scheme, improving training stability and accuracy.

Contribution

The paper proposes a novel Bayesian FL framework with distribution-level aggregation and a specialized AirComp scheme, addressing communication and heterogeneity challenges in wireless FL.

Findings

Significant improvement in test accuracy over traditional FL.

Enhanced calibration performance in data-scarce environments.

Robust convergence guarantees under wireless impairments.

Abstract

The conventional FL methods face critical challenges in realistic wireless edge networks, where training data is both limited and heterogeneous, often leading to unstable training and poor generalization. To address these challenges in a principled manner, we propose a novel wireless FL framework grounded in Bayesian inference. By virtue of the Bayesian approach, our framework captures model uncertainty by maintaining distributions over local weights and performs distribution-level aggregation of local distributions into a global distribution. This mitigates local overfitting and client drift, thereby enabling more reliable inference. Nevertheless, adopting Bayesian FL increases communication overhead due to the need to transmit richer model information and fundamentally alters the aggregation process beyond simple averaging. As a result, conventional Over-the-Air Computation (AirComp), widely used to improve communication efficiency in standard FL, is no longer directly applicable. To overcome this limitation, we design a dedicated AirComp scheme tailored to Bayesian FL, which efficiently aggregates local posterior distributions at the distribution level by exploiting the superposition property of wireless channels. In addition, we derive an optimal transmit power control strategy, grounded in rigorous convergence analysis, to accelerate training under power constraints. Our analysis explicitly accounts for practical wireless impairments such as fading and noise, and provides theoretical guarantees for convergence. Extensive simulations validate the proposed framework, demonstrating significant improvements in test accuracy and calibration performance over conventional FL methods, particularly in data-scarce and heterogeneous environments.