PaperlayerView exact revision on arXiv

Source-linked AI summary

Distributionally Robust Federated Averaging

Yuyang Deng, Mohammad Mahdi Kamani, Mehrdad Mahdavi

arXiv:2102.12660v1cs.LGcs.DCstat.ML

TL;DR

The paper addresses how to perform distributionally robust federated learning when communication is reduced and the mixture parameter is updated only periodically. It proposes DRFA with adaptive sampling and randomized snapshotting, analyzes its convergence across several objective settings, and provides supporting federated experiments. The resulting framework gives a communication-efficient approach to robust federated training, with convergence guarantees and extensions to regularized objectives.

  • Problem

    Distributionally robust federated learning requires updating a global mixture parameter, but frequent communication for that update limits scalability under heterogeneous data.

  • Method

    DRFA uses periodic local averaging, adaptive client sampling, and randomized snapshotting to approximate historical gradients for infrequent mixture-parameter updates.

  • Results

    The paper establishes convergence rates for convex-linear and nonconvex-linear objectives, extends the approach with DRFA-Prox to regularized objectives, and reports corroborating federated experiments.

  • Takeaways & Limitations

    The results support communication-efficient distributionally robust federated training with partial participation and theoretical analysis of distributed minimax optimization.

  • Takeaways & Limitations

    The method’s gradient approximation has variance O(τ^2), requires two communication rounds per synchronization step, and uses smaller rates to compensate for high variance.

Abstract

from arXiv · show

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.

1 Introduction

The paper targets distributionally robust federated learning under heterogeneous data while retaining communication efficiency through periodic averaging and partial participation. DRFA updates the model locally and the mixture parameter periodically, using adaptive sampling and randomized snapshotting to support convergence.

  • Motivation: Federated learning must address communication efficiency, low device participation, and heterogeneous local data shards.FedAvg reduces communication through periodic local updates and partial participation, but does not resolve data heterogeneity.
  • Motivation: FedAvg-style objectives can generalize poorly as local data distributions become more diverse because they assume a common underlying distribution.The paper motivates distributionally robust objectives as an alternative to this shared-distribution assumption.
  • Motivation: Distributionally robust learning minimizes performance over the worst-case combination of empirical local distributions, but prior federated methods update the mixing parameter through communication at every iteration.This communication requirement limits scalability despite the robustness objective.
  • Proposed approach: DRFA combines periodic averaging, partial participation, adaptive sampling based on λ, and periodic server updates of the mixing parameter.The model parameter w is trained locally, while λ is updated only at synchronization rounds.
  • Proposed approach: Randomized snapshotting approximates accumulated historical local gradients for updating λ despite its infrequent synchronization.This mechanism addresses the convergence cost introduced by delaying mixture-parameter updates.
  • Theory and evidence: DRFA establishes O(1/T^3/8) convergence for convex-linear losses and O(1/T^1/8) convergence for nonconvex losses with reduced communication rounds.The paper also proposes DRFA-Prox for regularized objectives and reports experimental evidence over competitive baselines.

2 Related Work

Related work spans communication-efficient federated optimization, distributionally robust objectives for non-IID data, and centralized or decentralized minimax optimization. The paper positions its contribution at their intersection by analyzing local descent methods for distributed minimax learning.

  • Federated averaging: FedAvg and local SGD reduce federated communication through periodic local updates, with later analyses covering nonconvex, heterogeneous, and decentralized settings.Prior work includes convergence results for strongly convex, nonconvex, general convex, and non-IID objectives.
  • Distributionally robust optimization: Distributionally robust optimization addresses imbalanced or non-IID data by formulating learning as a minimax problem.Prior federated work minimizes the maximum combination of empirical local losses to mitigate data heterogeneity.
  • Positioning: The paper studies local descent in a distributed minimax setting that combines federated communication constraints with distributionally robust optimization.This connects the federated averaging and distributionally robust optimization strands reviewed above.
  • Smooth minimax optimization: Smooth minimax optimization includes gradient descent-ascent, extragradient, optimistic methods, accelerated schemes, proximal methods, and variance reduction.Existing results cover nonconvex-concave, strongly-convex-concave, decentralized, and related settings.

