Smoothness Matrices Beat Smoothness Constants: Better Communication Compression Techniques for Distributed Optimization
Year: 2022

Abstract:
Large scale distributed optimization has become the default tool for the training of supervised machine learning models with a large number of parameters and training data. Recent advancements in the field provide several mechanisms for speeding up the training, including compressed communication, variance reduction and acceleration. However, none of these methods is capable of exploiting the inherently rich data-dependent smoothness structure of the local losses beyond standard smoothness constants. In this paper, we argue that when training supervised models, smoothness matrices—information-rich generalizations of the ubiquitous smoothness constants—can and should be exploited for further dramatic gains, both in theory and practice. In order to further alleviate the communication burden inherent in distributed optimization, we propose a novel communication sparsification strategy that can take full advantage of the smoothness matrices associated with local losses. To showcase the power of this tool, we describe how our sparsification technique can be adapted to three distributed optimization algorithms—DCGD [Khirirat et al., 2018], DIANA [Mishchenko et al., 2019] and ADIANA [Li et al., 2020]—yielding significant savings in terms of communication complexity. The new methods always outperform the baselines, often dramatically so.

heading: 1 Introduction
text: With the desire to build and train high quality machine learning models comes an increased appetite for larger models, both in terms of the number of parameters encoding them, and in the amount of data required to train them. In the big data regime, the data is partitioned among many parallel machines, which then cooperatively train a single global model, usually orchestrated by a central server. Distributed training is cast as the distributed optimization problem
min x∈Rd
f(x) +R(x), f(x) := 1n n∑ i=1 fi(x), (1)
where d is the number of parameters of model x ∈ Rd, n is the number of machines participating in the training, fi(x) is the loss associated with the data stored on machine i ∈ [n] := {1, 2, . . . , n}, f(x) is the empirical loss, and R(x) is a regularizer. Ample research over the past two decades has shown that first-order methods are highly scalable and as a result are the methods of choice for distributed optimization problems [Liu and Zhang, 2020]. In particular, a substantial amount of work
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
has been devoted to speeding up the training process by developing efficient methods empowered with techniques such as compressed communication, variance reduction and acceleration.
Compressed communication. In distributed training, compute nodes have to communicate with each other, often via a central server, in order to be able to maintain consensus and jointly train a global model. However, communication of the information pertaining to local progress, which is typically contained in gradient(s) distilled from local data, is almost invariably the key bottleneck in distributed training systems [Xu et al., 2020]. One popular way to address this issue is to reduce the number of bits encoding the vector/tensor to be transferred via the help of a lossy compression operator. Numerous unbiased gradient compression operators have been proposed for this purpose, including several types of sparsifications [Wang et al., 2018, Mishchenko et al., 2020, Alistarh et al., 2018] and quantizations [Alistarh et al., 2017, Zhang et al., 2017, Horváth et al., 2019a, Wu et al., 2018].
Variance reduction. A marked issue that needs to be addressed by successful distributed optimization methods has to do with the (potential) “dissimilarity” of the local loss functions f1, . . . , fn, which in turn is due to the heterogeneity of the training data defining these functions. The higher the dissimilarity, the harder it is for the devices to find the minimizer of (1). This issue exists even in the unregularized case (R ≡ 0). Indeed, while in this case 1n ∑ i∇fi(x∗) = 0 if x∗ is a minimizer of f , this does not mean that the individual gradients, ∇f1(x∗), . . . ,∇fn(x∗), are all zero. This shows that local gradient information alone is not enough for any node to “realize” that a solution has been found, which encourages further, in this case unnecessary, iterations. If unaddressed properly, an algorithm is forced to use smaller learning rates, and this leads to unnecessarily slow convergence. On the other hand, when a fixed learning rate is used, the rate is fast, but convergence stops in a potentially large neighborhood1 of the optimum x∗. This issue is exacerbated further by the extra noise coming from gradient compression. Indeed, this noise prevents methods such as Distributed Compressed Gradient Descent (DCGD) [Khirirat et al., 2018] from converging to x∗ with a constant learning rate even in the interpolation regime characterized by the identities ∇fi(x∗) = 0 for all i. Fortunately, these issues can be resolved via carefully designed variance reduction techniques [Gower et al., 2020]. In particular, the first variance reduction mechanism for removing the variance coming from compression operators in distributed training is due to Mishchenko et al. [2019], embodied in their DIANA algorithm. The method was initially analyzed for ternary quantization only [Wen et al., 2017], and later generalized to handle a general class of unbiased compression operators [Horváth et al., 2019b, Gorbunov et al., 2020b].
Acceleration. To speed up distributed training even further, it is often possible to employ Nesterov’s acceleration technique [Nesterov, 1983, 2004] in concert with gradient compression and variance reduction. For instance, Li et al. [2020] developed the ADIANA method, which adds acceleration on top of a variant of DIANA that relies on the computation of full-batch gradients on all nodes. The resulting method offers provable speedups in convex and strongly convex regimes. Another example is the method ECLK of Qian et al. [2020], which employs compressed communication via any (possibly biased) compressor satisfying a certain contraction property in combination with a slightly different variance reduction technique known as error compensation [Stich and Karimireddy, 2019, Karimireddy et al., 2019], while acceleration is offered by a loopless variant of the accelerated method Katyusha [Allen-Zhu, 2017, Kovalev et al., 2020].
Further tricks. Numerous other techniques are often used to improve some other aspects of distributed training, including implementing multiple local gradient steps before communication [Stich, 2020, Karimireddy et al., 2020, Woodworth et al., 2020a], asynchronous communication protocols [Agarwal and Duchi, 2011, Lian et al., 2015, Recht et al., 2011], in-network aggregation [Sapio et al., 2021], and performing the distributed training in a decentralized peer-to-peer manner without the reliance on an orchestrating server [Koloskova et al., 2019, Alghunaim et al., 2019]. However, in this work, we do not explore these directions and focus on the three techniques described before, namely, compressed communication, variance reduction and acceleration.

