GTAFL: Addressing Test-Agnostic Long-Tailed Federated Learning in Medical Image

Guangyu Chen; Jingyun Zeng; Hui Jiang

cs.LGNov 5, 2026

GTAFL: Addressing Test-Agnostic Long-Tailed Federated Learning in Medical Image

Authors: Guangyu Chen, Jingyun Zeng, Hui Jiang

Organizations: College of Engineering, Northeastern University, Boston, MA 02115, USA · College of Computer Science, Beijing Institute of Technology, Zhuhai, China · School of Computer Science and Engineering, Huizhou University, Huizhou, China

Abstract

In healthcare, protecting patient privacy is crucial due to the sensitivity of medical data and its extensive accessibility. Federated Learning (FL) offers a decentralized and privacy-preserving training paradigm, making it an ideal solution for healthcare applications. A critical challenge in healthcare is that real-world medical data often exhibits long-tailed distributions in local and global views. Existing methods addressing long-tailed FL problem typically assume that the model will be evaluated on uniform test data distribution. However, practical test data in healthcare systems is often agnostic and unpredictable, leading to potential model failures in real-world scenarios. In this paper, we introduce a novel task termed Test-Agnostic Long-Tailed Federated Learning and propose GTAFL, a comprehensive framework to address this challenge. During the training stage, GTAFL employs adaptive re-sampling, expert classifier retraining, and self-supervised learning to correct biased classifiers and distorted feature spaces caused by long-tailed training distributions. During the inference stage, an ensemble mechanism combines retrained expert classifiers to handle test data with unknown distributions. Extensive experiments on CIFAR10 and two medical datasets manifest that our framework outperforms other state-of-the-art methods. The code is available at https://github.com/cgy-lavine/GTAFL

Figures & tables

Fig. 1: The illustration of two practical challenges: class imbalance and test agnostics. Existing method is only focused on uniform tests in federated settings.

Symbol	Description
$K$	Number of federated clients
$C$	Number of classes
$\mathbf{w}=\{\mathbf{u},\mathbf{v}\}$	Global model parameters
$\mathbf{u}^{t}$	Feature extractor parameters at round $t$
$\mathbf{v}^{t}$	Original classifier parameters at round $t$
$\hat{\mathbf{v}}^{t}$	Expert classifiers parameters at round $t$

TABLE I: Notation Table

Fig. 2: CReFF performance degradation with reduced training data size across different test distributions. Performance drops significantly as training data decreases from 100% to 10%, with Backward50 test showing the most severe degradation

Fig. 3: Expert classifier specialization demonstrates that each classifier achieves optimal performance on its corresponding test distribution while maintaining competitive accuracy across other test distribution.

Fig. 4: GTAFL framework workflow: (1) Clients receive global model and expert parameters, then perform local training with GI-GDS sampling and gradient computation. (2) Clients transmit class-wise gradients and model parameters to the server. (3) Server performs aggregation, expert retraining using optimized global federated features, and self-supervised refinement with unlabeled data. (4) The server sends the updated global model and expert classifiers parameters to clients. (5) Expert ensemble inference for test-agnostic prediction.

Dataset	Modality	Task	# of TrainSet	# of Testset	# of Classes	Imbalance Ratio
CIFAR10-LT	RGB	Image Recognition	50,000	10,000	10	50
PathMNIST-LT	Colon Pathology	Med-Image Classification	89,996	7,180	9	50
OCTMNIST-LT	Retinal OCT	Med-Image Classification	97,477	1,000	4	50

TABLE II: Summary of datasets used in experiments.

CIFAR10-LT
Method	Forward-LT					Uni.	Backward-LT					Avg.
	50	25	10	5	2	1	2	5	10	25	50
