ATC-FL: Adaptive Thompson Sampling for Efficient Clustering in Federated Learning

Federated Learning (FL) enables collaborative model training across decentralized clients without sharing raw data. To improve scalability and address client heterogeneity, clustering-based FL methods group similar clients to train specialized models. However, under realistic network constraints, existing approaches often incur high communication overhead or fail to adapt to dynamic conditions. We propose ATC-FL (Adaptive Thompson-sampling based Clustering for Federated Learning), a framework that adaptively selects and maintains client clusters to balance communication cost and model performance. ATC-FL formulates cluster count selection as a multi-armed bandit problem, using fixed-budget Thompson sampling to explore candidate configurations. At each round, client updates are embedded and clustered via K-means, with the Davies-Bouldin Index (DBI) serving as a low-overhead reward signal. A re-clustering scheduler adapts to client participation dynamics while bounding server computation. We evaluate ATC-FL on Human Activity Recognition and CIFAR-10 tasks over two emulated backbone-style network topologies (NSFNET and fat-tree) under non-IID data distributions, and additionally on a MIMIC-III medical benchmark to assess cross-domain generalization. We further analyze how ATC-FL’s convergence behavior and communication cost to reach a target accuracy vary with the severity of data heterogeneity. Compared to resource-aware federated clustering methods, ATC-FL achieves faster convergence, lower bandwidth usage, and reduced delay. These results suggest that adaptive clustering and network-aware control can improve efficiency while maintaining stable performance in realistic federated deployments.