heading: 2 Mining for Smoothness Information
text: 2.1. One size fits all. Arguably, one of the most ubiquitous, if not the most ubiquitous, assumptions used in the literature on first-order optimization methods is that of L-smoothness [Nesterov, 2004]. A
1In theR ≡ 0 case, this neighborhood is proportional to the variance of the local gradients at the optimum: 1n ∑n i=1 ‖∇fi(x ∗)‖2.
differentiable function φ : Rd → R is said to be L-smooth if there exists a constant L ≥ 0 such that φ(x) ≤ φ(y) + 〈∇φ(x), x− y〉+ L2 ‖x− y‖ 2 (2)
holds for all x, y ∈ Rd. However, most works in the area of finite-sum distributed optimization use it very crudely: they assume that all local loss functions fi as well as their average, f = 1n ∑ i fi, share the same smoothness constant L [Tang et al., 2019, Woodworth et al., 2020b, Stich, 2020]. This is crude because much information is lost. Indeed, assuming that each fi is Li-smooth, it is well known that f is Lf -smooth with Lf satisfying Lf ≤ 1n ∑ i Li. In the light of this, the above assumption is crude as it effectively replaces the values L1, . . . , Ln and Lf with a single parameter L satisfying L ≥ max{L1, . . . , Ln}. Since the stepsizes and convergence rates of first-order methods depend on the smoothness constant(s) employed, convergence analysis relying on such crude approximation may be significantly suboptimal, and the methods too slow when implemented following the theory.
2.2. “According to the work of their hands” (Lam 3:64). Significant theoretical and practical improvement can often be obtained when taking account of all the smoothness constants involved, avoiding the practice of replacing them all with a single crude bound. Such analyses are more rare, but fairly common. For example, [Richtárik and Takáč, 2016a, Hanzely and Richtárik, 2019a].
2.3. “Like treasure hidden in a field, which a man found and covered up” (Mat 13:44). The starting point of this paper is the observation that there is a hitherto untapped richness of smoothness information that can be used to construct better distributed optimization algorithms and obtain better theory. This information is available, but hidden from sight, and is based on the notion of matrix smoothness. Definition 1 (Matrix Smoothness). We say that a differentiable function φ : Rd → R is L-smooth if there exists a symmetric positive semidefinite matrix L 0 such that
φ(x) ≤ φ(y) + 〈∇φ(y), x− y〉+ 12‖x− y‖ 2 L, ∀x, y ∈ Rd. (3)
The standard L-smoothness condition (2) is obtained as a special case of (3) for matrices of the form L = LI, where I is the identity matrix. In particular, if function fi appearing in (1) is often the average loss over the training data stored on node i, i.e.,
fi(x) = 1 mi ∑mi m=1 φim(Aimx), (4)
where Aim ∈ Rdim×d is a data matrix, and φim : Rdim → R is a differentiable function (e.g., the loss over all but the last linear layer of a NN). The following simple result from Qu and Richtárik [2016b], used therein in the context of randomized coordinate descent methods, states that if the loss functions φim are smooth in the standard scalar sense, then fi is smooth in the matrix sense. Lemma 1. If each φim is λim-smooth, then the function fi defined in (4) is Li-smooth with
Li = 1 mi ∑mi m=1 λimA > imAim. (5)
In cases where the local functions fi are of the form (4)2—and it is clear this structure is ubiquitous— there is a lot of potentially useful information contained in the matrix smoothness “constant” Li. If we were to use the scalar smoothness constant of fi instead, we would be effectively tossing this richness away, and replacing it with Li = λmax(Li); the largest eigenvalue of Li. This seems wasteful. As we show in this work, it is. However, we offer a fix.

