FedSKETCH: Communication-Efficient and Private Federated Learning via Sketching

Introduction

Federated Learning is a recently emerging setting for distributed large scale machine learning problems. In Federated Learning, data is distributed across devices (which could be any smartphone or IOT edge device) and due to privacy concerns, users are only allowed to communicate with parameter server. The parameter server orchestrates optimization among devices by aggregating gradient-related information of devices and broadcasts the average of received vectors. Additionally, moving data across the devices for the purpose of learning a global model can be impractical and could violate the privacy of users/devices .

There are a number of challenges to be addressed in Federated Learning to efficiently learn a global model that performs well in average for all devices. The first challenge is the communication-efficiency as there could be a million of devices communicating iteratively among them which can incur huge communication overhead. The second challenge is data heterogeneity. Since the data in smartphones or devices are generated locally in Federated Learning, generated data may come from various probability distributions. Thus it is supposed that data distribution is non-iid. It is known that non-iid data distribution can lead to poor convergence error in practice . The last, yet important, issue is device privacy . It is important to make sure that the privacy of the sensitive information on each device is preserved during the training.

Almost all of the previous studies consider addressing the aforementioned challenges separately. One approach to deal with communication cost is the idea of local SGD with periodic averaging which asserts that instead of taking the average within each iteration, like baseline SGD , one may take the average periodically and performs local update, see local SGD . It is shown that local SGD with periodic averaging benefits from the same convergence rate as baseline SGD, while requiring less communication rounds. The second approach to deal with communication cost is aiming at reducing the size of communicated message per each communication round. Available methods reduce the size of the message by communicating compressed local gradients or models to parameter server via quantization , sparsification .

There are a number of research efforts such as aiming at mitigating the effect of data heterogeneity by exploiting variance reduction or gradient tracking techniques in distributed optimization settings where data distribution is non-iid.

Solving the privacy issue has been widely performed by injecting an additional layer of random noise in order to respect differential-privacy property of the method or using cryptography based approaches under secure multi-party computation framework.

Another promising recent approach with a potential to tackle all major issues in Federated Learning setting is based on sketching algorithms . Sketches are built from independent hash tables (functions), needed to compress a high dimensional vector into a lower dimensional one and the corresponding estimation error of sketching are well studied. With the focus of communication-efficiency, proposes a distributed SGD algorithm using sketching and they provide the convergence analysis in homogeneous data distribution setting. Also with focus on privacy, in , the authors derive a single framework in order to tackle these issues jointly and introduce DiffSketch based on the Count Sketch operator. Compression and privacy are performed using random hash functions such that no third parties are able to access the original data. Yet, does not provide the convergence analysis for the DiffSketch in Federated setting, and additionally the estimation error of the DiffSketch is relatively higher than the sketching scheme in which might end up in poor convergence error. Finally, considers using sketching technique for Federated Learning in heterogeneous setting from a communication-efficiency perspective. The proposed sketching schemes in are based on a deterministic scheme which requires having access to the exact values of the gradient-related information, thus are not privacy-preserving.

In this work, we provide a thorough convergence analysis for the Federated Learning using sketching for both homogeneous and heterogeneous settings. Additionally, all of our sketching algorithms including a novel scheme, do not need to obtain exact values of gradient, hence are privacy preserving. Therefore, our proposed algorithms based on sketching addresses all the aforementioned three main challenges jointly.

The main contributions of this paper are summarized as follows:

Based on the current compression methods, we provide a new algorithm – HEAPRIX – that displays an unbiased estimator of the full gradient we ought to communicate to the central parameter server. We theoretically show that HEAPRIX jointly reduces the cost of communication between devices and server, preserves privacy and is unbiased.

We develop a general algorithm for communication-efficient and privacy preserving federated learning based on this novel compression algorithm. Those methods, namely FedSKETCH and FedSKETCHGATE, are derived under homogeneous and heterogeneous data distribution settings.

Non asymptotic analysis of our method is established for convex, Polyak-Łojasiewicz (generalization of strongly-convex) and nonconvex functions in Theorem 2 and Theorem 3 for respectively the i.i.d. and non i.i.d. case, and highlight an improvement in the number of iteration required to achieve a stationary point.

We illustrate the benefits of FedSKETCH and FedSKETCHGATE over baseline methods through a set of experiments. In particular, we plot training loss and accuracy curves depending on the method used for training, the size of the sketches employed and the number of local updates performed at each round of communication. Numerical experiments show the advantages of, in particular, FedSKETCH-HEAPRIX algorithm that achieves comparable test accuracy as Federated SGD (FedSGD) while compressing the information exchanged between devices and server.

Related Work

In this section, we provide a summary of the prior related research efforts as follows:

Local SGD with Periodic Averaging: Compared to baseline SGD where model averaging happens in every iteration, the main idea behind Local SGD with periodic averaging comes from the intuition of variance reduction by periodic model averaging with purpose of saving communication rounds. While Local SGD has been proposed in under the title of Federated Learning Setting, the convergence analysis of Local SGD is studied in . The convergence analysis of Local SGD is improved in the follow up works in majority for homogeneous data distribution setting. The convergence analysis is further extended to heterogeneous setting, wherein studied under the title of Federated Learning, with improved rates in . Additionally, a few recent Federated Learning/Local SGD with adaptive gradient methods can be found in .

