🤖 machine learning

Fed-BAC: Federated Bandit-Guided Additive Clustering in Hierarchical Federated Learning

This paper proposes Fed-BAC, a hierarchical federated learning framework that integrates additive cluster personalization with a two-level bandit mechanism to jointly optimize server-to-cluster assignments and client selection, achieving significant accuracy gains, faster convergence, and improved fairness under severe data heterogeneity.

Original authors: Satwat Bashir, Tasos Dagiuklas, Muddesar Iqbal

Published 2026-05-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Satwat Bashir, Tasos Dagiuklas, Muddesar Iqbal

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 massive global classroom where thousands of students (clients) want to learn a subject together, but they are scattered across different cities. In a traditional setup, every student would have to send their homework directly to the main teacher at the central school (the Cloud). This creates a traffic jam, and the teacher gets overwhelmed.

Hierarchical Federated Learning (HFL) is like hiring local tutors (Edge Servers) in each city. The local tutors collect homework from their students, grade it, and send a summary to the main teacher. This reduces the traffic.

However, there's a problem: The "One-Size-Fits-All" Trap.
In the real world, students in different cities learn differently. Some are great at math but bad at art; others are the opposite. If the local tutor forces all students to learn from the exact same "average" textbook, the math geniuses get bored, and the art lovers get confused. This is called data heterogeneity.

Existing methods try to fix this by grouping students into separate classes (clusters), but they often do two things poorly:

They create isolated classrooms where no knowledge is shared between groups.
They pick which students to call on randomly, wasting time on students who aren't learning well.

Enter Fed-BAC: The Smart, Adaptive Classroom

The paper proposes a new system called Fed-BAC. Think of it as a super-smart educational system that uses two types of "gambling strategies" (called Bandits) to make decisions, combined with a unique way of teaching.

1. The Two-Level "Gambler" System

Fed-BAC uses two different decision-makers to solve the classroom chaos:

The Cloud Teacher (The "Contextual Bandit"):
Imagine the main teacher has a magic crystal ball. Every few weeks, the teacher looks at the local tutors and asks: "Which group of tutors should work together?"
Instead of guessing, the teacher uses a strategy called LinUCB. It looks at four clues (like how well the current group is doing, how balanced the groups are, etc.) to decide which tutors should be in the same "study group." If a group isn't working, the teacher shuffles them into a new group. This ensures the right tutors are working together.
The Local Tutor (The "Thompson Sampling" Picker):
Inside each city, the local tutor has to decide which 8 out of 10 students to call on for the next lesson.
Instead of picking randomly, the tutor uses a strategy called Thompson Sampling. Think of it like a slot machine. The tutor keeps a mental score for every student. If a student's homework improves the class average, the tutor gives them a "win" and is more likely to pick them next time. If they struggle, the tutor gives them a "loss" and picks someone else. Over time, the tutor learns exactly which students contribute the most and focuses on them.

2. The "Shared + Special" Textbook (Additive Clustering)

This is the most creative part. Most systems either use one big textbook for everyone or a completely different textbook for every single group. Fed-BAC does both at once.

Imagine every student gets a base textbook (The Global Network) that covers the common facts everyone needs to know. On top of that, every student gets a custom notebook (The Cluster-Specific Network) where they write down the specific tricks and shortcuts that work for their local style.