3 Distributionally Robust Federated Averaging

DRFA combines periodic local model averaging with adaptive client sampling and infrequent mixture-parameter synchronization for distributionally robust federated learning. Random snapshotting approximates the history-gradient information needed to update the mixture parameter, but introduces variance and requires an additional communication round.

  • Problem and setup: DRFA solves min_w max_λ F(w, λ) with periodic averaging and partial client participation under heterogeneous local data distributions.The objective is linear in the mixture parameter λ, which controls the mixture of local losses.
  • Algorithm structure: The algorithm has two modules: local model updating and periodic mixture-parameter synchronization.Each stage uses τ local updates, followed by server aggregation and a separate λ update.
  • Periodic model averaging: Clients selected according to λ perform τ local SGD steps before sending models for server averaging.Adaptive sampling both addresses scalability and reduces communication relative to full participation.
  • Mixture-parameter synchronization: Random snapshotting estimates accumulated mixture-parameter gradients by averaging selected clients’ models at a randomly chosen local iteration.The server evaluates local losses at the snapshot model to construct a stochastic dual gradient.
  • Communication and variance: The snapshot-gradient estimator has variance O(τ^2), and DRFA requires two communication rounds at each synchronization step.The synchronization gap τ must balance convergence guarantees against total communication cost.

4 Convergence Analysis

The convergence analysis addresses the difficulty of periodic primal-dual updates under heterogeneous local objectives by introducing gradient-dissimilarity measures and deriving rates for convex-linear and nonconvex-linear settings. The rates expose communication–accuracy trade-offs and limitations from dual-gradient variance and client participation.

  • Technical challenge: The minimax analysis is harder than empirical-risk analysis because updating λ changes the minimizer of F(·, λ) between synchronization stages.The analysis introduces weighted gradient dissimilarity to formalize heterogeneity among local distributions.
  • Convex losses: Under convex local losses and linear F in λ, DRFA establishes a convergence rate using synchronization gap τ = T^1/4 and a dual stepsize γ = 1/T^5/8.The theorem assumes smoothness, bounded gradients, and bounded stochastic-gradient variance.
  • Convex losses: DRFA trades an O(T^1/8) convergence rate for O(T^1/4) communication rounds in the convex setting.Using constant τ recovers the O(1/T) rate reported for the compared agnostic federated-learning method, at higher communication cost.
  • Nonconvex losses: For nonconvex local objectives, convergence is measured through the gradient of the 1/(2L)-Moreau envelope of Φ.This measure is used because the primal-dual gap is not meaningful without convexity, and direct gradient analysis is difficult in nonconvex-concave settings.
  • Limitations: The leading convergence term does not exhibit linear speedup in the number of sampled clients, and dual-gradient variance motivates a smaller dual stepsize.The authors identify tightening client-count dependence as future work.

5 DRFA-Prox: Optimizing Regularized Objective

DRFA-Prox extends DRFA to regularized mixture objectives by combining history-gradient approximation with a proximal update for λ. It retains DRFA’s convergence rates, while alternative regularized optimization gives rates under stronger curvature conditions.

  • DRFA-Prox: DRFA-Prox handles regularized mixture objectives by using proximity operations for λ while retaining DRFA’s history-gradient approximation for the unregularized term.The proximal step addresses the λ-dependence introduced by the regularizer g(λ).
  • Convergence guarantees: DRFA-Prox achieves the same convergence rates as DRFA for both convex and nonconvex losses.The convex result matches O(1/T^3/8), while the nonconvex result recovers DRFA’s corresponding rate.
  • Convergence guarantees: The convex regularized analysis uses synchronization gap τ = T^1/4 and step parameter γ = 1/T^5/8.These choices yield the stated convex convergence rate.
  • Limitations: The proximal problem adds computation time that is not reflected in the convergence rate.This is an implementation cost of DRFA-Prox rather than a change to its stated rate.
  • Alternative method: The alternative regularized method obtains an Õ(τ/T) rate in both strongly-convex-strongly-concave and PL-condition-strongly-concave settings.For the strongly-convex-strongly-concave case, τ = 1 recovers Õ(1/T); the PL-condition case is slightly worse than the strongly convex case.