Gradient Compression Based Algorithms for Distributed Setting: develop a solution for leveraging sketches of full gradients in a distributed setting while training a global model using SGD . They introduce Sketched-SGD and establish a communication complexity of order $\mathcal{O}(\log(d))$ (per round) where $d$ is the dimension of the vector of parameters, i.e. the dimension of the gradient. Other recent solutions to reduce the communication cost include quantized gradient as developed in . Yet, their dependence on the number of devices $p$ makes them harder to be used in some practical settings. Additionally, there are other research efforts such as that exploit compression in Federated Learning or distributed communication-efficient optimization. Finally, the recent work in jointly exploits variance reduction technique with compression in distributed optimization.

Privacy-preserving Setting: Differentially private methods for federated learning have been extensively developed and studied in recently.

The remaining of the paper is organized as follows. Section 3 gives a formal presentation of the general problem. Section 4 describes the various compression algorithms used for communication efficiency and privacy preservation, and introduces our new compression method. The training algorithms are provided in Section 5 and their respective analysis in the strongly-convex or nonconvex cases are provided Section 6. Finally, in Section 7 we provide empirical results for our proposed algorithms.

Notation: For the rest of the paper we indicate the number of communication rounds and number of bits per round per device with $R$ and $B$ respectively. For the rest of the paper we indicate the count sketch of any vector $\boldsymbol{x}$ with $\mathbf{S}(\boldsymbol{x})$. We also denote $[p]=\{1,\dots,p\}$.

Problem Setting

The federated learning optimization problem across $p$ distributed devices is defined as follows:

We focus on solving the optimization problem in Eq. (1) for the homogeneous data distribution. In the heterogeneous setting we consider the special case of $q_{1}=\ldots=q_{p}=\frac{1}{p}$.

Count Sketch as a Compression Operation

A common sketching solution employed to tackle (1) called Count Sketch (for more detail see ) is described Algorithm 1.

The algorithm for generating count sketching is using two sets of functions that encode any input vector $\boldsymbol{x}$ into a hash table $\boldsymbol{S}_{m\times t}(\boldsymbol{x})$. We use hash functions $\{h_{j,1\leq j\leq t}:[d]\rightarrow m\}$ (which are pairwise independent) along with another set of pairwise independent sign hash functions $\{\text{sign}_{j,1\leq j\leq t}:[d]\rightarrow\{+1,-1\}\}$ to map every entry of $\boldsymbol{x}$ ($\boldsymbol{x}_{i},\>1\leq i\leq d$) into $t$ different columns of hash table $\mathbf{S}_{m\times t}$. These steps are summarized in Algorithm 1.

We note that this definition leads to the following property

Note that if $\Delta=1$ then our algorithm reduces to the case of no compression. This property allows us to control the noise of the compression.

2 An Example of Unbiased Compressor via Sketching

An instance of such unbiased compressor is PRIVIX which obtains an estimate of input $\boldsymbol{x}$ from a count sketch noted $\boldsymbol{S}(\boldsymbol{x})$. In this algorithm, to query the quantity $x_{i}$, the $i-th$ element of the vector, we compute the median of $t$ approximated values specified by the indices of $h_{j}(i)$ for $1\leq j\leq t$. These steps are summarized in Algorithm 2.

Next, we review a few properties of PRIVIX as follows:

Unbiased estimation: As it is also mentioned in , we have:

Bounded variance: With $m=\mathcal{O}\left(\frac{e}{\mu^{2}}\right)$ and $t=\mathcal{O}\left(\ln\left(\frac{d}{\delta}\right)\right)$, we have the following bound with probability $1-\delta$:

We note that $\Delta=1+\mu^{2}d$ implies that if $m\rightarrow d$, $\Delta\rightarrow 1+1=2$, which means that the case of no compression is not covered. Thus, the algorithms based on this may converges poorly.

In the following we provide a review of privacy property of count sketch:

A randomized mechanism $\mathcal{O}$ satisfies $\epsilon-$differential privacy, if for input data ${S}_{1}$ and ${S}_{2}$ differing by up to one element, and for any output $D$ of $\mathcal{O}$,

For smaller $\epsilon$, it will become more difficult to specify what is the input for the algorithm $\mathcal{O}$. Hence, smaller $\epsilon$ implies stronger privacy, and we desire to have $\epsilon$ as small as possible to impose stronger privacy guarantees. In the following, we review an assumption from to discuss a property regarding privacy.

For the purpose of privacy analysis, similar to , we suppose that for any input vector $S$ with length $|S|=l$, each element $s_{i}\in S$ is drawn i.i.d. from a Gaussian distribution: $s_{i}\sim\mathcal{N}(0,\sigma^{2})$, and bounded by a large probability: $|s_{i}|\leq C,1\leq i\leq p$ for some positive constant $C>0$.

Based on Assumption 1, the reference proves the following:

