PaperlayerView exact revision on arXiv

Source-linked AI summary

ProxSkip: Yes! Local Gradient Steps Provably Lead to Communication Acceleration! Finally!

Konstantin Mishchenko, Grigory Malinovsky, Sebastian Stich, Peter Richtárik

arXiv:2202.09357v2cs.LGmath.OC

TL;DR

The paper addresses optimization problems in which evaluating a proximable nonsmooth term is expensive, especially when federated-learning communication is costly. ProxSkip skips most proximal evaluations while retaining convergence, yielding fewer prox evaluations and accelerated communication complexity. Its federated-learning analysis has stated scope limitations, including no linear speedup in the number of clients and disregarded client sampling.

  • Problem

    The paper asks whether local gradient-type methods can achieve better than O(κ log 1/ε) communication complexity without homogeneity assumptions or explicit acceleration mechanisms.

  • Method

    ProxSkip randomly evaluates the proximity operator with probability p, allowing local gradient steps between expensive prox evaluations; Scaffnew applies it to federated learning.

  • Results

    ProxSkip has O(κ log 1/ε) iteration complexity and O(√κ log 1/ε) expected prox evaluations, while Scaffnew is optimal in communication rounds.

  • Takeaways & Limitations

    The results show a provable communication advantage for local update methods without similarity assumptions and match the communication-complexity lower bound.

  • Takeaways & Limitations

    The federated-learning analysis does not achieve linear speedup with the number of clients, and the application disregards client sampling.

Abstract

from arXiv · show

We introduce ProxSkip -- a surprisingly simple and provably efficient method for minimizing the sum of a smooth ($f$) and an expensive nonsmooth proximable ($ψ$) function. The canonical approach to solving such problems is via the proximal gradient descent (ProxGD) algorithm, which is based on the evaluation of the gradient of $f$ and the prox operator of $ψ$ in each iteration. In this work we are specifically interested in the regime in which the evaluation of prox is costly relative to the evaluation of the gradient, which is the case in many applications. ProxSkip allows for the expensive prox operator to be skipped in most iterations: while its iteration complexity is $\mathcal{O}\left(κ\log \frac{1}{\varepsilon}\right)$, where $κ$ is the condition number of $f$, the number of prox evaluations is $\mathcal{O}\left(\sqrtκ \log \frac{1}{\varepsilon}\right)$ only. Our main motivation comes from federated learning, where evaluation of the gradient operator corresponds to taking a local GD step independently on all devices, and evaluation of prox corresponds to (expensive) communication in the form of gradient averaging. In this context, ProxSkip offers an effective acceleration of communication complexity. Unlike other local gradient-type methods, such as FedAvg, SCAFFOLD, S-Local-GD and FedLin, whose theoretical communication complexity is worse than, or at best matching, that of vanilla GD in the heterogeneous data regime, we obtain a provable and large improvement without any heterogeneity-bounding assumptions.

1. Introduction

The paper studies composite optimization and distributed learning settings where proximal evaluation or communication is expensive. It asks whether local gradient methods can improve communication complexity without homogeneity assumptions.

  • Paper direction: The paper develops a ProxGD-like approach that targets the cost imbalance between gradient evaluation and expensive proximal evaluation.Its federated-learning application connects local gradient steps with delayed communication under consensus constraints.
  • Composite optimization: ProxGD alternates gradient and proximity steps, but its suitability depends on the proximity operator being cheap to evaluate.The paper focuses on the opposite regime, where computing proxγψ is costly relative to computing ∇f.
  • Expensive proximity operators: Expensive proximity operators can represent difficult projections, including projections onto polyhedral or spectral constraint sets.For consensus constraints, proximity reduces to averaging, which can still be communication-intensive when clients store high-dimensional vectors separately.
  • Federated learning: Federated learning uses multiple local gradient steps and delayed communication to reduce communication costs, but heterogeneous client data causes local-update drift.Existing analyses of methods such as Scaffold, S-Local-GD, and FedLin obtain O(κ log 1/ε) communication complexity in the strongly convex regime.
  • Federated learning: Existing local gradient-type methods do not theoretically improve upon the communication complexity of vanilla GD in the heterogeneous-data regime.This motivates asking whether simple local methods can achieve better than O(κ log 1/ε) communication complexity without explicit acceleration mechanisms.

2. Contributions

