FedTVD: Balancing Data Quality and Quantity for Robust Federated Learning
2026-08-10 • Machine Learning
Machine LearningArtificial IntelligenceComputer Vision and Pattern RecognitionDistributed, Parallel, and Cluster Computing
AI summaryⓘ
The authors identify that in federated learning, models trained across many devices can suffer when the data on each device is very different in size and label distribution. They propose FedTVD, a method that gives less importance to clients whose data labels are very uneven, while still considering how much data each client has. This helps create a better overall model that works well even if some clients have unusual or smaller datasets. Their tests show FedTVD performs better than existing methods on several popular image datasets.
Federated LearningData HeterogeneityLabel Distribution SkewTotal Variation DistanceClient WeightingFedAvgModel AggregationData ImbalanceConvergenceDataset Size
Authors
Radwan Selo, Majid Kundroo, Taehong Kim
Abstract
Federated Learning (FL) enables collaborative model training across distributed client devices while preserving data privacy. However, FL faces significant challenges due to data heterogeneity, particularly in terms of label distribution skewness and variations in dataset sizes, which can lead to biased model updates and hinder convergence. To address this, we propose FedTVD, a novel FL algorithm that weights client contributions during aggregation by considering both data quality and quantity. Unlike traditional FL approaches such as FedAvg, which rely solely on dataset size for client weighting, FedTVD integrates Total Variation Distance (TVD) to measure the divergence between each client's local label distribution and a uniform global distribution. Clients with highly skewed distributions receive lower weights, preventing unbalanced datasets with imbalances from disproportionately influencing the global model. At the same time, dataset size is incorporated to ensure scalability and fairness. This dual-weighting mechanism effectively mitigates the impact of data imbalance, leading to more stable and generalized global models. Experimental results show that FedTVD consistently outperforms state-of-the-art methods across all datasets (FMNIST, CIFAR-10, and CIFAR-100) and all levels of data heterogeneity. Notably, it achieves up to 10.6% improvement over FedAvg on CIFAR-10 under highly skewed data, while maintaining top performance even under moderate and IID settings.