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

arXiv:2608.07007 · cs.AI, cs.DC, cs.ET, cs.LG · Submitted 2026-08-07 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: Paper Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "FedLBW: A Loss-Based Weighting Strategy for Federated Learning on Non-IID Data in Wireless Networks".

Jane: The paper was written by Majid Kundroo, Tinku Singh and Taehong Kim from Chungbuk National University and Bennett University.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title and Authors: Tom: We've got a fresh paper today that's been making rounds, and it's called FedLBW. Big one in the federated learning space.

Jane: I know it. Kundroo, Singh, and Kim from Chungbuk National University and Bennett University. They tackle a problem we always feel in wireless networks used for federated learning: clients not agreeing on data.

Tom: Right. The classic scenario is, you have 100 phones, each with completely different photos. They all train a shared model, but they are, well, all over the place.

Jane: Exactly. Most federated learning algorithms treat every client as if their data is a neat, random sample. But wireless networks are messy. Some clients have tiny datasets, some have huge ones.

Tom: So the server just averages everything like a simple classroom average. That unfairly boosts the person with the thickest textbook.

Jane: Yes! And FedLBW instead says, "Let's grade the homework before mixing the grades." The server uses a small validation set to figure out which local models are actually good.

Tom: It's like when a coach watches players warm up and picks the ones who are actually landing shots, rather than just listening to who shouts the loudest.

Jane: The wild part is they show this simple idea can beat a bunch of complex setups. They improved accuracy by up to 7.6 percent on CIFAR-10 in the challenging non-IID case.

Tom: That's a pretty bold claim. Let's dig into how this actually works and whether it holds up when devices start dropping off the network.

Summary of the Paper: Tom: So we just set the stage for FedLBW. Now let's talk about the actual mechanism, because it's simpler than you'd expect.

Jane: Totally. Imagine you're the server. You hold this tiny, balanced sample of data, like a mini snapshot of the whole world. Let's call it the proxy dataset.

Tom: Got it. Each round, clients train on their own private data, then send back just their model, not the data itself.

Jane: Then the server takes each incoming model and checks it against the proxy dataset. It computes a loss value, which is basically a sanity check on how well that model generalizes.

Tom: So a client with great local data will have a low loss, and a client who overfit to two weird pictures will have a high loss.

Jane: Precisely. Then instead of weighting by dataset size, they weight each client's update by the inverse of that loss. Low loss means high influence. This shifts the whole aggregation strategy.

Tom: The key phrase in the paper is "inverse of its validation loss." It directly rewards better-performing models rather than larger datasets.

Jane: That's a much more sensible objective if your goal is accuracy and robust convergence. But does it introduce overhead? The server has to evaluate every client.

Tom: The paper acknowledges this and calculates the cost as the validation set size times the number of participants. And realistically, that is small.

Jane: They frame this as okay because servers usually have spare compute. The proxy dataset itself is small, like 100 samples per class.

Tom: That is the core idea, but what I find really interesting is that it keeps the local data completely private. The proxy data is only used for grading, not for training.

Jane: Which addresses a few privacy concerns while still letting the server do intelligent aggregation. The loss is only a scalar, a single number.

Tom: And that little number might be worth a lot. Let's see if those gains hold up when the network gets worse.

Improvements Suggested: Tom: So we know the idea behind FedLBW. But what's the actual payoff? You mentioned 7.6 percent accuracy earlier, but let's break it down.

Jane: It gets better as the data gets messier. At the extreme non-IID level, with Dirichlet alpha equal to 0.1, FedLBW really shines on CIFAR-10.

Tom: That's the infamous alpha 0.1. In case anyone wonders, that means every client has data from very few classes. Almost no overlap.

Jane: And that's where FedLBW beats FedAvg by 7.66 percent. It's a pretty significant jump. But the improvements don't stop at the final numbers.

Tom: The paper also emphasizes convergence speed. You can see in their figures that FedLBW gets to a higher accuracy much earlier in training.

Jane: In the beginning of training, standard FedAvg can oscillate. The gradient is noisy because some clients are sending terrible models.

Tom: FedLBW smooths that noise out because it effectively says, "You, the good model, lead the way." Others follow.

