FedLite: A Scalable Approach for Federated Learning on Resource-constrained Clients

One major limitation of SL is that the underlying model partitioning leads to an increased communication cost for the resulting framework. In order to train the split neural network, the activations and gradients at the layer where the model is split (referred to as cut layer) need to be communicated between the server and clients at each iteration. The additional message size is in proportion to the mini-batch size as well as the cut layer size. As a result, the communication cost for the model training can become prohibitive whenever the mini-batch size and the activation size of the cut layer are large (e.g. in one of our experiments, the additional message size can be $10$ times larger than that of the client-side model).

In this paper, we aim to make the split neural network training communication-efficient and enable its widespread adoption for FL in resource-constrained settings. Our proposed solution is based on the critical observation that, given a mini-batch of data, the client does not need to communicate per-example activation vectors if the activation vectors (at the cut layer) for different examples in the mini-batch exhibit enough similarity. Thus, we propose a training framework that performs clustering of the activation vectors and only communicates the cluster centroids to the server. Interestingly, this is equivalent to adding a vector quantization layer in the middle of the split neural network (see Figure 1 for an overview of our proposed method). Our main contributions are as follows.

• •

  We propose a communication-efficient splitting-based FL approach for resource-constrained clients (cf. Section 4), which improves previous SL solutions. The approach employs a novel clustering scheme that leverages product quantization to effectively compress the activations communicated between the clients and server.

• •

  After applying the activation quantization, the clients can only receive possibly noisy gradients from the server. The inaccurate client-side model updates lead to significant accuracy drops in our experiments. In order to mitigate this problem, we propose a gradient correction scheme, which plays a critical role in achieving a high compression ratio with minimal accuracy loss.

• •

  We empirically evaluate the performance of our approach on three standard FL datasets (cf. Section 5). We show that our approach allows for up to $490\times$ communication reduction without significant accuracy loss.

• •

  In addition, we present a convergence analysis for the proposed method (cf. Section 4.3), which reveals a trade-off between communication reduction and convergence speed. The analysis further helps explain why the proposed gradient correction technique is beneficial.

It is also worth mentioning that the proposed method does not expose any additional client-side information to the server than vanilla split neural network training approach. Similar to SplitFed (Thapa et al., 2020), our method can also leverage existing privacy preserving mechanisms such as differential privacy (Wei et al., 2020) or instance-hiding schemes (Huang et al., 2020) to provide formal privacy guarantees, which are out of the scope of this paper. Besides, our proposed method has potential applications beyond FL with resource-constrained clients as it can be applied in any learning framework that can benefit from a vector quantization layer. For instance, it can be used to reduce the communication overhead in two-party vertical FL (Romanini et al., 2021), where the model is naturally split across two institutions and the data labels are generated on the server.

Suppose we have a collection of $M$ clients ${\mathcal{I}}=\{1,2,\dots,M\}$. Each client $i\in{\mathcal{I}}$ has a local dataset ${\mathbb{D}}_{i}$ and a corresponding empirical loss function $F_{i}({\bm{w}})=\sum_{\xi\in{\mathbb{D}}_{i}}f({\bm{w}};\xi)/n_{i}$, where ${\bm{w}}$ denotes the model parameters, and $n_{i}=|{\mathbb{D}}_{i}|$ denotes the number of samples of the local dataset. The goal of FL is to find a shared model ${\bm{w}}$ that can minimize the averaged loss over all clients, defined as follows:

where $p_{i}=n_{i}/\sum_{i=1}^{M}n_{i}$ is the relative weight of local loss $F_{i}$. Motivated by the data privacy concerns, under a FL framework, the clients perform local training and only communicate the resulting models to a coordinating server as opposed to sharing their raw local data with the server (McMahan et al., 2017). In current FL algorithms, the model size is limited by the resource constraints of clients. Consequently, these algorithms are not feasible when dealing with large machine learning models that cannot fit in clients memory or require large compute.

