FedLBW: A Loss-Based Weighting Strategy for Federated Learning on Non-IID Data in Wireless Networks

2026-08-07

AI Summary

This work addresses the challenges of slow convergence and degraded accuracy in federated learning over wireless networks, which stem from non-independent and identically distributed (non-IID) client data and frequent client dropouts. To mitigate these issues, the authors propose FedLBW, a loss-based weighted aggregation method that leverages a small proxy dataset on the server to compute validation losses for each client’s model update. Instead of the conventional data-size-based weighting, FedLBW dynamically assigns weights inversely proportional to these validation losses, thereby prioritizing more reliable updates during aggregation. Experimental results on FashionMNIST, CIFAR-10, and CIFAR-100 demonstrate that FedLBW significantly enhances robustness: under extreme non-IID settings, it improves CIFAR-10 accuracy by up to 7.6%, and consistently outperforms baseline methods such as FedAvg even under high client dropout rates.

Abstract

Federated Learning (FL) enables collaborative machine learning (ML) across distributed clients while preserving privacy. However, efficient model convergence in FL remains challenging, especially in wireless networks where non-independent and identically distributed (non-IID) data and frequent client dropouts are common. Traditional FL algorithms, such as FedAvg, rely solely on dataset size to weight client updates. This introduces biases towards clients with larger datasets and makes the process sensitive to non-IID data, outliers, and client dropouts. To address these challenges, we propose Federated Learning with Loss-Based Weighting (FedLBW), a novel aggregation method that assigns each client's update a weight proportional to the inverse of its validation loss, computed using a small proxy dataset on the server, rather than its dataset size. This ensures that lower-loss models exert greater influence during aggregation, prioritizing the most reliable updates and boosting overall performance. Through extensive experiments across multiple datasets, including FashionMNIST (CNN), CIFAR-10 (ResNet-18), and CIFAR-100 (ResNet-34), we demonstrate that FedLBW achieves higher accuracy and faster convergence compared to baseline algorithms such as FedAvg, FedAvgM, FedProx, FedNova, FedLAW and FedDkw, with notable improvements of up to 7.6 % higher accuracy on CIFAR-10 in extreme non-IID cases. Moreover, FedLBW showcases exceptional resilience to increasing dropout probabilities, consistently maintaining significantly higher accuracy even in challenging conditions. These results establish FedLBW as an effective and resilient solution for FL in wireless network environments, offering marked improvements in model accuracy, convergence speed, and robustness to non-IID data and client dropouts.

Problem

• Federated Learning
• non-IID data
• client dropout
• model convergence
• wireless networks

Innovation

• loss-based weighting
• federated learning
• non-IID data
• client dropout resilience
• model aggregation