Towards Lightweight, Communication-Efficient, and Energy-Aware Federated Learning under Non-IID Settings

Federated learning (FL) has emerged as a promising privacy-preserving methodology for cooperatively training machine learning models across numerous machines without transferring sensitive data between them. However, the substantial communication overhead and diverse client contributions hinder its practical implementation. Recent approaches leverage reinforcement learning (RL) for adaptive client selection, but such methods often suffer from high complexity in real-world FL environments. This paper presents a greedy token-based client selection (GreedyTBCS) strategy that can surpass reinforcement learning-based methods in stable non-Independent and Identically Distributed (non-IID) settings. The priority is given to clients whose updates show high energy concentration in the dominant frequency components, which means that the learning signals are structured and aligned globally. We also introduced a contribution-aware Aggregation technique (CAAgg) to achieve early convergence, particularly in non-IID settings. Communication efficiency is further enhanced by introducing a novel energy-aware mechanism for adaptively controlling the number of local training epochs named Token Adaptive Local Training (TALT). Experimental results demonstrate that the proposed framework including GreedyTBCS, CAAgg and TALT can minimize the number of communication cycles by 59.7% on MNIST, 50.6% on CIFAR-10 and 10% on CIFAR-100 compared to FedAvg with random client selection. Moreover, the framework always achieves equivalent accuracy on MNIST and CIFAR-10 datasets and 3.7% increase in CIFAR-100 dataset in non-IID contexts, which shows a good trade-off between communication efficiency and model performance. The proposed framework can be combined with DCT-based structured sparsification to greatly minimize the communication payload while maintaining the model accuracy. Extensive experiments show that the proposed framework scales to high numbers of clients, and is robust to moderate heterogeneous networks, and can increase energy efficiency by reducing communication and adaptively adjusting local computation.