Year: 2021

Abstract:
Johnson-Lindenstrauss lemma is one of the most valuable tools in machine learning, since it enables the reduction to the dimension of various learning problems. In this paper, we exploit the power of Fast-JL transform or so-called sketching technique and apply it to federated learning settings. Federated learning is an emerging learning scheme which allows multiple clients to train models without data exchange. Though most federated learning frameworks only require clients and the server to send gradient information over the network, they still face the challenges of communication efficiency and data privacy. We show that by iteratively applying independent sketches combined with additive noises, one can achieve the above two goals simultaneously. In our designed framework, each client only passes a sketched gradient to the server, and de-sketches the averagegradient information received from the server to synchronize. Such framework enjoys several benefits: 1). Better privacy, since we only exchange randomly sketched gradients with low-dimensional noises, which is more robust against emerging gradient attacks; 2). Lower communication cost per round, since our framework only communicates low-dimensional sketched gradients, which is particularly valuable in a small-bandwidth channel; 3). No extra overall communication cost. We provably show that the introduced randomness does not increase the overall communication at all.

heading: 1 INTRODUCTION
text: Federated learning enables multiple parties to collaboratively train a machine learning model without directly exchanging training data. This has become particularly important in areas of artificial intelligence where users care data privacy, security, and access rights, including healthcare (Li et al., 2020b; 2019), internet of things (Chen et al., 2020), and fraud detection (Zheng et al., 2020).
Given the importance and popularity of federated learning, it has become an important research topic for academia and industry, mostly focusing on two central themes. One is on data privacy. Federated learning seemingly protects clients’ privacy, since it only communicates gradient information. Unfortunately, recent studies (Geiping et al., 2020; Zhu & Han, 2020; Wang et al., 2019) have demonstrated that attackers can recover the input data from the shared gradients. The reason why the attacks work is the gradients carry important information about the training data (Ateniese et al., 2015; Fredrikson et al., 2015). The second one is on communication efficiency. Machine learning models are becoming increasingly larger but client devices that carry private data only have limited network bandwidth. The size of the gradient is the same as the size of the model, and the amount of data to communication between the clients and the servers thus is large. This becomes even more problematic when conducting federated learning on mobile and edge devices, where the bandwidth of the network is low. The communication cost is one of the most important the key performance bottlenecks in federated learning systems Goga & Teixeira (2012); Konečnỳ et al. (2016). Many works try to address this challenge through local optimization method, such as local gradient descent (GD), local stochastic gradient descent (SGD) and their variants (Konečnỳ et al., 2016; McMahan et al., 2017; Stich, 2018).
Despite existing efforts, no work addresses both challenges simultaneously as far as we concern. Therefore, we ask the following question:
Is there a FL framework that protects the local privacy and tackles the communication challenge?
In this paper, we achieve these goals by using an old but powerful idea — the Johnson-Lindenstrauss transform and its fast variations (Fast-JL, see Ailon & Chazelle (2006)). If we view the transform as a sketch matrix, and its transpose as de-sketch, then our main idea is to iteratively apply the sketch and de-sketch matrices to gradients. Instead of running the vanilla gradient descent w(t+1) ← w(t) − η · g(t) using true gradient g(t) ∈ Rd, we apply sketch and de-sketch to the gradient:
w(t+1) ← w(t) − η ·R> ·R · g(t).
Here R ∈ Rbsketch×d denotes a sketching matrix that sketches the true gradient to a lower dimension andR> ∈ Rd×bsketch denotes the de-sketching process that maps the sketched gradient back to the true gradient dimension. The coordinate-wise embedding property (Song & Yu, 2021) ensures R>Rg(t) being an unbiased estimator of g(t) with bounded second moments, implying the new sketched gradient descent scheme preserving the original convergence properties. Hence, all clients will only communicate sketched gradients to the server, the server averages the sketched gradients and broacasts it back to all clients. Finally, each client de-sketches the received gradients and perform local updates. Since the sketching dimension is always small compared to original dimension, we save communication cost per iteration via Johnson-Lindenstrauss.
Though the communication problem has been addressed, such framework still faces privacy challenge: applying Johnson-Lindenstrauss “masks” the communicated gradients, but in order to give a provable privacy guarantee on this framework, we introduce additive low-dimensional Gaussian noises to make sure that the communicated vectors themselves are differential private.
We summarize the contributions in this work as follows:
Framework contribution: We propose a new federated learning framework that iteratively applies sketch and de-sketch matrices to gradients, which enjoys the following advantages:
• Privacy: it preserves the privacy via additive low-dimensional Gaussian noise. • Communication: it reduces communication per round, since at each synchronization step, only a
sketched gradient of lower-dimensional is communicated. • Convergence: it preserves the convergence rate of vanilla local gradient descent method.
Technical contribution: Our analysis technique is also of independent interest in the relating areas:
• Unlike classical sketch-and-solve paradigm, our iterative sketch and de-sketch method can be combined with gradient-based methods and extended to broader optimization problems. • We provide rigorous analysis on the impact of introducing sketching through coordinate-wise embedding, which can be generalized to other areas (Song & Yu, 2021). • As a by-product, we give a novel linear convergence result of local GD under the strongly-convex and smooth scheme.

