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

TL;DR

Fed-BAC introduces a novel federated learning approach combining bandit algorithms and additive clustering to optimize client selection and cluster assignment, significantly improving accuracy and convergence speed in heterogeneous data environments.

Contribution

It proposes Fed-BAC, integrating bandit frameworks with additive clustering for enhanced personalization and efficiency in hierarchical federated learning.

Findings

Achieves up to +35.5 percentage points accuracy gain over HierFAVG.

Converges 1.5 to 4.8 times faster depending on dataset and target accuracy.

Requires only 80% client participation, improving fairness.

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 ($\alpha = 0.5$) and severe ($\alpha = 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.