FedAvg	42.15	37.57	31.61	26.55	20.20	15.75	11.77	7.48	5.17	2.97	1.93	18.47
FedProx	38.21	34.90	29.73	25.41	19.50	15.29	11.34	7.28	4.95	2.85	1.82	17.39
BalanceFL	39.75	36.02	30.17	25.57	19.97	19.00	18.35	18.40	20.27	22.33	23.58	24.86
CReFF	79.64	78.26	75.24	73.61	71.99	71.60	71.52	72.05	72.38	71.91	72.34	73.69

TABLE III: Performance comparison of federated learning methods on CIFAR10-LT, PathMNIST-LT, and OCTMNIST-LT across various test distributions.

CIFAR10-LT Ablation Study
Components				Forward-LT			Uni.	Backward-LT			Avg.
CR	EN	CL	GS	50	25	10	1	10	25	50
		✓	✓	83.72	79.28	77.42	66.32	54.92	47.20	42.29	64.45
✓				79.64	78.26	75.24	71.60	72.38	71.91	72.34	74.48
✓			✓	79.25	77.98	75.51	70.56	71.35	71.14	72.02	73.97
✓	✓		✓	79.64	78.35	75.76	70.54	71.23	70.73	71.45	73.96

TABLE IV: Ablation study evaluating individual and combined effects of GTAFL components on CIFAR10-LT. Components: CR = Classifier Retraining, GS = Global Sampling, CL = Contrastive Learning, EN = Ensemble.

CIFAR10-LT	Forward 50				Uni.				Backward 50
Method	All	Head	Middle	Tail	All	Head	Middle	Tail	All	Head	Middle	Tail
FedAvg	42.15	56.83	0.53	0.00	15.75	52.50	0.00	0.00	1.93	55.67	0.00	0.00
FedProx	38.21	51.69	0.00	0.00	15.29	50.97	0.00	0.00	1.82	52.58	0.00	0.00
BalanceFL	39.75	53.73	0.00	0.59	19.00	31.63	9.07	16.98	23.58	19.59	0.00	27.39
CReFF	79.64	86.11	62.79	56.47	71.60	78.23	63.50	72.70	72.34	75.26	57.06	74.58
FedGraB	73.20	87.08	40.79	11.18	41.07	78.57	40.33	13.50	12.88	72.16	39.34	6.33

TABLE V: Top-1 test accuracy (%) achieved by compared FL methods and CReFF on CIFAR10-LT Forward 50, Uniform, and Backward 50 settings.

Fig. 5: Performance comparison of GTAFL and CReFF on varying CIFAR10 data fractions. GTAFL demonstrates greater robustness as the data size decreases. (FW50 = Forward50, Uni. = Uniform, BW50 = Backward50). When the data fraction reduces to 10% of original dataset, GTAFL outperforms CReFF by nearly 5%. For more details, please refer to VI-C2 .

Method	Client	Server	GFF Opt	CR	CL
FedAvg	1129s	43s	—	—	—
CReFF	1376s	597s	95s	502s	—
Ours	1712s	4687s	277s	1586s	2699s

TABLE VI: Computational overhead comparison between FedAvg, CReFF, and GTAFL on CIFAR10-LT. Client = Client Time, Server = Server Time, GFF Opt = GFF Optimization, CR = Classifier Retraining, CL = Contrastive Learning.

Method	Uplink	Downlink	Extra vs. FedAvg
FedAvg	7.333 GB	7.333 GB	—
CReFF	7.423 GB	7.349 GB	105 MB
GTAFL (Ours)	7.602 GB	7.380 GB	316 MB

TABLE VII: Communication overhead comparison over 200 training rounds on CIFAR10-LT.

Strategy	FW50	FW25	Uni.	BW25	BW50
Hard Voting (Ours)	0.7456	0.7321	0.6749	0.7684	0.7660
Soft Voting	0.7452	0.7342	0.6755	0.7680	0.7713
Conf.-Weighted	0.7456	0.7349	0.6753	0.7685	0.7717