To deploy large machine learning models on resource-constrained clients, a few recent works proposed methods to reduce the effective model size on clients. For example, Diao et al. (2020) varied the layer width on clients depending on their computational capacities; and Horvath et al. (2021) proposed ordered dropout to randomly removing neurons in the neural network. These methods are effective in reducing the width of the intermediate layers; however, they still place the full classification layer (which is in proportion to the size of the output space) on each client to compute the the client update. This may not be feasible for many clients as the parameters in the classification layer of a model can dominate the total number of model parameters (Krizhevsky, 2014).

There are a large amount of literature studying gradient compression in distributed training, such as QSGD (Alistarh et al., 2017), SignSGD (Bernstein et al., 2018), Atomo (Wang et al., 2018), PowerSGD (Vogels et al., 2019), etc. Although these works share the same goal as our work (i.e., reducing communciation by compressing message size), we focus on orthogonal aspects. In particular, all these previous gradient compression methods only compress one gradient vector using the inherent properties of the stochastic gradient (e.g., sparsity or low rankness). However, our method considers compressing a batch of activation vectors and utilizes the similarity among the activations. This is a new dimension that does not exist in distributed training scenarios and has not been studied in gradient compression literature.

Split learning (SL) is another way of minimizing 1 without explicitly sharing local data on clients (Vepakomma et al., 2018). In particular, SL splits the neural network model into two parts on a layer basis. The first few layers (called client-side model parameterized by ${\bm{w}}_{\text{c}}$) are shared across clients and the server, while the remaining layers (called server-side model parameterized by ${\bm{w}}_{\text{s}}$) are only stored and trained on the server. Under this splitting, the original loss function for a data sample $\xi$ can be re-written as follows:

where $u$ is the client-side function mapping the input data to the activation space and $h$ is the server-side function mapping the activation to a scalar loss value. Training both the client- and server-side models requires communicating client activations (i.e., the output of the client-side model, also called smashed data in some literature) between the clients and the server. We further elaborate on a concrete SL algorithm as a baseline in Section 3.

Recently, He et al. (2020); Han et al. (2021) proposed to add a classification layer on the client-side model in SL so that each client can locally update its parameters. This can effectively reduce the communication frequency between the clients and the server. However, this kind of method is not suitable for the settings where the classification layer dominates the entire model size such that the clients may not have enough memory space to store the entire classification layer.

Product quantization (PQ) is a compression method and has been widely used in approximate nearest neighbor search (Jegou et al., 2010; Ge et al., 2013b). Given a batch of vectors, PQ divides each vector into multiple chunks and performs K-means clustering separately on each chunk. Now, the resulting cluster centroids construct the codebook and the closest codeword will be used to represent each vector. Instead of directly computing the distance between two high-dimensional vectors, computing the distance between two codewords is orders of magnitude more efficient and faster. Note that the previous works used PQ to compress either learned features after training (see, e.g., Ge et al., 2013a) or the parameters of a neural network (see, e.g., Chen et al., 2020). In this paper, we show that PQ is also effective in compressing the output of intermediate layers and present a simple technique to enable backpropagation through the corresponding quantization layer.

Each iteration of SplitFed contains four steps detailed as follows. Wherever it simplifies the presentation, we assume that the mini-batch size $B$ on each client is $1$ but the derivations hold true for arbitrary batch sizes.

• 1.

  Client Forward Pass: Let ${\mathcal{S}}$ be a randomly selected subset of clients. For $i\in{\mathcal{S}}$, compute the output of the client-side model ${\bm{z}}_{i}=u({\bm{w}}_{\text{c}};\xi)$, where $\xi\in{\mathbb{D}}_{i}$ is a randomly chosen training sample, and ${\bm{z}}_{i}\in\mathbb{R}^{d}$ denotes the activations of the last layer of the client-side model. Then, each selected client sends ${\bm{z}}_{i}$ (together with its corresponding label if necessary¹¹1In vertical FL applications (Romanini et al., 2021), the labels do not need to be transferred.) to the server.

