🤖 machine learning

Multi-Level Analyzation of Imbalance to Resolve Non-IID-Ness in Federated Learning

This paper proposes FedBB, a federated learning framework that addresses data imbalance at inter-case, inter-class, and inter-client levels through a Positive Negative Balanced loss function and Client Balanced Reweighting strategy, thereby outperforming existing algorithms in accuracy and efficiency while preserving privacy.

Original authors: Haengbok Chung, Jae Sung Lee

Published 2026-06-10

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haengbok Chung, Jae Sung Lee

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a group of doctors from different hospitals trying to build a single, super-smart AI to diagnose diseases. They want to do this without sharing their actual patient files (to protect privacy), so they train their own local AI models and send just the "lessons learned" back to a central server to combine them. This is Federated Learning.

However, there's a big problem: The data is messy.

Some hospitals see mostly healthy patients, some see mostly sick ones, and some see a weird mix of rare diseases. In the paper, the authors call this "Non-IID" (data that isn't identical or evenly distributed). Because of this mess, the final combined AI often gets confused and performs poorly.

The authors of this paper, Haengbok Chung and Jae Sung Lee, say, "We need to look at this mess from three different angles." They call their solution FedBB (Federated Learning with Positive-Negative Balanced loss and Client Balanced Reweighting).

Here is how they break it down using simple analogies:

1. The Three Levels of Messiness

The authors realized that "imbalance" happens in three different ways, like looking at a pile of fruit from a microscope, a window, and a drone:

Inter-Case (The Microscope View):
- The Problem: Imagine a specific disease, like "Pneumonia." For every one patient who has pneumonia, there might be 100 who don't. The AI gets lazy and just guesses "No Pneumonia" for everyone because it's right 99% of the time.
- The Fix: The AI needs to be forced to pay extra attention to the rare "Yes" cases.
Inter-Class (The Window View):
- The Problem: Now look at the whole fruit basket. Maybe you have 1,000 apples, 100 oranges, but only 5 bananas. The AI becomes an "Apple Expert" but a terrible "Banana Expert" because it never saw enough bananas.
- The Fix: The AI needs to be told, "Hey, those 5 bananas are super important! Don't ignore them."
Inter-Client (The Drone View):
- The Problem: Imagine Hospital A has 10,000 patients but they are all "Apple" cases. Hospital B has only 100 patients, but they are a perfect mix of Apples, Oranges, and Bananas.
- The Old Way: The central server usually says, "Hospital A has more data, so we'll listen to them 100 times more."
- The Result: The final AI becomes an "Apple Expert" and forgets everything else.
- The Fix: The server should actually listen more to Hospital B, because their small dataset is more balanced and teaches the AI a wider variety of lessons.

2. The Solution: FedBB

The authors built a two-part tool to fix these problems:

Part A: The "Fair Teacher" (PNB Loss)

This is a new way of grading the AI during its local training.

How it works: Instead of just counting how many times the AI got a question right, this method acts like a strict teacher who gives extra credit for getting the rare or difficult cases right.
The Analogy: If a student gets 100 easy questions right but misses the one hard question, the teacher doesn't just say "Good job." The teacher says, "You missed the rare one! That counts for double points." This forces the AI to stop ignoring the minority cases (the rare diseases or the few bananas).
Bonus: This works for both "Multi-class" (pick one category) and "Multi-label" (pick many categories at once, like a patient having both pneumonia and a broken rib).

Part B: The "Smart Mixer" (CBR)

This is a new way for the central server to combine the models from all the hospitals.

How it works: Instead of just counting how many patients a hospital has, the server checks how balanced their data is.
The Analogy: Imagine mixing smoothies. If you have a giant bucket of just strawberry smoothie (Hospital A) and a small cup of a perfect fruit salad (Hospital B), you shouldn't just dump the whole bucket in. You should give the small cup of fruit salad a louder voice in the final mix because it represents a more diverse flavor profile.
The Result: The final global model becomes a "Generalist" that is good at diagnosing everything, not just the most common things.

3. Why This Matters (According to the Paper)

The authors tested this on two types of data:

Medical X-rays: Real chest X-rays from different hospitals (NIH and CheXpert datasets).
Natural Images: Standard computer vision datasets like CIFAR (cats, dogs, cars, etc.).

The Results:

Better Performance: FedBB beat all other popular methods (like FedAvg, FedProx, MOON) in accuracy. It was especially good at handling the "messy" data where some diseases or objects were very rare.
Efficiency: It didn't require a supercomputer. It used the same amount of computing power as the basic method (FedAvg) but got much better results.
Privacy: It only needs to know a little bit of statistics (like "how balanced is your data?") rather than seeing the actual patient files, which keeps privacy safe.

Summary

The paper argues that to build a truly smart AI in a privacy-friendly way, we can't just throw all the data together. We have to:

Teach the local AI to care about the rare cases (using the PNB Loss).
Listen to the local AI that has the most diverse data, even if it's smaller (using CBR).

By doing this, the final "Global Brain" becomes fairer, more accurate, and ready to handle the real world, where data is rarely perfect.

Technical Summary: Multi-Level Analyzation of Imbalance to Resolve Non-IID-Ness in Federated Learning

1. Problem Definition

Federated Learning (FL) enables collaborative model training across distributed clients while preserving data privacy. However, a critical barrier to its efficacy is non-identically distributed (non-IID) data, often driven by severe class imbalance. Existing solutions for non-IID data often overlook the nuanced nature of class imbalance, particularly in multi-label classification scenarios.

The authors identify that conventional approaches to class imbalance are insufficient because they:

Focus primarily on multi-class classification, failing to address the complexity of multi-label tasks where $2^n$ label combinations exist.
Ignore inter-case imbalance, where the disparity between positive and negative cases within a single category causes models to overfit to the majority class (typically the negative case).
Rely on dataset size for aggregation, assigning higher weights to clients with larger datasets. In FL, large datasets are often highly skewed, leading to a biased global model that favors majority classes seen frequently by those specific clients.

2. Methodology: FedBB

To address these issues, the authors propose FedBB (Federated Learning with Positive Negative Balanced loss and Client Balanced Reweighting). The core innovation is a multi-level analysis of imbalance divided into three distinct scales:

Inter-case Imbalance: The discrepancy between the number of positive and negative cases within a single category.
Inter-class Imbalance: The difference in data distribution among different categories within a single client.
Inter-client Imbalance: The heterogeneity of data skewness across different clients.

FedBB addresses these levels through two primary components:

A. Positive Negative Balanced (PNB) Loss Function

Designed for local training, the PNB loss function mitigates both inter-case and inter-class imbalances.

Mechanism: Unlike standard Cross-Entropy or Class-Balanced (CB) loss, PNB explicitly calculates weights based on the effective number of positive and negative cases for each category.
Multi-Label Adaptation: It acknowledges that in multi-label tasks, samples with multiple co-occurring rare conditions (e.g., multiple diseases in a chest X-ray) are statistically rarer. PNB assigns higher weights to these minority cases and categories to prevent the model from being dominated by negative or majority-class samples.
Formulation: The loss incorporates smoothing via a temperature parameter ( $\tau$ ) and a hyperparameter ( $\beta$ ) to calculate effective sample numbers, preventing extreme weighting when data is scarce. A scaling factor ( $\mu$ ) is used to regulate the overall loss magnitude.

B. Client Balanced Reweighting (CBR)

Designed for model aggregation, CBR adjusts the weight of each client's contribution to the global model.

Mechanism: Instead of weighting clients solely by dataset size (as in FedAvg), CBR evaluates the skewness of a client's local data distribution.
Calculation: It calculates the standard deviation of class counts per client to determine data balance. Clients with less skewed (more balanced) datasets are assigned higher weights, as their local models are deemed to have better generalization capabilities.
Hybrid Weighting: The final aggregation weight ( $W_k$ ) is a linear combination of the balance-based weight ( $\omega^a_k$ ) and the size-based weight ( $\omega^b_k$ ), controlled by a hyperparameter $\gamma$ . This ensures that clients with large but highly skewed datasets do not dominate the global model, while clients with small but balanced datasets are not ignored.

3. Key Contributions

The paper claims the following contributions:

Multi-Level Imbalance Analysis: A novel framework that decomposes class imbalance in FL into inter-case, inter-class, and inter-client levels, arguing that addressing all three is necessary for robust performance.
PNB Loss: A loss function that simultaneously addresses inter-case and inter-class imbalances, applicable to both multi-label and multi-class classification, assigning higher weights to minority cases and classes.
CBR Aggregation: A model aggregation strategy that prioritizes clients with less skewed data distributions, enhancing the generalizability of the global model without requiring access to raw local data (preserving privacy).
Efficiency and Privacy: The method achieves state-of-the-art (SOTA) performance with minimal computational overhead and relies only on limited statistical information (class counts), avoiding the leakage of raw data distributions.

4. Experimental Results

The authors evaluated FedBB on medical X-ray datasets (NIH CXR14, CheXpert) and natural image datasets (CIFAR-10/100, Tiny-ImageNet, and artificially created Long-Tailed variants).

Performance: FedBB consistently outperformed existing algorithms (FedAvg, FedProx, MOON, FedNova, FedAlign, FedLC) across all datasets and imbalance settings (controlled by parameter $\delta$ $δ$ ).
- On NIH CXR14, FedBB achieved an AUC of 72.47 (at $\delta=0.1$ ), surpassing FedAvg (67.77) and FedProx (68.51).
- On LT-CIFAR-100, FedBB achieved 45.13% accuracy, outperforming FedLC (41.08%) and FedAlign (41.37%).
Efficiency: Despite the added complexity of calculating imbalance factors, FedBB maintains computational efficiency comparable to FedAvg.
- Parameters: 0.61 million (same as FedAvg).
- MFLOPs: 87.3 million (significantly lower than MOON's 262.2 million).
Ablation Studies: Experiments confirmed that both PNB loss and CBR independently contribute to performance improvements. The method also demonstrated the ability to harmonize with other FL algorithms (e.g., combining FedBB with FedProx or MOON) to further boost accuracy.

5. Significance and Claims

The paper positions FedBB as a baseline for generic and personalized Federated Learning. Its significance lies in:

Addressing the Root Cause: By treating class imbalance as a multi-level problem rather than a single metric, it resolves the specific non-IID-ness that degrades FL performance.
Privacy Preservation: The reliance on aggregate statistics (class counts) rather than raw data distributions makes it suitable for privacy-sensitive domains like healthcare.
Versatility: The approach is effective for both multi-label (medical imaging) and multi-class (natural images) tasks, a versatility not commonly found in prior imbalance-specific FL methods.

The authors conclude that while FedBB requires a centralized server and is parameter-based (making it theoretically vulnerable to certain attacks), it represents a significant step forward in creating global models that accurately classify all classes, even in highly skewed federated environments. Future work is suggested to explore serverless, non-parameter-based methods to further mitigate security risks.