Self-Supervised Federated Learning under Data Heterogeneity for Label-Scarce Diatom Classification

Mingkun Tan; Xilu Wang; Michael Kloster; Tim W. Nattkemper

cs.CVMar 31, 2026

Self-Supervised Federated Learning under Data Heterogeneity for Label-Scarce Diatom Classification

Authors: Mingkun Tan, Xilu Wang, Michael Kloster, Tim W. Nattkemper

Organizations: Biodata Mining Group, Faculty of Technology, University of Bielefeld, Germany · Computer Science Research Centre, University of Surrey, United Kingdom · Phycology Group, Faculty of Biology, University of Duisburg-Essen, Germany

Abstract

Label-scarce visual classification under decentralized and heterogeneous data is a fundamental challenge in pattern recognition, especially when sites exhibit partially overlapping class sets. While self-supervised federated learning (SSFL) offers a promising solution, existing studies commonly assume the same data heterogeneity pattern throughout pre-training and fine-tuning. Moreover, current partitioning schemes often fail to generate pure partially class-disjoint data settings, limiting controllable simulation of real-world label-space heterogeneity. In this work, we introduce SSFL for diatom classification as a representative real-world instance and systematically investigate stage-specific data heterogeneity. We study cross-site variation in unlabeled data volume during pre-training and label-space misalignment during downstream fine-tuning. To study the latter in a controllable setting, we propose PreDi, a partitioning scheme that disentangles label-space heterogeneity into two orthogonal dimensions, namely class Prevalence and class-set size Disparity, enabling separate analysis of their effects. Guided by the resulting insights, we further propose PreP-WFL (Prevalence-based Personalized Weighted Federated Learning) to adaptively strengthen rare-class representations in low-prevalence scenarios. Extensive experiments show that SSFL consistently outperforms local-only training under both homogeneous and heterogeneous settings. The pronounced heterogeneity in unlabeled data volume is associated with improved representation pre-training, whereas under label-space heterogeneity, prevalence dominates performance and disparity has a smaller effect. PreP-WFL effectively mitigates this degradation, with gains increasing as prevalence decreases. These findings provide a mechanistic basis for characterizing label-space heterogeneity in decentralized recognition systems.

Figures & tables

Figure 1: Overview of the self-supervised federated learning framework with PreP-WFL. In the pre-training stage, masked image reconstruction is adopted as the self-supervised task: at each communication round $t$ , every client $k\in\{1,\dots,K\}$ updates its local encoder $f_{\theta_{k}}$ and decoder $g_{\phi_{k}}$ on the unlabeled dataset $\mathcal{D}^{u}_{k}$ , and uploads the parameters $(\theta^{t}_{k},\phi^{t}_{k})$ to the server. The server performs FedAvg to obtain a global encoder $(f_{\theta^{(*)}},g_{\phi^{(*)}})$ and broadcasts it to all clients. In the fine-tuning stage, the final global encoder $f_{\theta^{(*)}}$ initializes the local encoders, and a classifier head $h_{\psi_{k}}$ is attached on each client and trained on its labeled dataset $\mathcal{D}^{\ell}_{k}$ . The PreP - weight module (red box) performs a one-shot aggregation of the clients’ label sets ${\mathcal{C}_{k}}$ (label IDs only, no images) to compute the global class prevalence $\rho_{c}$ and the corresponding weights $w_{c}$ . These weights are then broadcast to all clients and incorporated into a weighted cross-entropy loss during local fine-tuning, thereby emphasizing rare, site-specific taxa in the federated model.

Figure 2: Overview of the dataset reconstruction pipeline. From the original dataset, we first derive the filtered labeled set $\mathcal{D}^{f}$ , which is partitioned into training and test sets $\mathcal{D}^{\mathrm{train}}$ and $\mathcal{D}^{\mathrm{test}}$ ( $80/20$ ). The remaining taxa form $\mathcal{D}^{r}$ . From $\mathcal{D}^{\mathrm{train}}$ , we sample balanced (50 images per class) labeled subsets $\{\mathcal{D}^{\ell}_{k}\}_{k=1}^{4}$ to serve as the labeled dataset on each client for downstream fine-tuning. The unlabeled pool $\mathcal{D}^{u}$ is constructed by combining $\mathcal{D}^{\mathrm{train}}$ with $\mathcal{D}^{r}$ and discarding all labels.

