Fast Convergence through Distributed Augmentation for Class-Imbalanced Federated Learning

Arathi Nair M; J. Harshan; Anwitaman Datta

cs.LGOct 4, 2026

Fast Convergence through Distributed Augmentation for Class-Imbalanced Federated Learning

Authors: Arathi Nair M, J. Harshan, Anwitaman Datta

Organizations: Amar Nath and Shashi Khosla School of Information Technology, Indian Institute of Technology Delhi, India · Department of Electrical Engineering, Indian Institute of Technology Delhi, India · College of Computing and Data Science, Nanyang Technological University, Singapore

Abstract

In federated learning, mitigating class imbalance is essential to improve minority-class performance. A common approach to address this problem is to augment minority-class samples to achieve local class balance. Existing approaches treat augmentation as a heuristic and do not establish how the amount of augmentation influences the convergence of federated learning, leading to excessive augmentation and increased training time. To address this limitation, we first establish the relationship between augmentation and the convergence behavior of federated learning. Leveraging this insight, we propose DAFL, a distributed augmentation framework that determines the minimum augmentation required for each client-class pair by jointly minimizing augmentation and training time while constraining global class imbalance, thereby improving minority-class F1-score. Experimental results demonstrate that DAFL consistently improves minority-class F1-score while substantially reducing training time, particularly under severe global class imbalance and high label proportion imbalance.

Figures & tables

Category / References		CI	MS	CD	LA	Limit.
Loss modification	Imbalance [ 6 , 7 ]	✔	✘	⚫	✘	Does not increase diversity of minority-class data
	Non-IID [ 10 , 11 , 12 , 13 ]	⚫	✘	✔	✘
Aggregation- based	Imbalance [ 8 , 9 ]	✔	✘	⚫	✘
	Non-IID [ 14 , 15 , 16 ]	⚫	✘	✔	✘
Data augmentation [ 17 , 18 , 19 , 20 , 21 , 23 ]		✔	✔	⚫	✘	Heuristic augmentation policies
DAFL (proposed)		✔	✔	✔	✔

TABLE I: Comparison of existing methods with DAFL. ✔, ✘, and ⚫ denote yes, no, and indirect support, respectively. Indirect: addressed by some works or achieved as a byproduct. CI: class imbalance; MS: minority samples; CD: client drift; LA: latency-aware; Limit.: severe classimbalance limitation.

	D1								D2
$\mathcal{I}$ / $J$	0.0002 / 0.4097								0.3982 / 0.4097
Desc.	Low $\mathcal{I}$ , High $J$								High $\mathcal{I}$ , High $J$
	C1	C2	C3	C4	C5	C6	C7	C8	C1	C2	C3	C4	C5	C6	C7	C8
Client 1	52	18	61	2	30	612	70	5	1	126	1	1	1	1	485	1
Client 2	63	21	74	3	28	701	65	6	1	1	1	10	248	1	1	1
Client 3	48	20	69	4	33	5800	60	4	395	2	1	1	1	1	1	1

TABLE II: Dataset details for D1–D6, including client-class counts (C1–C8 denote Class 1–Class 8) and the corresponding

\mathcal{I}

and

J

values. D1–D3 isolate extreme

\mathcal{I}

J

settings while D4–D6 cover intermediate settings. Here, the descriptors low/moderate/high are relative rather than specific thresholds. Also,

J\in\left[0,\frac{C-1}{C}\right]

Dataset	$\mathcal{I}$				$J$				Added Samples				Time for Augmentation (s)				Time at Server to find $\hat{\mathbf{A}}$ (s)
Dataset	LB	DAFL			LB	DAFL			LB	DAFL			LB	DAFL
		$q=2$	$q=3$	$q=4$		$q=2$	$q=3$	$q=4$		$q=2$	$q=3$	$q=4$		$q=2$	$q=3$	$q=4$	$q=2$	$q=3$	$q=4$
