Heterogeneity-Aware Client Sampling for Optimal and Efficient Federated Learning

Shudi Weng; Chao Ren; Ming Xiao; Mikael Skoglund

cs.LGMay 16, 2025

Heterogeneity-Aware Client Sampling for Optimal and Efficient Federated Learning

Authors: Shudi Weng, Chao Ren, Ming Xiao, Mikael Skoglund

Organizations: Division of Information Science and Engineering, School of Electrical Engineering and Computer Science (EECS), KTH Royal Institute of Technology, 11428 Stockholm, Sweden

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.

Figures & tables

Fig. 1: An FL scenario at the network edge with heterogeneous communication probabilities $q_{m}$ and computation epochs $T_{m}$ .

Fig. 2 : Comparison of distance to the optimum in FedAvg with SGD [ 30 ] , proximal SGD (

\mu=1

), SGD with decayed learning rate (decay rate

0.005

), momentum SGD (momentum

0.3

), and our proposed FedACS in Example 1 , where

\boldsymbol{E}_{m}\sim\mathcal{N}(0,\mathbf{I}_{10\times 10})

M=30

K=15

, and

\eta=0.001

. Heterogeneous setting where

T_{m}

is uniformly distributed from

1

30

and

q_{m}

from

0.01

0.3

across clients.

Fig. 3: Illustration of the proposed biased client sampling scheme FedACS to counterbalance the objective inconsistency induced by heterogeneous communication and computation.

Datasets	MNIST	CIFAR-10	CINIC-10
Neural network	CNN	CNN	CNN
Model architecture ∗	C (1,10) – C (10,20) – D – L (50) – L (10)	C (3,32) – R – M – C (32,32) – R – M – L (256) – R – L (64) – R – L (10)	C (3,32) – R – M – C (32,32) – R – M – D – L (512) – R – D – L (256) – R – D – L (10)
Loss function	Negative log likelihood loss		Cross-entropy loss
Learning rate $\eta_{1}$	$\eta_{1}=0.02$		$\eta_{1}=0.03$
Number of global rounds $T$	200
Batch size	512		1024

TABLE I: NN architecture, loss function, learning rate, and batch size specifications

Fig. 4: Verification of consistent communication-computation design principle. In a heterogeneous setting,

T_{m}

takes fixed values in

\{2,3\}

, where the first

20

clients have

T_{m}=2

and the last 10 clients have

T_{m}=3

, and

q_{m}

from

0.01

0.3

across clients. Other parameters remain the same as in Fig. 2 .

Panel of Figure 6 — Fig. 5: Test Accuracy of algorithms over various datasets in dynamic heterogeneous FL. The average test accuracy over multiple runs is smoothed and plotted for each algorithm. Vertical bars indicate the variance at the corresponding communication round, reflecting the stability of each algorithm.

Algorithm	MNIST		CIFAR-10		CINIC-10
Algorithm	Rounds for 70%	Time (s) for 70%	Rounds for 40%	Time (s) for 40%	Rounds for 35%	Time (s) for 35%
FedACS	68 (1.0 $\times$ )	3548.6 (1.0 $\times$ )	100 (1.0 $\times$ )	6928.9 (1.0 $\times$ )	55 (1.0 $\times$ )	11452.3 (1.0 $\times$ )
FedAvg	93 (1.37 $\times$ )	5470.4 (1.38 $\times$ )	143 (1.43 $\times$ )	12482.0 (1.80 $\times$ )	74 (1.34 $\times$ )	14550.0 (1.27 $\times$ )
c-a-FedAvg	84 (1.24 $\times$ )	4884.4 (1.51 $\times$ )	130 (1.30 $\times$ )	7917.0 (1.14 $\times$ )	67 (1.22 $\times$ )	13256.8 (1.16 $\times$ )
FedVarp	115 (1.69 $\times$ )	7289.9 (2.05 $\times$ )	165 (1.65 $\times$ )	13479.8 (1.95 $\times$ )	104 (1.89 $\times$ )	22200.5 (1.94 $\times$ )
FedNova	99 (1.46 $\times$ )	6043.4 (1.70 $\times$ )	161 (1.61 $\times$ )	13314.8 (1.92 $\times$ )	100 (1.81 $\times$ )	21802.2 (1.90 $\times$ )

TABLE II: Training efficiency of algorithms over various datasets in dynamic heterogeneous FL. We report the average number of communication rounds and total computation time required to reach the predefined accuracy thresholds for the first time.