heading: 3 Motivation and Contributions
text: To the best of our knowledge, none of the current distributed optimization methods, including the methods DCGD [Khirirat et al., 2018], DIANA [Mishchenko et al., 2019] and ADIANA [Li et al., 2020] discussed in Section 1, are capable of exploiting the inherently rich data-dependent smoothness structure of the local losses beyond standard smoothness constants. To this effect, we impose the following assumption throughout the paper: Assumption 1. The functions fi : Rd → R are differentiable, convex, lower bounded3 and Lismooth. Moreover, f is L-smooth with (standard) smoothness constant L := λmax(L).
2Our theoretical results hold for general loss functions fi and do not assume the structure (4) 3Lower boundedness of fi(x) can be dropped if Li 0 is positive definite. This part of the assumption is
not a restriction in applications as all loss function are lower bounded.
In this paper, we argue that when training supervised models, smoothness matrices (see Definition 1)— information-rich generalizations of the classical and ubiquitous smoothness constants—can and should be exploited for further dramatic gains, both in theory and practice.
3.1. Unbiased diagonal sketches. We study unbiased diagonal sketches, defined as follows: Definition 2 (Unbiased diagonal sketch). Let S be a random subset of the set of coordinates/features of the model x ∈ Rd we wish to train, i.e., S ⊆ [d] := {1, 2, . . . , d}. Let S be proper, i.e., pj := Prob(j ∈ S) > 0 for all coordinates j ∈ [d]. We now define a random diagonal matrix (sketch) C = CS ∈ Rd×d via
C = Diag(c1, . . . , cd), cj = {
1/pj if j∈S, 0 otherwise.
(6)
Note that given a vector x = (x1, . . . , xd) ∈ Rd, we have (Cx)j = { xj/pj if j∈S
0 if j /∈S . So, we can control the sparsity level of the product Cx by engineering the properties of the random set S. Also note that E[Cx] = x for all x. 3.2. Data-dependent sparsification operators. In order to further alleviate the communication burden inherent in distributed optimization, we further propose data-dependent sparsification operators that can take full advantage of the smoothness matrices Li associated with the local losses fi. To the best of our knowledge, this is in sharp contrast with the design of all existing tractable compression techniques used in distributed training, which are proposed independently of the training data, and typically based on intuitive or information-theoretic principles.
With each node i we associate an unbiased diagonal matrix Ci of the form (6). We use this and the smoothness matrix of fi to define a sparsification technique, described next. Definition 3 (Data-dependent sparsification). In situations when the i-th node wished to communicate local gradient ∇fi(x), we ask the node to send the sparse (=compressed) vector CiL† 1/2 i ∇fi(x) to the server instead. The server then constructs (=decompresses) an unbiased estimator of∇fi(x) as:
gi(x) = L 1/2 i CiL †1/2 i ∇fi(x), (7)
where L† 1/2 i denotes the square root of the Moore-Penrose pseudoinverse of Li.
Notable differences of our proposed communication protocol when compared with standard sparsification techniques are: i) we use the smoothness matrix Li, ii) the compressed vector CiL †1/2 i ∇fi(x) is not unbiased, iii) we devise a separate decompression mechanism (7), also involving Li, and this enforces effective unbiasedness.
3.3. Matrix-smoothness-aware redesign of 3 methods. To showcase the power of our approach, we demonstrate how our matrix-smoothness-aware sparsification technique (7) can be adapted to DCGD, DIANA and ADIANA, in each case leading to significant communication savings. By doing so, we show that matrix smoothness can be effectively used to speed up communication compression, variance reduction and acceleration, respectively. This results in three novel methods: DCGD+, DIANA+, and ADIANA+; see Table 1.
3.4. Dramatic improvements in complexity results. We perform complexity analyses for our methods and derive convergence rates under matrix smoothness4 (see Assumption 3) and strong convexity assumptions (see Theorems 2, 3 and 4). We show that new methods always outperform the originals/baselines, and often dramatically so.
To illustrate the potential of our sparsification technique (7) embedded in the new methods, let all machines i ∈ [n] use sketches Ci induced by independent5 samplings Si with probabilities pi;j := Prob(j ∈ Si). Then we show that, with optimized probabilities pi;j , DCGD+ can be O(min(n, d)) times faster then DCGD (see Remark 3) and DIANA+ can be O(min(n, d)) times faster than DIANA (see Remark 4), depending on matrices Li. For the accelerated method, we highlight improvements when condition numbers of subproblems are O(nd2). We show that ADIANA+ can be faster than the original ADIANA by a factor of O( √ d) in high compression regime, and by a
factor of O( √
min(n, d)) in low compression regime (see Remark 5). Main theoretical results are summarized in Table 2.
3.5. Single node case. Specializing our theory to the single machine setting (n = 1), we design new non-distributed algorithms providing an alternative viewpoint to randomized coordinate descent methods (see Appendix E).
3.6. Lower bounds. Using matrices as linear compression operators, we further investigate the trade-off between communicated bits and variance induced by the compression (see Appendix F).
3.7. Experiments. We conduct numerical experiments using LibSVM datasets [Chang and Lin, 2011], confirming the effectiveness and superiority of our sparsification protocol (7) over the standard sparsification scheme (see Section 6 and Section C).