• 2.

  Server Update: The server treats all the activations (also referred to as smashed data) $\{{\bm{z}}_{i}\}_{i\in{\mathcal{S}}}$ as inputs to perform one step of gradient descent on the server-side model: $\Delta{\bm{w}}_{\text{s}}=\eta_{\text{s}}\sum_{i\in{\mathcal{S}}}p_{i}\nabla_{{\bm{w}}_{\text{s}}}h({\bm{w}}_{\text{s}};{\bm{z}}_{i})/\sum_{i\in{\mathcal{S}}}p_{i}$, where $\eta_{\text{s}}$ denotes the server learning rate.²²2The server learning rate here has a different definition from previous FL literature, e.g., Reddi et al. (2020). In addition, the server also computes the gradient with respect to the activations $\nabla_{{\bm{z}}_{i}}h({\bm{w}}_{\text{s}};{\bm{z}}_{i})$ and sends it back to the corresponding client $i$.

• 3.

  Client Backward Pass: Each client computes the gradient with respect to the client-side model using the chain rule: ${\bm{g}}_{i}\triangleq\nabla_{{\bm{w}}_{\text{c}}}f=\nabla_{{\bm{z}}_{i}}h({\bm{w}}_{\text{s}};{\bm{z}}_{i})\nabla_{{\bm{w}}_{\text{c}}}u({\bm{w}}_{\text{c}};\xi)$.

• 4.

  Client-side Model Synchronization: At last, the client-side model is updated by synchronizing client gradients $\Delta{\bm{w}}_{\text{c}}=\eta_{\text{c}}\sum_{i\in{\mathcal{S}}}p_{i}{\bm{g}}_{i}/\sum_{i\in{\mathcal{S}}}p_{i}$, where $\eta_{\text{c}}$ denotes the learning rate for the client-side model.

One can easily validate that, with the above four steps, SplitFed is equivalent to mini-batch stochastic gradient descent (SGD) with a total batch size of $B|{\mathcal{S}}|$. No matter how the network is split, the performance of the algorithm remains unchanged and has the same iteration complexity as mini-batch SGD.

Besides, note that, a selected client needs to upload both the activations during “Client Forward Pass” and the gradients of the client-side model during “Client-side Model Synchronization”. Thus, the communication size per client per iteration can be written as $|{\bm{w}}_{\text{c}}|+Bd$. Since the up-link communication bandwidth is quite limited for clients in federated learning, when the mini-batch size $B$ or the dimension of the activation $d$ gets large, the up-link communication may become the main bottleneck that limits the scalability of SplitFed. It is critical to compress the uploaded message from clients to the server (Reisizadeh et al., 2019; Li et al., 2019).

In Table 1, we further compare the difference between FedAvg and SplitFed. One can observe that SplitFed always requires less computation and memory on clients, and could potentially communicate less than FedAvg. In terms of convergence rate, it is still under discussion which algorithm is better Woodworth et al. (2020). Note that the goal of this paper is to show that we can effectively reduce the additional communication in SplitFed rather than showing SplitFed is better than FedAvg.

The key idea of our method is to compress the redundant information in the client activations for each mini-batch. For instance, if we assume that all the activation vectors within a mini-batch are nearly identical to each other, then instead of sending all $B$ activation vectors, the client needs to send only one representative vector to the server, thereby, reducing the communication cost by a factor of $B$.

Building on the above observation, we let each client input their activations into a clustering algorithm to get $L$ centroids. Then, each activation vector is represented by the index of its closest centroid. Instead of sending a mini-batch of $B$ vectors, only $L$ centroids and the cluster assignments are transmitted to the server. This procedure is equivalent to vector quantization. Here, a simple choice of the clustering algorithm is K-means (Krishna & Murty, 1999). However, as we observe in Figure 3, vanilla K-means clustering leads to a high quantization error with very limited compression ratio. It is non-trivial to design a proper clustering algorithm that can minimize the communication between clients and server, while maintaining a low quantization error. In Section 4.1, we design a novel variant of product quantizer that can flexibly operate this trade-off.