Jane: The paper also has a convergence analysis, not just intuition. They derive bounds that break down the error from sampling, from noise, from local drift.

Tom: And they show the loss-based weighting reduces something they call the weighting bias. It's a more efficient descent direction.

Jane: So you're not just getting lucky with the test set. There is a mathematical justification for why the weighting helps.

Tom: But the really impressive part is the dropout resilience. Let's get to how it behaves when clients just disappear.

First Page of the Paper: Tom: We've covered the mechanism and the accuracy gains. Now I want to revisit the abstract and intro because they set the bar high.

Jane: Right. The paper kicks off by listing the flaws of the classic FedAvg: it's biased towards big datasets and sensitive to non-IID outliers.

Tom: And the authors identify something that's super relevant to wireless: client dropouts. The paper calls them "inevitable" in wireless networks.

Jane: The motivating section is actually one of the clearest parts. It says the traditional method gives no incentive for clients to train better.

Tom: That's the "economic" angle I really like. In FedLBW, if you want more influence on the global model, you need to train a better local model.

Jane: The first page also mentions the proxy dataset debate. Some people in FL say it violates the spirit of pure federated learning.

Tom: But the paper defends this by citing recent work like FedLAW and SA-FL, which show these small proxy sets are already common.

Jane: So they're shifting the discussion from "whether you can use a proxy" to "how to use it safely and efficiently."

Tom: And on the practical side, they test on FashionMNIST, CIFAR-10, and CIFAR-100. They use different models for each, which is a nice, robust setup.

Jane: For FashionMNIST it's a simple CNN, for CIFAR-10 it's ResNet-18, and for CIFAR-100 it's a heavier ResNet-34.

Tom: Everything is over 300 rounds too. That's a long training run, so we're not seeing some overnight miracle.

Jane: I also like that they test how sensitive FedLBW is to the proxy data distribution. Like, what if the proxy data is shifted?

Tom: That's a huge question. They have a whole experiment with SVHN as a proxy, which is a massive domain shift, and accuracy only drops from 72 percent to 70 percent.

Jane: Even in that extreme scenario, FedLBW still outperforms FedAvg. That suggests the grading is quite robust.

Tom: The first page tees up all of that beautifully. The claim isn't just "here's a new weighting scheme"; it's "here's a new reliability estimator for clients."

Jane: And that reliability comes at almost no extra communication cost, which is a big deal for wireless. The server just sends the same-sized model down.

Tom: Let's wrap up by tying all of this together and looking at where the field goes next.

Conclusion: Tom: We've had a good look at this paper. Before we wrap up, let's recap the biggest takeaways.

Jane: FedLBW is a surprisingly simple fix: weight client updates by the inverse of their validation loss on a proxy dataset, not by the number of samples.

Tom: And that small change pays off massively in non-IID settings and when clients drop out.

Jane: The magic is that it automatically trusts models that generalize well. In wireless networks, where connections are unstable, that's a lifesaver.

Tom: The numbers speak for themselves: up to 7.66 percent improvement over FedAvg on CIFAR-10 at extreme non-IID, and it holds up at 50 percent dropout where FedAvg crashes.

Jane: FedAvg drops to 39.81 percent accuracy at 50 percent dropout, while FedLBW holds at 64.40 percent. That's a huge cliff and FedLBW just doesn't fall off it.

Tom: It also gives clients an incentive to improve local training since their contribution is tied to performance. It aligns the goals.

Jane: The convergence analysis gives a solid theoretical foundation. It's not just tuning the weights and hoping for the best.

Tom: And the proxy robustness tests show that even if the proxy data is imperfect, you still get gains.

Jane: It's a really practical approach. The authors didn't try to make a complex optimization problem; they just found a sensible criterion.

Tom: I'd love to see this tested on more diverse tasks, like text or reinforcement learning, but for image classification in wireless, it's a strong claim.

Jane: We'll be watching for that. As a wrap-up, FedLBW is a nice reminder that the simple idea is sometimes the best one.

Tom: Alright folks, we've got more papers to get to, but that's a good look at this one. Stay tuned.

Jane: Thanks for listening in.

Majid Kundroo, Tinku Singh, Taehong Kim

Chungbuk National University · Bennett University

