PaperlayerView exact revision on arXiv

Source-linked AI summary

On Convergence of FedProx: Local Dissimilarity Invariant Bounds, Non-smoothness and Beyond

Xiao-Tong Yuan, Ping Li

arXiv:2206.05187v1stat.MLcs.LG

TL;DR

Prior FedProx convergence guarantees relied on non-standard local-dissimilarity assumptions and applied only to smooth losses, while population optimality remained unclear. The paper uses algorithmic stability to analyze FedProx and a minibatch stochastic extension, FedMSPP, for population optimization. The theory gives local-dissimilarity-invariant convergence bounds for smooth and non-smooth problems, with linear speedup in minibatch size and device sampling for smooth FedMSPP.

  • Problem

    Prior FedProx convergence guarantees relied on non-standard local-dissimilarity assumptions and applied only to smooth losses, while population optimality remained unclear.

    Source: §1.1 · p. 4, §2.1 · p. 8

  • Method

    The paper uses algorithmic stability to analyze FedProx and a minibatch stochastic extension, FedMSPP, for population optimization.

    Source: §1 · p. 2, §3 · p. 10

  • Results

    The theory gives local-dissimilarity-invariant convergence bounds for smooth and non-smooth problems, with linear speedup in minibatch size and device sampling for smooth FedMSPP.

    Source: §3.1 · p. 12, §6 · p. 17

  • Takeaways & Limitations

    On all three benchmark datasets, FedMSPP with varying minibatch sizes needs significantly fewer samples than FedProx and FedAvg to reach desired test accuracies.

    Source: §5.3 · p. 17

    Theorem 1’s rate is slower than Li et al. (2020b)’s bound but avoids the unrealistic (B, 0)-LGD and strong regularization conditions on I.

    Source: §2.1 · p. 8

Abstract

from arXiv · show

The FedProx algorithm is a simple yet powerful distributed proximal point optimization method widely used for federated learning (FL) over heterogeneous data. Despite its popularity and remarkable success witnessed in practice, the theoretical understanding of FedProx is largely underinvestigated: the appealing convergence behavior of FedProx is so far characterized under certain non-standard and unrealistic dissimilarity assumptions of local functions, and the results are limited to smooth optimization problems. In order to remedy these deficiencies, we develop a novel local dissimilarity invariant convergence theory for FedProx and its minibatch stochastic extension through the lens of algorithmic stability. As a result, we contribute to derive several new and deeper insights into FedProx for non-convex federated optimization including: 1) convergence guarantees independent on local dissimilarity type conditions; 2) convergence guarantees for non-smooth FL problems; and 3) linear speedup with respect to size of minibatch and number of sampled devices. Our theory for the first time reveals that local dissimilarity and smoothness are not must-have for FedProx to get favorable complexity bounds. Preliminary experimental results on a series of benchmark FL datasets are reported to demonstrate the benefit of minibatching for improving the sample efficiency of FedProx.

1 Introduction

Federated learning must handle heterogeneous local data and partial device participation, while FedProx’s prior convergence theory relied on restrictive assumptions and smooth losses. This work develops stability-based convergence analyses for FedProx and a minibatch extension, covering weaker conditions and reporting preliminary sample-efficiency results.

1 Introduction

  • Federated learning trains a shared model without transmitting raw data, but non-iid local data and partial device participation complicate optimization and convergence analysis.FedProx is studied because of its simplicity, low communication cost, and tolerance to partial participation.Source: §1 · p. 2

1.1 Review of FedProx

  • FedProx’s earlier convergence guarantees relied on a stringent local-dissimilarity condition, covered only smooth losses, and left population optimality unclear.The work also identifies full-batch local updates as potentially costly, especially when data are imbalanced across devices.Source: §1.1 · p. 4