While the above approach is appealing, it also introduces new challenges. First, note that, in the backward pass, due to the additional quantization layer, the server no longer knows the true client activations; thus, constraining the server to compute the gradients with respect to the (possibly noisy) quantized activations. Furthermore, in order to update the client-side model parameters, the clients need to receive the gradients with respect to the original activations, which is not available any more. In Section 4.2, we introduce a gradient correction technique to mitigate this gradient mismatch, which is critical for our method to achieve a high compression ratio with minimal accuracy loss.

In the following sections, we present the key components of our proposed method and highlight how they address the aforementioned challenges. In addition, a convergence analysis is provided in Section 4.3. Unless otherwise stated, for the ease of exposition, we will focus on the operations on a specific client $i$ and omit the client index. However, it is important to keep in mind that the client-side operations happen in parallel on all selected clients.

As discussed in Section 4.1, there is a trade-off between the final performance and compression ratio. Although the additional quantization layer helps reduce the communication cost, it also causes performance drop. While utilizing the quantization layer, it is important to understand the range of the compression ratio that does not degrade the final performance too much. To this end, we vary the number of clusters $L$ and number of subvectors $q$ in our proposed method and report the resulting accuracy vs. compression trade-off in Figure 4. Note that we define the compression ratio to be the ratio between the original activation size and compressed message (codebook + codewords) size.

One can observe that the proposed method can achieve at least $10\times$ compression ratio with almost no accuracy loss. Furthermore, if we allow for $5\%$ relative loss compared to SplitFed, then our method can achieve a $490\times$ compression ratio on FEMNIST. This suggests that $99.8\%$ information is redundant and can be dropped during the up-link communication. Surprisingly, on SO Tag, the Recall@5 even improves when adding the quantization layer. We conjecture that compressing the outputs of a intermediate layer might have similar effects to dropout.

Here, we provide a concrete example showing how the proposed method improves the communication and computation efficiencies over previous works. On FEMNIST, setting $q=1152$ and $L=2$ amounts to $490\times$ compression ratio, which amounts to a $490\times$ reduction in the additional communication introduced by the network splitting. Compared to SplitFed, the overall up-link commutation cost (including the client-side gradients synchronization) is about $10\times$ smaller. Compared to FedAvg, the up-link communication cost per round is reduced by $62\times$, with $64\times$ less trainable parameters on the clients. We further compare the training curves with respect to total communication costs of FedLite against SplitFed and FedAvg in Figure 6 in Appendix. Besides, the overall wall-clock time saving may depend on the characteristics of the training system. One can get an estimate using the analytical model in (Agarwal et al., 2021).

From Figure 4, it is easy to observe that the gradient correction technique (i.e., $\lambda>0$) is crucial for improving performance. Without correction, the algorithm can even diverge in the high compression ratio regime. While $\lambda$ is separately tuned for each point (i.e., each $(q,L)$ pair) in Figure 4, we found that fixing $\lambda$ for all $(q,L)$ pairs still leads to significant improvements. In Figure 5, we report the performance on FEMNIST for various choices of $q$ and $L$. In particular, with $q=288$, the accuracy improvement ranges from $3-72\%$. Furthermore, we found that a small $\lambda$ ranging from $10^{-5}$ to $10^{-3}$ works well on all training tasks in our experiments. Setting a larger $\lambda$ leads to near-zero quantization error. But the model may have poor performance, as it tends to output the same activation for all inputs and becomes meaningless.

In Section 4.1, we already show that stacking subvectors into groups (i.e., setting $R<q$) is critical to reach a high compression ratio. But how does it affect the model performance? To answer the question, we run the proposed method with $R=q>1$ (vanilla PQ), and report the results in Figure 5(c). Observe that the proposed method significantly improves the compression ratio with minimal loss of accuracy. In addition, any other quantizers that compress each individual vector but do not utilize the similarities among vectors (such as stochastic quantization Suresh et al. (2017)) can be combined with our method to further improve the compression ratio, as they compress a different dimension from ours.