Splits	Client 1	Client 2	Client 3	Client 4
$\text{Split}^{u}_{\text{IID}}$	18,115	18,115	18,115	18,116
$\text{Split}^{u}_{1}(\alpha=1.0)$	19,988	23,405	14,258	14,810
$\text{Split}^{u}_{2}(\alpha=0.5)$	18,696	32,209	10,653	10,903
$\text{Split}^{u}_{3}(\alpha=0.2)$	22,558	35,750	6,546	7,607
$\text{Split}^{u}_{4}(\alpha=0.1)$	18,982	39,156	6,975	7,348

Table 1: Unlabeled data splits. Each cell reports the number of unlabeled samples held by a given client.

Figure 3: Heterogeneity in unlabeled data volume. Bars show the number of unlabeled samples per client for the IID Split

{}^{u}_{\mathrm{IID}}

and the four non-IID partitions Split

{}^{u}_{1}

–Split

{}^{u}_{4}

(corresponding to

\alpha\!\in\!\{1.0,0.5,0.2,0.1\}

). See also Table 1 .

Figure 4: Illustration of class prevalence and class-set size disparity. The x-axis indicates prevalence (low to high), the y-axis indicates disparity (low to high), and each colored circle denotes a class (A–D).

Figure 5: Heatmap of the number of classes per client under different labeled-data splits. The x-axis enumerates the IID split $\text{Split}^{\ell}_{\mathrm{IID}}$ and PreDi non-IID splits $\text{Split}^{\ell}_{\bar{\rho}^{*},\sigma^{*}}$ obtained by varying $\bar{\rho}^{*}\in\{3.5,3.0,2.5,2.0,1.5\}$ and $\sigma^{*}\in\{0.0,1.0,2.0,3.0\}$ . The y-axis lists the clients, and each cell indicates the number of classes assigned to that client under the corresponding split.

Figure 6: SSFL performance under IID settings (see also Table 2 ).

	Method	Pre-train data	Fine-tune data	Test Accuracy (%)
Client 1	SSL	$\mathcal{D}^{u}_{1}$ in $\mathrm{Split}^{u}_{\mathrm{IID}}$	$\mathcal{D}^{\ell}_{1}$ in $\mathrm{Split}^{\ell}_{\mathrm{IID}}$	77.57
		ImageNet		78.85
Client 2	SSL	$\mathcal{D}^{u}_{2}$ in $\mathrm{Split}^{u}_{\mathrm{IID}}$	$\mathcal{D}^{\ell}_{2}$ in $\mathrm{Split}^{\ell}_{\mathrm{IID}}$	77.58
		ImageNet		78.28
Client 3	SSL	$\mathcal{D}^{u}_{3}$ in $\mathrm{Split}^{u}_{\mathrm{IID}}$	$\mathcal{D}^{\ell}_{3}$ in $\mathrm{Split}^{\ell}_{\mathrm{IID}}$	77.18
		ImageNet		78.84

Table 2: Accuracy results under IID settings.

\mathcal{D}_{k}^{u}

and

\mathcal{D}_{k}^{\ell}

denote the unlabeled and labeled datasets of client

k

, respectively.

Figure 7: Results on unlabeled splits. Bars show mean accuracy and error bars denote standard deviation over three runs. More severe quantity skew in the unlabeled data (smaller

\alpha

) consistently yields higher accuracy.

	Pre-train data	Fine-tune data	Test Accuracy (%)
Client 1	$\mathcal{D}^{u}_{1}$ in $\text{Split}^{u}_{4}$	$\mathcal{D}_{1}^{\ell}$ in $\mathrm{Split}^{\ell}_{\mathrm{IID}}$	79.69
Client 2	$\mathcal{D}^{u}_{2}$ in $\text{Split}^{u}_{4}$	$\mathcal{D}_{2}^{\ell}$ in Split ${}^{\ell}_{\mathrm{IID}}$	81.33
Client 3	$\mathcal{D}^{u}_{3}$ in $\text{Split}^{u}_{4}$	$\mathcal{D}_{3}^{\ell}$ in Split ${}^{\ell}_{\text{IID}}$	70.64
Client 4	$\mathcal{D}^{u}_{4}$ in $\text{Split}^{u}_{4}$	$\mathcal{D}_{4}^{\ell}$ in Split ${}^{\ell}_{\text{IID}}$	73.86
Federated	$\text{Split}^{u}_{4}$	$\mathrm{Split}^{\ell}_{\mathrm{IID}}$	86.33

Table 3: Accuracy results on the most heterogeneous unlabeled split Split

{}^{u}_{4}

. For each client

k

\mathcal{D}^{u}_{k}

and

\mathcal{D}^{\ell}_{k}