The paper introduces ProxSkip, a prox-skipping method with control variates, and applies it to federated learning as Scaffnew. It establishes accelerated communication complexity and extends the approach to stochastic gradients and decentralized training.

  • 2.1. ProxSkip: a general prox skipping algorithm: ProxSkip is designed for problems with expensive proximal operators and evaluates the prox only with probability p.When p=1, it recovers ProxGD; otherwise, the expected prox frequency is once every 1/p iterations.
  • 2.1. ProxSkip: a general prox skipping algorithm: Control variates shift the gradient step and remain constant between consecutive prox calls, stabilizing prox skipping.When a prox step occurs, both the iterate and control variate are adjusted.
  • 2.1. ProxSkip: a general prox skipping algorithm: The control variate converges to the gradient at the optimum, making the optimum a fixed point and allowing prox skips without impacting convergence.Without this adjustment, the gradient at the optimum need not vanish, which can cause drift away from the solution.
  • 2.1. ProxSkip: a general prox skipping algorithm: ProxSkip converges linearly with O(κ log 1/ε) iterations and O(√κ log 1/ε) expected prox evaluations when γ=1/L and p=1/√κ.Here κ=L/µ is the condition number.
  • 2.2. Scaffnew: ProxSkip applied to federated learning: Scaffnew applies ProxSkip to federated learning, interpreting prox evaluation as communication through averaging across clients.The method performs local steps between randomized communication rounds.
  • 2.2. Scaffnew: ProxSkip applied to federated learning: Scaffnew reaches O(√κ log 1/ε) communication complexity without data-similarity or stronger-smoothness assumptions and is optimal in communication rounds.Its average number of local steps per communication round is O(√κ).
  • 2.3. Extensions: The paper extends ProxSkip to stochastic gradients under expected smoothness and to fully decentralized training over communication graphs.The decentralized extension retains insensitivity to data heterogeneity.

3. Theory

The theory analyzes ProxSkip under smooth strong convexity of f and convexity of ψ, using a Lyapunov argument and firm nonexpansiveness. The resulting parameter choices preserve iteration complexity while reducing prox evaluations.

  • 3.1. Assumptions: Under L-smooth, µ-strongly convex f and proper, closed, convex ψ, the optimization problem has a unique minimizer x⋆.The minimizer is x⋆:=argmin f(x)+ψ(x).
  • 3.2. Convergence analysis: The analysis uses firm nonexpansiveness of the proximity operator in its first technical Lyapunov bound.That lemma bounds the Lyapunov function using iterate distance and control-variate error; the property is not used elsewhere in the proof.
  • 3.2. Convergence analysis: The control variates converge to h⋆:=∇f(x⋆), while gradient smoothness links convergence of x_t to convergence of the auxiliary sequence w_t.This convergence is part of the analysis of the method’s stabilization mechanism.
  • 3.3. Main theorem: Theorem 3.6 combines two lemmas into a recurrence whose convergence claim follows by unrolling it.The theorem assumes the stated smoothness and convexity conditions, γ≤1/L, and 0<p≤1.
  • 3.5. How often should one skip the prox?: Setting p=1 recovers ProxGD with ζ=1/κ and O(κ log 1/ε) iteration complexity.With p=1, ProxSkip is identical to ProxGD.
  • 3.5. How often should one skip the prox?: For fixed γ, decreasing p to 1/√(γµ) leaves ζ=min{γµ,p^2} unchanged, so prox evaluations can be skipped without rate deterioration.The reduction factor remains unchanged over this range of probabilities.
  • 3.5. How often should one skip the prox?: With γ=1/L and p=1/√κ, ProxSkip requires O(κ log 1/ε) iterations and O(√κ log 1/ε) prox evaluations.The iteration and prox-calculation complexities differ by a factor of √κ.

4. Application to Federated Learning

In federated learning, Scaffnew applies local gradient updates with control-variate corrections and probabilistic communication. Under smooth strong convexity without function-similarity assumptions, it achieves O(κ log 1/ε) iterations and O(√κ log 1/ε) communications.

  • Method: Scaffnew combines local client-model updates, control-variate updates, and probabilistic model averaging.Averaging occurs with probability p in each iteration, while local control variates correct client drift.
  • Assumptions: Each local function fi is assumed L-smooth and µ-strongly convex.The analysis does not require assumptions on similarity among the functions fi.
  • Convergence: O(κ log 1/ε) iterations follow with γ = 1/L and p = 1/√κ.The result is obtained for exact local gradients under the stated smooth strong-convexity assumption.
  • Convergence: O(√κ log 1/ε) communications are required.This improves the communication count relative to the iteration complexity under the same setting.
  • Implications: The method matches a communication-complexity lower bound and establishes an advantage for local updates without similarity assumptions.The paper describes this rate as optimal in communication complexity within the considered setting.