Why this is great: The "base textbook" allows different groups to share common knowledge (so they don't reinvent the wheel). The "custom notebook" lets them handle their unique local differences. It's like having a universal language plus a local dialect.

What Did They Find? (The Results)

The researchers tested this system on three different "subjects" (datasets: CIFAR-10, SVHN, Fashion-MNIST) with two levels of difficulty:

Moderate Difficulty: Students are somewhat different.
Severe Difficulty: Students are very different (some are total opposites).

Here is what happened:

It Gets Better When Things Get Messy: The more different the students were (Severe Difficulty), the more Fed-BAC crushed the competition. While other systems struggled and their scores dropped, Fed-BAC's scores actually went up. It proved that its "Shared + Special" approach is perfect for messy, real-world data.
It's a Speed Demon: Fed-BAC reached high accuracy 1.5 to 4.8 times faster than the other methods. It didn't need to wait for everyone to finish; by picking the best students (via the local tutor) and shuffling groups efficiently (via the cloud teacher), it learned quickly.
It Saves Energy (Sort of): Even though Fed-BAC only asked 80% of the students to participate in each round (saving energy and bandwidth), it still beat systems that forced 100% participation.
Fairness: In the old systems, some students (servers) were left behind with terrible grades. Fed-BAC made sure that even the "worst" student in the class did better than the "average" student in the other systems.

The Bottom Line

Fed-BAC is like a classroom that never stops optimizing itself. It constantly reorganizes study groups to find the best matches and picks the most helpful students to lead the discussion. By combining a shared foundation with personalized notes, it learns faster, handles chaos better, and ensures no one gets left behind, all while using fewer resources than the old ways.

Technical Summary: Fed-BAC

Problem Statement

Hierarchical Federated Learning (HFL) leverages edge servers to perform partial aggregation, reducing communication bottlenecks in Multi-Access Edge Computing (MEC) environments. However, existing HFL methods face three critical limitations when dealing with real-world data heterogeneity:

Static Assumptions: Most methods assume homogeneous data distributions across edge servers, leading to "client drift" and accuracy degradation (up to 20%) when aggregating updates from clients with divergent label distributions.
Isolated Clustering: While clustered FL approaches (e.g., IFCA) group similar clients, they typically train fully isolated models, preventing cross-cluster knowledge transfer.
Inefficient Selection: Client selection is often random or static. Existing bandit-based approaches operate in flat architectures and fail to jointly optimize cluster assignment and client selection across the three-tier topology (cloud-edge-client).

Furthermore, personalization methods that maintain per-client residuals (e.g., Ditto, APFL) incur $O(N)$ parameter overhead, which becomes impractical at scale and exacerbates communication costs in three-tier systems.

Methodology: Fed-BAC Framework

The paper proposes Fed-BAC (Federated Bandit-Guided Additive Clustering), a framework that integrates additive cluster personalization with a two-level bandit decision-making architecture.

1. Additive Cluster Personalization

Instead of training isolated cluster models, Fed-BAC employs an additive decomposition in the logit space. The predicted logits $\hat{y}$ for a client in cluster $k$ are defined as:
$\hat{y} = h(x; \Theta_{global}) + f(x; \Theta_k)$

$\Theta_{global}$ : A shared global network capturing common knowledge across all servers.
$\Theta_k$ : A cluster-specific network capturing distribution variations.
This architecture allows all $M$ edge servers to contribute to the global model every round, enabling cross-cluster knowledge transfer while preserving specialization through cluster-specific residuals.

2. Two-Level Bandit Framework

Fed-BAC decouples decision-making into two levels to address the three-tier topology:

Cloud Level (Cluster Assignment): Uses Contextual LinUCB to dynamically assign edge servers to clusters.
- Contextual Features: The bandit utilizes four FL-specific features: loss ratio (current vs. best alternative cluster), cluster balance, assignment stability, and training progress.
- Mechanism: Every $\tau_{re}$ rounds, the cloud evaluates the fit of each server to its current cluster versus alternatives. If a server's loss ratio indicates a better fit elsewhere, the assignment is updated. This prevents "herd behavior" and allows the system to discover the optimal number of active clusters ( $K \le M$ ).
Edge Level (Client Selection): Uses Thompson Sampling (TS) at each edge server to select a subset of clients ( $B$ out of $N_m$ ) for participation.
- Mechanism: Each server maintains a Beta distribution prior for every client. Based on the single-round accuracy change ( $r_{TS}$ ) of the server's cluster model, the posteriors are updated.
- Goal: To identify high-contributing clients dynamically, reducing communication overhead while maintaining model performance.

3. Training Protocol

The system operates in synchronous rounds involving:

Distribution: Cloud sends $\Theta_{global}$ and the specific $\Theta_k$ to each edge server.
Selection & Training: Edge servers use TS to select clients, who perform local SGD on both networks.
Aggregation:
- Edge Aggregation: Weighted averaging of client updates.
- Cloud Aggregation: Two-phase update. First, a global aggregation of $\Theta_{global}$ from all servers (using total server data sizes, not just selected subsets). Second, a cluster-specific aggregation of $\Theta_k$ for each active cluster.

Key Contributions

Two-Level Bandit Architecture: A novel hierarchical decision system that decouples cluster assignment (LinUCB at the cloud) from client selection (Thompson Sampling at the edge). This allows for different update frequencies suited to the timescales of each decision.
Integrated Additive Personalization: The first integration of additive cluster personalization within a three-tier HFL architecture. Unlike IFCA, which isolates clusters, Fed-BAC shares a global backbone, facilitating knowledge transfer between groups.
Joint Optimization: The framework jointly optimizes cluster assignment and client selection, addressing a gap where existing methods handle these tasks separately or assume static configurations.

Experimental Results

Experiments were conducted on three benchmarks (CIFAR-10, SVHN, Fashion-MNIST) under moderate ( $\alpha=0.5$ ) and severe ( $\alpha=0.1$ ) Dirichlet non-IID partitioning.

Accuracy Gains: Fed-BAC achieved distributed accuracy gains of up to +35.5 percentage points over HierFAVG and +8.4 percentage points over IFCA.
Heterogeneity Sensitivity: The performance advantage of Fed-BAC increased with the severity of data heterogeneity. Under severe heterogeneity ( $\alpha=0.1$ ), HierFAVG degraded significantly, while Fed-BAC maintained high accuracy.
Efficiency: Despite requiring only 80% client participation ( $p=0.8$ ), Fed-BAC converged 1.5–4.8× faster than baselines depending on the dataset and target accuracy. This speedup resulted in a 3.0× to 5.9× reduction in total communication bytes to reach representative accuracy targets.
Scalability: A scaling study on CIFAR-10 with 500 clients (5× scale) confirmed that Fed-BAC retains its lead over HierFAVG (+12.6pp) and IFCA (+3.3pp), validating its robustness at larger deployment scales.
Fairness: Fed-BAC demonstrated improved cross-server fairness, exhibiting lower variance in per-server accuracy compared to IFCA and HierFAVG, particularly under severe heterogeneity.

Significance and Claims

The paper claims that Fed-BAC addresses the critical challenges of statistical heterogeneity and communication efficiency in Hierarchical Federated Learning. By combining additive decomposition (enabling knowledge sharing) with bandit-guided adaptive decisions (optimizing both cluster structure and client participation), the framework outperforms existing state-of-the-art methods.

The authors emphasize that the additive personalization approach becomes increasingly valuable as data distributions diverge, confirming that sharing a global backbone while maintaining cluster-specific residuals is superior to fully isolated clustering. The results suggest that Fed-BAC provides a scalable solution for MEC deployments where data heterogeneity is severe and communication resources are constrained.

The paper concludes modestly, noting that the current evaluation uses a single participation rate and that future work should explore systematic sweeps of participation rates, real-world edge hardware testbeds, and formal convergence guarantees for the coupled non-stationary bandit system.