heading: 4 New Communication-Efficient Methods Exploiting Matrix Smoothness
text: Consider the problem (1) with the smoothness Assumption 1 and for strongly convex f . Assumption 2 (µ-convexity). f : Rd → R is µ-convex for some µ > 0, i.e., for all x, y ∈ Rd
f(x) ≥ f(y) + 〈∇f(x), x− y〉+ µ2 ‖x− y‖ 2.
We present our new distributed methods, redesigned for matrix smoothness, and their convergence guarantees. Each node i ∈ [n] generates diagonal sketches Ci independently from others via an
4The closest to our result is work of Hanzely and Richtárik [2019b] and their ISEGA method which is able to exploit diagonal smoothness matrices. To the best of our knowledge, we are the first to fully exploit smoothness matrices of arbitrary structure, and elevate them as a new tool at the disposal of algorithm designers.
5Sampling Si is called independent if pi;jl := Prob({j, l} ⊆ Si) = pi;jpi;l for all j, l ∈ [d] such that j 6= l.
arbitrary sampling Si and, togther with its smoothness matrix Li, composes the compression matrix CiL †1/2 i . Probability matrices Pi, P̃i associated with the sampling Si and sketch Ci are defined as
Pi = (pi;jl) d j,l=1, P̃i = (p̃i;jl) d j,l=1 pi;jl = Prob({j, l} ⊆ Si), p̃i;jl = pi;jl pi;jjpi;ll − 1. (8)
Next, we introduce the key quantity, L̃max, describing the joint contribution of our sparsification (7) to the complexities of the three proposed methods:
L̃max = max1≤i≤n L̃i, L̃i = λmax(P̃i ◦ Li), (9) Above, ◦ stands for the Hadamard (i.e. element-wise) product.