1.2 Our Contributions

  • FedProx’s smooth non-convex convergence analysis gives bounds independent of local-dissimilarity conditions, while its non-smooth analysis is invariant to the number of selected devices.The non-smooth bound is comparable to the smooth bound when participation is low, but slower for large networks when participation is high.Source: §1.2 · pp. 4–5
  • FedMSPP extends FedProx with minibatch stochastic proximal updates and has local-dissimilarity-invariant population convergence bounds.For sufficiently large T, the stated communication complexity is 1/(bIϵ^2), yielding linear speedup with minibatch size b and sampled-device count I.Source: §1.2 · p. 5
  • The authors report that their bounds are comparable to leading rates in identical settings while covering generic non-smooth, non-convex cases.Their comparison also states that the bounds are independent of local dissimilarity and robust to partial participation.Source: §1.2 · p. 6
  • Across all three benchmark datasets, FedMSPP with varying minibatch sizes needs significantly fewer samples than FedProx and FedAvg to reach desired test accuracies.The numerical study is described as preliminary and intended to corroborate the theoretical sample-efficiency finding.Source: §1.2 · p. 7

2 Convergence of FedProx

The section derives FedProx convergence guarantees for smooth and non-smooth federated objectives without requiring local dissimilarity conditions. It also identifies participation-dependent behavior and a proof technique for handling partial participation in the non-smooth analysis.

2.1 Results for Smooth Problems

  • For smooth non-convex objectives, Theorem 1 gives a convergence guarantee under Lipschitz and smoothness assumptions, with a slower rate than Li et al.'s O(1/T) bound but without its stringent dissimilarity and regularization conditions.The analysis also identifies a benefit of device sampling for linear speedup.Source: §2.1 · p. 8

2.2 Results for Non-smooth Problems

  • For non-smooth, weakly convex objectives, Theorem 2 establishes convergence for exactly solved local updates, with the Moreau-envelope view connecting the bound to stationarity.The authors assume a bounded objective gap for the Moreau envelope.Source: §2.2 · p. 9
  • Partial participation makes the randomly aggregated local objectives dependent and their average biased, so the proof uses that small learning rates keep local iterates close to the preceding global iterate.This closeness permits replacing a local iterate with the previous global iterate without introducing too much approximation error.Source: §2.2 · p. 9
  • Theorem 2 is presented as the first federated-learning convergence guarantee for generic non-smooth, weakly convex problems, beyond composite-convex or cardinality-constrained settings.The cited alternatives are FCO and Fed-HT, respectively.Source: §2.2 · p. 9
  • The non-smooth bound does not depend on the number of selected devices; it is comparable to the smooth bound for small participation but slower near-full participation when M is large.At full participation, the non-smooth O(1/√T) bound is inferior to the smooth O(1/T^2/3) rate.Source: §2.2 · p. 10

3 Convergence of FedProx with Stochastic Minibatching

FedMSPP replaces FedProx’s full-batch local proximal oracle with a stochastic minibatch version. Its analysis gives smooth and non-smooth convergence guarantees without local dissimilarity conditions, with linear speedup in minibatch size and device sampling under the stated smooth-problem bound.

3 Convergence of FedProx with Stochastic Minibatching

  • FedMSPP replaces the full-batch local empirical-risk objective with its minibatch counterpart, while vanilla FedProx is recovered with full-data batches.The method uses a stochastic minibatch proximal-point estimator for local updates.Source: §3 · p. 10

3.1 Results for Smooth Problems

  • The smooth, non-convex analysis uses algorithmic stability to show that averaged local directions align in expectation with the global gradient, without local dissimilarity conditions.For T ≥ (bI)^3, the bound gives communication complexity 1/(bIϵ^2), with linear speedup in minibatch size and device sampling.Source: §3.1 · pp. 11–12

3.2 Results for Non-smooth Problems

  • For weakly convex, non-smooth problems, FedMSPP has a convergence result assuming a bounded Moreau-envelope objective gap and exactly solved local updates.The theorem assumes G-Lipschitz, ν-weakly convex losses and chooses ρ < 1/(2ν).Source: §3.2 · p. 12

4 Additional Related Work

Prior FedProx guarantees relied on stringent gradient-similarity assumptions, while other heterogeneous federated methods addressed some—but not all—of the challenges considered here. The paper applies algorithmic stability to establish guarantees for FedProx and FedMSPP across those settings.