For a sketching algorithm $\mathcal{O}$ using Count Sketch $\mathbf{S}_{t\times m}$ with $t$ arrays of $m$ bins, for any input vector $S$ with length $l$ satisfying Assumption 1, $\mathcal{O}$ achieves $t.\ln\left(1+\frac{\alpha C^{2}m(m-1)}{\sigma^{2}(l-2)}(1+\ln(l-m))\right)-$differential privacy with high probability, where $\alpha$ is a positive constant satisfying $\frac{\alpha C^{2}m(m-1)}{\sigma^{2}(l-2)}(1+\ln(l-m))\leq\frac{1}{2}-\frac{1}{\alpha}$.

The proof of this theorem can be found in .

Theorem 1 implies that if we use smaller hash table either through using smaller $m$ or $t$, we will obtain stronger differential privacy. On the other hand, smaller hash table means bigger estimation error for a compression based on sketching. Therefore, there is an interesting trade-off between communication complexity and obtained privacy.

3 Biased Compressor

The following Lemma links these two definitions:

An instance of biased compressor based on sketching is given in Algorithm 3.

HEAVYMIX, with sketch size $\Theta\left(m\log\left(\frac{d}{\delta}\right)\right)$ is a biased compressor with $\alpha=1$ and $\Delta=d/m$ with probability $\geq 1-\delta$. In other words, with probability $1-\delta$, $\texttt{HEAVYMIX}\in C(\frac{d}{m},1)$.

We note that Algorithm 3 is a variation of the sketching algorithm developed in with distinction that HEAVYMIX does not require extra second round of communication to obtain the exact values of topm. Additionally, while sketching algorithm based on HEAVYMIX has smaller estimation error compared to PRIVIX, it requires having access to the exact values of topm, therefore such sketching does not benefit from differentially privacy similar to PRIVIX. In the following we introduce our sketching scheme which enjoys from privacy property as well as smaller estimation error.

4 Sketching Based on Induced Compressor

The following Lemma from shows that we can convert the biased compressor into an unbiased one:

We note that if $\Delta_{2}\geq 1$ and $\Delta_{1}\leq 1$, we have $\Delta=\Delta_{2}+\frac{1-\Delta_{2}}{\Delta_{1}}\leq\Delta_{2}$ .

Using this concept of the induced compressor we introduce HEAPRIX:

We highlight that in this case if $m\rightarrow d$, then $C(x)\rightarrow x$ which means that the algorithm convergence can be improved by decreasing the noise of compression (with choice of bigger $m$).

In the following we define two general framework for different sketching algorithms for homogeneous and heterogeneous data distributions.

Algorithms for Homogeneous and Heterogeneous Settings

In the following, we first present two algorithms for the homogeneous setting. Then, we present two other algorithms to deal with data heterogeneity. We emphasize that, for the sake of privacy in all of our algorithms, the query step is happening locally and the main task of the parameter server is to perform the average of the received messages from the devices and broadcast the average back to the devices.

In this section, we propose two algorithms for the setting where data across distributed devices are identically distributed. The proposed algorithms for Federated Learning leverage sketching techniques to compress communication. The main difference between the first suggested algorithm and the DiffSketch algorithm in is that we use distinct local and global learning rates. Additionally, unlike , we do not add local Gaussian noise to ensure privacy.

which is the aggregation of the consecutive stochastic gradients multiplied with local updates $\eta$.

Upon receiving all $\mathbf{S}_{j}^{(r)}$ from sampled devices, the server computes

and broadcasts it to all devices. Devices after receiving $\mathbf{S}^{(r)}$ from server update global model $\boldsymbol{x}^{(r)}$ using rule

We summarize these steps in FedSKETCH, see Algorithm 5. A variant of this algorithm which uses a different compression scheme, called HEAPRIX is also described in Algorithm 5. We note that for this variant we need to have an additional communication round between server and worker $j$ to aggregate $\delta_{j}^{(r)}\triangleq\mathbf{S}_{j}\left[\texttt{HEAVYMIX}(\mathbf{S}^{(r)})\right]$. Then, server averages all $\delta^{(r)}_{j}$ and broadcasts to all devices the following quantity:

where $\boldsymbol{S}^{(r)}$ is computed using Eq. (2) and then updates its global model using $\boldsymbol{x}^{(r+1)}=\boldsymbol{x}^{(r)}-\gamma{\mathbf{\Phi}}^{(r)}$.

2 Heterogeneous Setting

In this section, we focus on the optimization problem in Eq. (1) in special case of $q_{1}=\ldots=q_{p}=\frac{1}{p}$ with full device participation ($k=p$). We also note that these results can be extended to the scenario where devices are sampled, but for simplicity we do not analyze it in this section. In the previous section, we discussed algorithm FedSKETCH, which is originally designed for homogeneous setting where data distribution available at devices are identical. However, in a heterogeneous setting where data distribution could be different, the aforementioned algorithms may fail to perform well in practice. The main reason to cause this issue is that in Federated learning devices are using local stochastic descent direction which could be different than global descent direction when the data distribution are non-identical.

Convergence Analysis

