[Paper Review] Distributionally Robust Federated Averaging

Yuyang Deng, Mohammad Mahdi Kamani+1 more|arXiv (Cornell University)|Feb 25, 2021

Stochastic Gradient Optimization Techniques20 citations

TL;DR

This paper proposes Distributionally Robust Federated Averaging (DRFA), a communication-efficient algorithm for federated learning that achieves distributional robustness by minimizing the worst-case loss across heterogeneous client data distributions. DRFA uses periodic global averaging with adaptive sampling and a novel snapshotting scheme to approximate the history of mixing parameters, enabling provably fast convergence in both convex and nonconvex settings with reduced communication rounds.

ABSTRACT

In this paper, we study communication efficient distributed algorithms for distributionally robust federated learning via periodic averaging with adaptive sampling. In contrast to standard empirical risk minimization, due to the minimax structure of the underlying optimization problem, a key difficulty arises from the fact that the global parameter that controls the mixture of local losses can only be updated infrequently on the global stage. To compensate for this, we propose a Distributionally Robust Federated Averaging (DRFA) algorithm that employs a novel snapshotting scheme to approximate the accumulation of history gradients of the mixing parameter. We analyze the convergence rate of DRFA in both convex-linear and nonconvex-linear settings. We also generalize the proposed idea to objectives with regularization on the mixture parameter and propose a proximal variant, dubbed as DRFA-Prox, with provable convergence rates. We also analyze an alternative optimization method for regularized cases in strongly-convex-strongly-concave and non-convex (under PL condition)-strongly-concave settings. To the best of our knowledge, this paper is the first to solve distributionally robust federated learning with reduced communication, and to analyze the efficiency of local descent methods on distributed minimax problems. We give corroborating experimental evidence for our theoretical results in federated learning settings.

Motivation & Objective

• To address the challenge of data heterogeneity in federated learning, where standard FedAvg models generalize poorly on non-iid client data.
• To develop a communication-efficient algorithm that maintains distributional robustness without requiring frequent global updates of the mixing weights.
• To enable robust generalization across diverse client data distributions while minimizing communication overhead through periodic averaging and adaptive sampling.
• To theoretically analyze convergence rates of local descent methods on distributed minimax problems, a gap in prior work.
• To extend the framework to regularized objectives via a proximal variant, DRFA-Prox, with provable convergence guarantees.

Proposed method

• Proposes DRFA, a federated algorithm that periodically synchronizes local models and updates the global mixing parameter λ only at synchronization rounds.
• Introduces a snapshotting scheme to approximate the accumulation of history gradients for the mixing parameter λ, compensating for infrequent updates.
• Employs adaptive sampling of clients based on current λ values to improve convergence and robustness during averaging.
• Analyzes convergence in convex-linear and nonconvex-linear settings, establishing O(1/T) suboptimality bounds.
• Develops DRFA-Prox for regularized objectives using proximal updates, with convergence analysis under strongly-convex-strongly-concave and non-convex (PL condition)-strongly-concave settings.
• Uses novel technical tools to bound the error from delayed λ updates, including geometric decay of λ-estimation error and gradient variance control.

Experimental results

Research questions

• RQ1 Can we design a communication-efficient federated learning algorithm that maintains distributional robustness against data heterogeneity?
• RQ2 How can we achieve provable convergence for local descent methods in distributed minimax optimization with infrequent global updates of the mixing parameter?
• RQ3 What is the convergence rate of DRFA in convex and nonconvex settings under partial client participation and periodic averaging?
• RQ4 How does regularization affect the convergence of distributionally robust federated learning, and can we design a proximal variant with theoretical guarantees?
• RQ5 Can we theoretically justify the use of snapshotting to approximate history gradients of the mixing parameter in a communication-efficient setting?

Key findings

• DRFA achieves an O(1/T) convergence rate for the suboptimality gap, with explicit dependence on communication frequency τ and data heterogeneity.
• The algorithm's convergence is bounded by terms involving the initial error in λ estimation, the condition number κ, and gradient variance, with decay rates dependent on τ.
• For nonconvex settings under the Polyak-Łojasiewicz (PL) condition, DRFA achieves O(1/T) convergence, demonstrating robustness beyond convexity.
• The proximal variant DRFA-Prox achieves O(1/T) convergence under strongly-convex-strongly-concave and non-convex (PL)-strongly-concave settings, extending the framework to regularization.
• Empirical results corroborate the theoretical findings, showing improved robustness and generalization on heterogeneous federated data compared to FedAvg.
• The analysis reveals that communication efficiency is preserved even with infrequent λ updates, thanks to the snapshotting mechanism that captures historical gradient trends.