heading: 2 RELATED WORK
text: Federated Learning Federated learning (FL) is an emerging framework in distributed deep learning. FL allows multiple parties or clients collaboratively train a model without data sharing. Fl let the local client perform most of the computation and a central sever update the model parameters through aggregation then transfers the parameters to local models (Dean et al., 2012; Shokri & Shmatikov, 2015; McMahan et al., 2016; 2017). In this way, the details of the data are not disclosed in between each party.
Unlike the standard parallel setting, FL has three unique challenge (Li et al., 2020a), including communication cost, data heterogeneity and client robustness. In our work, we focus on the first two challenges. The training data are massively distributed over an incredibly large number of devices, and the connection between the central server and a device is slow. A direct consequence is the slow communication, which motivated communication-efficient FL algorithm. Federated average (FedAvg) (McMahan et al., 2017) firstly addressed the communication efficiency problem by introducing a global model to aggregate local stochastic gradient descent updates. Later, different variations and adaptations have arisen. This encompasses a myriad of possible approaches, including developing better optimization algorithms (Wang et al., 2020a) and generalizing model to heterogeneous clients under special assumptions (Zhao et al., 2018; Kairouz et al., 2019; Li et al., 2021).
Local GD and Local SGD To seek the communication efficiency of Federated learning, local SGD has been proposed (Konečnỳ et al., 2016; McMahan et al., 2017), where each client does a few SGD iterations locally before the server averages the local estimators. Different variants of local SGD algorithms has been explored, including with momentum (Yu et al., 2019b; Wang et al., 2018), with quantization (Basu et al., 2019; Reisizadeh et al., 2020), and with various variance-reduction methods (Liang et al., 2019; Karimireddy et al., 2020). Convergence analysis for local SGD mainly focuses on two regimes: identical data regime (Stich, 2018; Basu et al., 2019; Stich & Karimireddy, 2019; Haddadpour & Mahdavi, 2019; Khaled et al., 2020) and heterogeneous data regime (Jiang & Agrawal, 2018; Yu et al., 2019a; Basu et al., 2019; Haddadpour & Mahdavi, 2019; Khaled et al., 2019; 2020). In this work, we propose our framework based upon vanilla local GD and our analysis focus on the heterogeneous data regime.
Sketching Sketching has many applications in numerical linear, such as linear regression, lowrank approximation (Clarkson & Woodruff, 2013; Nelson & Nguyên, 2013; Meng & Mahoney, 2013; Boutsidis & Woodruff, 2014; Song et al., 2017; Andoni et al., 2018), distributed problems (Woodruff & Zhong, 2016; Boutsidis et al., 2016), reinforcement learning Wang et al. (2020b), tensor decomposition (Song et al., 2019), clustering (Esfandiari et al., 2021), cutting plane method (Jiang et al., 2020), generative adversarial networks (Xiao et al., 2018) and linear programming (Lee et al., 2019; Jiang et al., 2021; Song & Yu, 2021).
Notations For a positive integer n, we use [n] to denote the set {1, 2, · · · , n}. We use E[·] to denote expectation (if it exists), and use Pr[·] to denote probability. For a vector x, we use ‖x‖2 := ( ∑n i=1 x 2 i )
1/2 to denote its `2 norm. We denote 1{x=l} for l ∈ R to be the indicator function which equals to 1 if x = l and 0 otherwise. Let f : A → B and g : C → A be two functions, we use f ◦ g to denote the composition of functions f and g, i.e., for any x ∈ C, (f ◦ g)(x) = f(g(x)). We denote Id to be the identity mapping.

heading: 3 PROBLEM SETUP
text: Consider a federated learning scenario with N clients and corresponding local losses fc : Rd → R, our goal is to find
min x∈Rd
f(x) := 1
N N∑ c=1 fc(x) (1)
In this work, we consider the following classical convex and smooth setting for our objectives. Assumption 3.1. Assume that the set of minimizers of (1) is nonempty. Each fc is µ-strongly convex for µ ≥ 0 and L-smooth. That is, for all x, y ∈ Rd,
µ 2 ‖y − x‖22 ≤ fc(y)− fc(x) + 〈y − x,∇fc(x)〉 ≤ L 2 ‖y − x‖22.
Note in the case µ = 0, this assumption reduces back to convexity and smoothness.
Apart from the above assumption, we allow local losses to have arbitrary heterogeneity. In other words, we allow fc’s to be arbitrary functions.

heading: 4 OUR ALGORITHM
text: In this section, we propose a federated learning framework that addresses the communication efficiency issue. When the learning gradients are of high dimension, classical federated learning framework which sends the exact gradient could incur a heavy communication cost per round. Sketching technique, which emerges as an effective way to reduce the dimension of vector while preserving significant amount of information (Sarlós, 2006; Woodruff, 2014), is highly preferred in this setting. It enables us to compress the gradient vector into a lower dimension while preserving convergence rates, and greatly saves the communication cost per round.
Motivated by above discussion, we propose the iterative sketching-based federated learning algorithm, which builds upon vanilla local gradient descent: we start with a predetermined sequence of
independent sketching matrices shared across all clients. In each round, local clients accumulate and sketch its change over K local steps, then transmit the low-dimensional sketch to the server. Server then averages the sketches and transmits them back to all clients. Upon receiving, each client de-sketches to update the local model.
Algorithm 1 Iterative Sketching-based Federated Learning Algorithm with K local steps 1: procedure ITERATIVESKETCHINGFL 2: Each client initializes w0 using the same set of random seed 3: for t = 1→ T do . T denotes the total number of global steps 4: /* Client */ 5: parfor c = 1→ N do . N denotes the total number of clients 6: if t = 1 then 7: ut,0c ← w0 . ut,0c ∈ Rd 8: else 9: ut,0c ← wt−1 + deskt(∆w̃t−1) . deskt : Rbsketch → Rd de-sketch the change 10: end if 11: wt ← ut,0c 12: for k = 1→ K do 13: ut,kc ← ut,k−1c − ηlocal · ∇fc(w)|w=ut,k−1c 14: end for 15: ∆wc(t)← ut,Kc − wt 16: Client c sends skt(∆wc(t)) to server . skt : Rd → Rbsketch sketch the change 17: end parfor 18: /* Server */ 19: ∆w̃t ← ηglobal · 1N ∑N c=1 skt(∆wc(t)) . ∆w̃
t ∈ Rd 20: Server sends ∆w̃t to each client 21: end for 22: end procedure
We highlight several distinct features of our algorithm:
•Communication: In each sync step, we only communicates a low-dimensional sketched gradients, indicating a smaller communication cost per round. This property is particularly valuable in a smallbandwidth setting. • De-sketch:i: We emphasize that unlike the classical sketch-and-solve paradigm that decreases the problem dimension, our algorithm applies sketching in each round, combined with a de-sketching process which recovers back to the true gradient dimension. • Simple server task:: Server only needs to do simple averaging, indicating no need of a trustworthy party as the server. • Decentralization:: Our algorithm can be generalized to decentralized learning settings, where local clients can only communicate with neighboring nodes. In this case, it requiresO(diam) rounds to propagate the sketched local changes, where diam is the diameter of the network graph.
4.1 sk/desk VIA COORDINATE WISE EMBEDDING
In this section, we discuss the concrete realization of the skt/deskt operators in Algorithm 1 through random sketching matrices. Note we should require any processed gradient deskt ◦ skt(g) to "be close" to the true gradient g to avoid breaking the convergence property of the algorithm. To achieve this, we first introduce the following property for a broad family of sketching matrices, namely coordinate-wise embedding (Jiang et al., 2021; Song & Yu, 2021), that naturally connects with skt/deskt operators. Definition 4.1 (a-coordinate-wise embedding). We say a randomized matrix R ∈ Rbsketch×d satisfying a-coordinate wise embedding if for any vector g, h ∈ Rd, we have ER∼Π[h>R>Rg] = h>g and ER∼Π[(h>R>Rg)2] ≤ (h>g)2 + absketch ‖h‖ 2 2 · ‖g‖22.
iWe elaborate the difference of our iterative sketch and de-sketch approach compared to classical sketchand-solve approaches in Section 4.2.
In general, well-known sketching matrices have their coordinate-wise embedding parameter a being a small constant (See Section D). Note that if we choose h to be one-hot vector ei, then the above conditions translate to ER∼Π[R>Rg] = g and ER∼Π[‖R>Rg‖22] ≤ (1 + a · dbsketch ) · ‖g‖ 2 2.
This implies that by choosing
skt = Rt ∈ Rbsketch×d (sketching), deskt = R>t ∈ Rd×bsketch (de-sketching) (2)
for any iteration t ≥ 1, where Rt’s are independent random matrices with sketching dimension bsketch, we obtain an unbiased sketching/de-sketching scheme with bounded variance as state in the following Theorem 4.2.
Theorem 4.2. Let skt and deskt be defined by Eq. (2) using a sequence of independent sketching matrices Rt ∈ Rbsketch×d satisfying a-coordinate wise embedding property (Def. 4.1). Then the following properties hold:
1. Independence: For different iterations, (skt, deskt)’s are independent of each other.
2. Linearity: Both skt and deskt are linear operators.
3. Unbiased estimator: For any fixed vector h ∈ Rd, it holds E[deskt(skt(h))] = h.
4. Bounded second moment: For any fixed vector h ∈ Rd, it holds E[‖deskt(skt(h))‖22] ≤ (1 + a · d/bsketch) · ‖h‖22.
We will use the above property to instantiate the convergent proof and communication complexity in section 5. We remark that unlike traditional sketching matrix R, one can intuitively think of matrix R> as a “de-sketch” matrix, it undoes sketching and recovers the sketched vector to original dimension.

heading: 4.2 SKETCH-AND-SOLVE VS SKETCH-AND-DE-SKETCH
text: In this section, we provide a brief discussion of the difference between classical sketch-and-solve paradigm (Clarkson & Woodruff, 2013; Woodruff, 2014) and our sketch-and-de-sketch scheme. The general idea of sketch-and-solve is to apply sketching matrix R to the entire problem, and use certain black-box algorithm to solve the sketched-down low-dimensional version of the problem. Intuitively, sketching preserves the structure of the problem, therefore, the same algorithm can be exploited on a smaller problem, and to achieve a 1 ± guarantee. One downside of this method is the problem that is applicable needs to have certain structures. For example, it is not clear that given a non-convex objective, directly applying sketching will do any help.
In contrary, our sketch-and-de-sketch scheme only applies sketching to certain key component of a problem, for example, the gradient in a gradient descent method. Unlike sketch-and-solve, using sketch and de-sketch, we preserve the structure of the algorithm. This makes it feasible to a wider range of applications, such as showing the progress of gradient descent on non-convex objective, as we will show in appendix E and G. One drawback of this scheme is it gives a worse approximation guarantee (1 ± O( dbsketch )), but in gradient descent, this can be mitigated via choosing a smaller stepsize for compensation. In a distributed setting, where privacy of communicated message and its size are main concerns, sketch-and-de-sketch is particularly valuable.

heading: 4.3 PRIVACY-PRESERVED SKETCHING VIA LOW-DIMENSIONAL NOISES
text: In this section, we discuss how to add low-dimensional Gaussian noise to make Algorithm 1 differential private. Consider a gradient vector gc generated via training of a client c, and another vector g′c that differs from gc by exactly one entry. Further, we assume this difference has a bounded magnitude of γ in terms of absolute value. The goal is to protect the sketched gradient, Rgc. Our strategy is as follows: on line 16 of Algorithm 1, we add a random vector ∆c ∈ Rbsketch where each entry of ∆c is drawn fromN (0, σ2) and with σ ≥ Ω( √ d/bsketchγ −1 √
log(1/δ)), then this modification on line 16 will produce an ( , δ)-differential private guarantee for the sketched gradient.
We remark this is the most important step to preserve the privacy of the federated learning algorithm, since potential attackers might hack into a single client and have access to the skt/deskt operators for
each iteration. If they can further observe the communicated gradients from other clients, then the power of random masking with sketching is almost diminished, since they intuitively, they can desketch the sketched changes to obtain useful information. After adding low-dimensional Gaussian noises, attackers can no longer obtain useful information even they have the skt/deskt operators, which significantly improves the privacy of our system.
From an analytical perspective, adding this noise does not affect our analysis too much — since it’s sampled from a zero-mean Gaussian distribution, our estimator R>(Rg + ∆c) is still an unbiased estimator of vector g, it merely adds a variance term which can be factored into the variance of our un-modified estimator. Hence, in the analysis section, we present the theorems and proofs for the scenario where noises are not added. Similar analysis can be adapted to additive noise version.

heading: 5 CONVERGENCE THEORY AND COMMUNICATION COMPLEXITY
text: In this section, we analyze the convergence property of our proposed framework for smooth and convex objectives. Note the similarity shared by our framework and classical federated learning algorithms. We try to follow the existing analysis and focus on discussing the impact of the introduced randomness due to the sketching and de-sketching. We will show our approach enjoys benign convergence property and does not increase total communication at all.

heading: 5.1 SINGLE-STEP SCHEME
text: To start off, we consider a simple scenario where the number of local steps K = 1. In this case, by the linearity of skt and deskt operators, the update rule can be written as
wt+1 = wt − η · deskt(skt(∇f(wt))). (3)
We remark that in the case of skt and deskt being identity mappings, our framework is equivalent to a distributed implementation of the vanilla gradient descent algorithm. Therefore, we follow the classical analysis of gradient descent and have the following key lemma:
Lemma 5.1. If Assumption 3.1 holds and K = 1. Denote η := ηlocal · ηglobal. Then we have
E[‖wt+1 − w∗‖22] ≤ (1− ηµ)E[‖wt − w∗‖22]− 2η(1− η(1 + α)L) · E[f(wt)− f(w∗)].
where w∗ is a minimizer of problem (1).
Proof. Note the update rule (3) implies
‖wt+1 − w∗‖22 = ‖wt − w∗‖22 − 2η〈wt − w∗, deskt(skt(∇f(wt)))〉+ η2‖deskt(skt(∇f(wt)))‖22.
Taking the conditional expectation over last synchronization, we have by Theorem 4.2,
E[〈wt − w∗, deskt(skt(∇f(wt)))〉 | Ft] = 〈wt − w∗,∇f(wt)〉, E[‖deskt(skt(∇f(wt)))‖22 | Ft] ≤ (1 + α)‖∇f(wt)‖22,
Combining with µ-strongly convexity and L-smoothness of f , we obtain
E[‖wt+1 − w∗‖22 | Ft] ≤ ‖wt − w∗‖22 − 2η〈wt − w∗,∇f(wt)〉+ η2(1 + α)‖∇f(wt)‖22 ≤ (1− ηµ)‖wt − w∗‖22 − 2η(1− η(1 + α)L) · (f(wt)− f(w∗)).
Taking the expectation of both sides over Ft we complete the proof.
Lemma 5.1 implies the introduced randomness from gradients only influence the second order term by a multiplicative factor 1+α in expectation. Therefore, by scaling down the stepsize by a factor of 1+α, we can obtain the exact same convergence guarantee as the vanilla gradient descent algorithm. For the strongly convex and smooth objective case, we obtain the following linear convergence.
Theorem 5.2. If Assumption 3.1 holds with µ > 0. Let K = 1 and η := ηlocal · ηglobal ≤ 1(1+α)L , then we have E[f(wT )−f(w∗)] ≤ L2 E[‖w 0−w∗‖22]e−ηµT , wherew∗ is a minimizer of problem (1).
Proof. By the choice of stepsize and Lemma 5.1, we obtain
E[‖wT − w∗‖22] ≤ (1− ηµ)T E[‖w0 − w∗‖22].
Therefore, we have
E[f(wT )− f(w∗)] ≤ L 2 E[‖wT − w∗‖22] ≤ L 2 E[‖w0 − w∗‖22]e−ηµT .
For the convex case, we consider the average of iterations and obtain the sublinear convergence. Theorem 5.3. If Assumption 3.1 holds with µ = 0. Let K = 1 and η := ηlocal · ηglobal ≤ 12(1+α)L , then we have f(wT ) − f(w∗) ≤ 1η(T+1) E[‖w
0 − w∗‖22]. where w∗ is a minimizer of problem (1) and wT = 1T+1 ∑T t=0 w t.
Proof. By the choice of stepsize and Lemma 5.1, we obtain η(f(wt)− f(w∗)) ≤ E[‖wt − w∗‖22]− E[‖wt+1 − w∗‖22]. Take the telescope summation over t = 0, · · · , T , we have T∑ t=0 (f(wt)− f(w∗)) ≤ 1 η (E[‖w0 − w∗‖22]− E[‖wT+1 − w∗‖22]) ≤ 1 η E[‖w0 − w∗‖22]. By the convexity of f we complete the proof.
We point out that in the case of skt and deskt being identity mappings, the parameter α reduces back 0 and Theorem 5.2 and 5.3 matches the convergence property of vanilla gradient descent exactly.
Further, above convergence results imply that comparing to vanilla gradient descent, our approach needs to shrink the stepsize by a factor of O(α), thus enlarge the number of iterations by a factor of O(α) to achieve desired accuracy. Therefore, using sketching matrices with dimension bsketch, ours communicates O(bsketch/d ·α) as many bits in total compared to vanilla approaches. According to Theorem 4.2, we have α = O(d/bsketch) for commonly used sketching matrices, implying our approach does not introduce extra communication cost at all.

heading: 5.2 MULTI-STEP SCHEME
text: Now we are ready to move on to generalK local step scheme. In this section, we assume ηglobal = 1. For notation simplicity, we denote ut,−1c = u t−1,K−1 c for t ≥ 2. We also introduce the following notations of the average iterates, iterates variance, local gradients and average gradients to help with the analysis.
ut,k := 1
N N∑ c=1 ut,kc , V t,k := 1 N N∑ c=1 ‖ut,kc − ut,k‖22, gt,kc := ∇fc(ut,kc ), gt,k := 1 N N∑ c=1 gt,kc
Using the new set of notations, one key observation is the average iterates satisfies
ut,k = ut,k−1 − ηlocal · gt,k−1 + 1{k=0} · ηlocal · (Id − deskt ◦ skt)( K−1∑ i=0 gt−1,i), ∀(t, k) 6= (1, 0)
We remark again if skt and deskt being identity mappings, our updates reduce back to the vanilla local gradient descent algorithm with K local steps.
By considering the distance to optimal solution ‖ut,k − w∗‖22, we have the following intermediate lemma parallel to Lemma 5.1. Due to the space limitation, we defer the proof to appendix, see Lemma F.6. Lemma 5.4. If Assumption 3.1 holds and η := ηlocal ≤ 14L , then we have E[‖ut,k − w∗‖22] ≤ (1− µη)E[‖ut,k−1 − w∗‖22]− η E[f(ut,k−1)− f(w∗)] + 1.5ηLE[V t,k−1]
+ 1{k=0}η 2αK · ( 4L K−1∑ i=0 E[f(ut−1,i)− f(w∗)] + 2L2 K−1∑ i=0 E[V t−1,i] )
for any (t, k) 6= (1, 0), where w∗ is a minimizer of problem (1).
Follow upon the above lemma, next step is to capture the quantity of the iterates variance V t,k, we observe that in each round, we start off with V t,0 = 0 due to synchronization ut,0c = u
t,0. Then V t,k can be viewed as the accumulation of the variance of the next k local updates,
V t,k = 1
N N∑ c=1 ‖ut,0c − k−1∑ i=0 ηlocal · gt,ic − ut,0 + k−1∑ i=0 ηlocal · gt,i‖22 = η2local N N∑ c=1 ‖ k−1∑ i=0 (gt,i − gt,ic )‖22
Therefore, it naturally requires us to characterize V t,k through certain measure of dissimilarity of local gradients. To achieve so, we avoid the common Lipschitz assumption which only gives a blur upper bound of the difference of local gradients, but follow the approach of Khaled et al. (2019; 2020), which focus on the quantity
σ2 def =
1
N N∑ c=1 ‖∇fc(w∗)‖2 > 0
that is always finite and naturally characterize the degree of heterogeneity of local objectives. And we have the following observation of the summation of iterates variance over a single round: Lemma 5.5. If Assumption 3.1 holds and ηlocal ≤ 18LK . Then for any t ≥ 0,
K−1∑ k=0 V t,k ≤ 8η2localLK2 K−1∑ k=0 (f(ut,k)− f(w∗)) + 4η2localK3σ2
Combining Lemma 5.5 and Lemma 5.4, we notice that by choosing appropriately small stepsize, the first term in upper-bounding V t,k can be absorbed when we consider the average iterates over a single round. Therefore, We can obtain the following convergence result for the strongly convex and smooth losses. Theorem 5.6. If Assumption 3.1 holds with µ > 0. If ηlocal ≤ 18(1+α)LK ,
E[f(wT+1)− f(w∗)] ≤ L 2 E[‖w0 − w∗‖22]e−µηlocalT + 4η2localL2K3σ2/µ.
where w∗ is a minimizer of problem (1).
Proof. Telescoping sum up Lemma 5.4 as k varies from 0 to K − 1 and absorb the higher-order terms by the choice of the stepsize, we have for any t ≥ 1, E[‖ut+1,0 − w∗‖22] + K−1∑ k=1 E[‖ut,k − w∗‖22] ≤ (1− µηlocal) K−1∑ k=0 E[‖ut,k − w∗‖22]) + 8η3localLK3σ2. Rearranging the terms, we obtain for any t ≥ 1, E[‖ut+1,0 − w∗‖22] ≤ (1− µηlocal)E[‖ut,0 − w∗‖22] + 8η3localLK3σ2,
implying
E[‖wT+1 − w∗‖22] ≤ E[‖w0 − w∗‖22]e−µηlocalT + 8η2localLK3σ2/µ. We conclude by the L-smoothness of function f .
Corollary 5.7. If Assumption 3.1 holds with µ > 0. Then within Algorithm 1 outputs an -optimal solution wT ∈ Rd satisfying E[f(wT ) − f(w∗)] ≤ by using O((LN/µ) max{d, √ σ2/(µ )} log(LE[‖w0 − w∗‖22]/ )) bits of communication cost.
We observe the same phenomenon as in the single-step scheme again, that compared to vanilla approaches, ours shrinks the stepsize by a factor of O(α), thus enlarge the number of rounds approximately by a factor of O(α). Since our approach only communicates O(bsketch/d) as many bits per round due to sketching, the total communication cost does not increase at all for commonly used sketching matrices, according to Theorem 4.2.
We also point out that when ≥ σ2/(µd2), our analysis implies a linear convergence rate of local GD under only strongly-convex and smooth assumptions, which is new as far as we concern. We also have a similar observation in the convex losses case, as shown in the below theorem.
Theorem 5.8. If Assumption 3.1 holds with µ = 0. If ηlocal ≤ 18(1+α)LK ,
E[f(wT )− f(w∗)] ≤ 4E[‖w 0 − w∗‖22]
ηlocalKT + 32η2localLK 2σ2,
where wT = 1KT ( ∑T t=1 ∑K−1 k=0 u
t,k) is the average over parameters throughout the execution of Algorithm 1.
Proof. Telescoping sum up Lemma 5.4 as t varies from 0 to T −1 and k varies from 0 to K−1 and absorb the higher-order terms by the choice of the stepsize,
E[‖uT+1,0 − w∗‖22]− E[‖w0 − w∗‖22] ≤ − 1
4 ηlocal T∑ t=1 K−1∑ k=0 E[f(ut,k)− f(w∗)] + 8η3localLK3Tσ2.
Rearranging the above equation, we have
1
KT T∑ t=1 K−1∑ k=0 E[f(ut,k)− f(w∗)] ≤ 4E[‖w 0 − w∗‖22] ηlocalKT + 32η2localLK 2σ2.
By the convexity of f we complete the proof.
Corollary 5.9. If Assumption 3.1 holds with µ = 0. Then Algorithm 1 outputs an -optimal solution wT ∈ Rd satisfying E[f(wT )− f(w∗)] ≤ by using O(E[‖w0−w∗‖22]N max{Ld/ , σ √ L/ 3/2}) bits of communication cost.
We compare our communication cost with the work of Khaled et al. (2019), which analyzes the local gradient descent using the same assumption and framework. The result of Khaled et al. (2019) shows a communication cost of O ( E[‖w0 − w∗‖22]Ndmax{L , σ √ L 3/2 } )
, which is strictly no less than our results. This shows again our approach does not introduce extra overall communication cost.

heading: 6 DISCUSSION
text: In this work, we propose the iterative sketch-based federated learning framework, which only communicates the sketched gradients with noises. Such framework enjoys the benefits of both better privacy and lower communication cost per round. We also rigorously prove that the randomness from sketching will not introduce extra overall communication cost.
Though our framework is built upon the local gradient descent algorithm and our theoretical discussion follows the analysis framework of Khaled et al. (2019; 2020), we emphasize that our approach and results can be extended to other gradient-based optimization algorithms and analysis, including but not limited to gradient descent with momentum and local stochastic gradient descent. The key reason is due to the sketched and de-sketched gradient R>Rg is an unbiased estimator of the true gradient g with second moments being a multiplier of ‖g‖22. As the iterates approach the optimal solution, the second moments approaches 0 correspondingly, resulting an exact match of the vanilla approach. Therefore, by scaling down the stepsize appropriately, we are able to recover the same convergence guarantee as the original gradient-based algorithms.
By a simple modification to our algorithm with additive Gaussian noises on the sketched gradients, we can also prove the differential privacy of our learning system by “hiding” the most important component in the system for guarding the safety and privacy. This additive noise also does not affect the convergence behavior of our algorithm too much, since it does make the estimator biased, and the additive variance can be factored into our original analysis.
Despite the benefits in terms of privacy and communication cost, we point out our approach does have trade-off on computation complexity. Since we need to shrink stepsize and thus enlarge the number of iterations to achieve certain accuracy level, our approach requires O(d/bsketch) as many computational cost compared to vanilla approaches. However, in a privacy-and-communication focused distributed learning scenario, we hope this work provides a new solution and motivates future works.
Ethics Statement. This paper mainly focuses on the theoretical perspective of federated learning, and it proposes algorithm to improve the privacy of the learning system.
Reproducibility Statement. This paper contains several theoretical results, for discussions related to sketching, we refer readers to section D, for discussions related to the convergence analysis and communication costs of Algorithm 1, we refer readers to section F and G. For discussions related to differential privacy, we refer readers to section H.

heading: 1 Introduction 1

heading: 2 Related Work 2

heading: 3 Problem Setup 3

heading: 4 Our Algorithm 3
text: 4.1 sk/desk via coordinate wise embedding . . . . . . . . . . . . . . . . . . . . . . . 4
4.2 Sketch-and-Solve vs Sketch-and-De-sketch . . . . . . . . . . . . . . . . . . . . . 5
4.3 Privacy-preserved sketching via low-dimensional noises . . . . . . . . . . . . . . . 5

heading: 5 Convergence Theory and Communication Complexity 6
text: 5.1 Single-step scheme . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 6
5.2 Multi-step scheme . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7

heading: 6 Discussion 9

heading: A Preliminary 16

heading: B Probability 16

heading: C Optimization Backgrounds 17

heading: D Sketching 18
text: D.1 Definition . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18
D.2 Coordinate wise embedding . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 19
D.3 Expectation and variance . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 19
D.4 Bounding inner product . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23
D.5 Infinite norm bound . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25

heading: E Analysis of convergence: Single-step scheme 28
text: E.1 Preliminary . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 28
E.2 Strongly-convex f convergence analysis . . . . . . . . . . . . . . . . . . . . . . . 29
E.3 Convex f convergence analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . 29
E.4 Non-convex f convergence analysis . . . . . . . . . . . . . . . . . . . . . . . . . 30
F k-step convex & strongly-convex fc analysis 31
F.1 Preliminary . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 31
F.2 Unifying the update rule of Algorithm 1 . . . . . . . . . . . . . . . . . . . . . . . 32
F.3 Upper bounding ‖gt,k‖22 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 32 F.4 Lower bounding 〈ut,k − w∗, gt,k〉 . . . . . . . . . . . . . . . . . . . . . . . . . . 33
F.5 Upper bounding variance within K local steps . . . . . . . . . . . . . . . . . . . . 34
F.6 Bounding the expected gap between ut,k and w∗ . . . . . . . . . . . . . . . . . . 35
F.7 Main result: convex case . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 36
F.8 Main result: strongly-convex case . . . . . . . . . . . . . . . . . . . . . . . . . . 38
G k-step non-convex f convergence analysis 39
H Differential Privacy 42
Roadmap. We organize the appendix as follows. In section A, we introduce some notations and definitions that will be used across the appendix. In section B, we study several probability tools we will be using in the proof of cretain properties of various sketching matrices. In section C, we lay out some key assumptions on local objective function fc and global objective function f , in order to proceed our discussion of convergence theory. In section E, we give complete proofs for single-step scheme. We dedicate sections F and G to illustrate formal analysis of the convergence results of Algorithm 1 under k local steps, given different assumptions of objective function f . In section H, we introduce additive noise to make our sketched gradients differential private and prove it for specific AMS sketch matrix.

heading: A PRELIMINARY
text: For a positive integer n, we use [n] to denote the set {1, 2, · · · , n}. We use E[·] to denote expectation (if it exists), and use Pr[·] to denote probability. For a function f , we use Õ(f) to denote f poly(log f). For a vector x, For a vector x, we use ‖x‖1 := ∑ i |xi| to denote its `1 norm, we
use ‖x‖2 := ( ∑n i=1 x 2 i )
1/2 to denote its `2 norm, we use ‖x‖∞ := maxi∈[n] |xi| to denote its `∞ norm. For a matrix A and a vector x, we define ‖x‖A := √ x>Ax. For a full rank square matrix A, we use A−1 to denote its true inverse. For a matrix A, we use A† to denote its pseudo-inverse. For a matrix A, we use ‖A‖ to denote its spectral norm. We use ‖A‖F := ( ∑ i,j A 2 i,j)
1/2 to denote its Frobenius norm. We use A> to denote the transpose of A. We denote 1{x=l} for l ∈ R to be the indicator function which equals to 1 if x = l and 0 otherwise. Let f : A → B and g : C → A be two functions, we use f ◦ g to denote the composition of functions f and g, i.e., for any x ∈ C, (f ◦ g)(x) = f(g(x)).