TABLE VIII: Comparison of ensemble strategies on CIFAR10-LT. FW = Forward, BW = Backward, Conf.-Weighted = Confidence-Weighted.

Federated Learning

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

See all 987 papers

TrustFed: Enabling Trustworthy Medical AI under Data Privacy Constraints
Mar 23, 2026Vagish Kumar, Syed Bahauddin Alam, Souvik ChakrabortyFederated Learning Uncertainty Quantification
Department of Applied Mechanics Indian Institute of Technology Delhi · Grainger College of Engineering Nuclear, Plasma & Radiological Engineering Department National Center for Supercomputing Applications, Urbana, IL, USA · Department of Applied Mechanics Yardi School of Artificial Intelligence (ScAI) Grainger College of Engineering Indian Institute of Technology Delhi University of Illinois Urbana-Champaign, Urbana, IL, USA
Aegis: Generative Gradient Masking for Privacy-Preserving Medical Federated Learning
Sep 29, 2026Chaoyu Zhang, Shanghao Shi, Heng Jin +3Privacy-Preserving Federated Learning Model Inversion Attacks
Virginia Tech · Washington University in St. Louis · University of South Florida
One Round Is All You Need: Analytic Federated Learning for Task-Heterogeneous Multi-Label Medical Image Classification
Jul 22, 2026Afsaneh Mahanipour, Hana KhamfroushChest X-Ray Classification Federated Learning Aggregation
Department of Computer Science University of Kentucky Lexington, KY, USA

Mar 23, 2026J/K move · Enter open · S save

TrustFed: Enabling Trustworthy Medical AI under Data Privacy Constraints

Vagish Kumar, Syed Bahauddin Alam, Souvik Chakraborty

Department of Applied Mechanics Indian Institute of Technology Delhi · Grainger College of Engineering Nuclear, Plasma & Radiological Engineering Department National Center for Supercomputing Applications, Urbana, IL, USA · Department of Applied Mechanics Yardi School of Artificial Intelligence (ScAI) Grainger College of Engineering Indian Institute of Technology Delhi University of Illinois Urbana-Champaign, Urbana, IL, USA

Federated Learning Uncertainty Quantification Heterogeneous Federated Learning Conformal Prediction Medical Imaging

Open paper AskPDF arXiv

Protecting patient privacy remains a fundamental barrier to scaling machine learning across healthcare institutions, where centralizing sensitive data is often infeasible due to ethical, legal, and regulatory constraints. Federated learning offers a promising alternative by enabling privacy-preserving, multi-institutional training without sharing raw patient data; however, real-world deployments face severe challenges from data heterogeneity, site-specific biases, and class imbalance, which degrade predictive reliability and render existing uncertainty quantification methods ineffective. Here, we present TrustFed, a federated uncertainty quantification framework that provides distribution-free, finite-sample coverage guarantees under heterogeneous and imbalanced healthcare data, without requiring centralized access. TrustFed introduces a representation-aware client assignment mechanism that leverages internal model representations to enable effective calibration across institutions, along with a soft-nearest threshold aggregation strategy that mitigates assignment uncertainty while producing compact and reliable prediction sets. Using over 430,000 medical images across six clinically distinct imaging modalities, we conduct one of the most comprehensive evaluations of uncertainty-aware federated learning in medical imaging, demonstrating robust coverage guarantees across datasets with diverse class cardinalities and imbalance regimes. By validating TrustFed at this scale and breadth, our study advances uncertainty-aware federated learning from proof-of-concept toward clinically meaningful, modality-agnostic deployment, positioning statistically guaranteed uncertainty as a core requirement for next-generation healthcare AI systems.

Abstract

Federated Learning

Explore similar work

TrustFed: Enabling Trustworthy Medical AI under Data Privacy Constraints

Aegis: Generative Gradient Masking for Privacy-Preserving Medical Federated Learning

One Round Is All You Need: Analytic Federated Learning for Task-Heterogeneous Multi-Label Medical Image Classification