In this section we start with a few common assumptions, then we provide the convergence results.

We note that Assumption 2 is a common assumption in the literature of stochastic optimization. Additionally, it is shown in that PL condition implies strong convexity property with same module. Additionally, PL objectives could also be nonconvex, hence strong convexity does not imply PL condition necessarily.

2 Convergence of FEDSKETCH for Homogeneous Setting

Now we focus on the homogeneous case where data is distributed i.i.d. among local devices. In this case, the stochastic local gradient of each worker is an unbiased estimator of the global gradient. We will need the following additional common assumption on the stochastic gradients.

Suppose that the conditions in Assumptions 2-4 hold. Given $0<m=O\left(\frac{e}{\mu^{2}}\right)\leq d$, and Consider FedSKETCH in Algorithm 5 with sketch size $B=O\left(m\log\left(\frac{dR}{\delta}\right)\right)$. If the local data distributions of all users are identical (homogeneous setting), then with probability $1-\delta$ we have

For the FedSKETCH-PRIVIX algorithm, by choosing stepsizes as $\eta=\frac{1}{L\gamma}\sqrt{\frac{k}{R\tau\left(\frac{\mu^{2}d}{k}+1\right)}}$ and $\gamma\geq k$, the sequence of iterates satisfies $\frac{1}{R}\sum_{r=0}^{R-1}\left\|\nabla f({\boldsymbol{w}}^{(r)})\right\|_{2}^{2}\leq{\epsilon}$ if we set $R=O\left(\frac{1}{\epsilon}\right)$ and $\tau=O\left(\frac{\mu^{2}d+1}{{k}\epsilon}\right)$.

For FedSKETCH-HEAPRIX algorithm, by choosing stepsizes as $\eta=\frac{1}{L\gamma}\sqrt{\frac{k}{R\tau\left(\frac{\mu^{2}d-1}{k}+1\right)}}$ and $\gamma\geq k$, the sequence of iterates satisfies $\frac{1}{R}\sum_{r=0}^{R-1}\left\|\nabla f({\boldsymbol{w}}^{(r)})\right\|_{2}^{2}\leq{\epsilon}$ if we set $R=O\left(\frac{1}{\epsilon}\right)$ and $\tau=O\left(\frac{\mu^{2}d}{{k}\epsilon}\right)$.

As a consequence of Remark 7, the total communication cost per-worker becomes

We note that this result in addition to improving over the communication complexity of federated learning of the state-of-the-art from $O\left(\frac{d}{\epsilon}\right)$ in to $O\left(\frac{mk}{\epsilon}\log\left(\frac{dk}{\epsilon\delta}\right)\right)$, it also implies differential privacy. As a result, total communication cost is

We note that the state-of-the-art in the total communication cost is

We improve this result, in terms of dependency on $d$, to

In comparison to , we improve the total communication per worker from $RB=O\left(\frac{m}{\epsilon^{2}}\log\left(\frac{d}{\epsilon^{2}\delta}\right)\right)$ to $RB=O\left(\frac{m}{\epsilon}\log\left(\frac{d}{\epsilon\delta}\right)\right)$.

It is worthy to note that most of the available communication-efficient algorithm with quantization or compression only consider communication-efficiency from devices to server. However, Algorithm 5 also improves the communication efficiency from server to devices as well because of using lower dimensional sketching size and the fact that the average of sketching has also small dimension.

We note that it is not fair to compare our algorithms with algorithms without compression. However, in the following Corollary we share an interesting observation regarding our algorithm for PL and thus strongly convex objectives in homogeneous setting.

To achieve the convergence error of $\epsilon$, we need to have $R=O\left(\kappa(\frac{\mu^{2}d}{k}+1)\log\frac{1}{\epsilon}\right)$ and $\tau=\left(\frac{(\mu^{2}d+1)}{(\frac{\mu^{2}d}{k}+1)k\epsilon}\right)$. This leads to the total communication cost per worker of

As a consequence, the total communication cost becomes:

We note that the state-of-the-art in the total communication cost is

We improve this result, in terms of dependency on $d$, to

leading to an improvement from $kd$ to $k+d$. These results are summarized in Table 1.

3 Convergence of FedSKETCHGATE in Data Heterogeneous Setting

Suppose that the conditions in Assumptions 2 and 5 hold. Given $0<m=O\left(\frac{e}{\mu^{2}}\right)\leq d$, and Consider FedSKETCHGATE in Algorithm 6 with sketch size $B=O\left(m\log\left(\frac{dR}{\delta}\right)\right)$. If the local data distributions of all users are identical (homogeneous setting), then with probability $1-\delta$ we have

For the FedSKETCHGATE-PRIVIX algorithm, by choosing stepsizes as $\eta=\frac{1}{L\gamma}\sqrt{\frac{p}{R\tau\left(\mu^{2}d\right)}}$ and $\gamma\geq p$, the sequence of iterates satisfies $\frac{1}{R}\sum_{r=0}^{R-1}\left\|\nabla f({\boldsymbol{w}}^{(r)})\right\|_{2}^{2}\leq{\epsilon}$ if we set $R=O\left(\frac{\mu^{2}d+1}{\epsilon}\right)$ and $\tau=O\left(\frac{1}{{p}\epsilon}\right)$.