heading: B PROBABILITY
text: Lemma B.1 (Chernoff bound Chernoff (1952)). Let X = ∑n i=1Xi, where Xi = 1 with probability
pi and Xi = 0 with probability 1−pi, and all Xi are independent. Let µ = E[X] = ∑n i=1 pi. Then 1. Pr[X ≥ (1 + δ)µ] ≤ exp(−δ2µ/3), ∀δ > 0 ; 2. Pr[X ≤ (1− δ)µ] ≤ exp(−δ2µ/2), ∀0 < δ < 1. Lemma B.2 (Hoeffding bound Hoeffding (1963)). Let X1, · · · , Xn denote n independent bounded variables in [ai, bi]. Let X = ∑n i=1Xi, then we have
Pr[|X − E[X]| ≥ t] ≤ 2 exp ( − 2t
2∑n i=1(bi − ai)2
) .
Lemma B.3 (Bernstein inequality Bernstein (1924)). Let X1, · · · , Xn be independent zero-mean random variables. Suppose that |Xi| ≤M almost surely, for all i. Then, for all positive t,
Pr [ n∑ i=1 Xi > t ] ≤ exp ( − t 2/2∑n j=1 E[X2j ] +Mt/3 ) .
Lemma B.4 (Khintchine’s inequality, Khintchine (1923); Haagerup (1981)). Let σ1, · · · , σn be i.i.d. sign random variables, and let z1, · · · , zn be real numbers. Then there are constants C > 0 so that for all t > 0
Pr [∣∣∣ n∑ i=1 ziσi ∣∣∣ ≥ t‖z‖2] ≤ exp(−Ct2). Lemma B.5 (Hason-wright inequality Hanson & Wright (1971); Rudelson & Vershynin (2013)). Let x ∈ Rn denote a random vector with independent entries xi with E[xi] = 0 and |xi| ≤ K. Let A be an n× n matrix. Then, for every t ≥ 0, Pr[|x>Ax− E[x>Ax]| > t] ≤ 2 · exp(−cmin{t2/(K4‖A‖2F ), t/(K2‖A‖)}). Lemma B.6 (Lemma 1 on page 1325 of Laurent and Massart Laurent & Massart (2000)). Let X ∼ X 2k be a chi-squared distributed random variable with k degrees of freedom. Each one has zero mean and σ2 variance. Then
Pr[X − kσ2 ≥ (2 √ kt+ 2t)σ2] ≤ exp(−t),
Pr[kσ2 −X ≥ 2 √ ktσ2] ≤ exp(−t).
Lemma B.7 (Tail bound for sub-exponential distribution Foss et al. (2011)). We sayX ∈ SE(σ2, α) with parameters σ > 0, α > 0 if:
E[eλX ] ≤ exp(λ2σ2/2), ∀|λ| < 1/α.
Let X ∈ SE(σ2, α) and E[X] = µ, then:
Pr[|X − µ| ≥ t] ≤ exp(−0.5 min{t2/σ2, t/α}).
Lemma B.8 (Matrix Chernoff bound Tropp (2011); Lu et al. (2013)). Let X be a finite set of positive-semidefinite matrices with dimension d× d, and suppose that
max X∈X
λmax(X) ≤ B.
Sample {X1, · · · , Xn} uniformly at random from X without replacement. We define µmin and µmax as follows:
µmin := n · λmin( E X∼X [X]) and µmax := n · λmax( E X∼X [X]).
Then
Pr [ λmin( n∑ i=1 Xi) ≤ (1− δ)µmin ] ≤ d · exp(−δ2µmin/B) for δ ∈ [0, 1),
Pr [ λmax( n∑ i=1 Xi) ≥ (1 + δ)µmax ] ≤ d · exp (−δ2µmax/(4B)) for δ ≥ 0.

heading: C OPTIMIZATION BACKGROUNDS
text: Definition C.1. Let f : Rd → R be a function, we say f is L-smooth if for any x, y ∈ Rd, we have
‖∇f(x)−∇f(y)‖2 ≤ L‖x− y‖2
Equivalently, for any x, y ∈ Rd, we have
f(y) ≤ f(x) + 〈y − x,∇f(x)〉+ L 2 ‖y − x‖22
Definition C.2. Let f : Rd → R be a function, we say f is convex if for any x, y ∈ Rd, we have
f(x) ≥ f(y) + 〈x− y,∇f(y)〉
Definition C.3. Let f : Rd → R be a function, we say f is µ-strongly-convex if for any x, y ∈ Rd, we have
‖∇f(x)−∇f(y)‖2 ≥ µ‖x− y‖2
Equivalently, for any x, y ∈ Rd, we have
f(y) ≥ f(x) + 〈y − x,∇f(x)〉+ µ 2 ‖y − x‖22
Fact C.4. Let f : Rd → R be an L-smooth and convex function, then for any x, y ∈ Rd, we have
f(y)− f(x) ≥ 〈y − x,∇f(x)〉+ 1 2L · ‖∇f(y)−∇f(x)‖22
Fact C.5 (Inequality 4.12 in Bottou et al. (2018)). Let f : Rd → R be a µ-strongly convex function. Let x∗ be the minimizer of f . Then for any x ∈ Rd, we have
f(x)− f(x∗) ≤ 1 2µ ‖∇f(x)‖22

