D ISTRIBUTED M ETHODS WITH C OMPRESSED C OMMUNICATION FOR S OLVING V ARIATIONAL I NEQUALITIES , WITH T HEORETICAL G UARANTEES
Anonymous authors
Paper under double-blind review
A BSTRACT
Variational inequalities in general and saddle point problems in particular are increasingly relevant in machine learning applications, including adversarial learning,
GANs, transport and robust optimization.
With increasing data and problem sizes necessary to train high performing models across various applications, we need to rely on parallel and distributed computing.
However, in distributed training, communication among the compute nodes is a key bottleneck during training, and this problem is exacerbated for high dimensional and over-parameterized models.
Due to these considerations, it is important to equip existing methods with strategies that would allow to reduce the volume of transmitted information during training while obtaining a model of comparable quality.
In this paper, we present the ﬁrst theoretically grounded distributed methods for solving variational inequalities and saddle point problems using compressed communication: MASHA1 and MASHA2 .
Our theory and methods allow for the use of both unbiased (such as Rand k ; MASHA1 ) and contractive (such as Top k ; MASHA2 ) compressors.
We empirically validate our conclusions using two experimental setups: a standard bilinear min-max problem, and large-scale distributed adversarial training of transformers.
I NTRODUCTION 1. T HE EXPRESSIVE POWER OF VARIATIONAL INEQUALITIES
Due to their abstract mathematical nature and the associated ﬂexibility they offer in modeling various practical problems of interests, variational inequalities (VI) have been an active area of research in applied mathematics for more than half a century (Korpelevich, 1976; Harker & Pang., 1990;
Facchinei & Pang, 2003).
It is well known that VIs can be used to formulate and study convex optimization problems, convex-concave saddle point problems and games, for example, in an elegant unifying mathematical framework (Korpelevich, 1976; Bauschke & Combettes, 2017).
Recently, Gidel et al.
(2019) pointed out that multi-player games can be cast as VIs, and proposed to study mini-max or non-zero-sum games formulations of GANs (Goodfellow et al., 2014) in this fashion.
This allowed them to successfully transfer established insights and well-known techniques from the vast literature on VIs to the study of GANs.
In particular, oscillatory behavior of optimization methods (such as SGD) not originally designed to solve VI problems is well understood in the VI literature, and established tools, such as averaging and extrapolation, can be successfully applied to the training of GANs.
Besides their usefulness in studying GANs and alternative adversarial learning models (Madry et al., 2018), VIs have recently attracted considerable attention of the machine learning community due to their ability to model other situations where the minimization of a single loss function does not sufﬁce, such as auction theory (Syrgkanis et al., 2015) and robust and multi-agent reinforcement learning (Pinto et al., 2017).
In summary, VIs have recently become a potent tool enabling new advances in practical machine learning situations reaching beyond supervised learning where optimization problems and techniques, which can be seen as special instances of VIs and methods for solving them, reign supreme.
1. T RAINING OF SUPERVISED MODELS VIA DISTRIBUTED OPTIMIZATION
On the other hand, in the domain of classical, and hence also much better understood, supervised machine learning characterized by the fact that standard optimization techniques apply and work well, researchers and practitioners face other challenges that are currently beyond the reach of existing
VI methods.
Indeed, the training of modern supervised machine learning models in general, and deep neural networks in particular, is still extremely challenging.
Due to their desire to improve the generalization of deployed models, machine learning engineers need to rely on training datasets of ever increasing sizes and on elaborate over-parametrized models (Arora et al., 2018).
Supporting workloads of such unprecedented magnitudes would be impossible without combining the latest advances in hardware acceleration, distributed systems and distributed algorithm design (Verbraeken et al., 2019).
When training such modern supervised models in a distributed fashion, communication cost is often the bottleneck of the training system, and for this reason, a lot of effort was recently targeted at the design of communication efﬁcient distributed optimization methods (Koneˇcn ý et al., 2016; Smith et al., 2018; Ghosh et al., 2020; Gorbunov et al., 2021).
A particularly successful technique for improving the communication efﬁciency of distributed ﬁrst order optimization methods is communication compression.
The idea behind this technique is rooted in the observation that in practical implementations it often advantageous to communicate messages compressed via (often randomized) lossy compression techniques instead of communicating the full messages (Seide et al., 2014;
Alistarh et al., 2017).
If the number of parallel workers is large enough, the noise introduced by compression is reduced, and training with compressed communication will often lead to the comparable test error while reducing the amount of communicated bits, which results in faster training, both in theory and practice (Mishchenko et al., 2019; Gorbunov et al., 2021).
1. T WO CLASSES OF COMPRESSION OPERATORS
We say that a (possibly) stochastic mapping Q : R d !
R d is an unbiased compression operator if there exists a constant q  1 such that
[Equation]
Further, we say that a stochastic mapping C : R d !
R d is a contractive compression operator if there exists a constant   1 such that
[Equation]
If b is the number of bits needed to represent a single ﬂoat (e.g., b = 32 or b = 64 ), then the number of bits needed to represent a generic vector z 2 R d is k z k bits := bd .
To describe how much a compression operator reduces its input vector on average, we deﬁne the notion of expected density, deﬁned via  :
= 1 bd E k Q
( z ) k bits , where k Q ( z ) k bits denotes the number of bits needed to represent the quantized vector Q ( z ) .
Note that   1 .
For the Rand k operator (Alistarh et al., 2018;
Beznosikov et al., 2020)
we have q = d/k and  = k/d .
1. T OWARDS COMMUNICATION - EFFICIENT DISTRIBUTED METHODS FOR VI S
While classical VI algorithms, such as the extragradient method originally proposed by Korpelevich (1976) and later studied by many authors (Nemirovski, 2004; Juditsky et al., 2008), were not designed to work in a distributed environment, virtually all methods that were (Yuan et al., 2014; Hou et al., 2021; Deng & Mahdavi, 2021; Beznosikov et al., 2021b;c) do not consider the general VI problem, but tackle the special case of saddle point problems only.
Moreover, none of these distributed methods support communication compression, with the exception of the work of Yuan et al.
(2014), which relies on rounding to the nearest integer multiple of a certain quantity.
This compression mechanism does not offer theoretical beneﬁts and does not even lead to convergence to the solution due to the errors introduced through rounding persist and prevent the method from solving the problem.
S UMMARY OF C ONTRIBUTIONS
In this paper, we investigate whether it is possible to design communication-efﬁcient algorithms for solving distributed VI communication compression techniques
[Table]
(1) and (2) from the optimization literature (Seide et al., 2014; Alistarh et al., 2017; Mishchenko et al., 2019; Gorbunov et al., 2021;
Richtárik et al., 2021) and adapting and embedding them into established and efﬁcient methods for solving VIs (Korpelevich, 1976; Nemirovski, 2004; Juditsky et al., 2008; Alacaoglu & Malitsky, 2021).
Whether or not this is possible is an open problem.
In summary, we design the ﬁrst provably communication-efﬁcient algorithms for solving general distributed VI problems (see Section 3, Equation 3) in the deterministic (see (4)) and stochastic (see (5))
regimes, supporting both unbiased ( MASHA1 = Algorithm 1) and contractive ( MASHA2 = Algorithm 2) compressors.
Our methods are explicitly designed to be variance reduced to achieve better theoretical properties and better practical performance.
In Table 1 we give a high level overview of existing methods for VIs, and contrast them with our methods and results.
We now elaborate a bit more:
2. T WO DISTRIBUTED PROBLEMS : DETERMINISTIC AND STOCHASTIC
We study two distributed VI problems: i) deterministic , where the monotone operator F : R d !
R d featured in the VI is the average of M operators { F m } Mm =1 , where M is the number of devices/machines, which can be evaluated in each communication round, and ii) stochastic , where each monotone operator F m : R d !
R d has a ﬁnite-sum structure on its own, and only a single operator in the sum can be evaluated in each iteration.
In contrast to previous works, we study general constrained VIs in the distributed setup (see Section 3), and not merely saddle point problems.
2. T WO NEW METHODS WITH COMPRESSED COMMUNICATION : MASHA1 AND MASHA
We develop two extensions of the extragradient / extrastep method of Korpelevich (1976) to distributed
VIs depending on whether we use unbiased (1) or contractive (2) compressors, since each type of compressor demands a different algorithmic design and a different analysis.
In particular, contractive compressors are notoriously hard to analyze even for optimization problems (Karimireddy et al., 2019;
Richtárik et al., 2021).
Our method based on unbiased compressors is called MASHA1 (Algorithm 1), and our method based on contraction compressors is called MASHA2 (Algorithm 2).
Both are designed to handle the deterministic and also the stochastic setting, and both are enhanced with bespoke variance-reduction techniques for better theoretical and practical performance.
Due to space [Table] restrictions, we only describe MASHA1 in the main body of the paper, and relegate MASHA2 and the associated theory to Appendix B.
2. T HEORETICAL COMPLEXITY RESULTS
We establish a number of theoretical complexity results for our methods, which we summarize in
Table 2.
We consider the strongly convex - strongly concave regime as well as the more general convex - concave regime.
In the ﬁrst case we obtain linear convergence results ( O (log 1 / ✏ ) ) in terms of the distance to solution, and in the latter we obtain fast sublinear convergence results ( O (1 / ✏ ) ) in terms of the gap function.
To get an estimate for the number of information transmitted, one need to multiply the estimates from Table 1 by 1 /q and 1 /  , respectively. Then we get that from the point of view of the transmitted information, MASHA1 is better by a factor p 1 /q + 1 /M in comparison with the original extragradient.
P ROBLEM F ORMULATION AND A SSUMPTIONS
Let us ﬁrst introduce basic notation.
We write h x, y i :=
P di =1 x i y i to denote the standard Euclidean inner product of vectors x, y 2 R d .
This induces ` 2 -norm in R d as usual: k x k :=
p h x, x i .
We also introduce the proximal operator, deﬁned as prox g ( z ) :
= arg min u 2 Z { g ( u ) + 12 k u  z k 2 } , which is well deﬁned for proper lower semicontinuous convex functions g : R d !
R [ { + 1 } .
3. P ROBLEM FORMULATION
We study distributed variational inequality (VI) problems of the form
Find z ⇤ 2 Z such that h F ( z ⇤ ) , z  z
⇤ i
+ h ( z )  h (
[Equation]
where Z is a nonempty closed convex subset of R d , F : R d !
R d is an operator with certain favorable properties (e.g., Lipschitzness and monotonicity), and h : R d !
R [ { + 1 } is a proper lower semicontinuous convex function.
We assume that the training data describing F is distributed across M workers/nodes/clients
[Equation]
where F m : R d ! R d for all m 2 { 1 , 2 , . .
, M } , and for some results we further assume that F m is of a ﬁnite-sum structure as well:
[Equation]
3. A SSUMPTIONS
Next, we list two key assumptions - both are standard in the literature on VIs.
Assumption 1 (Lipschitzness in mean)
For all clients m = 1 , 2 , . .
, M , the operator F m ( z ) :
R d !
R d is Lipschitz in mean with constant L m  0 on Z .
That is,
[Equation]
For problem (4), Assumption 1 is to be interpreted to hold with r is equal 1 , i.e., with F m = F m, 1 .
Assumption 2 (Monotonicity and convexity) (SM)
Strong monotonicity/strong convexity.
There exist non-negative constants µ F , µ h such that µ h + µ F > 0 , and the following statements hold:
[Equation]
) Monotonicity/convexity
[Equation]
MASHA1 : H ANDLING U NBIASED C OMPRESSORS
In this section we present only one of our two new algorithms: MASHA1 (Algorithm 1) - the method that relies on unbiased compressors.
Due to lack of space, we include MASHA2 (Algorithm 2) - the method that relies on contractive compressors in the appendix.
Both algorithms are presented for two modes: deterministic (4) and stochastic (4)–(5).
We denote the lines related to the deterministic regime in blue, and in red - to the stochastic one.
Black lines are common to both modes.
4. T HE A LGORITHM
MASHA1 is a modiﬁcation of the extrastep method.
At the beginning of each iteration of Algorithm 1, each device knows the value of F ( w k ) , hence it can be calculated ¯ z k locally.
¯ z k is a convex combination of z k and w k with momentum parameter ⌧ (typically, ⌧ is close to 1).
Also Algorithm can compute z k +1 / 2 locally.
Then it sends the compressed version of the difference F m ( z k +1 / 2 )
F m ( w k ) to the server, and the server does a reverse compressed broadcast.
As a result, an unbiased estimate of F ( z k +1 / 2 )  F ( w k ) appears on each node.
Also, the nodes receive a bit of information b k .
This bit is generated randomly on the server and is equal to 0 with probability ⌧ .
Then, all devices make a ﬁnal update on z k +1 , and also update the w k +1 =
z k +1 point if the b k = 1 or save it from the previous iteration w k +1 =
w k .
In the case when the point w k +
=
z k +1 , we need to exchange the full values of F m ( w k +1 ) in order that at the beginning of the next iteration the value of F ( w k +1 ) is known to all nodes.
In the stochastic case, Algorithm 1 has the same form, the only thing that changes is that one need to generate a function number (batch number) from 1 to r .
We use a possibly difference compressor on each device and also on the server.
To distinguish between them, we will use the following notation: Q dev m , q dev m ,  dev m and Q serv , q serv , q serv .
Note that, if Q is the identity quantization, i.e., Q ( x ) = x , then MASHA1 is a distributed analogue of the method from Alacaoglu & Malitsky (2021).
It is important to note main differences from minimization problems.
For minimization problems, compressed gradient methods are constructed on the gradient descent or the accelerated gradient
Algorithm
Parameters: Stepsize  > 0 , number of iterations K .
Initialization: Choose z
=
w 0 2 Z .
Server sends to devices z 0 =
w 0 and devices compute F m ( w 0 ) and send to server and get F ( w 0 ) for k = 0 , 1 , 2 , . .
, K  1 do for each device m in parallel do
[Equation]
Generate ⇡ km from { 1 , . .
, r } independently
Compute F m ( z k +1 / 2 ) & send Q dev m ( F m ( z k +1 / 2 )  F m ( w k )) to server
Compute F m, ⇡ km ( z k +1 / 2 ) & send Q dev m ( F m, ⇡ km ( z k +1 / 2 )
F m ( w k )) to server end for for server do
Compute Q serv 
[Equation]
M
P
[Equation]
1 Q dev m ( F m ( z
[Equation]
+1 / 2 )
F m ( w k ))
& send to devices
Compute Q serv 
[Equation]
P m =1 Q dev m ( F m, ⇡ km ( z k +1 / 2 )  F m, ⇡ km ( w k ))
& send to devices
Sends to devices one bit b k : 1 with probability 1  ⌧ , 0 with with probability ⌧ end for for all devices in parallel do
[Equation]
P
[Equation]
)    F ( w k ) ◆ if b k
= 1 then w k
+
=
z k +
Compute F m ( w k +1 ) & send it to server; and get F ( w k +1 ) as a response from server else w k [-+
-]+1 =
w k end if end for end for descent methods. Here, the extragradient method is taken as a basis.
In the experiments, we will see the importance of this fact, i.e., gradient descent type methods will diverge even on simple problems.
The second key difference is that we need not compress F m ( z k +1 / 2 ) itself, but the difference
F m ( z k +1 / 2 )  F ( w k ) .
We will also see the importance of this approach in the experiments.
The later idea is similar to the approach used in the DIANA, VR-DIANA (Mishchenko et al., 2019; Horváth et al., 2019), MARINA (Gorbunov et al., 2021) and EF21 (Richtárik et al., 2021) methods in optimization.
4. T HEORY
We now establish convergence of MASHA1 in both regimes: deterministic and stochastic (ﬁnite-sum).
Our analysis relies on the following Lyapunov function :
[Equation]
[Equation]
This criterion is used in the strongly monotone case.
For the general monotone case, another convergence criterion is used - the gap function :
[Equation]
[Equation] Here we do not take the maximum over the entire set Z (as in the classical version), but over C – a compact subset of Z .
Thus, we can also consider unbounded sets Z .
This is permissible, since such a version of the criterion is valid if the solution z ⇤ lies in C ; for details see the work of Nesterov (2007).
Theorem 1 (Convergence of MASHA1 ) Let Assumption 1 be satisﬁed.
Then, if one of cases from
Assumption 2 is additionally fulﬁlled, the following estimates hold for the iterates of MASHA1 :
• for strongly-monotone/convex case with stepsize 0 <   min
[Equation]
• for monotone/convex case with   p
[Equation]
[Equation]
[Equation]
where
[Equation]
For proof, see Appendix A.
An important issue of convergence is the choice of ⌧ .
If ⌧ = 0 we have the fastest convergence rate, on the other hand, this means that we must send full F m in each iteration (because we update w k ).
This is disadvantageous from a communication point of view.
The next corollaries give the rules for the right choice of ⌧ , as well as the iterative (in deterministic case), oracle (in stochastic) and communication complexities of Algorithm 1 in both cases.
The method without quantization (for example, the ordinary distributed extragradient method) in one iteration transmits the number of information proportional to ⌦ ( bdM ) bits, then here we measure the communication complexity in terms of ⌦ ( bdM ) . In literature about method with compression, it is typical to consider only devices compression (Alistarh et al. ; Mishchenko et al. ; Horváth et al.
Gorbunov et al.
; Beznosikov et al.
(2020)), because it is easier for server to broadcast information. In the next corollary we consider the transfer of information in one direction only (from devices to the server). It means we choose q serv = 1 and  serv = 1 . Corollary 1 (Convergence of MASHA1 in the deterministic case) Let the deterministic problem (3)–(4) be solved by MASHA1 with precision " .
Let Assumption 1, and one of the two cases from
Assumption 2 be satisﬁed.
If we choose
[Equation]
then we have the following estimates for the total # of iterations and the total # of transferred bits from devices to the server:
• in the strongly-monotone/convex case:
[Equation]
• in the monotone/convex case
[Equation]
One can see that our Algorithm 1 can outperform the uncompressed extragradient method.
Let us compare them in the monotone case.
The communication complexity of the extragradient method is O
[Equation]
Let us consider the case when L m =
L for all m 2 { 1 , 2 , . .
, M } and q dev m = q dev .
Then MASHA1 has communication complexity
O ⇣
[Equation]
⇤
[Equation]
p
[Equation]
M
[Equation]
⌘ .
If  (1 + q dev /M ) < 1 , we outperform the standard uncompressed extragradient method.
For example, for if we consider RandK operator  = 1 /q dev .
This means that we have O
[Equation] Corollary 2 (Convergence of MASHA1 in the stochastic case) Let the stochastic problem (3)–(5) be solved by MASHA1 with precision " .
Let Assumption 1 and one of the two cases from Assumptionbe satisﬁed.
If we choose
[Equation]
then we have the following estimates for the total # of oracle calls and the total # of transferred bits:
• in the strongly-monotone/convex case:
[Equation]
i log [Equation] [Equation] ⇤ k
[Equation]
⌘ calls [Equation] C [Equation]
i [Equation] [Equation] z
[Equation]
z ⇤ k
[Equation]
⌘ bits ,
• in the monotone/convex case
[Equation]
bits . In the case of a ﬁnite sum, in addition to the number of transmitted information, it is not iterations are important to us, but the number of calls to the oracle for F m,r .
This is due to the fact that it is r times more expensive to calculate full F m , and the calculation of F m should be avoided.
For deterministic case (when r = 1 ) estimates from Corollary 2 is the same as in Corollary 1.
E XPERIMENTS 5. B ILINEAR S ADDLE P OINT P ROBLEM
We start our experiments with a distributed bilinear problem:
[Equation]
(13) where A m 2 R d ⇥ d , a m , b m 2 R d .
This is a saddle point problem, and F is written as follows:
F ( x, y ) :
= ( r x f ( x, y ); r y f ( x, y )) .
This operator is  strongly monotone and, moreover, all functions F m are ( k A m k 2 +  ) -Lipschitz.
Therefore, such a distributed problem is well suited for the primary comparison of our methods.
We take d = 100 and generate positive deﬁnite matrices A m and vectors a m , b m randomly.
The purpose of the experiment is to understand whether the MASHA1 and MASHA2 methods are superior to those in the literature.
As a comparison, we take QGD with natural dithering Horvath et al.
(2019), classical error feedback with Top 30% compression, as well as an extra step method, each step of which is used with natural rounding. In MASHA1 (Algorithm 1) we also used natural dithering, in MASHA2 (Algorithm 2) – Top30%.
See results in Figure 1.
We see that methods based on gradient descent ( QSG and EF ) diverge. This conﬁrms that one needs to use method speciﬁcally designed for saddle point problems (for example, the extragradient method), and not classical optimization methods.
The much slower convergence of the quantized extragradient method shows the efﬁciency of our approach in which we compress the differences
F m ( z k +1 / 2 )  F ( w k ) .
5. A DVERSARIAL T RAINING OF T RANSFORMERS
We now evaluate how compression performs for variational inequalities (and for saddle point problems, as a special case) in a more practically motivated scenario.
Indeed, saddle point problems (special case of variational inequalities) have sample applications in machine learning (see Appendix D), including adversarial training .
We consider one of these tasks and train a transformer-based masked language model (Vaswani et al., 2017; Devlin et al., 2019; Liu et al., 2019) using a ﬂeet oflow-cost preemptible workers with T4 GPU and low-bandwidth interconnect.
For this task, we use the compute-efﬁcient adversarial training regimen proposed for transformers by Zhu et al.
Liu et al.
Formally, the adversarial formulation of the problem is the min-max problem
[Equation]
where w are the weights of the model, { ( x n , y n ) }
Nn =1 are pairs of the training data, ⇢ is the so-called adversarial noise which introduces a perturbation in the data, and  and  are the regularization parameters.
To make our setup more realistic, we train ALBERT-large with layer sharing (Lan et al., 2020), which was recently shown to be much more communication-efﬁcient during training (Ryabinin et al., 2021; Diskin et al., 2021).
We train our model on a combination of Bookcorpus and Wikipedia datasets with the same optimizer (LAMB) and parameters as in the original paper (Lan et al., 2020), use the adversarial training conﬁguration of Zhu et al.
(2019), and follow system design considerations for preemptible instances (Ryabinin et al., 2021).
In terms of communication, we consider 4 different setups for gradient compression: the “baseline” strategy with uncompressed gradients, full 8-bit quantization (Dettmers, 2015; Lin et al., 2018), mixed 8-bit quantization, and Power compression (Vogels et al., 2019) with rank r =8 .
For mixed 8-bit quantization and Power we only apply compression to gradient tensors with more than 2 16 elements, sending smaller ones uncompressed.
These small tensors represent layer biases and LayerNorm scales (Ba et al., 2016) that collectively amount to  1% of the total gradient, but can be more difﬁcult to compress than regular weight tensors.
Finally, since Power is a biased compression algorithm, we use error feedback (Karimireddy et al., 2019;
Richtárik et al., 2021) with a modiﬁed formulation proposed by Vogels et al.
For all experimental setups, we report learning curves in terms of the model training objective, similarly to (Fedus et al., 2021;
Ryabinin et al., 2021).
To better quantify the differences in training loss, we also evaluate the downstream performance for each model on several popular tasks from (Wang et al., 2018) after each model was trained on approximately 80 billion tokens.
Finally, we measure the communication efﬁciency of each proposed strategy by measuring the average wall time per communication round when all 16 workers are active. [Table] N/A N/A
N/A
N/A
N/A N/A N/A N/A N/A N/A [Table]
[Table]
[Table]
[Table]
[Table]
[Table]
The learning curves in Figure 2 (upper) follow a predictable pattern, with more extreme compression techniques demonstrating slower per-iteration convergence.
One curious exception to that is full 8-bit quantization, which was unable to achieve competitive training loss.
The remaining three setups converge to similar loss values below 2.
Both the baseline and mixed 8-bit compression show similar values in terms of downstream performance, with Power compression showing mild degradation.
But in terms of information transfer time, methods using compression (especially Power) are signiﬁcantly superior to the method without compression.
This makes it possible to use such techniques to increase the training time without sacriﬁcing quality.