[Paper Review] Robust Federated Learning against both Data Heterogeneity and Poisoning Attack via Aggregation Optimization

Yueqi Xie, Weizhong Zhang+5 more|arXiv (Cornell University)|Nov 10, 2022

Privacy-Preserving Technologies in Data4 citations

TL;DR

SmartFL proposes a unified framework for robust federated learning by optimizing global model aggregation through subspace training using a small amount of proxy data. By tuning aggregation weights within the convex hull of client models, it achieves state-of-the-art robustness against both data heterogeneity and poisoning attacks with minimal proxy data and low computational cost.

ABSTRACT

Non-IID data distribution across clients and poisoning attacks are two main challenges in real-world federated learning (FL) systems. While both of them have attracted great research interest with specific strategies developed, no known solution manages to address them in a unified framework. To universally overcome both challenges, we propose SmartFL, a generic approach that optimizes the server-side aggregation process with a small amount of proxy data collected by the service provider itself via a subspace training technique. Specifically, the aggregation weight of each participating client at each round is optimized using the server-collected proxy data, which is essentially the optimization of the global model in the convex hull spanned by client models. Since at each round, the number of tunable parameters optimized on the server side equals the number of participating clients (thus independent of the model size), we are able to train a global model with massive parameters using only a small amount of proxy data (e.g., around one hundred samples). With optimized aggregation, SmartFL ensures robustness against both heterogeneous and malicious clients, which is desirable in real-world FL where either or both problems may occur. We provide theoretical analyses of the convergence and generalization capacity for SmartFL. Empirically, SmartFL achieves state-of-the-art performance on both FL with non-IID data distribution and FL with malicious clients. The source code will be released.

Motivation & Objective

• Address the dual challenge of data heterogeneity and poisoning attacks in real-world federated learning systems.
• Develop a unified framework that simultaneously mitigates both issues without relying on heuristics or task-specific designs.
• Reduce reliance on large proxy datasets by optimizing aggregation in a low-dimensional subspace spanned by client models.
• Ensure stable convergence and generalization in the presence of malicious clients and non-IID data distributions.
• Enable efficient, scalable aggregation with minimal computational overhead by limiting optimization to client model weights only.

Proposed method

• Optimize global model aggregation by finding the optimal convex combination of client models using server-collected proxy data.
• Train the global model within the convex hull of client model parameters, reducing the number of tunable parameters to the number of clients, independent of model size.
• Use a subspace training technique to fit the global model on proxy data while constraining updates to the span of client models.
• Apply proxy data to reweight client contributions dynamically in each communication round, improving robustness to outliers and non-IID data.
• Leverage mode connectivity insights to justify that the subspace contains high-quality global models.
• Integrate knowledge distillation with ensemble logits to extend applicability to heterogeneous model architectures.

Experimental results

Research questions

• RQ1 Can a single aggregation strategy robustly handle both data heterogeneity and poisoning attacks in federated learning?
• RQ2 Can robustness be achieved with minimal proxy data by optimizing within a low-dimensional subspace of client models?
• RQ3 Does subspace-based aggregation outperform full-space fine-tuning in terms of generalization and convergence stability?
• RQ4 How does the method perform under extreme data non-IID distributions and high attack rates?
• RQ5 Can the approach maintain stability and effectiveness when client models have different architectures?

Key findings

• SmartFL achieves state-of-the-art performance on both non-IID data and poisoning attack benchmarks, outperforming FedAVG, Sageflow, FLTrust, and FedDF.
• With only around 100 proxy samples, SmartFL avoids overfitting and maintains high generalization, unlike full-space fine-tuning methods.
• The method demonstrates stable convergence and robustness even under high attack rates (up to 70%) and extreme data heterogeneity (α=0.01).
• In experiments on MNIST and CIFAR-10, SmartFL consistently outperforms baselines across all attack types, including Omniscient and Label Flip attacks.
• On heterogeneous architectures (ResNet-8, MobileNet, ShuffleNet), SmartFL achieves higher ensemble accuracy than FedDF, demonstrating effective knowledge transfer via subspace optimization.
• The subspace-based approach ensures convergence and generalization, with theoretical guarantees supported by empirical validation.