denote its local unlabeled and labeled datasets, respectively.

Figure 8: Results on non-IID labeled data.The x-axis shows the target average class prevalence

\bar{\rho}^{*}

and the colored curves correspond to different levels of class-set size disparity

\sigma^{*}

. Accuracy is primarily driven by

\bar{\rho}

, whereas varying

\sigma

leads to only minor differences across settings. See also Table 4 .

	Test Accuracy (%)
	$\bar{\rho}^{*}=3.5$	$\bar{\rho}^{*}=3.0$	$\bar{\rho}^{*}=2.5$	$\bar{\rho}^{*}=2.0$	$\bar{\rho}^{*}=1.5$
$\sigma^{*}=0.0$	82.47	79.56	77.00	75.89	67.82
$\sigma^{*}=1.0$	82.53	79.64	77.47	75.61	67.50
$\sigma^{*}=2.0$	82.34	79.84	77.35	74.43	65.86
$\sigma^{*}=3.0$	82.81	79.67	77.97	74.71	66.94

Table 4: Accuracy results on non-IID labeled data under different (

\bar{\rho}^{*},\sigma^{*}

) configurations.

$\sigma^{*}=0.0$	Test Accuracy (%)
$\sigma^{*}=0.0$	$\bar{\rho}^{*}=3.5$	$\bar{\rho}^{*}=3.0$	$\bar{\rho}^{*}=2.5$	$\bar{\rho}^{*}=2.0$	$\bar{\rho}^{*}=1.5$
Client 1	69.55	59.79	50.18	42.08	11.04
Client 2	71.12	58.74	54.46	43.11	33.71
Client 3	68.63	67.66	50.28	38.30	30.50
Client 4	69.04	62.43	50.54	42.41	30.10
Federated	82.47	79.56	77.00	75.89	67.82

Table 5: Fine-tuning accuracy results under fixed disparity

\sigma^{*}=0.0

for varying prevalence

\bar{\rho}^{*}

$\bar{\rho}^{*}=3.5$	Test Accuracy (%)
$\bar{\rho}^{*}=3.5$	$\sigma^{*}=0.0$	$\sigma^{*}=1.0$	$\sigma^{*}=2.0$	$\sigma^{*}=3.0$
Client 1	69.55	69.37	71.96	72.24
Client 2	71.12	65.85	64.40	62.38
Client 3	68.63	68.71	66.34	66.48
Client 4	69.04	74.16	75.60	77.90
Federated	82.47	82.53	82.34	82.81

Table 6: Fine-tuning accuracy results under fixed prevalence

\bar{\rho}^{*}=3.5

for varying disparity

\sigma^{*}

Figure 9: Accuracy gain (percentage points) of PreP-WFL over the Fed-MAE baseline across $(\bar{\rho}^{*},\sigma^{*})$ configurations of labeled data.

		Test Accuracy (%)
		$\bar{\rho}^{*}=3.5$	$\bar{\rho}^{*}=3.0$	$\bar{\rho}^{*}=2.5$	$\bar{\rho}^{*}=2.0$	$\bar{\rho}^{*}=1.5$
$\sigma^{*}=0.0$	Fed-MAE	82.47	79.56	77.00	75.89	67.82
$\sigma^{*}=0.0$	Fed-MAE $+$ PreP-WFL (ours)	82.62	81.22	79.01	78.11	71.16
$\sigma^{*}=1.0$	Fed-MAE	82.53	79.64	77.47	75.61	67.50
$\sigma^{*}=1.0$	Fed-MAE $+$ PreP-WFL (ours)	82.56	81.00	79.51	77.30	70.21
$\sigma^{*}=2.0$	Fed-MAE	82.34	79.84	77.35	74.43	65.86

Table 7: Accuracy results of Fed-MAE (baseline) and Fed-MAE with PreP-WFL (ours) across

(\bar{\rho}^{*},\sigma^{*})

configurations.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

	Method	Pre-train data	Fine-tune data	Test F1-score
Client 1	SSL	$\mathcal{D}^{u}_{1}$ in $\mathrm{Split}^{u}_{\mathrm{IID}}$	$\mathcal{D}^{\ell}_{1}$ in $\mathrm{Split}^{\ell}_{\mathrm{IID}}$	0.7054
		ImageNet		0.7234
Client 2	SSL	$\mathcal{D}^{u}_{2}$ in $\mathrm{Split}^{u}_{\mathrm{IID}}$	$\mathcal{D}^{\ell}_{2}$ in $\mathrm{Split}^{\ell}_{\mathrm{IID}}$	0.7100
		ImageNet		0.7186
