ProxSkip for Stochastic Variational Inequalities: A Federated Learning Algorithm for Provable Communication Acceleration
Abstract
Recently Mishchenko et al.
[12] proposed and analyzed ProxSkip, a provably efficient method for minimizing the sum of a smooth ( f ) and an expensive nonsmooth proximable ( R ) function (i.e. min x ∈ R d f ( x ) + R ( x ) ).
The main advantage of ProxSkip, is that in the federated learning (FL) setting, offers provably an effective acceleration of communication complexity.
This work extends this approach to the more general regularized variational inequality problems (VIP).
In particular, we propose ProxSkip-VIP algorithm, which generalizes the original ProxSkip framework of [12] to VIP, and we provide convergence guarantees for a class of structured nonmonotone problems.
In the federated learning setting, we explain how our approach achieves acceleration in terms of the communication complexity over existing state-of-the-art FL algorithms.
1. Introduction
Min-max optimization and, more generally, variational inequality problems (VIPs) appear in a wide range of applications in machine learning including Generative Adversarial Networks (GANs) [7], adversarial training of neural networks [10, 15] and distributionally robust learning [16].
Motivated by these applications, in this work we consider the following regularized variational inequality problem (VIP): find x ∗ ∈ R d , such that
[Equation]
where F : R d → R and R : R d → R is the regularizer (a proper lower semicontinuous convex function).
This problem is quite general and covers a wide range of possible problem formulations.
For example, special cases of (1) are the regularized minimization problems [11] and minimax problems [1]:
[Equation]
In this work, we are interested in the situations when operator F is accessible through the calls of unbiased stochastic oracle.
This is natural when F has an expectation form F ( x ) =
E ξ ∼D [ F ξ ( x )] or a finite-sum form F ( x ) = 1 n (cid:80) ni
=1 F i ( x ) .
In this scenario, one of the most popular algorithms for solving (1) is the stochastic proximal method[13]
[Equation]
where prox γR ( x ) ≜ argmin z ∈ R d (cid:110) R ( z ) + 12 γ ∥
z − x ∥ 2 (cid:111) , g t is an unbiased estimator of F ( x t ) (i.e.
E [ g t ] =
F ( x t ) ) and γ > 0 is the step-size of the method.
Typically computing the proximal operator is easy and cheap.
However, in our work, following the approach of Mishchenko et al.
[12], we are interested in the situation when the evaluation of the proximity operator is expensive.
That is, we assume that the computation of prox γR is costly relative to the evaluation of the unbiased estimator g t .
It is in this scenario that ProxSkip-VIP (Alg. 1) thrives, as it skips the evaluation of the proximity operator and it requires its computation only once every few iterations.
In the federated learning setting (see Sec.
4 for more details) ProxSkip-VIP can be interpreted as a new distributed method performing local steps.
In that scenario, ProxSkip-VIP becomes equivalent to the update rule of Algorithm (2) (ProxSkip-VIP-FL) where the computation of the proximity operator becomes equivalent to communications between workers.
Thus, skipping proximity operator’s computation means that the algorithm performs local updates (it skips communication).
See [12] for the full exposition of this connection.
Main Contributions
Our main contributions are summarized below:
• We generalize the ProxSkip framework proposed in [12] for minimization problems into the VIP regime, and proposed the ProxSkip-VIP algorithm.
• We prove that ProxSkip-VIP converges linearly to a neighborhood of the optimal set when problem (1), has µ -quasi-strongly monotone and L -star-cocoercive operator F .
This is a class of structure non-monotone problems.
As a corollary of our results, for the deterministic regime where g t = F ( x t ) , ProxSkip-VIP converges linearly to the exact solution.
• We extend the ProxSkip-VIP method into the federated learning setting, and propose ProxSkipVIP-FL; we show that the algorithm enjoys an improved communication complexity over existing literature of local Stochastic Methods for solving VIPs.
2. Preliminaries
First, we introduce the main settings throughout the paper.
Assumption 1 (Main Settings)
We assume that problem (1) has a unique 1 solution x ∗ and that:
The operator F is µ -quasi-strongly monotone and ℓ -star-cocoercive with µ, L > 0 , i.e.
[Equation]
The function R is a proper lower semicontinuous convex function.
Note that given that an operator F is L ′ -Lipschitz continuous and µ ′ -strongly
monotone, it can be shown that the operator F is ( κL ′ )-star-cocoercive with κ = L ′
/µ The convergence results in this paper will depend on the following operator noise at x ∗ that is finite for any reasonable sampling: σ 2 ≜ Var ( g ( x ∗ ; ξ )) <
+ ∞ .
Regarding the inherent stochasticity, we further use the following expected cocoercivity assumption [5, 9] to characterize the behavior of the operator estimation
Assumption 2 (Expected Cocoercivity)
We assume that stochastic operator g ( x ; ξ ) is such that for all x ∈ R
d there is L > 0 :
[Equation]
See [5, 9] for more details about this assumption and why is weaker among other bounds on he noise.
3. Algorithm: ProxSkip-VIP
In this section, we incorporate the ProxSkip algorithm [12] into our problem (1), and propose the following ProxSkip-VIP algorithm.
Input: Initial point x 0 , parameters γ 1 , γ 2 , γ 3 , p , initial control variate h 0 , number of iterations T 1: for all t = 0 , 1 , ..., T do 2: ˆ x t +1 = x
t − γ 1 ( g ( x t ; ξ t ) − h t ) 3:
Flip a coin θ t , θ t = 1 w.p. p , otherwise4: if θ t = 1 then 5: x t +1 = prox γ 2 R (ˆ x t +1 − γ 2 h t ) 6: else 7: x t +1 =
ˆ x t + 8: end if 9: h t +1 =
h t + γ
( x t +1 − ˆ x t +1 ) 10: end for
Output: x T
Here the key step is the randomized prox-skipping and the control variate h t .
The proximal oracle is called very rarely if p is small, which helps to reduce the computational cost if the proximal oracle is expensive; the introducing of h t helps to stabilize the iterations toward the optimal point.
3.1. Convergence Analysis
Our main theorem on the convergence guarantee is presented below.
Theorem 1 (Convergence of ProxSkip-VIP) Let Assumption 1 and 2 hold, and let γ 1 = γ ∈ (cid:16) 0 , min (cid:110) 1 µ , 12 L (cid:111)(cid:17) , γ 1 = γ 2 p , γ 3 = 1 γ 2 .
Then the iterates of ProxSkip (Alg. 1) satisfy
[Equation]
γµ [Equation] p
[Equation]
where V t ≜ ∥ x t − x ∗ t ∥ 2 + γ 22 ∥
h t − F ( x ∗ t )
∥ 2 . We defer the proof to Appendix B. With the above theorem, we can get the corresponding complexity results in the following, whose proof is deferred to Appendix C.
Corollary 2 With the setting in Theorem 1, if we further set γ ≤ µϵ 2 σ 2 and p = √ γµ , we have
E [ V T ] ≤ ϵ with iteration complexity and the number of calls of the proximal oracle prox ( · ) as
[Equation]
Deterministic ProxSkip:
As a corollary of our results, for the deterministic regime where g t =
F ( x t ) , ProxSkip-VIP converges linearly to the exact solution since σ 2 = 0 .
In this scenario, under the same assumptions with Theorem 1, the iterates of ProxSkip (Alg. 1) satisfy:
[Equation]
In this setting, we get E [ V T ] ≤ ϵ with iteration complexity and number of calls of the proximal oracle prox ( · ) as [Equation] and [Equation]
4. Connection with Federated Learning
Now we specialize the discussion into the federated learning setting, i.e., find x ∗ ∈ R d such that
[Equation]
where F ( x ) ≜ 1 n (cid:80) ni =1 f i ( x ) and f i ( x ) ≜
E ξ i ∼D i [ f i ( x ; ξ i )] , the data ξ i follows an unknown distribution
D i ( i = 1 , 2 , · · · , n ).
For example, the federated learning minimization or minimax problems [12, 14] are special cases of (7).
As mentioned in [13], the problem (7) can be recast into the problem (1) while
[Equation]
where x i ∈ R d , x = ( x 1 , x 2 , · · · , x n ) ∈ R nd , and
[Equation]
Note that prox γR ( x ) =
(¯ x, ¯ x, · · · , ¯ x ) and ¯ x = 1 n (cid:80) ni
=1 x i , which is easy to compute [13].
4.1. Algorithm: ProxSkip-VIP-FL
With the above reformulation (8), we propose the following ProxSkip-VIP-FL algorithm based on
Algorithm 1 for the federated learning problem 7, which is presented below.
Our algorithm is a unified framework which can attain many variants by applying different (local) operator estimator g i ( x i,t ; ξ i,t ) . Different from the centralized setting we discussed in Section 3, the federated learning framework often incurs the heterogeneous setting, i.e., the distributions {D i } i are not identical. 4.2. Convergence Analysis
Let us now present the convergence guarantees of ProxSkip-VIP-FL (Alg. 2). Theorem 3 (Complexity of ProxSkip-VIP-FL) Lets assume the same setting as in Corollary 2. Then ProxSkip-VIP-FL achieves E [ V T ] ≤ ϵ (where V T is defined in Theorem 1), Input variate
1: for all ., T do 2: Server: Flip a t p p workers 3: for workers i ∈ [ n ] in parallel do 4: ˆ x i,t +1 = x
i,t − γ {+
+} g i ( x i,t ; ξ i,t ) − h i,t ) //
Local update with variate 5: if θ t = 1 then 6:
Worker: x i the server 7:
Server: computes x i t
1 n i,t workers //
Communication 8: else 9: x i,t +
=
ˆ x i,t + // Otherwise skip the communication step 10: end if 11: h i,t +1 = h i,t + γ 3 ( x i,t +1 − ˆ x i,t +1 ) 12: end for 13: end for
Output
T • complexity
O (cid:18) max (cid:26) L µ , σµ 2 ϵ (cid:27) ln
(cid:18)ϵ (cid:19)(cid:19) • and communication complexity as
O (cid:32)(cid:115) max (cid:26) L µ , σµ 2 ϵ (cid:27) ln (cid:18)ϵ (cid:19)(cid:33) . We further extend our discussion on the complexity of ProxSkip-VIP-FL regarding the heterogeneous setting in Appendix D.
Comparison with Existing Literature
The above theorem provided the complexity results of
Algorithm 2 in the federated learning setting.
Note that many algorithms in federated learning can be cast into our framework.
For example, in federated minimax problems, a well-known algorithm is the Local SGDA algorithm [6], if we apply their (mini-batch) gradient estimator into our Algorithm 2, our result will avoid the dependence on the condition number κ (when ϵ is small enough), and we can attain an improved communication and iteration complexity for solving (7).
We also compared our Algorithm 2 with existing literature in federated learning, as shown in
Table 1.
We can find that our algorithm extended the problem settings, also outperforms over other algorithms in terms of both the iteration and communication complexities (when ϵ is small enough).
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
Ours (Theorem 3)
SM, LS
[Table]
1 SM: strongly monotone, LS: (Lipschitz) smooth. κ ≜ L/µ , L and µ are the modulus of SM and LS.
The definition of ξ 2 follows that in Appendix D
[Table]
and σ 2 ≜ Var ( g ( x ∗ ; ξ )) < + ∞ (or uniform bound on Var ( g ( · ; ξ )) ).p is the probability of synchronization, by setting ϵ is small enough, we can take p =
O ( √ ϵ )
[Table] which recovers O (1 / √ ϵ ) communication complexity dependence on ϵ in our result.
˜ O ( · ) hides the logarithmic terms.
[Table]
Our algorithm works for quasi-strongly monotone and star-cocoercive, which is more general than the SM and LS setting, note that an L -LS and µ -SM operator can be shown to be ( κL )-star-cocoercive [9].
Supplementary Material Appendix A. Useful Lemmas
Lemma
For any optimal solution x ∗ ∈
X ∗ of (1) , we have
[Equation]
Proof Note that ⟨
[Equation]
next for any x ∈ R d ,
[Equation]
⇐⇒
[Equation]
⇐⇒
[Equation]
which concludes the proof
[Equation]
note that this conclusion indicates that the two parameters in the RHS above should be identical
[Equation]
Lemma 5 (Firm Nonexpansivity of the Proximal Operator [2]) Let f be a proper closed and convex function, then for any x,
y ∈
R
d we have
[Equation]
or equivalently,
[Equation]
The following result is helpful in the proof.
Lemma 6 With Assumption 1 and 2, we have
[Equation]
Proof
[Equation]
which concludes the proof.
Appendix B. Proof of Theorem
Proof Note that
[Equation]
probability p
[Equation]
with probability
[Equation]
and
[Equation]
with probability
[Equation]
For simplicity, we denote P ( x t ) ≜ prox γ 2 R (ˆ x t +1 − γ 2 h t ) , so we have
E ξ t [ V t +1 ] =
[Equation]
next note that x ∗ t = prox γ 2 R ( x ∗ t − γ 2 F ( x ∗ t )) , we have
[Equation]
so by Lemma 5, we have
E ξ t [ V t +1 ]
[Equation]
, (20) let w t ≜ x t − γ 1 g ( x t ; ξ t ) , w
[Equation]
t ≜ x
[Equation]
γ
[Equation]
F ( x
[Equation]
) , (21) recall that γ = γ 1 = γ 2 p , so we have
[Equation]
= (cid:13)(cid:13)
[Equation]
− w
[Equation]
+
[Equation]
x
[Equation]
(cid:1)(cid:13)(cid:13)
[Equation]
(cid:10)
[Equation]
− w
[Equation]
+1 + γ (cid:0) h
[Equation]
F (
[Equation]
) (cid:1) [Equation] h t − F ( x
[Equation]
+1 ) (cid:11) =
(cid:13) (cid:13) w t − w ∗ t
+
(cid:13) (cid:13)
[Equation]
γ
[Equation]
(cid:13) (cid:13) [Equation] t − F ( x
[Equation]
+1 ) (cid:13) (cid:13) 2 , (22) so we have
E
[Equation]
so we also have w ∗ t ≡ x ∗ − γF ( x ∗ ) ≜ w ∗ .
Then by the standard analysis on GDA
[Equation]
we have
[Equation]
take the expectation
[Equation]
we have
[Equation]
) so we have
E ξ t [ V t +1 ] ≤
[Equation]
≤ E ξ t (cid:34)
[Equation]
≤ E ξ t (cid:34) ( [Equation] − [Equation]
γµ
[Equation]
− γL ))
[Equation]
−
[Equation]
recall that γ ≤ 12 L , we have
[Equation]
by taking the full expectation
[Equation]
we have
[Equation]
which concludes the proof
[Equation]
Appendix C. Proof of Corollary
Proof With the above setting, we know that
[Equation]
and
[Equation]
so it is easy to see that by setting
[Equation]
we have
[Equation]
which induces the iteration complexity to be
[Equation]
and the corresponding number of calls to the proximal oracle is
[Equation]
which concludes the proof.
Appendix D. Additional Discussion of Section
In the heterogeneous setting, we often use the following assumption to characterize the difference in heterogeneity.
Assumption 3 (Bounded Heterogeneity)
We assume that g i ( x ; ξ i ) is an unbiased estimator of f i ( x ) for any i ∈ { 1 , 2 , · · · , n } , and ξ 2 i ( x ) ≜ sup x ∈ R d ∥ f i ( x ) − F ( x ) ∥ 2 ≤ ξ 2 ≤ + ∞ .
With the Assumption 3, we can easily get the following result for the stochastic operator g ( x ; ξ ) ≜ ( g 1 ( x ; ξ 1 ) , g 2 ( x ; ξ 2 ) , · · · , g n ( x ; ξ n )) .
Theorem 7 (Local Operator Estimator) Let Assumptions 1 and 2 hold.
Then if the local operator estimator g i ( x ; ξ i ) satisfies Assumption 3, we have
[Equation]
where D ≜ 4 n (cid:80) ni =1 (cid:0) Var ( g i ( x ∗ ; ξ i )) + ξ 2 i ( x ∗ ) (cid:1) ≤ 4( σ 2 + ξ 2 ) . With the above result, it can be easily shown that, if we have the same setting in Theorem 3 and additional Assumption 3, the iteration complexity and communication complexity of ProxSkip-VIPFL are
[Equation]
O
[Equation]
the proof is similar to that of Corollary 2, which we do not repeat here.