6 Experiments

Experiments evaluate DRFA on heterogeneous federated datasets, emphasizing Fashion MNIST with one class per device. Larger synchronization gaps and DRFA’s periodic averaging reduce communication while preserving worst-distribution accuracy.

  • Experimental setup: Experiments use Fashion MNIST, Adult, and Shakespeare, with Fashion MNIST evaluated using logistic regression across heterogeneous devices.In the Fashion MNIST setup, each of 10 devices accesses only one class of data.
  • Synchronization gap: Larger synchronization gaps require fewer communications and less time to reach 50% worst-distribution accuracy.The comparison uses synchronization gaps τ ∈ {5, 10, 15}.
  • Comparison with baselines: DRFA reaches the same worst-distribution accuracy as AFL, q-FedAvg, and FedAvg with fewer communications and lower wall-clock time.The baseline comparison uses Fashion MNIST, logistic regression, 10 devices, and synchronization gap τ = 10.
  • Comparison with baselines: FedAvg does not reach 50% accuracy even after 300 communication rounds, whereas DRFA matches the worst-distribution level achieved by AFL and q-FedAvg more efficiently.DRFA also performs as well as FedAvg in global training accuracy over 100 communication rounds.
  • Comparison with baselines: q-FedAvg has communication rounds close to AFL but requires substantially more local computation.This contrasts communication efficiency with local computational cost in the baseline comparison.

7 Conclusion

The paper proposes communication-efficient distributionally robust federated training and analyzes local SGD in several distributed minimax settings. Experiments support convergence and distributional robustness, while the authors identify faster convergence and variance reduction as future directions.

  • The analysis covers smooth convex-linear, nonconvex-linear, strongly-convex-strongly-concave, and PL-condition nonconvex-strongly-concave settings.
  • Experiments demonstrate convergence of the method and distributional robustness of the learned model.
  • Future work includes improving convergence rates and investigating variance reduction for faster mixing-parameter updates.

Broader Impact

The work advocates distributionally robust federated learning to handle highly heterogeneous client data. Its objective seeks a model with reasonably good performance for everyone rather than benefiting only a subset of clients.

  • The proposed algorithm minimizes the maximum loss over worst-case distributions across clients’ data.
  • The approach targets reasonably good performance for all users when data distributions are highly heterogeneous.
  • The stated impact includes preserving user privacy while training a high-quality model.

A Additional Experiments

Additional experiments evaluate DRFA on heterogeneous Adult and Shakespeare federated datasets and compare it with AFL, q-FedAvg, and FedAvg. The reported results emphasize communication efficiency, worst-case accuracy, client-performance variance, and convergence under several theoretical settings.

  • Adult dataset: gender partition: On Adult data partitioned by gender across 10 nodes, DRFA reaches the same worst accuracy with fewer communication rounds than AFL and q-FedAvg.The experiment uses logistic regression; AFL has lower computational cost, while FedAvg does not reach the same worst accuracy after equal communications.
  • Adult dataset: race partition: On Adult data partitioned by race, DRFA reaches the same worst-case accuracy with fewer communications than AFL and q-FedAvg.The setting uses a two-hidden-layer MLP, and DRFA reduces client-performance variance more effectively than q-FedAvg.
  • Shakespeare dataset: On Shakespeare with 100 clients and an RNN, DRFA and FedAvg outperform AFL and q-FedAvg in communication efficiency.AFL reaches the same level with lower computation cost, but requires more communication for average performance than FedAvg and DRFA.
  • Alternative algorithm: DRFA-GA updates the mixing parameter through full-batch gradient ascent after clients evaluate the current global model on local data.
  • Convergence theory: In the strongly-convex-strongly-concave case, suitable synchronization choices recover centralized-style speedup behavior and can yield linear speedup in sampled workers.The convergence discussion also states that data heterogeneity slows the rate without affecting its dominating term.
  • Convergence theory: Under the generalized PL and strongly concave setting, the convergence rate is slightly worse than in the strongly-convex-strongly-concave case but can still achieve linear worker speedup.
  • Convergence theory: The PL-condition analysis requires a large sampled-client count, which the authors identify as an imperfection of the analysis.
  • Technical analysis: The variance analysis shows quadratic dependence on the synchronization gap τ and possible linear speedup in the number of sampled workers.