heading: 4.1 DCGD+
text: We now present our matrix-smoothness-aware sparsification technique by adapting DCGD algorithm [Khirirat et al., 2018]. Upon receiving the current model xk from the server, each node computes L†
1/2 i ∇fi(xk) based on local training data and smoothness matrix. Next, sparsified up-
dates Cki L †1/2 i ∇fi(xk) are sent back to the server, which then averages decompressed updates L 1/2 i C k i L †1/2 i ∇fi(xk) and performs proximal step to get a new model xk+1.
Algorithm 1 DCGD+ 1: Input: Initial point x0 ∈ Rd, current point xk, step size γ, diagonal sketch Cki 2: on server 3: send xk to all nodes 4: get sparse updates Cki L †1/2 i ∇fi(xk) from each node
5: xk+1 = proxγR(x k − γgk), where gk = 1n ∑n i=1 L 1/2 i C k i L †1/2 i ∇fi(xk)
With this method we get convergence up to a neighborhood due to compressoin noise. Theorem 2 (see G.3). Let Assumptions 1 and 2 hold and assume that each node generates its own diagonal sketch Ci independently from others. Define σ∗ := 1n ∑n i=1 L̃i‖∇fi(x∗)‖2L†i . Then, for the step-size 0 < γ ≤ 1 L+ 2n L̃max , the iterates {xk} of Algorithm 1 satisfy
E [ ‖xk − x∗‖2 ] ≤ (1− γµ)k ‖x0 − x∗‖2 + 2γσ ∗
µn . (10)

heading: 4.2 Variance reduction: DIANA+
text: Next, we apply our sparsification technique to the variance reduced method DIANA [Mishchenko et al., 2019]. In this method, each node maintains an auxiliary control vector hki , called shift, which helps to reduce the variance coming from the sparsification. Moreover, the central server keeps track of only the averaged shift hk. Then, the model xk as well as control vectors hki , h
k are updated by decompressing sparse information ∆ki using matrices Li.
Algorithm 2 DIANA+ 1: Input: Initial point x0 ∈ Rd, initial shifts h0i ∈ Range(Li), current point xk, step size parameter γ and α, sketch Cki and C k i := L 1/2 i C k i L †1/2 i , current shifts h k 1 , . . . , h k n and h k := 1n ∑n i=1 h k i .
2: on each node 3: get xk from the server 4: send sparse update ∆ki = C k i L †1/2 i (∇fi(xk)− hki ) 5: update local gradient and shift ∆ki = L 1/2 i ∆ k i , g k i = h k i + ∆ k i , h k+1 i = h k i + α∆ k i 6: on server 7: get sparse updates ∆ki from each node 8: ∆k = 1n ∑n i=1 ∆ k i = 1 n ∑n i=1 L 1/2 i ∆ k i , g k = ∆k + hk = 1n ∑n i=1 C k i ( ∇fi(xk)− hki ) + hki 9: xk+1 = proxγR(x k − γgk), hk+1 = hk + α∆k
In this case we get rid of the neighborhood and provide linear convergence to the exact solution x∗. Throughout the paper, we use Õ notation to ignore log 1ε factors and constants. Theorem 3 (see G.4). Let Assumptions 1 and 2 hold and assume that each node generates its own diagonal sketch Ci independently from others. Then, for the step-size γ = 1L+ 6n L̃max , Algorithm 2
guarantees E [ ‖xk − x∗‖2 ] ≤ ε after
Õ ( ωmax + L µ + L̃max nµ ) (11)
iterations, where ωmax = max1≤i≤n ωi and ωi = max1≤j≤d p−1i;j −1 is the variance of compression operator induced by sketch Ci. Remark 1 (Variance Reduction: ISEGA+). In Appendix I we apply our redesign to another variance reduced method called ISEGA [Mishchenko et al., 2020, Hanzely and Richtárik, 2019b]. At the core of ISEGA, the mechanism for variance reduction is based on SEGA method [Hanzely et al., 2018]. The key difference between ISEGA and DIANA is that ISEGA updates the control variates h more aggressively using projection instead of the mere α-step towards the projection used in DIANA. The method is presented as Algorithm 7 in Appendix I. Theorem 22 provides the result – we can see that the worst case complexity is identical to DIANA+. However, in terms of the practical performance, we expect ISEGA+ to outperform DIANA+ due to the more aggressive update rule of control variates. Remark 2 (Variance Reduction with Bi-directional Compression: DIANA++). As an extension to DIANA+, in Appendix J we apply our sparsification technique both for nodes and for the central server, thus compressing gradients in both directions of communication. We develop and analyze DIANA++ (see Algorithm 8), for which the central server applies compression in its turn with sketch C independently. To converge in a linear rate, DIANA++ maintains an additional control vector to reduce the variance coming from the master’s sparsification. Our convergence theory (Theorem 23) recovers the same complexity (11) of DIANA+ if no compression is applied by the master.

