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

2026-05-12

AI Summary

This work addresses the challenge of jointly optimizing cluster assignment and client selection in hierarchical federated learning under data heterogeneity. The authors propose Fed-BAC, the first framework integrating additive clustering personalization with a two-level bandit mechanism: at the cloud level, a contextual bandit dynamically assigns servers to clusters, while at the edge, Thompson Sampling selects high-contribution clients. Fed-BAC models both shared and cluster-specific patterns through a global network and cluster-adapted networks operating in tandem. Evaluated on three benchmarks including CIFAR-10, Fed-BAC significantly outperforms HierFAVG and IFCA, achieving accuracy gains of up to 35.5 and 8.4 percentage points, respectively, with only 80% client participation. It accelerates convergence by 1.5–4.8× and maintains effectiveness and fairness even when scaled up fivefold.

Abstract

Hierarchical federated learning (HFL) leverages edge servers for partial aggregation in edge computing. Yet existing FL methods lack mechanisms for jointly optimizing cluster assignment and client selection under data heterogeneity. This paper proposes Fed-BAC, which integrates additive cluster personalization with a two-level bandit framework: contextual bandits at the cloud learn server-to-cluster assignments, while Thompson Sampling at each edge server identifies high-contributing clients. The additive decomposition enables the sharing of knowledge between groups through a globally aggregated network, while cluster-specific networks capture distribution variations. Across three classification benchmarks (CIFAR-10, SVHN, Fashion-MNIST) under moderate ($α= 0.5$) and severe ($α= 0.1$) Dirichlet non-IID partitioning, Fed-BAC achieves distributed accuracy gains of up to +35.5pp over HierFAVG and +8.4pp over IFCA, while requiring only 80% client participation, converging 1.5 to 4.8$\times$ faster depending on dataset and accuracy target, and improving cross-server fairness. These gains are further validated at 5$\times$ deployment scale on CIFAR-10. The advantage of Fed-BAC increases with heterogeneity severity, confirming that additive cluster personalization becomes increasingly valuable as data distributions diverge.

Problem

• Hierarchical Federated Learning
• Data Heterogeneity
• Cluster Assignment
• Client Selection
• Non-IID

Innovation

• Federated Learning
• Additive Clustering
• Contextual Bandits
• Thompson Sampling
• Hierarchical Federated Learning