Algorithm	$\alpha=0.05$		$\alpha=0.2$		$\alpha=0.35$
Algorithm	Rounds for 35%	Time (s) for 35%	Rounds for 35%	Time (s) for 35%	Rounds for 40%	Time (s) for 40%
FedACS	107 (1.0 $\times$ )	6186.2 (1.0 $\times$ )	51 (1.0 $\times$ )	4747.3 (1.0 $\times$ )	36 (1.0 $\times$ )	3391.5 (1.0 $\times$ )
FedAvg	149 (1.39 $\times$ )	9591.7 (1.55 $\times$ )	88 (1.72 $\times$ )	8930.0 (1.88 $\times$ )	88 (2.44 $\times$ )	8918.3 (2.62 $\times$ )
c-a-FedAvg	142 (1.33 $\times$ )	8982.4 (1.45 $\times$ )	88 (1.72 $\times$ )	8873.4 (1.87 $\times$ )	90 (2.5 $\times$ )	9293.9 (2.74 $\times$ )
FedVarp	–	–	156 (3.06 $\times$ )	17728.6 (3.73 $\times$ )	159 (4.42 $\times$ )	17733.9 (5.22 $\times$ )
FedNova	–	–	125 (2.45 $\times$ )	13690.4 (2.88 $\times$ )	94 (2.61 $\times$ )	9901.3 (2.92 $\times$ )

TABLE III: Training Efficiency of algorithms on CIFAR-10 under varying levels of data imbalance in dynamic heterogeneous FL. We report the average number of communication rounds and total computation time required to reach the predefined accuracy thresholds for the first time across different datasets.

Algorithm	$\mathrm{frac}=0.3,\alpha=0.05$		$\mathrm{frac}=0.5,\alpha=0.05$		$\mathrm{frac}=0.4,\alpha=0.2$
Algorithm	Rounds for 35%	Time (s) for 35%	Rounds for 35%	Time (s) for 35%	Rounds for 35%	Time (s) for 40%
FedACS	107 (1.0 $\times$ )	6186.2 (1.0 $\times$ )	103 (1.0 $\times$ )	1912.3 (1.0 $\times$ )	40 (1.0 $\times$ )	1881.0 (1.0 $\times$ )
FedAvg	149 (1.39 $\times$ )	9591.7 (1.55 $\times$ )	138 (1.34 $\times$ )	3609.2 (1.89 $\times$ )	80 (2 $\times$ )	4326.2 (2.30 $\times$ )
c-a-FedAvg	142 (1.33 $\times$ )	8982.4 (1.45 $\times$ )	133 (1.29 $\times$ )	3647.9 (1.91 $\times$ )	81 (2.03 $\times$ )	4108.4 (2.18 $\times$ )
FedVarp	–	–	–	–	128 (3.2 $\times$ )	6704.2 (3.56 $\times$ )
FedNova	–	–	–	–	98 (2.45 $\times$ )	5021.8 (2.67 $\times$ )

TABLE IV: Training Efficiency of algorithms on CIFAR-10 under varying client participation ratios in dynamic heterogeneous FL. We report the average number of communication rounds and total computation time required to reach the predefined accuracy thresholds for the first time across different datasets.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

No.	Bounded Dissimilarity Assumptions	Ref.
4	$\forall m\in[M]:\lVert\nabla F_{m}(x)-\nabla F(x)\rVert^{2}\leq D_{m}^{2}$	[ 9 ]
5	$\sum_{m=1}^{M}\omega_{m}\lVert\nabla F_{m}(x)-\nabla F(x)\rVert^{2}\leq D^{2}$	[ 24 ]
6	$\exists\beta^{2}\geq 0,\kappa^{2}\geq 0:\sum_{m=1}^{M}\omega_{m}\lVert\nabla F_{m}(x)-\nabla F(x)\rVert^{2}\leq\beta^{2}\lVert\nabla F(x)\rVert^{2}+\kappa^{2}$	[ 8 ]
7	$\forall m\in[M]:\lVert\nabla F_{m}(x)\rVert^{2}\leq G^{2}$	[ 35 ]
8	$\forall m\in[M]:\lVert\nabla F_{m}(x\|\xi)\rVert^{2}\leq G^{2}$	[ 23 ]

Appendix

TABLE V: Common assumptions on bounded gradient dissimilarity.

Panel of Figure 11 — Fig. 6: Test Accuracy of algorithms over various datasets in static heterogeneous FL. The average test accuracy over multiple runs is smoothed and plotted for each algorithm. Vertical bars indicate the variance of raw data at the corresponding communication round, reflecting the stability of each algorithm. Vertical bars indicate the variance of the distance to the optimum.

Algorithm	MNIST		CIFAR-10		CINIC-10
Algorithm	Rounds for 70%	Time (s) for 70%	Rounds for 40%	Time (s) for 40%	Rounds for 35%	Time (s) for 35%
FedACS	64 (1.0 $\times$ )	4307.7 (1.0 $\times$ )	68 (1.0 $\times$ )	4706.3 (1.0 $\times$ )	36 (1.0 $\times$ )	12837.2 (1.0 $\times$ )
FedAvg	92 (1.44 $\times$ )	6458.7 (1.50 $\times$ )	112 (1.64 $\times$ )	8609.0 (1.83 $\times$ )	57 (1.58 $\times$ )	19999.3 (1.56 $\times$ )
c-a-FedAvg	68 (1.07 $\times$ )	5039.2 (1.17 $\times$ )	101 (1.49 $\times$ )	8221.3 (1.77 $\times$ )	55 (1.52 $\times$ )	18334.6 (1.43 $\times$ )
FedVarp	92 (1.44 $\times$ )	6846.9 (1.59 $\times$ )	148 (2.18 $\times$ )	11862.5 (2.52 $\times$ )	73 (2.03 $\times$ )	25479.9 (1.98 $\times$ )
FedNova	103 (1.61 $\times$ )	7893.4 (1.83 $\times$ )	117 (1.72 $\times$ )	8536.8 (1.81 $\times$ )	83 (2.31 $\times$ )	29968.5 (2.33 $\times$ )