heading: 4.3 Acceleration with variance reduction: ADIANA+
text: Finally, we redesign the accelerated method ADIANA [Li et al., 2020] to effectively exploit local smoothness matrices. The algorithm develops four sequences {xk, yk, zk, wk} of models, which are layered via convex combinations, proximal steps and probabilistic assignments. In each iteration, nodes receive models xk and wk from the server, and send back sparse updates ∆ki and δ k i using local data and control vectors hki . Then, decompressing these sparse vectors with matrices Li, nodes update their shifts hki and the server updates all four models along with averaged shift h k.
Algorithm 3 ADIANA+ 1: Input: Initial points x0 = y0 = z0 = w0 ∈ Rd, initial shifts h0i ∈ Range(Li), current point xk, parameters γ, α, β, η, θ1, θ2, q, sketch Cki and C k i := L 1/2 i C k i L †1/2 i , current shifts h k 1 , . . . , h k n
and hk = 1n ∑n i=1 h k i
2: on server 3: xk = θ1zk + θ2wk + (1− θ1 − θ2)yk, send xk and wk to all nodes 4: on each node 5: send sparse updates ∆ki = C k i L †1/2 i (∇fi(xk)− hki ) and δki = Cki L †1/2 i (∇fi(wk)− hki ) 6: update local gradient ∆ki = L 1/2 i ∆ k i , g k i = h k i +∆ k i and shift δ k i = L 1/2 i δ k i , h k+1 i = h k i +αδ k i 7: on server 8: get sparse updates ∆ki and δ k i from each node
9: ∆k = 1n ∑n i=1 L 1/2 i ∆ k i , δ k = 1n ∑n i=1 L 1/2 i δ k i , g k = ∆k + hk
10: hk+1 = hk + αδk, yk+1 = proxηR(x k − ηgk)
11: zk+1 = βzk + (1− β)xk + γη (y k+1 − xk), wk+1 = { yk with probability q, wk with probability 1−q.
Clearly, ADIANA+ enjoys the accelerated rate, which is strictly better then the one for DIANA+. Theorem 4 (see G.5). Let Assumptions 1 and 2 hold and assume that each node generates its own diagonal sketch Ci independently from others. Then, the iteration complexity of Algorithm 3
guaranteeing E [ ‖zk − x∗‖2 ] ≤ ε is Õ ( ωmax + √ ωmax L̃max µn ) if nL ≤ L̃max Õ ( ωmax + √ L µ + √ ωmax √ L̃max µn √ L µ ) if nL > L̃max.
(12)

heading: 5 Improvements Over the Original Methods
text: To compare the proposed methods with originals and highlight improvement factors, we choose independent sampling for all nodes. For Algorithms 1 and 2, we optimize probabilities of the samplings based on the complexities we found.
5.1. Parameters describing matrices Li. Define parameters ν, νs describing local smoothness as
ν := ∑n i=1 Li
maxi∈[n] Li , νs := maxi∈[n]
∑d j=1 L 1/s i;j
maxj∈[d] L 1/s i;j
, (13)
where Li = λmax(Li) and s = 1 or s = 2. Let Lmax := max1≤i≤n Li and Li;j be the jth diagonal element of matrix Li. Note that parameters ν ∈ [1, n] and νs ∈ [1, d] describe the distribution over the nodes and coordinates respectively. If Li are distributed uniformly, then ν = n and νs = d. On the other extreme, when the distribution is extremely non-uniform, we have ν n and νs d. These parameters are used to highlight the range of iteration complexities new methods can provide.
5.2. Importance sampling for DCGD+. Let τ = E [|Si|] = ∑d j=1 pi;j be the expected mini-batch size for the samplings Si, where pi;j = pi;jj . Notice that convergence rate of Algorithm 1 depends on L̃max = max1≤i≤n L̃i. Since each node i ∈ [n] generates its own diagonal sketch Ci independently from others, each node can optimize L̃i = λmax(P̃i ◦ Li) independently based on local smoothness matrix Li. In general, minimizing λmax(P̃i ◦ Li) with respect to probability matrix P̃i is hard. However, when each node uses an independent sampling, which means pi;jl = pi;jpi;l if j 6= l, then
λmax(P̃i ◦ Li) = max1≤j≤d ( 1 pi;j − 1 ) Li;j , (14)
for which we can find the optimal probabilities pi;j . To minimize the maximum term in (14), we should have (1/pi;j − 1) Li;j = ρi for some ρi ≥ 0. Then the solution is
pi;j = Li;j
Li;j+ρi , (15) where ρi ≥ 0 is the unique solution to ∑d j=1 Li;j Li;j+ρi
= τ . While the latter does not allow closed form solution for ρi, it can be computed numerically using one dimensional solvers. Thus, we can efficiently compute the optimal probabilities (15). Proposition 5 (Optimality). The independent sampling with probabilities (15) is the optimal independent sampling for the rate (10). Remark 3 (Improvement over DCGD [Khirirat et al., 2018]). With probabilities (15) we show in Appendix H.1 that
L µ + L̃max nµ ≤ ( ν n + ν1 τn ) Lmax µ . (16)
In the interpolation regime (i.e. ∇fi(x∗) = 0 for all i ∈ [n]), the iteration complexity of DCGD is Õ(Lµ + ωLmax nµ ) for general compression operator with variance parameter ω. In case of sparsification, it is known that τ = d/n is the optimal choice under standard (scalar) smoothness assumption [Mishchenko et al., 2019]. If we specialize compression to sparsification with τ = d/n entries (which gives ω = d/τ − 1 = n − 1), we get Õ(Lmaxµ ). Notice that, in this regime, Theorem 2 also provides linear convergence with iteration complexity Õ(Lµ + L̃max nµ ). Based on (16), it is bounded by Õ(( νn + ν1 d ) Lmax µ ), which is always better than Õ( Lmax µ ) and can be as small as Õ( Lmax min(n,d)µ ). Hence, for mini-batch τ = d/n, DCGD+ (Algorithm 1) guarantees the same Õ(Lmaxµ ) complexity in the worst case, but could provide up to min(n, d) times speedup. An anologous observation can be made between standard sparsificaiton with uniform probabilites and our sparsification with uniform probabilities (see Remark 7 for the details).
5.3. Importance sampling for DIANA+. To find optimal probabilities for DIANA+, we minimize ωmax + L̃max µn part of the complexity (11), which is in turn equivalent to minimize
max1≤j≤d ( 1 pi;j − 1 ) L′i;j , L ′ i;j := Li;j µn + 1, (17)
which can be solved in the same way as (14) yielding
pi;j = L′i;j
L′i;j+ρ ′ i
= Li;j+µn
Li;j+(1+ρ′i)µn . (18)
Proposition 6 (Optimality). The independent sampling with probabilities (18) is the optimal6 independent sampling for the complexity (11).
Remark 4 (Improvement over DIANA [Mishchenko et al., 2019, Horváth et al., 2019b]). The iteration complexity of DIANA, with τ = d/n is Õ(n+ Lmaxµ ). With probabilities (18) we upper bound the iteration complexity (11) in Appendix H.2 as follows
ωmax + L µ + L̃max µn ≤ 2d τ + ( ν n + 2ν1 τn ) Lmax µ . (19)
Therefore, with τ = d/n, DIANA+ (Algorithm 2) guarantees the same Õ(n+ Lmaxµ ) complexity in the worst case, but could provide up to min(n, d) times speedup with complexity Õ(n+ Lmaxmin(n,d)µ ).
5.4. Independent sampling for ADIANA+. Clearly, if we sparsify with uniform probabilities pi;j = τ/d, then Algorithm 3 recovers the rate of ADIANA.
Remark 5 (Improvement over ADIANA [Li et al., 2020]). To show that the rate could be significantly better in some cases, consider the following choice pi;j = √
L′i;j L′i;j+ρ ′′ i , L′i;j = Li;j µn + 1, where ρ ′′ i
is determined uniquely from ∑d j=1 pi;j = τ . Then, with these probabilities and for Lmax/µ =
O(nd2), we show in Appendix H.3 that Lµ ≤ νLmax nµ , ωmax = O ( ν2d τ ) , Lmaxµn = O ( ν2d τ √ Lmax nµ ) .
Furthermore, assuming both ν and ν2 are O(1), choosing τ = d/n we get Lµ ≤ O ( Lmax nµ ) , ωmax =
O (n) , Lmaxµn = O (√ nLmax µ ) . Then, the complexity (12) of ADIANA+ reduces to{ n+n(Lmaxnµ ) 1/4 if nL≤L̃max,
n+ √ Lmax nµ +(n Lmax µ ) 3/8 if nL>L̃max,
which, compared to the complexity of ADIANA with ω = O(n) compression, gives √ d times improvement in the first case and √
min(n, d) times improvement in the second case (ignoring the first summand n of the complexities).

heading: 6 Experiments
text: We numerically compare the proposed matrix-smoothness-aware sparsification strategy (7) with the usual sparsification scheme. We devise three different experiments on logistic regression with LibSVM data [Chang and Lin, 2011]. Due to space limitations, only two of them is presented in the main part. Further numerical results alongside with their experimental details are in Appendix C.
Experiment: Proposed and usual sparsification techniques for the 3 distributed methods. We compare six different methods: well-established DCGD, DIANA, ADIANA and our methods DCGD+, DIANA+, ADIANA+, all with uniform sampling for τ = 1. In order to highlight the importance of the variance reduction, in this experiment we choose the starting point to be close to the optimum.
Figure 1 demonstrates that: i) methods with matrix-aware sparsification (i.e., DCGD+, DIANA+, ADIANA+) outperform their baselines (i.e., DCGD, DIANA, ADIANA) ii) acceleration almost
6In the sense that it minimizes a quantity, which is the complexity of DIANA+ up to some constant factor.
always outperforms the non-accelerated variant, often dramatically so and iii) variance reduction never hurts the convergence, but often stabilizes the oscillation of the non-variance reduced counterpart.
Experiment: Variance reduction with new sparsification and importance sampling. We now comment on the experiment illustrated in Figure 2. We examine three sparsification schemes (two variants of our strategy and the usual sparsification) and their influence on convergence using six different datasets. Considered schemes are i) DIANA+ with importance sampling (18), ii) DIANA+ with uniform sampling, and iii) DIANA with uniform sampling, i.e., uniform sparsification unaware of smoothness matrices. In all three cases we fixed the sampling size τ = 1.
As expected, Figure 2 confirms our theoretical findings. First, it demonstrates that our sparsification (7) always outperforms the naive/direct sparsification, sometimes by a large margin. Second, it shows the benefit of importance sampling (18) over the uniform sampling.