C.3 Proof of Technical Lemmas

The technical lemmas establish unbiasedness and variance bounds for stochastic updates, control local-model deviations, and combine primal and dual inequalities through telescoping arguments. The proof analysis highlights synchronization-gap variance and the resulting need to choose τ carefully.

  • Gradient bounds: The stochastic primal and dual gradients are unbiased with bounded variance under the sampling rules.
  • Model deviation: A general bounded-deviation lemma controls the expected squared distance between local and virtual global models for DRFA, DRFA-Prox, and DRFA-GA.
  • Proof construction: The analysis uses telescoping sums across local iterations and communication stages to obtain the convergence bounds.
  • Primal-dual analysis: The proof bounds primal and dual one-iteration terms before combining them and taking the minimax extrema.
  • Dual tracking: Because the optimal dual variable changes as the primal variable updates, the proof groups global stages to enable telescoping within each group.

D.2 Proof of Useful Lemmas

This section develops lemmas and proof steps establishing convergence guarantees for DRFA, DRFA-Prox, and DRFA-GA under progressively stronger structural assumptions. The analyses combine smoothness, convexity or concavity, proximal arguments, and recursive bounds.

  • DRFA: Choosing η = 1/(4LT^3/4), γ = 1/T^1/2, and τ = T^1/4 yields the stated DRFA convergence expression.The resulting bound includes terms scaling with T^-1/4 and T^-1/8, alongside problem-dependent constants.
  • DRFA-Prox: For DRFA-Prox, Lemmas 10–11 control proximal dual updates and snapshot-gradient errors in convex and nonconvex analyses.The proof bounds changes in λ and the discrepancy between averaged and instantaneous gradient quantities.
  • DRFA-Prox: Choosing η = 1/(4LT^3/4), γ = 1/T^1/2, and τ = T^1/4 gives the stated DRFA-Prox convergence rate.The bound contains T^-1/4 and T^-1/8 terms together with constants involving L, G_w, λ, and D_Λ.
  • DRFA-GA: Under strong convexity, strong concavity, or a generalized PL condition, DRFA-GA analysis uses properties of Φ and recursive bounds on λ and projected iterates.The proof invokes smoothness, strong concavity, projection inequalities, and decreasing optimality gaps.

F.3 Proof of Theorem 5

This section proves Theorem 5 by combining bounds for the primal objective and the mixing-parameter optimality gap. The argument relies on smoothness, strong convexity, strong concavity, and recursive control of delayed λ updates.

  • Theorem 5: Theorem 5 analysis combines a bound for DRFA-GA’s primal iterates with a recursive bound for the mixing parameter λ.The proof plugs Lemma 14 into the recursion and unrolls it before applying telescoping sums.
  • Mixing-parameter analysis: The λ-analysis controls the discrepancy between delayed λ updates and λ∗(w(t)) using the κ-Lipschitz property of λ∗(·).This connects changes in the primal iterate to the optimality gap of the mixing parameter.
  • Final bound: The final theorem bound is obtained after substituting the recursive λ estimate and simplifying logarithmic terms.The displayed result includes factors involving µT and log T.
Loading 2102.12660v1…