For FedSKETCHGATE-HEAPRIX algorithm, by choosing stepsizes as $\eta=\frac{1}{L\gamma}\sqrt{\frac{p}{R\tau\left(\mu^{2}d\right)}}$ and $\gamma\geq p$, the sequence of iterates satisfies $\frac{1}{R}\sum_{r=0}^{R-1}\left\|\nabla f({\boldsymbol{w}}^{(r)})\right\|_{2}^{2}\leq{\epsilon}$ if we set $R=O\left(\frac{\mu^{2}d}{\epsilon}\right)$ and $\tau=O\left(\frac{1}{{p}\epsilon}\right)$.

4 Comparison with Prior Methods [30], [45] and [41]

Comparison to . We note that our convergence analysis does not rely on the bounded gradient assumption and it can be seen that we improve both the number of communication rounds $R$ and the size of vector $B$ per communication round while preserving the privacy property. Additionally, we highlight that, while provides a convergence analysis for convex objectives, our analysis holds for PL (thus strongly convex case), general convex and general nonconvex objectives.

Comparison with . Consider two versions of FetchSGD in this reference. First while in our schemes we do not to have access to the exact entries of gradients, since the approaches in is based on $top_{m}$ queries, both of the proposed algorithms (in ) require to have access to the exact value of $top_{k}$ gradients, hence they do not preserve privacy. Second, both of the convergence results in rely on the bounded gradient assumption and it is known that this assumption is not in consistent with $L$-smoothness when data distribution is heterogeneous which is the case in Federated Learning (see for more detail). However, our convergence results do not need any bounded gradient assumption. Third, Theorem 1 is based on an Assumption that Contraction Holds for the sequence of gradients encountered during the optimization which may not hold necessarily in practice, yet based on this strong assumption their total communication cost ($RB$) to achieve $\epsilon$ error is $BR=O\left(m\max(\frac{1}{\epsilon^{2}},\frac{d^{2}-dm}{m^{2}\epsilon})\log\left(\frac{d}{\delta}\max(\frac{1}{\epsilon^{2}},\frac{d^{2}-dm}{m^{2}\epsilon})\right)\right)$ (Note for the sake of comparison we let the compression ration in to be $\frac{m}{d}$). In contrast, without any extra assumptions, our results in Theorem 3 for PRIVIX and HEAPRIX are respectively $BR=O\left(\frac{m\left(\mu^{2}d+1\right)}{\epsilon}\log\left(\frac{\mu^{2}d^{2}+d}{\epsilon\delta}\log\left(\frac{1}{\epsilon}\right)\right)\right)$ and $BR=O\left(\frac{m\left(\mu^{2}d\right)}{\epsilon}\log\left(\frac{\mu^{2}d^{2}}{\epsilon\delta}\log\left(\frac{1}{\epsilon}\right)\right)\right)$ which improves total communication cost in Theorem 1 in in regimes where $\frac{1}{\epsilon}\geq d$ or $d>>m$. Theorem 2 in is based on another assumption of Sliding Window Heavy Hitters, which is similar to gradient diversity assumption in (but it is weaker assumption of contraction in Theorem 1 in ), and they showed that the total communication cost is $BR=O\left(\frac{m\max(I^{2/3},2-\alpha)}{\epsilon^{3}\alpha}\log\left(\frac{d\max(I^{2/3},2-\alpha)}{\epsilon^{3}\delta}\right)\right)$ ($I$ is constant comes from the extra assumption over the window of gradients which similar to bounded gradient diversity) which is again worse than obtained result in this paper with weaker assumptions in a regime where $\frac{I^{2/3}}{\epsilon^{2}}\geq d$. Next, unlike which only focuses on nonconvex objectives, in this work we provide the convergence analysis for PL (thus strongly convex case), general convex and general nonconvex objectives. Finally, although the algorithm in requires additional memory for the server to store the compression error correction vector, our algorithm does not need such additional storage.

Comparison with . The reference considers two-way compression from parameter server to devices and vice versa. They provide the convergence rate of $R=O\left(\frac{\omega^{\text{Up}}\omega^{\text{Down}}}{\epsilon^{2}}\right)$ for strongly-objective functions where $\omega^{\text{Up}}$ and $\omega^{\text{Down}}$ are uplink and downlink’s compression noise (specializing to our case for the sake of comparison $\omega^{\text{Up}}=\omega^{\text{Down}}=\theta\left(d\right)$) for general heterogeneous data distribution. In contrast, while as pointed out in Remark 5 that our algorithms are using bidirectional compression due to use of sketching for communication, our convergence rate for strongly-convex objective is $R=O(\kappa\mu^{2}d\log\left(\frac{1}{\epsilon}\right))$ with probability $1-\delta$.

Numerical Example

In this section, we provide empirical results on MNIST dataset to demonstrate the effectiveness of our proposed algorithms. The model we use is the LeNet-5 Convolutional Neural Network (CNN) architecture introduced in , with $60\,000$ model parameters in total.