cs.AI, cs.DC, cs.ET, cs.LG

Submitted: 2026-08-07

Updated: 2026-08-10

Journal ref: Kundroo, Majid, Tinku Singh, and Taehong Kim. "FedLBW: A loss-based weighting strategy for federated learning on non-IID data in wireless networks." Expert Systems with Applications (2025): 130487

DOI: 10.1016/j.eswa.2025.130487

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 60/100

The gist: The paper introduces FedLBW (Federated Learning with Loss-Based Weighting), a novel aggregation strategy for federated learning (FL) designed to address challenges in wireless networks where non-IID

Key concepts

Federated Learning
A machine learning approach where multiple clients (e.g., phones) train a shared model using their own private data, without sending the data to a central server. The server aggregates model updates from clients to improve the global model.
Non-IID Data
Data that is not identically and independently distributed across clients. In wireless networks, clients often have very different data distributions (e.g., one client has only cat photos, another only dogs), which can hurt standard federated learning performance.
Proxy Dataset
A small, balanced dataset held by the server, used only to evaluate incoming client models. The server computes each model's loss on this proxy to gauge how well it generalizes, without using it for training or accessing client data.
Loss-Based Weighting
Instead of weighting client updates by dataset size, FedLBW weights them by the inverse of the validation loss on the proxy dataset. Lower loss means higher influence, rewarding models that generalize better and improving aggregation in non-IID settings.

Terminology

Summary

The paper introduces FedLBW (Federated Learning with Loss-Based Weighting), a novel aggregation strategy for federated learning (FL) designed to address challenges in wireless networks where non-IID data and frequent client dropouts are common. The authors are from Chungbuk National University (Republic of Korea) and Bennett University (India), with Taehong Kim as the corresponding author.

The paper identifies significant limitations in traditional FL weighting approaches, particularly FedAvg, which weights client updates solely based on dataset size. The authors state:

Classical FL algorithms, such as FedAvg [5], have weighted client updates based on the size of each client's dataset. However, this approach has significant limitations.

These limitations include:

  • Bias toward larger datasets: in scenarios involving non-IID data, it may not prioritize updates from clients with high-quality local models

  • Failure to account for non-IID nature: this method fails to account for the non-IID nature of data, a common characteristic in FL scenarios, especially in wireless networks where data heterogeneity is prevalent

  • Sensitivity to outliers: the data size-based weighting strategy is susceptible to outliers, allowing poorly trained local models from clients with large datasets to influence the global model disproportionately

  • Client dropout vulnerability: client dropouts, an inevitable characteristic of wireless networks... can lead to oscillations around stationary points of the global loss function, potentially causing severe performance degradation

  • Lack of incentive: this approach provides no incentive for clients to improve their local training processes, as their contributions are solely determined by the size of their datasets rather than the quality of their trained models

FedLBW replaces the traditional sample-count-based weighting with a performance-based weighting scheme:

Unlike the traditional approach, our performance-based weighting scheme assigns higher weights to clients with lower validation loss values, effectively giving more importance to better-performing local models.

The aggregation formula is:

"Our proposed FedLBW modifies these weights based on the inverse of local model's validation loss values as:

w(t+1) ← Σ(c∈S t) [(1/L c(t+1)) / (Σ(j∈S t) 1/L j(t+1))] w c(t+1)"

where L c(t+1) is the validation loss of client c's returned model computed on a small proxy dataset on the server:

L c(t+1) = (1/V) Σ(x,y)∈V l(f w c(t+1)(x), y)

The paper justifies the use of proxy datasets:

"FedLBW employs a small, class-balanced proxy dataset on the server to quantitatively assess client model contributions... This aligns with established research in FL. Li et al. demonstrated in FedLAW that a minimal proxy dataset (10 samples per class) effectively determines client contribution importance."

The privacy guarantees are preserved:

client training data remains distributed while the proxy dataset serves exclusively for evaluation—representing a balanced approach between performance optimization and privacy preservation