Prior approaches

  • Prior FedProx guarantees for convex and non-convex problems required a stringent gradient-similarity assumption, while other approaches addressed heterogeneity through primal-dual optimization or variance reduction.FedPD relaxed the assumption in a primal-dual framework; SCAFFOLD and VRL-SGD use variance reduction to correct local update directions.Source: §4 · p. 12
  • For non-smooth federated problems, FCO addressed composite optimization, but prior methods otherwise required dissimilarity conditions, full participation, or smooth objectives.The paper positions its stability-based analysis as overcoming these shortcomings simultaneously for FedProx.Source: §4 · p. 13

Minibatch stochastic proximal point methods

  • FedMSPP combines FedProx with minibatch stochastic proximal-point updates, whose improved iteration stability and adaptivity had been studied in large-scale learning but not for non-smooth heterogeneous federated problems.The paper also builds its analysis on algorithmic stability, previously used to establish generalization bounds and convergence rates.Source: §4 · p. 13

5 Experimental Results

The experiments compare FedMSPP with FedProx and FedAvg on three heterogeneous federated-learning benchmarks, assessing samples and communication rounds needed to reach target test accuracies. FedMSPP uses fewer samples across all three datasets, while usually requiring only slightly more communication rounds.

5 Experimental Results

  • The preliminary study uses FedAvg as a baseline to assess whether minibatching speeds FedProx convergence while retaining comparable test performance.The comparison covers both sample use and communication rounds.Source: §5 · p. 13, §5.2 · p. 14

5.1 Data and Models

  • The benchmarks are MNIST and FEMNIST image-classification tasks and Sent140 tweet-sentiment classification, distributed across devices with heterogeneous label or account data.The image tasks use two-layer CNNs; Sent140 uses a two-layer LSTM classifier.Source: §5.1 · p. 14

5.2 Implementation Details and Performance Metrics

  • The implementation compares FedProx, FedMSPP, and FedAvg using SGD locally, and measures samples accessed and communication rounds needed to reach dataset-specific target test accuracies.FedMSPP uses three minibatch sizes per dataset; the listed targets are 80%, 90%, and 95% for MNIST; 80%, 85%, and 91% for FEMNIST; and 68%, 70%, and 73% for Sent140.Source: §5.2 · p. 14, Table 3 · p. 14

5.3 Results

  • FedMSPP with varying minibatch sizes needs significantly fewer samples than FedProx and FedAvg on all three datasets, while usually requiring only slightly more communication rounds for comparable accuracy.The authors describe it as a computationally more efficient replacement for FedProx on the considered heterogeneous federated-learning tasks.Source: Figure 1 · p. 15, Figure 2 · p. 16, §5.3 · p. 17

6 Conclusion

The conclusion presents FedProx and FedMSPP convergence guarantees that avoid local-dissimilarity assumptions, cover smooth and non-smooth problems, and provide linear scaling for smooth problems. The preliminaries introduce stability-based tools for generalization analysis.

6 Conclusion

  • FedProx and FedMSPP achieve favorable convergence rates invariant to local dissimilarity, applicable to smooth and non-smooth problems, and linearly scaling with minibatch size and device-sampling ratio for smooth problems.The paper identifies these as three shortcomings addressed together by its theory.Source: §6 · p. 17

A Preliminaries

  • Uniform argument stability is introduced for algorithms trained on datasets, with stability analysis framed around datasets differing in one data point.The text describes stability as a tool for analyzing generalization bounds.Source: A Preliminaries · p. 22
  • The preliminaries use Efron-Stein inequalities to derive gradient-based generalization bounds for uniformly stable algorithms under Lipschitz and smooth loss assumptions.The bounds are stated for losses that are G-Lipschitz and L-smooth.Source: A Preliminaries · pp. 22–23

B.1 Proof of Theorem 1

The proof of Theorem 1 combines bounds for device sampling and inexact local updates with smoothness-based descent, then averages the resulting inequalities across iterations.

