[Paper Review] Robust Federated Learning in a Heterogeneous Environment

Avishek Ghosh, Justin Hong+2 more|arXiv (Cornell University)|Jun 16, 2019

Machine Learning and Algorithms38 references150 citations

TL;DR

The paper presents a three-stage modular framework for robust federated learning in heterogeneous data with Byzantine workers, including robust clustering of local models and robust distributed optimization, achieving near-optimal error rates. It also analyzes a robust Lloyd-type clustering algorithm and demonstrates substantial practical gains over non-robust methods.

ABSTRACT

We study a recently proposed large-scale distributed learning paradigm, namely Federated Learning, where the worker machines are end users' own devices. Statistical and computational challenges arise in Federated Learning particularly in the presence of heterogeneous data distribution (i.e., data points on different devices belong to different distributions signifying different clusters) and Byzantine machines (i.e., machines that may behave abnormally, or even exhibit arbitrary and potentially adversarial behavior). To address the aforementioned challenges, first we propose a general statistical model for this problem which takes both the cluster structure of the users and the Byzantine machines into account. Then, leveraging the statistical model, we solve the robust heterogeneous Federated Learning problem \emph{optimally}; in particular our algorithm matches the lower bound on the estimation error in dimension and the number of data points. Furthermore, as a by-product, we prove statistical guarantees for an outlier-robust clustering algorithm, which can be considered as the Lloyd algorithm with robust estimation. Finally, we show via synthetic as well as real data experiments that the estimation error obtained by our proposed algorithm is significantly better than the non-Byzantine-robust algorithms; in particular, we gain at least by 53\% and 33\% for synthetic and real data experiments, respectively, in typical settings.

Motivation & Objective

• Motivate robust federated learning under data heterogeneity and Byzantine faults in a practical, scalable setting.
• Develop a modular three-stage framework that combines local ERMs, robust clustering, and robust distributed optimization.
• Provide theoretical guarantees showing near-optimal estimation error in dimension and data samples.
• Introduce and analyze a robust Lloyd-based clustering approach suitable for high dimensions.
• Demonstrate empirical gains over non-robust baselines on synthetic and real datasets.

Proposed method

• Stage I: Each worker computes a local empirical risk minimizer (ERM) erhatw(i) = argmin_w F(i)(w) where F(i)(w) = (1/n) sum_j f(w; x^{i,j}).
• Stage II: The center clusters ERMs into K clusters using an outlier-robust method (threshold-based clustering when n is large, or a robust Lloyd/K-means variant with robust mean estimates).
• Stage III: Within each cluster, run a robust, distributed optimization algorithm (e.g., trimmed mean for low/moderate dimension or iterative filtering for high dimension) to obtain cluster-specific models.
• Theoretical results establish error bounds: ||hatw_i - w_i*|| = O~( (alpha_hat_i * d)/sqrt(n) + d/sqrt(n M_i) ) under certain assumptions, with high probability.
• For the high-dimensional/outlier setting, analysis extends to robust clustering with trimmed K-means and to iterative filtering approaches.
• The clustering analyses include a robust Lloyd-style algorithm with guarantees on misclustered points decaying exponentially with proper initialization.

Experimental results

Research questions

• RQ1 How can federated learning be made robust to Byzantine workers while handling heterogeneous (clustered) data distributions?
• RQ2 Can a modular three-stage framework (local ERMs, robust clustering, robust distributed optimization) achieve near-optimal estimation error under Byzantine and heterogeneity?
• RQ3 What are the theoretical guarantees for robust clustering of local models in the presence of adversaries, and how do they scale with dimension?
• RQ4 How do robust aggregation and clustering choices impact practical performance on synthetic and real datasets?
• RQ5 What are the performance gains over non-robust federated learning methods in heterogeneous environments?

Key findings

• The proposed three-stage algorithm achieves near-optimal error rates: ||hatw_i - w_i*|| = O~( (alpha_hat_i * d)/sqrt(n) + d/sqrt(n M_i) ).
• Robust clustering (trimmed K-means / K-geomedians) reduces misclustering compared to standard K-means, with exponential decay in misclustered points under proper initialization.
• Empirical results show significant improvements over non-robust baselines: estimation error reductions of at least 53% on synthetic data and 33% on Yahoo! Learning to Rank real data experiments.
• A robust Lloyd-style clustering analysis provides theoretical guarantees in high-dimensional settings, making the approach practical for large d.
• The framework remains modular, allowing different robust subroutines in each stage while maintaining overall performance guarantees.
• The Oracle-like comparison indicates the method matches the performance of an oracle that knows non-Byzantine cluster identities in many regimes.