Four methods are compared in our experiments: Federated SGD (FedSGD), SketchSGD , FedSketch-PRIVIX (FS-PRIVIX) and FedSketch-HEAPRIX (FS-HEAPRIX). We implement the algorithms by simulating the distributed and federated environment. Note that in Algorithm 5, FS-PRIVIX with global learning rate $\gamma=1$ is equivalent to the DiffSketch algorithm proposed in . In the following experiments, we set the number of workers to $50$. For federated learning algorithms, we use different number of local updates $\tau$. For SketchedSGD which is under synchronous distributed learning framework, $\tau$ is fixed and equal to $1$. For all methods, we tune the learning rates (both local, i.e. $\eta$ and global, i.e. $\gamma$, if applicable) over the log-scale and report the best results.

In each round of local update, we randomly choose half of the local devices to be active, which is the common practice in real-world applications. For the data distribution on each device, we test both homogeneous and heterogeneous setting. In the former case, each device receives uniformly drawn data samples (each class has equal probability to be selected). In the latter case, each device only receives samples from one or two classes among ten digits in the MNIST dataset. Since data is not distributed i.i.d. among local devices, training is expected to be harder in the heterogeneous case.

Homogeneous case. In Figure 1 first column, we provide the training loss and test accuracy for the four algorithms mentioned above, with $\tau=1$ (since SketchSGD requires single local update per round). We also test different sizes of sketching matrix, $(t,k)=(20,40)$ and $(50,100)$. Note that these two choices of sketch size correspond to a $75\times$ and $12\times$ compression ratio, respectively. In general, as one would expect, higher compression ratio leads to worse learning performance. In both cases, FS-HEAPRIX performs the best in terms of both training objective and test accuracy. FS-PRIVIX is better when sketch size is large (i.e. when the estimation from sketches are more accurate), while SketchSGD performs better with small sketch size.

The results for multiple local updates are given in column 2 and column 3 in Figure 1, where we set $\tau=2,5$. We see that FS-HEAPRIX is significantly better than FS-PRIVIX, either with small or large sketching matrix. In both cases, FS-HEAPRIX yields acceptable extra test error compared to FedSGD, especially when considering the high compression ratio (e.g. $75\times$). However, FS-PRIVIX performs poorly with small sketch size $(20,40)$, and even diverges with $\tau=5$. We also observe that the performances of FS-HEAPRIX improve when the number of local updates increases. That is, the proposed method is able to further reduce the communication cost by reducing the number of rounds required for communication. This is also consistent with our theoretical claims established in this paper. For $\tau=1,2,5$, we see that a sketch size of $(50,100)$ is sufficient to give similar test accuracy as the Federated SGD (FedSGD) algorithm.

Heterogeneous case. We plot similar sets of results in Figure 2 for non-i.i.d. data distribution (heterogeneous setting). This setting leads to more twists and turns in the training curves. From the first column ($\tau=1$), we see that SketchSGD performs very poorly in the heterogeneous case, while both our proposed FedSketchGATE methods, see Algorithm 6, achieve similar generalization accuracy as the Federated SGD (FedSGD) algorithm, even with fairly small sketch size (i.e. $75\times$ compression ratio). Note that, the slow convergence of federated SGD in non-i.i.d. data distribution case has also been reported in literature, e.g. . In addition, FS-HEAPRIX is again better than FS-PRIVIX in terms of both training loss and test accuracy.

Furthermore, we notice in column 2 and 3 of Figure 2 the advantage of FS-HEAPRIX over FS-PRIVIX with multiple local updates. However, empirically we see that in the heterogeneous setting, more local updates $\tau$ tend to undermine the learning performance, especially with small sketch size. Nevertheless, we see that when sketch size is large, i.e. $(50,100)$, FS-HEAPRIX can still provide comparable test accuracy as FedSGD with $\tau=5$.

Our empirical study demonstrates that our proposed FedSketch (and FedSketchGATE) frameworks are able to perform well in homogeneous (resp. heterogeneous) learning setting, with high compression rate. In particular, FedSketch methods are advantageous over prior SketchedSGD method in both cases. FS-HEAPRIX performs the best among all the tested compressed optimization algorithms, which in many cases achieves similar generalization accuracy as Federated SGD with small sketch size. In general, in any tested case, we can at least achieve $12\times$ compression ratio with very little loss in test accuracy.

Conclusion

In this paper, we introduced FedSKETCH and FedSKETCHGATE algorithms for homogeneous and heterogeneous data distribution setting respectively for Federated Learning wherein communication between server and devices is only performed using count sketch. Our algorithms, thus, provide communication-efficiency and privacy. We analyze the convergence error for nonconvex, Polyak-Łojasiewicz and general convex objective functions in the scope of Federated Optimization. We provide insightful numerical experiments showcasing the advantages of our FedSKETCH and FedSKETCHGATE methods over current federated optimization algorithm. The proposed algorithms outperform competing compression method and can achieve comparable test accuracy as Federated SGD, with high compression ratio.

References

Notation.

We will use the following fact (which is also used in ) in proving results.

