FedTVD: Balancing Data Quality and Quantity for Robust Federated Learning

FedTVD is a novel federated learning (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.

The paper introduces a TVD-based weighting mechanism that assigns adaptive weights to clients based on both label distribution skewness and dataset size. This ensures that clients with heavily skewed data have a moderated influence on the global model, while those with substantial and balanced data are appropriately emphasized, thereby promoting better convergence, improved generalization, and enhanced fairness. Through extensive empirical evaluations across diverse datasets and application scenarios, FedTVD consistently outperforms existing methods, particularly in environments with highly heterogeneous data. By dynamically adjusting client contributions based on data quality and quantity, FedTVD ensures stable convergence and enhanced performance, making it particularly suitable for large-scale, non-IID environments.

The FedTVD algorithm follows the typical client-server paradigm in FL. On the client side, each client trains locally on its private dataset, computes its label distribution P by counting examples per class, constructs a balanced reference distribution Q (uniform across classes), calculates the TVD between P and Q, and sends back the updated model parameters, TVD value, and local sample size to the server. On the server side, the server broadcasts the global model, collects client updates, applies a softmax function to the TVD values to create normalized weights (lower TVD yields higher weight), computes the total number of samples across clients, and combines the softmax-derived weight with the relative data volume using a tunable parameter λ to produce a composite scaling factor. The final global model is aggregated as a weighted sum of client models using these scaling factors.

The convergence analysis shows that FedTVD maintains the same convergence rate as FedAvg under standard convexity and smoothness assumptions. By incorporating distributional information into the aggregation process, it reduces the impact of client drift, which is often the dominant source of error in non-IID FL. The modified client drift term for FedTVD is expected to be smaller than that of FedAvg in typical non-IID scenarios because TVD captures the degree of label imbalance, and lower weights are assigned to clients with highly skewed distributions.

The complexity analysis demonstrates that FedTVD preserves the same overall order of server-side complexity as FedAvg, with only minor constant-factor overhead. The client-side complexity remains dominated by the training phase, as label counting and TVD computation are lightweight operations.

Experimental results show that under extreme data heterogeneity (α = 0.1), FedTVD achieves 58.82% test accuracy on CIFAR-10, surpassing FedAvg (48.24%) and FedDkw (56.30%). Under moderate heterogeneity (α = 0.5 and 1.0), FedTVD maintains its advantage, achieving 74.69% and 76.22% on CIFAR-10 respectively. On CIFAR-100, FedTVD achieves 37.68%, 39.58%, and 39.79% for α = 0.1, 0.5, and 1.0 respectively, outperforming other methods. Even in IID settings, FedTVD maintains competitive accuracy, achieving 78.04% on CIFAR-10 and 40.51% on CIFAR-100.

The paper also analyzes the impact of the number of clients and client participation rate (CPR). With 100 clients and a CPR of 0.1, FedTVD achieves 58.82% accuracy, outperforming FedDkw (56.30%) and FedAvg (48.24%). As CPR increases to 0.5, FedTVD reaches 62.93%, maintaining its lead. With 200 clients, FedTVD achieves 57.39% at CPR 0.1 and 61.12% at CPR 0.5, consistently outperforming other methods.

An ablation study on the weighting parameter λ shows that under extreme non-IID conditions (α = 0.1), even small increases in λ lead to notable accuracy gains—for example, λ = 0.25 improves performance by over 7% compared to λ = 0.00. While λ = 1.00 achieves the highest accuracy in this extreme setting, it performs slightly worse in the moderate skewed setting (α = 1.0). In contrast, λ = 0.50 delivers consistently strong results across both scenarios, making it a robust default choice for diverse FL environments.

The paper concludes that FedTVD advances the state of the art in FL by incorporating a lightweight yet effective metric, the TVD, into the aggregation procedure. By combining TVD-derived weights with sample-based proportions, it captures both distributional alignment and data volume, helping to mitigate the common challenges of non-IID data. This leads to more robust convergence, particularly in heterogeneous settings. FedTVD’s ability to handle outliers and imbalanced data makes it well-suited for real-world applications where client data is naturally diverse. The paper notes that while the approach specifically targets label distribution skew, it does not directly address feature distribution skew, concept drift, or fairness across subgroups defined by non-label features, which are identified as important future directions.

Improvements to AI Systems:

1. Dynamic Client Weighting in Federated Aggregation: Integrate FedTVD’s dual-weighting mechanism (TVD-based label skewness + dataset size) into existing FL frameworks. The improved system can automatically down-weight clients with highly imbalanced local label distributions, reducing model bias and preventing a few skewed clients from dominating the global model—leading to faster convergence and higher accuracy in non-IID environments.

2. Robustness to Data Heterogeneity: Enhance the aggregation layer of any distributed learning system to compute Total Variation Distance between each client’s label distribution and a uniform reference. The improved system can maintain stable performance across extreme skew (e.g., α=0.1), moderate skew, and even IID settings, with up to 10.6% accuracy gains over FedAvg on CIFAR-10 under heavy imbalance.

3. Scalable and Fair Model Training: Incorporate the tunable parameter λ to balance between distributional quality and data volume. The improved system can adaptively prioritize either balanced data or large datasets depending on the deployment scenario, ensuring fairness for clients with small but well-distributed datasets while still leveraging large, diverse datasets—making it suitable for large-scale FL with hundreds of clients and low participation rates.

4. Lightweight Client-Side Preprocessing: Add label counting and TVD computation as a low-cost client-side step before model upload. The improved system can operate with minimal overhead (constant-factor complexity) while providing the server with actionable distributional metadata, enabling real-time weighting without significant communication or computation costs.

5. Improved Convergence in Non-IID Settings: Replace static aggregation weights with TVD-derived softmax weights. The improved system can reduce client drift—a primary error source in non-IID FL—by explicitly penalizing clients whose local distributions deviate most from uniformity, leading to more stable and generalized global models without sacrificing convergence rate.

6. Adaptive Handling of Extreme Imbalance: Use the ablation findings to set λ=0.50 as a default. The improved system can achieve robust performance across both extreme and moderate heterogeneity, avoiding the need for manual tuning per dataset—making it a plug-and-play enhancement for real-world applications where data skew is unknown or varies over time.

7. Outlier and Skewed-Client Mitigation: Automatically identify and down-weight outlier clients based on high TVD scores. The improved system can prevent malicious or unintentionally biased clients from corrupting the global model, improving reliability in federated deployments with heterogeneous data sources (e.g., mobile devices, hospitals, IoT sensors).

8. Performance Under Low Client Participation: Leverage TVD weighting to maintain high accuracy even with a client participation rate of 0.1 or 0.5. The improved system can achieve up to 62.93% accuracy on CIFAR-10 with 100 clients at 0.5 CPR, outperforming FedAvg and FedDkw, making it viable for resource-constrained or intermittently connected networks.

9. Cross-Dataset Generalization: Apply FedTVD’s aggregation logic to any classification task with label distributions. The improved system can consistently outperform state-of-the-art methods on FMNIST, CIFAR-10, and CIFAR-100 across all heterogeneity levels, enabling deployment in diverse domains like medical imaging, finance, or edge analytics without re-engineering.

10. Foundation for Future Extensions: Use FedTVD’s framework as a baseline to extend toward feature distribution skew, concept drift, and subgroup fairness. The improved system can be designed to incorporate additional distributional metrics beyond labels, allowing future iterations to address broader non-IID challenges while retaining the proven TVD-based weighting core.

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.