5. Extensions

The paper extends ProxSkip to stochastic gradients and decentralized communication. The resulting rates depend on stochastic-gradient assumptions, client-communication probability, and network connectivity, while the federated-learning analysis lacks linear client-count speedup.

  • Stochastic gradients: The stochastic analysis assumes unbiased gradient estimates and expected smoothness.Expected smoothness covers settings such as subsampling and compression; bounded variance implies it with A = L and C = σ^2 under the stated lemma.
  • Stochastic gradients: Theorem 5.5 provides convergence guarantees for SProxSkip under the stated stepsize and probability conditions.The theorem assumes 0 < γ ≤ 1/A and 0 < p ≤ 1.
  • Decentralized communication: Scaffnew is obtained by applying ProxSkip to the consensus formulation, with local gradient steps replacing prox evaluations between communications.Algorithm 3 uses control-variate-adjusted client updates and skips communication when the coin toss indicates no averaging.
  • Limitations: The federated-learning analysis does not achieve linear speedup in the number of clients.The paper attributes this limitation to the analysis technique and notes that deeper investigation is needed.
  • Decentralized communication: The decentralized rate depends on the spectral gap δ of the mixing matrix and the communication probability p.The analysis assumes symmetric, doubly stochastic, positive semidefinite mixing and uses γ ≤ 1/L and τ ≤ p/γ.
  • Decentralized communication: Choosing p = 1/√(δκ) is optimal for sufficiently well-connected networks, whereas p = 1 is optimal when δ ≤ 1/κ.The resulting rate matches a lower bound in matrix-vector multiplications.

6. Experiments

Experiments on logistic regression compare Scaffnew with local-update baselines in deterministic and stochastic settings. Scaffnew performs best under theoretical parameters, while communication frequency strongly affects convergence.

  • Setup: The experiments use classical logistic regression with regularization λ = 10^-4L on the w8a LIBSVM dataset.Implementations use Python and RAY on a 24-core Intel Xeon workstation.
  • Deterministic setting: Figure 1 compares deterministic communication rounds against local-update methods and LocalGD under tuned and theoretical parameter settings.Its third subplot varies parameter p using the best theoretical stepsize.
  • Stochastic setting: Figure 2 compares stochastic communication rounds against local-update methods and LocalSGD under tuned parameters, theoretical parameters, and different client counts.The comparisons are organized across three subplots corresponding to these settings.
  • Deterministic setting: Scaffold and Scaffnew share the same communication-round convergence rate, which is better than the other methods in the tuned deterministic experiments.LocalGD converges only to a neighborhood of the solution because of data heterogeneity.
  • Deterministic setting: Scaffnew dramatically outperforms other methods with theoretical hyper-parameters because its theory permits large stepsizes.This comparison uses communication rounds as the performance measure.
  • Parameter sensitivity: The best tested local-phase length is 1/p = 300; both 1/p = 100 and 1/p = 1000 cause worse convergence.The predicted choice p = 1/√κ is close to the experimental result, and experiments report linear speedup with more clients.

B.1. Proof of Lemma 3.4

The proof derives a fixed-point relation for the optimum, recalls ProxSkip’s stochastic prox-skipping update, and bounds one-step Lyapunov progress using firm nonexpansiveness and algebraic expansion.

  • Proof strategy: The proof begins from first-order optimality conditions to obtain a fixed-point identity for x⋆.It uses h⋆ := ∇f(x⋆).
  • Algorithmic update: The method updates xt and ht according to Algorithm 1, applying P with probability p and otherwise retaining the gradient-type iterate.The displayed update encodes randomized prox application and skipping.
  • One-step analysis: The proof takes the conditional expectation of the Lyapunov function at time t + 1 over the iteration’s coin toss.This converts the randomized update into a one-step expected-progress relation.
  • One-step analysis: Firm nonexpansiveness of P supplies the key inequality used in the one-step Lyapunov bound.The remaining terms are handled by expanding squared norms and collecting terms.
  • Algebraic reduction: The identity x̂t+1 = wt + γht rewrites the first terms of the bound in the proof’s auxiliary variables.The proof then applies smoothness and strong convexity properties to complete the estimate.
  • Stochastic extension: The same framework extends to stochastic gradients through SProxSkip, whose update replaces ∇f(xt) with gt(xt).The stochastic algorithm retains the randomized prox decision and control-variate update.

