[Paper Review] Agnostic Federated Learning

Mehryar Mohri, Gary Sivek+1 more|arXiv (Cornell University)|Feb 1, 2019

Privacy-Preserving Technologies in Data38 references259 citations

TL;DR

AFL optimizes a centralized model for any mixture of client distributions, offers data-dependent generalization guarantees, and provides a fast minimax SGD algorithm with convergence proofs, plus empirical validation and fairness insights.

ABSTRACT

A key learning scenario in large-scale applications is that of federated learning, where a centralized model is trained based on data originating from a large number of clients. We argue that, with the existing training and inference, federated models can be biased towards different clients. Instead, we propose a new framework of agnostic federated learning, where the centralized model is optimized for any target distribution formed by a mixture of the client distributions. We further show that this framework naturally yields a notion of fairness. We present data-dependent Rademacher complexity guarantees for learning with this objective, which guide the definition of an algorithm for agnostic federated learning. We also give a fast stochastic optimization algorithm for solving the corresponding optimization problem, for which we prove convergence bounds, assuming a convex loss function and hypothesis set. We further empirically demonstrate the benefits of our approach in several datasets. Beyond federated learning, our framework and algorithm can be of interest to other learning scenarios such as cloud computing, domain adaptation, drifting, and other contexts where the training and test distributions do not coincide.

Motivation & Objective

• Motivate federated learning when the target distribution is unknown or mixtures of client distributions may occur.
• Define the agnostic objective as a worst-case loss over mixture weights across clients.
• Develop data-dependent generalization bounds using weighted Rademacher complexity to guide algorithm design.
• Propose a fast stochastic minimax optimization algorithm with convergence guarantees under convex loss and hypothesis spaces.
• Illustrate benefits through empirical evaluation and discuss extensions to related domains such as cloud learning and domain adaptation.

Proposed method

• Formulate AFL where the target distribution is any mixture of client distributions and define the agnostic loss as a maximum over mixture-induced losses.
• Introduce a data-dependent Rademacher complexity framework to bound the agnostic loss across mixtures.
• Derive a skewness-based regularization term from the theory to guide learning.
• Develop a fast stochastic gradient descent algorithm for the minimax AFL objective with convergence guarantees.
• Provide a convex-optimization-based solution path (Section 5.2) and analyze convergence via stochastic gradient properties.
• Demonstrate the approach empirically on multiple datasets and discuss extensions to clustering clients and other domains.

Experimental results

Research questions

• RQ1 What is the impact of optimizing for a mixture of client distributions versus the standard uniform (aggregate) distribution in federated learning?
• RQ2 How can one quantify and control the generalization risk when test distributions are mixtures of client domains?
• RQ3 Can AFL yield a fairness-oriented training objective that mitigates Good-Intent fairness concerns across protected groups?
• RQ4 What algorithms and guarantees enable scalable optimization for the agnostic objective in federated settings?

Key findings

• AFL can outperform the uniform-distribution solution in terms of agnostic loss by a constant additive term under certain setups.
• The paper provides data-dependent Rademacher complexity bounds for the agnostic objective, guiding algorithm design.
• A minimax stochastic optimization algorithm with convergence guarantees is proposed for convex losses and hypothesis sets.
• A natural notion of good-intent fairness emerges, where the predictor performs similarly across protected subgroups while minimizing the worst-case subgroup loss.
• Experiments demonstrate benefits of AFL on several datasets and highlight applicability to domain adaptation, drifting, and cloud learning scenarios.