heading: D SKETCHING
text: In this section, we discuss the (α, β, δ)-coordinate wise embedding property we proposed in this work through several commonly used sketching matrices.
We consider several standard sketching matrices:
1. Random Gaussian matrices.
2. Subsampled randomized Hadamard/Fourier transform matrices Lu et al. (2013).
3. AMS sketch matrices Alon et al. (1999), random {−1,+1} per entry. 4. Count-Sketch matrices Charikar et al. (2002), each column only has one non-zero entry,
and is −1,+1 half probability each. 5. Sparse embedding matrices Nelson & Nguyên (2013), each column only has s non-zero
entries, and each entry is − 1√ s ,+ 1√ s half probability each.
6. Uniform sampling matrices.
We list the definitions and results of above sketching matrices for coordinate-wise embedding in Table 1.
Definition D.1 (k-wise independence). H = {h : [m] → [l]} is a k-wise independent hash family if ∀i1 6= i2 6= · · · 6= ik ∈ [n] and ∀j1, · · · , jk ∈ [l],
Pr h∈H
[h(i1) = j1 ∧ · · · ∧ h(ik) = jk] = 1
lk .
Definition D.2 (Random Gaussian matrix). We say R ∈ Rb×n is a random Gaussian matrix if all entries are sampled from N (0, 1/b) independently. Definition D.3 (Subsampled randomized Hadamard/Fourier transform matrix Lu et al. (2013)). We say R ∈ Rb×n is a subsampled randomized Hadamard transform matrixii if it is of the form R =√ n/bSHD, where S ∈ Rb×n is a random matrix whose rows are b uniform samples (without replacement) from the standard basis of Rn, H ∈ Rn×n is a normalized Walsh-Hadamard matrix, and D ∈ Rn×n is a diagonal matrix whose diagonal elements are i.i.d. Rademacher random variables.
Definition D.4 (AMS sketch matrix Alon et al. (1999)). Let h1, h2, · · · , hb be b random hash functions picking from a 4-wise independent hash family H = {h : [n] → {− 1√
b ,+ 1√ b }}. Then
R ∈ Rb×n is a AMS sketch matrix if we set Ri,j = hi(j). Definition D.5 (Count-sketch matrix Charikar et al. (2002)). Let h : [n] → [b] be a random 2-wise independent hash function and σ : [n]→ {−1,+1} be a random 4-wise independent hash function. Then R ∈ Rb×n is a count-sketch matrix if we set Rh(i),i = σ(i) for all i ∈ [n] and other entries to zero.
iiIn this case, we require logn to be an integer.
Definition D.6 (Sparse embedding matrix I Nelson & Nguyên (2013)). We say R ∈ Rb×n is a sparse embedding matrix with parameter s if each column has exactly s non-zero elements being ±1/ √ s uniformly at random, whose locations are picked uniformly at random without replacement (and independent across columns) iii. Definition D.7 (Sparse embedding matrix II Nelson & Nguyên (2013)). Let h : [n]× [s]→ [b/s] be a a ramdom 2-wise independent hash function and σ : [n]× [s]→ {−1, 1} be a 4-wise independent. Then R ∈ Rb×n is a sparse embedding matrix II with parameter s if we set R(j−1)b/s+h(i,j),i = σ(i, j)/ √ s for all (i, j) ∈ [n]× [s] and all other entries to zero.iv Definition D.8 (Uniform sampling matrix). We say R ∈ Rb×n is a uniform sampling matrix if it is of the form R = √ n/bSD, where S ∈ Rb×n is a random matrix whose rows are b uniform samples (without replacement) from the standard basis of Rn, and D ∈ Rn×n is a diagonal matrix whose diagonal elements are i.i.d. Rademacher random variables.
D.2 COORDINATE WISE EMBEDDING
We define coordinate-wise embedding as follows Definition D.9 ((α, β, δ)-coordinate wise embedding). We say a randomized matrix R ∈ Rb×n satisfying (α, β, δ)-coordinate wise embedding if
1. E R∼Π
[g>R>Rh] = g>h,
2. E R∼Π [(g>R>Rh)2] ≤ (g>h)2 + α b ‖g‖22‖h‖22,
3. Pr R∼Π
[ |g>R>Rh− g>h| ≥ β√
b ‖g‖2‖h‖2
] ≤ δ.
Remark D.10. Given a randomized matrix R ∈ Rb×n satisfying (α, β, δ)-coordinate wise embedding and any orthogonal projection P ∈ Rn×n, above definition implies
1. E R∼Π
[PR>Rh] = Ph,
2. E R∼Π
[(PR>Rh)2i ] ≤ (Ph)2i + α
b ‖h‖22,
3. Pr R∼Π
[ |(PR>Rh)i − (Ph)i| ≥
β√ b ‖h‖2
] ≤ δ.
since ‖P‖2 ≤ 1 implies ‖Pi,:‖2 ≤ 1 for all i ∈ [n].
D.3 EXPECTATION AND VARIANCE
Lemma D.11. Let R ∈ Rb×n denote any of the random matrix in Definition D.2, D.3, D.4, D.6, D.7, D.8. Then for any fixed vector h ∈ Rn and any fixed vector g ∈ Rn, the following properties hold:
E R∼Π
[g>R>Rh] = g>h
Proof.
E R∼Π [g>R>Rh] = g> E R∼Π [R>R]h = g>Ih = g>h.
Lemma D.12. LetR ∈ Rb×n denote a subsampled randomized Hadamard transform or AMS sketch matrix as in Definition D.3, D.4. Then for any fixed vector h ∈ Rn and any fixed vector g ∈ Rn, the following properties hold:
E R∼Π [(g>R>Rh)2] ≤ (g>h)2 + 2 b ‖g‖22 · ‖h‖22.
iiiFor our purposes the signs need only be O(log d)-wise independent, and each column can be specified by a O(log d)-wise independent permutation, and the seeds specifying the permutations in different columns need only be O(log d)-wise independent.
ivThis definition has the same behavior as sparse embedding matrix I for our purpose.
Proof. If Ea[a] = b, it is easy to see that
E a [(a− b)2] = E a [a2 − 2ab+ b2] = E a [a2 − b2]
We can rewrite it as follows:
E R∼Π [(g>R>Rh)2 − (g>h)2] = E R∼Π [(g>(R>R− I)h)2],
It can be bounded as follows:
E R∼Π
[(g>(R>R− I)h)2]
= E R∼Π ( b∑ k=1 (Rg)k(Rh)k − g>h )2 = E
R∼Π   b∑ k=1 n∑ i=1 Rk,igi · ∑ j∈[n]\{i} Rk,jhj 2 
= E R∼Π  b∑ k=1 n∑ i=1 Rk,igi · ∑ j∈[n]\{i} Rk,jhj  ·  b∑ k′=1 n∑ i′=1 Rk′,i′gi′ · ∑ j′∈[n]\{i′} Rk′,j′hj′  = E
R∼Π  b∑ k=1 n∑ i=1 R2k,ig 2 i · ∑ j∈[n]\{i} R2k,jh 2 j +  b∑ k=1 n∑ i=1 R2k,igihi · ∑ j∈[n]\{i} R2k,jgjhj  = 1
b  n∑ i=1 g2i ∑ j∈[n]\{i} h2j + 1 b  n∑ i=1 gihi ∑ j∈[n]\{i} gjhj  ≤ 2 b ‖g‖22‖h‖22,
where the second step follows from R2k,i = 1/b, ∀k, i ∈ [b] × [n], the forth step follows from E[Rk,iRk,jRk′,i′Rk′,j′ ] 6= 0 only if i = i′, j = j′, k = k′ or i = j′, j = i′, k = k′, the fifth step follows from Rk,i and Rk,j are independent if i 6= j and R2k,i = R2k,j = 1/b, and the last step follows from Cauchy-Schwartz inequality.
Therefore,
E R∼Π [(g>R>Rh)2 − (g>h)2] = E R∼Π [(g>(R>R− I)h)2] ≤ 2 b ‖g‖22‖h‖22.
Lemma D.13. Let R ∈ Rb×n denote a random Gaussian matrix as in Definition D.2. Then for any fixed vector h ∈ Rn and any fixed vector g ∈ Rn, the following properties hold:
E R∼Π [(g>R>Rh)2] ≤ (g>h)2 + 3 b ‖g‖22 · ‖h‖22.
Proof. Note
E R∼Π
[(g>R>Rh)2]
= E R∼Π   b∑ k=1 n∑ i=1 Rk,igi · n∑ j=1 Rk,jhj 2 
= E R∼Π  b∑ k=1 n∑ i=1 Rk,igi · n∑ j=1 Rk,jhj  ·  b∑ k′=1 n∑ i′=1 Rk′,i′gi′ · n∑ j′=1 Rk′,j′hj′ 
= E R∼Π [ b∑ k=1 ∑ k′∈[b]\{k} n∑ i=1 n∑ i′=1 R2k,iR 2 k′,i′gihigi′hi′ +( b∑ k=1 n∑ i=1 R4k,ig 2 i h 2 i )
+  b∑ k=1 n∑ i=1 ∑ j∈[n]\{i} R2k,iR 2 k,jg 2 i h 2 j +  n∑ k=1 n∑ i=1 n∑ i′∈[n]\{i} R2k,iR 2 k,i′gihigi′hi′  +
 b∑ k=1 n∑ i=1 ∑ j∈[n]\{i} R2k,iR 2 k,jgihjgjhi ]
= b− 1 b n∑ i=1 n∑ i′=1 gihigi′hi′ + 3 b n∑ i=1 g2i h 2 i
+ 1
b n∑ i=1 ∑ j∈[n]\[i] g2i h 2 j + 1 b n∑ i=1 ∑ i′∈[n]\[i] gihigi′hi′ + 1 b n∑ i=1 ∑ j∈[n]\[i] gihjgjhi
≤ (g>h)2 + 3 b ‖g‖22‖h‖22,
where the third step follows from that for independent entries of a random Gaussian matrix, E[Rk,iRk,jRk′,i′Rk′,j′ ] 6= 0 only if 1. k 6= k′, i = j, i′ = j′, or 2. k = k′, i = i′ = j = j′, or 3. k = k′, i = i′ 6= j = j′, or 4. k = k′, i = j 6= i′ = j′, or 5. k = k′, i = j′ 6= i′ = j, the fourth step follows from E[R2k,i] = 1/b and E[R4k,i] = 3/b2, and the last step follows from Cauchy-Schwartz inequality.
Lemma D.14. Let R ∈ Rb×n denote a count-sketch matrix as in Definition D.5. Then for any fixed vector h ∈ Rn and any fixed vector g ∈ Rn, the following properties hold:
E R∼Π [(g>R>Rh)2] ≤ (g>h)2 + 3 b ‖g‖22‖h‖22.
Proof. Note
E R∼Π
[(g>R>Rh)2]
= E R∼Π   b∑ k=1 n∑ i=1 Rk,igi n∑ j=1 Rk,jhj 2 
= E R∼Π  b∑ k=1 n∑ i=1 Rk,igi n∑ j=1 Rk,jhj  ·  b∑ k′=1 n∑ i′=1 Rk′,i′gi′ n∑ j′=1 Rk′,j′hj′  = E
R∼Π [ b∑ k=1 ∑ k′∈[b]\{k} n∑ i=1 n∑ i′∈[n]\{i} R2k,iR 2 k′,i′gihigi′hi′ +( b∑ k=1 n∑ i=1 R4k,ig 2 i h 2 i )
+  b∑ k=1 n∑ i=1 ∑ j∈[n]\{i} R2k,iR 2 k,jg 2 i h 2 j +  n∑ k=1 n∑ i=1 n∑ i′∈[n]\{i} R2k,iR 2 k,i′gihigi′hi′  +
 b∑ k=1 n∑ i=1 ∑ j∈[n]\{i} R2k,iR 2 k,jgihjgjhi ]