The authors list four main contributions:

  1. Novel weighting strategy: We propose a novel weighting strategy (FedLBW) for the aggregation step in FL that uses the inverse of local model's validation loss values as weights instead of number of data samples.

  2. Theoretical convergence analysis: We provide a theoretical convergence analysis for FedLBW.

  3. Extensive empirical evidence: "We provide extensive empirical evidence demonstrating FedLBW's robustness across diverse datasets (FashionMNIST, CIFAR-10, CIFAR-100) and model architectures (CNN, ResNet-18, ResNet-34), achieving consistent performance improvements and faster convergence, particularly in extreme non-IID scenarios (Dirichlet parameter α=0.1) and under low client participation ratios."

  4. Resilience to client dropouts: "We highlight FedLBW's exceptional resilience to client dropouts, especially in wireless networks, where it maintains stable performance even under extreme dropout conditions, outperforming both conventional methods, such as FedAvg, and SOTA methods, like FedLAW."

The paper provides a nonconvex convergence analysis with key assumptions including smoothness, bounded stochastic gradients, client heterogeneity, partial participation, and loss-based weight properties. The main theorem provides:

E∥∇F(w t)∥2 ≤ (4(F(w0)−F*))/(ηET) + 8L f η2E [(ρσ2)/(S tB) + (C−S t)/(S t(C−1)) ζ2 + G2(ρ−1)] + c d L f3 η3 G2 E3

with a convergence rate of O(1/√T) under appropriate stepsize selection.

The analysis shows FedLBW's advantage over FedAvg:

"if FedLBW weights satisfy a bounded loss ratio L min/L max ≥ δ > 0, then Σ(β c t)2 ≤ 1/(δS t) and hence ρ ≤ 1/δ. Thus, when δ is not too small, FedLBW yields smaller ρ (and a smaller bias term G2(ρ−1)) than FedAvg."

  • Framework: Python 3, PyTorch, and FedEasy

  • System: 1 server, 100 clients (AMD Ryzen Threadripper 3960X, Nvidia RTX 3060)

  • Datasets: FashionMNIST (CNN), CIFAR-10 (ResNet-18), CIFAR-100 (ResNet-34)

  • Data partitioning: Dirichlet distribution with α = 0.1 (extreme non-IID), 0.3 (moderate), 0.6 (mild)

  • Proxy dataset: 100 samples per class for FashionMNIST and CIFAR-10 (1,000 total), and 10 samples per class for CIFAR-100 (1,000 total)

  • Hyper-parameters: CPR = 0.1, SGD with learning rate 0.01, batch size 32, 3 local epochs, 300 rounds

  • Baselines: FedAvg, FedAvgM, FedProx, FedNova, FedLAW, and FedDkw

FedLBW consistently outperforms all baselines. Key highlights:

  • FashionMNIST (α=0.1): FedLBW achieves 83.51%, surpassing FedAvg (81.19%) and FedLAW (82.02%)

  • CIFAR-10 (α=0.1): FedLBW achieves 58.63%, with a 7.66% improvement over FedAvg (50.97%) and 9.44% advantage over FedLAW (49.19%)

  • CIFAR-100: FedLBW achieves 39.11%, 41.29%, and 42.01% for α = 0.1, 0.3, 0.6 respectively

The authors note:

"FedLBW exhibits its strongest gains under highly non-IID conditions (e.g., α = 0.1), where data heterogeneity among clients is most pronounced... By down-weighting these clients during aggregation, FedLBW prevents the global model from being dominated by class-specific or biased updates, effectively stabilizing learning and preserving a balanced global representation."

Under varying dropout probabilities on CIFAR-10 with ResNet-18 (α=0.3):

"Under severe dropout (p = 0.5), FedAvg drops sharply to 39.81%, whereas FedLBW and FedLAW remain stable at 64.40% and 63.72%. Thus, FedLBW preserves over 90% of its baseline performance even when half the clients disconnect."

The paper explains: FedAvg's 28.03% drop from p=0.0 to p=0.5 compared to FedLBW's small 6.25% decline.

FedLBW shows particular strength at low participation:

In the most challenging scenario, involving 100 total clients and a 0.1 CPR, FedLBW achieves an accuracy of 49.89%, significantly outperforming FedAvg (43.80%) and FedLAW (43.32%)

With 50 clients at full participation, FedLBW reaches 65.49% versus FedAvg's 59.88%.

FedLBW demonstrates robustness to various proxy distribution shifts:

"FedLBW retains high performance across a wide range of proxy–client distribution shifts. While closer alignment between client and proxy datasets generally improves the reliability of loss-based aggregation, FedLBW demonstrates robustness even under substantial divergence, including severe domain shifts."

Even with SVHN as a domain-shifted proxy (FID=138.37, Spearman ρ=0.4219), accuracy remains at 70.03% compared to 72.20% with the default CIFAR-10 proxy.

The paper concludes:

"This study presents FedLBW, a novel loss-based weighting strategy for model aggregation in FL, which effectively addresses the limitations of traditional weighting approaches. Our comprehensive empirical evaluations demonstrate FedLBW's superior performance across diverse datasets, model architectures, and challenging non-IID data distributions. The proposed approach not only enhances model accuracy and convergence but also exhibits remarkable resilience to client dropouts—maintaining stability even under extreme dropout rates, where conventional methods like FedAvg experience catastrophic degradation."

Future work directions include:

"Investigation of alternative weighting functions or combinations of weighting factors could lead to even further improvements in the performance and robustness of FL systems, particularly in challenging scenarios with highly non-IID data or specific application domains especially in wireless network environments."

  • Majid Kundroo: Writing – original draft, Software, Methodology, Conceptualization

  • Tinku Singh: Writing – original draft, Data curation, Methodology

  • Taehong Kim: Writing – review & editing, Supervision, Conceptualization

The work was supported by an IITP grant funded by the Korean government (MSIT) under grant RS-2025-02217801.

Improvements for AI systems

Improved AI systems can now:

  • Weight client models by quality, not dataset size: Instead of averaging updates proportionally to sample counts, FL systems can assign aggregation weights based on the inverse of each client’s validation loss computed on a small, class-balanced proxy dataset. This means the global model prioritizes updates from clients with genuinely better local models, reducing the influence of large but poorly trained datasets.

  • Maintain high accuracy under extreme non-IID data: With Dirichlet α=0.1, the improved system outperforms FedAvg by up to 7.66 percentage points on CIFAR-10 and surpasses state-of-the-art FedLAW by 9.44 points. It prevents class-specific or biased client updates from dominating the global model, preserving a balanced global representation.

  • Stay stable when clients drop out: In wireless networks with dropout probabilities up to 0.5, the improved FL system retains over 90% of its baseline accuracy (e.g., 64.40% vs FedAvg’s 39.81% on CIFAR-10), avoiding the catastrophic degradation that sample-count weighting suffers.

  • Learn faster with low participation: With only 10% client participation and 100 clients, accuracy reaches 49.89% vs FedAvg’s 43.80%, enabling practical FL deployments with limited connectivity and bandwidth.

  • Assess model contributions without violating privacy: The server’s small proxy dataset (e.g., 10–100 samples per class) evaluates each returned model’s validation loss, eliminating the need to access raw client data while still measuring model quality. This makes the system more privacy-preserving than approaches relying on shared statistics.

  • Remain robust to proxy data distribution shifts: Even when the proxy dataset is significantly shifted from client data (e.g., SVHN as proxy for CIFAR-10), the improved system still achieves 70.03% accuracy, retaining 97% of its performance with a well-aligned proxy. This ensures reliability in real-world deployments where a perfect proxy distribution is unavailable.

  • Provide theoretical convergence guarantees: The system’s aggregation strategy provably converges at O(1/√T) for nonconvex objectives, and outperforms FedAvg under bounded loss-ratio conditions. This gives engineers confidence in stable training dynamics and predictable behavior across rounds.

  • Integrate seamlessly with existing FL frameworks: The improvement is a drop-in replacement for the aggregation step in standard FL systems (FedAvg, FedProx, FedNova, etc.), requiring no changes to local training, communication protocol, or model architecture. It works with CNNs, ResNet-18, ResNet-34, and datasets like FashionMNIST, CIFAR-10, and CIFAR-100.

  • Extend to multi-agent and ensemble learning: Any AI system that aggregates multiple independently-trained models—such as distributed ensemble learning, split learning, or decentralized model fusion—can adopt loss-based weighting to assign confidence to each model based on its validation performance, improving robustness to heterogeneous data sources and unreliable contributors.

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.

Sources

Related papers