Let $\{x_{i}\}_{i=1}^{p}$ denote any fixed deterministic sequence. We sample a multiset $\mathcal{P}$ (with size $K$) uniformly at random where $x_{j}$ is sampled with probability $q_{j}$ for $1\leq j\leq p$ with replacement. Let $\mathcal{P}=\{i_{1},\ldots,i_{K}\}\subset[p]$ (some $i_{j}$’s may have the same value). Then

Appendix A Results for the Homogeneous Setting

In this section, we study the convergence properties of our FedSKETCH method presented in Algorithm 5. Before stating the proofs for FedSKETCH in the homogeneous setting, we first mention the following intermediate lemmas.

Using unbiased compression and under Assumption 4, we have the following bound:

Next we show that from Assumptions 5, we have

where the last inequality is due to $\left\|\sum_{j=1}^{n}\mathbf{a}_{i}\right\|^{2}\leq n\sum_{j=1}^{n}\left\|\mathbf{a}_{i}\right\|^{2}$, which together with (10) leads to the following bound:

Under Assumption 2, and according to the FedCOM algorithm the expected inner product between stochastic gradient and full batch gradient can be bounded with:

where ➀ is due to $2\langle\mathbf{a},\mathbf{b}\rangle=\|\mathbf{a}\|^{2}+\|\mathbf{b}\|^{2}-\|\mathbf{a}-\mathbf{b}\|^{2}$, and ➁ follows from Assumption 2. ∎

The following lemma bounds the distance of local solutions from global solution at $r$th communication round.

Now we are ready to present our result for the homogeneous setting. We first state and prove the result for the general nonconvex objectives.

For FedSKETCH$(\tau,\eta,\gamma)$, for all $0\leq t\leq R\tau-1$, under Assumptions 2 to 4, if the learning rate satisfies

and all local model parameters are initialized at the same point ${\boldsymbol{w}}^{(0)}$, then the average-squared gradient after $\tau$ iterations is bounded as follows:

where $\boldsymbol{w}^{(*)}$ is the global optimal solution with function value $f(\boldsymbol{w}^{(*)})$.

Before proceeding to the proof of Theorem 5, we would like to highlight that

From the updating rule of Algorithm 5 we have

In what follows, we use the following notation to denote the stochastic gradient used to update the global model at $r$th communication round

Then using the unbiased estimation property of sketching we have:

By taking expectation on both sides of above inequality over sampling, we get:

where in ➀ we incorporate outer summation $\sum_{c=0}^{\tau-1}$, and ➁ follows from condition

Summing up for all $R$ communication rounds and rearranging the terms gives:

From above inequality, is it easy to see that in order to achieve a linear speed up, we need to have $\eta\gamma=O\left(\frac{\sqrt{k}}{\sqrt{R\tau}}\right)$. ∎

In Eq. (17) for the choice of $\eta\gamma=O\left(\frac{1}{L}\sqrt{\frac{k}{R\tau\left(\omega+1\right)}}\right)$, and $\gamma\geq k$ the convergence rate reduces to:

Note that according to Eq. (22), if we pick a fixed constant value for $\gamma$, in order to achieve an $\epsilon$-accurate solution, $R=O\left(\frac{1}{\epsilon}\right)$ communication rounds and $\tau=O\left(\frac{\omega+1}{k\epsilon}\right)$ local updates are necessary. We also highlight that Eq. (22) also allows us to choose $R=O\left(\frac{\omega+1}{\epsilon}\right)$ and $\tau=O\left(\frac{1}{k\epsilon}\right)$ to get the same convergence rate.

Condition in Eq. (16) can be rewritten as

So based on Eq. (23), if we set $\eta=O\left(\frac{1}{L\gamma}\sqrt{\frac{p}{R\tau\left(\omega+1\right)}}\right)$, it implies that:

We note that $\gamma^{2}\left(\sqrt{\left(\frac{\omega}{k}+1\right)^{2}\gamma^{2}+4}-\left(\frac{\omega}{k}+1\right)\gamma\right)^{2}=\Theta(1)\leq 5$ therefore even for $\gamma\geq m$ we need to have

Therefore, for the choice of $\tau=O\left(\frac{\omega+1}{k\epsilon}\right)$, due to condition in Eq. (25), we need to have $R=O\left(\frac{1}{\epsilon}\right)$. Similarly, we can have $R=O\left(\frac{\omega+1}{\epsilon}\right)$ and $\tau=O\left(\frac{1}{k\epsilon}\right)$.

By letting $\gamma=1$, $\omega=0$ and $k=p$ the convergence rate in Eq. (17) reduces to

which matches the rate obtained in . In this case the communication complexity and the number of local updates become

This simply implies that in this special case the convergence rate of our algorithm reduces to the rate obtained in , which indicates the tightness of our analysis.

A.2 Main result for the PL/Strongly convex setting

We now turn to stating the convergence rate for the homogeneous setting under PL condition which naturally leads to the same rate for strongly convex functions.

For FedSKETCH$(\tau,\eta,\gamma)$, for all $0\leq t\leq R\tau-1$, under Assumptions 2 to 4 and 3,if the learning rate satisfies