C.2. Two lemmas

The section extends two deterministic lemmas to stochastic gradients and derives expectation-based inequalities under unbiasedness and expected smoothness assumptions.

  • Lemma C.1 extends Lemma 3.4 to the stochastic case under Assumptions 3.1 and 3.2, with expectation over Algorithm 4's randomness.
  • Lemma C.2 extends Lemma 3.5 to the stochastic case for any μ ≥ 0 and step size 0 < γ ≤ 1.
  • Taking expectation and applying unbiasedness and expected smoothness yields an inequality combining squared distance and Bregman-divergence terms.
  • The strong-convexity relation μ∥x_t − x⋆∥2 ≤ 2D_f(x⋆, x_t) is substituted into the preceding bound to obtain the desired estimate.

C.3. Proof of Theorem 5.5

Combining the two stochastic lemmas produces a contraction recurrence for the Lyapunov function, with a rate governed by ζ := min{γμ, p^2} and an additive stochastic term.

  • Combining Lemmas C.1 and C.2 gives ζ := min{γμ, p^2} and the expected recurrence E[Ψ_{t+1}] ≤ (1 − ζ)E[Ψ_t] + γ^2C.

C.4. Proof of Corollary 5.6

The parameter choices balance convergence and prox usage: select the smallest nondegrading prox probability and the largest admissible step size, then derive iteration and expected prox-evaluation bounds.

  • The prox probability p is chosen as the smallest value that does not reduce the rate ζ := min{γμ, p^2}.
  • The step size γ is selected as the largest value consistent with the bounds controlling both terms in the error estimate.
  • The resulting substitution yields an iteration-complexity bound.
  • With prox evaluated independently each iteration with probability p, the expected number of prox evaluations equals pT.

D. Decentralized Analysis

The decentralized analysis embeds the method in a broader primal–dual algorithm with randomly skipped prox updates, then establishes convergence under range and parameter assumptions.

  • Algorithm: Algorithm 5 takes a gradient-type step using a dual control variate, flips a Bernoulli coin, and applies the prox update only when the coin succeeds.
  • Algorithm: When the prox update is skipped, the dual variable remains unchanged and the primal iterate is not corrected by the dual difference.
  • Connection to Algorithm 3: Algorithm 3 is a special case of Algorithm 5 with L = (I − W)^1/2 and ψ the indicator function of zero, yielding prox_{τψ*}(y) = y.
  • Connections: For p = 1, Algorithm 5 reduces to a previously studied primal–dual algorithm, including formulations known as PDFP2O and PAPC.
  • Convergence theorem: Theorem D.1 analyzes a Lyapunov function under Assumptions 3.1 and 3.2, the range condition on ∂ψ*, and a rate ζ = min{γμ, pγτλ+ min(LL⊤)}.

D.1. Proof of Theorem 5.7

The proof specializes the proximity operator to communication by choosing an indicator-function penalty, then shows the resulting updates match the distributed algorithm. It verifies the required operator conditions and derives theorem parameters for the convergence bound.

  • Proof of Theorem 5.7: Choosing ψ as the indicator of {0} makes the proximity operator represent the communication step.The indicator is zero at 0 and +∞ elsewhere.
  • Proof of Theorem 5.7: The conjugate ψ∗ is identically zero, so prox_τψ∗(y) = y and ∂ψ∗(y) = {0}.These properties establish the conditions required by Theorem D.1.
  • Proof of Theorem 5.7: Defining h_t = −L^⊤y_t rewrites Algorithm 5’s first step as the first step of Algorithm 3.The second step likewise reduces using prox_τψ∗(y_t + τL x̂_t+1) = y_t + τL x̂_t+1.
  • Proof of Theorem 5.7: Setting L = (I − W)^1/2 gives L^⊤L = I − W, recovering Algorithm 3’s second step in equivalent form.The proof also records min(I − W) = 1 − λ_2(W) = δ and ∥L^⊤L∥ = ∥I − W∥ < 1.
  • Proof of Theorem 5.7: Theorem D.1 applies with y_i,0 = 0, while the optimal dual norm satisfies ∥y_i,∗∥_2 = ∥L†∇f_i(x∗)∥_2 ≤ 1/λ+.This supplies the final bound used in the theorem’s specialization.
Loading 2202.09357v2…