Heterogeneity-Aware Client Sampling: A Unified Solution for Consistent Federated Learning
Shudi Weng, Chao Ren, Ming Xiao, Mikael Skoglund
TL;DR
The paper tackles objective inconsistency in federated learning arising from joint heterogeneity in communication and computation. It develops a unified non-convex analysis that separates structural and statistical effects of heterogeneity on surrogate and true objectives, proving convergence of the surrogate with rate $O(1/\sqrt{R})$ and bounding the true objective by $\chi_{\boldsymbol{\omega}\Vert \boldsymbol{\Omega}}^2\kappa^2$. To address the issue, the authors propose FedACS, an adaptive client-sampling scheme that aligns surrogate weights with target client weights via $p_m^{(r)}=\frac{\omega_m/( (1-q_m^{(r)})\|\boldsymbol{a}_m^{(r)}\|_1 )}{\sum_{j=1}^M \omega_j/( (1-q_j^{(r)})\|\boldsymbol{a}_j^{(r)}\|_1 )}$, ensuring unbiased aggregation and convergence to the true optimum when $\chi_{\boldsymbol{\omega}\Vert \boldsymbol{\Omega}}^2=0$. Theoretical results are complemented by extensive experiments on MNIST, CIFAR-10, and CINIC-10 under dynamic heterogeneity, showing substantial accuracy gains and reduced communication and computation costs. Overall, the work provides a principled framework and a universal method for making federated learning robust to both communication and computation heterogeneity, with practical implications for privacy-preserving deployments.
Abstract
Federated learning (FL) commonly involves clients with diverse communication and computational capabilities. Such heterogeneity can significantly distort the optimization dynamics and lead to objective inconsistency, where the global model converges to an incorrect stationary point potentially far from the pursued optimum. Despite its critical impact, the joint effect of communication and computation heterogeneity has remained largely unexplored, due to the intrinsic complexity of their interaction. In this paper, we reveal the fundamentally distinct mechanisms through which heterogeneous communication and computation drive inconsistency in FL. To the best of our knowledge, this is the first unified theoretical analysis of general heterogeneous FL, offering a principled understanding of how these two forms of heterogeneity jointly distort the optimization trajectory under arbitrary choices of local solvers. Motivated by these insights, we propose Federated Heterogeneity-Aware Client Sampling, FedACS, a universal method to eliminate all types of objective inconsistency. We theoretically prove that FedACS converges to the correct optimum at a rate of $O(1/\sqrt{R})$, even in dynamic heterogeneous environments. Extensive experiments across multiple datasets show that FedACS outperforms state-of-the-art and category-specific baselines by 4.3%-36%, while reducing communication costs by 22%-89% and computation loads by 14%-105%, respectively.