B.1 Proof of Theorem 1

  • The proof uses uniform sampling without replacement, independence, and Lipschitz losses to establish bounds for sampled-device updates.The cited proof fragments describe these ingredients but do not provide the bound values.Source: §B.1 · p. 23
  • A lemma quantifies how local-solver precision affects the gradient norm, assuming smooth losses and an inexact FedProx update with ηt < 1/L.The proof then uses Lipschitzness and smoothness to apply the lemma.Source: §B.1 · p. 24
  • The main proof applies smoothness under a step-size restriction, takes conditional expectations, and averages the resulting inequality over T iterations.The derivation also imposes a condition on the inexactness εt.Source: §B.1 · pp. 25–26

B.2 Proof of Theorem 2

The proof of Theorem 2 handles weakly convex losses by using an inexact three-point inequality and the Moreau envelope to obtain a stationarity bound.

B.2 Proof of Theorem 2

  • An inexact extension of the three-point lemma for weakly convex functions underpins the proof of Theorem 2.The proof assumes weak convexity and uses a proximal step with a parameter constrained by that curvature.Source: §B.2 · pp. 26–27
  • For exactly solved FedProx updates, the proof assumes losses are Lipschitz and weakly convex, with ηt < 1/ν.These conditions support the local-update bounds used in the argument.Source: §B.2 · pp. 26–27
  • The proof converts update bounds into a Moreau-envelope bound, then selects a proximal parameter subject to ρ < 1/(2ν) and averages across iterations.The envelope links the proximal-point distance to a gradient measure.Source: §B.2 · pp. 29–30

C.1 Proof of Theorem 3

The proof of Theorem 3 analyzes FedMSPP with randomness from device and minibatch sampling, using smoothness and local-update bounds before averaging across iterations.

C.1 Proof of Theorem 3

  • The proof tracks randomness from partial device participation and minibatch sampling on each selected device.These sources jointly generate the filtration of random iterates used in the analysis.Source: §C.1 · p. 31
  • The preliminary FedMSPP bounds combine independent sampling of data and devices with Lipschitz and smooth loss assumptions.The proof also invokes an analogous local-update lemma and a condition on solver inexactness.Source: §C.1 · pp. 30–32
  • Theorem 3's proof applies smoothness, takes conditional expectations, and averages the resulting inequality over T iterations to obtain its bound.The proof identifies the preceding inequalities as sufficient for the desired bound.Source: §C.1 · pp. 33–35

C.2 Proof of Theorem 4

Theorem 4’s proof concludes that the iterate is close to the previous iterate when the learning rate is small enough, under the stated weak-convexity and oracle conditions.

C.2 Proof of Theorem 4

  • The proof’s stated consequence is that the iterate is close to the previous iterate when the learning rate is small enough.The proof also states that its argument is almost identical to Theorem 2’s proof.Source: §C.2 · p. 35
  • The auxiliary bound assumes G-Lipschitz, ν-weakly convex losses, an exactly solved FedMSPP local update, and η_t < 1/ν.The proof invokes Lemma 12 and Moreau-envelope reasoning in deriving the bound.Source: §C.2 · pp. 35, 36, 38

D Proofs of Preliminary Lemmas

D.1 Proof of Lemma 1

The auxiliary proofs establish uniform stability through single-example replacement and derive Lemma 3 inequalities using stability and smoothness.

D.1 Proof of Lemma 1

  • Replacing one sample and using regularity yields uniform stability for any pair of datasets differing in that sample.The proof invokes the optimality of the replaced-sample solution, Lipschitz continuity of loss, and strong convexity of the regularizer.Source: §D.1 · p. 39

D.2 Proof of Lemma 3

  • For Lemma 3, the proof compares i.i.d. samples differing in one example, averages over indices, and uses uniform stability for the empirical case.Source: §D.2 · p. 40
  • The proof derives its second inequality using loss smoothness and Lemma 2.Source: §D.2 · p. 40
Loading 2206.05187v1…