D1	0	0.0001	0.0001	0.0001	0.001	0.001	0.001	0.001	24913	24909	24909	24909	0.993	0.390	0.291	0.373	33.666	35.219	34.900
D2	0.0005	0.1479	0.1479	0.0798	0.0005	0.001	0.001	0.001	35571	24913	24913	24911	7.258	0.351	0.343	0.364	34.782	38.081	40.032
D3	0.0009	0.0558	0.2836	0.3492	0.0001	0.001	0.001	0.001	81622	39060	3080	7	5.526	0.682	0.042	0.297	55.485	4.238	0.012
D4	0.0003	0.0795	0.0795	0.0461	0.0007	0.001	0.001	0.001	28020	24309	24309	24311	4.406	0.334	0.395	0.318	32.878	32.897	34.163

TABLE III: Updated

\mathcal{I}

J

, number of augmented samples, and time associated with augmentation for LB and DAFL

(\lambda_{2}=10^{q})

Panel of Figure 4 — Fig. 1: Heatmap for client-class counts for D2 (a) before augmentation, (b) after LB, and (c) after DAFL.

Panel of Figure 5 — Fig. 2: Performance on datasets D1 (a)–(c), D2 (d)–(f), D3 (g)–(i), D4 (j)–(l), D5 (m)–(o), and D6 (p)–(r), with each triplet showing minority-class F1-score, accuracy, and macro F1-score, respectively.

Table 6

	$\displaystyle\mathbb{E}_{t}[F(\mathbf{w}^{(t+1)})]\leq F(\mathbf{w}^{(t)})-\eta\langle\nabla F(\mathbf{w}^{(t)}),\mathbb{E}_{t}[\mathbf{g}^{(t)}]\rangle+\frac{L\eta^{2}}{2}\mathbb{E}_{t}[\\|\mathbf{g}^{(t)}\\|^{2}].$		(18)
	$\displaystyle\nabla F_{k}(\mathbf{w})-\nabla F(\mathbf{w})=\underbrace{\textstyle\sum_{c}(\pi_{k}(c)-\pi(c))\nabla\bar{F}_{c}(\mathbf{w})}_{=:\,u_{k}(\mathbf{w})\ \text{(label proportion imbalance)}}+\underbrace{\textstyle\sum_{c}\pi_{k}(c)(\nabla F_{k,c}(\mathbf{w})-\nabla\bar{F}_{c}(\mathbf{w}))}_{=:\,v_{k}(\mathbf{w})\ \text{(conditional shift)}}.$		(19)
	$\displaystyle\mathbb{E}_{t}\\|\mathbf{z}^{(t)}\\|^{2}\;\leq\;(E-1)^{2}[8\eta^{2}E^{2}L^{2}\\|{\nabla F(\mathbf{w}^{(t)})}\\|^{2}+16\eta^{2}E^{2}L^{2}(CG^{2}\mathcal{I}+\Gamma)+2\eta^{2}E^{2}L^{2}\sigma^{2}].$		(20)
	$\displaystyle\frac{L\eta^{2}}{2}\mathbb{E}_{t}\\|{\mathbf{g}^{(t)}}\\|^{2}\leq\frac{3L\eta^{2}E^{2}}{2}\\|{\nabla F(\mathbf{w}^{(t)})}\\|^{2}+L^{3}\eta^{4}E^{2}(E-1)^{2}(12\\|{\nabla F(\mathbf{w}^{(t)})}\\|^{2}+24(CG^{2}\mathcal{I}+\Gamma)+3\sigma^{2})+\frac{3L\eta^{2}E}{2}\sigma_{g}^{2}.$		(21)
	$\displaystyle\mathbb{E}_{t}[F(\mathbf{w}^{(t+1)})]\leq F(\mathbf{w}^{(t)})\;-\;[\eta E-\frac{\eta}{8}-16A-\frac{3L\eta^{2}E^{2}}{2}-12B]\\|{\nabla F(\mathbf{w}^{(t)})}\\|^{2}+(32A+24B)(CG^{2}\mathcal{I}+\Gamma)+(4A+3B)\sigma^{2}+\frac{3L\eta^{2}E}{2}\sigma_{g}^{2}.$		(22)
	$\displaystyle\mathbb{E}_{t}[F(\mathbf{w}^{(t+1)})]\;\leq\;\;F(\mathbf{w}^{(t)})-\frac{\eta E}{2}\\|{\nabla F(\mathbf{w}^{(t)})}\\|^{2}+\frac{3L\eta^{2}E}{2}\sigma_{g}^{2}+40\,\eta^{3}E^{2}(E-1)^{2}L^{2}(CG^{2}\mathcal{I}+\Gamma+\sigma^{2}).$		(23)
	$\displaystyle\frac{1}{T}{\textstyle\sum_{t=0}^{T-1}\mathbb{E}\\|{\nabla F(\mathbf{w}^{(t)})}\\|^{2}}\;\leq\;\frac{1}{\eta ET}2(F(\mathbf{w}^{0})-F(\mathbf{w}^{*}))+3L\eta\,\sigma_{g}^{2}+80\,\eta^{2}E(E-1)^{2}L^{2}(CG^{2}\mathcal{I}+\Gamma+\sigma^{2}).$		(24)