Appendix

TABLE VI: Training Efficiency of algorithms over various datasets in static heterogeneous FL. We report the average number of communication rounds and total computation time required to reach the predefined accuracy thresholds for the first time across different datasets.

Panel of Figure 13 — Fig. 7: Test accuracy of algorithms on CIFAR-10 under varying levels of data imbalance in dynamic heterogeneous FL. The average test accuracy over multiple runs is smoothed and plotted for each algorithm. Vertical bars indicate the variance of raw data at the corresponding communication round, reflecting the stability of each algorithm. Vertical bars indicate the variance of the distance to the optimum.

Panel of Figure 14 — Fig. 8: Test accuracy of algorithms on CIFAR-10 under varying participation rates in dynamic heterogeneous FL. The average test accuracy over multiple runs is smoothed and plotted for each algorithm. Vertical bars indicate the variance of raw data at the corresponding communication round, reflecting the stability of each algorithm. Vertical bars indicate the variance of the distance to the optimum.

Explore similar work

May 4, 2026cs.DC

FedPLT: Scalable, Resource-Efficient, and Heterogeneity-Aware Federated Learning via Partial Layer Training

Federated Learning (FL) has gained significant attention in distributed machine learning by enabling collaborative model training across decentralized system while preserving data privacy. Although extensive research has addressed statistical data heterogeneity, FL still faces several challenges, including high communication and computation overheads and severe device heterogeneity, which require further investigation. Prior work has addressed these issues through sub-model training and partial parameter training. However, such methods often suffer from inconsistent parameter distributions across clients, inaccurate global loss estimation, and increased bias and variance. Guided by our empirical analysis, we propose FedPLT (Federated Learning with Partial Layer Training), an innovative and structured partial parameter training approach that exhibits training behavior similar to full model training while assigning client-specific portions of the model according to their communication and computational capabilities. In addition, we evaluate the performance of FedPLT when combined with optimal client sampling under communication constraints. We show that this integration improves FL performance by reducing sampling variance under the same communication budget. Through extensive experiments, we demonstrate that FedPLT achieves performance comparable to, or even surpassing, that of full-model training (i.e., FedAvg), while requiring significantly fewer trainable parameters per client. Moreover, FedPLT outperforms existing methods in highly heterogeneous environments, effectively adapts to client resource constraints, and reduces the number of straggling clients. In particular, FedPLT reduces the number of trainable parameters by 71%-82% while achieving performance on par with full-model training.

Ahmad Dabaja, Rachid El-Azouzi

Laboratoire Informatique d’Avignon (LIA), Avignon University, France

Communication-Efficient Distributed Training Heterogeneous Federated Learning

Ask PDF

Apr 28, 2026cs.LG

Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

Federated learning increasingly operates in a large-model regime where communication, memory, and computation are all scarce. Typically, non-IID client data induce drift that degrades the stability and performance of local training. Existing remedies such as SCAFFOLD introduce heterogeneity-correction mechanisms to address this challenge, but they incur substantial extra communication and memory overhead. This paper proposes a subspace optimization method for federated learning (SSF), which performs heterogeneity-corrected optimization in a low-dimensional subspace using only projected quantities, while preserving full-dimensional control information through a backfill-style update that retains residual components whenever the active subspace changes. Under standard smoothness and bounded-variance assumptions, SSF attains a non-asymptotic rate of order

\widetilde{\mathcal{O}}(1/T+1/\sqrt{NKT})

. Experiments show favorable accuracy--efficiency trade-offs under heterogeneous data.

Shuchen Zhu, Zhengyang Huang, Yuqi Xu +1

Peking University · Beihang University

Non-IID Federated Learning Federated Learning Heterogeneous Federated Learning

Ask PDF

Nov 6, 2023cs.LG

Convergence Analysis of Sequential Federated Learning on Heterogeneous Data

There are two categories of methods in Federated Learning (FL) for joint training across multiple clients: (i) parallel FL (PFL), where clients train models in a parallel manner; and (ii) sequential FL (SFL), where clients train models in a sequential manner. In contrast to that of PFL, the convergence theory of SFL on heterogeneous data is still lacking. In this paper, we establish the convergence guarantees of SFL for strongly/general/non-convex objectives on heterogeneous data. The convergence guarantees of SFL are better than that of PFL on heterogeneous data with both full and partial client participation. Experimental results validate the counterintuitive analysis result that SFL outperforms PFL on extremely heterogeneous data in cross-device settings.

Yipeng Li, Xinchen Lyu

National Engineering Research Center for Mobile Network Technologies Beijing University of Posts and Telecommunications Beijing, 100876, China

Non-IID Federated Learning Stochastic Optimization Convergence Federated Learning Heterogeneous Federated Learning

Ask PDF