and if the all the models are initialized with $\boldsymbol{w}^{(0)}$ we obtain:

By setting $\Delta=1-\eta\mu\gamma{\tau}$ we obtain the following bound:

If we let $\eta\gamma\mu\tau\leq\frac{1}{2}$, $\eta=\frac{1}{2L\left(\frac{\omega}{k}+1\right)\tau\gamma}$ and $\kappa=\frac{L}{\mu}$ the convergence error in Theorem 6, with $\gamma\geq k$ results in:

which indicates that to achieve an error of $\epsilon$, we need to have $R=O\left(\left(\frac{\omega}{k}+1\right)\kappa\log\left(\frac{1}{\epsilon}\right)\right)$ and $\tau=\frac{\left(\omega+1\right)}{k\left(\frac{\omega}{k}+1\right)\epsilon}$. Additionally, we note that if $\gamma\rightarrow\infty$, yet $R=O\left(\left(\frac{\omega}{k}+1\right)\kappa\log\left(\frac{1}{\epsilon}\right)\right)$ and $\tau=\frac{\left(\omega+1\right)}{k\left(\frac{\omega}{k}+1\right)\epsilon}$ will be necessary.

A.3 Main result for the general convex setting

and if the all the models initiate with $\boldsymbol{w}^{(0)}$, with $\phi=\frac{1}{\sqrt{k\tau}}$ and $\eta=\frac{1}{2L\gamma\tau\left(1+\frac{\omega}{k}\right)}$ we obtain:

We note that above theorem implies that to achieve a convergence error of $\epsilon$ we need to have $R=O\left(L\left(1+\frac{\omega}{k}\right)\frac{1}{\epsilon}\log\left(\frac{1}{\epsilon}\right)\right)$ and $\tau=O\left(\frac{\left(\omega+1\right)^{2}}{k\left(\frac{\omega}{k}+1\right)^{2}\epsilon}\right)$.

Next rearranging Eq. (30) and replacing $\mu$ with $\phi$ leads to the following error bound:

Next, if we set $\phi=\frac{1}{\sqrt{k\tau}}$ and $\eta=\frac{1}{2\left(1+\frac{\omega}{k}\right)L\gamma\tau}$, we obtain that

Appendix B Proof of Main Theorems

The proof of Theorem 2 follows directly from the results in . For the sake of the completeness we review an assumptions from this reference for the quantization with their notation.

Consider FedCOM in . Suppose that the conditions in Assumptions 2, 4 and 6 hold. If the local data distributions of all users are identical (homogeneous setting), then we have

Nonconvex: By choosing stepsizes as $\eta=\frac{1}{L\gamma}\sqrt{\frac{p}{R\tau\left(\frac{\omega}{p}+1\right)}}$ and $\gamma\geq p$, the sequence of iterates satisfies $\frac{1}{R}\sum_{r=0}^{R-1}\left\|\nabla f({\boldsymbol{w}}^{(r)})\right\|_{2}^{2}\leq{\epsilon}$ if we set $R=O\left(\frac{1}{\epsilon}\right)$ and $\tau=O\left(\frac{\frac{\omega}{p}+1}{{p}\epsilon}\right)$.

Since the sketching PRIVIX and HEAPRIX, satisfy Assumption 6 with $\omega=\mu^{2}d$ and $\omega=\mu^{2}d-1$ respectively with probability $1-\delta$. Therefore, all the results in Theorem 2, conclude from Theorem 8 with probability $1-\delta$ and plugging $\omega=\mu^{2}d$ and $\omega=\mu^{2}d-1$ respectively into the corresponding convergence bounds. ∎

B.2 Proof of Theorem 3

For the heterogeneous setting, the results in requires the following extra assumption that naturally holds for the sketching:

We note that since sketching is a linear compressor, in the case of our algorithms for heterogeneous setting we have $G_{q}=0$.

Next, we restate the Theorem in here as follows:

Consider FedCOMGATE in . If Assumptions 2, 5, 6 and 7 hold, then even for the case the local data distribution of users are different (heterogeneous setting) we have

Nonconvex: By choosing stepsizes as $\eta=\frac{1}{L\gamma}\sqrt{\frac{p}{R\tau\left(\omega+1\right)}}$ and $\gamma\geq p$, we obtain that the iterates satisfy $\frac{1}{R}\sum_{r=0}^{R-1}\left\|\nabla f({\boldsymbol{w}}^{(r)})\right\|_{2}^{2}\leq\epsilon$ if we set $R=O\left(\frac{\omega+1}{\epsilon}\right)$ and $\tau=O\left(\frac{1}{p\epsilon}\right)$.

Since the sketching methods PRIVIX and HEAPRIX, satisfy the Assumption 6 with $\omega=\mu^{2}d$ and $\omega=\mu^{2}d-1$ respectively with probablity $1-\delta$, we conclude the proofs of Theorem 3 using Theorem 9 with probability $1-\delta$ and plugging $\omega=\mu^{2}d$ and $\omega=\mu^{2}d-1$ respectively into the convergence bounds. ∎