Table 7

Explore similar work

FedReLa: Imbalanced Federated Learning via Re-Labeling
Jun 24, 2026Guangzheng Hu, Patricia Menéndez, Feng Liu +3Federated Learning Imbalanced Classification
School of Mathematics and Statistics, University of Melbourne, Victoria, Australia · School of Computing and Information Systems, University of Melbourne, Victoria, Australia · School of Statistics and Data Science, LPMC, KLMDASR, and LEBPS, Nankai University
Class-Grouped Normalized Momentum and Faster Hyperparameter Exploration to Tackle Class Imbalance in Federated Learning
Jul 1, 2026Haemin Park, Diego Klabjan, Martin W. Braun +2Federated Learning Imbalanced Classification
Department of Industrial Engineering & Management Sciences, Northwestern University, Evanston, IL, USA · Intel Corporation, Chandler, AZ, USA
WHERE to Generate Matters: Budget-Aware Synthetic Augmentation for Label Skewed Federated Learning
Jul 7, 2026Sangwoo Lee, Sunghwan Park, Jaewoo LeeFederated Learning Fedavg
Chung-Ang University, Seoul, South Korea

Jun 24, 2026, stat.MLJ/K move · Enter open · S save

FedReLa: Imbalanced Federated Learning via Re-Labeling

Guangzheng Hu, Patricia Menéndez, Feng Liu, Mingming Gong+2

School of Mathematics and Statistics, University of Melbourne, Victoria, Australia · School of Computing and Information Systems, University of Melbourne, Victoria, Australia · School of Statistics and Data Science, LPMC, KLMDASR, and LEBPS, Nankai University

Federated Learning Imbalanced Classification Fedavg Heterogeneity Label Propagation

Open paper AskPDF arXiv

Federated learning has emerged as the foremost approach for decentralized model training with privacy preservation. The global class imbalance and cross-client data heterogeneity naturally coexist, and the mismatch between local and global imbalances exacerbates the performance degradation of the aggregated model. The agnosticism of global class distribution poses significant challenges for data-level methods, especially under extreme conditions with severe class absence across clients. In this paper, we propose FedReLa, a novel data-level approach that tackles the coexistence of data heterogeneity and class imbalance in federated learning. By re-labeling samples with a feature-dependent label re-allocator, FedReLa corrects biased global decision boundaries without requiring knowledge of the global class distribution. This modular, model-agnostic approach can be integrated with algorithmic methods to deliver consistent improvements without additional communication overhead. Through extensive experiments, our method significantly improves the accuracy of minority classes and the overall accuracy on stepwise-imbalanced and long-tailed datasets, outperforming the previous state of the art.

Abstract

Explore similar work

FedReLa: Imbalanced Federated Learning via Re-Labeling

Class-Grouped Normalized Momentum and Faster Hyperparameter Exploration to Tackle Class Imbalance in Federated Learning

WHERE to Generate Matters: Budget-Aware Synthetic Augmentation for Label Skewed Federated Learning