= b− 1 b n∑ i=1 ∑ i′∈[n]\i gihigi′hi′ + n∑ i=1 g2i h 2 i
+ 1
b n∑ i=1 ∑ j∈[n]\{i} g2i h 2 j + 1 b n∑ i=1 ∑ i′∈[n]\{i} gihigi′hi′ + 1 b n∑ i=1 ∑ j∈[n]\{i} gihjgjhi
≤ (g>h)2 + 3 b ‖g‖22‖h‖22,
where in the third step we are again considering what values of k, k′, i, i′, j, j′ that makes E[Rk,iRk,jRk′,i′Rk′,j′ ] 6= 0. Since the hash function σ(·) of the count-sketch matrix is 4-wise independent, ∀k, k′, when i 6= i′ 6= j 6= j′, or i = i′ = j 6= j′ (and the other 3 symmetric cases), we have that E[Rk,iRk,jRk′,i′Rk′,j′ ] = 0. Since the count-sketch matrix has only one non-zero entry in every column, when k 6= k′, if i = i′ or i = j′ or j = i′ or j = j′, we also have E[Rk,iRk,jRk′,i′Rk′,j′ ] = 0. Thus we only need to consider the cases: 1. k 6= k′, i = j 6= i′ = j′, or 2. k = k′, i = i′ = j = j′, or 3. k = k′, i = i′ 6= j = j′, or 4. k = k′, i = j 6= i′ = j′, or 5. k = k′, i = j′ 6= i′ = j. And the fourth step follows from E[R2k,i] = 1/b and E[R4k,i] = 1/b, and the last step follows from Cauchy-Schwartz inequality.
Lemma D.15. Let R ∈ Rb×n denote a sparse embedding matrix as in Definition D.6, D.7. Then for any fixed vector h ∈ Rn and any fixed vector g ∈ Rn, the following properties hold:
2. E R∼Π [(g>R>Rh)2] ≤ (g>h)2 + 2 b ‖g‖22 · ‖h‖22.
Proof. Note
E R∼Π
[(g>R>Rh)2]
= E R∼Π   b∑ k=1 n∑ i=1 Rk,igi n∑ j=1 Rk,jhj 2 
= E R∼Π  b∑ k=1 n∑ i=1 Rk,igi n∑ j=1 Rk,jhj  ·  b∑ k′=1 n∑ i′=1 Rk′,i′gi′ n∑ j′=1 Rk′,j′hj′  = E
R∼Π [ b∑ k=1 n∑ i=1 R2k,ig 2 i ∑ j∈[n]\{i} R2k,jh 2 j +  b∑ k=1 n∑ i=1 R2k,igihi ∑ j∈[n]\{i} R2k,jgjhj  +
∑ k ∑ i 6=i′ R2k,iR 2 k,i′gihigi′hi′ +(∑ k ∑ i R4k,ig 2 i h 2 i ) + ∑ k 6=k′ ∑ i 6=i′ R2k,iR 2 k′,i′gihigi′hi′  +
∑ k 6=k′ ∑ i R2k,iR 2 k′,ig 2 i h 2 i ] = 1
b ∑ i 6=j g2i h 2 j + 1 b ∑ i 6=j gihigjhj + 1 b ∑ i6=i′ gihigi′hi′ + 1 s ∑ i g2i h 2 i + b− 1 b ∑ i6=i′ gihigi′hi′ + s− 1 s ∑ i g2i h 2 i
≤ (g>h)2 + 2 b ‖g‖22‖h‖22,
where the third step follows from the fact that the sparse embedding matrix has independent columns and s non-zero entry in every column, the fourth step follows from E[R2k,i] = 1/b, E[R4k,i] = 1/(bs), and E[R2k,iR2k′,i] = s(s−1) b(b−1) · 1 s2 ,∀k 6= k
′ and the last step follows from Cauchy-Schwartz inequality.
Lemma D.16. Let R ∈ Rb×n denote a uniform sampling matrix as in Definition D.8. Then for any fixed vector h ∈ Rn and any fixed vector g ∈ Rn, the following properties hold:
2. E R∼Π [(g>R>Rh)2] ≤ (g>h)2 + n b ‖g‖22‖h‖22.
Proof. Note
E R∼Π
[(g>R>Rh)2]
= E R∼Π   b∑ k=1 n∑ i=1 Rk,igi n∑ j=1 Rk,jhj 2 
= E R∼Π  b∑ k=1 n∑ i=1 Rk,igi n∑ j=1 Rk,jhj  ·  b∑ k′=1 n∑ i′=1 Rk′,i′gi′ n∑ j′=1 Rk′,j′hj′  = E
R∼Π (∑ k ∑ i R4k,ig 2 i h 2 i ) + ∑ k 6=k′ ∑ i 6=i′ R2k,iR 2 k′,i′gihigi′hi′  = n
b ∑ i g2i h 2 i + (b− 1)n (n− 1)b ∑ i 6=i′ gihigi′hi′
≤ (g>h)2 + n b ‖g‖22‖h‖22,
where the third step follows from the fact that the random sampling matrix has one non-zero entry in every row, the fourth step follows from E[R2k,iR2k′,i′ ] = n/((n − 1)b2) for k 6= k′, i 6= i′ and E[R4k,i] = n/b2.
Remark D.17. Lemma D.16 indicates that uniform sampling fails in bounding variance in some sense, since the upper bound give here involves n.
D.4 BOUNDING INNER PRODUCT
Lemma D.18 (Gaussian). Let R ∈ Rb×n be a random Gaussian matrix (Definition D.2). Then we have:
Pr [
max i6=j |〈R∗,i, R∗,j〉| ≥
√ log(n/δ)√
b
] ≤ Θ(δ).
Proof. Note for i 6= j, R∗,i, R∗,j ∼ N (0, 1b Ib) are two independent Gaussian vectors. Let zk = Rk,iRk,j and z = 〈R∗,i, R∗,j〉. Then we have for any |λ| ≤ b/2,
E[eλzk ] = 1√
1− λ2/b2 ≤ exp(λ2/b2),
where the first step follows from zk = 14 (Rk,i + Rk,j) 2 + 14 (Rk,i − Rk,j) 2 = b2 (Q1 −Q2) where Q1, Q2 ∼ χ21, and E[eλQ] = 1√1−2λ for any Q ∼ χ 2 1. This implies zk ∈ SE(2/b2, 2/b) is a sub-exponential random variable. Thus, we have z =∑b k=1 zk ∈ SE(2/b, 2/b), by sub-exponential concentration Lemma B.7 we have
Pr[|z| ≥ t] ≤ 2 exp(−bt2/4)
for 0 < t < 1. Picking t = √ log(n2/δ)/b, we have
Pr [ |〈R∗,i, R∗,j〉| ≥ c √
log(n/δ)√ b
] ≤ δ/n2.
Taking the union bound over all (i, j) ∈ [n]× [n] and i 6= j, we complete the proof.
Lemma D.19 (SRHT). Let R ∈ Rb×n be a subsample randomized Hadamard transform (Definition D.3). Then we have:
Pr [
max i6=j |〈R∗,i, R∗,j〉| ≥
√ log(n/δ)√
b
] ≤ Θ(δ).
Proof. For fixed i 6= j, let X = [R∗,i, R∗,j ] ∈ Rb×2. Then X>X = ∑b k=1Gk, where
Gk = [Rk,i, Rk,j ] >[Rk,i, Rk,j ] =
[ 1 b Rk,iRk,j
Rk,iRk,j 1 b
] .
Note the eigenvalues of Gk are 0 and 2b and E[X >X] = b · E[Gk] = I2 for all k ∈ [b]. Thus, applying matrix Chernoff bound B.8 to X>X we have
Pr [ λmax(X >X) ≤ 1− t ] ≤ 2 exp (−t2b/2) for t ∈ [0, 1), and
Pr [ λmax(X >X) ≥ 1 + t ] ≤ 2 exp (−t2b/8) for t ≥ 0.
which implies the eigenvalues of X>X are between [1− t, 1 + t] with probability 1−4 exp (− t 2b 8 ). So the eigenvalues of X>X − I2 are between [−t, t] with probability 1 − 4 exp (− t 2b 8 ). Picking
t = c √
log(n/δ)√ b , we have
Pr [ ‖X>X − I2‖ ≥ c √
log(n/δ)√ b ] ≤ δ n2 .
Note
X>X − I2 = [
0 〈R∗,i, R∗,j〉 〈R∗,i, R∗,j〉 0
] ,
whose spectral norm is |〈R∗,i, R∗,j〉|. Thus, we have
Pr [ |〈R∗,i, R∗,j〉| ≥ c √
log(n/δ)√ b
] ≤ δ/n2.
Taking a union bound over all pairs (i, j) ∈ [n]× [n] and i 6= j, we complete the proof.
Lemma D.20 (AMS). Let R ∈ Rb×n be a random AMS matrix (Definition D.4). Let {σi, i ∈ [n]} be independent Rademacher random variables and R ∈ Rb×n with R∗,i = σiR∗,i, ∀i ∈ [n]. Then we have:
Pr [
max i6=j |〈R∗,i, R∗,j〉| ≥
√ log(n/δ)√
b
] ≤ Θ(δ).
Proof. Note for any fixed i 6= j, R∗,i and R∗,j are independent. By Hoeffding inequality (Lemma B.2), we have
Pr [ |〈R∗,i, R∗,j〉| ≥ t ] ≤ 2 exp ( − 2t
2∑b i=1( 1 b − (− 1 b )) 2
) ≤ 2e−t 2b/2
Choosing t = √ 2 log(2n2/δ)/ √ b, we have
Pr [ |〈R∗,i, R∗,j〉| ≥ √ 2 log(2n2/δ)/ √ b ] ≤ δ n2 .
Taking a union bound over all pairs (i, j) ∈ [n]× [n] and i 6= j, we complete the proof.
Lemma D.21 (Count-Sketch). Let R ∈ Rb×n be a count-sketch matrix (Definition D.5). Let {σi, i ∈ [n]} be independent Rademacher random variables and R ∈ Rb×n with R∗,i = σiR∗,i, ∀i ∈ [n]. Then we have:
max i 6=j |〈R∗,i, R∗,j〉| ≤ 1.
Proof. Directly follow the definition of count-sketch matrices.
Lemma D.22 (Sparse embedding). Let R ∈ Rb×n be a sparse embedding matrix with parameter s (Definition D.6 and D.7). Let {σi, i ∈ [n]} be independent Rademacher random variables and R ∈ Rb×n with R∗,i = σiR∗,i, ∀i ∈ [n]. Then we have:
Pr [
max i 6=j |〈R∗,i, R∗,j〉| ≥
c √
log(n/δ)√ s
] ≤ Θ(δ).
Proof. Note for fixed i 6= j, R∗,i and R∗,j are independent. Assume R∗,i and R∗,j has u non-zero elements at the same positions, where 0 ≤ u ≤ s, then by Hoeffding inequality (Lemma B.2), we have
Pr[|〈R∗,i, R∗,j〉| ≥ t] ≤ 2 exp ( − 2t
2∑u i=1( 1 s − (− 1 s )) 2
) ≤ 2 exp(−t2s2/(2u)) (4)
Let t = √ (2u/s2) log(2n2/δ), we have
Pr [ |〈R∗,i, R∗,j〉| ≥ √ 2s−1 log(2n2/δ) ] ≤ Pr [ |〈R∗,i, R∗,j〉| ≥ √ 2us−2 log(2n2/δ) ] ≤ δ/n2 (5)
since u ≤ s. By taking a union bound over all (i, j) ∈ [n] × [n] and i 6= j, we complete the proof.
D.5 INFINITE NORM BOUND
Lemma D.23 (SRHT and AMS). Let R ∈ Rb×n denote a subsample randomized Hadamard transform (Definition D.3) or AMS sketching matrix (Definition D.4). Then for any fixed vector h ∈ Rn and any fixed vector g ∈ Rn, the following properties hold:
Pr R∼Π
[ |(g>R>Rh)− (g>h)| > log
1.5(n/δ)√ b
‖g‖2‖h‖2 ] ≤ Θ(δ).
Proof. We can rewrite (g>R>Rh)− (g>h) as follows:,
(g>R>Rh)− (g>h) = n∑ i=1 n∑ j∈[n]\i gihj〈R∗,i, R∗,j〉+ n∑ i=1 gihi(‖R∗,i‖22 − 1)
= n∑ i=1 n∑ j∈[n]\i gihj〈σiR∗,i, σjR∗,j〉.
where σi’s are independent Rademacher random variables and R∗,i = σiR∗,i, ∀i ∈ [n], and the second step follows from ‖R∗,i‖22 = 1,∀i ∈ [n]. We define matrix A ∈ Rn×n and B ∈ Rn×n as follows:
Ai,j = gihj · 〈R∗,i, R∗,j〉, ∀i ∈ [n], j ∈ [n] Bi,j = gihj ·max
i′ 6=j′ |〈R∗,i′ , R∗,j′〉| ∀i ∈ [n], j ∈ [n]
We defineA◦ ∈ Rn×n to be the matrixA ∈ Rn×n with removing diagonal entries, applying Hasonwright inequality (Lemma B.5), we have
Pr σ
[|σ>A◦σ| ≥ τ ] ≤ 2 · exp(−cmin{τ2/‖A◦‖2F , τ/‖A◦‖})
We can upper bound ‖A◦‖ and ‖A◦‖F . ‖A◦‖ ≤ ‖A◦‖F
≤ ‖A‖F ≤ ‖B‖F = ‖g‖2 · ‖h‖2 ·max
i 6=j |〈R∗,i, R∗,j〉|
≤ ‖g‖2 · ‖h‖2 ·max i 6=j |〈R∗,i, R∗,j〉|.
where the forth step follows from B is rank-1.
For SRHT, noteR has the same distribution as R. By Lemma D.19 (for AMS, we use Lemma D.20) with probability at least 1−Θ(δ), we have :
max i 6=j |〈R∗,i, R∗,j〉| ≤
√ log(n/δ)√
b .
Conditioning on the above event holds. Choosing τ = ‖g‖2 · ‖h‖2 · log1.5(n/δ)/ √ b, we can show that
Pr [∣∣∣(g>R>Rh)− (g>h)∣∣∣ ≥ ‖g‖2 · ‖h‖2 log1.5(n/δ)√ b ] ≤ Θ(δ).
Thus, we complete the proof.
Lemma D.24 (Random Gaussian). Let R ∈ Rb×n denote a random Gaussian matrix (Definition D.2). Then for any fixed vector h ∈ Rn and any fixed vector g ∈ Rn, the following properties hold:
Pr R∼Π
[ |(g>R>Rh)− (g>h)| > log
1.5(n/δ)√ b
‖g‖2‖h‖2 ] ≤ Θ(δ).
Proof. We follow the same procedure as proving Lemma D.23.
We can rewrite (g>R>Rh)− (g>h) as follows:,
(g>R>Rh)− (g>h) = n∑ i=1 n∑ j∈[n]\i gihj〈R∗,i, R∗,j〉+ n∑ i=1 gihi(‖R∗,i‖22 − 1)
= n∑ i=1 n∑ j∈[n]\i gihj〈σiR∗,i, σjR∗,j〉+ n∑ i=1 gihi(‖R∗,i‖22 − 1). (6)
where σi’s are independent Rademacher random variables and R has the same distribution as R. To bound the first term ∑n i=1 ∑n j∈[n]\i gihj〈σiR∗,i, σjR∗,j〉, we define matrix A ∈ Rn×n and B ∈ Rn×n as follows:
Ai,j = gihj · 〈R∗,i, R∗,j〉, ∀i ∈ [n], j ∈ [n] Bi,j = gihj ·max
i′ 6=j′ |〈R∗,i′ , R∗,j′〉| ∀i ∈ [n], j ∈ [n]
We defineA◦ ∈ Rn×n to be the matrixA ∈ Rn×n with removing diagonal entries, applying Hasonwright inequality (Lemma B.5), we have
Pr σ
[|σ>A◦σ| ≥ τ ] ≤ 2 · exp(−cmin{τ2/‖A◦‖2F , τ/‖A◦‖})
We can upper bound ‖A◦‖ and ‖A◦‖F .
‖A◦‖ ≤ ‖A◦‖F ≤ ‖A‖F ≤ ‖B‖F = ‖g‖2 · ‖h‖2 ·max
i 6=j |〈R∗,i, R∗,j〉|
≤ ‖g‖2 · ‖h‖2 ·max i 6=j |〈R∗,i, R∗,j〉|.
where the forth step follows from B is rank-1.
Using Lemma D.18 with probability at least 1−Θ(δ), we have :
max i 6=j |〈R∗,i, R∗,j〉| ≤
√ log(n/δ)√
b .
Conditioning on the above event holds. Choosing τ = ‖g‖2 · ‖h‖2 · log1.5(n/δ)/ √ b, we can show that
Pr ∣∣∣ n∑ i=1 n∑ j∈[n]\i gihj〈σiR∗,i, σjR∗,j〉 ∣∣∣ ≥ ‖g‖2 · ‖h‖2 log1.5(n/δ)√ b  ≤ Θ(δ). (7) To bound the second term ∑n i=1 gihi(‖R∗,i‖22 − 1), note that b‖R∗,i‖22 ∼ χ2b for every i ∈ [n]. Applying Lemma B.6, we have
Pr [∣∣∣‖R∗,i‖22 − 1∣∣∣ ≥ c√log(n/δ)√ b ] ≤ δ/n.
which implies
Pr [ n∑ i=1 gihi ∣∣∣‖R∗,i‖22 − 1∣∣∣ ≥ ‖g‖2‖h‖2 c√log(n/δ)√ b ] ≤ Θ(δ). (8)
Plugging the bounds Eq. (7) and (8) back to Eq. (6), we complete the proof.
Lemma D.25 (Count-sketch). Let R ∈ Rb×n denote a count-sketch matrix (Definition D.5). Then for any fixed vector h ∈ Rn and any fixed vector g ∈ Rn, the following properties hold:
Pr R∼Π
[ |(g>R>Rh)− (g>h)| ≥ log(1/δ)‖g‖2‖h‖2 ] ≤ Θ(δ).
Proof. We follow the identical procedure as proving Lemma D.23 to apply Hason-wright inequality (Lemma B.5).
Then note Lemma D.21 shows
max i6=j |〈R∗,i, R∗,j〉| ≤ 1
Thus, choosing τ = c‖g‖2 · ‖h‖2 · log(1/δ), we can show that Pr [ |(g>R>Rh)− (g>h)| ≥ c‖g‖2 · ‖h‖2 log(1/δ) ] ≤ δ.
which completes the proof.
Lemma D.26 (Count-sketch 2). Let R ∈ Rb×n denote a count-sketch matrix (Definition D.5). Then for any fixed vector h ∈ Rn and any fixed vector g ∈ Rn, the following properties hold:
Pr R∼Π
[ |(g>R>Rh)− (g>h)| ≥ 1√
bδ ‖g‖2‖h‖2
] ≤ Θ(δ).
Proof. It is known that a count-sketch matrix with b = −2δ−1 rows satisfies the ( , δ, 2)-JL moment property (see e.g. Theorem 14 of Woodruff (2014)). Using Markov’s inequality, ( , δ, 2)-JL moment property implies
Pr R∼Π
[ |(g>R>Rh)− (g>h)| ≥ ‖g‖2‖h‖2 ] ≤ Θ(δ),
where = 1√ bδ .
Remark D.27. In LP solver, we need δ = 1/ poly(n), thus Lemma D.25 is stronger than Lemma D.26.
Lemma D.28 (Sparse embedding). Let R ∈ Rb×n denote a sparse-embedding matrix (Definition D.6 and D.7). Then for any fixed vector h ∈ Rn and any fixed vector g ∈ Rn, the following properties hold:
3. Pr R∼Π
[ |(g>R>Rh)− (g>h)| > log
1.5(n/δ)√ s
‖g‖2‖h‖2 ] ≤ Θ(δ).
Proof. We follow the identical procedure as proving Lemma D.23 to apply Hason-wright inequality (Lemma B.5).
Then note Lemma D.22 shows with probability at least 1− δ we have
max i 6=j |〈R∗,i, R∗,j〉| ≤
c √
log(n/δ)√ s .
Conditioning on the above event holds, choosing τ = c′‖g‖2 · ‖h‖2 · log1.5(1/δ), we can show that
Pr [ |(g>R>Rh)− (g>h)| ≥ c
′ log1.5(n/δ)√ s
‖g‖2 · ‖h‖2 ] ≤ Θ(δ).
Thus, we complete the proof.
Lemma D.29 (Uniform sampling). Let R ∈ Rb×n denote a uniform sampling matrix (Definition D.8). Then for any fixed vector h ∈ Rn and any fixed vector g ∈ Rn, the following properties hold:
3.|(g>R>Rh)− (g>h)| ≤ (1 + n b )‖g‖2‖h‖2
where I ⊂ [n] be the subset of indexes chosen by the uniform sampling matrix.
Proof. We can rewrite (g>R>Rh)− (g>h) as follows:,
(g>R>Rh)− (g>h) = n∑ i=1 n∑ j∈[n]\i gihj〈R∗,i, R∗,j〉+ n∑ i=1 gihi(‖R∗,i‖22 − 1)
= n
b ∑ i∈I gihi − n∑ i=1 gihi.
where the second step follows from the uniform sampling matrix has only one nonzero entry in each row.
Let I ⊂ [n] be the subset chosen by the uniform sampling matrix, then ‖R∗,i‖22 = n/b for i ∈ I and ‖R∗,i‖22 = 0 for i ∈ [n] \ I . So we have
|(g>R>Rh)− (g>h)| = ∣∣∣∑ i∈I gihi( n b − 1)− ∑ i∈[n]\I gihi ∣∣∣ ≤ (1 + n
b )‖g‖2‖h‖2.