Review summaries:
The papers consider the standard distributed optimization scenarios, where a server wants to minimize the average of $n$ loss functions, which are evenly distributed across $n$ participating worker nodes. Typically, optimization problems designed for such distributed optimization scenarios assume a blanket smoothness constant for the average and the individual functions involved. Instead, the authors consider the notion of matrix smoothness, where the distance between the points in defining the smoothness is measured with matrix norm (of a smoothness matrix). Under the assumption that the individual loss function satisfies matrix smoothness with different matrices, the authors design compression algorithms using these smoothness matrices and then use them in conjunction with various distributed optimization algorithms. The authors show that as a result of their smoothness-aware compression, they improve over the convergence guarantees of the standard baseline algorithms considered, namely, DCGD, DIANA, ADIANA.

This paper proposes a new notion of smoothness called matrix smoothness which leads to better theory bounds and a data-dependent compression scheme. The compression scheme is applied to three existing algorithms resulting in improved convergence bounds in smooth and strongly convex setting. Simple experiments are done to show the effectiveness of the proposed scheme.

This paper analyzes several optimization concepts: compression, variance reduction, and acceleration, using the assumption of smoothness matrices rather than the smoothness parameter. Sketching matrices are designed as sum_{i\in S} L^{1/2}_i c_i * L^{1/2,dagger}_i where L_i is the "smoothness matrix" associated with the ith data sample. Taking this additional information into account, which for generalized linear models is not too hard to compute, it is suggested that there can be a factor of min(n,d) improvement in unaccelerated methods and sqrt(min(m,n)) benefit in accelerated methods.