Client 3	SSL	$\mathcal{D}^{u}_{3}$ in $\mathrm{Split}^{u}_{\mathrm{IID}}$	$\mathcal{D}^{\ell}_{3}$ in $\mathrm{Split}^{\ell}_{\mathrm{IID}}$	0.7051
		ImageNet		0.7552

Appendix

Table 8: F1-score results under IID settings.

\mathcal{D}_{k}^{u}

and

\mathcal{D}_{k}^{\ell}

denote the unlabeled and labeled datasets of client

k

, respectively.

Pre-train data	Test F1-score
$\text{Split}^{u}_{\text{IID}}$	0.8008
$\text{Split}^{u}_{1}(\alpha=1.0)$	0.7934
$\text{Split}^{u}_{2}(\alpha=0.5)$	0.7975
$\text{Split}^{u}_{3}(\alpha=0.2)$	0.8063
$\text{Split}^{u}_{4}(\alpha=0.1)$	0.8168

Appendix

Table 9: F1-score results on unlabeled splits.

	Test F1-score
	$\bar{\rho}^{*}=3.5$	$\bar{\rho}^{*}=3.0$	$\bar{\rho}^{*}=2.5$	$\bar{\rho}^{*}=2.0$	$\bar{\rho}^{*}=1.5$
$\sigma^{*}=0.0$	0.7853	0.7570	0.7310	0.7217	0.6277
$\sigma^{*}=1.0$	0.7833	0.7589	0.7357	0.7163	0.6244
$\sigma^{*}=2.0$	0.7827	0.7611	0.7312	0.7059	0.6010
$\sigma^{*}=3.0$	0.7882	0.7579	0.7404	0.7038	0.6158

Appendix

Table 10: F1-score results on non-IID labeled data under different (

\bar{\rho}^{*},\sigma^{*}

) configurations.

		Test F1-score
		$\bar{\rho}^{*}=3.5$	$\bar{\rho}^{*}=3.0$	$\bar{\rho}^{*}=2.5$	$\bar{\rho}^{*}=2.0$	$\bar{\rho}^{*}=1.5$
$\sigma^{*}=0.0$	Fed-MAE	0.7853	0.7570	0.7310	0.7217	0.6277
$\sigma^{*}=0.0$	Fed-MAE $+$ PreP-WFL (ours)	0.7822	0.7696	0.7426	0.7364	0.6611
$\sigma^{*}=1.0$	Fed-MAE	0.7833	0.7589	0.7357	0.7163	0.6244
$\sigma^{*}=1.0$	Fed-MAE $+$ PreP-WFL (ours)	0.7838	0.7687	0.7492	0.7279	0.6473
$\sigma^{*}=2.0$	Fed-MAE	0.7827	0.7611	0.7312	0.7059	0.6010

Appendix

Table 11: F1-score results of Fed-MAE (baseline) and Fed-MAE with PreP-WFL (ours) across

(\bar{\rho}^{*},\sigma^{*})

configurations.

Federated Learning

Federated learning trains shared models across data holders without requiring them to pool their raw data.

See all 943 papers

FedReLa: Imbalanced Federated Learning via Re-Labeling
Jun 24, 2026Guangzheng Hu, Patricia Menéndez, Feng Liu +3Class-Imbalanced Learning Federated Learning
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
FedCC: Towards Addressing Label Distribution Skews in Distillation-Based Federated Learning
Aug 24, 2026Wenxuan Ye, Onur Ayan, Xueli An +1Federated Knowledge Distillation Non-IID Federated Learning
Huawei Heisenberg Research Center, Huawei Technologies Duesseldorf GmbH · TUM School of Computation, Information and Technology, Technical University of Munich
Label Granularity Skew in Federated Learning with Hierarchical Image Classification
Aug 10, 2026Jaeheon Kim, Hokeun Kim, Bong Jun ChoiHierarchical Classification Federated Fine-Tuning
School of Computer Science and Engineering, Soongsil University, Seoul, Republic of Korea · School of Computing and Augmented Intelligence, Arizona State University, Tempe, United States

Jun 24, 2026J/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

Class-Imbalanced Learning Federated Learning Heterogeneous Federated Learning

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

Federated Learning

Explore similar work

FedReLa: Imbalanced Federated Learning via Re-Labeling

FedCC: Towards Addressing Label Distribution Skews in Distillation-Based Federated Learning

Label Granularity Skew in Federated Learning with Hierarchical Image Classification