heading: E ANALYSIS OF CONVERGENCE: SINGLE-STEP SCHEME
text: E.1 PRELIMINARY
Throughout the proof of convergence, we will use Ft to denote the sequence wt−1, wt−2, . . . , w0. Also, we use η as a shorthand for ηglobal · ηlocal.
E.2 STRONGLY-CONVEX f CONVERGENCE ANALYSIS
Theorem E.1. Let f : Rd → R satisfying Assumption 3.1 with µ > 0. Let w∗ ∈ Rd be the optimal solution to f and assume sk/desk functions satisfying Theorem 4.2. Suppose η := ηglobal ·ηlocal has the property that η ≤ 1(1+α)L , then
E[f(wt+1)]− f(w∗) ≤ (1− µη)t · (f(w0)− f(w∗))
Proof. We shall first bound f(wt+1)− f(wt):
f(wt+1)− f(wt) ≤〈wt+1 − wt,∇f(wt)〉+ L 2 ‖wt+1 − wt‖22
=〈deskt(∆w̃t),∇f(wt)〉+ L
2 ‖deskt(∆w̃t)‖22
=− 〈ηglobal · deskt( 1
N N∑ c=1 skt(ηlocal · ∇fc(wt))),∇f(wt)〉
+ L
2 ‖ηglobal · deskt(
1
N N∑ c=1 skt(ηlocal · ∇fc(wt)))‖22
=− ηglobal · ηlocal · 〈deskt(skt(∇f(wt))),∇f(wt)〉 + (ηglobal · ηlocal)2 · ‖deskt(skt(∇f(wt)))‖22
where the first step uses the L-smoothness condition of f , and the last step uses the linearity property of sk/desk functions.
Taking expectation over iteration t conditioning on Ft and note that only wt+1 depends on randomness at t, we get
E[f(wt+1)− f(wt) | Ft]
≤ − η · 〈E[deskt(skt(∇f(wt))) | Ft],∇f(wt)〉+ Lη2
2 E[‖deskt(skt(∇f(wt)))‖22 | Ft]
≤ − η · 〈∇f(wt),∇f(wt)〉+ Lη 2
2 (1 + α) · ‖∇f(wt)‖22
≤ − η 2 · ‖∇f(wt)‖22 ≤ − µη · (f(wt)− f(w∗)) (9)
where the second step comes from the fact that deskt(skt(h)) is an unbiased estimator for any fixed h ∈ Rd and the bound on its variance, the third step comes from η ≤ 1(1+α)L , and the last step comes from Fact C.5.
Upon rearranging and subtracting both sides by f(w∗), we get
E[f(wt+1)]− f(w∗) | Ft] ≤ (1− µη) · (f(wt)− f(w∗)) (10)
Note that if we apply expectation over Ft on both sides of Eq. (10) we can get
E[f(wt+1)]− f(w∗) ≤ (1− µη) · (E[f(wt)]− f(w∗)) (11)
Notice since 1 − µη ≤ 1, this is a contraction map, if we iterate this recurrence relation, we will finally get
E[f(wt+1)− f(w∗)] ≤ (1− µη)t · (f(w0)− f(w∗)). (12)
E.3 CONVEX f CONVERGENCE ANALYSIS
Assume f is a convex function, we obtain a convergence bound in terms of the average of all parameters.
Theorem E.2. Let f : Rd → R satisfying Assumption 3.1 with µ = 0. Suppose sk/desk functions satisfying Theorem 4.2. If η := ηglobal · ηlocal ≤ 12(1+α)L , then
E[f(wT )− f(w∗)] ≤ E[‖w 0 − w∗‖22]
η · (T + 1) where wT := 1T+1 ∑T t=0 w t and w∗ ∈ Rd is the optimal solution.
Proof. We shall first compute the gap between wt+1 and w∗:
‖wt+1 − w∗‖22 = ‖wt − deskt(∆w̃t)− w∗‖22 = ‖wt − η · deskt(skt(∇f(wt)))− w∗‖22 = ‖wt − w∗‖22 + η2 · ‖deskt(skt(∇f(wt)))‖22 − 2η · 〈wt − w∗, deskt(skt(∇f(wt)))〉 (13)
By unbiasedness of deskt ◦ skt, we have
E[〈wt − w∗, deskt(skt(∇f(wt)))〉 | Ft] = E[〈wt − w∗,∇f(wt)〉 | Ft] (14)
Taking total expectation of Eq. (13) and plug in Eq. (14), we get
E[‖wt+1 − w∗‖22 | Ft] = E[‖wt − w∗‖22 | Ft] + η2 · E[‖deskt(skt(∇f(wt)))‖22 | Ft]− 2η · E[〈wt − w∗,∇f(wt)〉 | Ft] ≤ E[‖wt − w∗‖22 | Ft] + η2 · (1 + α) · E[‖∇f(wt)‖22 | Ft] + 2η · E[〈w∗ − wt,∇f(wt)〉 | Ft] ≤ E[‖wt − w∗‖22 | Ft] + η2 · (1 + α) · E[‖∇f(wt)‖22 | Ft] + 2η · E[f(w∗)− f(wt) | Ft] (15)
where the second step follows from the variance of deskt ◦ skt, and the last step follows from the convexity of f . Taking the expectation over Ft and re-organizing the above equation, we can get
2η · E[f(wt)− f(w∗)] ≤ E[‖wt − w∗‖22]− E[‖wt+1 − w∗‖22] + η2 · (1 + α) · E[‖∇f(wt)‖22] ≤ E[‖wt − w∗‖22]− E[‖wt+1 − w∗‖22] + η2 · (1 + α) · 2L · E[f(wt)− f(w∗)]
where the second step follows from the convexity and L-smoothness of f . Rearrange the above inequality, we have
(2η − η2 · (1 + α) · 2L) · E[f(wt)− f(w∗)] ≤ E[‖wt − w∗‖22]− E[‖wt+1 − w∗‖22]
Note η ≤ 12(1+α)L , we have
η · E[f(wt)− f(w∗)] ≤ E[‖wt − w∗‖22]− E[‖wt+1 − w∗‖22]
Sum over all T iterations, we arrive at
η · T∑ t=0 E[f(wt)− f(w∗)] ≤E[‖w0 − w∗‖22]− E[‖wT+1 − w∗‖22] ≤ E[‖w0 − w∗‖22] (16)
Let wT = 1T+1 ∑T t=0 w
t denote the average of parameters across iterations, then by convexity of f , we conclude:
E[f(wT )− f(w∗)] ≤ E[‖w 0 − w∗‖22]
η · (T + 1)
E.4 NON-CONVEX f CONVERGENCE ANALYSIS
Next, we prove a version when f is not even a convex function, due to loss of convexity, we can no longer bound the gap between E[f(wt)] and f(w∗), but we can instead bound the minimum (or average) expected gradient.
Theorem E.3. Let f : Rd → R be an L-smooth function (Def. C.1) and sk/desk functions satisfying Theorem 4.2, let w∗ ∈ Rd be the optimal solution to f . Suppose η := ηlocal · ηglobal ≤ 1(1+α)L , then
min t∈[T ]
E[‖∇f(wt)‖22] ≤ 2
η(T + 1) (E[f(w0)]− f(w∗))
Proof. Note that the only place we used strongly-convex assumption in the proof of Theorem E.1 is Eq. (9), so by the same analysis, we can get
E[f(wt+1)− f(wt) | Ft] ≤− η
2 · ‖∇f(wt)‖22
Rearranging and taking total expectation over Ft, we get
E[‖∇f(wt)‖22] ≤ 2
η (E[f(wt)]− E[f(wt+1)])
Averaging over all T iterations, we get
1
T + 1 T∑ t=0 E[‖∇f(wt)‖22] ≤ 2 η(T + 1) T∑ t=0 (E[f(wt)]− E[f(wt+1)])
= 2
η(T + 1) (E[f(w0)]− E[f(wT )])
≤ 2 η(T + 1) (E[f(w0)]− f(w∗))
This implies our final result:
min t∈[T ]
E[‖∇f(wt)‖22] ≤ 2
η(T + 1) (E[f(w0)]− f(w∗))
Remark E.4. Notice due to the structure of sk/desk functions, i.e., their variance is bounded in terms of true gradient, the convergence rate depends completely on the term 1(1+α)L . If it’s a constant, then we essentially recover a convergence rate of gradient descent. On the other hand, if
1 (1+α)L ≤ 1√ T , then we get a similar convergence rate as SGD. One clear advantage of our sk/desk functions is they don’t introduce extra noise term as in SGD, since we can choose appropriate step size to absorb the variance term.
F k-STEP CONVEX & STRONGLY-CONVEX fc ANALYSIS
F.1 PRELIMINARY
In this section, we assume each fc satisfies Assumption 3.1 and ηglobal = 1. For notation simplicity, we also denote ut,−1c = u t−1,K−1 c for t ≥ 2. Definition F.1. Let (t, k) ∈ {1, · · · , T + 1}×{−1, 0, 1, · · · ,K− 1}, we define the following terms for iteration (t, k):
ut,k := 1
N N∑ c=1 ut,kc , r t,k := ut,k − w∗
to be the average of local parameters and its distance to the optimal solution,
gt,kc := ∇fc(ut,kc ), gt,k := 1
N N∑ c=1 ∇fc(ut,kc )
to be the local gradient and its average,
V t,k := 1
N N∑ c=1 ‖ut,kc − ut,k‖22
to be the variances of local updates,
σ2 = 1
N N∑ c=1 ‖∇fc(w∗)‖2
to be a finite constant that characterize the heterogeneity of local objectives.
We also define the following indicator function: let l ∈ R, then we define 1{x=l} to be
1{x=l} = { 1 if x = l, 0 otherwise.
F.2 UNIFYING THE UPDATE RULE OF ALGORITHM 1
Lemma F.2. We have the following facts for ut,kc and ũt,k:
ut,0c = u t,0 ut,kc = u t,k−1 c − ηlocal · gt,k−1c , ∀k ≥ 1
ut,k = ut,k−1 − ηlocal · gt,k−1 + 1{k=0} · ηlocal · (Id − deskt ◦ skt)( K−1∑ i=0 gt−1,i), ∀(t, k) 6= (1, 0)
where Id : Rd → Rd is the identity function.
Proof. First two equation directly follows from the update rule of Algorithm 1. For k = 1, 2, · · · ,K − 1, taking the average of the second equation we obtain:
ut,k = ut,k−1 − ηlocal · gt,k−1
For k = 0 and t ≥ 2, we have
ut,0 = ut−1,0 − ηlocal · deskt(skt( K−1∑ i=0 gt−1,i))
= ut−1,0 − ηlocal K−1∑ i=0 gt−1,i + ηlocal K−1∑ i=0 gt−1,i − ηlocal · deskt(skt( K−1∑ i=0 gt−1,i))
= ut−1,K−1 − ηlocal · gt−1,K−1 + ηlocal · (Id − deskt ◦ skt)( K−1∑ i=0 gt−1,i)
Combining above results together, we prove the third equation.
F.3 UPPER BOUNDING ‖gt,k‖22
Lemma F.3. Suppose for any c ∈ [N ], fc : Rd → R is convex and L-smooth. Then
‖gt,k‖22 ≤ 2L2V t,k + 4L(f(ut,k)− f(w∗))
Proof. By triangle inequality and Cauchy-Schwartz inequality, we have
‖gt,k‖22 = ‖gt,k −∇f(ut,k) +∇f(ut,k)‖22 ≤ 2‖gt,k −∇f(ut,k)‖22 + 2‖∇f(ut,k)‖22
where the first term can be bounded as
‖gt,k −∇f(ut,k)‖22 = ‖ 1
N N∑ c=1 ∇fc(ut,kc )− 1 N N∑ c=1 ∇fc(ut,k)‖22
≤ 1 N N∑ c=1 ‖∇fc(ut,kc )− fc(ut,k)‖22
≤ L 2
N N∑ c=1 ‖ut,kc − ut,k‖22
and the second term can be bounded as follows:
‖∇f(ut,k)‖22 = ‖∇f(ut,k)−∇f(w∗)‖22 ≤ 2L(f(ut,k)− f(w∗))
where the last step follows from that f is L-smooth and Fact C.4.
Combining bounds on these two terms, we get
‖gt,k‖22 ≤ 2L2
N N∑ c=1 ‖ut,kc − ut,k‖22 + 2L2‖ut,k − w∗‖22
≤ 2L2V t,k + 4L(f(ut,k)− f(w∗))
F.4 LOWER BOUNDING 〈ut,k − w∗, gt,k〉
Lemma F.4. Suppose each fc satisfies Assumption 3.1 with µ ≥ 0, then
〈ut,k − w∗, gt,k〉 ≥ f(ut,k)− f(w∗)− L 2 V t,k + µ 2 ‖ut,k − w∗‖22
Proof. We will provide a lower bound on this inner product:
〈ut,k − w∗, gt,k〉 = 1 N N∑ c=1 〈ut,k − w∗,∇fc(ut,kc )〉
It suffices to consider each term separately:
〈ut,k − w∗,∇fc(ut,kc )〉 = 〈ut,k − ut,kc + ut,kc − w∗,∇fc(ut,kc )〉 = 〈ut,k − ut,kc ,∇fc(ut,kc )〉+ 〈ut,kc − w∗,∇fc(ut,kc )〉
The first term can be lower bounded via L-smoothness:
〈ut,k − ut,kc ,∇fc(ut,kc )〉 ≥ fc(ut,k)− fc(ut,kc )− L
2 ‖ut,k − ut,kc ‖22
The second term can be lower bounded via convexity:
〈ut,kc − w∗,∇fc(ut,kc )〉 ≥ fc(ut,kc )− fc(w∗) + µ
2 ‖ut,kc − w∗‖22
Combining these two bounds and average them, we get a lower bound:
〈ut,k − w∗, gt,k〉 ≥ 1 N N∑ c=1 (fc(u t,k)− fc(w∗)− L 2 ‖ut,k − ut,kc ‖22 + µ 2 ‖ut,kc − w∗‖22)
≥ 1 N N∑ c=1 (fc(u t,k)− fc(w∗))− L 2 V t,k + µ 2 ‖ut,k − w∗‖22 = f(ut,k)− f(w∗)− L 2 V t,k + µ 2 ‖ut,k − w∗‖22
F.5 UPPER BOUNDING VARIANCE WITHIN K LOCAL STEPS
Lemma F.5. Suppose each fc is convex and L-smooth. Assume ηlocal ≤ 18LK . Then for any t ≥ 0, K−1∑ k=0 V t,k ≤ 8η2localLK2 K−1∑ k=0 (f(ut,k)− f(w∗)) + 4η2localK3σ2
Proof. By Lemma F.2, we know V t,0 = 0 for any t ≥ 0. Consider k ∈ {1, 2, · · · ,K − 1}, we have
V t,k = 1
N N∑ c=1 ‖ut,kc − ut,k‖22
= 1
N N∑ c=1 ‖ut,0c − k−1∑ i=0 ηlocal · gt,ic − ut,0 + k−1∑ i=0 ηlocal · gt,i‖22
= η2local N N∑ c=1 ‖ k−1∑ i=0 (gt,i − gt,ic )‖22
≤ η 2 localk
N N∑ c=1 k−1∑ i=0 ‖gt,i − gt,ic ‖22
≤ η 2 localK
N N∑ c=1 k−1∑ i=0 ‖gt,ic ‖22 (17)
where the second step follows from Lemma F.2, the last step follows from gt,i being the average of gt,ic . By Cauchy-Schwartz inequality, we further have:
‖gt,ic ‖22 ≤ 3‖gt,ic −∇fc(ut,i)‖22 + 3‖∇fc(ut,i)−∇fc(w∗)‖22 + 3‖∇fc(w∗)‖22 ≤ 3L2‖ut,ic − ut,i‖22 + 6L(fc(ut,i)− fc(w∗) + 〈w∗ − ut,0,∇fc(w∗)〉) + 3‖∇fc(w∗)‖22.
where the last step follows from applying L-smoothness to the first and second term.
Averaging with respect to c,
1
N N∑ c=1 ‖gt,ic ‖22 ≤ 3L2V t,i + 6L(f(ut,i)− f(w∗)) + 3σ2.
Note that the inner product term vanishes since 1N ∑N c=1∇fc(w∗) = ∇f(w∗) = 0.
Plugging back to Eq. (17), we obtain
V t,k ≤ η 2 localK
N N∑ c=1 k−1∑ i=0 ‖gt,ic ‖22
≤ η2localK k−1∑ i=0 (3L2V t,i + 6L(f(ut,i)− f(w∗)) + 3σ2).
Summing up above inequality as k varies from 0 to K − 1, K−1∑ k=0 V t,k ≤ η2localK K−1∑ k=0 k−1∑ i=0 (3L2V t,i + 6L(f(ut,i)− f(w∗)) + 3σ2)
≤ η2localK K−1∑ k=0 K−1∑ i=0 (3L2V t,i + 6L(f(ut,i)− f(w∗)) + 3σ2)
= 3η2localL 2K2 K−1∑ i=0 V t,i + 6η2localLK 2 K−1∑ i=0 (f(ut,i)− f(w∗)) + 3η2localK3σ2
Rearranging terms we obtain:
(1− 3η2localL2K2) K−1∑ k=0 V t,k ≤ 6η2localLK2 K−1∑ i=0 (f(ut,i)− f(w∗)) + 3η2localK3σ2
Since ηlocal ≤ 18LK , we have 1− 3η 2 localL 2K2 ≥ 34 , implying K−1∑ k=0 V t,k ≤ 8η2localLK2 K−1∑ i=0 (f(ut,i)− f(w∗)) + 4η2localK3σ2
F.6 BOUNDING THE EXPECTED GAP BETWEEN ut,k AND w∗
Lemma F.6. Suppose each fc satisfies Assumption 3.1 with µ ≥ 0. If sk/desk satisfying Theorem 4.2 and ηlocal ≤ 14L , then for any (t, k) 6= (1, 0), we have
E[‖ut,k − w∗‖22] ≤ (1− µηlocal)E[‖ut,k−1 − w∗‖22] + 3
2 ηlocalLE[V t,k−1]− ηlocal E[f(ut,k−1)− f(w∗)]
+ 1{k=0}η 2 localαK
( 2L2 K−1∑ i=0 E[V t−1,i] + 4L K−1∑ i=0 E[f(ut−1,i)− f(w∗)] )
Proof. By Lemma F.2, we have for any (t, k) 6= (1, 0),
ut,k = ut,k−1 − ηlocal · gt,k−1 + 1{k=0} · ηlocal · (Id − deskt ◦ skt)( K−1∑ i=0 gt−1,i)
Therefore, denoting ht := (Id − deskt ◦ skt)( ∑K−1 i=0 g t−1,i), we have
‖ut,k − w∗‖22 = ‖ut,k−1 − w∗ − ηlocal · gt,k−1 + 1{k=0}ηlocal · ht‖22 = ‖ut,k−1 − w∗‖22 + η2local · ‖gt,k−1‖22 − 2ηlocal〈ut,k−1 − w∗, gt,k−1〉
+ 2ηlocal1{k=0}〈ut,k−1 − w∗, ht〉 − 2η2local1{k=0}〈gt,k−1, ht〉 + η2local1{k=0} · ‖ht‖22 (18)
Note by Theorem 4.2, we have:
E[deskt(skt(h))] = h, E[‖deskt(skt(h))‖22] ≤ (1 + α) · ‖h‖22 hold for any vector h. Hence, by taking expectation over Eq. (18),
E[‖ut,k − w∗‖22|Ft] = E[‖ut,k−1 − w∗‖22|Ft] + η2local · E[‖gt,k−1‖22|Ft] − 2ηlocal E[〈ut,k−1 − w∗, gt,k−1〉|Ft] + 1{k=0} · η2local · E[‖ht‖22|Ft]
Note that since E[ht | Ft] = 0, so the two inner products involving ht vanishes. Since
E[‖ht‖22|Ft] = E[‖(Id − deskt ◦ skt)( K−1∑ i=0 gt−1,i)‖22|Ft]
≤ αE[‖ K−1∑ i=0 gt−1,i‖22|Ft]
≤ αK K−1∑ i=0 E[‖gt−1,i‖22|Ft]
Taking total expectation, we have
E[‖ut,k − w∗‖22]
≤ E[‖ut,k−1 − w∗‖22] + η2local · E[‖gt,k−1‖22]− 2ηlocal E[〈ut,k−1 − w∗, gt,k−1〉]
+ 1{k=0} · η2local · αK K−1∑ i=0 E[‖gt−1,i‖22]
≤ E[‖ut,k−1 − w∗‖22] + η2local · E[2L2V t,k−1 + 4L(f(ut,k−1)− f(w∗))]
− 2ηlocal E[f(ut,k−1)− f(w∗)− L
2 V t,k−1 +
µ 2 ‖ut,k−1 − w∗‖22]
+ 1{k=0} · η2local · αK K−1∑ i=0 E[2L2V t−1,i + 4L(f(ut−1,i)− f(w∗))]
≤ (1− µηlocal)E[‖ut,k−1 − w∗‖22] + ηlocal · L · (1 + 2ηlocalL) · E[V t,k−1] − 2ηlocal · (1− 2ηlocalL) · E[f(ut,k−1)− f(w∗)]
+ 1{k=0} · η2local · αK · ( 2L2 K−1∑ i=0 E[V t−1,i] + 4L K−1∑ i=0 E[f(ut−1,i)− f(w∗)] )
where the second step follows from Lemma F.3 and Lemma F.4. Since ηlocal ≤ 14L , we have E[‖ut,k − w∗‖22] ≤ (1− µηlocal)E[‖ut,k−1 − w∗‖22] + 3
2 ηlocalLE[V t,k−1]− ηlocal E[f(ut,k−1)− f(w∗)]
+ 1{k=0}η 2 localαK
( 2L2 K−1∑ i=0 E[V t−1,i] + 4L K−1∑ i=0 E[f(ut−1,i)− f(w∗)] )
F.7 MAIN RESULT: CONVEX CASE
Theorem F.7. Assume each fc is convex and L-smooth. If Theorem 4.2 holds and ηlocal ≤ 1 8(1+α)LK ,
E[f(wT )− f(w∗)] ≤ 4E[‖w 0 − w∗‖22]
ηlocalKT + 32η2localLK 2σ2,
where wT = 1KT ( ∑T t=1 ∑K−1 k=0 u
t,k) is the average over parameters throughout the execution of Algorithm 1.
Proof. Summing up Lemma F.6 as t varies from 1 to T and k varies from 0 to K − 1, E[‖uT+1,0 − w∗‖22]− E[‖w0 − w∗‖22]
≤ 3 2 ηlocalL T∑ t=1 K−1∑ k=0 E[V t,k]− ηlocal T∑ t=1 K−1∑ k=0 E[f(ut,k)− f(w∗)]
+ T∑ t=1 K−1∑ k=0 1{k=0}η 2 localαK ( 2L2 K−1∑ i=0 E[V t,i] + 4L K−1∑ i=0 E[f(ut,i)− f(w∗)] )
= 3
2 ηlocalL T∑ t=1 K−1∑ k=0 E[V t,k]− ηlocal T∑ t=1 K−1∑ k=0 E[f(ut,k)− f(w∗)]
+ η2localαK ( 2L2 T∑ t=1 K−1∑ i=0 E[V t,i] + 4L T∑ t=1 K−1∑ i=0 E[f(ut,i)− f(w∗)] )
= ηlocalL( 3
2 + 2ηlocalαLK) T∑ t=1 K−1∑ k=0 E[V t,k]
− ηlocal(1− 4ηlocalαLK) T∑ t=1 K−1∑ k=0 E[f(ut,k)− f(w∗)]
≤ 2ηlocalL T∑ t=1 K−1∑ k=0 E[V t,k]− 1 2 ηlocal T∑ t=1 K−1∑ k=0 E[f(ut,k)− f(w∗)]
≤ 2ηlocalL T∑ t=1 (8η2localLK 2 K−1∑ i=0 E[f(ut,i)− f(w∗)] + 4η2localK3σ2)− 1 2 ηlocal T∑ t=1 K−1∑ k=0 E[f(ut,k)− f(w∗)]
≤ − 1 4 ηlocal T∑ t=1 K−1∑ k=0 E[f(ut,k)− f(w∗)] + 8η3localLK3Tσ2
where the fourth step follows from ηlocal ≤ 18αLK , the last step follows from ηlocal ≤ 1 8LK . Rearranging the terms, we obtain
1
KT T∑ t=1 K−1∑ k=0 E[f(ut,k)− f(w∗)] ≤ 4E[‖w 0 − w∗‖22] ηlocalKT + 32η2localLK 2σ2
Finally, by the convexity of f we complete the proof.
Now we are ready to answer the question: how much communication cost is sufficient to guarantee E[f(wT )− f(w∗)] ≤ ? we have the following communication cost result: Theorem F.8. Assume each fc is convex and L-smooth. If Theorem 4.2 holds. With O ( E[‖w0 − w∗‖22]N max{Ld , σ √ L 3/2 } ) bits of communication cost, Algorithm 1 outputs an - optimal solution wT satisfying:
E[f(wT )− f(w∗)] ≤ ,
where wT = 1KT ( ∑T t=1 ∑K−1 k=0 u t,k).
Proof. To calculate the communication complexity, we first note communication only happens in sync steps. Specifically, in each sync step, the algorithm requires O(Nbsketch) bits of communication cost, where bsketch denotes the sketching dimension. Therefore, the total cost of communication is given by O(NbsketchT ). To obtain the optimal communication cost for -optimal solution, we choose T,K, ηlocal and bsketch by solving the following optimization problem:
min T,K,ηlocal,bsketch,α NbsketchT
s.t. 0 < ηlocal ≤ 1
8(1 + α)LK
4E[‖w0 − w∗‖22] ηlocalKT ≤ 2
32η2localLK 2σ2 ≤
2
d ≥ bsketch = O( d
α ) ≥ 1
where d is the parameter dimension and the last constraint is due to Theorem 4.2. Above constraints imply:
T ≥ 8E[‖w 0 − w∗‖22]
ηlocalK , ,Kηlocal ≤ min{
1
8(1 + α)L ,
1
8σ
√
L }
Therefore, when ≥ σ 2
(1+α)2L , the optimal solution is given by
Kηlocal = 1
8(1 + α)L , T =
64E[‖w0 − w∗‖22](1 + α)L
, bsketch = O( d
α )
and the corresponding optimal communication cost is O(E[‖w 0−w∗‖22LNd ).
when < σ 2
(1+α)2L , the optimal solution is given by
Kηlocal = 1
8σ
√
L , T =
64E[‖w0 − w∗‖22σ √ L]
3/2 , bsketch = O(
d α )
and the corresponding optimal communication cost is O(E[‖w 0−w∗‖22σ
√ LNd
α 3/2 ).
Combining above two cases, the optimal α is given by O(d), and the corresponding optimal communication cost will be O(E[‖w0 − w∗‖22]N max{Ld , σ √ L 3/2 }).
F.8 MAIN RESULT: STRONGLY-CONVEX CASE
Theorem F.9. Assume each fc is µ-strongly convex and L-smooth. If Theorem 4.2 holds and ηlocal ≤ 18(1+α)LK ,
E[f(wT+1)− f(w∗)] ≤ L 2 E[‖w0 − w∗‖22]e−µηlocalT + 4η2localL2K3σ2/µ.
Proof. Summing up Lemma F.6 as k varies from 0 to K − 1, then we have for any t ≥ 1,
(E[‖ut+1,0 − w∗‖22] + K−1∑ k=1 E[‖ut,k − w∗‖22])− (1− µηlocal) K−1∑ k=0 E[‖ut,k − w∗‖22])
≤ 3 2 ηlocalL K−1∑ k=0 E[V t,k]− ηlocal K−1∑ k=0 E[f(ut,k)− f(w∗)]
+ K−1∑ k=0 1{k=0}η 2 localαK ( 2L2 K−1∑ i=0 E[V t,i] + 4L K−1∑ i=0 E[f(ut,i)− f(w∗)] )
= 3
2 ηlocalL K−1∑ k=0 E[V t,k]− ηlocal K−1∑ k=0 E[f(ut,k)− f(w∗)]
+ η2localαK ( 2L2 K−1∑ i=0 E[V t,i] + 4L K−1∑ i=0 E[f(ut,i)− f(w∗)] )
= ηlocalL( 3
2 + 2ηlocalαLK) K−1∑ k=0 E[V t,k]− ηlocal(1− 4ηlocalαLK) K−1∑ k=0 E[f(ut,k)− f(w∗)]
≤ 2ηlocalL K−1∑ k=0 E[V t,k]− 1 2 ηlocal K−1∑ k=0 E[f(ut,k)− f(w∗)]
≤ 2ηlocalL(8η2localLK2 K−1∑ i=0 E[f(ut,i)− f(w∗)] + 4η2localK3σ2)− 1 2 ηlocal K−1∑ k=0 E[f(ut,k)− f(w∗)]
≤ − 1 4 ηlocal K−1∑ k=0 E[f(ut,k)− f(w∗)] + 8η3localLK3σ2
where the fourth step follows from ηlocal ≤ 18αLK , the last step follows from ηlocal ≤ 1 8LK . Rearranging the terms, we obtain
E[‖ut+1,0 − w∗‖22] ≤ (1− µηlocal)E[‖ut,0 − w∗‖22] + 8η3localLK3σ2
implying
E[‖ut+1,0 − w∗‖22]− 8η2localLK3σ2/µ ≤ (1− µηlocal)(E[‖ut,0 − w∗‖22]− 8η2localLK3σ2/µ).
Therefore, we have
E[‖wT+1 − w∗‖22]− 8η2localLK3σ2/µ ≤ (1− µηlocal)T (E[‖w0 − w∗‖22]− 8η2localLK3σ2/µ)
≤ E[‖w0 − w∗‖22]e−µηlocalT
Finally, by L-smoothness of function f , we obtain
E[f(wT+1)− f(w∗)] ≤ L 2 E[‖wT+1 − w∗‖22] ≤ L 2 E[‖w0 − w∗‖22]e−µηlocalT + 4η2localL2K3σ2/µ.
Theorem F.10. Assume each fc is µ-strongly convex and L-smooth. If Theorem 4.2 holds. With O ( LN µ max{d, √ σ2 µ } log( LE[‖w0−w∗‖22] ) ) bits of communication cost, Algorithm 1 outputs an - optimal solution wT satisfying:
E[f(wT )− f(w∗)] ≤ .
Proof. To calculate the communication complexity, we first note communication only happens in sync steps. Specifically, in each sync step, the algorithm requires O(Nbsketch) bits of communication cost, where bsketch denotes the sketching dimension. Therefore, the total cost of communication is given by O(NbsketchT ). To obtain the optimal communication cost for -optimal solution, we choose T,K, ηlocal and bsketch by solving the following optimization problem:
min T,K,ηlocal,bsketch,α NbsketchT
s.t. 0 < ηlocal ≤ 1
8(1 + α)LK
L 2 E[‖w0 − w∗‖22]e−µηlocalT ≤ 2 4η2localL 2K3σ2/µ ≤
2
d ≥ bsketch = O( d
α ) ≥ 1
where d is the parameter dimension and the last constraint is due to Theorem 4.2. Above constraints imply:
T ≥ 1 µηlocal
log( LE[‖w0 − w∗‖22] ), ηlocal ≤ min{
1
8(1 + α)LK ,
1
2LKσ
√ µ
2K }
Therefore, the optimal value is given when K = 1. When ≥ σ 2
16(1+α)2µ , the optimal solution is given by
ηlocal = 1
8(1 + α)L , T =
8(1 + α)L
µ log(
LE[‖w0 − w∗‖22]
), bsketch = O( d
α )
and the corresponding optimal communication cost is O(LNdµ log( LE[‖w0−w∗‖22] )).
when < σ 2
16(1+α)2µ , the optimal solution is given by
ηlocal = 1
2Lσ
√ µ
2 , T =
2Lσ µ3/2
√ 2
log(
LE[‖w0 − w∗‖22]
), bsketch = O( d
α )
and the corresponding optimal communication cost is O( σLNd αµ3/2 √
log( LE[‖w0−w∗‖22] )).
Combining above two cases, the optimal α is given by O(d), and the corresponding optimal communication cost will be O(LNµ max{d, √ σ2 µ } log( LE[‖w0−w∗‖22] )).
G k-STEP NON-CONVEX f CONVERGENCE ANALYSIS
In this section, we present convergence result for non-convex f case in the k-local-step regime. In order for the proof to go through, we assume that for any c ∈ [N ] and any w ∈ Rd, there exists a universal constant G such that
‖∇fc(w)‖2 ≤ G.
Throughout the proof, we will use Ft to denote the sequence wt−1, wt−2, . . . , w0. Also, we use η as a shorthand for ηglobal · ηlocal. Note that in k-local-step scheme, the average of local gradients is no longer the true gradient, therefore, we can no longer bound everything using the true gradients. This means it’s necessary to introduce the gradient norm upper bound assumption. Lemma G.1. Let f : Rd → R satisfying Assumption 3.1 and sk/desk functions satisfying Theorem 4.2. Further, assume ηlocal ≤ 12LK . Then
E[f(wt+1)− f(wt) | Ft] ≤ − ηglobal · ‖∇f(wt)‖22 + η · L ·K2 ·G2 · ( ηlocal + η
2 · (1 + α) ) Proof. We start by bounding f(wt+1)− f(wt) without taking conditional expectation:
f(wt+1)− f(wt)
≤ 〈wt+1 − wt,∇f(wt)〉+ L 2 ‖wt+1 − wt‖22 = 〈deskt(∆w̃t),∇f(wt)〉+ L
2 ‖deskt(∆w̃t)‖22
= A+ L
2 B
where
A := − 〈ηglobal · deskt( 1
N N∑ c=1 skt( K−1∑ k=0 ηlocal · ∇fc(ut,kc ))),∇f(wt)〉
B := ‖ηglobal · deskt( 1
N N∑ c=1 skt( K∑ k=1 ηlocal · ∇fc(ut,kc )))‖22
Bounding E[A | Ft] Using the fact that skt/deskt are linear functions and E[deskt(skt(h))] = h, we get
E[A | Ft] = − 〈ηglobal · 1
N N∑ c=1 K−1∑ k=0 ηlocal · ∇fc(ut,kc ),∇f(wt)〉
= − ηglobal · 〈 1
N N∑ c=1 (K−1∑ k=0 ηlocal · ∇fc(ut,kc )−∇fc(wt) +∇fc(wt) ) ,∇f(wt)〉
= − ηglobal · ‖∇f(wt)‖22 + ηglobal · ηlocal · 1
N N∑ c=1 K−1∑ k=0 〈∇fc(ut,kc )−∇fc(wt),∇f(wt)〉
It suffices to bound the inner product, notice for k = 0, the inner product is 0, so assume k ≥ 1:
〈∇fc(ut,kc )−∇fc(wt),∇f(wt)〉 ≤ ‖∇fc(ut,kc )−∇fc(wt)‖2 · ‖∇f(wt)‖2 ≤ L · ‖ut,kc − wt‖2 · ‖∇f(wt)‖2 (19)
where the gap between ut,kc and w t can be further expanded:
‖ut,kc − wt‖2 = ‖ut,kc − u t,k 0 ‖2
= ‖ηlocal k−1∑ i=0 ∇fc(ut,ic )‖2
≤ ηlocal k−1∑ i=0 ‖∇fc(ut,ic )‖2
≤ ηlocal · k ·G (20)
Plug in Eq. (20) to Eq. (19), we get
〈∇fc(ut,kc )−∇fc(wt),∇f(wt)〉 ≤ L · ηlocal · k ·G2
Recall that η = ηglobal · ηlocal. Put things together, we finally obtain a bound on E[A | Ft]:
E[A | Ft] ≤ − ηglobal · ‖∇f(wt)‖22 + η · ηlocal · L · ( K−1∑ k=0 k) ·G2
≤ − ηglobal · ‖∇f(wt)‖22 + η · ηlocal · L ·K2 ·G2 (21)
Bounding E[B | Ft] Using the fact that skt/deskt are linear functions, we get
B = η2global · η2local · 1
N2 · ‖ N∑ c=1 K−1∑ k=0 deskt(skt(∇fc(ut,kc )))‖22
≤ η2global · η2local · 1
N2 ·N ·K N∑ c=1 K−1∑ k=0 ·‖deskt(skt(∇fc(ut,kc )))‖22
= η2 · K N · N∑ c=1 K−1∑ k=0 ‖deskt(skt(∇fc(ut,kc )))‖22
Using variance bound of deskt(skt(h)), we get
E[B | Ft] ≤ η2 · K
N · (1 + α) · N∑ c=1 K−1∑ k=0 ‖∇fc(ut,kc )‖22
≤ η2 · K N · (1 + α) · N∑ c=1 K−1∑ k=0 G2 = η2 ·K2 · (1 + α) ·G2 (22)
Put things together Put the bound on E[A | Ft] and the bound on E[B | Ft], we get
E[f(wt+1)− f(wt) | Ft]
≤ − ηglobal · ‖∇f(wt)‖22 + η · ηlocal · L ·K2 ·G2 + L
2 · η2 ·K2 · (1 + α) ·G2 = − ηglobal · ‖∇f(wt)‖22 + η · L ·K2 ·G2 · ( ηlocal + η
2 · (1 + α) ) Theorem G.2. Let f : Rd → R be L-smooth. Let w∗ ∈ Rd be the optimal solution to f and assume sk/desk functions satisfying Theorem 4.2. Then
min t∈[T ]
E[‖∇f(wt)‖22] ≤ 1
(T + 1)ηglobal · (E[f(w0)]− f(w∗)) + ηlocal · LK2G2 ·
( ηlocal + η
2 · (1 + α) ) Proof. By Lemma G.1, we know that
E[f(wt+1) | Ft]− f(wt) ≤ − ηglobal · ‖∇f(wt)‖22 + η · L ·K2 ·G2 · ( ηlocal + η
2 · (1 + α) ) Rearranging the inequality and taking total expectation, we get
E[‖∇f(wt)‖22] ≤ 1
ηglobal · (E[f(wt)]− E[f(wt+1)]) + ηlocal · LK2G2 ·
( ηlocal + η
2 · (1 + α) ) Sum over all T iterations and averaging, we arrive at
1
T + 1 T∑ t=0 E[‖∇f(wt)‖22]
≤ 1 (T + 1)ηglobal
· (E[f(w0)]− E[f(wT )]) + ηlocal · LK2G2 · ( ηlocal + η
2 · (1 + α) ) ≤ 1
(T + 1)ηglobal · (E[f(w0)]− f(w∗)) + ηlocal · LK2G2 ·
( ηlocal + η
2 · (1 + α)
)

heading: H DIFFERENTIAL PRIVACY
text: We define ( , δ)-differential privacy Dwork et al. (2006b;a) as
Definition H.1. Let , δ be positive real number and A be a randomized algorithm that takes a dataset as input (representing the actions of the trusted party holding the data). Let im(A) denote the image of A. The algorithm A is said to provide , δ-differential privacy if, for all datasets D1 and D2 that differ on a single element (i.e., the data of one person), and all subsets S of im(A):
Pr[A(D1) ∈ S] ≤ exp( ) · Pr[A(D2) ∈ S] + δ
where the probability is taken over the randomness used by the algorithm.
We state a lemma from prior work Kenthapadi et al. (2013), Lemma H.2 (Kenthapadi et al. (2013)). Let δ ∈ (0, 1/4) and > 0. Let Y and Y ′ be points in Rd such that ‖Y − Y ′‖2 ≤ w. Then for any D ⊂ Rd, and any ∆ drawn from {N (0, σ2)}d, where σ ≥ 4w −1 √ log(1/δ), the following inequality holds:
Pr[Y ′ + ∆ ∈ D] ≤ e · Pr[Y + ∆ ∈ D] + δ.
Throughout this section, we assume that for any gradient gc and any one-entry perturbation g′c of gc, the magnitude of the one-entry perturbation is upper bounded by γ.
Let R ∈ Rbsketch×d denote our sketching matrix. Let Y ′, Y ∈ Rbsketch . Let g, g′ ∈ Rd such that g and g′ differ by exactly one entry and the magnitude of this difference is bounded by γ. Let Y = Rg and Y ′ = Rg′.
We demonstrate how to pick the corresponding variance of noise for AMS sketch. We remark that for other sketching matrices where entries are bounded and hence its spectral norm is bounded, one can run similar argument to obtain such results.
Fact H.3. If R is an AMS sketching matrices, then √ d/b ≤ ‖R‖ ≤ √ d/ √ b.
Proof. Recall that each entry of an AMS sketching matrix is ± 1√ b , hence its max row or column `1 norm is d√ b . Recall that max row or column `1 norm is an upper bound on the spectral norm, hence we know that ‖R‖ ≤ d√ b . On the other hand, we know its spectral norm is lower bounded by its max row or column `2 norm, hence,
‖R‖ ≥ ‖R‖F /b ≥ √ d/b.
Put things together, we have √ d/b ≤ ‖R‖ ≤ √ d/ √ b.
Putting it all together, we have
Lemma H.4. Let ∆ ∈ Rbsketch denote the noise vector that sampled from Gaussian distribution N (0, σ2Ib) with σ ≥ 4
√ d
bsketch γ −1
√ log(1/δ), then we have Rg + ∆ is ( , δ)-differential private.
Proof. By Lemma H.2, it suffices to to bound ‖Y ′ − Y ‖2. Note that
‖Y − Y ′‖2 = ‖Rg −Rg′‖2 ≤ ‖R‖ · ‖g − g′‖2
≤
√ d
bsketch ‖g − g′‖2
≤
√ d
bsketch γ.
This means if we pick σ ≥ 4 √
d bsketch
γ −1 √
log(1/δ), then our sketched gradient with noise is ( , δ)differential private.

Review summaries:
This paper proposes a sketching-based federated learning algorithm to address the communication and privacy concerns in the federated learning setting. In the proposed algorithm, each client projects its local update (gradient) to a lower dimension via a sketching matrix before communicating the update to the server. The server aggregates the sketched local updates from different clients and communicates the aggregated update to the clients. The clients then de-sketch the message received from the server to update their local model.   The paper shows that if the sketching matrix satisfies the "coordinate-wise embedding" property, then the proposed algorithm leads to provable convergence for smooth and convex objective functions.   The paper also discusses the issue of privacy and proposes to add Gaussian noise to the sketched gradients to claim differential privacy of the resulting learning algorithm.

This paper proposes a novel federated learning framework where the client only submits a sketched gradient to the server and de-sketches the average gradient received from the server. The proposed framework can preserve the privacy of the data as well as reduce the communication cost. The paper theoretically analyzes the convergence rate and shows that the proposed method requires less communication cost than existing methods. 

The paper studies the application of linear sketching algorithms to reduce the communication complexity in federated Learning problems. In particular, the paper proposes that the clients send a low-dimensional linear sketch of their updates to the server. Convergence rate of local SGD when deployed with such a method for aggregating the updates are discussed for various function classes. Finally, the paper also discusses a method for integrating differentially privacy guarantees in their method.