Balancing Energy Efficiency and Distributional Robustness in Over-the-Air Federated Learning

The increasing number of edge devices generates massive amounts of private data, creating storage, capacity, and processing challenges. Federated learning (FL) addresses these by enabling collaborative learning among devices without transferring raw data to a central server. However, FL faces challenges such as data heterogeneity across devices, leading to biases and fairness issues. Robust min-max formulations have been introduced to address this heterogeneity, ensuring model effectiveness across various data distributions. These methods often select a subset of clients in each communication round for efficiency. AirComp leverages the superposition property of wireless channels for efficient model aggregation, offering scalability and latency advantages. Energy consumption is a critical concern, especially for battery-powered devices. Existing dynamic client scheduling methods, while improving energy efficiency, may rely on several tuning parameters and lack predictability. This paper proposes a distributed learning approach balancing energy efficiency and distributional robustness, focusing on communication cost reduction. The key contribution is a novel algorithm, Channel-Aware Agnostic Federated Learning (CA-AFL), which is the first work to jointly address energy efficiency and robustness to data heterogeneity in terms of communication costs. Theoretically, the algorithm is shown to default to either a distributionally robust FL algorithm (ignoring channel conditions) or a fully energy-conservative client selection algorithm at the extremes of its tuning parameter. Intermediate values strike a balance between both aspects. Simulation results demonstrate significant energy savings (up to one-third) compared to conventional AFL, with negligible performance reduction and outperforming an existing energy-efficient algorithm (GCA) in terms of worst client accuracy and global standard deviation (STD).

The paper references several key works. [1] introduces the concept of communication-efficient learning in FL. [2] presents an agnostic federated learning approach using a robust min-max formulation to handle data heterogeneity. Subsequent works [3]-[5] extend this formulation to centralized and decentralized settings. The challenges of scalability and latency in wireless edge-centric systems, and AirComp as a solution, are discussed, citing works [7]-[9]. Existing works on dynamic client scheduling, such as [10], are reviewed, highlighting their limitations concerning tuning parameters and unpredictability. The current work builds on these foundations by addressing the joint problem of energy efficiency and robustness in AirComp-based FL.

The paper proposes a system model with a parameter server (PS) coordinating N edge devices. Unlike orthogonal digital transmission, AirComp uses analog transmission, leveraging the superposition property of the wireless channel for model aggregation. To reduce energy, only K devices are selected per round based on a probability distribution p( ) combining energy and distributional robustness metrics. The model aggregation process using AirComp over a given sub-carrier is mathematically defined. The distributionally robust optimization problem is formulated as a min-max problem, aiming for robustness to data heterogeneity. The energy consumption for each client is modeled based on transmission power and time, considering channel inversion techniques. The effective channel is defined based on the total channel inversion power. The core contribution lies in the novel client selection mechanism. It utilizes a bias-configurable probability mass function (PMF) controlled by a tuning parameter C. This PMF incorporates both the energy consumption (based on channel conditions) and the probability distribution from the robust optimization problem. Proposition 1 proves the expression for the bias-configurable PMF, which becomes unbiased when C=0 and fully biased towards energy efficiency as C approaches infinity. The final client selection probability is derived as the product of the energy-aware PMF and the probability distribution from the robust optimization problem, then normalized. This is similar in concept to the product of experts (PoE) in machine learning. Proposition 2 shows that as C approaches infinity, the algorithm defaults to a greedy energy-efficient client selection strategy. The proposed CA-AFL algorithm is detailed, outlining the descent step (client model updates and upload) and the ascent step (updating the probability simplex vector). Algorithm 1 summarizes the steps.

The simulation results show that CA-AFL performs comparably to existing algorithms in terms of average global accuracy, reaching approximately 80% convergence across all methods. However, CA-AFL significantly outperforms baselines in terms of worst-client accuracy and standard deviation (STD) of global accuracy, demonstrating improved robustness and fairness. Figures 2(a), 2(b), and 2(c) illustrate these performance metrics against the number of communication rounds. CA-AFL, especially with C=2 and C=8, shows faster convergence in worst-client accuracy and STD compared to FedAvg and GCA. Furthermore, Figure 3 presents the performance metrics against total energy consumption. CA-AFL demonstrates substantial energy savings compared to FedAvg and AFL, achieving comparable accuracy with significantly less energy. Specifically, with C=8, CA-AFL matches the energy efficiency of GCA but achieves superior worst-client accuracy. The results demonstrate that CA-AFL achieves the best attainable worst client accuracy and STD of GCA and FedAvg with significantly less energy consumption. The paper highlights a key finding: CA-AFL achieves comparable performance to the AFL algorithm while using only about one-third of the energy.

The results validate the effectiveness of the proposed CA-AFL algorithm in balancing energy efficiency and distributional robustness. The ability to control the balance via the tuning parameter C offers flexibility to adjust the algorithm's behavior based on the specific requirements of the application. The superior performance compared to baselines, especially in terms of worst-client accuracy and STD, indicates improved model fairness and reliability. The substantial energy savings are significant for resource-constrained edge devices. The comparison with GCA shows that CA-AFL can outperform a gradient and channel-aware algorithm while maintaining superior energy efficiency, illustrating the advantage of the proposed robust client selection strategy. The findings suggest that incorporating distributional robustness in the context of energy-aware client selection leads to a more balanced and efficient FL system.

The paper successfully introduces CA-AFL, a novel algorithm for energy-efficient and distributionally robust FL with AirComp. Simulation results show substantial energy savings and improved robustness compared to existing methods. Future work could explore adaptive tuning of the parameter C based on real-time system conditions and investigate the algorithm's performance under various network conditions and different datasets.

The simulation environment might not perfectly represent real-world scenarios. The impact of different channel models and more complex network topologies should be further investigated. The current study uses a specific dataset and model; generalizability to other applications and model types could be further assessed.